Update app.py

app.py

@@ -6,29 +6,29 @@ from diffusers import StableCascadeDecoderPipeline, StableCascadePriorPipeline
 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):
+    prior_pipeline.enable_model_cpu_offload()
     prior_output = prior_pipeline(
         prompt=prompt,
+        negative_prompt=negative_prompt,
+        width=width,
         height=height,
         guidance_scale=4.0,
         num_images_per_prompt=1,
-        width=width,
-        num_inference_steps=steps,
-        negative_prompt=negative_prompt
+        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,
         guidance_scale=0.0,
         output_type="pil",
-        num_inference_steps=steps,
+        num_inference_steps=10,
         negative_prompt=negative_prompt
     ).images[0]
-    return torch.clamp(decoder_output, 0, 1)
+    return decoder_output
 
 with gr.Blocks() as demo:
     with gr.Row():