updated

app.py

@@ -4,48 +4,50 @@ from diffusers import AutoPipelineForText2Image
 from transformers import pipeline
 import json
 
-# ---------- TEXT MODEL ----------
+# ---------- DEVICE ----------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# ---------- TEXT MODEL (SAFE + FAST) ----------
 text_model = pipeline(
-    "text2text-generation",
-    model="google/flan-t5-base"
+    "text-generation",
+    model="google/flan-t5-small"
 )
 
-# ---------- IMAGE MODEL (FAST) ----------
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
+# ---------- IMAGE MODEL (FAST SDXL TURBO) ----------
 pipe = AutoPipelineForText2Image.from_pretrained(
     "stabilityai/sdxl-turbo",
     torch_dtype=torch.float16 if device == "cuda" else torch.float32
 )
 pipe = pipe.to(device)
 
-# ---------- DECISION MODEL ----------
+# ---------- DECISION (AGENT) ----------
 def decide_action(user_input):
     prompt = f"""
-Classify the user request.
-
-Rules:
-- If user wants image → return JSON: {{"action": "image"}}
-- Otherwise → return JSON: {{"action": "text"}}
+Decide action:
+- If user wants image → return JSON: {{"action":"image"}}
+- Otherwise → return JSON: {{"action":"text"}}
 
 User: {user_input}
 """
 
-    response = text_model(prompt, max_length=50)[0]["generated_text"]
-
     try:
-        decision = json.loads(response)
+        result = text_model(prompt, max_length=50)[0]["generated_text"]
+        decision = json.loads(result)
         return decision.get("action", "text")
     except:
         return "text"
 
 # ---------- TOOLS ----------
 def text_tool(prompt):
-    response = text_model(prompt, max_length=200)
-    return response[0]["generated_text"]
+    result = text_model(prompt, max_length=200)
+    return result[0]["generated_text"]
 
 def image_tool(prompt):
-    image = pipe(prompt, num_inference_steps=2, guidance_scale=0.0).images[0]
+    image = pipe(
+        prompt,
+        num_inference_steps=2,
+        guidance_scale=0.0
+    ).images[0]
     return image
 
 # ---------- AGENT ----------
@@ -66,10 +68,10 @@ def chat(user_input, history):
     return history, image
 
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Agentic AI (Fast Text + Image)")
+    gr.Markdown("# 🤖 Agentic AI (Text + Image)")
 
     chatbot = gr.Chatbot()
-    image_output = gr.Image()
+    image_output = gr.Image(label="Generated Image")
 
     with gr.Row():
         inp = gr.Textbox(placeholder="Ask anything or generate image...")