Update app.py

app.py

@@ -17,9 +17,12 @@ db_diffusion_model = from_pretrained_keras("kadirnar/dreambooth_diffusion_model_
 sd_dreambooth_model._diffusion_model = db_diffusion_model
 
 # generate images
-def infer(prompt,negative_prompt):
+def infer(prompt,negative_prompt,batch_size):
     generated_images = sd_dreambooth_model.text_to_image(
-        prompt,negative_prompt=negative_prompt
+        prompt,
+        negative_prompt=negative_prompt,
+        batch_size=batch_size
+        
     )
     return generated_images 
     
@@ -27,6 +30,7 @@ output = gr.Gallery(show_label=False)
 inputs = [
     gr.Textbox(lines=1, value="sks traditional furniture", label="Base Prompt"),
     gr.Textbox(lines=1, value="deformed", label="Negative Prompt"),
+    gr.Slider(minimum=1, maximum=10, default=1, step=1, label="Number of Image")
 
 ]
 # customize interface