Update src/app_pages/image_generation.py

src/app_pages/image_generation.py

@@ -76,8 +76,18 @@ MODEL_REGISTRY: Dict[str, Dict[str, Any]] = {
         ],
         "param_name": "aspect_ratio",
     },
-}
+       "z-image-turbo": {
+        "id": "prunaai/z-image-turbo",
+        "aspect_ratios": [
+            "match_input_image", "1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9", "21:9"
+        ],
+            "z-image-turbo": {
+                "id": "prunaai/z-image-turbo",
+                "type": "custom",
+            },
 
+    },
+}
 
 def retrieve_from_store(vector_store_id: str, user_query: str, category: str) -> str:
     client = get_openai_client()
@@ -248,50 +258,45 @@ def generate_ad_images_replicate(
     prompt: str,
     replicate_model_key: str,
     aspect_ratio: str,
+    z_width: int | None = None,
+    z_height: int | None = None,
+    z_output_quality: int | None = None,
 ) -> List[Image.Image]:
-    if not REPLICATE_API_TOKEN:
-        raise RuntimeError("REPLICATE_API_TOKEN is missing; cannot use Replicate models.")
-    if replicate_model_key not in MODEL_REGISTRY:
-        raise KeyError(f"Unknown replicate model key: {replicate_model_key}")
 
     model_cfg = MODEL_REGISTRY[replicate_model_key]
     model_id = model_cfg["id"]
-    param_name = model_cfg.get("param_name", "aspect_ratio")
 
-    try:
-        output = replicate.run(
-            model_id,
-            input={"prompt": prompt, param_name: aspect_ratio},
-        )
-    except Exception as e:
-        raise RuntimeError(f"Replicate run failed ({model_id}): {e}") from e
+    if replicate_model_key == "z-image-turbo":
+        input_payload = {
+            "prompt": prompt,
+            "width": z_width or 1024,
+            "height": z_height or 1024,
+            "output_quality": z_output_quality or 80,
+        }
+    else:
+        param_name = model_cfg.get("param_name", "aspect_ratio")
+        input_payload = {
+            "prompt": prompt,
+            param_name: aspect_ratio,
+        }
 
-    items: List[Any] = output if isinstance(output, list) else [output]
+    output = replicate.run(model_id, input=input_payload)
 
+    items = output if isinstance(output, list) else [output]
     images: List[Image.Image] = []
+
     for item in items[:1]:
-        try:
-            if hasattr(item, "read"):
-                data = item.read()
-                images.append(Image.open(BytesIO(data)))
-                continue
-
-            url = str(item)
-            if url.startswith(("http://", "https://")):
-                with urllib.request.urlopen(url) as resp:
-                    data = resp.read()
-                images.append(Image.open(BytesIO(data)))
-                continue
-
-            raise RuntimeError(f"Unsupported Replicate output item: {type(item)} / {item}")
-        except Exception as e:
-            raise RuntimeError(f"Failed to read Replicate output: {e}") from e
+        if hasattr(item, "read"):
+            images.append(Image.open(BytesIO(item.read())))
+        else:
+            with urllib.request.urlopen(str(item)) as resp:
+                images.append(Image.open(BytesIO(resp.read())))
 
-    if not images:
-        raise RuntimeError("Replicate returned no usable image outputs.")
     return images
 
 
+
+
 def _image_task(
     idx: int,
     prompt: str,
@@ -300,13 +305,29 @@ def _image_task(
     size: str,
     quality: str,
     aspect_ratio: str,
+    z_width: int,
+    z_height: int,
+    z_output_quality: int,
 ) -> Tuple[int, str, Optional[Image.Image], Optional[bytes], Optional[str], Optional[str]]:
     try:
         if model_key in OPENAI_IMAGE_MODELS:
             openai_model_id = OPENAI_IMAGE_MODELS[model_key]
-            img = generate_ad_images_openai(prompt, model_id=openai_model_id, n=1, size=size, quality=quality)[0]
+            img = generate_ad_images_openai(
+                prompt,
+                model_id=openai_model_id,
+                n=1,
+                size=size,
+                quality=quality
+            )[0]
         else:
-            img = generate_ad_images_replicate(prompt, replicate_model_key=model_key, aspect_ratio=aspect_ratio)[0]
+            img = generate_ad_images_replicate(
+                prompt,
+                replicate_model_key=model_key,
+                aspect_ratio=aspect_ratio,
+                z_width=z_width,
+                z_height=z_height,
+                z_output_quality=z_output_quality,
+            )[0]
 
         buf = BytesIO()
         img.save(buf, format="PNG", optimize=True)
@@ -318,6 +339,7 @@ def _image_task(
             app_type="adgen",
             format="PNG",
         )
+
         if not url:
             return idx, prompt, img, png_bytes, None, "R2 upload failed"
 
@@ -339,6 +361,7 @@ def render_image_generation_page(uid: str | None = None):
         "Imagen 4 Ultra": "imagegen-4-ultra",
         "Imagen 4": "imagen-4",
         "gpt-image-1.5": "gpt-image-1.5",
+        "z-image-turbo":"z-image-turbo"
     }
 
     try:
@@ -352,9 +375,25 @@ def render_image_generation_page(uid: str | None = None):
 
     model_label = st.selectbox("Image Model Name (required)", list(UI_MODEL_OPTIONS.keys()))
     model_key = UI_MODEL_OPTIONS[model_label]
-
     aspect_ratio = "1:1"
-    if model_key not in OPENAI_IMAGE_MODELS:
+    z_width = 1024
+    z_height = 1024
+    z_output_quality = 80
+
+    if model_key == "z-image-turbo":
+        col1, col2, col3 = st.columns(3)
+
+        with col1:
+            z_width = st.slider("Width", 64, 1440, 1024, step=64)
+
+        with col2:
+            z_height = st.slider("Height", 64, 1440, 1024, step=64)
+
+        with col3:
+            z_output_quality = st.slider("Output Quality", 0, 100, 80)
+
+
+    elif model_key not in OPENAI_IMAGE_MODELS:
         ratios = MODEL_REGISTRY[model_key]["aspect_ratios"]
         aspect_ratio = st.selectbox(
             "Aspect ratio",
@@ -442,7 +481,7 @@ def render_image_generation_page(uid: str | None = None):
         try:
             with ThreadPoolExecutor(max_workers=max_workers) as ex:
                 futures = [
-                    ex.submit(_image_task, i, p, category, model_key, size, quality, aspect_ratio)
+                    ex.submit(_image_task, i, p, category, model_key, size, quality, aspect_ratio,z_width,z_height,z_output_quality)
                     for i, p in enumerate(output.prompt, start=1)
                 ]
                 results = [f.result() for f in as_completed(futures)]
@@ -518,4 +557,4 @@ def render_image_generation_page(uid: str | None = None):
             file_name="ad_images.zip",
             mime="application/zip",
             key="zip_dl_last",
-        )
+        )