| import aiohttp | |
| import asyncio | |
| import json | |
| from typing import Dict, Any | |
| class ModalImageGenerator: | |
| def __init__(self, session: aiohttp.ClientSession): | |
| self.session = session | |
| self.base_url = ( | |
| "https://allanwatts705--kolors-app-model-web-inference.modal.run/" | |
| ) | |
| async def generate_image( | |
| self, | |
| prompt: str, | |
| reference_image_url: str = "https://image.lexica.art/full_webp/d6ddd5c5-060c-4aba-b9d0-cf0e02dc65bd", | |
| ip_adapter_weight: float = 0.4, | |
| ) -> Dict[str, Any]: | |
| params = { | |
| "ip_adapter_weight": ip_adapter_weight, | |
| "reference_image_url": reference_image_url, | |
| "prompt": prompt, | |
| } | |
| async with self.session.get(self.base_url, params=params) as response: | |
| if response.status == 200: | |
| return await response.json() | |
| else: | |
| raise Exception(f"Failed to generate image: {response.status}") | |
| async def wait_for_image( | |
| self, result: Dict[str, Any], max_attempts: int = 60, delay: int = 5 | |
| ) -> Dict[str, Any]: | |
| attempts = 0 | |
| while attempts < max_attempts: | |
| if result.get("image", {}).get("url"): | |
| return result | |
| await asyncio.sleep(delay) | |
| attempts += 1 | |
| raise Exception("Image generation timed out") | |
| async def generate_and_wait( | |
| self, prompt: str, reference_image_url: str, ip_adapter_weight: float = 0.4 | |
| ) -> str: | |
| result = await self.generate_image( | |
| prompt, reference_image_url, ip_adapter_weight | |
| ) | |
| final_result = await self.wait_for_image(result) | |
| return final_result["image"]["url"] | |