Update app.py

app.py

@@ -158,30 +158,30 @@ class ModelManager:
         return pipe
 
     @timer_func
-    def process_image(self, input_image, resolution, num_inference_steps, strength, hdr, guidance_scale):
-        condition_image = self.prepare_image(input_image, resolution, hdr)
+    def process_image_batch(self, input_images, resolution, num_inference_steps, strength, hdr, guidance_scale):
+        condition_images = [self.prepare_image(img, resolution, hdr) for img in input_images]
 
         prompt = "masterpiece, best quality, highres"
         negative_prompt = "low quality, normal quality, ugly, blurry, blur, lowres, bad anatomy, bad hands, cropped, worst quality, verybadimagenegative_v1.3, JuggernautNegative-neg"
 
         options = {
-            "prompt": prompt,
-            "negative_prompt": negative_prompt,
-            "image": condition_image,
-            "control_image": condition_image,
-            "width": condition_image.size[0],
-            "height": condition_image.size[1],
+            "prompt": [prompt] * len(input_images),
+            "negative_prompt": [negative_prompt] * len(input_images),
+            "image": condition_images,
+            "control_image": condition_images,
+            "width": condition_images[0].size[0],
+            "height": condition_images[0].size[1],
             "strength": strength,
             "num_inference_steps": num_inference_steps,
             "guidance_scale": guidance_scale,
             "generator": torch.Generator(device=device).manual_seed(0),
         }
 
-        print("Running inference...")
-        result = self.pipe(**options).images[0]
-        print("Image processing completed successfully")
+        print("Running inference on batch...")
+        results = self.pipe(**options).images
+        print("Batch processing completed successfully")
 
-        return result
+        return results
 
     def prepare_image(self, input_image, resolution, hdr):
         condition_image = self.resize_and_upscale(input_image, resolution)
@@ -320,27 +320,23 @@ def process_video(input_video, resolution, num_inference_steps, strength, hdr, g
 
     try:
         progress(0.2, desc="Processing frames...")
-        for i, frame_file in enumerate(tqdm(frame_files[:frames_to_process], 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.")
                 break
 
-            output_frame_path = os.path.join(processed_frames_folder, frame_file)
-            if not preserve_frames or not os.path.exists(output_frame_path):
-                if i % frame_interval == 0:
-                    # Process this frame
-                    input_image = Image.open(os.path.join(frames_folder, frame_file))
-                    processed_image = model_manager.process_image(input_image, resolution, num_inference_steps, strength, hdr, guidance_scale)
+            batch_frames = frame_files[i:min(i+batch_size, frames_to_process)]
+            input_images = [Image.open(os.path.join(frames_folder, frame)) for frame in batch_frames]
+            
+            processed_images = model_manager.process_image_batch(input_images, resolution, num_inference_steps, strength, hdr, guidance_scale)
+
+            for frame_file, processed_image in zip(batch_frames, processed_images):
+                output_frame_path = os.path.join(processed_frames_folder, frame_file)
+                if not preserve_frames or not os.path.exists(output_frame_path):
                     processed_image.save(output_frame_path)
-                else:
-                    # Copy the previous processed frame or the original frame
-                    prev_frame = f"frame_{int(frame_file.split('_')[1].split('.')[0]) - 1:06d}.png"
-                    prev_frame_path = os.path.join(processed_frames_folder, prev_frame)
-                    if os.path.exists(prev_frame_path):
-                        shutil.copy2(prev_frame_path, output_frame_path)
-                    else:
-                        shutil.copy2(os.path.join(frames_folder, frame_file), output_frame_path)
-            progress((0.2 + 0.7 * (i + 1) / frames_to_process), desc=f"Processing frame {i+1}/{frames_to_process}")
+                
+            progress((0.2 + 0.7 * (i + batch_size) / frames_to_process), desc=f"Processed batch {i//batch_size + 1}/{(frames_to_process-1)//batch_size + 1}")
 
         # Always attempt to reassemble video
         progress(0.9, desc="Reassembling video...")