From 407e68816d81ebd062b861e7d124ed90fa391478 Mon Sep 17 00:00:00 2001
From: robertvoy <hello@robertvoy.com>
Date: Thu, 31 Jul 2025 13:37:59 +1000
Subject: [PATCH] Add index preservation to fix distributed batch ordering

---
 distributed.py | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/distributed.py b/distributed.py
index c1d5c55..18d1353 100644
--- a/distributed.py
+++ b/distributed.py
@@ -1740,6 +1740,9 @@ async def job_complete_endpoint(request):
             image_data_list.sort(key=lambda x: x[0])
             tensors = [tensor for _, tensor in image_data_list]
             
+            # Keep the indices for later use
+            indices = [idx for idx, _ in image_data_list]
+            
             # Validate final order
             if metadata:
                 debug_log(f"Reordered {len(tensors)} images based on metadata indices")
@@ -1768,10 +1771,11 @@ async def job_complete_endpoint(request):
             debug_log(f"Current pending jobs: {list(prompt_server.distributed_pending_jobs.keys())}")
             if multi_job_id in prompt_server.distributed_pending_jobs:
                 if batch_size > 0:
-                    # Put batch as single item
+                    # Put batch as single item with indices
                     await prompt_server.distributed_pending_jobs[multi_job_id].put({
                         'worker_id': worker_id,
                         'tensors': tensors,
+                        'indices': indices,
                         'is_last': is_last
                     })
                     debug_log(f"Received batch result for job {multi_job_id} from worker {worker_id}, size={len(tensors)}")
@@ -1934,6 +1938,7 @@ class DistributedCollectorNode:
                     
                     # Check if batch mode
                     tensors = result.get('tensors', [])
+                    indices = result.get('indices', [])  # Get the indices
                     if tensors:
                         # Batch mode
                         debug_log(f"Master - Got batch from worker {worker_id}, size={len(tensors)}, is_last={is_last}")
@@ -1941,9 +1946,15 @@ class DistributedCollectorNode:
                         if worker_id not in worker_images:
                             worker_images[worker_id] = {}
                         
-                        # Add all tensors with sequential indices
-                        for idx, tensor in enumerate(tensors):
-                            worker_images[worker_id][idx] = tensor
+                        # Use actual indices if available, otherwise fall back to sequential
+                        if indices:
+                            for i, tensor in enumerate(tensors):
+                                actual_idx = indices[i]
+                                worker_images[worker_id][actual_idx] = tensor
+                        else:
+                            # Fallback for backward compatibility
+                            for idx, tensor in enumerate(tensors):
+                                worker_images[worker_id][idx] = tensor
                             
                         
                         collected_count += len(tensors)