Update app.py

app.py

@@ -124,7 +124,7 @@ def enhance_prompt(input_prompt, model_choice):
     
     return enhanced_text
 
-def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt):
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     
@@ -137,13 +137,14 @@ def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height,
         num_inference_steps=num_inference_steps, 
         width=width, 
         height=height,
+        num_images_per_prompt=num_images_per_prompt,
         generator=generator
     ).images[0]
     
     return image, seed
 
 @spaces.GPU
-def process_workflow(image, text_prompt, vlm_model_choice, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps):
+def process_workflow(image, text_prompt, vlm_model_choice, use_enhancer, model_choice, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt):
     if image is not None:
         # Convert image to PIL if it's not already
         if not isinstance(image, Image.Image):
@@ -159,7 +160,7 @@ def process_workflow(image, text_prompt, vlm_model_choice, use_enhancer, model_c
     if use_enhancer:
         prompt = enhance_prompt(prompt, model_choice)
     
-    generated_image, used_seed = generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps)
+    generated_image, used_seed = generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images_per_prompt)
     
     return generated_image, prompt, used_seed
 
@@ -197,7 +198,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
         with gr.Column(scale=1):
             with gr.Group(elem_classes="input-group"):
                 input_image = gr.Image(label="Input Image (VLM Captioner)")
-                vlm_model_choice = gr.Radio(["Long Captioner", "Florence"], label="VLM Model", value="Long Captioner")
+                vlm_model_choice = gr.Radio(["Florence-2", "Long Captioner"], label="VLM Model", value="Florence-2")
             
             with gr.Accordion("Advanced Settings", open=False):
                 text_prompt = gr.Textbox(label="Text Prompt (optional, used if no image is uploaded)")
@@ -210,6 +211,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
                 height = gr.Slider(label="Height", minimum=512, maximum=2048, step=64, value=1024)
                 guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=5.0)
                 num_inference_steps = gr.Slider(label="Inference Steps", minimum=20, maximum=50, step=1, value=20)
+                num_images_per_prompt = gr.Slider(1, 4, 1, step=1, label="Number of images per prompt")
             
             generate_btn = gr.Button("Generate Image", elem_classes="submit-btn")
         
@@ -223,7 +225,8 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="blue", secondar
         fn=process_workflow,
         inputs=[
             input_image, text_prompt, vlm_model_choice, use_enhancer, model_choice,
-            negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps
+            negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps,
+            num_images_per_prompt
         ],
         outputs=[output_image, final_prompt, used_seed]
     )