aljaz/Auto-GPT
1
0
Fork 0
mirror of https://github.com/aljazceru/Auto-GPT.git synced 2025-12-24 01:14:22 +01:00
Code Issues Packages Projects Releases Wiki Activity

Huggingface retry generate_image with delay (#2745)

Browse Source 
Co-authored-by: Media <12145726+rihp@users.noreply.github.com>
Co-authored-by: Nicholas Tindle <nick@ntindle.com>
Co-authored-by: Nicholas Tindle <nicktindle@outlook.com>
Co-authored-by: k-boikov <64261260+k-boikov@users.noreply.github.com>
Co-authored-by: merwanehamadi <merwanehamadi@gmail.com>
Co-authored-by; lc0rp
This commit is contained in:
Kory Becker
2023-05-16 11:02:55 -04:00
committed by GitHub
parent c1cd54d1ea
commit f424fac1d8
2 changed files with 143 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
autogpt/commands/image_gen.py
View File

@@ -1,5 +1,7 @@
""" Image Generation Module for AutoGPT."""
import io
import json
import time
import uuid
from base64 import b64decode
@@ -61,20 +63,42 @@ def generate_image_with_hf(prompt: str, filename: str) -> str:
"X-Use-Cache": "false",
}
response = requests.post(
API_URL,
headers=headers,
json={
"inputs": prompt,
},
)
retry_count = 0
while retry_count < 10:
response = requests.post(
API_URL,
headers=headers,
json={
"inputs": prompt,
},
)
image = Image.open(io.BytesIO(response.content))
logger.info(f"Image Generated for prompt:{prompt}")
if response.ok:
try:
image = Image.open(io.BytesIO(response.content))
logger.info(f"Image Generated for prompt:{prompt}")
image.save(filename)
return f"Saved to disk:{filename}"
except Exception as e:
logger.error(e)
break
else:
try:
error = json.loads(response.text)
if "estimated_time" in error:
delay = error["estimated_time"]
logger.debug(response.text)
logger.info("Retrying in", delay)
time.sleep(delay)
else:
break
except Exception as e:
logger.error(e)
break
image.save(filename)
retry_count += 1
return f"Saved to disk:{filename}"
return f"Error creating image."
def generate_image_with_dalle(prompt: str, filename: str, size: int) -> str: