feat: refactor offset widget to use custom canvas; update README and version to 3.1.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 == 'kuschanow' }}
+    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

@@ -19,8 +19,8 @@ This node can flip latent and merge original and flipped version.
 - `latent`
 
 **Usage:**  
-![sample](https://i.imgur.com/YMyYorQ.png)
-![sample](https://i.imgur.com/W5BasCO.png)
+![sample](assets/latent-mirror-usage-1.png)
+![sample](assets/latent-mirror-usage-2.png)
 
 ## Latent shift
 This node can shift latent along x and y-axis.
@@ -36,39 +36,10 @@ This node can shift latent along x and y-axis.
 - `latent`
 
 **Usage:**  
-![sample](https://i.imgur.com/1Dp5dSw.png)
+![sample](assets/latent-shift-usage.png)
 
-## TSampler with transforms (Latent Control)
-This node can multiply, mirror and shift latent during generation.
-
-**Input:**  
-exactly matches the base KSampler
-
-**Fields:**
-- base KSampler fields
-- `start_mirror_at` – a number between 0 and 1 that indicates at what point the sampler will start mirroring
-- `stop_mirror_at` – a number between 0 and 1 that indicates at what point the sampler will stop mirroring
-- `mirror_mode` – can be `replace` or `combine`. `replace` will replace the latent with the transformed one, `combine` will add the original and the transformed latent and divide by 2
-- `mirror_direction` – can be `none`, `vertically`, `horizontally`, `both`, `90 degree rotation` or `180 degree rotation`
-- `start_shift_at` – a number between 0 and 1 that indicates at what point the sampler will start shifting
-- `stop_shift_at` – a number between 0 and 1 that indicates at what point the sampler will stop shifting
-- `shift_mode` – can be `replace` or `combine`. `replace` will replace the latent with the transformed one, `combine` will add the original and the transformed latent and divide by 2
-- `x_shift` – a number between -1 and 1 that indicates how much the latent should be shifted
-- `y_shift` – a number between -1 and 1 that indicates how much the latent should be shifted
-- `start_multiplier_at` – a number between 0 and 1 that indicates at what point the sampler will start multiplying
-- `stop_multiplier_at` – a number between 0 and 1 that indicates at what point the sampler will stop multiplying
-- `multiplier_mode` – can be `replace` or `combine`. `replace` will replace the latent with the transformed one, `combine` will add the original and the transformed latent
-- `multiplier` – multiply latent by specified number
-
-**Output:**
-exactly matches the base KSampler
-
-**Usage:**  
-**You also can use those params together**  
-![sample](https://i.imgur.com/RMJTnWF.png)  
-![sample](https://i.imgur.com/fQ7UWuS.png)  
-![sample](https://i.imgur.com/pxWupAx.png)  
-![sample](https://i.imgur.com/1YkERDu.png)  
+## ~~TSampler with transforms (Latent Control)~~ 
+Removed from version 2.0.0
 
 ## TSampler (Latent Control)
 This node allows to combine a lot of transforms with different parameters.
@@ -84,8 +55,8 @@ exactly matches the base KSampler
 exactly matches the base KSampler
 
 **Usage:**  
-![sample](https://i.imgur.com/PlGnAtA.png)  
-![sample](https://i.imgur.com/CtrBRPn.png)  
+![sample](assets/tsampler-usage-1.png)  
+![sample](assets/tsampler-usage-2.png)  
 
 Multiply, Mirror and Shift transform nodes parameters exactly match the corresponding `KSampler with transforms (Latent Control)` parameters.
 
@@ -107,9 +78,9 @@ You can apply specific offset for transform nodes.
 - `offset`
 
 **Usage:**  
-![sample](https://i.imgur.com/ExZacqG.png)
-![sample](https://i.imgur.com/tR6KSmI.png)
-![sample](https://i.imgur.com/MGVLfve.png)
+![sample](assets/offset-usage-1.png)
+![sample](assets/offset-usage-2.png)
+![sample](assets/offset-usage-3.png)
 
 You can combine different offsets to achieve interesting patterns. For example:  
 **0 0 0 1** and **0 0 1** give this pattern: **0 0 1 1 0 1 0 1 1 0 0 1**.
@@ -119,35 +90,57 @@ You can combine different offsets to achieve interesting patterns. For example:
 Each transform node has own one-time version. They allow to make one transform action at specified step.
 
 **Usage:**  
-![sample](https://i.imgur.com/Q1Vyob0.png)
+![sample](assets/one-time-nodes-usage.png)
 
 ## Latent normalize
 
 Fixes some issues when sampling modified latent space.
 
-**Input**  
+**Input:**  
 exactly matches the `VAE Decode` node
 
-**Output**
+**Output:**
 - latent
 
 When you multiply latent by negative or big positive (bigger than 2) number and paste this latent in sampler, you can see that the
 image will be generated very poorly. This is because stable diffusion cannot work with such set of numbers (meaning the numbers contained in latent).
 
-![sample](https://i.imgur.com/3FXk8n7.png)
+![sample](assets/latent-normalize-poor-result.png)
 
 But you can prevent this behavior by sequential decode and encode latent using vae. Node `Latent normalize` make this process easier.
 
-![sample](https://i.imgur.com/hkFYYVh.png)
+![sample](assets/latent-normalize-fix.png)
 
 This node also change some results even if output without this node looks good.
 
-![sample](https://i.imgur.com/kP0f6vh.png)
-![sample](https://i.imgur.com/YI8ZqLd.png)
+![sample](assets/latent-normalize-comparison-1.png)
+![sample](assets/latent-normalize-comparison-2.png)
 
 And it very slightly changes results from latent, which have not been modified.
 
-![sample](https://i.imgur.com/xTU08xm.png)
-![sample](https://i.imgur.com/yzgW7QT.png)
+![sample](assets/latent-normalize-unmodified-1.png)
+![sample](assets/latent-normalize-unmodified-2.png)
 
+## Transform hijack
+
+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:**
+- model
+- transforms
+
+**Outputs:**
+- model
+
+**Usage:**
+![sample](assets/transform-hijack-usage.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

@@ -1,11 +1,15 @@
 from .nodes import *
 
+
+WEB_DIRECTORY = "js"
+
 NODE_CLASS_MAPPINGS = {
     "LatentMirror": LatentMirror,
     "LatentShift": LatentShift,
     "LatentNormalize": LatentNormalize,
-    "TSamplerWithTransform": TSamplerWithTransform,
-    "TransformSampler": TransformSampler,
+    "TransformSampler": TSampler,
+    "TransformSamplerAdvanced": TSamplerAdvanced,
+    "TransformHijack": TransformHijack,
     "MirrorTransform": MirrorTransform,
     "ShiftTransform": ShiftTransform,
     "MultiplyTransform": MultiplyTransform,
@@ -25,16 +29,17 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "LatentMirror": "Latent mirror",
     "LatentShift": "Latent shift",
     "LatentNormalize": "Latent normalize",
-    "TSamplerWithTransform": "TSampler with transforms (Latent Control)",
     "TransformSampler": "TSampler (Latent Control)",
+    "TransformSamplerAdvanced": "TSampler Advanced (Latent Control)",
+    "TransformHijack": "Transform Hijack",
     "MirrorTransform": "Mirror transform",
     "ShiftTransform": "Shift transform",
     "MultiplyTransform": "Multiply transform",
     "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",

js/offset.js

@@ -0,0 +1,132 @@
+import { app } from "/scripts/app.js";
+import {addCustomCanvasWidget, generatePattern, recursiveLinkUpstream, renameNodeInputs, removeNodeInputs} from "./utils.js";
+
+function drawSquares(ctx, startX, startY, squareSize, pattern) {
+	const inset = Math.max(1, Math.floor(squareSize / 16));
+	const cell = squareSize - inset;
+
+	pattern.forEach((value, index) => {
+		const x = startX + index * squareSize + inset / 2;
+		const y = startY + inset / 2;
+
+		ctx.lineWidth = inset;
+		ctx.strokeStyle = "#222223";
+
+		if (value === 1) {
+			ctx.fillStyle = "#222223";
+			ctx.fillRect(x, y, cell, cell);
+		}
+
+		ctx.strokeRect(x, y, cell, cell);
+
+		if (squareSize >= 24) {
+			ctx.font = `bold ${squareSize / 3}px Arial`;
+			ctx.textAlign = "center";
+			ctx.textBaseline = "middle";
+			ctx.fillStyle = value === 1 ? "#dbdbdc" : "#222223";
+			ctx.fillText(value.toString(), x + cell / 2, y + cell / 2);
+		}
+	});
+}
+
+// Builds the 0/1 preview pattern from the node's own widgets, or (for
+// OffsetCombine) from every upstream offset node feeding into it.
+function getOffsetPattern(node) {
+	let patterns = []
+
+	if (node.type === "OffsetCombine") {
+		const connectedNodes = recursiveLinkUpstream(node, node.inputs[0].type, node.type, 0)
+		for (let [node_ID] of connectedNodes) {
+			const connectedNode = node.graph._nodes_by_id[node_ID]
+			if (connectedNode && connectedNode.type !== "OffsetCombine") {
+				patterns.push({
+					process_every: connectedNode.widgets[0].value,
+					offset: connectedNode.widgets[1].value + node.widgets[0].value,
+					mode: connectedNode.widgets[2].value,
+				})
+			}
+		}
+	} else {
+		patterns.push({
+			process_every: node.widgets[0].value,
+			offset: node.widgets[1].value,
+			mode: node.widgets[2].value,
+		})
+	}
+
+	if (patterns.length === 0) return []
+
+	return generatePattern(patterns)
+}
+
+app.registerExtension({
+    name: "Comfy.LatentControl.TransformOffset",
+    async beforeRegisterNodeDef (nodeType, nodeData, app){
+        if (nodeData.name === "TransformOffset") {
+            const onNodeCreated = nodeType.prototype.onNodeCreated;
+            nodeType.prototype.onNodeCreated = function () {
+                const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+
+                addCustomCanvasWidget(this, "Offset-Canvas", getOffsetPattern, drawSquares)
+
+				return r;
+			}
+        }
+    },
+});
+
+app.registerExtension({
+    name: "Comfy.LatentControl.OffsetCombine",
+    async beforeRegisterNodeDef (nodeType, nodeData, app){
+        if (nodeData.name === "OffsetCombine") {
+            const onNodeCreated = nodeType.prototype.onNodeCreated;
+            nodeType.prototype.onNodeCreated = function () {
+                const r = onNodeCreated ? onNodeCreated.apply(this, arguments) : undefined;
+
+                addCustomCanvasWidget(this, "Offset-Canvas", getOffsetPattern, drawSquares)
+
+				this.getExtraMenuOptions = function(_, options) {
+					options.unshift(
+						{
+							content: `add offset`,
+							callback: () => {
+								this.addInput("offset", "OFFSET")
+
+								renameNodeInputs(this, "offset")
+
+								this.setDirtyCanvas(true);
+							},
+						},
+						{
+							content: `remove offset`,
+							callback: () => {
+								removeNodeInputs(this, [this.inputs.length-1])
+								renameNodeInputs(this, "offset")
+							},
+						},
+						{
+							content: "remove all unconnected offsets",
+							callback: () => {
+								let indexesToRemove = []
+
+								for (let i = 0; i < this.inputs.length; i++) {
+									if (!this.inputs[i].link) {
+										indexesToRemove.push(i)
+									}
+								}
+
+								if (indexesToRemove.length) {
+									removeNodeInputs(this, indexesToRemove)
+								}
+								renameNodeInputs(this, "offset")
+							},
+						},
+					);
+				}
+
+				return r;
+			}
+        }
+    },
+});
+

js/utils.js

@@ -0,0 +1,163 @@
+function gcd(a, b) {
+    // Функция для вычисления наибольшего общего делителя (НОД)
+    while (b !== 0) {
+        let t = b;
+        b = a % b;
+        a = t;
+    }
+    return a;
+}
+
+function lcm(a, b) {
+    // Функция для вычисления наименьшего общего кратного (НОК)
+    return (a * b) / gcd(a, b);
+}
+
+function findPatternLength(rules) {
+    // Вычисление длины цикла как НОК всех process_every
+    return rules.map(rule => rule.process_every).reduce((acc, val) => lcm(acc, val), 1);
+}
+
+export function generatePattern(rules) {
+    let length = findPatternLength(rules); // Определение длины паттерна
+    let pattern = new Array(length).fill(0);
+
+    rules.forEach(rule => {
+        let offset = rule.offset % rule.process_every;
+
+        for (let i = 0; i < length; i++) {
+			let value = ((i + offset) % rule.process_every === 0) === (rule.mode === "process_every") ? 1 : 0;
+			pattern[i] = pattern[i] || value;
+        }
+    });
+
+    return pattern;
+}
+
+export function recursiveLinkUpstream(node, slot_type, node_type, depth) {
+	depth += 1
+	let connections = []
+	const inputList = [...Array(node.inputs.length).keys()]
+	for (let i of inputList) {
+		const link = node.inputs[i].link
+		if (link) {
+			const nodeID = node.graph.links[link].origin_id
+			const slotID = node.graph.links[link].origin_slot
+			const connectedNode = node.graph._nodes_by_id[nodeID]
+
+			if (connectedNode.outputs[slotID].type === slot_type) {
+
+				connections.push([connectedNode.id, depth])
+
+				if (connectedNode.inputs) {
+					const index = (connectedNode.type === node_type) ? 0 : null
+					connections = connections.concat(recursiveLinkUpstream(connectedNode, slot_type, node_type, depth))
+				}
+			}
+		}
+	}
+
+	return connections
+}
+
+export function renameNodeInputs(node, name) {
+	for (let i=0; i < node.inputs.length; i++) {
+		node.inputs[i].name = `${name}${i + 1}`
+	}
+}
+
+export function removeNodeInputs(node, indexesToRemove) {
+	indexesToRemove.sort((a, b) => b - a);
+
+	for (let i of indexesToRemove) {
+		if (node.inputs.length <= 2) { console.log("too short"); continue } // if only 2 left
+		node.removeInput(i)
+	}
+
+	node.setSize(node.computeSize())
+	node.graph?.setDirtyCanvas(true, true)
+}
+
+const MARGIN_X = 15;         // side padding — matches the node's standard widgets
+const MARGIN_Y = 10;         // top/bottom padding around the grid (node-local px)
+const MIN_SQUARE = 24;       // smallest square that still shows the row on tiny/long patterns
+const MAX_SQUARE = 72;       // largest square, so a short pattern doesn't blow up
+const MIN_NODE_WIDTH = 220;  // enough width for the squares to read
+
+// Square side that fits the given (full node) width and column count. Squares
+// shrink to fit a long pattern (down to 1px) and are capped so a short one stays
+// reasonable. Guarantees squareSize * cols <= width - 2*MARGIN_X, i.e. no overflow.
+function fitSquareSize(width, cols) {
+	const availWidth = width - MARGIN_X * 2;
+	return Math.max(1, Math.min(availWidth / Math.max(cols, 1), MAX_SQUARE));
+}
+
+/**
+ * Adds a self-contained canvas widget that draws a pattern preview directly onto
+ * the LiteGraph node context. Works with the current ComfyUI frontend (Nodes v2):
+ * it paints in node-local coordinates, so it scales and positions correctly with
+ * zoom and DPI — no DOM overlay and no manual ctx.getTransform() math.
+ *
+ * The preview auto-sizes to the space available inside the node: computeSize()
+ * reports exactly the height the grid needs for the current width and pattern
+ * length, so the node reserves the right amount of room, the squares grow/shrink
+ * as the node gets wider/narrower or the pattern longer/shorter, and the drawing
+ * always stays inside the node body.
+ *
+ * @param node        the LiteGraph node
+ * @param name        widget name
+ * @param getPattern  (node) => number[]  the 0/1 pattern to render
+ * @param drawSquares (ctx, startX, startY, squareSize, pattern) => void
+ */
+export function addCustomCanvasWidget(node, name, getPattern, drawSquares) {
+	const widget = {
+		type: "customCanvas",
+		name,
+		value: undefined,
+		options: { serialize: false },
+		computeSize() {
+			// Use the real node width, not the passed value: LiteGraph feeds a
+			// widened width to selected nodes, which would inflate the reserve.
+			const nodeWidth = node.size[0];
+			const cols = getPattern(node).length || 1;
+			const rowHeight = Math.max(fitSquareSize(nodeWidth, cols), MIN_SQUARE);
+			return [nodeWidth, rowHeight + MARGIN_Y * 2];
+		},
+		draw(ctx, node, width, widgetY) {
+			const pattern = getPattern(node);
+			if (!pattern || !pattern.length) return;
+
+			// Use node.size[0] rather than the `width` argument: on a selected node
+			// LiteGraph passes an enlarged width, which would push the grid right and
+			// peg the square size at its cap. node.size[0] is stable in both states.
+			const nodeWidth = node.size[0];
+			const cols = pattern.length;
+			const squareSize = fitSquareSize(nodeWidth, cols);
+
+			// Same band height computeSize reserved -> the grid stays inside the node.
+			const bandHeight = Math.max(squareSize, MIN_SQUARE) + MARGIN_Y * 2;
+			const gridWidth = squareSize * cols;
+
+			const startX = (nodeWidth - gridWidth) / 2;
+			const startY = widgetY + (bandHeight - squareSize) / 2;
+
+			// Light backing drawn first so empty cells stay readable on dark nodes.
+			ctx.fillStyle = "#ffffffcc";
+			ctx.fillRect(startX, startY, gridWidth, squareSize);
+
+			drawSquares(ctx, startX, startY, squareSize, pattern);
+		},
+	};
+
+	node.addCustomWidget(widget);
+
+	// Widen a touch and reserve the preview height on creation. On workflow load
+	// LiteGraph applies the saved size after onNodeCreated, so this does not fight
+	// persisted sizes.
+	node.size[0] = Math.max(node.size[0], MIN_NODE_WIDTH);
+	node.setSize(node.computeSize());
+
+	return widget;
+}
+
+

nodes/KSamplerNodes/TSamplerWithTransform.py

@@ -1,88 +0,0 @@
-import comfy.samplers
-from .TransformSampler import TransformSampler
-from .Transforms import MirrorTransform, ShiftTransform, MultiplyTransform
-
-
-MIRROR_DIRECTIONS = ["none", "vertically", "horizontally", "both", "90 degree rotation", "180 degree rotation"]
-MODE = ["replace", "combine"]
-
-class TSamplerWithTransform:
-    @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
-                "positive": ("CONDITIONING", ),
-                "negative": ("CONDITIONING", ),
-                "latent_image": ("LATENT", ),
-                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "start_mirror_at": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "stop_mirror_at": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "mirror_mode": (MODE,),
-                "mirror_direction": (MIRROR_DIRECTIONS, {"default": "none"}),
-                "start_shift_at": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "stop_shift_at": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "shift_mode": (MODE, {"default": "replace"}),
-                "x_shift": ("FLOAT", {"default": 0, "min": -1, "max": 1, "step": 0.01}),
-                "y_shift": ("FLOAT", {"default": 0, "min": -1, "max": 1, "step": 0.01}),
-                "start_multiplier_at": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "stop_multiplier_at": ("FLOAT", {"default": 0, "min": 0.0, "max": 1.0, "step": 0.01}),
-                "multiplier_mode": (MODE, {"default": "combine"}),
-                "multiplier": ("FLOAT", {"default": 1, "min": -10, "max": 10, "step": 0.01}),
-             }
-        }
-
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
-
-    CATEGORY = "sampling"
-
-    def sample(self,
-               model,
-               seed,
-               steps,
-               cfg,
-               sampler_name,
-               scheduler,
-               positive,
-               negative,
-               latent_image,
-               denoise=1.0,
-               start_mirror_at=0,
-               stop_mirror_at=0,
-               mirror_mode="replace",
-               mirror_direction="none",
-               start_shift_at=0,
-               stop_shift_at=0,
-               shift_mode="replace",
-               x_shift=0,
-               y_shift=0,
-               start_multiplier_at=0,
-               stop_multiplier_at=0,
-               multiplier_mode="combine",
-               multiplier=1):
-
-        transforms = (
-            MirrorTransform().process(start_mirror_at, stop_mirror_at, mirror_mode, mirror_direction) +
-            ShiftTransform().process(start_shift_at, stop_shift_at, shift_mode, x_shift, y_shift) +
-            MultiplyTransform().process(start_multiplier_at, stop_multiplier_at, multiplier_mode, multiplier))[0]
-
-        return TransformSampler().sample(
-            model,
-            seed,
-            steps,
-            cfg,
-            sampler_name,
-            scheduler,
-            positive,
-            negative,
-            latent_image,
-            transform_optional=transforms,
-            denoise=denoise)
-
-

nodes/KSamplerNodes/TransformHijack.py

@@ -0,0 +1,20 @@
+from .transform_apply import attach_transforms
+
+
+class TransformHijack:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model": ("MODEL",),
+                "transforms": ("TRANSFORM",),
+            },
+        }
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "func"
+
+    CATEGORY = "sampling/transforms"
+
+    def func(self, model, transforms):
+        return (attach_transforms(model, transforms),)

nodes/KSamplerNodes/TransformSampler.py

@@ -1,86 +1,31 @@
-import torch
-import comfy.samplers
-from latent_preview import prepare_callback as preview_callback
+from .transform_apply import attach_transforms
+from nodes import KSampler, KSamplerAdvanced
 
 
-def prepare_callback(model, steps, transforms, x0_output_dict=None):
-    def transform_callback(step, x0, x, total_steps):
-        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)
+def insert_transform_input(input_types):
+    input_types["optional"] = {"transform_optional": ("TRANSFORM",)}
+    return input_types
 
 
-    return callback
+class Transforms:
+    clazz = None
 
-
-def sample_common(model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent, transform, 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, transform)
-    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
-    return (out,)
-
-class TransformSampler:
     @classmethod
-    def INPUT_TYPES(s):
-        return {
-            "required": {
-                "model": ("MODEL",),
-                "seed": ("INT", {"default": 0, "min": 0, "max": 0xffffffffffffffff}),
-                "steps": ("INT", {"default": 20, "min": 1, "max": 10000}),
-                "cfg": ("FLOAT", {"default": 8.0, "min": 0.0, "max": 100.0, "step":0.1, "round": 0.01}),
-                "sampler_name": (comfy.samplers.KSampler.SAMPLERS, ),
-                "scheduler": (comfy.samplers.KSampler.SCHEDULERS, ),
-                "positive": ("CONDITIONING", ),
-                "negative": ("CONDITIONING", ),
-                "latent_image": ("LATENT", ),
-                "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}),
-            },
-            "optional":{
-                "transform_optional": ("TRANSFORM",),
-            }
-        }
+    def INPUT_TYPES(cls):
+        return insert_transform_input(cls.clazz.INPUT_TYPES())
 
-    RETURN_TYPES = ("LATENT",)
-    FUNCTION = "sample"
+    FUNCTION = "func"
 
-    CATEGORY = "sampling"
+    def func(self, **kwargs):
+        transforms = kwargs.pop("transform_optional", None)
+        if transforms:
+            kwargs["model"] = attach_transforms(kwargs["model"], transforms)
+        return getattr(self, self.clazz.FUNCTION)(**kwargs)
 
-    def sample(self, model, seed, steps, cfg, sampler_name, scheduler, positive, negative, latent_image, denoise=1.0, transform_optional=None):
 
-        if transform_optional is None:
-            transform_optional = []
+def variations_factory(original_class: type, name=None) -> type:
+    name = name or original_class.__name__ + "Transform"
+    return type(name, (Transforms, original_class), {'clazz': original_class})
 
-        return sample_common(
-            model,
-            seed,
-            steps,
-            cfg,
-            sampler_name,
-            scheduler,
-            positive,
-            negative,
-            latent_image,
-            transform=transform_optional,
-            denoise=denoise)
+TSampler = variations_factory(KSampler)
+TSamplerAdvanced = variations_factory(KSamplerAdvanced)

nodes/KSamplerNodes/Transforms/OffsetCombine.py

@@ -1,4 +1,4 @@
-
+from itertools import chain
 
 
 class OffsetCombine:
@@ -8,6 +8,7 @@ class OffsetCombine:
             "required": {
                 "offset1": ("OFFSET", ),
                 "offset2": ("OFFSET", ),
+                "offset": ("INT", {"default": 0, "min": -10000, "max": 10000}),
             }
         }
 
@@ -16,5 +17,10 @@ class OffsetCombine:
 
     CATEGORY = "sampling/transforms"
 
-    def combine(self, offset1, offset2):
-        return (offset1 + offset2,)
+    def combine(self, offset, **kwargs):
+        offsets = sum(chain([v for k, v in kwargs.items()]), [])
+
+        for o in offsets:
+            o["offset"] += offset
+
+        return (offsets,)

nodes/KSamplerNodes/Transforms/utils.py

@@ -17,14 +17,14 @@ def shift_transform(x0, params):
 
     if params["mode"] == "replace":
         if params["x_shift"] != 0:
-            x = torch.roll(x, shifts=int(x.size()[2] * params["x_shift"]), dims=[2])
+            x = torch.roll(x, shifts=int(x.size()[2] * params["x_shift"]), dims=[3])
         if params["y_shift"] != 0:
-            x = torch.roll(x, shifts=int(x.size()[1] * params["y_shift"]), dims=[1])
+            x = torch.roll(x, shifts=int(x.size()[1] * params["y_shift"]), dims=[2])
     elif params["mode"] == "combine":
         if params["x_shift"] != 0:
-            x = (torch.roll(x, shifts=int(x.size()[2] * params["x_shift"]), dims=[2]) + x) / 2
+            x = (torch.roll(x, shifts=int(x.size()[2] * params["x_shift"]), dims=[3]) + x) / 2
         if params["y_shift"] != 0:
-            x = (torch.roll(x, shifts=int(x.size()[1] * params["y_shift"]), dims=[1]) + x) / 2
+            x = (torch.roll(x, shifts=int(x.size()[1] * params["y_shift"]), dims=[2]) + x) / 2
 
     return x
 
@@ -34,26 +34,26 @@ def mirror_transform(x0, params):
 
     if params["mode"] == "replace":
         if params["direction"] == "vertically":
-            x = torch.flip(x, [1])
-        elif params["direction"] == "horizontally":
             x = torch.flip(x, [2])
+        elif params["direction"] == "horizontally":
+            x = torch.flip(x, [3])
         elif params["direction"] == "both":
-            x = torch.flip(x, [1, 2])
+            x = torch.flip(x, [2, 3])
         elif params["direction"] == "90 degree rotation":
-            x = torch.rot90(x, dims=[1, 2])
+            x = torch.rot90(x, dims=[2, 3])
         elif params["direction"] == "180 degree rotation":
-            x = torch.rot90(torch.rot90(x, dims=[1, 2]), dims=[1, 2])
+            x = torch.rot90(torch.rot90(x, dims=[2, 3]), dims=[2, 3])
     elif params["mode"] == "combine":
         if params["direction"] == "vertically":
-            x = (torch.flip(x, [1]) + x) / 2
-        elif params["direction"] == "horizontally":
             x = (torch.flip(x, [2]) + x) / 2
+        elif params["direction"] == "horizontally":
+            x = (torch.flip(x, [3]) + x) / 2
         elif params["direction"] == "both":
-            x = (torch.flip(x, [1, 2]) + x) / 2
+            x = (torch.flip(x, [2, 3]) + x) / 2
         elif params["direction"] == "90 degree rotation":
-            x = (torch.rot90(x, dims=[1, 2]) + x) / 2
+            x = (torch.rot90(x, dims=[2, 3]) + x) / 2
         elif params["direction"] == "180 degree rotation":
-            x = (torch.rot90(torch.rot90(x, dims=[1, 2]), dims=[1, 2]) + x) / 2
+            x = (torch.rot90(torch.rot90(x, dims=[2, 3]), dims=[2, 3]) + x) / 2
 
     return x
 
@@ -62,11 +62,9 @@ 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 = x0 * params["factor"] + latent * (1 - params["factor"])
+    x = latent * params["factor"] + x0 * (1 - params["factor"])
     x *= params["multiplier"]
 
     return x
@@ -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

nodes/LatentMirror.py

@@ -1,4 +1,5 @@
 import torch
+from nodes import PreviewImage
 
 MIRROR_DIRECTIONS = ["vertically", "horizontally", "both"]
 
@@ -15,6 +16,9 @@ class LatentMirror:
                     "max": 10.0,
                     "step": 0.01
                 })
+            },
+            "optional": {
+                "vae_optional": ("VAE",)
             }
         }
 
@@ -23,7 +27,7 @@ class LatentMirror:
 
     CATEGORY = "latent/advanced"
 
-    def mirror(self, latent, direction, multiplier):
+    def mirror(self, latent, direction, multiplier, vae_optional = None):
         l = latent.copy()
         if direction == "vertically" or direction == "both":
             l["samples"] = torch.flip(l["samples"], dims=[2]) + l["samples"]
@@ -31,4 +35,8 @@ class LatentMirror:
             l["samples"] = torch.flip(l["samples"], dims=[3]) + l["samples"]
 
         l["samples"] *= multiplier
+
+        if vae_optional:
+            return {"result": (l,), "ui": PreviewImage().save_images(vae_optional.decode(l["samples"]))["ui"]}
+
         return (l,)

nodes/LatentShift.py

@@ -1,4 +1,5 @@
 import torch
+from nodes import PreviewImage
 
 
 class LatentShift:
@@ -19,6 +20,9 @@ class LatentShift:
                     "max": 1,
                     "step": 0.01
                 }),
+            },
+            "optional": {
+                "vae_optional": ("VAE",)
             }
         }
 
@@ -27,7 +31,7 @@ class LatentShift:
 
     CATEGORY = "latent/advanced"
 
-    def shift(self, latent, x_shift, y_shift):
+    def shift(self, latent, x_shift, y_shift, vae_optional = None):
         l = latent.copy()
 
         if x_shift != 0:
@@ -35,4 +39,8 @@ class LatentShift:
         if y_shift != 0:
             l["samples"] = torch.roll(l["samples"], shifts=int(l["samples"].size()[2] * y_shift), dims=[2])
 
+        if vae_optional:
+            return {"result": (l,), "ui": PreviewImage().save_images(vae_optional.decode(l["samples"]))["ui"]}
+
         return (l,)
+

nodes/__init__.py

@@ -1,8 +1,9 @@
 from .LatentMirror import LatentMirror
 from .LatentShift import LatentShift
 from .LatentNormalize import LatentNormalize
-from .KSamplerNodes.TSamplerWithTransform import TSamplerWithTransform
-from .KSamplerNodes.TransformSampler import TransformSampler
+from .KSamplerNodes.TransformSampler import TSampler
+from .KSamplerNodes.TransformSampler import TSamplerAdvanced
+from .KSamplerNodes.TransformHijack import TransformHijack
 from .KSamplerNodes.Transforms import MirrorTransform
 from .KSamplerNodes.Transforms import MultiplyTransform
 from .KSamplerNodes.Transforms import ShiftTransform

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.1.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 = ""