From 1200d1cf25fe6fb55a1e90ecc1e3b91c01b4060c Mon Sep 17 00:00:00 2001 From: YacratesWyh <1410669639@qq.com> Date: Fri, 12 Jan 2024 15:19:36 +0800 Subject: [PATCH 1/3] Latent Batch --- WAS_Node_Suite.py | 43 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index efb00f9..0be4b2c 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -4715,6 +4715,48 @@ class WAS_Image_Batch: self._check_image_dimensions(batched_tensors, image_names) batched_tensors = torch.cat(batched_tensors, dim=0) return (batched_tensors,) +# IMAGE 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.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) + batched_tensors = torch.cat(batched_tensors, dim=0) + return (batched_tensors,) # MASK TO BATCH @@ -13723,6 +13765,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, From 7f155247b0e026b2e2309230d4a517e64a21f77a Mon Sep 17 00:00:00 2001 From: YacratesWyh <1410669639@qq.com> Date: Fri, 12 Jan 2024 17:10:39 +0800 Subject: [PATCH 2/3] fix logic --- WAS_Node_Suite.py | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 0be4b2c..7a18010 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -4738,11 +4738,22 @@ class WAS_Latent_Batch: RETURN_NAMES = ("latent",) FUNCTION = "latent_batch" CATEGORY = "WAS Suite/Latent" + # def batch(self, samples1, samples2): + # samples_out = samples1.copy() + # s1 = samples1["samples"] + # s2 = samples2["samples"] + # if s1.shape[1:] != s2.shape[1:]: + # s2 = comfy.utils.common_upscale(s2, s1.shape[3], s1.shape[2], "bilinear", "center") + # s = torch.cat((s1, s2), dim=0) + # samples_out["samples"] = s + # samples_out["batch_index"] = samples1.get("batch_index", [x for x in range(0, s1.shape[0])]) + samples2.get("batch_index", [x for x in range(0, s2.shape[0])]) + # return (samples_out,) + def _check_latent_dimensions(self, tensors, names): - dimensions = [tensor.shape for tensor in tensors] + 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_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}") @@ -4755,8 +4766,14 @@ class WAS_Latent_Batch: raise ValueError("At least one input latent must be provided.") self._check_latent_dimensions(batched_tensors, latent_names) - batched_tensors = torch.cat(batched_tensors, dim=0) - return (batched_tensors,) + samples_out = {} + # s = [(tensor["samples"]) for tensor in batched_tensors] + 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 From 29e7f446b0c92fe4c3573c737b4e3399352ea8b2 Mon Sep 17 00:00:00 2001 From: YacratesWyh <1410669639@qq.com> Date: Fri, 12 Jan 2024 17:15:38 +0800 Subject: [PATCH 3/3] delete comment --- WAS_Node_Suite.py | 16 ++-------------- 1 file changed, 2 insertions(+), 14 deletions(-) diff --git a/WAS_Node_Suite.py b/WAS_Node_Suite.py index 7a18010..974c122 100644 --- a/WAS_Node_Suite.py +++ b/WAS_Node_Suite.py @@ -4715,7 +4715,7 @@ class WAS_Image_Batch: self._check_image_dimensions(batched_tensors, image_names) batched_tensors = torch.cat(batched_tensors, dim=0) return (batched_tensors,) -# IMAGE TO BATCH +# Latent TO BATCH class WAS_Latent_Batch: def __init__(self): @@ -4738,18 +4738,7 @@ class WAS_Latent_Batch: RETURN_NAMES = ("latent",) FUNCTION = "latent_batch" CATEGORY = "WAS Suite/Latent" - # def batch(self, samples1, samples2): - # samples_out = samples1.copy() - # s1 = samples1["samples"] - # s2 = samples2["samples"] - - # if s1.shape[1:] != s2.shape[1:]: - # s2 = comfy.utils.common_upscale(s2, s1.shape[3], s1.shape[2], "bilinear", "center") - # s = torch.cat((s1, s2), dim=0) - # samples_out["samples"] = s - # samples_out["batch_index"] = samples1.get("batch_index", [x for x in range(0, s1.shape[0])]) + samples2.get("batch_index", [x for x in range(0, s2.shape[0])]) - # return (samples_out,) - + def _check_latent_dimensions(self, tensors, names): dimensions = [(tensor["samples"].shape) for tensor in tensors] if len(set(dimensions)) > 1: @@ -4767,7 +4756,6 @@ class WAS_Latent_Batch: self._check_latent_dimensions(batched_tensors, latent_names) samples_out = {} - # s = [(tensor["samples"]) for tensor in batched_tensors] samples_out["samples"] = torch.cat([tensor["samples"] for tensor in batched_tensors], dim=0) samples_out["batch_index"] = [] for tensor in batched_tensors: