Manjushri commited on
Commit
8763fae
·
1 Parent(s): 08025ce

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +47 -0
app.py ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from torchvision.transforms import ToPILImage, ToTensor
4
+ from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionUpscalePipeline
5
+
6
+ device = "cuda" if torch.cuda.is_available() else "cpu"
7
+
8
+ # Define the models
9
+ model_2x = "stabilityai/sd-x2-latent-upscaler"
10
+ model_4x = "stabilityai/stable-diffusion-x4-upscaler"
11
+ # Load the models
12
+ sd_2_0_2x = .from_pretrained(model_2x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x)
13
+ sd_2_1_4x = StableDiffusionUpscalePipeline.from_pretrained(model_4x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionUpscalePipeline.from_pretrained(model_4x)
14
+
15
+ # Define the input and output components for the Gradio interface
16
+ input_image = gr.inputs.Image()
17
+ output_image = gr.outputs.Image()
18
+
19
+ # Define the function that will be called when the interface is used
20
+ def upscale_image(model, image):
21
+ # Convert the image to a PyTorch tensor
22
+ image_tensor = ToTensor()(image)
23
+
24
+ # Upscale the image using the selected model
25
+ if model == "SD 2.0 2x Latent Upscaler":
26
+ upscaled_tensor = sd_2_0_2x(image_tensor)
27
+ else:
28
+ upscaled_tensor = sd_2_1_4x(image_tensor)
29
+
30
+ # Convert the upscaled tensor back to a PIL image
31
+ upscaled_image = ToPILImage()(upscaled_tensor)
32
+
33
+ # Return the upscaled image
34
+ return upscaled_image
35
+
36
+ # Define the Gradio interface
37
+ iface = gr.Interface(
38
+ fn=upscale_image,
39
+ inputs=["radio", input_image],
40
+ outputs=output_image,
41
+ radio_options=["SD 2.0 2x Latent Upscaler", "SD 2.1 4x Upscaler"],
42
+ title="Image Upscaler",
43
+ description="Upscale an image using either the SD 2.0 2x Latent Upscaler or the SD 2.1 4x Upscaler."
44
+ )
45
+
46
+ # Launch the interface
47
+ iface.launch(debug=True)