Update app.py

app.py

@@ -1,17 +1,17 @@
 import base64
 import io
-import os
 from typing import List, Tuple, Optional
 
 import gradio as gr
 from PIL import Image
 
+
 # -----------------------
-# Optional imports guarded at call-time
+# OpenAI + Google helpers
 # -----------------------
 
 def _get_openai_client(api_key: str):
-    from openai import OpenAI  # imported here to avoid issues if lib missing
+    from openai import OpenAI  # local import so app still loads if lib missing
     return OpenAI(api_key=api_key)
 
 
@@ -22,7 +22,7 @@ def _configure_google(api_key: str):
 
 
 # -----------------------
-# Helper: apply visual presets
+# Prompt / preset logic
 # -----------------------
 
 def apply_preset_to_prompt(
@@ -31,7 +31,6 @@ def apply_preset_to_prompt(
     style: str,
     content_type: str,
 ) -> str:
-    """Augment the prompt with preset & style language."""
     base_prompt = base_prompt.strip()
 
     preset_addons = {
@@ -68,7 +67,6 @@ def apply_preset_to_prompt(
         "Cinematic": " cinematic lighting, dramatic composition, filmic contrast",
     }
 
-    ct_addon = ""
     if content_type == "Image":
         ct_addon = " high-resolution concept art,"
     elif content_type == "Infographic Spec":
@@ -76,6 +74,8 @@ def apply_preset_to_prompt(
             " detailed infographic design specification, including layout regions, "
             "sections, labels, and visual hierarchy,"
         )
+    else:
+        ct_addon = ""
 
     extra = " ".join(
         x
@@ -92,11 +92,12 @@ def apply_preset_to_prompt(
             return f"{base_prompt}, {extra}"
         else:
             return extra.strip()
+
     return base_prompt or "high quality image"
 
 
 # -----------------------
-# OpenAI: Text + Image
+# OpenAI text + images
 # -----------------------
 
 def generate_text_openai(
@@ -105,6 +106,7 @@ def generate_text_openai(
     mode: str,
 ) -> str:
     client = _get_openai_client(api_key)
+
     system_msg = (
         "You are an expert creator for the ZEN AI ecosystem. "
         "Write clear, concise, high-leverage content. "
@@ -114,7 +116,8 @@ def generate_text_openai(
 
     if mode == "Infographic Spec":
         user_prompt = (
-            f"Create a Palantir/Anduril-level infographic specification based on:\n\n{prompt}\n\n"
+            "Create a Palantir/Anduril-level infographic specification based on this topic:\n\n"
+            f"{prompt}\n\n"
             "Return:\n"
             "1) Title options\n"
             "2) 3–5 main sections\n"
@@ -125,7 +128,6 @@ def generate_text_openai(
     else:
         user_prompt = prompt
 
-    # Using Chat Completions interface
     resp = client.chat.completions.create(
         model="gpt-4.1-mini",
         messages=[
@@ -157,7 +159,6 @@ def generate_image_openai(
 ) -> List[Image.Image]:
     client = _get_openai_client(api_key)
 
-    # Map size choices to OpenAI-supported ones
     size_map = {
         "Square (1024x1024)": "1024x1024",
         "Portrait (1024x1792)": "1024x1792",
@@ -172,7 +173,6 @@ def generate_image_openai(
         "quality": quality,
         "n": n_images,
     }
-    # seed is optional on some models; safe to include conditionally
     if seed is not None:
         kwargs["seed"] = seed
 
@@ -191,7 +191,6 @@ def generate_text_google(
     mode: str,
 ) -> str:
     genai = _configure_google(api_key)
-    # Default to a strong text model
     model = genai.GenerativeModel("gemini-1.5-pro")
 
     if mode == "Infographic Spec":
@@ -221,10 +220,9 @@ def generate_image_google(
     seed: Optional[int],
 ) -> List[Image.Image]:
     """
-    NOTE: Model & output handling may need adjustment depending on
-    the exact Nano-Banana / Nano-Banana-Pro API you use in Google AI Studio.
-
-    This assumes a GenerativeModel that returns inline image data.
+    This assumes your Nano-Banana / Nano-Banana-Pro image model in
+    Google AI Studio returns inline image bytes in the response.
+    Adjust parsing if your model behaves differently.
     """
     genai = _configure_google(api_key)
     model = genai.GenerativeModel(google_image_model)
@@ -232,8 +230,6 @@ def generate_image_google(
     images: List[Image.Image] = []
 
     for i in range(n_images):
-        # Some image models support generation_config with a seed;
-        # here we pass it if present.
         generation_config = {}
         if seed is not None:
             generation_config["seed"] = seed + i
@@ -243,11 +239,12 @@ def generate_image_google(
             generation_config=generation_config or None,
         )
 
-        # Try to pull image bytes from response parts
-        for cand in resp.candidates:
-            for part in cand.content.parts:
-                if hasattr(part, "inline_data") and getattr(part.inline_data, "data", None):
-                    raw = base64.b64decode(part.inline_data.data)
+        # Extract images from candidates
+        for cand in getattr(resp, "candidates", []):
+            for part in getattr(cand, "content", {}).parts:
+                inline = getattr(part, "inline_data", None)
+                if inline and getattr(inline, "data", None):
+                    raw = base64.b64decode(inline.data)
                     img = Image.open(io.BytesIO(raw)).convert("RGB")
                     images.append(img)
 
@@ -255,7 +252,7 @@ def generate_image_google(
 
 
 # -----------------------
-# Core Gradio callback
+# Core callback
 # -----------------------
 
 def run_generation(
@@ -273,11 +270,8 @@ def run_generation(
     seed: int,
     use_seed: bool,
     google_image_model: str,
-    google_text_model_hint: str,
+    google_text_model_hint: str,  # currently just logged
 ) -> Tuple[str, List[Image.Image], str]:
-    """
-    Returns: (text_output, images, debug_info)
-    """
     text_output = ""
     images: List[Image.Image] = []
     debug_lines = []
@@ -285,7 +279,6 @@ def run_generation(
     if not base_prompt.strip():
         return "Please enter a prompt.", [], "No prompt provided."
 
-    # Build full prompt for images
     content_type = "Image" if task_type == "Image" else task_type
     full_prompt = apply_preset_to_prompt(
         base_prompt=base_prompt,
@@ -300,7 +293,7 @@ def run_generation(
     debug_lines.append(f"Task: {task_type}")
     debug_lines.append(f"Provider: {provider}")
     debug_lines.append(f"Preset: {preset}, Style: {style}")
-    debug_lines.append(f"OpenAI model size: {size}, quality: {quality}")
+    debug_lines.append(f"OpenAI size: {size}, quality: {quality}")
     debug_lines.append(f"Google image model: {google_image_model}")
     debug_lines.append(f"Google text model hint: {google_text_model_hint}")
     debug_lines.append(f"Seed enabled: {use_seed}, seed: {seed if use_seed else 'None'}")
@@ -308,7 +301,7 @@ def run_generation(
     seed_val: Optional[int] = seed if use_seed else None
 
     try:
-        # TEXT or INFOGRAPHIC
+        # TEXT / INFOGRAPHIC
         if task_type in ["Text", "Infographic Spec"]:
             if provider == "OpenAI":
                 if not openai_key.strip():
@@ -332,11 +325,9 @@ def run_generation(
             if provider == "OpenAI":
                 if not openai_key.strip():
                     return "Missing OpenAI API key.", [], "OpenAI key not provided."
-                # Decide OpenAI image model based on preset selection or UI (we can infer)
-                # We'll expose choice via size/style; model stays fixed to user-facing dropdown externally
-                # but here we assume they want GPT-Image-1 by default
+
+                # Default to GPT-Image-1; for Palantir preset, swap to DALL·E 3
                 image_model = "gpt-image-1"
-                # For Palantir/Anduril preset, sometimes DALL·E 3 is good – user can switch later by editing code.
                 if "Palantir" in preset:
                     image_model = "dall-e-3"
 
@@ -355,7 +346,7 @@ def run_generation(
                     return "Missing Google API key.", [], "Google key not provided."
                 images = generate_image_google(
                     api_key=google_key.strip(),
-                    google_image_model=google_image_model.strip(),
+                    google_image_model=google_image_model.strip() or "nano-banana-pro",
                     prompt=full_prompt,
                     n_images=n_images,
                     seed=seed_val,
@@ -363,8 +354,8 @@ def run_generation(
 
         if not text_output and task_type == "Image":
             text_output = (
-                "Image(s) generated successfully. Add 'Text' or 'Infographic Spec' mode "
-                "to generate descriptive copy or specs."
+                "Image(s) generated. Use Text / Infographic Spec mode to "
+                "generate captions, copy, or layout specs."
             )
 
         if not images and task_type == "Image":
@@ -373,29 +364,31 @@ def run_generation(
         return text_output, images, "\n".join(debug_lines)
 
     except Exception as e:
-        return f"Error: {e}", [], "\n".join(debug_lines + [f"Exception: {e}"])
+        debug_lines.append(f"Exception: {e}")
+        return f"Error: {e}", [], "\n".join(debug_lines)
 
 
 # -----------------------
 # UI
 # -----------------------
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+with gr.Blocks() as demo:  # <- no theme arg
     gr.Markdown(
         """
 # 🧬 ZEN Omni Studio — Text • Images • Infographics
 
-A unified creator for the ZEN ecosystem.
+Multi-provider creator for the ZEN ecosystem:
 
-- 🔑 **Bring your own keys** for OpenAI + Google (Gemini / Nano-Banana).
-- 🧠 Generate **text** and **infographic specs** for ZEN, AI literacy, dashboards, and more.
-- 🎨 Generate **high-quality images** with advanced controls and presets.
+- 🔑 Bring your own OpenAI + Google (Gemini / Nano-Banana / Nano-Banana-Pro) keys  
+- 🎨 Generate **images** with presets + fine-grained controls  
+- 🧠 Generate **text** and **infographic specs** for ZEN dashboards, posters, and more  
         """
     )
 
     with gr.Row():
         with gr.Column():
-            gr.Markdown("### 🔐 API Keys (kept only in your browser session)")
+            gr.Markdown("### 🔐 API Keys (local to this session)")
+
             openai_key = gr.Textbox(
                 label="OPENAI_API_KEY",
                 type="password",
@@ -419,16 +412,15 @@ A unified creator for the ZEN ecosystem.
                 label="Primary Provider",
             )
 
-            # Prompt region
             base_prompt = gr.Textbox(
                 label="Main Prompt",
                 lines=5,
-                placeholder="Describe what you want to create for ZEN (image, copy, infographic, etc.)",
+                placeholder="Describe the ZEN image, text, or infographic you want.",
             )
             negative_prompt = gr.Textbox(
                 label="Negative Prompt (optional)",
                 lines=2,
-                placeholder="Things to avoid: low-res, cluttered, distorted text, etc.",
+                placeholder="Things to avoid: low-res, clutter, warped text, etc.",
             )
 
             with gr.Row():
@@ -457,8 +449,7 @@ A unified creator for the ZEN ecosystem.
                     label="Style Accent",
                 )
 
-            # OpenAI image options
-            gr.Markdown("### 🎛 OpenAI Image Controls (DALL·E / GPT-Image)")
+            gr.Markdown("### 🎛 OpenAI Image Controls")
             with gr.Row():
                 size = gr.Dropdown(
                     [
@@ -497,14 +488,14 @@ A unified creator for the ZEN ecosystem.
 
             gr.Markdown("### 🧪 Google Image / Text Model Hints")
             google_image_model = gr.Textbox(
-                label="Google Image Model (default: Nano-Banana-Pro)",
+                label="Google Image Model (default: nano-banana-pro)",
                 value="nano-banana-pro",
-                placeholder="e.g. nano-banana-pro (adjust to your actual model id in Google AI Studio)",
+                placeholder="e.g. nano-banana-pro or your exact model id",
             )
             google_text_model_hint = gr.Textbox(
-                label="Google Text Model Hint (for future tweaking)",
+                label="Google Text Model Hint",
                 value="gemini-1.5-pro",
-                placeholder="Used in code as default Gemini text model.",
+                placeholder="Used internally as default text model.",
             )
 
             generate_btn = gr.Button("🚀 Generate", variant="primary")
@@ -515,19 +506,17 @@ A unified creator for the ZEN ecosystem.
 
             gr.Markdown("### 🖼 Image Output")
             image_gallery = gr.Gallery(
-                label="Generated Images",
                 show_label=False,
                 columns=2,
                 height=500,
             )
 
-            gr.Markdown("### 🧾 Debug / Logs (for you, not end users)")
+            gr.Markdown("### 🧾 Debug / Logs")
             debug_output = gr.Textbox(
                 label="Debug Info",
                 lines=10,
             )
 
-    # Wire up callback
     generate_btn.click(
         fn=run_generation,
         inputs=[