|
import comfy.utils
|
|
import torch
|
|
|
|
def reshape_latent_to(target_shape, latent):
|
|
if latent.shape[1:] != target_shape[1:]:
|
|
latent = comfy.utils.common_upscale(latent, target_shape[3], target_shape[2], "bilinear", "center")
|
|
return comfy.utils.repeat_to_batch_size(latent, target_shape[0])
|
|
|
|
|
|
class LatentAdd:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "op"
|
|
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def op(self, samples1, samples2):
|
|
samples_out = samples1.copy()
|
|
|
|
s1 = samples1["samples"]
|
|
s2 = samples2["samples"]
|
|
|
|
s2 = reshape_latent_to(s1.shape, s2)
|
|
samples_out["samples"] = s1 + s2
|
|
return (samples_out,)
|
|
|
|
class LatentSubtract:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "op"
|
|
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def op(self, samples1, samples2):
|
|
samples_out = samples1.copy()
|
|
|
|
s1 = samples1["samples"]
|
|
s2 = samples2["samples"]
|
|
|
|
s2 = reshape_latent_to(s1.shape, s2)
|
|
samples_out["samples"] = s1 - s2
|
|
return (samples_out,)
|
|
|
|
class LatentMultiply:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "samples": ("LATENT",),
|
|
"multiplier": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "op"
|
|
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def op(self, samples, multiplier):
|
|
samples_out = samples.copy()
|
|
|
|
s1 = samples["samples"]
|
|
samples_out["samples"] = s1 * multiplier
|
|
return (samples_out,)
|
|
|
|
class LatentInterpolate:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "samples1": ("LATENT",),
|
|
"samples2": ("LATENT",),
|
|
"ratio": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
|
|
}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "op"
|
|
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def op(self, samples1, samples2, ratio):
|
|
samples_out = samples1.copy()
|
|
|
|
s1 = samples1["samples"]
|
|
s2 = samples2["samples"]
|
|
|
|
s2 = reshape_latent_to(s1.shape, s2)
|
|
|
|
m1 = torch.linalg.vector_norm(s1, dim=(1))
|
|
m2 = torch.linalg.vector_norm(s2, dim=(1))
|
|
|
|
s1 = torch.nan_to_num(s1 / m1)
|
|
s2 = torch.nan_to_num(s2 / m2)
|
|
|
|
t = (s1 * ratio + s2 * (1.0 - ratio))
|
|
mt = torch.linalg.vector_norm(t, dim=(1))
|
|
st = torch.nan_to_num(t / mt)
|
|
|
|
samples_out["samples"] = st * (m1 * ratio + m2 * (1.0 - ratio))
|
|
return (samples_out,)
|
|
|
|
class LatentBatch:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "samples1": ("LATENT",), "samples2": ("LATENT",)}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "batch"
|
|
|
|
CATEGORY = "latent/batch"
|
|
|
|
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,)
|
|
|
|
class LatentBatchSeedBehavior:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "samples": ("LATENT",),
|
|
"seed_behavior": (["random", "fixed"],{"default": "fixed"}),}}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
FUNCTION = "op"
|
|
|
|
CATEGORY = "latent/advanced"
|
|
|
|
def op(self, samples, seed_behavior):
|
|
samples_out = samples.copy()
|
|
latent = samples["samples"]
|
|
if seed_behavior == "random":
|
|
if 'batch_index' in samples_out:
|
|
samples_out.pop('batch_index')
|
|
elif seed_behavior == "fixed":
|
|
batch_number = samples_out.get("batch_index", [0])[0]
|
|
samples_out["batch_index"] = [batch_number] * latent.shape[0]
|
|
|
|
return (samples_out,)
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"LatentAdd": LatentAdd,
|
|
"LatentSubtract": LatentSubtract,
|
|
"LatentMultiply": LatentMultiply,
|
|
"LatentInterpolate": LatentInterpolate,
|
|
"LatentBatch": LatentBatch,
|
|
"LatentBatchSeedBehavior": LatentBatchSeedBehavior,
|
|
}
|
|
|
