Update app.py

app.py

@@ -1,92 +1,111 @@
-# yellowflash_chat_with_image.py
-# Native ChatInterface + Gemini image gen in-chat
+# yellowflash_with_gemini25_images.py
+# Hardcoded API keys for TESTING ONLY (do not publish).
 
-import time, requests, io, base64
+import time, io, base64, mimetypes, traceback
+import requests
 import gradio as gr
 from PIL import Image
 
 # ---------------------------
-# HARDCODED KEYS (TESTING)
+# HARDCODED KEYS
 # ---------------------------
-GEMINI_KEY = "AIzaSyAPfDiu2V_aD6un00qHt5bkISm6C0Pkx7o"
-GEMINI_TEXT_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
-
-GROQ_KEY = "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"
-GROQ_URL = "https://api.groq.com/openai/v1/chat/completions"
+GEMINI_KEY = "YOUR_GEMINI_KEY"
+GROQ_KEY   = "YOUR_GROQ_KEY"
+GROQ_URL   = "https://api.groq.com/openai/v1/chat/completions"
 GROQ_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
 
 # ---------------------------
-# IMAGE GENERATION CALL (using google-genai style API)
+# Lazy import of google-genai
 # ---------------------------
-def call_gemini_image(prompt_text):
+GENAI_AVAILABLE = False
+GENAI_CLIENT = None
+
+def ensure_genai_client():
+    global GENAI_AVAILABLE, GENAI_CLIENT
+    if GENAI_AVAILABLE:
+        return True
     try:
-        # NOTE: requires google-genai installed and correct endpoint
         from google import genai
-        client = genai.Client(api_key=GEMINI_KEY)
-        result = client.models.generate_images(
-            model="models/imagen-4.0-generate-001",
-            prompt=prompt_text,
-            config=dict(number_of_images=1, output_mime_type="image/jpeg"),
-        )
-        if not result.generated_images:
-            return "⚠️ No image generated."
-
-        gi = result.generated_images[0]
-        if hasattr(gi, "image"):
-            return gi.image  # PIL.Image -> Gradio can render inline
-        if hasattr(gi, "dataURI"):
-            header, b64 = gi.dataURI.split(",", 1)
-            img_bytes = io.BytesIO(base64.b64decode(b64))
-            return Image.open(img_bytes).convert("RGB")
-        return "⚠️ Unknown image response."
-    except Exception as e:
-        return f"Image generation error: {e}"
+        GENAI_CLIENT = genai.Client(api_key=GEMINI_KEY)
+        GENAI_AVAILABLE = True
+        return True
+    except Exception:
+        return False
 
 # ---------------------------
-# TEXT CALLS
+# Text-only calls
 # ---------------------------
-def call_gemini_text(message, history):
-    headers = {"Content-Type":"application/json","x-goog-api-key": GEMINI_KEY}
-    contents = []
-    for u, m in history:
-        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 = requests.post(GEMINI_TEXT_URL, headers=headers, json=payload, timeout=20)
-    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"}
+    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})
+        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=20)
+    r = requests.post(GROQ_URL, headers=headers, json=payload, timeout=25)
     data = r.json()
     if "choices" in data and data["choices"]:
-        ch = data["choices"][0]
-        return ch.get("message",{}).get("content","")
+        return data["choices"][0]["message"]["content"]
     return str(data)
 
 # ---------------------------
-# CHAT FUNCTION
+# Gemini 2.5 Image/Text 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"])
+
+    messages = []
+    image = None
+
+    try:
+        for chunk in GENAI_CLIENT.models.generate_content_stream(
+            model=model, contents=contents, config=config
+        ):
+            if not chunk.candidates:
+                continue
+            parts = chunk.candidates[0].content.parts
+            if not parts:
+                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}]")
+    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
+
+# ---------------------------
+# Chat fn for ChatInterface
 # ---------------------------
 def chat_fn(message, history, model_choice):
     try:
-        if model_choice == "Google Gemini 2.0 Flash":
-            return call_gemini_text(message, history)
+        if model_choice == "Gemini Imagen 2.5":
+            return call_gemini25_image_text(message)
         elif model_choice == "Meta LLaMA 4":
             return call_llama_text(message, history)
-        elif model_choice == "Gemini Imagen 2.0":
-            # Return image directly in chat
-            img = call_gemini_image(message)
-            return img
         else:
-            return "⚠️ Unknown model selected."
+            return "Unknown model."
     except Exception as e:
         return f"Error: {e}"
 
@@ -94,24 +113,20 @@ def chat_fn(message, history, model_choice):
 # CSS
 # ---------------------------
 css = """
-#topbar { display:flex; justify-content:space-between; align-items:center;
-          padding:16px 22px; background:#0f0f0f; border-bottom:1px solid #1f1f1f; }
-#title { font-weight:800; color:#ffcc33; font-size:20px; }
-#model_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; }
-.gradio-container .chat-interface .chatbot { min-height: calc(100vh - 220px); }
-.gr-button { border-radius:10px !important; }
+#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; }
 """
 
 # ---------------------------
-# BUILD UI
+# Build UI
 # ---------------------------
 with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
     with gr.Row(elem_id="topbar"):
         model_dropdown = gr.Dropdown(
-            choices=["Google Gemini 2.0 Flash", "Meta LLaMA 4", "Gemini Imagen 2.0"],
-            value="Google Gemini 2.0 Flash",
+            choices=["Meta LLaMA 4", "Gemini Imagen 2.5"],
+            value="Meta LLaMA 4",
             show_label=False,
             elem_id="model_dropdown"
         )
@@ -120,7 +135,7 @@ with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
     gr.ChatInterface(
         fn=chat_fn,
         title="⚡",
-        description="FAST AS LIGHTNING — text + image in one UI",
+        description="Text + Image with Gemini 2.5",
         additional_inputs=[model_dropdown],
     )