diff --git a/stoat_retarded_ai_tree/python/README.md b/stoat_retarded_ai_tree/python/README.md index b2eff00..c0db057 100644 --- a/stoat_retarded_ai_tree/python/README.md +++ b/stoat_retarded_ai_tree/python/README.md @@ -1,3 +1,2 @@ -replicate requests Pillow diff --git a/stoat_retarded_ai_tree/python/main.py b/stoat_retarded_ai_tree/python/main.py index 1f8994d..8ac7d5c 100644 --- a/stoat_retarded_ai_tree/python/main.py +++ b/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()