Update stable_cascade.py

stable_cascade.py

@@ -4,8 +4,17 @@ import gradio as gr
 
 
 # Initialize the prior and decoder pipelines
-prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16).to("cuda")
-decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.float16).to("cuda")
+prior = StableCascadePriorPipeline.from_pretrained("stabilityai/stable-cascade-prior", torch_dtype=torch.bfloat16)
+decoder = StableCascadeDecoderPipeline.from_pretrained("stabilityai/stable-cascade", torch_dtype=torch.bfloat16)
+
+prior_pipeline.enable_model_cpu_offload()
+decoder_pipeline.enable_model_cpu_offload()
+
+prior_pipeline.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
+decoder_pipeline.decoder = torch.compile(decoder_pipeline.decoder, mode="max-autotune", fullgraph=True)
+
+prior_pipeline.to("cuda")
+decoder_pipeline.to("cuda")
 
 def generate_images(
     prompt="a photo of a girl",