Create app.py

app.py

@@ -0,0 +1,82 @@
+import torch
+import gradio as gr
+from diffusers import AutoPipelineForText2Image
+from transformers import pipeline
+import json
+
+# ---------- TEXT MODEL ----------
+text_model = pipeline(
+    "text2text-generation",
+    model="google/flan-t5-base"
+)
+
+# ---------- IMAGE MODEL (FAST) ----------
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+pipe = AutoPipelineForText2Image.from_pretrained(
+    "stabilityai/sdxl-turbo",
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+)
+pipe = pipe.to(device)
+
+# ---------- DECISION MODEL ----------
+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"}}
+
+User: {user_input}
+"""
+
+    response = text_model(prompt, max_length=50)[0]["generated_text"]
+
+    try:
+        decision = json.loads(response)
+        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"]
+
+def image_tool(prompt):
+    image = pipe(prompt, num_inference_steps=2, guidance_scale=0.0).images[0]
+    return image
+
+# ---------- AGENT ----------
+def agent(user_input):
+    action = decide_action(user_input)
+
+    if action == "image":
+        img = image_tool(user_input)
+        return "🖼️ Image generated", img
+    else:
+        text = text_tool(user_input)
+        return text, None
+
+# ---------- UI ----------
+def chat(user_input, history):
+    text, image = agent(user_input)
+    history = history + [(user_input, text)]
+    return history, image
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 Agentic AI (Fast Text + Image)")
+
+    chatbot = gr.Chatbot()
+    image_output = gr.Image()
+
+    with gr.Row():
+        inp = gr.Textbox(placeholder="Ask anything or generate image...")
+        btn = gr.Button("Send")
+
+    state = gr.State([])
+
+    btn.click(chat, [inp, state], [chatbot, image_output])
+
+demo.launch()