Update app.py

app.py

@@ -1,22 +1,26 @@
-# yellowflash_with_gemini25_images.py
-# Hardcoded API keys for TESTING ONLY (do not publish).
+# yellowflash_chat_with_separate_image_dropdown_keys.py
+# TEST ONLY: hardcoded API keys included (do NOT publish)
 
-import time, io, base64, mimetypes, traceback
-import requests
+import time, io, base64, traceback, requests
 import gradio as gr
 from PIL import Image
 
 # ---------------------------
-# HARDCODED KEYS
+# HARDCODED KEYS (TESTING)
 # ---------------------------
-GEMINI_KEY = "YOUR_GEMINI_KEY"
-GROQ_KEY   = "YOUR_GROQ_KEY"
-GROQ_URL   = "https://api.groq.com/openai/v1/chat/completions"
+# Gemini text
+GEMINI_KEY = "AIzaSyAPfDiu2V_aD6un00qHt5bkISm6C0Pkx7o"
+GEMINI_TEXT_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
+
+# Groq / LLaMA
+GROQ_KEY = "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"
+GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"
 GROQ_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
 
-# ---------------------------
-# Lazy import of google-genai
-# ---------------------------
+# Gemini image streaming (2.5)
+GEMINI_IMAGE_MODEL = "gemini-2.5-flash-image-preview"  # streaming model name
+
+# lazy genai client
 GENAI_AVAILABLE = False
 GENAI_CLIENT = None
 
@@ -30,93 +34,130 @@ def ensure_genai_client():
         GENAI_AVAILABLE = True
         return True
     except Exception:
+        GENAI_AVAILABLE = False
+        GENAI_CLIENT = None
         return False
 
 # ---------------------------
-# Text-only calls
+# helpers
 # ---------------------------
+def post_with_retries(url, headers, payload, timeout=18, max_retries=2):
+    for i in range(max_retries):
+        try:
+            r = requests.post(url, headers=headers, json=payload, timeout=timeout)
+            r.raise_for_status()
+            return r
+        except Exception:
+            if i == max_retries - 1:
+                raise
+            time.sleep(0.5 + i)
+    raise Exception("Max retries exceeded")
+
+# ---------------------------
+# Text model calls
+# ---------------------------
+def call_gemini_text(message, history):
+    headers = {"Content-Type":"application/json","x-goog-api-key": GEMINI_KEY}
+    contents = []
+    for u, m in (history or []):
+        contents.append({"role":"user","parts":[{"text":u}]})
+        contents.append({"role":"model","parts":[{"text":m}]})
+    contents.append({"role":"user","parts":[{"text":message}]})
+    payload = {"contents": contents}
+    r = post_with_retries(GEMINI_TEXT_URL, headers, payload, timeout=30)
+    data = r.json()
+    return data.get("candidates",[{}])[0].get("content",{}).get("parts",[{}])[0].get("text","")
+
 def call_llama_text(message, history):
     headers = {"Authorization": f"Bearer {GROQ_KEY}", "Content-Type": "application/json"}
     msgs = []
-    for u, m in history:
-        msgs.append({"role": "user", "content": u})
-        msgs.append({"role": "assistant", "content": m})
-    msgs.append({"role": "user", "content": message})
+    for u, m in (history or []):
+        msgs.append({"role":"user","content":u})
+        msgs.append({"role":"assistant","content":m})
+    msgs.append({"role":"user","content":message})
     payload = {"model": GROQ_MODEL, "messages": msgs}
-    r = requests.post(GROQ_URL, headers=headers, json=payload, timeout=25)
+    r = post_with_retries(GROQ_URL, headers, payload, timeout=30)
     data = r.json()
     if "choices" in data and data["choices"]:
-        return data["choices"][0]["message"]["content"]
+        ch = data["choices"][0]
+        if isinstance(ch.get("message"), dict):
+            return ch["message"].get("content","")
+        return ch.get("text","")
     return str(data)
 
 # ---------------------------
-# Gemini 2.5 Image/Text call
+# Gemini 2.5 image streaming call
 # ---------------------------
 def call_gemini25_image_text(prompt):
     if not ensure_genai_client():
-        return "ERROR: google-genai not installed. Run `pip install google-genai`."
-
-    from google.genai import types
-
-    model = "gemini-2.5-flash-image-preview"
-    contents = [types.Content(role="user", parts=[types.Part.from_text(text=prompt)])]
-    config = types.GenerateContentConfig(response_modalities=["IMAGE", "TEXT"])
+        return "ERROR: google-genai client not installed. Add `google-genai` to requirements and rebuild."
+    try:
+        from google.genai import types
+        messages = []
+        image = None
 
-    messages = []
-    image = None
+        contents = [types.Content(role="user", parts=[types.Part.from_text(text=prompt)])]
+        cfg = types.GenerateContentConfig(response_modalities=["IMAGE", "TEXT"])
 
-    try:
-        for chunk in GENAI_CLIENT.models.generate_content_stream(
-            model=model, contents=contents, config=config
-        ):
+        for chunk in GENAI_CLIENT.models.generate_content_stream(model=GEMINI_IMAGE_MODEL, contents=contents, config=cfg):
             if not chunk.candidates:
                 continue
-            parts = chunk.candidates[0].content.parts
-            if not parts:
+            cand = chunk.candidates[0]
+            if not getattr(cand, "content", None):
                 continue
