Update app.py

app.py

@@ -3,8 +3,11 @@ import spaces
 import gradio as gr
 from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 
-prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.float16).to("cuda")
-decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16).to("cuda")
+prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.float16)
+decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.float16)
+
+prior_pipeline.enable_model_cpu_offload()
+decoder_pipeline.enable_model_cpu_offload()
 
 @spaces.GPU
 def generate(prompt, negative_prompt, width, height, steps):