Update app.py

app.py

@@ -32,15 +32,25 @@ def genie (Model, Prompt, negative_prompt, height, width, scale, steps, seed, re
                 upscaler.enable_xformers_memory_efficient_attention()
                 upscaler = upscaler.to(device)
                 torch.cuda.empty_cache()
-                upscaled = upscaler(prompt=prompt, negative_prompt=negative_prompt, image=image, num_inference_steps=15, guidance_scale=0).images[0]
+                upscaled = upscaler(prompt=Prompt, negative_prompt=negative_prompt, image=image, num_inference_steps=15, guidance_scale=0).images[0]
                 torch.cuda.empty_cache()
                 return upscaled
             else:
                 return image
         else:
-            image = pipe(Prompt, negative_prompt=negative_prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=scale).images[0]
-            torch.cuda.empty_cache()
-            return image
+            if upscale == "Yes":
+                upscaler = DiffusionPipeline.from_pretrained("stabilityai/sd-x2-latent-upscaler", torch_dtype=torch.float16, use_safetensors=True)
+                upscaler.enable_xformers_memory_efficient_attention()
+                upscaler = upscaler.to(device)
+                torch.cuda.empty_cache()
+                upscaled = upscaler(prompt=Prompt, negative_prompt=negative_prompt, image=image, num_inference_steps=15, guidance_scale=0).images[0]
+                torch.cuda.empty_cache()
+                return upscaled
+            else:
+                
+               image = pipe(Prompt, negative_prompt=negative_prompt, height=height, width=width, num_inference_steps=steps, guidance_scale=scale).images[0]
+               torch.cuda.empty_cache()
+               return image
     
     if Model == "Anime":
         anime = DiffusionPipeline.from_pretrained("circulus/canvers-anime-v3.8.1", torch_dtype=torch.float16, safety_checker=None) if torch.cuda.is_available() else DiffusionPipeline.from_pretrained("circulus/canvers-anime-v3.8.1")