otmanheddouch commited on
Commit
af6ce36
·
verified ·
1 Parent(s): 3bc946a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +36 -0
app.py ADDED
@@ -0,0 +1,36 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from diffusers import StableDiffusionPipeline
4
+ from torch import autocast
5
+ def generate_image(prompt, model_id, guidance_scale):
6
+ """Generate image"""
7
+ img_path = "image.png"
8
+ with autocast(device):
9
+ img = pipes[model_id](prompt, guidance_scale=guidance_scale)["images"][0]
10
+ img.save(img_path)
11
+ return img_path
12
+
13
+ # the model id
14
+ model_ids = ["CompVis/stable-diffusion-v1-4"]
15
+ # check for device
16
+ device = 'cuda' if torch.cuda.is_available() else 'cpu'
17
+ print(f"The code is runing on {device}")
18
+ # load the p model pipeline
19
+ pipes = {
20
+ model_id:StableDiffusionPipeline.from_pretrained(model_id, variant="fp16", torch_dtype=torch.float32).to(device)
21
+ for model_id in model_ids
22
+ }
23
+
24
+ example_prompt = "a photo of an astronaut riding a horse on mars"
25
+ # gradio interface
26
+ interface = gr.Interface(
27
+ fn=generate_image,
28
+ inputs= [gr.Textbox(label="Prompt", default=example_prompt),
29
+ gr.Dropdown(model_ids,label="Select the model ID"),
30
+ gr.Slider(minimum=0.1,maximum=10, default=7.5, label="Guidance Scale")],
31
+ outputs="image",
32
+ live=False,
33
+ title="Text to Image",
34
+ description="Enter the prompt in text input, click generate to generete the image")
35
+ # lauch the gradio
36
+ interface.launch(share=True, debug=True) # set share to False if you dont want to make public