Update app.py

app.py

@@ -275,7 +275,7 @@ def frames_to_video(input_folder, output_path, fps, original_video_path):
     os.remove(temp_output_path)
 
 @timer_func
-def process_video(input_video, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames=None, frame_interval=1, preserve_frames=False, progress=gr.Progress()):
+def process_video(input_video, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames=None, frame_interval=1, preserve_frames=False, batch_size=4, progress=gr.Progress()):
     abort_event.clear()  # Clear the abort flag at the start of a new job
     print("Starting video processing...")
     model_manager.load_models(progress)  # Ensure models are loaded
@@ -287,14 +287,15 @@ def process_video(input_video, resolution, num_inference_steps, strength, hdr, g
 
     # Save job config
     config = {
-        "resolution": resolution,
-        "num_inference_steps": num_inference_steps,
-        "strength": strength,
-        "hdr": hdr,
-        "guidance_scale": guidance_scale,
-        "max_frames": max_frames,
-        "frame_interval": frame_interval,
-        "preserve_frames": preserve_frames
+    "resolution": resolution,
+    "num_inference_steps": num_inference_steps,
+    "strength": strength,
+    "hdr": hdr,
+    "guidance_scale": guidance_scale,
+    "max_frames": max_frames,
+    "frame_interval": frame_interval,
+    "preserve_frames": preserve_frames,
+    "batch_size": batch_size
     }
     with open(os.path.join(job_folder, "config.json"), "w") as f:
         json.dump(config, f)
@@ -320,7 +321,6 @@ def process_video(input_video, resolution, num_inference_steps, strength, hdr, g
 
     try:
         progress(0.2, desc="Processing frames...")
-        batch_size = 8
         for i in tqdm(range(0, frames_to_process, batch_size), desc="Processing batches"):
             if abort_event.is_set():
                 print("Job aborted. Stopping processing of new frames.")
@@ -359,7 +359,7 @@ def process_video(input_video, resolution, num_inference_steps, strength, hdr, g
         return None
 
 @spaces.GPU(duration=400)
-def gradio_process_media(input_media, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, progress=gr.Progress()):
+def gradio_process_media(input_media, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, batch_size, progress=gr.Progress()):
     abort_event.clear()  # Clear the abort flag at the start of a new job
     if input_media is None:
         return None, "No input media provided."
@@ -385,7 +385,7 @@ def gradio_process_media(input_media, resolution, num_inference_steps, strength,
     video_extensions = ('.mp4', '.avi', '.mov', '.mkv')
     if file_path.lower().endswith(video_extensions):
         print("Processing video...")
-        result = process_video(file_path, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, progress)
+        result = process_video(file_path, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, batch_size, progress)
         if result:
             return result, "Video processing completed successfully."
         else:
@@ -427,6 +427,7 @@ with gr.Blocks(css=css, theme=gr.themes.Default(primary_hue="blue")) as iface:
             max_frames = gr.Number(label="Max Frames to Process (leave empty for full video)", precision=0)
             frame_interval = gr.Slider(1, 30, 1, step=1, label="Frame Interval (process every nth frame)")
             preserve_frames = gr.Checkbox(label="Preserve Existing Processed Frames", value=True)
+            batch_size = gr.Slider(1, 16, 1, step=1, label="Batch Size")
 
         with gr.Column(scale=1):
             submit_button = gr.Button("Process Media")
@@ -436,7 +437,7 @@ with gr.Blocks(css=css, theme=gr.themes.Default(primary_hue="blue")) as iface:
 
     submit_button.click(
         gradio_process_media,
-        inputs=[input_media, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames],
+        inputs=[input_media, resolution, num_inference_steps, strength, hdr, guidance_scale, max_frames, frame_interval, preserve_frames, batch_size],
         outputs=[output, status]
     )