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

@@ -1,10 +1,12 @@
 # ComfyUI-Advanced-Latent-Control
 
-**This custom node helps to transform latent in different ways.**
+**This custom nodes helps to transform latent in different ways.**
 
-## Custom Nodes
-### Latent mirror
-This node can flip latent and merge original and flipped version
+>You can access new features earlier by switching from the master branch to dev,
+but you need to remember that there may be some issues on the dev branch and some nodes' behavior may change after release.
+
+## Latent mirror
+This node can flip latent and merge original and flipped version.
 
 **Input:** 
 - `latent`
@@ -17,11 +19,11 @@ 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 axes
+## Latent shift
+This node can shift latent along x and y-axis.
 
 **Input:** 
 - `latent`
@@ -34,42 +36,13 @@ This node can shift latent along x and y axes
 - `latent`
 
 **Usage:**  
-![sample](https://i.imgur.com/1Dp5dSw.png)
+![sample](assets/latent-shift-usage.png)
 
-### KSampler with transforms (Latent Control)
-This node can multiply, mirror and shift latent during generation
+## ~~TSampler with transforms (Latent Control)~~ 
+Removed from version 2.0.0
 
-**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)  
-
-### KSampler (Latent Control)
-This node allows to combine a lot of transforms with different parameters
+## TSampler (Latent Control)
+This node allows to combine a lot of transforms with different parameters.
 
 **Input:**
 - base KSampler fields
@@ -82,10 +55,10 @@ 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
+Multiply, Mirror and Shift transform nodes parameters exactly match the corresponding `KSampler with transforms (Latent Control)` parameters.
 
 There are two new transform nodes:
 - Latent add
@@ -93,7 +66,7 @@ There are two new transform nodes:
 
 They work exactly the same as LatentAdd and LatentBlend nodes from standard node pack, but also, can multiply result by specified number.
 
-### Offset
+## Offset
 You can apply specific offset for transform nodes.
 
 **Fields:**
@@ -105,16 +78,69 @@ 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**
+**0 0 0 1** and **0 0 1** give this pattern: **0 0 1 1 0 1 0 1 1 0 0 1**.
 
-### One time nodes
+## One time nodes
 
 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:**  
