From a94bcea667c11a4e0d1005f95f648ed70cac87f5 Mon Sep 17 00:00:00 2001 From: blepping Date: Tue, 29 Jul 2025 23:11:06 -0600 Subject: [PATCH] Add probinject and probsubtract_b blend modes Add BlehLatentAsImage and BlehImageAsLatent nodes --- py/latent_utils.py | 29 ++++++++++++++ py/nodes/__init__.py | 10 +++-- py/nodes/misc.py | 92 +++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 125 insertions(+), 6 deletions(-) diff --git a/py/latent_utils.py b/py/latent_utils.py index e88682a..6735447 100644 --- a/py/latent_utils.py +++ b/py/latent_utils.py @@ -50,6 +50,27 @@ if USE_ORIG_NORMALIZE: normalize = normalize_orig +def normalize_to_scale( + latent: torch.Tensor, + target_min: float, + target_max: float, + *, + dim=(-3, -2, -1), + eps: float = 1e-07, +) -> torch.Tensor: + min_val, max_val = ( + latent.amin(dim=dim, keepdim=True), + latent.amax(dim=dim, keepdim=True), + ) + normalized = latent - min_val + normalized /= (max_val - min_val).add_(eps) + return ( + normalized.mul_(target_max - target_min) + .add_(target_min) + .clamp_(target_min, target_max) + ) + + def hslerp(a, b, t): if a.shape != b.shape: raise ValueError("Input tensors a and b must have the same shape.") @@ -813,6 +834,14 @@ BLENDING_MODES = { kernel_size=9, sigma=3.0, ), + "probinject": BlendMode( + lambda a, b, t, **kwargs: prob_blend(torch.zeros_like(b), b, t, **kwargs).add_( + a, + ), + ), + "probsubtract_b": BlendMode( + lambda a, b, t, **kwargs: a - prob_blend(torch.zeros_like(b), b, t, **kwargs), + ), "gradient": BlendMode(gradient_blend), # Adds tensor b to tensor a, scaled by t. "inject": BlendMode(lambda a, b, t: (b * t).add_(a)), diff --git a/py/nodes/__init__.py b/py/nodes/__init__.py index 52c4334..9456516 100644 --- a/py/nodes/__init__.py +++ b/py/nodes/__init__.py @@ -16,25 +16,27 @@ _blepping_integrations = None NODE_CLASS_MAPPINGS = { "BlehBlockCFG": blockCFG.BlockCFGBleh, "BlehBlockOps": ops.BlehBlockOps, + "BlehCast": misc.BlehCast, + "BlehCFGInitSampler": samplers.BlehCFGInitSampler, "BlehDeepShrink": deepShrink.DeepShrinkBleh, "BlehDisableNoise": misc.BlehDisableNoise, "BlehDiscardPenultimateSigma": misc.DiscardPenultimateSigma, + "BlehEnsurePreviewer": misc.BlehEnsurePreviewer, "BlehForceSeedSampler": samplers.BlehForceSeedSampler, "BlehGlobalSageAttention": sageAttention.BlehGlobalSageAttention, "BlehHyperTile": hyperTile.HyperTileBleh, + "BlehImageAsLatent": misc.BlehImageAsLatent, "BlehInsaneChainSampler": samplers.BlehInsaneChainSampler, - "BlehCFGInitSampler": samplers.BlehCFGInitSampler, + "BlehLatentAsImage": misc.BlehLatentAsImage, + "BlehLatentBlend": ops.BlehLatentBlend, "BlehLatentOps": ops.BlehLatentOps, "BlehLatentScaleBy": ops.BlehLatentScaleBy, - "BlehLatentBlend": ops.BlehLatentBlend, "BlehModelPatchConditional": modelPatchConditional.ModelPatchConditionalNode, "BlehPlug": misc.BlehPlug, "BlehRefinerAfter": refinerAfter.BlehRefinerAfter, "BlehSageAttentionSampler": sageAttention.BlehSageAttentionSampler, "BlehSetSamplerPreset": samplers.BlehSetSamplerPreset, - "BlehCast": misc.BlehCast, "BlehSetSigmas": misc.BlehSetSigmas, - "BlehEnsurePreviewer": misc.BlehEnsurePreviewer, "BlehTAEVideoDecode": taevid.TAEVideoDecode, "BlehTAEVideoEncode": taevid.TAEVideoEncode, } diff --git a/py/nodes/misc.py b/py/nodes/misc.py index 3406285..94199f2 100644 --- a/py/nodes/misc.py +++ b/py/nodes/misc.py @@ -1,7 +1,6 @@ # ruff: noqa: TID252 from __future__ import annotations -import itertools import operator import random from decimal import Decimal @@ -9,8 +8,8 @@ from decimal import Decimal import torch from comfy import model_management -from .. import latent_utils from ..better_previews.previewer import ensure_previewer +from ..latent_utils import normalize_to_scale class DiscardPenultimateSigma: @@ -323,6 +322,95 @@ class BlehEnsurePreviewer: return (any_input,) +class BlehImageAsLatent: + DESCRIPTION = "This node allows you to rearrange an IMAGE to look like a LATENT. Can be useful if you want to apply some latent operations to an IMAGE. Can be reversed with the BlehLatentAsImage node." + FUNCTION = "go" + CATEGORY = "latent/advanced" + RETURN_TYPES = ("LATENT",) + + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "audio": ("IMAGE",), + "rescale": ( + "BOOLEAN", + { + "default": True, + "tooltip": "When enabled, will rescale the image values which usually are from 0 to 1 to -1 to 1.", + }, + ), + }, + } + + @classmethod + def go(cls, *, image: torch.Tensor, rescale: bool) -> tuple: + image = image.to(device="cpu", dtype=torch.float32, copy=True) + if image.ndim == 3: + image = image[None] + elif image.ndim != 4: + raise ValueError("Unexpected number of dimensions in image") + image = image.movedim(-1, 1) + if rescale: + image = image.sub_(0.5).mul_(2.0) + return ({"samples": image},) + + +class BlehLatentAsImage: + DESCRIPTION = "This node lets you rearrange a LATENT to look like an IMAGE. Note: It does not respect anything like masks or latent selection metadata (from nodes like LatentFromBatch) that might exist." + FUNCTION = "go" + CATEGORY = "latent/advanced" + RETURN_TYPES = ("IMAGE",) + + @classmethod + def INPUT_TYPES(cls) -> dict: + return { + "required": { + "latent": ("LATENT",), + "values_mode": ( + ("rescale", "rescale_perchannel", "clamp"), + {"default": "rescale"}, + ), + "channels_into_batch": ( + "BOOLEAN", + { + "default": False, + "tooltip": "When enabled, will create a greyscale image for each latent channel.", + }, + ), + }, + } + + @classmethod + def go( + cls, + *, + latent: dict, + values_mode: str, + channels_into_batch: bool, + ) -> tuple: + samples = latent["samples"] + if samples.ndim != 4: + raise ValueError("Expected a 4D latent but didn't get one") + if channels_into_batch: + samples = samples.reshape(-1, *samples.shape[2:]).unsqueeze(1) + samples = samples.expand(samples.shape[0], 3, *samples.shape[2:]) + image = samples.movedim(1, -1).to( + device="cpu", + dtype=torch.float32, + copy=True, + )[..., :4] + if values_mode == "clamp": + return (image.clamp(0.0, 1.0),) + image = normalize_to_scale( + image, + 0.0, + 1.0, + dim=(2, 3) if values_mode == "rescale_perchannel" else (1, 2, 3), + ) + return (image,) + + # class BlehConditioningBlend: # DESCRIPTION = "TBD" # FUNCTION = "go"