god bless that every first http request send fails and needs a retry where the second then works
This commit is contained in:
@@ -1,3 +1,2 @@
|
|||||||
replicate
|
|
||||||
requests
|
requests
|
||||||
Pillow
|
Pillow
|
||||||
|
|||||||
@@ -6,7 +6,6 @@ import time
|
|||||||
import base64
|
import base64
|
||||||
import re, json
|
import re, json
|
||||||
import os
|
import os
|
||||||
import replicate as replicate_client
|
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
import io
|
import io
|
||||||
|
|
||||||
@@ -182,12 +181,40 @@ def send_claude_message(reacted_emotes, metrics, custom_emotes_json, attachment_
|
|||||||
return parsed
|
return parsed
|
||||||
|
|
||||||
def replicate_fetch_image(image_prompt):
|
def replicate_fetch_image(image_prompt):
|
||||||
os.environ["REPLICATE_API_TOKEN"] = replicate_api_key
|
headers = {
|
||||||
output = replicate_client.run(
|
"Authorization": f"Bearer {replicate_api_key}",
|
||||||
"black-forest-labs/flux-schnell",
|
"Content-Type": "application/json"
|
||||||
input={"prompt": image_prompt}
|
}
|
||||||
|
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():
|
def main():
|
||||||
custom_emotes_json = get_custom_emotes()
|
custom_emotes_json = get_custom_emotes()
|
||||||
|
|||||||
Reference in New Issue
Block a user