refactor: improve result collection in BatchProcessor for parallel processing

batch_sample.py

@@ -86,9 +86,10 @@ class BatchProcessor:
             ]
         )
 
-        content_tensors = []
-        style_tensors = []
-        content_pil_images = []
+        # Initialize ordered lists for results
+        content_tensors = [None] * batch_size
+        style_tensors = [None] * batch_size
+        content_pil_images = [None] * batch_size
 
         # Process in parallel using ThreadPoolExecutor for I/O operations
         with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
@@ -107,37 +108,41 @@ class BatchProcessor:
                         content_input,
                         content_inference_transforms,
                     )
-                content_futures.append(future)
+                content_futures.append((i, future))
 
             # Submit style processing tasks
             style_futures = []
-            for style_input in style_inputs:
+            for i, style_input in enumerate(style_inputs):
                 future = executor.submit(
                     self._process_style_image, style_input, style_inference_transforms
                 )
-                style_futures.append(future)
+                style_futures.append((i, future))
 
-            # Collect results
-            for future in as_completed(content_futures):
+            # Collect results in order
+            for i, future in content_futures:
                 try:
                     content_tensor, content_pil = future.result()
                     if content_tensor is not None:
-                        content_tensors.append(content_tensor)
-                        content_pil_images.append(content_pil)
+                        content_tensors[i] = content_tensor
+                        content_pil_images[i] = content_pil
                 except Exception as e:
-                    print(f"Error processing content: {e}")
+                    print(f"Error processing content at index {i}: {e}")
                     continue
 
-            for future in as_completed(style_futures):
+            for i, future in style_futures:
                 try:
                     style_tensor = future.result()
                     if style_tensor is not None:
-                        style_tensors.append(style_tensor)
+                        style_tensors[i] = style_tensor
                 except Exception as e:
-                    print(f"Error processing style: {e}")
+                    print(f"Error processing style at index {i}: {e}")
                     continue
 
-        # Stack tensors into batches
+        # Filter out None values and stack tensors
+        content_tensors = [t for t in content_tensors if t is not None]
+        style_tensors = [t for t in style_tensors if t is not None]
+        content_pil_images = [img for img in content_pil_images if img is not None]
+
         if content_tensors and style_tensors:
             content_batch = torch.stack(content_tensors)
             style_batch = torch.stack(style_tensors)