67 lines
4.2 KiB
Python
67 lines
4.2 KiB
Python
import logging
|
|
import asyncio
|
|
import aiohttp
|
|
from .ai_responder import exponential_backoff, AIResponderBase
|
|
from io import BytesIO
|
|
|
|
|
|
class LeonardoAIDrawMixIn(AIResponderBase):
|
|
async def draw_leonardo(self, description: str) -> BytesIO:
|
|
error_backoff = exponential_backoff(max_attempts=12)
|
|
generation_id = None
|
|
image_url = None
|
|
image_bytes = None
|
|
while True:
|
|
error_sleep = next(error_backoff)
|
|
try:
|
|
async with aiohttp.ClientSession() as session:
|
|
if generation_id is None:
|
|
async with session.post("https://cloud.leonardo.ai/api/rest/v1/generations",
|
|
json={"prompt": description,
|
|
"modelId": "6bef9f1b-29cb-40c7-b9df-32b51c1f67d3",
|
|
"num_images": 1,
|
|
"sd_version": "v2",
|
|
"promptMagic": True,
|
|
"unzoomAmount": 1,
|
|
"width": 512,
|
|
"height": 512},
|
|
headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
|
|
"Accept": "application/json",
|
|
"Content-Type": "application/json"},
|
|
) as response:
|
|
response = await response.json()
|
|
if "sdGenerationJob" not in response:
|
|
logging.warning(f"No 'sdGenerationJob' found in response, sleep for {error_sleep}s: {repr(response)}")
|
|
await asyncio.sleep(error_sleep)
|
|
continue
|
|
generation_id = response["sdGenerationJob"]["generationId"]
|
|
if image_url is None:
|
|
async with session.get(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
|
|
headers={"Authorization": f"Bearer {self.config['leonardo-token']}",
|
|
"Accept": "application/json"},
|
|
) as response:
|
|
response = await response.json()
|
|
if "generations_by_pk" not in response:
|
|
logging.warning(f"Unexpected response, sleep for {error_sleep}s: {repr(response)}")
|
|
await asyncio.sleep(error_sleep)
|
|
continue
|
|
if len(response["generations_by_pk"]["generated_images"]) == 0:
|
|
await asyncio.sleep(error_sleep)
|
|
continue
|
|
image_url = response["generations_by_pk"]["generated_images"][0]["url"]
|
|
if image_bytes is None:
|
|
async with session.get(image_url) as response:
|
|
image_bytes = BytesIO(await response.read())
|
|
async with session.delete(f"https://cloud.leonardo.ai/api/rest/v1/generations/{generation_id}",
|
|
headers={"Authorization": f"Bearer {self.config['leonardo-token']}"},
|
|
) as response:
|
|
await response.json()
|
|
logging.info(f'Drawed a picture with leonardo AI on this description: {repr(description)}')
|
|
return image_bytes
|
|
except Exception as err:
|
|
logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}\n{repr(err)}")
|
|
else:
|
|
logging.warning(f"Failed to generate image, sleep for {error_sleep}s: {repr(description)}")
|
|
await asyncio.sleep(error_sleep)
|
|
raise RuntimeError(f"Failed to generate image {repr(description)}")
|