Spaces:
Running
Running
from diffusers import StableDiffusionPipeline | |
import torch | |
import gradio as gr | |
# Load the model | |
model_id = "CompVis/stable-diffusion-v-1-4-original" | |
pipe = StableDiffusionPipeline.from_pretrained(model_id) | |
pipe = pipe.to("cuda") # Use GPU if available | |
def generate_image(prompt): | |
image = pipe(prompt).images[0] | |
return image | |
# Define the chatbot function | |
def chatbot(prompt): | |
image = generate_image(prompt) | |
return image | |
# Create the Gradio interface | |
interface = gr.Interface( | |
fn=chatbot, | |
inputs="text", | |
outputs="image", | |
title="Text to Image Chatbot", | |
description="Generate images from text using Stable Diffusion" | |
) | |
# Launch the interface | |
interface.launch() |