Update app.py

app.py

@@ -257,7 +257,7 @@ def preprocess_openpose(image):
     image = cv2.resize(image, (new_width, new_height))
     return Image.fromarray(image)
 
-def process_image_batch(images, pipe, prompt, negative_prompt, progress, batch_size=2):
+def process_image_batch(images, pipe, prompt, negative_prompt, num_inference_steps, progress, batch_size=2):
     all_processed_images = []
     for i in range(0, len(images), batch_size):
         batch = images[i:i+batch_size]
@@ -271,7 +271,7 @@ def process_image_batch(images, pipe, prompt, negative_prompt, progress, batch_s
                     prompt=prompt,
                     negative_prompt=negative_prompt,
                     ref_image=img,
-                    num_inference_steps=30
+                    num_inference_steps=num_inference_steps
                 ).images
                 processed_batch.extend(result)
         else:
@@ -279,7 +279,7 @@ def process_image_batch(images, pipe, prompt, negative_prompt, progress, batch_s
                 prompt=batch_prompt,
                 negative_prompt=batch_negative_prompt,
                 image=batch,
-                num_inference_steps=30
+                num_inference_steps=num_inference_steps
             ).images
 
         all_processed_images.extend(processed_batch)
@@ -288,7 +288,7 @@ def process_image_batch(images, pipe, prompt, negative_prompt, progress, batch_s
     return all_processed_images
 
 # Define the function to generate images
-def generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, control_images, progress=gr.Progress(track_tqdm=True)):
+def generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, control_images, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
     global controlnet_pipe, pipe, reference_pipe
 
     clear_memory()
@@ -322,9 +322,9 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
         for i in range(0, len(preprocessed_images), chunk_size):
             chunk = preprocessed_images[i:i+chunk_size]
             if controlnet_type == "Reference":
-                images_chunk = process_image_batch(chunk, reference_pipe, prompt, negative_prompt, progress)
+                images_chunk = process_image_batch(chunk, reference_pipe, prompt, negative_prompt, num_inference_steps, progress)
             else:
-                images_chunk = process_image_batch(chunk, controlnet_pipe, prompt, negative_prompt, progress)
+                images_chunk = process_image_batch(chunk, controlnet_pipe, prompt, negative_prompt, num_inference_steps, progress)
             images.extend(images_chunk)
             clear_memory()
 
@@ -347,7 +347,7 @@ def generate_images_with_progress(prompt, negative_prompt, batch_count, use_cont
 
         images = []
         for i in tqdm(range(batch_count), desc="Generating images"):
-            generated = pipe(prompt=[prompt], negative_prompt=[negative_prompt], num_inference_steps=30, width=1024, height=1024).images
+            generated = pipe(prompt=[prompt], negative_prompt=[negative_prompt], num_inference_steps=num_inference_steps, width=1024, height=1024).images
             images.extend(generated)
             progress((i + 1) / batch_count)  # Update progress bar
             clear_memory()  # Clear memory after each image, even in single image mode
@@ -391,6 +391,7 @@ with gr.Blocks() as demo:
                     lines=5
                 )
             batch_count = gr.Slider(minimum=1, maximum=10, step=1, label="Batch Count", value=1)
+            num_inference_steps = gr.Slider(minimum=1, maximum=50, step=1, label="Number of Inference Steps", value=30)
             use_controlnet = gr.Checkbox(label="Use ControlNet", value=False)
             controlnet_type = gr.Dropdown(choices=["Canny", "Depth", "OpenPose", "Reference"], label="ControlNet Type")
             controlnet_status = gr.Textbox(label="ControlNet Status", value="", interactive=False)
@@ -470,7 +471,7 @@ with gr.Blocks() as demo:
     folder_images_gallery.select(fn=select_folder_image, inputs=[selected_folder_images], outputs=selected_folder_images)
     clear_selection_button.click(fn=clear_selected_folder_images, inputs=[], outputs=selected_folder_images)
 
-    def generate_images_with_folder_images(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, use_control_folder, selected_folder_images, batch_images_input, progress=gr.Progress(track_tqdm=True)):
+    def generate_images_with_folder_images(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, use_control_folder, selected_folder_images, batch_images_input, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
         if mode == "Batch":
             if use_control_folder:
                 selected_images = [img[1] for img in loaded_images]
@@ -480,12 +481,12 @@ with gr.Blocks() as demo:
                 selected_images = [resize_image(Image.open(img).convert("RGB")) for img in batch_images_input]
         else:
             selected_images = [img[1] for img in selected_folder_images]
-        return generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, selected_images, progress)
+        return generate_images_with_progress(prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, selected_images, num_inference_steps, progress)
 
     generate_button = gr.Button("Generate Images")
     generate_button.click(
         generate_images_with_folder_images,
-        inputs=[prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, use_control_folder, selected_folder_images, batch_images_input],
+        inputs=[prompt, negative_prompt, batch_count, use_controlnet, controlnet_type, mode, use_control_folder, selected_folder_images, batch_images_input, num_inference_steps],
         outputs=gallery
     )