-
-            part = parts[0]
-            if hasattr(part, "text") and part.text:
-                messages.append(part.text)
-            elif hasattr(part, "inline_data") and part.inline_data:
-                data = part.inline_data.data
-                mime = part.inline_data.mime_type
-                ext = mimetypes.guess_extension(mime) or ".jpg"
-                try:
-                    img_bytes = io.BytesIO(data)
-                    image = Image.open(img_bytes).convert("RGB")
-                except Exception as e:
-                    messages.append(f"[Image decoding failed: {e}]")
+            parts = cand.content.parts or []
+            for part in parts:
+                if getattr(part, "text", None):
+                    messages.append(part.text)
+                inline = getattr(part, "inline_data", None)
+                if inline and getattr(inline, "data", None):
+                    try:
+                        img_bytes = io.BytesIO(inline.data)
+                        image = Image.open(img_bytes).convert("RGB")
+                    except Exception:
+                        data_uri = getattr(inline, "dataURI", None) or getattr(inline, "data_uri", None)
+                        if data_uri and data_uri.startswith("data:"):
+                            header, b64 = data_uri.split(",", 1)
+                            img_bytes = io.BytesIO(base64.b64decode(b64))
+                            image = Image.open(img_bytes).convert("RGB")
+
+        caption = " ".join([m for m in messages if m]).strip()
+        if not caption:
+            caption = f"Generated image for: {prompt}"
+        if image:
+            return [(caption, image)]
+        return caption
     except Exception as e:
-        return f"Error: {e}\n{traceback.format_exc()}"
-
-    caption = " ".join(messages).strip() or f"Generated image for: {prompt}"
-    if image:
-        return [(caption, image)]
-    return caption
+        return f"Image generation error: {e}\n{traceback.format_exc()}"
 
 # ---------------------------
-# Chat fn for ChatInterface
+# Chat fn
 # ---------------------------
-def chat_fn(message, history, model_choice):
+def chat_fn(message, history, text_model_choice, image_model_choice):
     try:
-        if model_choice == "Gemini Imagen 2.5":
+        if image_model_choice == "Gemini Imagen 2.5":
             return call_gemini25_image_text(message)
-        elif model_choice == "Meta LLaMA 4":
+        if text_model_choice == "Google Gemini 2.0 Flash":
+            return call_gemini_text(message, history)
+        elif text_model_choice == "Meta LLaMA 4":
             return call_llama_text(message, history)
         else:
-            return "Unknown model."
+            return "Unknown model selected."
     except Exception as e:
-        return f"Error: {e}"
+        return f"Error: {e}\n{traceback.format_exc()}"
 
 # ---------------------------
 # CSS
 # ---------------------------
 css = """
-#topbar { display:flex; justify-content:space-between; align-items:center; padding:14px 20px; background:#0f0f0f; border-bottom:1px solid #1f1f1f; position:fixed; top:0; left:0; right:0; z-index:999; }
-#title { font-weight:800; color:#ffcc33; font-size:18px; }
-#model_dropdown .gr-dropdown { background:transparent !important; border:1px solid #2b2b2b !important; color:#ddd !important; padding:8px 10px !important; border-radius:8px !important; width:260px !important; }
-.gradio-container { padding-top: 72px !important; }
+:root{ --topbar-h:68px; --accent:#2cc3ff; --title:#ffcc33;}
+body { background:#0d0d0d; color:#fff; font-family:Inter, Arial, sans-serif; margin:0; }
+#topbar { display:flex; justify-content:space-between; align-items:center; padding:12px 20px; background:#0f0f0f; border-bottom:1px solid #1f1f1f; position:fixed; left:0; right:0; top:0; z-index:999; height:var(--topbar-h); }
+#title { font-weight:800; color:var(--title); font-size:18px; padding-left:12px; }
+#left_controls { display:flex; gap:12px; align-items:center; }
+#model_dropdown .gr-dropdown, #image_dropdown .gr-dropdown { background:transparent !important; border:1px solid #2b2b2b !important; color:#ddd !important; padding:8px 12px !important; border-radius:8px !important; width:260px !important; box-shadow:none !important; }
+.gradio-container { padding-top: calc(var(--topbar-h) + 8px) !important; }
+.gradio-container .chat-interface .chatbot { min-height: calc(100vh - 220px); }
 """
 
 # ---------------------------
@@ -124,19 +165,26 @@ css = """
 # ---------------------------
 with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
     with gr.Row(elem_id="topbar"):
-        model_dropdown = gr.Dropdown(
-            choices=["Meta LLaMA 4", "Gemini Imagen 2.5"],
-            value="Meta LLaMA 4",
-            show_label=False,
-            elem_id="model_dropdown"
-        )
+        with gr.Row(elem_id="left_controls"):
+            text_model_dropdown = gr.Dropdown(
+                choices=["Google Gemini 2.0 Flash", "Meta LLaMA 4"],
+                value="Google Gemini 2.0 Flash",
+                show_label=False,
+                elem_id="model_dropdown"
+            )
+            image_model_dropdown = gr.Dropdown(
+                choices=["None", "Gemini Imagen 2.5"],
+                value="None",
+                show_label=False,
+                elem_id="image_dropdown"
+            )
         gr.HTML("<div id='title'>⚡ YellowFlash.ai</div>")
 
     gr.ChatInterface(
         fn=chat_fn,
-        title="⚡",
-        description="Text + Image with Gemini 2.5",
-        additional_inputs=[model_dropdown],
+        title="⚡ YellowFlash.ai",
+        description="Select a text model or switch image dropdown to 'Gemini Imagen 2.5' for inline image generation.",
+        additional_inputs=[text_model_dropdown, image_model_dropdown],
     )
 
 app.launch(share=True)