Update app.py

app.py

@@ -350,18 +350,27 @@ def genie (Model, Prompt, negative_prompt, height, width, scale, steps, seed, re
                torch.cuda.empty_cache()
                return image
 
-    if Model == 'Cascade':
+    if Model == 'CascadeXL':
         torch.cuda.empty_cache()
         from diffusers import StableCascadeCombinedPipeline
         pipe = StableCascadeCombinedPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.bfloat16)
         pipe.enable_xformers_memory_efficient_attention()
         pipe = pipe.to(device)
         image = pipe(prompt=Prompt, negative_prompt=negative_prompt, num_inference_steps=10, prior_num_inference_steps=20, prior_guidance_scale=3.0, width=width, height=height).images[0]
-        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=5, guidance_scale=0).images[0]
+               return upscaled
+            else:
+                return image
 
     return image
     
-gr.Interface(fn=genie, inputs=[gr.Radio(['PhotoReal', 'Anime', 'Disney', 'StoryBook', 'SemiReal', 'Animagine XL 3.0', 'SDXL 1.0', 'FusionXL', 'Cascade'], value='PhotoReal', label='Choose Model'),
+gr.Interface(fn=genie, inputs=[gr.Radio(['PhotoReal', 'Anime', 'Disney', 'StoryBook', 'SemiReal', 'Animagine XL 3.0', 'SDXL 1.0', 'FusionXL', 'CascadeXL'], value='PhotoReal', label='Choose Model'),
                                gr.Textbox(label='What you want the AI to generate. 77 Token Limit.'), 
                                gr.Textbox(label='What you Do Not want the AI to generate. 77 Token Limit'),
                                gr.Slider(512, 1280, 768, step=128, label='Height'),