Merge pull request #328 from YacratesWyh/mygo

Latent Batch

WAS_Node_Suite.py

@@ -4715,6 +4715,53 @@ class WAS_Image_Batch:
         self._check_image_dimensions(batched_tensors, image_names)
         batched_tensors = torch.cat(batched_tensors, dim=0)
         return (batched_tensors,)
+# Latent TO BATCH
+
+class WAS_Latent_Batch:
+    def __init__(self):
+        pass
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+            },
+            "optional": {
+                "latent_a": ("LATENT",),
+                "latent_b": ("LATENT",),
+                "latent_c": ("LATENT",),
+                "latent_d": ("LATENT",),
+            },
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    RETURN_NAMES = ("latent",)
+    FUNCTION = "latent_batch"
+    CATEGORY = "WAS Suite/Latent"
+        
+    def _check_latent_dimensions(self, tensors, names):
+        dimensions = [(tensor["samples"].shape) for tensor in tensors]
+        if len(set(dimensions)) > 1:
+            mismatched_indices = [i for i, dim in enumerate(dimensions) if dim[1] != dimensions[0][1]]
+            mismatched_latents = [names[i] for i in mismatched_indices]
+            if mismatched_latents:
+                raise ValueError(f"WAS latent Batch Warning: Input latent dimensions do not match for latents: {mismatched_latents}")
+
+    def latent_batch(self, **kwargs):
+        batched_tensors = [kwargs[key] for key in kwargs if kwargs[key] is not None]
+        latent_names = [key for key in kwargs if kwargs[key] is not None]
+
+        if not batched_tensors:
+            raise ValueError("At least one input latent must be provided.")
+
+        self._check_latent_dimensions(batched_tensors, latent_names)
+        samples_out = {} 
+        samples_out["samples"]  = torch.cat([tensor["samples"] for tensor in batched_tensors], dim=0)            
+        samples_out["batch_index"] = []
+        for tensor in batched_tensors:
+            cindex = tensor.get("batch_index", list(range(tensor["samples"].shape[0])))
+            samples_out["batch_index"] += cindex        
+        return (samples_out,)
 
 
 # MASK TO BATCH
@@ -13723,6 +13770,7 @@ NODE_CLASS_MAPPINGS = {
     "Image Voronoi Noise Filter": WAS_Image_Voronoi_Noise_Filter,
     "KSampler (WAS)": WAS_KSampler,
     "KSampler Cycle": WAS_KSampler_Cycle,
+    "Latent Batch": WAS_Latent_Batch,
     "Latent Noise Injection": WAS_Latent_Noise,
     "Latent Size to Number": WAS_Latent_Size_To_Number,
     "Latent Upscale by Factor (WAS)": WAS_Latent_Upscale,