From 69f3af35caef763efb05b00a829146729a9d9385 Mon Sep 17 00:00:00 2001
From: Regi Ellis <regi@bynine.io>
Date: Thu, 26 Jun 2025 11:37:32 -0400
Subject: [PATCH] Hotfix for OOM in Batch Node

---
 .gitignore         |  3 ++-
 pyproject.toml     |  2 +-
 src/__init__.py    | 44 +++++++++++++++++++++++---------------------
 src/__version__.py |  2 +-
 4 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5906c96..2290103 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,4 +28,5 @@ backup
 node.zip
 CLAUDE.md
 .claude
-last_session.yaml
\ No newline at end of file
+last_session.yaml
+ideas.yaml
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 2325538..6e26bc7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,7 +2,7 @@
 [project]
 name = "ComfyUI-EasyColorCorrector"
 description = "ComfyUI custom node for flexible and efficient image color correction and post-processing."
-version = "1.2.3"
+version = "1.2.4"
 
 license = {text = "MIT"}
 dependencies  = [
diff --git a/src/__init__.py b/src/__init__.py
index 7d6a2d8..cfffb4a 100644
--- a/src/__init__.py
+++ b/src/__init__.py
@@ -2283,11 +2283,15 @@ class BatchColorCorrection:
         )
         print(f"📊 Processing in batches of {frames_per_batch} frames")
 
-        # Prepare output containers on GPU
-        processed_frames = []
+        # Pre-allocate output tensor to avoid OOM from torch.cat
+        _, height, width, channels = images.shape
+        final_images = torch.zeros((total_frames, height, width, channels), device=device, dtype=images.dtype)
+        
+        # Lists for metadata (much smaller memory footprint)
         all_palette_data = []
         all_histograms = []
         all_palette_images = []
+        processed_count = 0
 
         # Process frames in GPU-optimized batches with interruption support
         try:
@@ -2298,18 +2302,17 @@ class BatchColorCorrection:
                 if model_management.interrupt_processing:
                     print("🛑 Batch processing interrupted by user")
                     # Return partially processed results
-                    if processed_frames:
-                        partial_images = torch.cat(processed_frames, dim=0)
-                        partial_count = partial_images.shape[0]
+                    if processed_count > 0:
+                        partial_images = final_images[:processed_count]
                         print(
-                            f"⚠️ Partial result: {partial_count}/{total_frames} frames processed"
+                            f"⚠️ Partial result: {processed_count}/{total_frames} frames processed"
                         )
                         return (
                             partial_images,
                             "",
                             torch.zeros((1, 512, 768, 3), device=device),
                             torch.zeros((1, 120, 600, 3), device=device),
-                            partial_count,
+                            processed_count,
                         )
                     else:
                         return (
@@ -2354,7 +2357,9 @@ class BatchColorCorrection:
                     device=device,
                 )
 
-                processed_frames.append(batch_processed)
+                # Copy batch results directly into pre-allocated tensor
+                final_images[batch_start:batch_end] = batch_processed
+                processed_count = batch_end
 
                 # Only extract palette from middle frame to avoid CPU overhead
                 if extract_palette and batch_start <= total_frames // 2 < batch_end:
@@ -2374,18 +2379,17 @@ class BatchColorCorrection:
 
         except KeyboardInterrupt:
             print("🛑 Batch processing interrupted by KeyboardInterrupt")
-            if processed_frames:
-                partial_images = torch.cat(processed_frames, dim=0)
-                partial_count = partial_images.shape[0]
+            if processed_count > 0:
+                partial_images = final_images[:processed_count]
                 print(
-                    f"⚠️ Partial result: {partial_count}/{total_frames} frames processed"
+                    f"⚠️ Partial result: {processed_count}/{total_frames} frames processed"
                 )
                 return (
                     partial_images,
                     "",
                     torch.zeros((1, 512, 768, 3), device=device),
                     torch.zeros((1, 120, 600, 3), device=device),
-                    partial_count,
+                    processed_count,
                 )
             else:
                 return (
@@ -2397,18 +2401,17 @@ class BatchColorCorrection:
                 )
         except Exception as e:
             print(f"❌ Error during batch processing: {e}")
-            if processed_frames:
-                partial_images = torch.cat(processed_frames, dim=0)
-                partial_count = partial_images.shape[0]
+            if processed_count > 0:
+                partial_images = final_images[:processed_count]
                 print(
-                    f"⚠️ Partial result after error: {partial_count}/{total_frames} frames processed"
+                    f"⚠️ Partial result after error: {processed_count}/{total_frames} frames processed"
                 )
                 return (
                     partial_images,
                     "",
                     torch.zeros((1, 512, 768, 3), device=device),
                     torch.zeros((1, 120, 600, 3), device=device),
-                    partial_count,
+                    processed_count,
                 )
             else:
                 return (
@@ -2419,9 +2422,8 @@ class BatchColorCorrection:
                     0,
                 )
 
-        # Combine all processed frames on GPU
-        if processed_frames:
-            final_images = torch.cat(processed_frames, dim=0)
+        # Processing completed successfully
+        if processed_count == total_frames:
 
             # Use middle frame data for representation
             representative_palette = all_palette_data[0] if all_palette_data else ""
diff --git a/src/__version__.py b/src/__version__.py
index 10aa336..b3f9ac7 100644
--- a/src/__version__.py
+++ b/src/__version__.py
@@ -1 +1 @@
-__version__ = "1.2.3"
+__version__ = "1.2.4"