customizable images

app.py

@@ -21,7 +21,7 @@ hf_hub_download("XLabs-AI/flux-controlnet-canny", "controlnet.safetensors")
 print("downloaded!")
 
 @spaces.GPU(duration=240)
-def process_image(image, prompt, steps, use_lora, use_controlnet, use_depth, use_hed, use_ip, lora_weight, negative_image, neg_prompt, true_gs, guidance, cfg):
+def process_image(number_of_images, image, prompt, steps, use_lora, use_controlnet, use_depth, use_hed, use_ip, lora_weight, negative_image, neg_prompt, true_gs, guidance, cfg):
     from src.flux.xflux_pipeline import XFluxPipeline
     def run_xflux_pipeline(
         prompt, image, repo_id, name, device,
@@ -129,7 +129,7 @@ def process_image(image, prompt, steps, use_lora, use_controlnet, use_depth, use
           height=1024,
           timestep_to_start_cfg=cfg,
           num_steps=steps,
-          num_images_per_prompt=2,
+          num_images_per_prompt=number_of_images,
           use_lora=use_lora,
           true_gs=true_gs,
           use_ip=use_ip,
@@ -204,6 +204,7 @@ with gr.Blocks(css=custom_css) as demo:
                 use_lora = gr.Checkbox(label="Use LORA", value=True)
                 lora_weight = gr.Slider(step=0.1, minimum=0, maximum=1, value=0.7, label="Lora Weight")
                 
+                number_of_images = gr.Slider(step=1, minimum=0, maximum=4, value=2, label="Number of Images")
                 true_gs = gr.Slider(step=0.1, minimum=0, maximum=10, value=3.5, label="TrueGs")
                 guidance = gr.Slider(minimum=1, maximum=10, value=4, label="Guidance")
                 cfg = gr.Slider(minimum=1, maximum=10, value=1, label="CFG")
@@ -214,7 +215,7 @@ with gr.Blocks(css=custom_css) as demo:
         with gr.Column(scale=2, elem_classes="app"):
             output = gr.Gallery(label="Galery output", elem_classes="galery")
 
-    submit_btn.click(process_image, inputs=[input_image, prompt, steps, use_lora, controlnet, use_depth, use_hed, use_ip, lora_weight, negative_image, neg_prompt, true_gs, guidance, cfg], outputs=output)
+    submit_btn.click(process_image, inputs=[number_of_images, input_image, prompt, steps, use_lora, controlnet, use_depth, use_hed, use_ip, lora_weight, negative_image, neg_prompt, true_gs, guidance, cfg], outputs=output)
 
 if __name__ == '__main__':
     demo.launch(share=True, debug=True)