Add probinject and probsubtract_b blend modes

Add BlehLatentAsImage and BlehImageAsLatent nodes

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)),

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,
 }

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"