stabe-diffusion / app.py
otmanheddouch's picture
Update app.py
bbcd3bc verified
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from torch import autocast
def generate_image(prompt, model_id, guidance_scale):
"""Generate image"""
img_path = "image.png"
with autocast(device):
img = pipes[model_id](prompt, guidance_scale=guidance_scale)["images"][0]
img.save(img_path)
return img_path
# the model id
model_ids = ["CompVis/stable-diffusion-v1-4"]
# check for device
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(f"The code is runing on {device}")
# load the p model pipeline
pipes = {
model_id:StableDiffusionPipeline.from_pretrained(model_id, variant="fp16", torch_dtype=torch.float32).to(device)
for model_id in model_ids
}
example_prompt = "a photo of an astronaut riding a horse on mars"
# gradio interface
interface = gr.Interface(
fn=generate_image,
inputs= [gr.Textbox(label="Prompt", value=example_prompt),
gr.Dropdown(model_ids,label="Select the model ID"),
gr.Slider(minimum=0.1,maximum=10, value=7.5, label="Guidance Scale")],
outputs="image",
live=False,
title="Text to Image",
description="Enter the prompt in text input, click generate to generete the image")
# lauch the gradio
interface.launch(share=True, debug=True) # set share to False if you dont want to make public