Update app.py

app.py

@@ -1,5 +1,5 @@
 # Thank AK. https://huggingface.co/spaces/akhaliq/cool-japan-diffusion-2-1-0/blob/main/app.py
-from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler
+from diffusers import StableDiffusionPipeline, StableDiffusionImg2ImgPipeline, EulerAncestralDiscreteScheduler,StableDiffusionLatentUpscalePipeline
 from transformers import CLIPFeatureExtractor
 import gradio as gr
 import torch
@@ -27,6 +27,10 @@ pipe_i2i = StableDiffusionImg2ImgPipeline.from_pretrained(
 )
 pipe_i2i.enable_xformers_memory_efficient_attention()
 
+upscaler = StableDiffusionLatentUpscalePipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16)
+upscaler.enable_xformers_memory_efficient_attention()
+upscaler.to("cuda")
+
 if torch.cuda.is_available():
   pipe = pipe.to("cuda")
   pipe_i2i = pipe_i2i.to("cuda")
@@ -45,27 +49,36 @@ def inference(prompt, guidance, steps, image_size="Square", seed=0, img=None, st
   if(image_size=="Portrait"):
       height=1024
       width=768
+      superreso=False
       #pipe.enable_attention_slicing()
   elif(image_size=="Landscape"):
       height=768
       width=1024  
+      superreso=False
       #pipe.enable_attention_slicing()
   elif(image_size=="Highreso."):
       height=1024
-      width=1024  
+      width=1024
+      superreso=False
+      #pipe.enable_attention_slicing()
+  elif(image_size=="Superreso."):
+      height=1024
+      width=1024
+      superreso=True
       #pipe.enable_attention_slicing()
   else:
       height=768
       width=768
+      superreso=False
       #pipe.enable_attention_slicing()
 
   print(prompt,neg_prompt)
     
   try:
     if img is not None:
-      return img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator), None
+      return img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator,superreso), None
     else:
-      return txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator), None
+      return txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator,superreso), None
   except Exception as e:
     return None, error_str(e)
 def auto_prompt_correction(prompt_ui,neg_prompt_ui,cool_japan_type_ui,disable_auto_prompt_correction):
@@ -123,22 +136,65 @@ def auto_prompt_correction(prompt_ui,neg_prompt_ui,cool_japan_type_ui,disable_au
 
     return prompt,neg_prompt
     
-def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator):
-    result = pipe(
-      prompt,
-      negative_prompt = neg_prompt,
-      num_inference_steps = int(steps),
-      guidance_scale = guidance,
-      width = width,
-      height = height,
-      generator = generator)
+def txt_to_img(prompt, neg_prompt, guidance, steps, width, height, generator,superreso=False):
+    if(superreso):
+      low_res_latents = pipe(
+        prompt,
+        negative_prompt = neg_prompt,
+        num_inference_steps = int(steps),
+        guidance_scale = guidance,
+        width = width,
+        height = height,
+        output_type="latent",
+        generator = generator)
+      low_res_latents = pipeline(prompt, generator=generator, output_type="latent").images
+
+      result = upscaler(
+        prompt=prompt,
+        negative_prompt = neg_prompt,
+        image=low_res_latents,
+        num_inference_steps=20,
+        guidance_scale=guidance,
+        generator=generator,
+      )
+    else:  
+      result = pipe(
+        prompt,
+        negative_prompt = neg_prompt,
+        num_inference_steps = int(steps),
+        guidance_scale = guidance,
+        width = width,
+        height = height,
+        generator = generator)
     
     return result.images[0]
 
-def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator):
+def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height, generator,superreso=False):
     ratio = min(height / img.height, width / img.width)
     img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
-    result = pipe_i2i(
+    if(superreso):
+      low_res_latents = pipe_i2i(
+        prompt,
+        negative_prompt = neg_prompt,
+        init_image = img,
+        num_inference_steps = int(steps),
+        strength = strength,
+        guidance_scale = guidance,
+        #width = width,
+        #height = height,
+        output_type="latent",
+        generator = generator)
+
+      result = upscaler(
+        prompt=prompt,
+        negative_prompt = neg_prompt,
+        image=low_res_latents,
+        num_inference_steps=20,
+        guidance_scale=guidance,
+        generator=generator,
+      )
+    else:  
+      result = pipe_i2i(
         prompt,
         negative_prompt = neg_prompt,
         init_image = img,
@@ -148,6 +204,7 @@ def img_to_img(prompt, neg_prompt, img, strength, guidance, steps, width, height
         #width = width,
         #height = height,
         generator = generator)
+
         
     return result.images[0]
 
@@ -199,7 +256,7 @@ with gr.Blocks(css=css) as demo:
               neg_prompt = gr.Textbox(label="Negative prompt", placeholder="What to exclude from the image")
               disable_auto_prompt_correction = gr.Checkbox(label="Disable auto prompt corretion.")
               with gr.Row():
-                image_size=gr.Radio(["Portrait","Landscape","Square","Highreso."])
+                image_size=gr.Radio(["Portrait","Landscape","Square","Highreso.","Superreso."])
                 image_size.show_label=False
                 image_size.value="Square"