Update app.py

app.py

@@ -1,19 +1,41 @@
 import torch #needed only for GPU
 from PIL import Image
 from io import BytesIO
-from diffusers import StableDiffusionUpscalePipeline
+import numpy as np
+from diffusers import StableDiffusionLatentUpscalePipeline, StableDiffusionUpscalePipeline
 import gradio as gr
+import modin.pandas as pd
+
 # load model for CPU or GPU
-model_id = "stabilityai/stable-diffusion-x4-upscaler"
+
+model_2x = "stabilityai/sd-x2-latent-upscaler"
+model_4x = "stabilityai/stable-diffusion-x4-upscaler"
 device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionUpscalePipeline.from_pretrained(model_id)
-pipe = pipe.to(device)
+upscaler2x = StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionLatentUpscalePipeline.from_pretrained(model_2x, safety_checker=None)
+upscaler4x = StableDiffusionUpscalePipeline.from_pretrained(model_4x, torch_dtype=torch.float16, revision="fp16") if torch.cuda.is_available() else StableDiffusionUpscalePipeline.from_pretrained(model_4x)
+
 #define interface 
-def upscale(low_res_img, prompt, negative_prompt, scale, steps):
- low_res_img = Image.open(low_res_img).convert("RGB")
- low_res_img = low_res_img.resize((128, 128))
- upscaled_image = pipe(prompt=prompt, negative_prompt=negative_prompt, image=low_res_img, guidance_scale=scale, num_inference_steps=steps).images[0]
- #upscaled_image.save("upsampled.png")
- return upscaled_image
+
+def upscale(raw_img, model, prompt, negative_prompt, scale, steps):
+    generator = torch.manual_seed(999999)
+    low_res_img = Image.open(raw_img).convert("RGB")
+    if model == "Upscaler 4x": 
+        low_res_img = low_res_img.resize((128, 128)) 
+    else: 
+        low_res_img
+    image = upscaler2x(prompt=prompt, negative_prompt=negative_prompt, image=low_res_img, guidance_scale=scale, num_inference_steps=steps).images[0] if model == "Upscaler 2x" else upscaler4x(prompt=prompt, negative_prompt=negative_prompt, image=low_res_img, guidance_scale=scale, num_inference_steps=steps).images[0]
+    return image
+    
 #launch interface
-gr.Interface(fn=upscale, inputs=[gr.Image(type='filepath', label='Low Resolution Image (less than 512x512, i.e. 128x128, 256x256, ect., ect..)'), gr.Textbox(label='Optional: Enter a Prompt to Slightly Guide the AI'), gr.Textbox(label='Experimental: Slightly influence What you do not want the AI to generate.'), gr.Slider(2, 15, 7, step=1, label='Guidance Scale: How much the AI influences the Upscaling.'), gr.Slider(10, 75, 50, step=1, label='Number of Iterations')], outputs=gr.Image(type='filepath'), title='SD 2.0 4x Upscaler', description='A 4x Low Resolution Upscaler using SD 2.0. <br>Expects a Lower than 512x512 image. <br><br><b>Warning: Images 512x512 or Higher Resolution WILL NOT BE UPSCALED and may result in Quality Loss!', article = "Code Monkey: <a href=\"https://huggingface.co/Manjushri\">Manjushri</a>").launch(max_threads=True, debug=True)
+    
+gr.Interface(fn=upscale, inputs=[
+    gr.Image(type="filepath", label='Lower Resolution Image'), 
+    gr.Radio(['Upscaler 2x','Upscaler 4x'], label="Models"),
+    gr.Textbox(label="Optional: Enter a Prompt to Slightly Guide the AI's Enhancement"), 
+    gr.Textbox(label='Experimental: Slightly influence What you do not want the AI to Enhance.'), 
+    gr.Slider(2, 15, 7, step=1, label='Guidance Scale: How much the AI influences the Upscaling.'), 
+    gr.Slider(5, 25, 10, step=1, label='Number of Iterations')], 
+    outputs=gr.Image(type="filepath", label = 'Upscaled Image'), 
+    title='SD Upscaler', 
+    description='2x Latent Upscaler using SD 2.0 And 4x Upscaler using SD 2.1. This version runs on CPU or GPU and is currently running on a T4 GPU. For 4x Upscaling use images lower than 512x512. For 2x Upscaling use 512x512 to 768x768 images.<br><br><b>Notice: Largest Accepted Resolution is 768x768', 
+    article = "Code Monkey: <a href=\"https://huggingface.co/Manjushri\">Manjushri</a>").launch(max_threads=True, debug=True)