From 610c52c0b86b80801b02ae10fbb01351e1aed4c5 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Mon, 19 Jan 2026 12:00:54 +0100 Subject: [PATCH] Update parallel_loader.py --- parallel_loader.py | 31 +++++++++++++++---------------- 1 file changed, 15 insertions(+), 16 deletions(-) diff --git a/parallel_loader.py b/parallel_loader.py index d9fa4da..1019495 100644 --- a/parallel_loader.py +++ b/parallel_loader.py @@ -19,14 +19,13 @@ class ParallelSharpnessLoader: "frame_scan_step": ("INT", {"default": 5, "min": 1, "step": 1, "label": "Analyze Every Nth Frame"}), "return_count": ("INT", {"default": 4, "min": 1, "max": 1024, "step": 1, "label": "Best Frames to Return"}), "min_distance": ("INT", {"default": 24, "min": 0, "max": 10000, "step": 1, "label": "Min Distance (Frames)"}), - - # MANUAL OFFSET (Optional: e.g. skip the first 2000 frames always) "manual_skip_start": ("INT", {"default": 0, "min": 0, "max": 10000000, "step": 1, "label": "Global Start Offset"}), }, } - RETURN_TYPES = ("IMAGE", "STRING", "INT") - RETURN_NAMES = ("images", "scores_info", "current_batch_index") + # Added a 4th output: STRING (The status sentence) + RETURN_TYPES = ("IMAGE", "STRING", "INT", "STRING") + RETURN_NAMES = ("images", "scores_info", "batch_int", "batch_status") FUNCTION = "load_video" CATEGORY = "BetaHelper/Video" @@ -42,19 +41,21 @@ class ParallelSharpnessLoader: if not os.path.exists(video_path): raise FileNotFoundError(f"Video not found: {video_path}") - # 2. Calculate Actual Start Frame - # Formula: (Batch Number * Frames Per Batch) + Global Offset + # 2. Calculate Offsets current_skip = (batch_index * scan_limit) + manual_skip_start + range_end = current_skip + scan_limit - print(f"xx- Parallel Loader | Batch: {batch_index} | Start Frame: {current_skip} | Range: {current_skip} -> {current_skip + scan_limit}") + # Create the Status String + status_msg = f"Batch {batch_index}: Skipped {current_skip} frames. Scanning range {current_skip} -> {range_end}." + print(f"xx- Parallel Loader | {status_msg}") cap = cv2.VideoCapture(video_path) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) if current_skip >= total_frames: print("xx- End of video reached.") - # Return a black frame to prevent crashing, or handle as you wish - return (torch.zeros((1, 64, 64, 3)), "End of Video", batch_index) + empty_img = torch.zeros((1, 64, 64, 3)) + return (empty_img, "End of Video", batch_index, "End of Video Reached") # 3. Scanning (Pass 1) if current_skip > 0: @@ -68,14 +69,12 @@ class ParallelSharpnessLoader: futures = [] while True: - # Stop if we hit the batch limit if scanned_count >= scan_limit: break ret, frame = cap.read() - if not ret: break # End of file + if not ret: break - # Submit to thread future = executor.submit(self.calculate_sharpness, frame) futures.append((current_frame, future)) scanned_count += 1 @@ -95,7 +94,7 @@ class ParallelSharpnessLoader: # 4. Selection if not frame_scores: - return (torch.zeros((1, 64, 64, 3)), "No frames in batch", batch_index) + return (torch.zeros((1, 64, 64, 3)), "No frames in batch", batch_index, status_msg + " (No frames found)") frame_scores.sort(key=lambda x: x[1], reverse=True) selected = [] @@ -106,7 +105,6 @@ class ParallelSharpnessLoader: selected.append((idx, score)) selected.sort(key=lambda x: x[0]) - print(f"xx- Selected Frames: {[f[0] for f in selected]}") # 5. Extraction (Pass 2) cap = cv2.VideoCapture(video_path) @@ -125,6 +123,7 @@ class ParallelSharpnessLoader: cap.release() if not output_tensors: - return (torch.zeros((1, 64, 64, 3)), "Extraction Failed", batch_index) + return (torch.zeros((1, 64, 64, 3)), "Extraction Failed", batch_index, status_msg) - return (torch.stack(output_tensors), ", ".join(info_log), batch_index) \ No newline at end of file + # Return all 4 outputs + return (torch.stack(output_tensors), ", ".join(info_log), batch_index, status_msg) \ No newline at end of file