Text-to-Image / app.py
MUTSC's picture
Update app.py
56914ca verified
raw
history blame
697 Bytes
import gradio as gr
from diffusers import StableDiffusionPipeline
from PIL import Image
import torch
# Load the Stable Diffusion model
model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
model = model.to("cuda" if torch.cuda.is_available() else "cpu")
def generate_image(prompt):
# Generate image from the model
with torch.no_grad():
image = model(prompt).images[0]
# Convert to PIL Image to display in Gradio
image = Image.fromarray(image.numpy())
return image
# Create a Gradio interface
interface = gr.Interface(fn=generate_image, inputs='text', outputs='image')
# Launch the interface
interface.launch()