Update app.py

app.py

@@ -1,20 +1,35 @@
-import gradio as gr, json, plotly.graph_objects as go
+import gradio as gr, json, plotly.graph_objects as go, torch
 from transformers import pipeline
-from diffusers import StableDiffusionPipeline
-import torch
 
 # ----------------------------
-# Load models once on startup
+# Detect device
 # ----------------------------
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Text model (fast chat)
-chat_model = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", device=0 if device=="cuda" else -1)
+# ----------------------------
+# Load text model (fallbacks)
+# ----------------------------
+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)
 
-# Image model (stable diffusion)
-sd_model = StableDiffusionPipeline.from_pretrained(
-    "stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16 if device=="cuda" else torch.float32
-).to(device)
+# ----------------------------
+# Try to load Stable Diffusion (only if GPU)
+# ----------------------------
+sd_model = None
+if device == "cuda":
+    try:
+        from diffusers import StableDiffusionPipeline
+        sd_model = StableDiffusionPipeline.from_pretrained(
+            "stabilityai/stable-diffusion-2-1",
+            torch_dtype=torch.float16
+        ).to(device)
+    except Exception as e:
+        print("⚠️ Could not load Stable Diffusion:", e)
+        sd_model = None
 
 
 SYSTEM_PROMPT = """You are ZEN Research Assistant.
@@ -22,16 +37,14 @@ You can respond in ONE of these forms:
 - Image → {"type":"image","prompt":"<prompt>"}
 - Chart → {"type":"chart","title":"<chart title>","data":[{"x":[...], "y":[...], "label":"<series>"}]}
 - Simulation → {"type":"simulation","topic":"<title>","steps":["...", "..."]}
-- Text → plain conversation, explanation, or reasoning.
+- Text → plain conversation.
 
 Rules:
 - Use JSON ONLY for image, chart, or simulation.
 - Simulation = imaginative thought experiment, 3–6 steps.
-- If not sure, default to conversational text.
 """
 
 def query_llm(prompt, history, persona):
-    # Construct conversation
     input_text = SYSTEM_PROMPT
     if persona != "Default":
         input_text += f"\nPersona: {persona}\n"
@@ -39,7 +52,7 @@ 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=400, do_sample=True, temperature=0.7)
+    out = chat_model(input_text, max_new_tokens=300, do_sample=True, temperature=0.7)
     return out[0]["generated_text"].split("Assistant:")[-1].strip()
 
 
@@ -52,8 +65,11 @@ def multimodal_chat(user_msg, history, persona):
         parsed = json.loads(assistant_content)
 
         if parsed.get("type") == "image":
-            img = sd_model(parsed["prompt"]).images[0]
-            history.append([user_msg, "🖼️ Generated image below."])
+            if sd_model is not None:
+                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."])
 
         elif parsed.get("type") == "chart":
             fig = go.Figure()
@@ -65,7 +81,7 @@ def multimodal_chat(user_msg, history, persona):
             history.append([user_msg, parsed.get("title","Chart below")])
 
         elif parsed.get("type") == "simulation":
-            steps = "\n".join([f"→ {s}" for s in parsed["steps"]])
+            steps = "\n".join([f"→ {s}" for s in parsed.get("steps",[])])
             history.append([user_msg, f"🔮 Simulation: {parsed.get('topic','Exploration')}\n{steps}"])
 
         else:
@@ -78,7 +94,7 @@ def multimodal_chat(user_msg, history, persona):
 
 
 with gr.Blocks(css="style.css") as demo:
-    gr.Markdown("🧠 **ZEN Research Lab (API-free Edition)** — Explore, simulate, and create", elem_id="zen-header")
+    gr.Markdown("🧠 **ZEN Research Lab (Adaptive Edition)** — works everywhere, GPU unlocks extra powers", elem_id="zen-header")
 
     persona = gr.Dropdown(["Default","Analyst","Artist","Futurist","Philosopher"], label="Mode", value="Default")
     chatbot = gr.Chatbot(label="Conversation", height=400)