Update app.py

app.py

@@ -1,118 +1,111 @@
-# yellowflash_full.py
-# ⚡ YellowFlash.ai - Multi Model Chat + Image Gen
-# TEST ONLY: Hardcoded API keys (do not push public)
+# yellowflash_chat_with_image.py
+# Native ChatInterface + Gemini image gen in-chat
 
-import time
-import requests
+import time, requests, io, base64
 import gradio as gr
-from google import genai
+from PIL import Image
 
 # ---------------------------
 # HARDCODED KEYS (TESTING)
 # ---------------------------
 GEMINI_KEY = "AIzaSyAPfDiu2V_aD6un00qHt5bkISm6C0Pkx7o"
-GEMINI_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
+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"
 GROQ_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct"
 
-# Gemini Image Gen (Imagen 2.0)
-GENAI_CLIENT = genai.Client(api_key=GEMINI_KEY)
-GEMINI_IMAGE_MODEL = "models/imagen-4.0-generate-001"
-
 # ---------------------------
-# Helpers
+# IMAGE GENERATION CALL (using google-genai style API)
 # ---------------------------
-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 as e:
-            if i == max_retries - 1:
-                raise
-            time.sleep(0.5 + i)
-    raise Exception("Max retries exceeded")
+def call_gemini_image(prompt_text):
+    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."
 
-def call_gemini(api_key, message, history):
-    headers = {"Content-Type": "application/json", "x-goog-api-key": api_key}
+        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}"
+
+# ---------------------------
+# TEXT 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 = post_with_retries(GEMINI_URL, headers, payload)
+    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_via_groq(api_key, model, message, history):
-    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}
+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})
-    payload = {"model": model, "messages": msgs}
-    r = post_with_retries(GROQ_URL, headers, payload)
+    payload = {"model": GROQ_MODEL, "messages": msgs}
+    r = requests.post(GROQ_URL, headers=headers, json=payload, timeout=20)
     data = r.json()
     if "choices" in data and data["choices"]:
         ch = data["choices"][0]
-        if isinstance(ch.get("message"), dict):
-            return ch["message"].get("content","")
-        return ch.get("text","")
+        return ch.get("message",{}).get("content","")
     return str(data)
 
-def call_gemini_image(prompt):
-    try:
-        result = GENAI_CLIENT.models.generate_images(
-            model=GEMINI_IMAGE_MODEL,
-            prompt=prompt,
-            config=dict(
-                number_of_images=1,
-                output_mime_type="image/jpeg",
-                person_generation="ALLOW_ALL",
-                aspect_ratio="1:1",
-                image_size="1K",
-            ),
-        )
-        if not result.generated_images:
-            return None
-        img = result.generated_images[0].image
-        return img
-    except Exception as e:
-        return str(e)
-
 # ---------------------------
-# Chat function
+# CHAT FUNCTION
 # ---------------------------
 def chat_fn(message, history, model_choice):
     try:
         if model_choice == "Google Gemini 2.0 Flash":
-            ans = call_gemini(GEMINI_KEY, message, history)
+            return call_gemini_text(message, history)
         elif model_choice == "Meta LLaMA 4":
-            ans = call_llama_via_groq(GROQ_KEY, GROQ_MODEL, message, history)
+            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:
-            ans = "Error: For image generation, use the special Image Gen UI below."
+            return "⚠️ Unknown model selected."
     except Exception as e:
-        ans = f"Error: {e}"
-    return ans
+        return f"Error: {e}"
 
 # ---------------------------
 # CSS
 # ---------------------------
 css = """
-#topbar { display:flex; justify-content:space-between; align-items:center; padding:18px 28px; background:#0f0f0f; border-bottom:1px solid #1f1f1f; }
+#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:10px 12px !important; border-radius:8px !important; width:260px !important; box-shadow:none !important; }
+#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; }
 """
 
 # ---------------------------
-# Build UI
+# BUILD UI
 # ---------------------------
 with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
     with gr.Row(elem_id="topbar"):
@@ -122,22 +115,13 @@ with gr.Blocks(css=css, title="⚡ YellowFlash.ai") as app:
             show_label=False,
             elem_id="model_dropdown"
         )
+        gr.HTML("<div id='title'>⚡ YellowFlash.ai</div>")
 
-    # Normal Chat
-    chat = gr.ChatInterface(
+    gr.ChatInterface(
         fn=chat_fn,
         title="⚡",
-        description="FAST AS LIGHTNING",
+        description="FAST AS LIGHTNING — text + image in one UI",
         additional_inputs=[model_dropdown],
     )
 
-    # Image Gen UI
-    with gr.Group():
-        gr.Markdown("### 🖼️ Gemini Imagen 2.0 - Generate Images")
-        with gr.Row():
-            img_prompt = gr.Textbox(label="Enter your prompt", scale=4)
-            gen_btn = gr.Button("Generate Image", scale=1)
-        img_output = gr.Image(label="Generated Image", type="pil")
-        gen_btn.click(fn=call_gemini_image, inputs=img_prompt, outputs=img_output)
-
 app.launch(share=True)