Update app_pages/text_image_generation.py

app_pages/text_image_generation.py

@@ -14,12 +14,61 @@ def render_multimodel_ad_generator(prefix: str = "img_text"):
     colA, colB = st.columns([1, 1])
 
     with colA:
-        model_key = st.selectbox("Model", list(MODEL_REGISTRY.keys()), index=0, key=f"{prefix}_model")
-        aspect_options = get_model_config(model_key)["aspect_ratios"] if get_model_config(model_key) else ["1:1"]
-        aspect_ratio = st.selectbox("Aspect Ratio", aspect_options, index=0, key=f"{prefix}_ar")
-        num_images = st.slider("Number of images", 1, 10, 1, key=f"{prefix}_count")
+        model_key = st.selectbox("Model",list(MODEL_REGISTRY.keys()),index=0, key=f"{prefix}_model")
+        cfg = get_model_config(model_key)
+        if model_key == "z-image-turbo":
+            width = st.slider(
+                "Width",
+                min_value=256,
+                max_value=2048,
+                value=1024,
+                step=64,
+                key=f"{prefix}_width",
+            )
+            height = st.slider(
+                "Height",
+                min_value=256,
+                max_value=2048,
+                value=1024,
+                step=64,
+                key=f"{prefix}_height",
+            )
+            output_quality = st.slider(
+                "Output Quality",
+                min_value=50,
+                max_value=100,
+                value=90,
+                step=1,
+                key=f"{prefix}_quality",
+            )
+
+            aspect_ratio = ""  # not used for z-turbo
+        else:
+            aspect_options = cfg.get("aspect_ratios", ["1:1"]) if cfg else ["1:1"]
+
+            aspect_ratio = st.selectbox(
+                "Aspect Ratio",
+                aspect_options,
+                index=0,
+                key=f"{prefix}_ar",
+            )
+
+            width = height = output_quality = None
+
+        num_images = st.slider(
+            "Number of images",
+            1,
+            10,
+            1,
+            key=f"{prefix}_count",
+        )
+
+        category = st.text_input(
+            "Category",
+            key=f"{prefix}_cat",
+            placeholder="e.g. skincare, supplements",
+        )
 
-        category = st.text_input("Category", key=f"{prefix}_cat", placeholder="e.g. skincare, supplements")
         platform = st.selectbox(
             "Platform",
             ["Facebook", "Instagram", "TikTok", "Google", "Newsbreak"],
@@ -38,6 +87,15 @@ def render_multimodel_ad_generator(prefix: str = "img_text"):
         if not uid:
             st.error("Please log in to generate images.")
             return
+        params = {}
+
+        if model_key == "z-image-turbo":
+            params = {
+                "width": width,
+                "height": height,
+                "output_quality": output_quality,
+            }
+
         handle_image_generation_optimized(
             uid=uid,
             model_key=model_key,
@@ -47,4 +105,5 @@ def render_multimodel_ad_generator(prefix: str = "img_text"):
             category=(category.strip() or None),
             platform=platform or None,
             private_mode=private_mode,
+            params=params,
         )