From 04503c6f132f8850751fc17fcc35ac9c98f6af88 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jukka=20Sepp=C3=A4nen?=
 <40791699+kijai@users.noreply.github.com>
Date: Wed, 6 Mar 2024 06:00:00 +0200
Subject: [PATCH] Allow decoding in sub-batches (#4)

* Allow decoding in sub-batches

* Update layered_diffusion.py
---
 layered_diffusion.py | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/layered_diffusion.py b/layered_diffusion.py
index 37ea74c..3eb219d 100644
--- a/layered_diffusion.py
+++ b/layered_diffusion.py
@@ -93,7 +93,13 @@ class LayeredDiffusionDecode:
 
     @classmethod
     def INPUT_TYPES(s):
-        return {"required": {"samples": ("LATENT",), "images": ("IMAGE",)}}
+        return {"required": 
+                {
+                "samples": ("LATENT",), 
+                "images": ("IMAGE",),
+                "per_batch": ("INT", {"default": 16, "min": 1, "max": 4096, "step": 1}),
+                },
+            }
 
     RETURN_TYPES = ("IMAGE", "MASK")
     FUNCTION = "decode"
@@ -102,7 +108,7 @@ class LayeredDiffusionDecode:
     def __init__(self) -> None:
         self.vae_transparent_decoder = None
 
-    def decode(self, samples, images):
+    def decode(self, samples, images, per_batch):    
         if self.vae_transparent_decoder is None:
             model_path = load_file_from_url(
                 url="https://huggingface.co/LayerDiffusion/layerdiffusion-v1/resolve/main/vae_transparent_decoder.safetensors",
@@ -118,14 +124,17 @@ class LayeredDiffusionDecode:
                     else torch.float32
                 ),
             )
-        latent = samples["samples"]
         pixel = images.movedim(-1, 1)  # [B, H, W, C] => [B, C, H, W]
-        pixel_with_alpha = self.vae_transparent_decoder.decode_pixel(pixel, latent)
+        decoded = []
+        for start_idx in range(0, samples["samples"].shape[0], per_batch):
+            decoded.append(self.vae_transparent_decoder.decode_pixel(pixel[start_idx:start_idx+per_batch], samples["samples"][start_idx:start_idx+per_batch]))
+        pixel_with_alpha = torch.cat(decoded, dim=0)
+        
         # [B, C, H, W] => [B, H, W, C]
         pixel_with_alpha = pixel_with_alpha.movedim(1, -1)
         image = pixel_with_alpha[..., 1:]
         alpha = pixel_with_alpha[..., 0]
-        return (image, alpha)
+        return (image, alpha,)
 
 
 class LayeredDiffusionDecodeRGBA(LayeredDiffusionDecode):