From c8bd4f1470ed2ee9ff0eacdd339f9607727cc461 Mon Sep 17 00:00:00 2001 From: Slava Cherkasov Date: Mon, 23 Mar 2026 23:53:25 +0100 Subject: [PATCH] more nodes --- __init__.py | 27 +- ...rtPrompt.png => Wispr_Flow_CZJi9pVooM.png} | Bin js/nodes.js | 341 +++++++++++ nodes.py | 572 +++++++++++++++++- smart_mobius_wan_loop.py | 404 +++++++++++++ smart_wan_first_middle_last_high_noise_end.py | 374 ++++++++++++ 6 files changed, 1699 insertions(+), 19 deletions(-) rename img/{SmartPrompt.png => Wispr_Flow_CZJi9pVooM.png} (100%) create mode 100644 smart_mobius_wan_loop.py create mode 100644 smart_wan_first_middle_last_high_noise_end.py diff --git a/__init__.py b/__init__.py index 4570448..6a0d61d 100644 --- a/__init__.py +++ b/__init__.py @@ -1,4 +1,29 @@ -from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS +from .nodes import NODE_CLASS_MAPPINGS as BASE_NODE_CLASS_MAPPINGS +from .nodes import NODE_DISPLAY_NAME_MAPPINGS as BASE_NODE_DISPLAY_NAME_MAPPINGS +from .smart_wan_first_middle_last_high_noise_end import ( + NODE_CLASS_MAPPINGS as WAN_NODE_CLASS_MAPPINGS, +) +from .smart_wan_first_middle_last_high_noise_end import ( + NODE_DISPLAY_NAME_MAPPINGS as WAN_NODE_DISPLAY_NAME_MAPPINGS, +) +from .smart_mobius_wan_loop import ( + NODE_CLASS_MAPPINGS as MOBIUS_NODE_CLASS_MAPPINGS, +) +from .smart_mobius_wan_loop import ( + NODE_DISPLAY_NAME_MAPPINGS as MOBIUS_NODE_DISPLAY_NAME_MAPPINGS, +) + +NODE_CLASS_MAPPINGS = { + **BASE_NODE_CLASS_MAPPINGS, + **WAN_NODE_CLASS_MAPPINGS, + **MOBIUS_NODE_CLASS_MAPPINGS, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + **BASE_NODE_DISPLAY_NAME_MAPPINGS, + **WAN_NODE_DISPLAY_NAME_MAPPINGS, + **MOBIUS_NODE_DISPLAY_NAME_MAPPINGS, +} WEB_DIRECTORY = "js" diff --git a/img/SmartPrompt.png b/img/Wispr_Flow_CZJi9pVooM.png similarity index 100% rename from img/SmartPrompt.png rename to img/Wispr_Flow_CZJi9pVooM.png diff --git a/js/nodes.js b/js/nodes.js index 5bdcdfd..8f5309d 100644 --- a/js/nodes.js +++ b/js/nodes.js @@ -12,6 +12,263 @@ const nodeWidgetHandlers = { } }; +const SMART_MODEL_LOADER_VISIBILITY = { + SmartModelLoader: { + safetensors: ["unet_name", "weight_dtype"], + gguf: ["gguf_name"], + }, + SmartDualModelLoader: { + safetensors: ["unet_name_1", "unet_name_2", "weight_dtype"], + gguf: ["gguf_name_1", "gguf_name_2"], + }, +}; + +const SMART_BUS_DEFAULT_LABELS = [ + "any 1", + "any 2", + "any 3", + "any 4", + "any 5", + "any 6", + "any 7", + "any 8", + "any 9", + "any 10", +]; + +const SMART_BUS_NODE_CONFIG = { + SmartBusIn: { + inputOffset: 1, + outputOffset: null, + editable: true, + slotCount: 10, + }, + SmartBusOut: { + inputOffset: null, + outputOffset: 1, + editable: false, + slotCount: 10, + }, + SmartBusIn5: { + inputOffset: 1, + outputOffset: null, + editable: true, + slotCount: 5, + }, + SmartBusOut5: { + inputOffset: null, + outputOffset: 1, + editable: false, + slotCount: 5, + }, +}; + +function isSmartBusInNodeType(nodeType) { + return nodeType === "SmartBusIn" || nodeType === "SmartBusIn5"; +} + +function isSmartBusOutNodeType(nodeType) { + return nodeType === "SmartBusOut" || nodeType === "SmartBusOut5"; +} + +function getSmartBusPropertyNames(config) { + return SMART_BUS_DEFAULT_LABELS.slice(0, config.slotCount); +} + +function ensureSmartBusProperties(node, config) { + if (!node.properties) { + node.properties = {}; + } + if (config.editable && !node.properties_info) { + node.properties_info = []; + } + + const propertyNames = getSmartBusPropertyNames(config); + + if (config.editable && node.properties_info) { + node.properties_info = node.properties_info.filter( + (info) => !SMART_BUS_DEFAULT_LABELS.includes(info.name) || propertyNames.includes(info.name) + ); + } + + propertyNames.forEach((defaultLabel) => { + const propertyName = defaultLabel; + const existingValue = node.properties[propertyName]; + if (config.editable && !node.properties_info.some((info) => info.name === propertyName)) { + node.addProperty(propertyName, existingValue ?? defaultLabel, "string"); + } + if (existingValue === undefined || existingValue === null || existingValue === "") { + node.properties[propertyName] = defaultLabel; + } else { + node.properties[propertyName] = existingValue; + } + }); +} + +function getSmartBusLabel(node, index) { + const fallback = SMART_BUS_DEFAULT_LABELS[index]; + const value = node.properties?.[fallback]; + if (typeof value !== "string") { + return fallback; + } + + const trimmed = value.trim(); + return trimmed || fallback; +} + +function copySmartBusLabels(sourceNode, targetNode) { + const sourceConfig = SMART_BUS_NODE_CONFIG[sourceNode.type]; + const targetConfig = SMART_BUS_NODE_CONFIG[targetNode.type]; + if (!sourceConfig || !targetConfig) { + return; + } + + ensureSmartBusProperties(sourceNode, sourceConfig); + ensureSmartBusProperties(targetNode, targetConfig); + + const sharedSlotCount = Math.min(sourceConfig.slotCount, targetConfig.slotCount); + SMART_BUS_DEFAULT_LABELS.slice(0, sharedSlotCount).forEach((defaultLabel, index) => { + targetNode.properties[defaultLabel] = getSmartBusLabel(sourceNode, index); + }); +} + +function getSmartBusSourceNode(node) { + const linkId = node.inputs?.[0]?.link; + if (!linkId || !node.graph?.links) { + return null; + } + + const link = node.graph.links[linkId]; + if (!link) { + return null; + } + + return node.graph.getNodeById?.(link.origin_id) ?? null; +} + +/** Walk bus input chain (BusOut passthrough …) until a SmartBusIn that owns labels. */ +function getSmartBusLabelSourceNode(node) { + let current = getSmartBusSourceNode(node); + const visited = new Set(); + while (current) { + if (visited.has(current.id)) { + return null; + } + visited.add(current.id); + if (isSmartBusInNodeType(current.type)) { + return current; + } + if (isSmartBusOutNodeType(current.type)) { + current = getSmartBusSourceNode(current); + continue; + } + break; + } + return null; +} + +function syncSmartBusOutFromSource(node) { + const labelSource = getSmartBusLabelSourceNode(node); + if (labelSource) { + copySmartBusLabels(labelSource, node); + } else { + const immediate = getSmartBusSourceNode(node); + if (immediate && isSmartBusOutNodeType(immediate.type)) { + copySmartBusLabels(immediate, node); + } + } + + const outLinkIds = node.outputs?.[0]?.links ?? []; + if (!node.graph || !outLinkIds.length) { + return; + } + for (const linkId of outLinkIds) { + const link = node.graph.links?.[linkId]; + if (!link) { + continue; + } + if (link.target_slot != null && link.target_slot !== 0) { + continue; + } + const targetNode = node.graph.getNodeById?.(link.target_id); + if (!targetNode || !isSmartBusOutNodeType(targetNode.type)) { + continue; + } + const targetCfg = SMART_BUS_NODE_CONFIG[targetNode.type]; + syncSmartBusOutFromSource(targetNode); + applySmartBusLabels(targetNode, targetCfg); + } +} + +function syncConnectedSmartBusOutputs(node) { + const linkIds = node.outputs?.[0]?.links ?? []; + if (!node.graph || !linkIds.length) { + return; + } + + for (const linkId of linkIds) { + const link = node.graph.links?.[linkId]; + if (!link) { + continue; + } + + const targetNode = node.graph.getNodeById?.(link.target_id); + if (!targetNode || !isSmartBusOutNodeType(targetNode.type)) { + continue; + } + + const targetCfg = SMART_BUS_NODE_CONFIG[targetNode.type]; + syncSmartBusOutFromSource(targetNode); + applySmartBusLabels(targetNode, targetCfg); + } +} + +function applySmartBusLabels(node, config) { + ensureSmartBusProperties(node, config); + + if (node.inputs?.[0]) { + node.inputs[0].label = "bus"; + } + if (node.outputs?.[0]) { + node.outputs[0].label = "bus"; + } + + SMART_BUS_DEFAULT_LABELS.slice(0, config.slotCount).forEach((_, index) => { + const label = getSmartBusLabel(node, index); + + if (config.inputOffset !== null) { + const input = node.inputs?.[config.inputOffset + index]; + if (input) { + input.label = label; + } + } + + if (config.outputOffset !== null) { + const output = node.outputs?.[config.outputOffset + index]; + if (output) { + output.label = label; + } + } + }); + + if (node.computeSize) { + node.size = node.computeSize(); + } + node.setDirtyCanvas?.(true, true); +} + +function applySmartModelLoaderGgufVisibility(node, config) { + const ggufW = findWidgetByName(node, "gguf"); + if (!ggufW || !config) return; + const isGguf = ggufW.value === true; + for (const name of config.safetensors) { + toggleWidget(node, findWidgetByName(node, name), !isGguf); + } + for (const name of config.gguf) { + toggleWidget(node, findWidgetByName(node, name), isGguf); + } +} + function handleSmartHVLoraStackCount(node, widget) { handleVisibility(node, widget.value, SMART_HV_LORA_STACK_NAME); } @@ -139,5 +396,89 @@ app.registerExtension({ return result; }; } + + const smartLoaderCfg = SMART_MODEL_LOADER_VISIBILITY[nodeData.name]; + if (smartLoaderCfg) { + const origComputeSize = nodeType.prototype.computeSize; + nodeType.prototype.computeSize = function () { + const size = origComputeSize ? origComputeSize.call(this) : [200, 100]; + let height = 50; + for (const w of this.widgets || []) { + if (w.hidden) continue; + height += w.computeSize ? w.computeSize()[1] + 4 : 24; + } + return [size[0], height]; + }; + + const onNodeCreatedLoader = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreatedLoader?.apply(this, arguments); + const ggufW = findWidgetByName(this, "gguf"); + const sync = () => applySmartModelLoaderGgufVisibility(this, smartLoaderCfg); + if (ggufW) { + const prev = ggufW.callback; + ggufW.callback = () => { + prev?.apply(ggufW, arguments); + sync(); + }; + } + sync(); + const onConfigure = this.onConfigure; + this.onConfigure = function () { + const r = onConfigure?.apply(this, arguments); + sync(); + return r; + }; + return result; + }; + } + + const smartBusCfg = SMART_BUS_NODE_CONFIG[nodeData.name]; + if (smartBusCfg) { + const onNodeCreatedBus = nodeType.prototype.onNodeCreated; + nodeType.prototype.onNodeCreated = function () { + const result = onNodeCreatedBus?.apply(this, arguments); + if (isSmartBusOutNodeType(nodeData.name)) { + syncSmartBusOutFromSource(this); + } + applySmartBusLabels(this, smartBusCfg); + return result; + }; + + const onConfigureBus = nodeType.prototype.onConfigure; + nodeType.prototype.onConfigure = function () { + const result = onConfigureBus?.apply(this, arguments); + if (isSmartBusOutNodeType(nodeData.name)) { + syncSmartBusOutFromSource(this); + } + applySmartBusLabels(this, smartBusCfg); + return result; + }; + + const onPropertyChangedBus = nodeType.prototype.onPropertyChanged; + nodeType.prototype.onPropertyChanged = function (property, value, prevValue) { + const result = onPropertyChangedBus?.apply(this, arguments); + const propertyNames = getSmartBusPropertyNames(smartBusCfg); + if (isSmartBusInNodeType(nodeData.name) && propertyNames.includes(property)) { + this.properties[property] = (typeof value === "string" && value.trim()) ? value.trim() : property; + applySmartBusLabels(this, smartBusCfg); + syncConnectedSmartBusOutputs(this); + } + return result; + }; + + const onConnectionsChangeBus = nodeType.prototype.onConnectionsChange; + nodeType.prototype.onConnectionsChange = function () { + const result = onConnectionsChangeBus?.apply(this, arguments); + if (isSmartBusOutNodeType(nodeData.name)) { + syncSmartBusOutFromSource(this); + } + applySmartBusLabels(this, smartBusCfg); + if (isSmartBusInNodeType(nodeData.name)) { + syncConnectedSmartBusOutputs(this); + } + return result; + }; + } } }); diff --git a/nodes.py b/nodes.py index 8b80a7b..4b12794 100644 --- a/nodes.py +++ b/nodes.py @@ -1,6 +1,9 @@ from __future__ import annotations import os +import importlib +import importlib.util +import torch import folder_paths import comfy import comfy.sd # Added import for CLIP type access @@ -13,6 +16,7 @@ import platform from inspect import cleandoc from typing import Any, Dict from server import PromptServer +from comfy_execution.graph_utils import is_link # Custom Nodes for ComfyUI # Note: All nodes in this file should use the "Smart" prefix in their names @@ -319,6 +323,183 @@ class SmartFormatString: class SmartFormatString10(SmartFormatString): max_param_num = 10 # Override with 10 parameters + +SMART_BUS_TYPE = "SMART_BUS" +SMART_BUS_SLOT_LABELS = ( + "any 1", + "any 2", + "any 3", + "any 4", + "any 5", + "any 6", + "any 7", + "any 8", + "any 9", + "any 10", +) + + +def _normalize_smart_bus(bus, size=10): + values = [None] * size + labels = list(SMART_BUS_SLOT_LABELS[:size]) + + if isinstance(bus, dict): + source_values = bus.get("values", []) + source_labels = bus.get("labels", []) + elif isinstance(bus, (list, tuple)): + source_values = bus + source_labels = [] + else: + source_values = [] + source_labels = [] + + for index, value in enumerate(source_values[:size]): + values[index] = value + + for index, label in enumerate(source_labels[:size]): + if isinstance(label, str) and label.strip(): + labels[index] = label.strip() + + return { + "values": values, + "labels": labels, + } + + +def _get_smart_bus_node(unique_id=None, extra_pnginfo=None): + if not unique_id or not extra_pnginfo: + return None + + try: + info = extra_pnginfo[0] if isinstance(extra_pnginfo, (list, tuple)) else extra_pnginfo + workflow = info.get("workflow") + node_value = unique_id[0] if isinstance(unique_id, (list, tuple)) else unique_id + node_id = str(node_value) + except (IndexError, AttributeError, TypeError): + return None + + if not workflow: + return None + + return next((node for node in workflow.get("nodes", []) if str(node.get("id")) == node_id), None) + + +def _get_smart_bus_labels(unique_id=None, extra_pnginfo=None, size=10): + labels = list(SMART_BUS_SLOT_LABELS[:size]) + node = _get_smart_bus_node(unique_id, extra_pnginfo) + if not node: + return labels + + properties = node.get("properties", {}) + for index, default_label in enumerate(labels): + value = properties.get(default_label) + if isinstance(value, str) and value.strip(): + labels[index] = value.strip() + + return labels + + +def _smart_bus_param_is_linked(prompt, unique_id, param_key): + """True if param_* has a wire in the prompt; False if absent/null (keep bus slot).""" + if prompt is None or unique_id is None: + return None + node_value = unique_id[0] if isinstance(unique_id, (list, tuple)) else unique_id + node_id = str(node_value) + node = prompt.get(node_id) + if not isinstance(node, dict): + return None + raw = (node.get("inputs") or {}).get(param_key) + if raw is None: + return False + return is_link(raw) + + +class SmartBusIn: + max_param_num = 10 + + @classmethod + def INPUT_TYPES(cls): + inputs = { + "required": {}, + "optional": { + "bus": (SMART_BUS_TYPE,), + }, + "hidden": { + "unique_id": "UNIQUE_ID", + "extra_pnginfo": "EXTRA_PNGINFO", + "prompt": "PROMPT", + }, + } + + for i in range(1, cls.max_param_num + 1): + inputs["optional"][f"param_{i}"] = (any_type,) + + return inputs + + RETURN_TYPES = (SMART_BUS_TYPE,) + RETURN_NAMES = ("bus",) + FUNCTION = "build_bus" + CATEGORY = "SmartHelperNodes/Bus" + DESCRIPTION = "Merge into a Smart Bus: only wired param slots replace values; unwired slots keep the incoming bus." + + def build_bus(self, bus=None, unique_id=None, extra_pnginfo=None, prompt=None, **kwargs): + bus_data = _normalize_smart_bus(bus, self.max_param_num) + values = bus_data["values"] + + for i in range(1, self.max_param_num + 1): + key = f"param_{i}" + linked = _smart_bus_param_is_linked(prompt, unique_id, key) + if linked is True: + if key in kwargs: + values[i - 1] = kwargs[key] + elif linked is False: + pass + else: + if key in kwargs: + values[i - 1] = kwargs[key] + + labels = _get_smart_bus_labels(unique_id, extra_pnginfo, self.max_param_num) + + return ({ + "values": values, + "labels": labels, + },) + + +class SmartBusOut: + max_param_num = 10 + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "bus": (SMART_BUS_TYPE,), + } + } + + RETURN_TYPES = (SMART_BUS_TYPE,) + (any_type,) * max_param_num + RETURN_NAMES = ("bus",) + SMART_BUS_SLOT_LABELS + FUNCTION = "expand_bus" + CATEGORY = "SmartHelperNodes/Bus" + DESCRIPTION = "Expand a Smart Bus back into its 10 stored passthrough values." + + def expand_bus(self, bus): + bus_data = _normalize_smart_bus(bus, self.max_param_num) + values = bus_data["values"] + return (bus_data, *values) + + +class SmartBusIn5(SmartBusIn): + max_param_num = 5 + DESCRIPTION = "Merge into a Smart Bus: only wired param slots replace values; unwired slots keep the incoming bus." + + +class SmartBusOut5(SmartBusOut): + max_param_num = 5 + RETURN_TYPES = (SMART_BUS_TYPE,) + (any_type,) * max_param_num + RETURN_NAMES = ("bus",) + SMART_BUS_SLOT_LABELS[:max_param_num] + DESCRIPTION = "Expand a Smart Bus back into its 5 stored passthrough values." + class SmartSaveText: """ A node that saves text to a file, creating directories as needed and appending to existing files. @@ -347,16 +528,18 @@ class SmartSaveText: directory = os.path.dirname(full_path) if directory: os.makedirs(directory, exist_ok=True) - - # Append to file (or create if doesn't exist) - with open(full_path, "a+", encoding="utf-8") as f: - # Add newline if file is not empty and doesn't end with one - if f.tell() != 0: - f.seek(0) - content = f.read() - if content and not content.endswith('\n'): - f.write('\n') - # Write the new text + + # Check if we need a leading newline (file exists, non-empty, doesn't end with newline) + needs_newline = False + if os.path.exists(full_path) and os.path.getsize(full_path) > 0: + with open(full_path, "rb") as f: + f.seek(-1, 2) # Seek to last byte from end + needs_newline = f.read(1) != b'\n' + + # Append to file + with open(full_path, "a", encoding="utf-8") as f: + if needs_newline: + f.write('\n') f.write(text) return (text,) @@ -642,19 +825,21 @@ class SmartLoadLoRA: "strength_clip": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01, "tooltip": "How strongly to modify the CLIP model. This value can be negative."}), }, "optional": { + "model2": ("MODEL", {"tooltip": "Optional second diffusion model; the same LoRA is applied to both models (UNet only on this input)."}), "lora_string": ("STRING", {"forceInput": True, "default": ""}) } } - RETURN_TYPES = ("MODEL", "CLIP", "STRING") - RETURN_NAMES = ("model", "clip", "lora_string") + RETURN_TYPES = ("MODEL", "CLIP", "STRING", "MODEL") + RETURN_NAMES = ("model", "clip", "lora_string", "model2") FUNCTION = "load_lora" CATEGORY = "SmartHelperNodes" - DESCRIPTION = "Load a LoRA model (just like the core LoRA loader), but also outputs the model name and strength to a formatted string." + DESCRIPTION = "Load a LoRA model (just like the core LoRA loader), but also outputs the model name and strength to a formatted string. Optional second model applies the same UNet LoRA." - def load_lora(self, model, clip, lora_name, strength_model, strength_clip, lora_string=""): + def load_lora(self, model, clip, lora_name, strength_model, strength_clip, model2=None, lora_string=""): if strength_model == 0 and strength_clip == 0: - return (model, clip, lora_string) + m2_out = model2 if model2 is not None else model + return (model, clip, lora_string, m2_out) lora_path = folder_paths.get_full_path("loras", lora_name) lora = None @@ -669,7 +854,12 @@ class SmartLoadLoRA: self.loaded_lora = (lora_path, lora) model_lora, clip_lora = comfy.sd.load_lora_for_models(model, clip, lora, strength_model, strength_clip) - + + if model2 is not None: + model2_lora, _ = comfy.sd.load_lora_for_models(model2, None, lora, strength_model, 0) + else: + model2_lora = model_lora + # Create formatted string output lora_name_without_ext = lora_name.split(".")[0] new_lora_string = f"{lora_name_without_ext}, model_str:{strength_model}, clip_str:{strength_clip}" @@ -677,8 +867,79 @@ class SmartLoadLoRA: final_lora_string = f"{lora_string}\n{new_lora_string}" else: final_lora_string = new_lora_string - - return (model_lora, clip_lora, final_lora_string) + + return (model_lora, clip_lora, final_lora_string, model2_lora) + + +class SmartLoadDoubleLoRA: + def __init__(self): + self.loaded_lora_1 = None + self.loaded_lora_2 = None + + @classmethod + def INPUT_TYPES(s): + lora_list = folder_paths.get_filename_list("loras") + return { + "required": { + "model1": ("MODEL", {"tooltip": "First diffusion model."}), + "model2": ("MODEL", {"tooltip": "Second diffusion model."}), + "clip": ("CLIP", {"tooltip": "The CLIP model both LoRAs will be applied to."}), + "lora_name_1": (lora_list, {"tooltip": "LoRA applied to model1 and CLIP."}), + "strength_model_1": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), + "strength_clip_1": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), + "lora_name_2": (lora_list, {"tooltip": "LoRA applied to model2 and CLIP."}), + "strength_model_2": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), + "strength_clip_2": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01}), + }, + "optional": { + "lora_string": ("STRING", {"forceInput": True, "default": ""}) + } + } + + RETURN_TYPES = ("MODEL", "MODEL", "CLIP", "STRING") + RETURN_NAMES = ("model1", "model2", "clip", "lora_string") + FUNCTION = "load_double_lora" + CATEGORY = "SmartHelperNodes" + DESCRIPTION = "Apply two different LoRAs: lora1 to model1, lora2 to model2, and both to CLIP sequentially." + + def _get_lora(self, cache_attr, lora_name): + lora_path = folder_paths.get_full_path("loras", lora_name) + cached = getattr(self, cache_attr) + if cached is not None and cached[0] == lora_path: + return cached[1] + lora = comfy.utils.load_torch_file(lora_path, safe_load=True) + setattr(self, cache_attr, (lora_path, lora)) + return lora + + def load_double_lora(self, model1, model2, clip, + lora_name_1, strength_model_1, strength_clip_1, + lora_name_2, strength_model_2, strength_clip_2, + lora_string=""): + # LoRA 1 → model1 + clip + if strength_model_1 == 0 and strength_clip_1 == 0: + m1_out, clip_out = model1, clip + else: + lora1 = self._get_lora("loaded_lora_1", lora_name_1) + m1_out, clip_out = comfy.sd.load_lora_for_models(model1, clip, lora1, strength_model_1, strength_clip_1) + + # LoRA 2 → model2 + clip (already patched by lora1) + if strength_model_2 == 0 and strength_clip_2 == 0: + m2_out = model2 + else: + lora2 = self._get_lora("loaded_lora_2", lora_name_2) + m2_out, clip_out = comfy.sd.load_lora_for_models(model2, clip_out, lora2, strength_model_2, strength_clip_2) + + # Build lora_string + parts = [] + if lora_string: + parts.append(lora_string) + n1 = lora_name_1.split(".")[0] + parts.append(f"{n1}, model_str:{strength_model_1}, clip_str:{strength_clip_1}") + n2 = lora_name_2.split(".")[0] + parts.append(f"{n2}, model_str:{strength_model_2}, clip_str:{strength_clip_2}") + + return (m1_out, m2_out, clip_out, "\n".join(parts)) + class SmartPrompt: """ @@ -883,6 +1144,145 @@ class SmartShowAnything: else: return {"ui": {"text": values}, "result": (values,), } +class SmartSetInt: + """ + A node that sets the output to a specific integer value. + When bypassed, it passes through the input value instead. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input": ("INT", {"default": 0, "tooltip": "Input value (used when node is bypassed)"}), + "value": ("INT", {"default": 0, "tooltip": "Output value (used when node is active)"}), + }, + } + + RETURN_TYPES = ("INT",) + RETURN_NAMES = ("output",) + FUNCTION = "set_value" + CATEGORY = "SmartHelperNodes" + DESCRIPTION = "Sets the output to a specific integer value. When bypassed, passes through the input value." + + def set_value(self, input, value): + return (value,) + + @classmethod + def check_lazy_status(cls, input, value): + # This indicates which inputs are needed + # When bypassed, ComfyUI will pass through the input + return ["value"] + +class SmartSetFloat: + """ + A node that sets the output to a specific float value. + When bypassed, it passes through the input value instead. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "input": ("FLOAT", {"default": 0.0, "tooltip": "Input value (used when node is bypassed)"}), + "value": ("FLOAT", {"default": 0.0, "tooltip": "Output value (used when node is active)"}), + }, + } + + RETURN_TYPES = ("FLOAT",) + RETURN_NAMES = ("output",) + FUNCTION = "set_value" + CATEGORY = "SmartHelperNodes" + DESCRIPTION = "Sets the output to a specific float value. When bypassed, passes through the input value." + + def set_value(self, input, value): + return (value,) + + @classmethod + def check_lazy_status(cls, input, value): + # This indicates which inputs are needed + # When bypassed, ComfyUI will pass through the input + return ["value"] + +class SmartModelInfo: + """ + A node that displays information about a MODEL object and passes it through. + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL", {"tooltip": "Model to inspect"}), + }, + } + + RETURN_TYPES = ("MODEL", "STRING") + RETURN_NAMES = ("model", "info") + FUNCTION = "get_model_info" + CATEGORY = "SmartHelperNodes" + DESCRIPTION = "Displays information about a MODEL object including config details" + OUTPUT_NODE = True + + def get_model_info(self, model): + info_lines = [] + + # Try to get model config info + try: + if hasattr(model, 'model') and hasattr(model.model, 'model_config'): + config = model.model.model_config + info_lines.append("=== Model Config ===") + + # Get model type/class name + if hasattr(config, '__class__'): + info_lines.append(f"Type: {config.__class__.__name__}") + + # Get unet config if available + if hasattr(config, 'unet_config'): + unet_config = config.unet_config + if isinstance(unet_config, dict): + info_lines.append("\n=== UNet Config ===") + for key, value in unet_config.items(): + if isinstance(value, (int, float, str, bool)): + info_lines.append(f"{key}: {value}") + + # Check for latent format + if hasattr(config, 'latent_format'): + latent_format = config.latent_format + if hasattr(latent_format, '__class__'): + info_lines.append(f"\nLatent Format: {latent_format.__class__.__name__}") + except Exception as e: + info_lines.append(f"Error reading model config: {str(e)}") + + # Check model_options + try: + if hasattr(model, 'model_options'): + options = model.model_options + if options and isinstance(options, dict): + info_lines.append("\n=== Model Options ===") + for key, value in options.items(): + if key != 'transformer_options' and isinstance(value, (int, float, str, bool)): + info_lines.append(f"{key}: {value}") + except Exception as e: + info_lines.append(f"Error reading model options: {str(e)}") + + # Check attachments (where custom data might be stored) + try: + if hasattr(model, 'attachments'): + attachments = model.attachments + if attachments and isinstance(attachments, dict): + info_lines.append("\n=== Attachments ===") + for key, value in attachments.items(): + info_lines.append(f"{key}: {value}") + except Exception as e: + info_lines.append(f"Error reading attachments: {str(e)}") + + if not info_lines: + info_lines.append("No model information available") + + info_text = "\n".join(info_lines) + return (model, info_text) + class SmartHWMonitor: """ A node that monitors hardware metrics including CPU/GPU temperature and RAM usage. @@ -1207,6 +1607,122 @@ class SmartHWMonitor: import time return int(time.time() / refresh_rate) + +_SMARTHELPER_GGUF_PKG = "SmartHelperGGUFPkg" + + +def _get_unet_gguf_filenames(): + try: + return folder_paths.get_filename_list("unet_gguf") + except KeyError: + return [] + + +def _ensure_comfyui_gguf_package(): + if _SMARTHELPER_GGUF_PKG in sys.modules: + return + root = None + for custom_root in folder_paths.get_folder_paths("custom_nodes"): + for name in ("ComfyUI-GGUF", "comfyui-gguf"): + candidate = os.path.join(custom_root, name) + if os.path.isfile(os.path.join(candidate, "__init__.py")): + root = candidate + break + if root: + break + if not root: + raise RuntimeError( + "ComfyUI-GGUF is required for GGUF mode. Install from https://github.com/city96/ComfyUI-GGUF" + ) + init_py = os.path.join(root, "__init__.py") + spec = importlib.util.spec_from_file_location( + _SMARTHELPER_GGUF_PKG, + init_py, + submodule_search_locations=[root], + ) + mod = importlib.util.module_from_spec(spec) + sys.modules[_SMARTHELPER_GGUF_PKG] = mod + spec.loader.exec_module(mod) + + +def _load_gguf_model(unet_name: str): + _ensure_comfyui_gguf_package() + nodes_mod = importlib.import_module(f"{_SMARTHELPER_GGUF_PKG}.nodes") + return nodes_mod.UnetLoaderGGUF().load_unet(unet_name)[0] + + +def _load_safetensors_diffusion_model(unet_name: str, weight_dtype: str): + model_options = {} + if weight_dtype == "fp8_e4m3fn": + model_options["dtype"] = torch.float8_e4m3fn + elif weight_dtype == "fp8_e4m3fn_fast": + model_options["dtype"] = torch.float8_e4m3fn + model_options["fp8_optimizations"] = True + elif weight_dtype == "fp8_e5m2": + model_options["dtype"] = torch.float8_e5m2 + unet_path = folder_paths.get_full_path_or_raise("diffusion_models", unet_name) + return comfy.sd.load_diffusion_model(unet_path, model_options=model_options) + + +class SmartModelLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "gguf": ("BOOLEAN", {"default": False}), + "unet_name": (folder_paths.get_filename_list("diffusion_models"),), + "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],), + "gguf_name": (_get_unet_gguf_filenames(),), + } + } + + RETURN_TYPES = ("MODEL", "STRING") + RETURN_NAMES = ("model", "model_name") + FUNCTION = "load" + CATEGORY = "SmartHelperNodes" + DESCRIPTION = "Load a diffusion model from .safetensors (Load Diffusion Model) or .gguf (Unet Loader GGUF)." + + def load(self, gguf, unet_name, weight_dtype, gguf_name): + if gguf: + model = _load_gguf_model(gguf_name) + return (model, gguf_name) + model = _load_safetensors_diffusion_model(unet_name, weight_dtype) + return (model, unet_name) + + +class SmartDualModelLoader: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "gguf": ("BOOLEAN", {"default": False}), + "unet_name_1": (folder_paths.get_filename_list("diffusion_models"),), + "unet_name_2": (folder_paths.get_filename_list("diffusion_models"),), + "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e4m3fn_fast", "fp8_e5m2"],), + "gguf_name_1": (_get_unet_gguf_filenames(),), + "gguf_name_2": (_get_unet_gguf_filenames(),), + } + } + + RETURN_TYPES = ("MODEL", "MODEL", "STRING") + RETURN_NAMES = ("model_1", "model_2", "model_names") + FUNCTION = "load" + CATEGORY = "SmartHelperNodes" + DESCRIPTION = "Load two diffusion models (safetensors or GGUF) with a combined name summary." + + def load(self, gguf, unet_name_1, unet_name_2, weight_dtype, gguf_name_1, gguf_name_2): + if gguf: + m1 = _load_gguf_model(gguf_name_1) + m2 = _load_gguf_model(gguf_name_2) + n1, n2 = gguf_name_1, gguf_name_2 + else: + m1 = _load_safetensors_diffusion_model(unet_name_1, weight_dtype) + m2 = _load_safetensors_diffusion_model(unet_name_2, weight_dtype) + n1, n2 = unet_name_1, unet_name_2 + text = f"Model #1: {n1}\nModel #2: {n2}" + return (m1, m2, text) + + NODE_CLASS_MAPPINGS = { "SmartHVLoraSelect": SmartHVLoraSelect, "SmartHVLoraStack": SmartHVLoraStack, @@ -1216,10 +1732,20 @@ NODE_CLASS_MAPPINGS = { "SmartSaveAnimatedGIF": SmartSaveAnimatedGIF, "SmartRemoveComments": SmartRemoveComments, "SmartLoadLoRA": SmartLoadLoRA, + "SmartLoadDoubleLoRA": SmartLoadDoubleLoRA, "SmartPrompt": SmartPrompt, "SmartModelOrLoraToString": SmartModelOrLoraToString, "SmartShowAnything": SmartShowAnything, "SmartHWMonitor": SmartHWMonitor, + "SmartSetInt": SmartSetInt, + "SmartSetFloat": SmartSetFloat, + "SmartBusIn": SmartBusIn, + "SmartBusOut": SmartBusOut, + "SmartBusIn5": SmartBusIn5, + "SmartBusOut5": SmartBusOut5, + "SmartModelInfo": SmartModelInfo, + "SmartModelLoader": SmartModelLoader, + "SmartDualModelLoader": SmartDualModelLoader, } NODE_DISPLAY_NAME_MAPPINGS = { @@ -1231,8 +1757,18 @@ NODE_DISPLAY_NAME_MAPPINGS = { "SmartSaveAnimatedGIF": "Smart Save Animated GIF", "SmartRemoveComments": "Smart Remove Comments", "SmartLoadLoRA": "Smart Load LoRA", + "SmartLoadDoubleLoRA": "Smart Load Double LoRA", "SmartPrompt": "Smart Prompt", "SmartModelOrLoraToString": "Smart Model or Lora to String", "SmartShowAnything": "Smart Show Anything", "SmartHWMonitor": "Smart Hardware Monitor", + "SmartSetInt": "Smart Set Int", + "SmartSetFloat": "Smart Set Float", + "SmartBusIn": "Smart Bus In", + "SmartBusOut": "Smart Bus Out", + "SmartBusIn5": "Smart Bus In 5", + "SmartBusOut5": "Smart Bus Out 5", + "SmartModelInfo": "Smart Model Info", + "SmartModelLoader": "Smart Model Loader", + "SmartDualModelLoader": "Smart Dual Model Loader", } diff --git a/smart_mobius_wan_loop.py b/smart_mobius_wan_loop.py new file mode 100644 index 0000000..b882c9f --- /dev/null +++ b/smart_mobius_wan_loop.py @@ -0,0 +1,404 @@ +import math +import torch +import torch.nn.functional as F +import comfy.clip_vision +import comfy.model_management as mm +import comfy.utils +import node_helpers + + +def _normalize_image_output(decoded): + """Convert decoded VAE output to ComfyUI's standard IMAGE shape.""" + if decoded.ndim == 5 and decoded.shape[0] == 1: + return decoded[0] + return decoded + + +def _find_coprime_skip(n): + """Find a shift_skip value > n//2 that is coprime with n, ensuring full position coverage.""" + candidate = n // 2 + 1 + while candidate < n: + if math.gcd(candidate, n) == 1: + return candidate + candidate += 1 + return 1 + + +class MobiusLatentShiftWrapper: + """ + Mobius latent shift adapted for WAN FLF conditioning. + + Rolls the noisy latent each step so the model sees varying temporal RoPE positions. + Optional concat-conditioning rolling can be enabled experimentally when desired. + """ + + def __init__(self, shift_skip, latent_frames, roll_concat_conditioning=False): + self.shift_skip = shift_skip + self.latent_frames = latent_frames + self.shift_idx = 0 + self.active_shift = 0 + self.last_sigma = None + self.roll_concat_conditioning = roll_concat_conditioning + + def __call__(self, apply_model_func, args_dict): + sigma = args_dict["timestep"] + + if self.last_sigma is None or not torch.equal(sigma, self.last_sigma): + self.last_sigma = sigma.clone() + self.active_shift = self.shift_idx + self.shift_idx = (self.shift_idx + self.shift_skip) % self.latent_frames + + if self.active_shift == 0: + return apply_model_func(args_dict["input"], args_dict["timestep"], **args_dict["c"]) + + shift = -self.active_shift + x = torch.roll(args_dict["input"], shift, dims=2) + c = dict(args_dict["c"]) + + # Rolling FLF concat conditioning makes the fixed-frame constraints move + # through time, which can destabilize the tail of the clip. + if self.roll_concat_conditioning and "c_concat" in c and c["c_concat"] is not None: + c_concat = c["c_concat"] + if hasattr(c_concat, "cond"): + rolled = torch.roll(c_concat.cond, shift, dims=2) + c["c_concat"] = c_concat.__class__(rolled) + else: + c["c_concat"] = torch.roll(c_concat, shift, dims=2) + + output = apply_model_func(x, args_dict["timestep"], **c) + return torch.roll(output, -shift, dims=2) + + +class LoopBoundaryBlend: + """ + Post-CFG callback that blends end-frame predictions toward start-frame predictions. + + For a seamless loop (start == end image), the denoised predictions at the temporal + boundaries should converge. This corrects the known WAN end-frame noise artifact + by gradually replacing the last K frames' predictions with the first K frames'. + """ + + def __init__(self, blend_frames, blend_strength): + self.blend_frames = blend_frames + self.blend_strength = blend_strength + + def __call__(self, args): + denoised = args["denoised"] + T = denoised.shape[2] + K = min(self.blend_frames, T // 4) + + if K <= 0 or self.blend_strength <= 0: + return denoised + + ramp = torch.linspace(0.0, self.blend_strength, K, device=denoised.device, dtype=denoised.dtype) + ramp = ramp.view(1, 1, K, 1, 1) + + result = denoised.clone() + result[:, :, -K:] = denoised[:, :, -K:] * (1.0 - ramp) + denoised[:, :, :K] * ramp + return result + + +class SmartMobiusWanLoop: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_high": ("MODEL",), + "model_low": ("MODEL",), + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), + "width": ("INT", {"default": 832, "min": 16, "max": 8192, "step": 16}), + "height": ("INT", {"default": 480, "min": 16, "max": 8192, "step": 16}), + "length": ("INT", {"default": 81, "min": 1, "max": 8192, "step": 4}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096}), + }, + "optional": { + "start_image": ("IMAGE",), + "end_image": ("IMAGE",), + "clip_vision_start_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_end_image": ("CLIP_VISION_OUTPUT",), + "motion_amplitude": ( + "FLOAT", + {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05, + "tooltip": "1.0 = official FLF, 2.0 = max structural repulsion boost"}, + ), + "enable_latent_shift": ( + "BOOLEAN", + {"default": True, + "tooltip": "Mobius-style circular latent rolling each denoising step"}, + ), + "shift_skip": ( + "INT", + {"default": 0, "min": 0, "max": 100, "step": 1, + "tooltip": "Roll stride per step. 0 = auto (coprime with latent frame count)"}, + ), + "enable_loop_blend": ( + "BOOLEAN", + {"default": True, + "tooltip": "Blend end-frame predictions toward start-frame predictions after CFG"}, + ), + "blend_frames": ( + "INT", + {"default": 5, "min": 1, "max": 20, "step": 1, + "tooltip": "Number of end frames to correct via boundary blending"}, + ), + "blend_strength": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "0 = no correction, 1 = fully replace end with start predictions"}, + ), + "enable_circular_rope": ( + "BOOLEAN", + {"default": True, + "tooltip": "Compress temporal RoPE so last frame is closer to first"}, + ), + "roll_concat_conditioning": ( + "BOOLEAN", + { + "default": False, + "tooltip": "Experimental. Roll concat_mask/concat_latent_image together with the latent cycle." + }, + ), + "circular_rope_strength": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05, + "tooltip": "0 = no change, 1 = maximum temporal compression"}, + ), + }, + } + + RETURN_TYPES = ("MODEL", "MODEL", "CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("model_high", "model_low", "positive", "negative", "latent") + FUNCTION = "execute" + CATEGORY = "SmartHelperNodes/Wan" + DESCRIPTION = ( + "Mobius seamless loop for WAN 2.2 dual-pass (high/low noise). " + "Three independent mechanisms: (1) Latent shift — cycles frame positions across " + "denoising steps to distribute RoPE bias. (2) Loop blend — corrects end-frame noise " + "by blending toward start-frame predictions. (3) Circular RoPE — compresses temporal " + "positions so the last frame is closer to the first in attention space." + ) + + def execute( + self, + model_high, + model_low, + positive, + negative, + vae, + width, + height, + length, + batch_size, + start_image=None, + end_image=None, + clip_vision_start_image=None, + clip_vision_end_image=None, + motion_amplitude=1.0, + enable_latent_shift=True, + shift_skip=0, + enable_loop_blend=True, + blend_frames=5, + blend_strength=0.5, + enable_circular_rope=True, + roll_concat_conditioning=False, + circular_rope_strength=0.5, + ): + spacial_scale = vae.spacial_compression_encode() + latent_frames = ((length - 1) // 4) + 1 + + latent = torch.zeros( + [batch_size, vae.latent_channels, latent_frames, height // spacial_scale, width // spacial_scale], + device=mm.intermediate_device(), + ) + + # --- Image preprocessing --- + if start_image is not None: + start_image = comfy.utils.common_upscale( + start_image[:length].movedim(-1, 1), width, height, "bilinear", "center" + ).movedim(1, -1) + if end_image is not None: + end_image = comfy.utils.common_upscale( + end_image[-length:].movedim(-1, 1), width, height, "bilinear", "center" + ).movedim(1, -1) + + # --- Build reference image sequence (gray fill + start/end frames) --- + official_image = torch.ones((length, height, width, 3), device=mm.intermediate_device()) * 0.5 + mask = torch.ones( + (1, 1, latent_frames * 4, height // spacial_scale, width // spacial_scale), + device=mm.intermediate_device(), + ) + + if start_image is not None: + official_image[:start_image.shape[0]] = start_image + mask[:, :, :start_image.shape[0] + 3] = 0.0 + if end_image is not None: + official_image[-end_image.shape[0]:] = end_image + mask[:, :, -end_image.shape[0]:] = 0.0 + + official_latent = vae.encode(official_image[:, :, :, :3]) + + # --- Inverse structural repulsion (PainterFLF2V algorithm) --- + if length > 2 and motion_amplitude > 1.001 and start_image is not None and end_image is not None: + start_l = official_latent[:, :, 0:1] + end_l = official_latent[:, :, -1:] + t = torch.linspace(0.0, 1.0, official_latent.shape[2], device=official_latent.device).view(1, 1, -1, 1, 1) + linear_latent = start_l * (1 - t) + end_l * t + + diff = official_latent - linear_latent + h, w = diff.shape[-2], diff.shape[-1] + low_freq_diff = F.interpolate( + diff.view(-1, vae.latent_channels, h, w), size=(max(1, h // 8), max(1, w // 8)), mode="area" + ) + low_freq_diff = F.interpolate(low_freq_diff, size=(h, w), mode="bilinear") + low_freq_diff = low_freq_diff.view_as(diff) + high_freq_diff = diff - low_freq_diff + + boost_scale = (motion_amplitude - 1.0) * 4.0 + concat_latent_image = official_latent + (high_freq_diff * boost_scale) + else: + concat_latent_image = official_latent + + # --- Mask reshape: [1,1,T*4,H,W] -> [1,4,T,H,W] --- + mask = mask.view(1, mask.shape[2] // 4, 4, mask.shape[3], mask.shape[4]).transpose(1, 2) + + # --- Inject conditioning --- + positive = node_helpers.conditioning_set_values( + positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask} + ) + negative = node_helpers.conditioning_set_values( + negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask} + ) + + # --- CLIP vision merge --- + clip_vision_output = None + if clip_vision_start_image is not None: + clip_vision_output = clip_vision_start_image + if clip_vision_end_image is not None: + if clip_vision_output is not None: + states = torch.cat( + [clip_vision_output.penultimate_hidden_states, + clip_vision_end_image.penultimate_hidden_states], + dim=-2, + ) + clip_vision_output = comfy.clip_vision.Output() + clip_vision_output.penultimate_hidden_states = states + else: + clip_vision_output = clip_vision_end_image + if clip_vision_output is not None: + positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output}) + negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output}) + + # --- Model patching (shared state across both models) --- + patched_high = model_high.clone() + patched_low = model_low.clone() + + if enable_latent_shift: + effective_skip = shift_skip if shift_skip > 0 else _find_coprime_skip(latent_frames) + if math.gcd(effective_skip, latent_frames) != 1: + effective_skip = _find_coprime_skip(latent_frames) + wrapper = MobiusLatentShiftWrapper( + effective_skip, + latent_frames, + roll_concat_conditioning=roll_concat_conditioning, + ) + patched_high.set_model_unet_function_wrapper(wrapper) + patched_low.set_model_unet_function_wrapper(wrapper) + + if enable_loop_blend and blend_strength > 0 and blend_frames > 0: + blend_fn = LoopBoundaryBlend(blend_frames, blend_strength) + patched_high.set_model_sampler_post_cfg_function(blend_fn) + patched_low.set_model_sampler_post_cfg_function(blend_fn) + + if enable_circular_rope and circular_rope_strength > 0.0: + scale_t = 1.0 - (circular_rope_strength / latent_frames) + for m in (patched_high, patched_low): + to = m.model_options.get("transformer_options", {}) + existing_rope = to.get("rope_options", {}) + existing_rope["scale_t"] = scale_t + to["rope_options"] = existing_rope + m.model_options["transformer_options"] = to + + out_latent = {"samples": latent} + return (patched_high, patched_low, positive, negative, out_latent) + + +class SmartMobiusVAEDecode: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "samples": ("LATENT",), + "vae": ("VAE",), + }, + "optional": { + "enable_frame_invariance": ( + "BOOLEAN", + {"default": True, "tooltip": "Decode boundary region with wrapped context to fix 3D VAE seam artifacts"}, + ), + "boundary_frames": ( + "INT", + {"default": 3, "min": 2, "max": 5, "step": 1, + "tooltip": "Latent frames to wrap at each boundary side"}, + ), + }, + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image",) + FUNCTION = "execute" + CATEGORY = "SmartHelperNodes/Wan" + DESCRIPTION = ( + "VAE decode with optional frame-invariance for seamless loops. " + "Wraps boundary latent frames and decodes them with full temporal context " + "to eliminate 3D VAE seam artifacts at the loop point." + ) + + def execute(self, samples, vae, enable_frame_invariance=True, boundary_frames=3): + latent = samples["samples"] + + if not enable_frame_invariance or latent.shape[2] < boundary_frames * 2 + 1: + decoded = _normalize_image_output(vae.decode(latent)) + return (decoded,) + + k = boundary_frames + + main_pixels = vae.decode(latent) + + boundary_latent = torch.cat([latent[:, :, -k:], latent[:, :, :k]], dim=2) + boundary_pixels = vae.decode(boundary_latent) + frame_dim = 1 if boundary_pixels.ndim == 5 else 0 + + left_pixel_count = 1 + (k - 1) * 4 + tail_slice = [slice(None)] * boundary_pixels.ndim + tail_slice[frame_dim] = slice(0, left_pixel_count) + tail_good = boundary_pixels[tuple(tail_slice)] + + head_slice = [slice(None)] * boundary_pixels.ndim + head_slice[frame_dim] = slice(left_pixel_count, None) + head_good = boundary_pixels[tuple(head_slice)] + + result = main_pixels.clone() + frame_dim = 1 if result.ndim == 5 else 0 + + tail_slice = [slice(None)] * result.ndim + tail_slice[frame_dim] = slice(-tail_good.shape[frame_dim], None) + result[tuple(tail_slice)] = tail_good + + head_slice = [slice(None)] * result.ndim + head_slice[frame_dim] = slice(0, head_good.shape[frame_dim]) + result[tuple(head_slice)] = head_good + + return (_normalize_image_output(result),) + + +NODE_CLASS_MAPPINGS = { + "SmartMobiusWanLoop": SmartMobiusWanLoop, + "SmartMobiusVAEDecode": SmartMobiusVAEDecode, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SmartMobiusWanLoop": "Smart Mobius WAN Loop", + "SmartMobiusVAEDecode": "Smart Mobius VAE Decode", +} diff --git a/smart_wan_first_middle_last_high_noise_end.py b/smart_wan_first_middle_last_high_noise_end.py new file mode 100644 index 0000000..5886cb5 --- /dev/null +++ b/smart_wan_first_middle_last_high_noise_end.py @@ -0,0 +1,374 @@ +import comfy +import comfy.clip_vision +import comfy.model_management +import comfy.utils +import node_helpers +import torch +import torch.nn.functional as F + + +class SmartWanFirstMiddleLastFrameToVideoHighNoiseEnd: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "positive": ("CONDITIONING",), + "negative": ("CONDITIONING",), + "vae": ("VAE",), + "width": ("INT", {"default": 832, "min": 16, "max": 8192, "step": 16, "display": "number"}), + "height": ("INT", {"default": 480, "min": 16, "max": 8192, "step": 16, "display": "number"}), + "length": ("INT", {"default": 81, "min": 1, "max": 8192, "step": 4, "display": "number"}), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 4096, "step": 1, "display": "number"}), + }, + "optional": { + "mode": (["NORMAL", "SINGLE_PERSON"], {"default": "NORMAL"}), + "start_image": ("IMAGE",), + "middle_image": ("IMAGE",), + "end_image": ("IMAGE",), + "middle_frame_ratio": ( + "FLOAT", + {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01, "round": 0.01, "display": "slider"}, + ), + "high_noise_mid_strength": ( + "FLOAT", + {"default": 0.8, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.01, "display": "slider"}, + ), + "high_noise_end_strength": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.01, "display": "slider"}, + ), + "low_noise_start_strength": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.01, "display": "slider"}, + ), + "low_noise_mid_strength": ( + "FLOAT", + {"default": 0.2, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.01, "display": "slider"}, + ), + "low_noise_end_strength": ( + "FLOAT", + {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05, "round": 0.01, "display": "slider"}, + ), + "structural_repulsion_boost": ( + "FLOAT", + {"default": 1.0, "min": 1.0, "max": 2.0, "step": 0.05, "round": 0.01, "display": "slider"}, + ), + "clip_vision_start_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_middle_image": ("CLIP_VISION_OUTPUT",), + "clip_vision_end_image": ("CLIP_VISION_OUTPUT",), + }, + } + + RETURN_TYPES = ("CONDITIONING", "CONDITIONING", "CONDITIONING", "LATENT") + RETURN_NAMES = ("positive_high", "positive_low", "negative", "latent") + FUNCTION = "execute" + CATEGORY = "SmartHelperNodes/Wan" + DESCRIPTION = "Standalone SmartHelper copy of Wan First-Middle-Last Frame to Video with separate high-noise end-frame strength." + + def execute( + self, + positive, + negative, + vae, + width, + height, + length, + batch_size, + mode="NORMAL", + start_image=None, + middle_image=None, + end_image=None, + middle_frame_ratio=0.5, + high_noise_mid_strength=0.8, + high_noise_end_strength=1.0, + low_noise_start_strength=1.0, + low_noise_mid_strength=0.2, + low_noise_end_strength=1.0, + structural_repulsion_boost=1.0, + clip_vision_start_image=None, + clip_vision_middle_image=None, + clip_vision_end_image=None, + ): + spacial_scale = vae.spacial_compression_encode() + latent_channels = vae.latent_channels + latent_t = ((length - 1) // 4) + 1 + + device = comfy.model_management.intermediate_device() + + latent = torch.zeros( + [batch_size, latent_channels, latent_t, height // spacial_scale, width // spacial_scale], + device=device, + ) + + if start_image is not None: + start_image = comfy.utils.common_upscale( + start_image[:length].movedim(-1, 1), + width, + height, + "bilinear", + "center", + ).movedim(1, -1) + + if middle_image is not None: + middle_image = comfy.utils.common_upscale( + middle_image[:1].movedim(-1, 1), + width, + height, + "bilinear", + "center", + ).movedim(1, -1) + + if end_image is not None: + end_image = comfy.utils.common_upscale( + end_image[-length:].movedim(-1, 1), + width, + height, + "bilinear", + "center", + ).movedim(1, -1) + + image = torch.ones((length, height, width, 3), device=device) * 0.5 + mask_base = torch.ones( + (1, 1, latent_t * 4, latent.shape[-2], latent.shape[-1]), + device=device, + ) + + middle_idx = self._calculate_aligned_position(middle_frame_ratio, length) + middle_idx = max(4, min(middle_idx, length - 5)) + + mask_high_noise = mask_base.clone() + mask_low_noise = mask_base.clone() + + if start_image is not None: + image[: start_image.shape[0]] = start_image + mask_high_noise[:, :, : start_image.shape[0] + 3] = 0.0 + + low_start_mask_value = 1.0 - low_noise_start_strength + mask_low_noise[:, :, : start_image.shape[0] + 3] = low_start_mask_value + + if middle_image is not None: + image[middle_idx : middle_idx + 1] = middle_image + + start_range = max(0, middle_idx) + end_range = min(length, middle_idx + 4) + + high_noise_mask_value = 1.0 - high_noise_mid_strength + mask_high_noise[:, :, start_range:end_range] = high_noise_mask_value + + low_middle_mask_value = 1.0 - low_noise_mid_strength + mask_low_noise[:, :, start_range:end_range] = low_middle_mask_value + + if end_image is not None: + image[-end_image.shape[0] :] = end_image + + high_noise_end_mask_value = 1.0 - high_noise_end_strength + mask_high_noise[:, :, -end_image.shape[0] :] = high_noise_end_mask_value + + low_end_mask_value = 1.0 - low_noise_end_strength + mask_low_noise[:, :, -end_image.shape[0] :] = low_end_mask_value + + # Build a separate high-noise image when the end-frame influence is reduced. + # The mask alone is not enough because the model still sees end-frame pixels in + # concat_latent_image; blend toward neutral gray instead of hard-disabling it. + needs_separate_high_latent = (end_image is not None and high_noise_end_strength < 1.0) + if needs_separate_high_latent: + image_high = image.clone() + image_high[-end_image.shape[0] :] = self._blend_with_neutral( + end_image, + high_noise_end_strength, + ) + concat_latent_image_high = vae.encode(image_high[:, :, :, :3]) + else: + concat_latent_image_high = None + + concat_latent_image = vae.encode(image[:, :, :, :3]) + + if concat_latent_image_high is None: + concat_latent_image_high = concat_latent_image + + if structural_repulsion_boost > 1.001 and length > 4: + mask_h, mask_w = mask_high_noise.shape[-2], mask_high_noise.shape[-1] + boost_factor = structural_repulsion_boost - 1.0 + + def create_spatial_gradient(img1, img2): + if img1 is None or img2 is None: + return None + + motion_diff = torch.abs(img2[0] - img1[0]).mean(dim=-1, keepdim=False) + motion_diff_4d = motion_diff.unsqueeze(0).unsqueeze(0) + motion_diff_scaled = F.interpolate( + motion_diff_4d, + size=(mask_h, mask_w), + mode="bilinear", + align_corners=False, + ) + + motion_normalized = (motion_diff_scaled - motion_diff_scaled.min()) / ( + motion_diff_scaled.max() - motion_diff_scaled.min() + 1e-8 + ) + + spatial_gradient = 1.0 - motion_normalized * boost_factor * 2.5 + spatial_gradient = torch.clamp(spatial_gradient, 0.02, 1.0) + return spatial_gradient[0, 0] + + if start_image is not None and middle_image is not None: + start_img = start_image[0:1].to(device) + mid_img = middle_image[0:1].to(device) + + spatial_gradient_1 = create_spatial_gradient(start_img, mid_img) + + if spatial_gradient_1 is not None: + start_end = start_image.shape[0] + 3 + mid_protect_start = max(start_end, middle_idx - 4) + transition_end = min(mid_protect_start, length) + + for frame_idx in range(start_end, transition_end): + current_mask = mask_high_noise[:, :, frame_idx, :, :] + mask_high_noise[:, :, frame_idx, :, :] = current_mask * spatial_gradient_1 + + if middle_image is not None and end_image is not None: + mid_img = middle_image[0:1].to(device) + end_img = end_image[-1:].to(device) + + spatial_gradient_2 = create_spatial_gradient(mid_img, end_img) + + if spatial_gradient_2 is not None: + mid_protect_end = middle_idx + 5 + transition_start = mid_protect_end + end_start = length - end_image.shape[0] + + for frame_idx in range(transition_start, end_start): + current_mask = mask_high_noise[:, :, frame_idx, :, :] + mask_high_noise[:, :, frame_idx, :, :] = current_mask * spatial_gradient_2 + + if start_image is not None and end_image is not None and middle_image is None: + start_img = start_image[0:1].to(device) + end_img = end_image[-1:].to(device) + + spatial_gradient = create_spatial_gradient(start_img, end_img) + + if spatial_gradient is not None: + start_end = start_image.shape[0] + 3 + end_start = length - end_image.shape[0] + + for frame_idx in range(start_end, end_start): + current_mask = mask_high_noise[:, :, frame_idx, :, :] + mask_high_noise[:, :, frame_idx, :, :] = current_mask * spatial_gradient + + needs_separate_low_latent = ( + mode == "SINGLE_PERSON" + or low_noise_start_strength == 0.0 + or low_noise_mid_strength == 0.0 + or (end_image is not None and low_noise_end_strength < 1.0) + ) + + if needs_separate_low_latent: + image_low_only = torch.ones((length, height, width, 3), device=device) * 0.5 + + if mode == "SINGLE_PERSON": + if start_image is not None: + image_low_only[: start_image.shape[0]] = start_image + else: + if start_image is not None and low_noise_start_strength > 0.0: + image_low_only[: start_image.shape[0]] = start_image + if middle_image is not None and low_noise_mid_strength > 0.0: + image_low_only[middle_idx : middle_idx + 1] = middle_image + if end_image is not None and low_noise_end_strength > 0.0: + image_low_only[-end_image.shape[0] :] = self._blend_with_neutral( + end_image, + low_noise_end_strength, + ) + + concat_latent_image_low = vae.encode(image_low_only[:, :, :, :3]) + else: + concat_latent_image_low = concat_latent_image + + mask_high_reshaped = mask_high_noise.view( + 1, + mask_high_noise.shape[2] // 4, + 4, + mask_high_noise.shape[3], + mask_high_noise.shape[4], + ).transpose(1, 2) + + mask_low_reshaped = mask_low_noise.view( + 1, + mask_low_noise.shape[2] // 4, + 4, + mask_low_noise.shape[3], + mask_low_noise.shape[4], + ).transpose(1, 2) + + positive_high_noise = node_helpers.conditioning_set_values( + positive, + {"concat_latent_image": concat_latent_image_high, "concat_mask": mask_high_reshaped}, + ) + + positive_low_noise = node_helpers.conditioning_set_values( + positive, + {"concat_latent_image": concat_latent_image_low, "concat_mask": mask_low_reshaped}, + ) + + negative_out = node_helpers.conditioning_set_values( + negative, + {"concat_latent_image": concat_latent_image_high, "concat_mask": mask_high_reshaped}, + ) + + clip_vision_output = self._merge_clip_vision_outputs( + clip_vision_start_image, + clip_vision_middle_image, + clip_vision_end_image, + ) + + if clip_vision_output is not None: + positive_low_noise = node_helpers.conditioning_set_values( + positive_low_noise, + {"clip_vision_output": clip_vision_output}, + ) + + negative_out = node_helpers.conditioning_set_values( + negative_out, + {"clip_vision_output": clip_vision_output}, + ) + + out_latent = {"samples": latent} + return (positive_high_noise, positive_low_noise, negative_out, out_latent) + + @classmethod + def _calculate_aligned_position(cls, ratio, total_frames): + desired_idx = int(total_frames * ratio) + latent_idx = desired_idx // 4 + aligned_idx = latent_idx * 4 + aligned_idx = max(0, min(aligned_idx, total_frames - 1)) + return aligned_idx + + @classmethod + def _merge_clip_vision_outputs(cls, *outputs): + valid_outputs = [output for output in outputs if output is not None] + + if not valid_outputs: + return None + + if len(valid_outputs) == 1: + return valid_outputs[0] + + all_states = [output.penultimate_hidden_states for output in valid_outputs] + combined_states = torch.cat(all_states, dim=-2) + + result = comfy.clip_vision.Output() + result.penultimate_hidden_states = combined_states + return result + + @staticmethod + def _blend_with_neutral(image, strength): + strength = max(0.0, min(float(strength), 1.0)) + return 0.5 + (image - 0.5) * strength + + +NODE_CLASS_MAPPINGS = { + "SmartWanFirstMiddleLastFrameToVideoHighNoiseEnd": SmartWanFirstMiddleLastFrameToVideoHighNoiseEnd, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SmartWanFirstMiddleLastFrameToVideoHighNoiseEnd": "Smart Wan First-Middle-Last Frame to Video (High Noise End)", +}