File size: 1,712 Bytes
31d928d fe3d026 31d928d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
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"]
|