diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index efb00f9..974c122 100644 --- a/WAS_Node_Suite.py +++ b/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,