wifix199's picture
Update app.py
e2f8f8a verified
raw
history blame
858 Bytes
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()