Update stable_cascade.py

stable_cascade.py

@@ -7,14 +7,14 @@ import gradio as gr
 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.enable_model_cpu_offload()
+decoder.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.prior = torch.compile(prior_pipeline.prior, mode="reduce-overhead", fullgraph=True)
+decoder.decoder = torch.compile(decoder_pipeline.decoder, mode="max-autotune", fullgraph=True)
 
-prior_pipeline.to("cuda")
-decoder_pipeline.to("cuda")
+prior.to("cuda")
+decoder.to("cuda")
 
 def generate_images(
     prompt="a photo of a girl",
@@ -22,6 +22,7 @@ def generate_images(
     height=1024,
     width=1024, 
     guidance_scale=4.0, 
+    num_images_per_prompt=1,
     prior_inference_steps=20, 
     decoder_inference_steps=10
     ):
@@ -46,7 +47,7 @@ def generate_images(
         width=width,
         negative_prompt=negative_prompt,
         guidance_scale=guidance_scale,
-        num_images_per_prompt=1,
+        num_images_per_prompt=num_images_per_prompt,
         num_inference_steps=prior_inference_steps
     )
 
@@ -80,6 +81,14 @@ def web_demo():
                 )
                 with gr.Row():
                     with gr.Column():
+                        text2image_num_images_per_prompt = gr.Slider(
+                            minimum=1,
+                            maximum=4,
+                            step=1,
+                            value=1,
+                            label="Number Image",
+                        )
+                        
                         text2image_height = gr.Slider(
                             minimum=128,
                             maximum=1280,
@@ -136,6 +145,7 @@ def web_demo():
                     text2image_height,
                     text2image_width,
                     text2image_guidance_scale,
+                    text2image_num_images_per_prompt,
                     text2image_prior_inference_step,
                     text2image_decoder_inference_step
                 ],