Spaces:
Runtime error
Runtime error
| 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 |