Spaces:
Runtime error
Runtime error
File size: 1,454 Bytes
fd5528a 8fdb8a0 fd5528a 8fdb8a0 eba3780 1eb5cb6 78c692b 8fdb8a0 39172b7 78c692b 8fdb8a0 39172b7 78c692b 39172b7 8fdb8a0 78c692b fd5528a eba3780 fd5528a 78c692b 8fdb8a0 78c692b 8fdb8a0 fd5528a 78c692b eba3780 fd5528a 78c692b fd5528a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import gradio as gr
from transformers import pipeline
from diffusers import AutoPipelineForText2Image
import torch
# -------- TEXT MODEL --------
text_model = pipeline("text-generation", model="distilgpt2")
# -------- IMAGE MODEL --------
pipe = AutoPipelineForText2Image.from_pretrained(
"stabilityai/sd-turbo",
torch_dtype=torch.float32
)
pipe = pipe.to("cpu")
# -------- FUNCTION --------
def respond(message, history):
user_input = message
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.append({"role": "user", "content": user_input})
history.append({"role": "assistant", "content": "🖼️ Image generated"})
return history, image
else:
result = text_model(user_input, max_length=80)[0]["generated_text"]
history.append({"role": "user", "content": user_input})
history.append({"role": "assistant", "content": result})
return history, None
# -------- UI --------
with gr.Blocks() as demo:
gr.Markdown("# AI Chat + Image (Gradio 6 Fixed)")
chatbot = gr.Chatbot(type="messages")
image_output = gr.Image()
msg = gr.Textbox(placeholder="Ask or generate image...")
state = gr.State([])
msg.submit(respond, [msg, state], [chatbot, image_output])
demo.launch() |