diff --git a/ai_agent.py b/ai_agent.py index 93f3992..8259faa 100644 --- a/ai_agent.py +++ b/ai_agent.py @@ -92,6 +92,13 @@ def _get_resolution_for_aspect_ratio(aspect_ratio: str) -> Tuple[int, int]: return aspect_ratio_resolution_map.get(aspect_ratio, (1280, 1024)) +def _convert_image_to_jpeg(image: bytes) -> bytes: + img = Image.open(BytesIO(image)).convert("RGB") + output = BytesIO() + img.save(output, format='JPEG', quality=95, optimize=True) + return output.getvalue() + + class AiAgent: def __init__(self, openrouter_token: str, openrouter_model: str, @@ -307,14 +314,10 @@ class AiAgent: async with aiohttp.ClientSession() as session: async with session.get(image_url) as response: if response.status == 200: - image_bytes = await response.read() + image = await response.read() if not image_url.endswith(".jpg"): - image = Image.open(BytesIO(image_bytes)).convert("RGB") - output = BytesIO() - image.save(output, format="JPEG", quality=80, optimize=True) - image_bytes = output.getvalue() - - return Ok(image_bytes) + image = _convert_image_to_jpeg(image) + return Ok(image) else: raise RuntimeError(f"Не удалось загрузить изображение ({response.status}).") @@ -362,13 +365,8 @@ class AiAgent: try: outputs = await self.replicate_client.async_run(REPLICATE_MODEL, input=arguments) - image_bytes = await outputs[0].aread() - - image = Image.open(BytesIO(image_bytes)).convert("RGB") - output = BytesIO() - image.save(output, format="JPEG", quality=80, optimize=True) - image_bytes = output.getvalue() - return Ok(image_bytes) + image = _convert_image_to_jpeg(await outputs[0].aread()) + return Ok(image) except Exception as e: print(f"Ошибка генерации изображения: {e}") return Err(str(e))