Spaces:
Sleeping
Sleeping
File size: 1,326 Bytes
af6ce36 bbcd3bc af6ce36 bbcd3bc af6ce36 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 |
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 |