Update app.py

app.py

@@ -7,19 +7,10 @@ from transformers import pipeline
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # ----------------------------
-# Load text model (fallbacks)
+# Load lightweight text model
 # ----------------------------
-if device == "cuda":
-    text_model_name = "HuggingFaceH4/zephyr-7b-beta"
-else:
-    text_model_name = "google/flan-t5-base"  # CPU-friendly
-
-chat_model = pipeline(
-    "text-generation",
-    model=text_model_name,
-    device=0 if device=="cuda" else -1,
-    return_full_text=False
-)
+text_model_name = "google/flan-t5-small"  # tiny, CPU-friendly
+chat_model = pipeline("text2text-generation", model=text_model_name, device=0 if device=="cuda" else -1)
 
 # ----------------------------
 # Try to load Stable Diffusion (only if GPU)
@@ -29,24 +20,21 @@ if device == "cuda":
     try:
         from diffusers import StableDiffusionPipeline
         sd_model = StableDiffusionPipeline.from_pretrained(
-            "stabilityai/stable-diffusion-2-1",
-            torch_dtype=torch.float16
+            "stabilityai/stable-diffusion-2-1-base"
         ).to(device)
     except Exception as e:
         print("⚠️ Could not load Stable Diffusion:", e)
         sd_model = None
 
-
+# ----------------------------
+# Core assistant logic
+# ----------------------------
 SYSTEM_PROMPT = """You are ZEN Research Assistant.
 You can respond in ONE of these forms:
 - Image → {"type":"image","prompt":"<prompt>"}
-- Chart → {"type":"chart","title":"<chart title>","data":[{"x":[...], "y":[...], "label":"<series>"}]}
+- Chart → {"type":"chart","title":"<chart title>","data":[{"x":[1,2,3], "y":[2,4,6], "label":"Example"}]}
 - Simulation → {"type":"simulation","topic":"<title>","steps":["...", "..."]}
 - Text → plain conversation.
-
-Rules:
-- Use JSON ONLY for image, chart, or simulation.
-- Simulation = imaginative thought experiment, 3–6 steps.
 """
 
 def query_llm(prompt, history, persona):
@@ -57,10 +45,9 @@ def query_llm(prompt, history, persona):
         input_text += f"User: {u}\nAssistant: {a}\n"
     input_text += f"User: {prompt}\nAssistant:"
 
-    out = chat_model(input_text, max_new_tokens=300, do_sample=True, temperature=0.7)
+    out = chat_model(input_text, max_new_tokens=256)
     return out[0]["generated_text"].strip()
 
-
 def multimodal_chat(user_msg, history, persona):
     history = history or []
     assistant_content = query_llm(user_msg, history, persona)
@@ -74,7 +61,7 @@ def multimodal_chat(user_msg, history, persona):
                 img = sd_model(parsed["prompt"]).images[0]
                 history.append([user_msg, "🖼️ Generated image below."])
             else:
-                history.append([user_msg, "⚠️ Image generation not available on this hardware."])
+                history.append([user_msg, "⚠️ Image generation requires GPU."])
 
         elif parsed.get("type") == "chart":
             fig = go.Figure()
@@ -97,11 +84,12 @@ def multimodal_chat(user_msg, history, persona):
 
     return history, img, fig
 
-
+# ----------------------------
+# Gradio UI
+# ----------------------------
 with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("🧠 **ZEN Research Lab (Adaptive Edition)**", elem_id="zen-header")
+    gr.Markdown("🧠 **ZEN Research Lab (Light Mode)**", elem_id="zen-header")
 
-    # Capabilities banner
     cap_text = "✅ Text  ✅ Charts  ✅ Simulation"
     if sd_model is not None:
         cap_text += "  ✅ Images"
@@ -132,18 +120,17 @@ with gr.Blocks(css="style.css") as demo:
     user_msg.submit(respond, inputs=[user_msg, chatbot, persona],
                     outputs=[chatbot, img_out, chart_out])
 
-    # Example starter buttons
+    # Examples
     with gr.Accordion("✨ Try these examples"):
-        with gr.Row():
-            gr.Examples(
-                examples=[
-                    ["Draw a futuristic city skyline at night"],
-                    ["Simulate first contact with an alien civilization"],
-                    ["Make a chart of AI adoption from 2010 to 2030"],
-                    ["Explain quantum entanglement in simple terms"],
-                ],
-                inputs=[user_msg]
-            )
+        gr.Examples(
+            examples=[
+                ["Draw a futuristic city skyline at night"],
+                ["Simulate first contact with an alien civilization"],
+                ["Make a chart of AI adoption from 2010 to 2030"],
+                ["Explain quantum entanglement in simple terms"],
+            ],
+            inputs=[user_msg]
+        )
 
 if __name__ == "__main__":
     demo.queue(max_size=50).launch()