+exactly matches the `VAE Decode` node
+
+**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](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](assets/latent-normalize-fix.png)
+
+This node also change some results even if output without this node looks good.
+
+![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](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,10 +1,15 @@
 from .nodes import *
 
+
+WEB_DIRECTORY = "js"
+
 NODE_CLASS_MAPPINGS = {
     "LatentMirror": LatentMirror,
     "LatentShift": LatentShift,
-    "TSamplerWithTransform": TSamplerWithTransform,
-    "TransformSampler": TransformSampler,
+    "LatentNormalize": LatentNormalize,
+    "TransformSampler": TSampler,
+    "TransformSamplerAdvanced": TSamplerAdvanced,
+    "TransformHijack": TransformHijack,
     "MirrorTransform": MirrorTransform,
     "ShiftTransform": ShiftTransform,
     "MultiplyTransform": MultiplyTransform,
@@ -23,16 +28,18 @@ NODE_CLASS_MAPPINGS = {
 NODE_DISPLAY_NAME_MAPPINGS = {
     "LatentMirror": "Latent mirror",
     "LatentShift": "Latent shift",
-    "TSamplerWithTransform": "TSampler with transforms (Latent Control)",
+    "LatentNormalize": "Latent normalize",
     "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], 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/LatentAddTransform.py

@@ -1,4 +1,6 @@
 from .utils import latent_add_transform, get_offset_list
+import comfy
+import torch
 
 
 class LatentAddTransform:
@@ -22,14 +24,14 @@ class LatentAddTransform:
     CATEGORY = "sampling/transforms"
 
     def process(self,
-                offset_optional,
                 latent,
                 start_at=0,
                 stop_at=0,
-                multiplier=1):
+                multiplier=1,
+                offset_optional=None):
         return ([{
             "params": {
-                "latent": latent["samples"][0],
+                "latent": latent["samples"][0].unsqueeze(0),
                 "start_at": start_at,
                 "stop_at": stop_at,
                 "multiplier": multiplier,

nodes/KSamplerNodes/Transforms/LatentIntrpolateTransform.py

@@ -1,4 +1,6 @@
 from .utils import latent_interpolate_transform, get_offset_list
+import comfy
+import torch
 
 
 class LatentInterpolateTransform:
@@ -23,15 +25,15 @@ class LatentInterpolateTransform:
     CATEGORY = "sampling/transforms"
 
     def process(self,
-                offset_optional,
                 latent,
                 start_at=0,
                 stop_at=0,
                 factor=0.5,
-                multiplier=1):
+                multiplier=1,
+                offset_optional=None):
         return ([{
             "params": {
-                "latent": latent["samples"][0],
+                "latent": latent["samples"][0].unsqueeze(0),
                 "start_at": start_at,
                 "stop_at": stop_at,
                 "factor": factor,

nodes/KSamplerNodes/Transforms/MirrorTransform.py

@@ -24,11 +24,11 @@ class MirrorTransform:
     CATEGORY = "sampling/transforms"
 
     def process(self,
-                offset_optional,
                 start_at=0,
                 stop_at=0,
                 mode="replace",
-                direction="horizontally",):
+                direction="horizontally",
+                offset_optional=None):
         return ([{
             "params": {
                 "start_at": start_at,

nodes/KSamplerNodes/Transforms/MultiplyTransform.py

@@ -22,11 +22,11 @@ class MultiplyTransform:
     CATEGORY = "sampling/transforms"
 
     def process(self,
-                offset_optional,
                 start_at=0,
                 stop_at=0,
                 mode="combine",
-                multiplier=1):
+                multiplier=1,
+                offset_optional=None):
         return ([{
             "params": {
                 "start_at": start_at,

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/ShiftTransform.py

@@ -23,12 +23,12 @@ class ShiftTransform:
     CATEGORY = "sampling/transforms"
 
     def process(self,
-                offset_optional,
                 start_at=0,
                 stop_at=0,
                 mode="replace",
                 x_shift=0,
-                y_shift=0):
+                y_shift=0,
+                offset_optional=None):
         return ([{
             "params": {
                 "start_at": start_at,

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,51 +34,47 @@ 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
 
 
 def latent_interpolate_transform(x0, params):
-    latent = params["latent"]
+    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')
-        latent.permute(0, 2, 3, 1)
+        latent = comfy.utils.common_upscale(latent, x0.shape[3], x0.shape[2], 'bicubic', crop='center')
 
-    x = x0 * params["factor"] + latent * (1 - params["factor"])
+    x = latent * params["factor"] + x0 * (1 - params["factor"])
     x *= params["multiplier"]
 
     return x
 
 
 def latent_add_transform(x0, params):
-    latent = params["latent"]
+    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')
-        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/LatentNormalize.py

@@ -0,0 +1,22 @@
+import torch
+from nodes import PreviewImage
+
+class LatentNormalize:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "latent": ("LATENT",),
+                "vae": ("VAE",)
+            }
+        }
+
+    RETURN_TYPES = ("LATENT",)
+    FUNCTION = "normalize"
+
+    CATEGORY = "latent/advanced"
+
+    def normalize(self, latent, vae):
+        image = vae.decode(latent["samples"])
+        sample = vae.encode(image[:,:,:,:3])
+        return {"result": ({"samples": sample},), "ui": PreviewImage().save_images(image)["ui"]}

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,7 +1,9 @@
 from .LatentMirror import LatentMirror
 from .LatentShift import LatentShift
-from .KSamplerNodes.TSamplerWithTransform import TSamplerWithTransform
-from .KSamplerNodes.TransformSampler import TransformSampler
+from .LatentNormalize import LatentNormalize
+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 = ""