updated

app.py

@@ -1,31 +1,52 @@
 import gradio as gr
 from transformers import pipeline
+from diffusers import AutoPipelineForText2Image
+import torch
 
-# -------- MODEL --------
+# -------- TEXT MODEL --------
 text_model = pipeline("text-generation", model="distilgpt2")
 
+# -------- IMAGE MODEL (LIGHT) --------
+pipe = AutoPipelineForText2Image.from_pretrained(
+    "stabilityai/sd-turbo",
+    torch_dtype=torch.float32
+)
+pipe = pipe.to("cpu")
+
 # -------- FUNCTION --------
 def respond(user_input, history):
-    result = text_model(user_input, max_length=80)[0]["generated_text"]
+    if any(word in user_input.lower() for word in ["draw", "image", "generate", "picture"]):
+        image = pipe(
+            user_input,
+            num_inference_steps=1,
+            guidance_scale=0.0
+        ).images[0]
 
-    history = history + [
-        {"role": "user", "content": user_input},
-        {"role": "assistant", "content": result}
-    ]
+        history.append((user_input, "🖼️ Image generated"))
+        return history, history, image
 
-    return history, history
+    else:
+        result = text_model(user_input, max_length=80)[0]["generated_text"]
+        history.append((user_input, result))
+        return history, history, None
 
 # -------- UI --------
 with gr.Blocks() as demo:
-    gr.Markdown("# Simple AI Chat (Fixed)")
+    gr.Markdown("# AI Chat + Image (Stable Version)")
+
+    chatbot = gr.Chatbot()
+    image_output = gr.Image()
 
-    chatbot = gr.Chatbot(type="messages")  # ✅ IMPORTANT
     state = gr.State([])
 
     with gr.Row():
-        inp = gr.Textbox(placeholder="Type something...")
+        inp = gr.Textbox(placeholder="Ask or generate image...")
         btn = gr.Button("Send")
 
-    btn.click(respond, [inp, state], [chatbot, state])
+    btn.click(
+        respond,
+        [inp, state],
+        [chatbot, state, image_output]
+    )
 
 demo.launch()