Update app.py

app.py

@@ -8,7 +8,7 @@ lora_path = "OedoSoldier/detail-tweaker-lora"
 vae = AutoencoderKL.from_pretrained("stabilityai/sd-vae-ft-mse").to("cuda")
 
 @spaces.GPU
-def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_scale=7.0,model="Real6.0"):
+def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_scale=7.0,model="Real6.0",num_images_slider=1):
         
     if model == "Real5.0":
         model_id = "SG161222/Realistic_Vision_V5.0_noVAE"
@@ -35,17 +35,18 @@ def generate_image(prompt, negative_prompt, num_inference_steps=30, guidance_sca
 
     
     # Generate the image
-    image = pipe(
+    result = pipe(
         prompt = prompt,
         negative_prompt = negative_prompt,
         cross_attention_kwargs = {"scale":1},
         num_inference_steps = num_inference_steps,
         guidance_scale = guidance_scale,
         width = 720,
-        height = 720
+        height = 720,
+        num_images_per_prompt=num_images
     ).images[0]
     
-    return image
+    return result.images
 
 title = """<h1 align="center">ProFaker</h1>"""
 # Create the Gradio interface
@@ -72,6 +73,13 @@ with gr.Blocks() as demo:
                     value="Real6.0",
                     label="Model",
                 )
+                num_images_slider = gr.Slider(  # New slider for number of images
+                        minimum=1,
+                        maximum=4,
+                        value=1,
+                        step=1,
+                        label="Number of Images to Generate"
+                )
                 steps_slider = gr.Slider(
                     minimum=1,
                     maximum=100,
@@ -88,13 +96,19 @@ with gr.Blocks() as demo:
                 )
         with gr.Column():
             # Output component
-            image_output = gr.Image(label="Generated Image")
+            gallery = gr.Gallery(
+                label="Generated Images",
+                show_label=True,
+                elem_id="gallery",
+                columns=2,
+                rows=2
+            )
             
     
     # Connect the interface to the generation function
     generate_button.click(
         fn=generate_image,
-        inputs=[prompt, negative_prompt, steps_slider, guidance_slider, model],
+        inputs=[prompt, negative_prompt, steps_slider, guidance_slider, model, num_images_slider],
         outputs=image_output
     )