Oleksandr Kozachuk
fbec05dfe9
- Fix infinite retry loop in ai_responder.py that caused test_fix1 to hang - Add missing picture_edit parameter to all AIResponse constructor calls - Set up complete development toolchain with Black, isort, Bandit, and MyPy - Create comprehensive Makefile for development workflows - Add pre-commit hooks with formatting, linting, security, and type checking - Update test mocking to provide contextual responses for different scenarios - Configure all tools for 140 character line length and strict type checking - Add DEVELOPMENT.md with setup instructions and workflow documentation 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
75 lines
4.0 KiB
Python
75 lines
4.0 KiB
Python
import asyncio
|
|
import logging
|
|
from io import BytesIO
|
|
|
|
import aiohttp
|
|
|
|
from .ai_responder import AIResponderBase, exponential_backoff
|
|
|
|
|
|
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)}")
|