god bless that every first http request send fails and needs a retry where the second then works

stoat_retarded_ai_tree/python/README.md

@@ -1,3 +1,2 @@
-replicate
 requests
 Pillow

stoat_retarded_ai_tree/python/main.py

@@ -6,7 +6,6 @@ import time
 import base64
 import re, json
 import os
-import replicate as replicate_client
 from PIL import Image
 import io
 
@@ -182,12 +181,40 @@ def send_claude_message(reacted_emotes, metrics, custom_emotes_json, attachment_
     return parsed 
 
 def replicate_fetch_image(image_prompt):
-    os.environ["REPLICATE_API_TOKEN"] = replicate_api_key
-    output = replicate_client.run(
-        "black-forest-labs/flux-schnell",
-        input={"prompt": image_prompt}
+    headers = {
+        "Authorization": f"Bearer {replicate_api_key}",
+        "Content-Type": "application/json"
+    }
+    resp = requests.post(
+        "https://api.replicate.com/v1/models/black-forest-labs/flux-schnell/predictions",
+        headers=headers,
+        json={"input": {"prompt": image_prompt}}
     )
-    return output[0].read()
+    resp.raise_for_status()
+    prediction = resp.json()
+    get_url = prediction["urls"]["get"]
+
+    max_retries = 5
+    retry_count = 0
+    while prediction["status"] not in ("succeeded", "failed", "canceled"):
+        time.sleep(1)
+        r = requests.get(get_url, headers=headers)
+        if r.status_code >= 500:
+            retry_count += 1
+            print(f"Replicate poll got {r.status_code} (retry {retry_count}/{max_retries}): {r.text}")
+            if retry_count > max_retries:
+                r.raise_for_status()
+            continue
+        r.raise_for_status()
+        prediction = r.json()
+        retry_count = 0
+
+    if prediction["status"] != "succeeded":
+        raise RuntimeError(f"Replicate prediction failed: {prediction['status']} - {prediction.get('error')}")
+
+    output = prediction["output"]
+    image_url = output[0] if isinstance(output, list) else output
+    return requests.get(image_url).content
 
 def main():
     custom_emotes_json = get_custom_emotes()