Spaces:
Sleeping
Sleeping
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 |