feat: update transforms to operate on model instead of latent; bug fixes; version bump to 3.0.0

Update Github Action for Publishing to Comfy Registry

.github/workflows/publish.yml

@@ -0,0 +1,26 @@
+name: Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    branches:
+      - main
+      - master
+    paths:
+      - "pyproject.toml"
+
+permissions:
+  issues: write
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    if: ${{ github.repository_owner == 'RomanKuschanow' }}
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+      - name: Publish Custom Node
+        uses: Comfy-Org/publish-node-action@v1
+        with:
+          ## Add your own personal access token to your Github Repository secrets and reference it here.
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

README.md

@@ -125,13 +125,22 @@ And it very slightly changes results from latent, which have not been modified.
 
 Allow you to use transforms with any samplers that you like.
 
+Instead of patching the latent, this node patches the **model**: it attaches the transforms
+to a cloned model via a sampler post-cfg hook. Connect the returned model to any sampler
+(`KSampler`, `KSamplerAdvanced`, custom sampler nodes, etc.) and the transforms will be applied
+during sampling.
+
 **Inputs:**
-- latent
+- model
 - transforms
 
 **Outputs:**
-- latent
+- model
 
 **Usage:**
 ![sample](https://i.imgur.com/YwVhHYF.png)  
 
+> **Breaking change in 3.0.0:** `Transform hijack` now takes and returns a `MODEL` instead of a
+> `LATENT`. This replaces the old global `common_ksampler` monkey-patch, which conflicted with the
+> stock `KSampler` and other custom nodes. Rewire this node to your model input/output after updating.
+

__init__.py

@@ -38,8 +38,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "LatentInterpolateTransform": "Latent interpolate transform",
     "LatentAddTransform": "Latent add transform",
     "OneTimeMirrorTransform": "Mirror transform (one time)",
-    "OneTimeMultiplyTransform": "Shift transform (one time)",
-    "OneTimeShiftTransform": "Multiply transform (one time)",
+    "OneTimeMultiplyTransform": "Multiply transform (one time)",
+    "OneTimeShiftTransform": "Shift transform (one time)",
     "OneTimeLatentInterpolateTransform": "Latent interpolate transform (one time)",
     "OneTimeLatentAddTransform": "Latent add transform (one time)",
     "TransformsCombine": "Combine transforms",

nodes/KSamplerNodes/TransformContext.py

@@ -1,65 +0,0 @@
-import torch
-import nodes
-import comfy
-from latent_preview import prepare_callback as preview_callback
-
-
-class TransformContext:
-    original_sample_function = nodes.common_ksampler
-
-    def get_transform_sample_function(self):
-        def prepare_callback(model, steps, x0_output_dict=None, transforms=None):
-            def transform_callback(step, x0, x, total_steps):
-                if transforms is None:
-                    return
-
-                for transform in transforms:
-                    for i in range(x0.size()[0]):
-                        x0[i] = transform["function"](step, x0[i].unsqueeze(0), total_steps, transform["params"])
-
-            preview = preview_callback(model, steps, x0_output_dict)
-
-            def callback(step, x0, x, total_steps):
-                transform_callback(step, x0, x, total_steps)
-                preview(step, x0, x, total_steps)
-
-            return callback
-
-        def sample(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, denoise=1.0,
-                          disable_noise=False, start_step=None, last_step=None, force_full_denoise=False):
-            latent_image = latent["samples"]
-            if disable_noise:
-                noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
-            else:
-                batch_inds = latent["batch_index"] if "batch_index" in latent else None
-                noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
-
-            noise_mask = None
-            if "noise_mask" in latent:
-                noise_mask = latent["noise_mask"]
-
-            callback = prepare_callback(model, steps, transforms=latent["transforms"])
-            disable_pbar = not comfy.utils.PROGRESS_BAR_ENABLED
-            samples = comfy.sample.sample(model, noise, steps, cfg, sampler_name, scheduler, positive, negative, latent_image,
-                                          denoise=denoise, disable_noise=disable_noise, start_step=start_step,
-                                          last_step=last_step,
-                                          force_full_denoise=force_full_denoise, noise_mask=noise_mask, callback=callback,
-                                          disable_pbar=disable_pbar, seed=seed)
-            out = latent.copy()
-            out["samples"] = samples
-            self.unhijack()
-            return (out,)
-
-        return sample
-
-    def hijack(self):
-        nodes.common_ksampler = self.get_transform_sample_function()
-
-    def unhijack(self):
-        nodes.common_ksampler = TransformContext.original_sample_function
-
-    def __enter__(self):
-        self.hijack()
-
-    def __exit__(self, exc_type, exc_value, exc_traceback):
-        self.unhijack()

nodes/KSamplerNodes/TransformHijack.py

@@ -1,32 +1,20 @@
-from .TransformContext import TransformContext
+from .transform_apply import attach_transforms
 
 
 class TransformHijack:
     @classmethod
     def INPUT_TYPES(cls):
         return {
-            "required" : {
-                "latent": ("LATENT",),
-                "transforms": ("TRANSFORM",)
+            "required": {
+                "model": ("MODEL",),
+                "transforms": ("TRANSFORM",),
             },
         }
 
-    RETURN_TYPES = ("LATENT",)
+    RETURN_TYPES = ("MODEL",)
     FUNCTION = "func"
 
     CATEGORY = "sampling/transforms"
 
-    _context = None
-    _hijack_node_id = None
-
-    def func(self, latent, transforms):
-        latent["transforms"] = transforms
-
-        if TransformHijack._context is None:
-            TransformHijack._hijack_node_id = id
-            TransformHijack._context = TransformContext()
-        else:
-            return (latent,)
-
-        TransformHijack._context.hijack()
-        return (latent,)
+    def func(self, model, transforms):
+        return (attach_transforms(model, transforms),)

nodes/KSamplerNodes/TransformSampler.py

@@ -1,4 +1,4 @@
-from .TransformContext import TransformContext
+from .transform_apply import attach_transforms
 from nodes import KSampler, KSamplerAdvanced
 
 
@@ -16,17 +16,11 @@ class Transforms:
 
     FUNCTION = "func"
 
-    def __init__(self):
-        self.original_function_name = self.clazz.FUNCTION
-
     def func(self, **kwargs):
-        ctx = TransformContext()
-        ctx.hijack()
-        latent = kwargs["latent_image"]
-        latent["transforms"] = kwargs.pop("transform_optional")
-        kwargs["latent_image"] = latent
-        out = getattr(self, self.clazz.FUNCTION)(**kwargs)
-        return out
+        transforms = kwargs.pop("transform_optional", None)
+        if transforms:
+            kwargs["model"] = attach_transforms(kwargs["model"], transforms)
+        return getattr(self, self.clazz.FUNCTION)(**kwargs)
 
 
 def variations_factory(original_class: type, name=None) -> type:
@@ -34,4 +28,4 @@ def variations_factory(original_class: type, name=None) -> type:
     return type(name, (Transforms, original_class), {'clazz': original_class})
 
 TSampler = variations_factory(KSampler)
-TSamplerAdvanced = variations_factory(KSamplerAdvanced)
+TSamplerAdvanced = variations_factory(KSamplerAdvanced)

nodes/KSamplerNodes/Transforms/utils.py

@@ -62,9 +62,7 @@ def latent_interpolate_transform(x0, params):
     latent = params["latent"].to(x0.device)
 
     if x0.shape != latent.shape:
-        latent.permute(0, 3, 1, 2)
         latent = comfy.utils.common_upscale(latent, x0.shape[3], x0.shape[2], 'bicubic', crop='center')
-        latent.permute(0, 2, 3, 1)
 
     x = latent * params["factor"] + x0 * (1 - params["factor"])
     x *= params["multiplier"]
@@ -76,9 +74,7 @@ def latent_add_transform(x0, params):
     latent = params["latent"].to(x0.device)
 
     if x0.shape != latent.shape:
-        latent.permute(0, 3, 1, 2)
-        latent = comfy.utils.common_upscale(latent, x0.shape[3], x0.shape[2], 'bicubic',  crop='center')
-        latent.permute(0, 2, 3, 1)
+        latent = comfy.utils.common_upscale(latent, x0.shape[3], x0.shape[2], 'bicubic', crop='center')
 
     x = x0 + latent
     x *= params["multiplier"]

nodes/KSamplerNodes/transform_apply.py

@@ -0,0 +1,51 @@
+import torch
+
+
+def apply_transforms_to_x0(x0, step, total_steps, transforms):
+    x = x0.clone()
+
+    for transform in transforms:
+        for i in range(x.size()[0]):
+            x[i] = transform["function"](step, x[i].unsqueeze(0), total_steps, transform["params"])
+
+    return x
+
+
+def _find_step(sigma, sigmas):
+    # sigma is a scalar tensor for the current model evaluation, sigmas is the full
+    # schedule. High order samplers evaluate the model at intermediate sigmas that are
+    # not part of the schedule; for those we return None so the transform is applied
+    # exactly once per step (parity with the old per-step callback).
+    diff = torch.abs(sigmas - sigma.to(sigmas.device))
+    idx = int(torch.argmin(diff).item())
+    if diff[idx] <= 1e-4 * max(1.0, float(sigmas[idx].abs())):
+        return idx
+    return None
+
+
+def make_post_cfg_function(transforms):
+    def post_cfg_function(args):
+        denoised = args["denoised"]
+
+        if not transforms:
+            return denoised
+
+        sigmas = args["model_options"].get("transformer_options", {}).get("sample_sigmas", None)
+        if sigmas is None:
+            return denoised
+
+        step = _find_step(args["sigma"], sigmas)
+        if step is None:
+            return denoised
+
+        total_steps = len(sigmas) - 1
+
+        return apply_transforms_to_x0(denoised, step, total_steps, transforms)
+
+    return post_cfg_function
+
+
+def attach_transforms(model, transforms):
+    m = model.clone()
+    m.set_model_sampler_post_cfg_function(make_post_cfg_function(transforms))
+    return m

pyproject.toml

@@ -0,0 +1,14 @@
+[project]
+name = "comfyui-advanced-latent-control"
+description = "This custom node helps to transform latent in different ways."
+version = "3.0.0"
+license = "LICENSE"
+
+[project.urls]
+Repository = "https://github.com/RomanKuschanow/ComfyUI-Advanced-Latent-Control"
+#  Used by Comfy Registry https://comfyregistry.org
+
+[tool.comfy]
+PublisherId = "kuschanow"
+DisplayName = "ComfyUI-Advanced-Latent-Control"
+Icon = ""