added upsacle function to upscale outputs automatically

app.py

@@ -1,4 +1,4 @@
-from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler
+from diffusers import AutoencoderKL, UNet2DConditionModel, StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, DPMSolverMultistepScheduler, DiffusionPipeline
 import gradio as gr
 import torch
 from PIL import Image
@@ -14,6 +14,12 @@ is_colab = utils.is_google_colab()
 state = None
 current_steps = 25
 
+UPSCALER = DiffusionPipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
+if torch.cuda.is_available():
+    UPSCALER.to("cuda")
+    UPSCALER.enable_xformers_memory_efficient_attention()
+
+
 class Model:
     def __init__(self, name, path="", prefix=""):
         self.name = name
@@ -119,7 +125,6 @@ models = [
      Model("Funko-Diffusion", "prompthero/funko-diffusion", "funko style"),
      Model("DreamShaper", "Lykon/DreamShaper", "dreamshaper"),
      Model("Realistic_Vision_V1.4", "SG161222/Realistic_Vision_V1.4", ""),
-     Model("upscalerx4","stabilityai/stable-diffusion-x4-upscaler","")
 
 
 
@@ -260,7 +265,7 @@ def txt_to_img(model_path, prompt, n_images, neg_prompt, guidance, steps, width,
 
     # update_state(f"Done. Seed: {seed}")
     
-    return replace_nsfw_images(result)
+    return replace_nsfw_images(upscale(result, prompt, neg_prompt, generator))
 
 def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance, steps, width, height, generator, seed):
 
@@ -313,7 +318,7 @@ def img_to_img(model_path, prompt, n_images, neg_prompt, img, strength, guidance
 
     # update_state(f"Done. Seed: {seed}")
         
-    return replace_nsfw_images(result)
+    return replace_nsfw_images(upscale(result, prompt, neg_prompt, generator))
 
 def replace_nsfw_images(results):
 
@@ -325,6 +330,11 @@ def replace_nsfw_images(results):
         results.images[i] = Image.open("nsfw.png")
     return results.images
 
+def upscale(results, prompt, neg_prompt, generator):
+   for i in range(len(results.images)):
+         results.images.append(UPSCALER(prompt=prompt, negative_prompt=neg_prompt, image=results.images[i], num_inference_steps=20, guidance_scale=0, generator=generator).images[0])
+   return results
+    
 # css = """.finetuned-diffusion-div div{display:inline-flex;align-items:center;gap:.8rem;font-size:1.75rem}.finetuned-diffusion-div div h1{font-weight:900;margin-bottom:7px}.finetuned-diffusion-div p{margin-bottom:10px;font-size:94%}a{text-decoration:underline}.tabs{margin-top:0;margin-bottom:0}#gallery{min-height:20rem}
 # """
 with gr.Blocks(css="style.css") as demo: