Spaces:
Running
Running
import torch | |
from diffusers import StableDiffusionPipeline | |
import gradio as gr | |
# Use the CompVis stable-diffusion-v1-4 model | |
model_id = "CompVis/stable-diffusion-v1-4" | |
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32) # float32 for CPU | |
pipe.to("cpu") # Ensure it runs on CPU | |
def generate_image(prompt, pipe): | |
# Generate the image using the pipeline | |
image = pipe(prompt).images[0] | |
return image | |
def chatbot(prompt): | |
# Generate the image based on the user's input | |
image = generate_image(prompt, pipe) | |
return image | |
# Create the Gradio interface | |
interface = gr.Interface( | |
fn=chatbot, | |
inputs="text", | |
outputs="image", | |
title="Stable Diffusion v1-4 Text-to-Image Chatbot", | |
description="Enter a text prompt and get an AI-generated image." | |
) | |
# Launch the interface | |
interface.launch() | |