amazing stuff

App/Generate/database/Modal.py

@@ -0,0 +1,48 @@
+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, 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"]

App/Generate/database/Model.py

@@ -14,6 +14,7 @@ from typing import List
 from pydantic import BaseModel
 import tempfile
 import json
+from .Modal import ModalImageGenerator
 
 SUPABASE = os.environ.get("SUPABASE", "RANDOM_STRING")
 database_url = SUPABASE
@@ -259,17 +260,18 @@ class Scene(orm.Model):
     async def generate_images(self):
         self.images = []
         async with aiohttp.ClientSession() as session:
-            image_generator = AsyncImageGenerator(session)
-            for payload in self.image_prompts:
-                while True:
-                    result = await image_generator.generate_image(payload)
-                    status = await image_generator.fetch_image_status(result["id"])
-
-                    if type(status["output"]) == list:
-                        self.images.extend(status["output"])
-                    else:
-                        self.images.append(status["output"])
-                        break
+            image_generator = ModalImageGenerator(session)
+            for prompt in self.image_prompts:
+                # You may want to define a method to get a suitable reference image URL
+                reference_image_url = "https://image.lexica.art/full_webp/11caada0-c3c3-4b70-b7a5-d3172c07fc14"
+                try:
+                    image_url = await image_generator.generate_and_wait(
+                        prompt, reference_image_url
+                    )
+                    self.images.append(image_url)
+                except Exception as e:
+                    print(f"Failed to generate image for prompt '{prompt}': {str(e)}")
+                await asyncio.sleep(1)  # Add a small delay between requests
 
 
 class Transition(orm.Model):