| 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"] |
|
|
