Update stable_cascade.py

stable_cascade.py

@@ -5,7 +5,10 @@ import gradio as gr
 
 # Initialize the prior and decoder pipelines
 prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to("cuda")
+prior.enable_xformers_memory_efficient_attention()
+
 decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float16).to("cuda")
+decoder.enable_xformers_memory_efficient_attention()
 
 def generate_images(
     prompt="a photo of a girl",
@@ -53,7 +56,7 @@ def generate_images(
         num_inference_steps=decoder_inference_steps
     ).images
 
-    return decoder_output[0]
+    return decoder_output
 
 
 def web_demo():