Update app.py

app.py

@@ -6,9 +6,11 @@ from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 prior_pipeline = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", variant="bf16", torch_dtype=torch.bfloat16)
 decoder_pipeline = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", variant="bf16", torch_dtype=torch.bfloat16)
 
+prior_pipeline.enable_model_cpu_offload()
+decoder_pipeline.enable_model_cpu_offload()
+
 @spaces.GPU
 def generate(prompt, negative_prompt, steps):
-    prior_pipeline.enable_model_cpu_offload()
     prior_output = prior_pipeline(
         prompt=prompt,
         negative_prompt=negative_prompt,
@@ -19,7 +21,6 @@ def generate(prompt, negative_prompt, steps):
         num_inference_steps=steps
     )
     
-    decoder_pipeline.enable_model_cpu_offload()
     decoder_output = decoder_pipeline(
         image_embeddings=prior_output.image_embeddings.to(torch.float16),
         prompt=prompt,