Merge branch 'master' into matt/be-1899-listassets-contract-fields

.github/workflows/ci-cursor-review.yml

@@ -23,9 +23,9 @@ jobs:
     # SHA-pinned per zizmor `unpinned-uses: hash-pin`. Bump this SHA to pick up
     # upstream changes; keep `workflows_ref` matching so prompts/scripts load
     # from the same commit as the workflow definition.
-    uses: Comfy-Org/github-workflows/.github/workflows/cursor-review.yml@047ca48febe3a6647608ed2e0c4331b491cb9d6a # github-workflows#9
+    uses: Comfy-Org/github-workflows/.github/workflows/cursor-review.yml@964d5aad37cbfb57c5b23961d42c2fd85868bf1d # github-workflows main (964d5aa)
     with:
-      workflows_ref: 047ca48febe3a6647608ed2e0c4331b491cb9d6a
+      workflows_ref: 964d5aad37cbfb57c5b23961d42c2fd85868bf1d
       diff_excludes: >-
         :!**/.claude/**
         :!**/dist/**

.github/workflows/cla.yml

@@ -32,9 +32,11 @@ jobs:
           PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }}
           PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }}
           BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot]
+        # For each commit emit the GitHub login when the author/committer email resolves to a GitHub account
+        # otherwise fall back to the raw git name.
         run: |
           others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
-            --jq '.[] | (.author.login // empty), (.committer.login // empty)' \
+            --jq '.[] | (.author.login // .commit.author.name // empty), (.committer.login // .commit.committer.name // empty)' \
             | sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
           if [ -n "$others" ]; then
             echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
@@ -43,7 +45,7 @@ jobs:
           fi
 
       - name: CLA Assistant
-        # Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase.
+        # Run on PR events, on "recheck" comment, or when someone posts the signing phrase.
         # IMPORTANT: this phrase must match `custom-pr-sign-comment` below.
         if: >
           github.event_name == 'pull_request_target' ||

.github/workflows/release-stable-all.yml

@@ -20,7 +20,7 @@ jobs:
       git_tag: ${{ inputs.git_tag }}
       cache_tag: "cu130"
       python_minor: "13"
-      python_patch: "12"
+      python_patch: "14"
       rel_name: "nvidia"
       rel_extra_name: ""
       test_release: true
@@ -71,7 +71,7 @@ jobs:
       git_tag: ${{ inputs.git_tag }}
       cache_tag: "xpu"
       python_minor: "13"
-      python_patch: "12"
+      python_patch: "14"
       rel_name: "intel"
       rel_extra_name: ""
       test_release: true

AGENTS.md

@@ -19,6 +19,9 @@
   better to remove a broken feature path than keep a complicated partial fix.
 - Preserve existing APIs, node names, model-loading behavior, file layout, and
   workflow compatibility unless the change is explicitly about replacing them.
+- When compatibility is explicitly out of scope, remove compatibility-only
+  aliases, duplicate nodes, legacy entry points, and preset wrappers instead of
+  retaining parallel ways to perform the same operation.
 - Code must look hand-written for this repository. Changes that read like
   generic AI-generated code will be rejected automatically: unnecessary helper
   layers, vague names, boilerplate comments, defensive branches without a real
@@ -96,6 +99,13 @@
   unless they are read by current code and change current behavior. Remove
   pass-through or stored-but-unused values instead of preserving upstream or
   deprecated API baggage.
+- Do not add a model-specific option to a shared helper when only one caller
+  needs it. Keep one-off behavior at the model integration boundary, or extend
+  the shared helper only when the option is a coherent reusable capability.
+- Implementations of shared model interfaces should accept the standard caller
+  contract without model-specific rejection branches for optional capabilities
+  they do not consume. Let supported behavior be determined by implementation
+  paths that actually use those inputs.
 - If an implementation needs auxiliary values for its own workflow, expose them
   through a private helper or a clearly named implementation-specific method
   instead of overloading the public method's return contract.
@@ -152,8 +162,30 @@
   adding parallel code paths. Use `comfy.quant_ops`, `comfy.model_management`,
   `comfy.memory_management`, `comfy.pinned_memory`, `comfy_aimdo`, and
   `comfy-kitchen` helpers where they already solve the problem.
-- Use optimized comfy-kitchen ops in places where they improve performance
-  without changing the expected dtype, device, memory, or interface behavior.
+- Model implementations must use an existing optimized Comfy Kitchen or
+  ComfyUI operation whenever one supports the required math and tensor layout
+  without changing expected dtype, device, memory, or interface behavior. This
+  is the default implementation requirement, not an optional follow-up
+  optimization.
+- Before implementing model math, inspect the operations already exposed by
+  Comfy Kitchen, `comfy.quant_ops`, and existing ComfyUI model helpers. Check
+  for optimized single, paired, fused, layout-specific, and quantized variants
+  before writing a local implementation or composing lower-level torch ops.
+- Use the compatible optimized operation first and adapt the model's inputs to
+  its documented layout while preserving the model's exact math. If several
+  optimized variants apply, benchmark representative model shapes and select
+  the fastest valid path.
+- Add or retain a local implementation only when no existing optimized
+  operation supports the required math, layout, dtype, device, autograd, or
+  patch contract. Keep differentiable or patch-compatible fallbacks when the
+  optimized inference operation does not provide those contracts.
+- Use the existing ComfyUI cast, offload, and cleanup helpers for parameters
+  passed to optimized operations. Preserve model-specific epsilon, scaling,
+  layout, dtype, device, and output-shape behavior.
+- Prefer ComfyUI's shared optimized kernels and backend dispatchers over
+  handwritten implementations of the same operation. Remove duplicate local
+  kernels and adapt inputs to the shared operation's documented layout while
+  preserving the model's original math and output contract.
 - All models should use the optimized attention function selected by ComfyUI.
   Treat optimized backend functions, dispatch helpers, and capability-selected
   callables as opaque. Higher-level code must not inspect function identity,
@@ -176,6 +208,12 @@
 - Model detection code that inspects linear weight shapes should only use the
   first dimension. The second dimension may be half the original size for
   NVFP4 or other 4-bit quantized models.
+- A model-detection signature must guard every state-dict key it dereferences.
+  Do not partially match a format and then raise an incidental `KeyError` while
+  extracting its configuration.
+- Order model-detection checks from established or more-specific signatures to
+  newer or broader signatures. Put a broad new detector near the generic
+  fallback when giving it higher precedence could steal another model family.
 - Avoid adding `einops` usage in core inference code. Use native torch tensor
   ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
   `unsqueeze`, and `squeeze` instead.
@@ -192,11 +230,23 @@
   methods for scalar or structural calculations.
 - Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
   storage dtype, bias dtype, and original tensor shape metadata.
+- Do not cast the result of an optimized backend operation back to its input
+  dtype unless that backend's documented result contract requires normalization.
+  In particular, trust the selected optimized-attention implementation to honor
+  its dtype contract.
 - Keep model-native latent layout handling inside the model or latent-format
   owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
   dimensions in nodes or other caller-side adapters just to satisfy a model
   forward; the model path should consume and return the native latent shape for
   that model family.
+- DiT models should accept latent dimensions that are not exact patch-size
+  multiples. Use `comfy.ldm.common_dit.pad_to_patch_size` on every patchified
+  target or reference input, then crop only the target output back to its
+  original dimensions.
+- Avoid defensive shape and configuration checks that merely replace the clear
+  failure from the tensor operation immediately below them. Add explicit
+  validation only when it provides materially better context at a real boundary
+  or prevents silent incorrect output.
 - Assume inputs to the main model forward are already in the compute dtype by
   default, except integer inputs such as some model timestep tensors. Do not add
   defensive or convenience casts in model code; it is better for invalid dtype
@@ -260,6 +310,15 @@
 - Model implementations should add the minimal number of ComfyUI nodes required
   to run the model. Reuse existing nodes as much as possible; adapting the model
   to work with existing nodes is strongly preferred over creating new nodes.
+- Use `io.Autogrow` for a variable number of repeated inputs instead of a fixed
+  series of numbered optional sockets. Set its minimum to zero when the model
+  has a valid no-item path, and cap it only when the model has a real limit.
+- Mark inputs optional when execution has a valid path that does not read them.
+  If one optional input is needed only to process another optional input, do not
+  force users on the path that supplies neither to connect it.
+- Conditioning nodes should normally output conditioning only. Do not expose
+  input or intermediate images as convenience outputs for downstream sizing or
+  routing; use the existing image path or a dedicated image operation instead.
 - Nodes should output only values they own. Do not add pass-through outputs for
   workflow convenience unless the node is explicitly an output node. Existing
   models, latents, conditioning, or other inputs should flow directly to the

CODEOWNERS

@@ -1,5 +1,6 @@
 * @comfyanonymous @kosinkadink @guill @alexisrolland @rattus128 @kijai
 
 /CODEOWNERS @comfyanonymous
+/AGENTS.md @comfyanonymous
 /.ci/ @comfyanonymous
 /.github/ @comfyanonymous

app/logger.py

@@ -2,9 +2,12 @@ from collections import deque
 from datetime import datetime
 import io
 import logging
+import os
 import sys
 import threading
 
+import comfy.logging
+
 ANSI_NAMED_COLORS = {
     'black':   '\033[30m',
     'red':     '\033[31m',
@@ -18,6 +21,7 @@ ANSI_NAMED_COLORS = {
 
 ANSI_LEVEL_COLORS = {
     'DEBUG':    ANSI_NAMED_COLORS['cyan'],
+    'DETAIL':   ANSI_NAMED_COLORS['blue'],
     'INFO':     ANSI_NAMED_COLORS['green'],
     'WARNING':  ANSI_NAMED_COLORS['yellow'],
     'ERROR':    ANSI_NAMED_COLORS['red'],
@@ -85,7 +89,12 @@ def on_flush(callback):
     if stderr_interceptor is not None:
         stderr_interceptor.on_flush(callback)
 
-def setup_logger(log_level: str = 'INFO', capacity: int = 300, use_stdout: bool = False):
+
+def get_log_level(level):
+    return comfy.logging.DETAIL if level == "DETAIL" else logging.getLevelName(level)
+
+
+def setup_logger(log_level: str = 'INFO', file_outputs=None, capacity: int = 300, use_stdout: bool = False):
     global logs
     if logs:
         return
@@ -99,13 +108,18 @@ def setup_logger(log_level: str = 'INFO', capacity: int = 300, use_stdout: bool
     stderr_interceptor = sys.stderr = LogInterceptor(sys.stderr)
 
     # Setup default global logger
+    if file_outputs is None:
+        file_outputs = [('DETAIL', 'comfyui_detail.log')]
     logger = logging.getLogger()
-    logger.setLevel(log_level)
+    console_level = get_log_level(log_level)
+    file_levels = [get_log_level(level) for level, _ in file_outputs]
+    logger.setLevel(min(console_level, *file_levels))
 
     formatter = ColoredFormatter("%(message)s")
 
     stream_handler = logging.StreamHandler()
     stream_handler.setFormatter(formatter)
+    stream_handler.setLevel(console_level)
 
     if use_stdout:
         # Only errors and critical to stderr
@@ -114,11 +128,24 @@ def setup_logger(log_level: str = 'INFO', capacity: int = 300, use_stdout: bool
         # Lesser to stdout
         stdout_handler = logging.StreamHandler(sys.stdout)
         stdout_handler.setFormatter(formatter)
+        stdout_handler.setLevel(console_level)
         stdout_handler.addFilter(lambda record: record.levelno < logging.ERROR)
         logger.addHandler(stdout_handler)
 
     logger.addHandler(stream_handler)
 
+    for output_level, output_path in file_outputs:
+        output_path = os.path.abspath(output_path)
+        try:
+            output_handler = logging.FileHandler(output_path, encoding="utf-8")
+        except OSError as e:
+            logging.warning("Could not open %s log %s: %s", output_level, output_path, e)
+            continue
+        output_handler.setLevel(get_log_level(output_level))
+        output_handler.setFormatter(logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s"))
+        logger.addHandler(output_handler)
+        logging.info("%s log: %s", output_level.title(), output_path)
+
 
 STARTUP_WARNINGS = []
 

app/model_manager.py

@@ -35,7 +35,11 @@ class ModelFileManager:
             for folder in model_types:
                 if folder in folder_black_list:
                     continue
-                output_folders.append({"name": folder, "folders": folder_paths.get_folder_paths(folder)})
+                output_folders.append({
+                    "name": folder,
+                    "folders": folder_paths.get_folder_paths(folder),
+                    "extensions": sorted(folder_paths.folder_names_and_paths[folder][1]),
+                })
             return web.json_response(output_folders)
 
         # NOTE: This is an experiment to replace `/models/{folder}`

comfy/cli_args.py

@@ -33,6 +33,31 @@ class EnumAction(argparse.Action):
         setattr(namespace, self.dest, value)
 
 
+LOG_LEVELS = ('DEBUG', 'DETAIL', 'INFO', 'WARNING', 'ERROR', 'CRITICAL')
+
+
+class VerboseAction(argparse.Action):
+    def __call__(self, parser, namespace, values, option_string=None):
+        if len(values) == 0:
+            output = ('DEBUG', None)
+        elif len(values) == 1 and values[0] in LOG_LEVELS:
+            output = (values[0], None)
+        elif len(values) == 2 and values[0] in LOG_LEVELS:
+            output = tuple(values)
+        else:
+            parser.error(f"{option_string} expects no values, a console LEVEL, or LEVEL FILE")
+        setattr(namespace, self.dest, [*getattr(namespace, self.dest, []), output])
+
+
+def get_console_log_level(outputs):
+    console_levels = [level for level, path in outputs if path is None]
+    return min(console_levels, key=LOG_LEVELS.index, default='INFO')
+
+
+def get_file_log_outputs(outputs):
+    return [(level, path) for level, path in outputs if path is not None]
+
+
 parser = argparse.ArgumentParser()
 
 parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0,::", help="Specify the IP address to listen on (default: 127.0.0.1). You can give a list of ip addresses by separating them with a comma like: 127.2.2.2,127.3.3.3 If --listen is provided without an argument, it defaults to 0.0.0.0,:: (listens on all ipv4 and ipv6)")
@@ -92,6 +117,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE"
 parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
 parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.")
 parser.add_argument("--enable-triton-backend", action="store_true", help="ComfyUI will enable the use of Triton backend in comfy-kitchen. Is disabled at launch by default.")
+parser.add_argument("--disable-triton-backend", action="store_true", help="Force-disable the comfy-kitchen Triton backend, overriding the automatic ROCm/AMD default and --enable-triton-backend.")
 
 class LatentPreviewMethod(enum.Enum):
     NoPreviews = "none"
@@ -111,7 +137,7 @@ parser.add_argument("--preview-method", type=LatentPreviewMethod, default=Latent
 parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
 
 cache_group = parser.add_mutually_exclusive_group()
-cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 10%% of system RAM (min 2GB, max 10GB), inactive 100%% of system RAM (max 96GB).")
+cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 10%% of system RAM (min 2GB, max 10GB), inactive 100%% of system RAM (max 128GB).")
 cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
 cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
 cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")
@@ -186,7 +212,7 @@ parser.add_argument("--disable-api-nodes", action="store_true", help="Disable lo
 
 parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
 
-parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level')
+parser.add_argument("--verbose", action=VerboseAction, nargs='*', default=[], metavar='LEVEL FILE', help='Set console logging with no values or LEVEL, or add a LEVEL FILE log output. May be repeated.')
 parser.add_argument("--log-stdout", action="store_true", help="Send normal process output to stdout instead of stderr (default).")
 
 

comfy/comfy_api_env.py

@@ -0,0 +1,46 @@
+"""Runtime config the frontend reads from /features to follow --comfy-api-base.
+
+For a non-prod comfy.org backend (staging or an ephemeral preview env), "/features" exposes the api and
+platform base so the frontend talks to it without a rebuild, plus the Firebase environment it should use.
+Prod bases are left alone and keep their build-time defaults.
+"""
+
+from typing import Any
+from urllib.parse import urlparse
+
+from comfy.cli_args import args
+
+_STAGING_API_HOST = "stagingapi.comfy.org"
+_TESTENV_HOST_SUFFIX = ".testenvs.comfy.org"
+_STAGING_PLATFORM_BASE_URL = "https://stagingplatform.comfy.org"
+
+
+def _is_staging_tier(host: str) -> bool:
+    return host == _STAGING_API_HOST or host.endswith(_TESTENV_HOST_SUFFIX)
+
+
+def normalize_comfy_api_base(url: str) -> str:
+    """Rewrite a testenv's friendly main host to its comfy-api '-registry' sibling."""
+    parsed = urlparse(url)
+    host = parsed.hostname or ""
+    if not host.endswith(_TESTENV_HOST_SUFFIX):
+        return url
+    label = host[: -len(_TESTENV_HOST_SUFFIX)]
+    if label.endswith("-registry"):
+        return url
+    return f"{parsed.scheme or 'https'}://{label}-registry{_TESTENV_HOST_SUFFIX}"
+
+
+def environment_overrides_for_base(base_url: str) -> dict[str, Any] | None:
+    """The /features overrides for a staging-tier base, or None for prod."""
+    if not _is_staging_tier(urlparse(base_url).hostname or ""):
+        return None
+    return {
+        "comfy_api_base_url": normalize_comfy_api_base(base_url).rstrip("/"),
+        "comfy_platform_base_url": _STAGING_PLATFORM_BASE_URL,
+        "firebase_env": "dev",
+    }
+
+
+def get_environment_overrides() -> dict[str, Any] | None:
+    return environment_overrides_for_base(getattr(args, "comfy_api_base", "") or "")

comfy/latent_formats.py

@@ -779,6 +779,10 @@ class ACEAudio(LatentFormat):
     latent_channels = 8
     latent_dimensions = 2
 
+class SeedVR2(LatentFormat):
+    latent_channels = 16
+    latent_dimensions = 3
+
 class ACEAudio15(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1

comfy/ldm/anima/lllite.py

@@ -0,0 +1,278 @@
+import re
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+import comfy.ops
+import comfy.utils
+
+
+MODULE_PATTERN = re.compile(r"lllite_dit_blocks_(\d+)_(self_attn_[qkv]_proj|cross_attn_q_proj|mlp_layer1)$")
+
+
+def _group_norm(channels, device=None, dtype=None, operations=None):
+    groups = 8
+    while groups > 1 and channels % groups != 0:
+        groups //= 2
+    return operations.GroupNorm(groups, channels, device=device, dtype=dtype)
+
+
+class AnimaLLLiteResBlock(nn.Module):
+    def __init__(self, channels, device=None, dtype=None, operations=None):
+        super().__init__()
+        self.norm1 = _group_norm(channels, device=device, dtype=dtype, operations=operations)
+        self.conv1 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype)
+        self.norm2 = _group_norm(channels, device=device, dtype=dtype, operations=operations)
+        self.conv2 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype)
+
+    def forward(self, x):
+        h = self.conv1(F.silu(self.norm1(x)))
+        h = self.conv2(F.silu(self.norm2(h)))
+        return x + h
+
+
+class AnimaLLLiteASPP(nn.Module):
+    def __init__(self, channels, dilations, device=None, dtype=None, operations=None):
+        super().__init__()
+        branches = []
+        for dilation in dilations:
+            if dilation == 1:
+                conv = operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype)
+            else:
+                conv = operations.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation, device=device, dtype=dtype)
+            branches.append(nn.Sequential(conv, _group_norm(channels, device=device, dtype=dtype, operations=operations), nn.SiLU()))
+        self.branches = nn.ModuleList(branches)
+        self.global_pool = nn.AdaptiveAvgPool2d(1)
+        self.global_conv = nn.Sequential(
+            operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype),
+            _group_norm(channels, device=device, dtype=dtype, operations=operations),
+            nn.SiLU(),
+        )
+        self.proj = nn.Sequential(
+            operations.Conv2d(channels * (len(dilations) + 1), channels, kernel_size=1, device=device, dtype=dtype),
+            _group_norm(channels, device=device, dtype=dtype, operations=operations),
+            nn.SiLU(),
+        )
+
+    def forward(self, x):
+        height, width = x.shape[-2:]
+        outputs = [branch(x) for branch in self.branches]
+        pooled = self.global_conv(self.global_pool(x))
+        outputs.append(F.interpolate(pooled, size=(height, width), mode="bilinear", align_corners=False))
+        return self.proj(torch.cat(outputs, dim=1))
+
+
+class AnimaLLLiteConditioning(nn.Module):
+    def __init__(self, cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations, device=None, dtype=None, operations=None):
+        super().__init__()
+        half_dim = cond_dim // 2
+        self.conv1 = operations.Conv2d(cond_in_channels, half_dim, kernel_size=4, stride=4, device=device, dtype=dtype)
+        self.norm1 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations)
+        self.conv2 = operations.Conv2d(half_dim, half_dim, kernel_size=3, padding=1, device=device, dtype=dtype)
+        self.norm2 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations)
+        self.conv3 = operations.Conv2d(half_dim, cond_dim, kernel_size=4, stride=4, device=device, dtype=dtype)
+        self.norm3 = _group_norm(cond_dim, device=device, dtype=dtype, operations=operations)
+        self.resblocks = nn.ModuleList([
+            AnimaLLLiteResBlock(cond_dim, device=device, dtype=dtype, operations=operations)
+            for _ in range(cond_resblocks)
+        ])
+        self.aspp = AnimaLLLiteASPP(cond_dim, aspp_dilations, device=device, dtype=dtype, operations=operations) if aspp_dilations else None
+        self.proj = operations.Conv2d(cond_dim, cond_emb_dim, kernel_size=1, device=device, dtype=dtype)
+        self.out_norm = operations.LayerNorm(cond_emb_dim, device=device, dtype=dtype)
+
+    def forward(self, x):
+        x = F.silu(self.norm1(self.conv1(x)))
+        x = F.silu(self.norm2(self.conv2(x)))
+        x = F.silu(self.norm3(self.conv3(x)))
+        for block in self.resblocks:
+            x = block(x)
+        if self.aspp is not None:
+            x = self.aspp(x)
+        x = self.proj(x).flatten(2).transpose(1, 2).contiguous()
+        return self.out_norm(x)
+
+
+class AnimaLLLiteModule(nn.Module):
+    def __init__(self, in_dim, cond_emb_dim, mlp_dim, device=None, dtype=None, operations=None):
+        super().__init__()
+        self.down = operations.Linear(in_dim, mlp_dim, device=device, dtype=dtype)
+        self.mid = operations.Linear(mlp_dim + cond_emb_dim, mlp_dim, device=device, dtype=dtype)
+        self.cond_to_film = operations.Linear(cond_emb_dim, 2 * mlp_dim, device=device, dtype=dtype)
+        self.up = operations.Linear(mlp_dim, in_dim, device=device, dtype=dtype)
+        self.depth_embed = nn.Parameter(torch.empty(cond_emb_dim, device=device, dtype=dtype), requires_grad=False)
+
+    def forward(self, x, cond_emb, strength):
+        original_shape = x.shape
+        if x.ndim == 5:
+            x = x.flatten(1, 3)
+
+        if x.shape[0] != cond_emb.shape[0]:
+            if x.shape[0] % cond_emb.shape[0] != 0:
+                raise ValueError(f"Anima LLLite batch mismatch: model input batch {x.shape[0]}, control batch {cond_emb.shape[0]}")
+            cond_emb = cond_emb.repeat(x.shape[0] // cond_emb.shape[0], 1, 1)
+        if x.shape[1] != cond_emb.shape[1]:
+            raise ValueError(f"Anima LLLite sequence mismatch: model input has {x.shape[1]} tokens, control has {cond_emb.shape[1]}")
+
+        cond_local = cond_emb + comfy.ops.cast_to_input(self.depth_embed, cond_emb)
+        hidden = F.silu(self.down(x))
+        gamma, beta = self.cond_to_film(cond_local).chunk(2, dim=-1)
+        hidden = self.mid(torch.cat((cond_local, hidden), dim=-1))
+        hidden = F.silu(hidden * (1 + gamma) + beta)
+        x = x + self.up(hidden) * strength
+
+        if len(original_shape) == 5:
+            x = x.reshape(original_shape)
+        return x
+
+
+class AnimaLLLite(nn.Module):
+    def __init__(self, state_dict, metadata, device=None, dtype=None, operations=None):
+        super().__init__()
+        metadata = metadata or {}
+        version = metadata.get("lllite.version", "2")
+        if version != "2":
+            raise ValueError(f"Unsupported Anima LLLite version {version!r}; only named-key v2 checkpoints are supported")
+
+        module_names = sorted({key.split(".", 1)[0] for key in state_dict if key.startswith("lllite_dit_blocks_")})
+        if not module_names:
+            raise ValueError("Anima LLLite checkpoint has no lllite_dit_blocks_* modules")
+
+        cond_in_channels = state_dict["lllite_conditioning1.conv1.weight"].shape[1]
+        cond_dim = state_dict["lllite_conditioning1.conv3.weight"].shape[0]
+        cond_emb_dim = state_dict["lllite_conditioning1.proj.weight"].shape[0]
+        resblock_ids = {int(key.split(".")[2]) for key in state_dict if key.startswith("lllite_conditioning1.resblocks.")}
+        cond_resblocks = max(resblock_ids) + 1 if resblock_ids else 0
+        use_aspp = any(key.startswith("lllite_conditioning1.aspp.") for key in state_dict)
+        dilation_string = metadata.get("lllite.aspp_dilations", "1,2,4,8")
+        aspp_dilations = tuple(int(value) for value in dilation_string.split(",") if value.strip()) if use_aspp else ()
+
+        self.cond_in_channels = cond_in_channels
+        self.inpaint_masked_input = metadata.get("lllite.inpaint_masked_input", "false").lower() == "true"
+        self.lllite_conditioning1 = AnimaLLLiteConditioning(
+            cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations,
+            device=device, dtype=dtype, operations=operations,
+        )
+
+        self.module_names = set()
+        self.block_count = 0
+        self.model_dim = None
+        for name in module_names:
+            match = MODULE_PATTERN.fullmatch(name)
+            if match is None:
+                raise ValueError(f"Unsupported Anima LLLite module name: {name}")
+            down_shape = state_dict[f"{name}.down.weight"].shape
+            mlp_dim, in_dim = down_shape
+            module_cond_dim = state_dict[f"{name}.cond_to_film.weight"].shape[1]
+            if module_cond_dim != cond_emb_dim:
+                raise ValueError(f"Anima LLLite conditioning dimension mismatch in {name}: {module_cond_dim} != {cond_emb_dim}")
+            if self.model_dim is None:
+                self.model_dim = in_dim
+            elif self.model_dim != in_dim:
+                raise ValueError(f"Anima LLLite model dimension mismatch in {name}: {in_dim} != {self.model_dim}")
+            self.add_module(name, AnimaLLLiteModule(in_dim, cond_emb_dim, mlp_dim, device=device, dtype=dtype, operations=operations))
+            self.module_names.add(name)
+            self.block_count = max(self.block_count, int(match.group(1)) + 1)
+
+    def encode_conditioning(self, image):
+        return self.lllite_conditioning1(image)
+
+    def apply(self, x, cond_emb, block_index, target, strength):
+        name = f"lllite_dit_blocks_{block_index}_{target}"
+        if name not in self.module_names:
+            return x
+        return self.get_submodule(name)(x, cond_emb, strength)
+
+
+class AnimaLLLitePatch:
+    def __init__(self, model_patch, image, mask, strength, sigma_start, sigma_end):
+        self.model_patch = model_patch
+        self.image = image
+        self.mask = mask
+        self.strength = strength
+        self.sigma_start = sigma_start
+        self.sigma_end = sigma_end
+
+    def __call__(self, args):
+        x = args["x"]
+        transformer_options = args["transformer_options"]
+        if self.strength == 0.0:
+            return args
+        sigmas = transformer_options.get("sigmas")
+        if sigmas is not None:
+            sigma = float(sigmas.max().item())
+            if not self.sigma_end <= sigma <= self.sigma_start:
+                return args
+        if x.shape[2] != 1:
+            raise ValueError(f"Anima LLLite only supports T=1, got T={x.shape[2]}")
+
+        target_height = x.shape[-2] * 8
+        target_width = x.shape[-1] * 8
+        image = comfy.utils.common_upscale(
+            self.image.movedim(-1, 1), target_width, target_height, "bicubic", crop="center"
+        ).clamp(0.0, 1.0)
+        image = image.to(device=x.device, dtype=x.dtype) * 2.0 - 1.0
+
+        if self.model_patch.model.cond_in_channels == 4:
+            mask = self.mask
+            if mask.ndim == 3:
+                mask = mask.unsqueeze(1)
+            if mask.ndim != 4 or mask.shape[1] != 1:
+                raise ValueError(f"Anima LLLite mask must have one channel, got shape {tuple(mask.shape)}")
+            mask = comfy.utils.common_upscale(
+                mask.float(), target_width, target_height, "nearest-exact", crop="center"
+            )
+            if mask.shape[0] != image.shape[0]:
+                if image.shape[0] % mask.shape[0] != 0:
+                    raise ValueError(
+                        f"Anima LLLite mask batch {mask.shape[0]} cannot be broadcast to image batch {image.shape[0]}"
+                    )
+                mask = mask.repeat(image.shape[0] // mask.shape[0], 1, 1, 1)
+            mask = (mask >= 0.5).to(device=x.device, dtype=x.dtype)
+            if self.model_patch.model.inpaint_masked_input:
+                image = image * (mask < 0.5).to(image.dtype)
+            image = torch.cat((image, mask * 2.0 - 1.0), dim=1)
+
+        cond_emb = self.model_patch.model.encode_conditioning(image)
+        transformer_options["model_patch_data"][self] = cond_emb
+        return args
+
+    def to(self, device_or_dtype):
+        return self
+
+    def models(self):
+        return [self.model_patch]
+
+
+class AnimaLLLiteAttentionPatch:
+    def __init__(self, patch, targets):
+        self.patch = patch
+        self.targets = targets
+
+    def __call__(self, q, k, v, pe=None, attn_mask=None, extra_options=None):
+        cond_emb = extra_options["model_patch_data"].get(self.patch)
+        if cond_emb is None:
+            return {"q": q, "k": k, "v": v, "pe": pe, "attn_mask": attn_mask}
+
+        block_index = extra_options["block_index"]
+        values = {"q": q, "k": k, "v": v}
+        for value_name, target in self.targets.items():
+            values[value_name] = self.patch.model_patch.model.apply(
+                values[value_name], cond_emb, block_index, target, self.patch.strength
+            )
+
+        return {"q": values["q"], "k": values["k"], "v": values["v"], "pe": pe, "attn_mask": attn_mask}
+
+
+class AnimaLLLiteMLPPatch:
+    def __init__(self, patch):
+        self.patch = patch
+
+    def __call__(self, args):
+        cond_emb = args["transformer_options"]["model_patch_data"].get(self.patch)
+        if cond_emb is None:
+            return args
+        args["x"] = self.patch.model_patch.model.apply(
+            args["x"], cond_emb, args["transformer_options"]["block_index"], "mlp_layer1", self.patch.strength
+        )
+        return args

comfy/ldm/cosmos/predict2.py

@@ -14,6 +14,7 @@ from torchvision import transforms
 import comfy.patcher_extension
 from comfy.ldm.modules.attention import optimized_attention
 import comfy.ldm.common_dit
+import comfy.ops
 import comfy.quant_ops
 
 
@@ -148,11 +149,29 @@ class Attention(nn.Module):
         x: torch.Tensor,
         context: Optional[torch.Tensor] = None,
         rope_emb: Optional[torch.Tensor] = None,
+        transformer_options: Optional[dict] = {},
     ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-        q = self.q_proj(x)
         context = x if context is None else context
-        k = self.k_proj(context)
-        v = self.v_proj(context)
+        q_input = x
+        k_input = context
+        v_input = context
+
+        transformer_patches = transformer_options.get("patches", {})
+        patch_name = "attn1_patch" if self.is_selfattn else "attn2_patch"
+        if patch_name in transformer_patches:
+            extra_options = transformer_options.copy()
+            extra_options["n_heads"] = self.n_heads
+            extra_options["dim_head"] = self.head_dim
+            for patch in transformer_patches[patch_name]:
+                out = patch(q_input, k_input, v_input, pe=rope_emb, attn_mask=None, extra_options=extra_options)
+                q_input = out.get("q", q_input)
+                k_input = out.get("k", k_input)
+                v_input = out.get("v", v_input)
+                rope_emb = out.get("pe", rope_emb)
+
+        q = self.q_proj(q_input)
+        k = self.k_proj(k_input)
+        v = self.v_proj(v_input)
         q, k, v = map(
             lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
             (q, k, v),
@@ -161,11 +180,16 @@ class Attention(nn.Module):
         def apply_norm_and_rotary_pos_emb(
             q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
         ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-            q = self.q_norm(q)
-            k = self.k_norm(k)
             v = self.v_norm(v)
             if self.is_selfattn and rope_emb is not None:  # only apply to self-attention!
-                q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rope_emb)
+                q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, q, offloadable=True)
+                k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, k, offloadable=True)
+                q, k = comfy.quant_ops.ck.rms_rope_split_half(q, k, rope_emb, q_scale, k_scale, self.q_norm.eps)
+                comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
+                comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
+            else:
+                q = self.q_norm(q)
+                k = self.k_norm(k)
             return q, k, v
 
         q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
@@ -188,7 +212,7 @@ class Attention(nn.Module):
             x (Tensor): The query tensor of shape [B, Mq, K]
             context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
         """
-        q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb)
+        q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb, transformer_options=transformer_options)
         return self.compute_attention(q, k, v, transformer_options=transformer_options)
 
 
@@ -555,8 +579,14 @@ class Block(nn.Module):
             self.layer_norm_mlp,
             scale_mlp_B_T_1_1_D,
             shift_mlp_B_T_1_1_D,
-        )
-        result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.to(compute_dtype))
+        ).to(compute_dtype)
+        patches = transformer_options.get("patches", {})
+        if "mlp_patch" in patches:
+            args = {"x": normalized_x_B_T_H_W_D, "transformer_options": transformer_options}
+            for patch in patches["mlp_patch"]:
+                args = patch(args)
+            normalized_x_B_T_H_W_D = args["x"]
+        result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
         x_B_T_H_W_D = torch.addcmul(x_B_T_H_W_D, gate_mlp_B_T_1_1_D.to(residual_dtype), result_B_T_H_W_D.to(residual_dtype))
         return x_B_T_H_W_D
 
@@ -863,11 +893,22 @@ class MiniTrainDIT(nn.Module):
                 x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape
             ), f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
 
+        transformer_options = kwargs.get("transformer_options", {})
+        patches = transformer_options.get("patches", {})
+        if "post_input" in patches:
+            transformer_options = transformer_options.copy()
+            transformer_options["model_patch_data"] = {}
+
+        if "post_input" in patches:
+            for patch in patches["post_input"]:
+                out = patch({"img": x_B_T_H_W_D, "x": x_B_C_T_H_W, "transformer_options": transformer_options})
+                x_B_T_H_W_D = out["img"]
+
         block_kwargs = {
             "rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0),
             "adaln_lora_B_T_3D": adaln_lora_B_T_3D,
             "extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
-            "transformer_options": kwargs.get("transformer_options", {}),
+            "transformer_options": transformer_options,
         }
 
         # The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream
@@ -877,7 +918,8 @@ class MiniTrainDIT(nn.Module):
         if x_B_T_H_W_D.dtype == torch.float16:
             x_B_T_H_W_D = x_B_T_H_W_D.float()
 
-        for block in self.blocks:
+        for block_index, block in enumerate(self.blocks):
+            transformer_options["block_index"] = block_index
             x_B_T_H_W_D = block(
                 x_B_T_H_W_D,
                 t_embedding_B_T_D,

comfy/ldm/ernie/model.py

@@ -5,6 +5,7 @@ import torch.nn.functional as F
 
 from comfy.ldm.modules.attention import optimized_attention
 import comfy.model_management
+import comfy.ops
 import comfy.quant_ops
 
 def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
@@ -111,11 +112,17 @@ class ErnieImageAttention(nn.Module):
         query = q_flat.view(B, S, self.heads, self.head_dim)
         key = k_flat.view(B, S, self.heads, self.head_dim)
 
-        query = self.norm_q(query)
-        key = self.norm_k(key)
-
-        if image_rotary_emb is not None:
-            query, key = comfy.quant_ops.ck.apply_rope_split_half(query, key, image_rotary_emb)
+        if image_rotary_emb is not None and not comfy.model_management.in_training:
+            q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.norm_q, query, offloadable=True)
+            k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.norm_k, key, offloadable=True)
+            query, key = comfy.quant_ops.ck.rms_rope_split_half(query, key, image_rotary_emb, q_scale, k_scale, self.norm_q.eps)
+            comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_offload_stream)
+            comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_offload_stream)
+        else:
+            query = self.norm_q(query)
+            key = self.norm_k(key)
+            if image_rotary_emb is not None:
+                query, key = comfy.quant_ops.ck.apply_rope_split_half(query, key, image_rotary_emb)
 
         q_flat = query.reshape(B, S, -1)
         k_flat = key.reshape(B, S, -1)

comfy/ldm/hidream_o1/attention.py

@@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None):
     The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
     """
 
-    def two_pass_attention(q, k, v, heads, **kwargs):
+    def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs):
         B, H, T, D = q.shape
 
         if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
-            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
+            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
         elif ar_len >= T:
-            out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
+            out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
         elif ar_len <= 0:
-            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
+            out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
         else:
             out_ar = comfy.ops.scaled_dot_product_attention(
                 q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
-                attn_mask=None, dropout_p=0.0, is_causal=True,
+                attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
             )
             out_gen = optimized_attention(
                 q[:, :, ar_len:], k, v, heads,
                 mask=None, skip_reshape=True, skip_output_reshape=True,
-                transformer_options=transformer_options,
+                transformer_options=transformer_options, enable_gqa=enable_gqa,
             )
             out = torch.cat([out_ar, out_gen], dim=2)
 

comfy/ldm/ideogram4/model.py

@@ -12,10 +12,13 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import comfy.model_management
+import comfy.ops
 import comfy.patcher_extension
+import comfy.quant_ops
 from comfy.ldm.lumina.model import FeedForward
 from comfy.ldm.modules.attention import optimized_attention_masked
-from comfy.text_encoders.llama import apply_rope, precompute_freqs_cis
+from comfy.text_encoders.llama import precompute_freqs_cis
 
 # Per-token role indicators
 SEQUENCE_PADDING_INDICATOR = -1
@@ -25,6 +28,22 @@ LLM_TOKEN_INDICATOR = 3
 IMAGE_POSITION_OFFSET = 65536
 
 
+def _split_half_rope_matrix(freqs_cis):
+    cos, sin, neg_sin = freqs_cis
+    half_dim = sin.shape[-1]
+    matrix = torch.stack(
+        (cos[..., :half_dim], neg_sin, sin, cos[..., half_dim:]), dim=-1
+    )
+    return matrix.reshape(*matrix.shape[:-1], 2, 2).unsqueeze(2)
+
+
+def _apply_rope_split_half1(x, freqs_cis):
+    x_dtype = x.dtype
+    x = x.reshape(*x.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).to(freqs_cis.dtype)
+    output = freqs_cis[..., 0] * x[..., 0] + freqs_cis[..., 1] * x[..., 1]
+    return output.movedim(-1, -2).reshape(*x.shape[:-3], -1).to(x_dtype)
+
+
 class Ideogram4Attention(nn.Module):
     def __init__(self, hidden_size, num_heads, eps=1e-5, dtype=None, device=None, operations=None):
         super().__init__()
@@ -42,16 +61,23 @@ class Ideogram4Attention(nn.Module):
         qkv = self.qkv(x).view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
         q, k, v = qkv.unbind(dim=2)
 
-        q = self.norm_q(q)
-        k = self.norm_k(k)
+        if comfy.model_management.in_training:
+            q = _apply_rope_split_half1(self.norm_q(q), freqs_cis)
+            k = _apply_rope_split_half1(self.norm_k(k), freqs_cis)
+        else:
+            q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.norm_q, q, offloadable=True)
+            k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.norm_k, k, offloadable=True)
+            q, k = comfy.quant_ops.ck.rms_rope_split_half(
+                q, k, freqs_cis, q_scale, k_scale, self.norm_q.eps
+            )
+            comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_offload_stream)
+            comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_offload_stream)
 
         # (B, heads, L, head_dim)
         q = q.transpose(1, 2)
         k = k.transpose(1, 2)
         v = v.transpose(1, 2)
 
-        q, k = apply_rope(q, k, freqs_cis)
-
         out = optimized_attention_masked(q, k, v, self.num_heads, attn_mask, skip_reshape=True, transformer_options=transformer_options)
         return self.o(out)
 
@@ -181,6 +207,7 @@ class Ideogram4Transformer(nn.Module):
             self.head_dim, position_ids[0].transpose(0, 1), self.rope_theta,
             rope_dims=self.mrope_section, interleaved_mrope=True, device=position_ids.device,
         )
+        freqs_cis = _split_half_rope_matrix(freqs_cis)
 
         if attn_mask is not None and attn_mask.dtype == torch.bool:
             attn_mask = torch.zeros_like(attn_mask, dtype=h.dtype).masked_fill_(~attn_mask, -torch.finfo(h.dtype).max)

comfy/ldm/joyimage/model.py

@@ -0,0 +1,454 @@
+# https://github.com/jdopensource/JoyAI-Image-Edit (Apache 2.0)
+import math
+from typing import Optional, Tuple
+
+import comfy_kitchen
+import torch
+import torch.nn as nn
+
+import comfy.ldm.common_dit
+import comfy.ops
+import comfy.patcher_extension
+from comfy.ldm.lightricks.model import GELU_approx, PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
+from comfy.ldm.modules.attention import optimized_attention
+
+
+class JoyImageModulate(nn.Module):
+    def __init__(self, hidden_size: int, factor: int, dtype=None, device=None):
+        super().__init__()
+        self.factor = factor
+        self.modulate_table = nn.Parameter(
+            torch.empty(1, factor, hidden_size, dtype=dtype, device=device)
+        )
+
+    def forward(self, x: torch.Tensor) -> list:
+        if x.ndim != 3:
+            x = x.unsqueeze(1)
+        table = comfy.ops.cast_to_input(self.modulate_table, x)
+        return [o.squeeze(1) for o in (table + x).chunk(self.factor, dim=1)]
+
+
+class JoyImageFeedForward(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        inner_dim: int,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.net = nn.ModuleList([
+            GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations),
+            nn.Identity(),
+            operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device),
+        ])
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            x = module(x)
+        return x
+
+
+class JoyImageAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        eps: float = 1e-6,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.num_attention_heads = num_attention_heads
+        inner_dim = num_attention_heads * attention_head_dim
+
+        self.img_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device)
+        self.img_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
+        self.img_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
+        self.img_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device)
+
+        self.txt_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device)
+        self.txt_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
+        self.txt_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
+        self.txt_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device)
+
+    def forward(
+        self,
+        img: torch.Tensor,
+        txt: torch.Tensor,
+        image_rotary_emb: torch.Tensor,
+        transformer_options=None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        heads = self.num_attention_heads
+
+        img_q, img_k, img_v = self.img_attn_qkv(img).chunk(3, dim=-1)
+        txt_q, txt_k, txt_v = self.txt_attn_qkv(txt).chunk(3, dim=-1)
+
+        img_q = img_q.unflatten(-1, (heads, -1))
+        img_k = img_k.unflatten(-1, (heads, -1))
+        img_v = img_v.unflatten(-1, (heads, -1))
+        txt_q = txt_q.unflatten(-1, (heads, -1))
+        txt_k = txt_k.unflatten(-1, (heads, -1))
+        txt_v = txt_v.unflatten(-1, (heads, -1))
+
+        txt_q = self.txt_attn_q_norm(txt_q)
+        txt_k = self.txt_attn_k_norm(txt_k)
+
+        img_q_scale, _, img_q_offload_stream = comfy.ops.cast_bias_weight(self.img_attn_q_norm, img_q, offloadable=True)
+        img_k_scale, _, img_k_offload_stream = comfy.ops.cast_bias_weight(self.img_attn_k_norm, img_k, offloadable=True)
+        img_q, img_k = comfy_kitchen.rms_rope(
+            img_q,
+            img_k,
+            image_rotary_emb,
+            img_q_scale,
+            img_k_scale,
+            self.img_attn_q_norm.eps,
+        )
+        comfy.ops.uncast_bias_weight(self.img_attn_q_norm, img_q_scale, None, img_q_offload_stream)
+        comfy.ops.uncast_bias_weight(self.img_attn_k_norm, img_k_scale, None, img_k_offload_stream)
+
+        joint_q = torch.cat([img_q, txt_q], dim=1)
+        joint_k = torch.cat([img_k, txt_k], dim=1)
+        joint_v = torch.cat([img_v, txt_v], dim=1)
+
+        joint_q = joint_q.flatten(2, 3)
+        joint_k = joint_k.flatten(2, 3)
+        joint_v = joint_v.flatten(2, 3)
+
+        joint_out = optimized_attention(joint_q, joint_k, joint_v, heads=heads, transformer_options=transformer_options)
+
+        seq_img = img.shape[1]
+        img_out = joint_out[:, :seq_img, :]
+        txt_out = joint_out[:, seq_img:, :]
+
+        img_out = self.img_attn_proj(img_out)
+        txt_out = self.txt_attn_proj(txt_out)
+        return img_out, txt_out
+
+
+class JoyImageTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        mlp_width_ratio: float = 4.0,
+        eps: float = 1e-6,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        mlp_hidden_dim = int(dim * mlp_width_ratio)
+
+        self.img_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device)
+        self.img_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+        self.img_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+        self.img_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
+
+        self.txt_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device)
+        self.txt_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+        self.txt_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
+        self.txt_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
+
+        self.attn = JoyImageAttention(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            attention_head_dim=attention_head_dim,
+            eps=eps,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor,
+        temb: torch.Tensor,
+        image_rotary_emb: torch.Tensor,
+        transformer_options=None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        (
+            img_mod1_shift,
+            img_mod1_scale,
+            img_mod1_gate,
+            img_mod2_shift,
+            img_mod2_scale,
+            img_mod2_gate,
+        ) = self.img_mod(temb)
+        (
+            txt_mod1_shift,
+            txt_mod1_scale,
+            txt_mod1_gate,
+            txt_mod2_shift,
+            txt_mod2_scale,
+            txt_mod2_gate,
+        ) = self.txt_mod(temb)
+
+        img_normed = self.img_norm1(hidden_states)
+        txt_normed = self.txt_norm1(encoder_hidden_states)
+        img_modulated = img_normed * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1)
+        txt_modulated = txt_normed * (1 + txt_mod1_scale.unsqueeze(1)) + txt_mod1_shift.unsqueeze(1)
+
+        img_attn, txt_attn = self.attn(img_modulated, txt_modulated, image_rotary_emb, transformer_options=transformer_options)
+
+        hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1)
+
+        img_ffn_normed = self.img_norm2(hidden_states)
+        txt_ffn_normed = self.txt_norm2(encoder_hidden_states)
+        img_ffn_input = img_ffn_normed * (1 + img_mod2_scale.unsqueeze(1)) + img_mod2_shift.unsqueeze(1)
+        txt_ffn_input = txt_ffn_normed * (1 + txt_mod2_scale.unsqueeze(1)) + txt_mod2_shift.unsqueeze(1)
+        hidden_states = hidden_states + self.img_mlp(img_ffn_input) * img_mod2_gate.unsqueeze(1)
+        encoder_hidden_states = encoder_hidden_states + self.txt_mlp(txt_ffn_input) * txt_mod2_gate.unsqueeze(1)
+
+        return hidden_states, encoder_hidden_states
+
+
+class JoyImageTimeTextImageEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        time_freq_dim: int,
+        time_proj_dim: int,
+        text_embed_dim: int,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
+        self.time_embedder = TimestepEmbedding(
+            in_channels=time_freq_dim,
+            time_embed_dim=dim,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+        self.act_fn = nn.SiLU()
+        self.time_proj = operations.Linear(dim, time_proj_dim, bias=True, dtype=dtype, device=device)
+        self.text_embedder = PixArtAlphaTextProjection(
+            text_embed_dim, dim, act_fn="gelu_tanh", dtype=dtype, device=device, operations=operations,
+        )
+
+    def forward(self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor):
+        timestep = self.timesteps_proj(timestep)
+        temb = self.time_embedder(timestep.to(dtype=encoder_hidden_states.dtype)).type_as(encoder_hidden_states)
+        timestep_proj = self.time_proj(self.act_fn(temb))
+        encoder_hidden_states = self.text_embedder(encoder_hidden_states)
+        return temb, timestep_proj, encoder_hidden_states
+
+
+class JoyImageTransformer3DModel(nn.Module):
+    def __init__(
+        self,
+        patch_size: list = [1, 2, 2],
+        in_channels: int = 16,
+        out_channels: Optional[int] = None,
+        hidden_size: int = 3072,
+        num_attention_heads: int = 24,
+        text_dim: int = 4096,
+        mlp_width_ratio: float = 4.0,
+        num_layers: int = 20,
+        rope_dim_list: list = [16, 56, 56],
+        theta: int = 256,
+        image_model=None,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.dtype = dtype
+        self.out_channels = out_channels or in_channels
+        self.patch_size = list(patch_size)
+        self.rope_dim_list = list(rope_dim_list)
+        self.theta = theta
+
+        attention_head_dim = hidden_size // num_attention_heads
+
+        self.img_in = operations.Conv3d(
+            in_channels,
+            hidden_size,
+            kernel_size=tuple(self.patch_size),
+            stride=tuple(self.patch_size),
+            dtype=dtype,
+            device=device,
+        )
+
+        self.condition_embedder = JoyImageTimeTextImageEmbedding(
+            dim=hidden_size,
+            time_freq_dim=256,
+            time_proj_dim=hidden_size * 6,
+            text_embed_dim=text_dim,
+            dtype=dtype,
+            device=device,
+            operations=operations,
+        )
+
+        self.double_blocks = nn.ModuleList([
+            JoyImageTransformerBlock(
+                dim=hidden_size,
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=attention_head_dim,
+                mlp_width_ratio=mlp_width_ratio,
+                dtype=dtype,
+                device=device,
+                operations=operations,
+            )
+            for _ in range(num_layers)
+        ])
+
+        self.norm_out = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
+        self.proj_out = operations.Linear(
+            hidden_size,
+            self.out_channels * math.prod(self.patch_size),
+            bias=True,
+            dtype=dtype,
+            device=device,
+        )
+
+    def _get_rotary_pos_embed_for_range(
+        self,
+        start: Tuple[int, int, int],
+        stop: Tuple[int, int, int],
+        device=None,
+    ) -> torch.Tensor:
+        # 3D RoPE for the patch grid range [start, stop) over (t, h, w). Token order after
+        # reshape(-1) is (t, h, w), matching the img_in Conv3d flatten.
+        rope_dim_list = self.rope_dim_list
+
+        grids = [torch.arange(start[i], stop[i], dtype=torch.float32, device=device) for i in range(3)]
+        mesh = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=0)
+
+        angles_parts = []
+        for i, dim in enumerate(rope_dim_list):
+            pos = mesh[i].reshape(-1)
+            freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device)[: (dim // 2)] / dim))
+            angles_parts.append(torch.outer(pos, freqs))
+
+        angles = torch.cat(angles_parts, dim=1)
+        cos = angles.cos()
+        sin = angles.sin()
+        return torch.stack((cos, -sin, sin, cos), dim=-1).unflatten(-1, (2, 2))
+
+    def get_rotary_pos_embed_for_components(
+        self,
+        component_sizes,
+        device=None,
+    ) -> torch.Tensor:
+        # Per-component 3D RoPE. component_sizes is a list of (t, h, w) patch grid sizes in
+        # sequence order [target, ref0, ref1, ...]; h/w restart at 0 for each component while t
+        # continues from the running offset, giving every image its own temporal position band.
+        freqs_parts = []
+        t_offset = 0
+        for (t, h, w) in component_sizes:
+            freqs = self._get_rotary_pos_embed_for_range(
+                start=(t_offset, 0, 0),
+                stop=(t_offset + t, h, w),
+                device=device,
+            )
+            freqs_parts.append(freqs)
+            t_offset += t
+        return torch.cat(freqs_parts, dim=0).unsqueeze(0).unsqueeze(2)
+
+    def unpatchify(self, x: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
+        c = self.out_channels
+        pt, ph, pw = self.patch_size
+        x = x.reshape(x.shape[0], t, h, w, pt, ph, pw, c)
+        x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)
+        return x.reshape(x.shape[0], c, t * pt, h * ph, w * pw)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        context: torch.Tensor = None,
+        ref_latents=None,
+        control=None,
+        transformer_options=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        transformer_options = {} if transformer_options is None else transformer_options.copy()
+        return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+            self._forward,
+            self,
+            comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+        ).execute(hidden_states, timestep, context, ref_latents, transformer_options, **kwargs)
+
+    def _forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: torch.Tensor,
+        context: torch.Tensor,
+        ref_latents=None,
+        transformer_options=None,
+        **kwargs,
+    ) -> torch.Tensor:
+        pt, ph, pw = self.patch_size
+        _, _, ot, oh, ow = hidden_states.shape
+
+        components = [hidden_states, *(ref_latents or [])]
+        component_sizes = []
+        img_tokens = []
+        for comp in components:
+            comp = comfy.ldm.common_dit.pad_to_patch_size(comp, self.patch_size)
+            _, _, ct, ch, cw = comp.shape
+            component_sizes.append((ct // pt, ch // ph, cw // pw))
+            tokens = self.img_in(comp).flatten(2).transpose(1, 2)  # (B, n_i, D)
+            img_tokens.append(tokens)
+
+        img = torch.cat(img_tokens, dim=1)
+
+        _, vec, txt = self.condition_embedder(timestep, context)
+        vec = vec.unflatten(1, (6, -1))
+
+        image_rotary_emb = self.get_rotary_pos_embed_for_components(
+            component_sizes,
+            device=hidden_states.device,
+        )
+
+        patches_replace = transformer_options.get("patches_replace", {})
+        blocks_replace = patches_replace.get("dit", {})
+        transformer_options["total_blocks"] = len(self.double_blocks)
+        transformer_options["block_type"] = "double"
+        for i, block in enumerate(self.double_blocks):
+            transformer_options["block_index"] = i
+            if ("double_block", i) in blocks_replace:
+                def block_wrap(args):
+                    out = {}
+                    out["img"], out["txt"] = block(
+                        hidden_states=args["img"],
+                        encoder_hidden_states=args["txt"],
+                        temb=args["vec"],
+                        image_rotary_emb=args["pe"],
+                        transformer_options=args.get("transformer_options"),
+                    )
+                    return out
+
+                out = blocks_replace[("double_block", i)]({"img": img,
+                                                           "txt": txt,
+                                                           "vec": vec,
+                                                           "pe": image_rotary_emb,
+                                                           "transformer_options": transformer_options},
+                                                          {"original_block": block_wrap})
+                txt = out["txt"]
+                img = out["img"]
+            else:
+                img, txt = block(
+                    hidden_states=img,
+                    encoder_hidden_states=txt,
+                    temb=vec,
+                    image_rotary_emb=image_rotary_emb,
+                    transformer_options=transformer_options,
+                )
+
+        tt, th, tw = component_sizes[0]
+        target_tokens = tt * th * tw
+        img = img[:, :target_tokens, :]
+        img = self.proj_out(self.norm_out(img))
+        img = self.unpatchify(img, tt, th, tw)
+        return img[:, :, :ot, :oh, :ow]

comfy/ldm/krea2/model.py

@@ -15,6 +15,7 @@ from einops import rearrange
 import comfy.model_management
 import comfy.patcher_extension
 import comfy.ldm.common_dit
+import comfy.utils
 from comfy.ldm.flux.layers import EmbedND, timestep_embedding
 from comfy.ldm.flux.math import apply_rope
 from comfy.ldm.modules.attention import optimized_attention_masked
@@ -73,11 +74,20 @@ class Attention(nn.Module):
         self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype)
 
     def forward(self, x, freqs=None, mask=None, transformer_options={}):
+        transformer_patches = transformer_options.get("patches", {})
+        extra_options = transformer_options.copy()
         q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x)
         q = rearrange(q, "B L (H D) -> B H L D", H=self.heads)
         k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads)
         v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads)
         q, k = self.qknorm(q, k)
+
+        if "block_index" in transformer_options and "attn1_patch" in transformer_patches:
+            for p in transformer_patches["attn1_patch"]:
+                out = p(q, k, v, pe=freqs, attn_mask=mask, extra_options=extra_options)
+                q, k, v = out.get("q", q), out.get("k", k), out.get("v", v)
+                freqs, mask = out.get("pe", freqs), out.get("attn_mask", mask)
+
         if freqs is not None:
             q, k = apply_rope(q, k, freqs)
         if self.kvheads != self.heads:
@@ -86,6 +96,11 @@ class Attention(nn.Module):
             v = v.repeat_interleave(rep, dim=1)
         out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True,
                                          transformer_options=transformer_options)
+
+        if "block_index" in transformer_options and "attn1_output_patch" in transformer_patches:
+            for p in transformer_patches["attn1_output_patch"]:
+                out = p(out, extra_options)
+
         return self.wo(out * F.sigmoid(gate))
 
 
@@ -158,8 +173,44 @@ class SingleStreamBlock(nn.Module):
         self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations)
         self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations)
 
-    def forward(self, x, vec, freqs, mask=None, transformer_options={}):
+    def forward(self, x, vec, freqs, mask=None, timestep_zero_index=None, transformer_options={}):
         prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
+        if timestep_zero_index is not None:
+            bs = x.shape[0]
+            ref_prescale = prescale[bs:]
+            ref_preshift = preshift[bs:]
+            ref_pregate = pregate[bs:]
+            ref_postscale = postscale[bs:]
+            ref_postshift = postshift[bs:]
+            ref_postgate = postgate[bs:]
+            prescale = prescale[:bs]
+            preshift = preshift[:bs]
+            pregate = pregate[:bs]
+            postscale = postscale[:bs]
+            postshift = postshift[:bs]
+            postgate = postgate[:bs]
+
+            pre = self.prenorm(x)
+            pre[:, :timestep_zero_index].mul_(1 + prescale).add_(preshift)
+            pre[:, timestep_zero_index:].mul_(1 + ref_prescale).add_(ref_preshift)
+            attn = self.attn(pre, freqs, mask, transformer_options=transformer_options)
+            del pre
+            attn[:, :timestep_zero_index].mul_(pregate)
+            attn[:, timestep_zero_index:].mul_(ref_pregate)
+            x = x + attn
+            del attn
+
+            post = self.postnorm(x)
+            post[:, :timestep_zero_index].mul_(1 + postscale).add_(postshift)
+            post[:, timestep_zero_index:].mul_(1 + ref_postscale).add_(ref_postshift)
+            mlp = self.mlp(post)
+            del post
+            mlp[:, :timestep_zero_index].mul_(postgate)
+            mlp[:, timestep_zero_index:].mul_(ref_postgate)
+            x = x + mlp
+            del mlp
+            return x
+
         x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options)
         x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
         return x
@@ -181,7 +232,7 @@ class LastLayer(nn.Module):
 class SingleStreamDiT(nn.Module):
     def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4,
                  layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12,
-                 txtheads=20, txtkvheads=20, image_model=None,
+                 txtheads=20, txtkvheads=20, default_ref_method=None, image_model=None,
                  device=None, dtype=None, operations=None, **kwargs):
         super().__init__()
         self.dtype = dtype
@@ -191,6 +242,7 @@ class SingleStreamDiT(nn.Module):
         self.heads = heads
         self.txtdim = txtdim
         self.txtlayers = txtlayers
+        self.default_ref_method = default_ref_method
 
         headdim = features // heads
         axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
@@ -221,61 +273,110 @@ class SingleStreamDiT(nn.Module):
             operations.Linear(features, features * 6, device=device, dtype=dtype),
         )
 
-    def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
+    def forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
         return comfy.patcher_extension.WrapperExecutor.new_class_executor(
             self._forward,
             self,
             comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options),
-        ).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs)
+        ).execute(x, timesteps, context, attention_mask, ref_latents, transformer_options, **kwargs)
 
-    def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
+    def process_img(self, x, index=0):
+        patch = self.patch
+        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
+        h, w = x.shape[-2] // patch, x.shape[-1] // patch
+        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
+
+        img_ids = torch.zeros(h, w, 3, device=x.device, dtype=torch.float32)
+        img_ids[..., 0] = index
+        img_ids[..., 1] = torch.arange(h, device=x.device, dtype=torch.float32)[:, None]
+        img_ids[..., 2] = torch.arange(w, device=x.device, dtype=torch.float32)[None, :]
+        return img, img_ids.reshape(1, h * w, 3).repeat(x.shape[0], 1, 1), h, w
+
+    def _forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
+        transformer_options = transformer_options.copy()
         temporal = x.ndim == 5
         if temporal:
             b5, c5, t5, h5, w5 = x.shape
             x = x.reshape(b5 * t5, c5, h5, w5)
-        bs, c, H_orig, W_orig = x.shape
+        bs, _, h_orig, w_orig = x.shape
         patch = self.patch
-        # Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end.
-        x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
-        H, W = x.shape[-2], x.shape[-1]
-        h_, w_ = H // patch, W // patch
 
         # context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim).
         context = self._unpack_context(context)
 
-        img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
+        img, imgpos, h_, w_ = self.process_img(x)
+        img_tokens = img.shape[1]
+        timestep_zero_index = None
+        ref_method = kwargs.get("ref_latents_method", self.default_ref_method)
+        if ref_method is not None and ref_latents is not None and len(ref_latents) > 0:
+            ref_tokens = []
+            ref_pos = []
+            ref_num_tokens = []
+            for index, ref in enumerate(ref_latents, 1):
+                if ref.ndim == 5:
+                    rb, rc, rt, rh5, rw5 = ref.shape
+                    ref = ref.reshape(rb * rt, rc, rh5, rw5)
+                ref = comfy.utils.repeat_to_batch_size(ref, bs)
+                kontext, kontext_ids, _, _ = self.process_img(ref, index=index)
+                ref_tokens.append(kontext)
+                ref_pos.append(kontext_ids)
+                ref_num_tokens.append(kontext.shape[1])
+            img = torch.cat([img] + ref_tokens, dim=1)
+            imgpos = torch.cat([imgpos] + ref_pos, dim=1)
+            del ref_tokens, ref_pos
+            if ref_method == "index_timestep_zero":
+                timestep_zero_index = img_tokens
+            transformer_options["reference_image_num_tokens"] = ref_num_tokens
+
         img = self.first(img)
 
         t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype))
         tvec = self.tproj(t)
+        if timestep_zero_index is not None:
+            t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype))
+            tvec = torch.cat((tvec, self.tproj(t0)), dim=0)
 
         context = self.txtfusion(context, mask=None, transformer_options=transformer_options)
         context = self.txtmlp(context)
 
-        txtlen, imglen = context.shape[1], img.shape[1]
+        txtlen = context.shape[1]
+        device = context.device
+        txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
+
+        patches = transformer_options.get("patches", {})
+        if "post_input" in patches:
+            for p in patches["post_input"]:
+                out = p({"img": img, "txt": context, "img_ids": imgpos, "txt_ids": txtpos, "transformer_options": transformer_options})
+                img, context = out["img"], out["txt"]
+                imgpos, txtpos = out["img_ids"], out["txt_ids"]
+
         combined = torch.cat((context, img), dim=1)
+        del context, img
+        if timestep_zero_index is not None:
+            timestep_zero_index += txtlen
 
         # Position ids: text at 0, image at (0, h_idx, w_idx).
-        device = combined.device
-        txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
-        imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
-        imgids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
-        imgids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
-        imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
         pos = torch.cat((txtpos, imgpos), dim=1)
+        del txtpos, imgpos
 
         freqs = self.pe_embedder(pos)
+        del pos
 
-        for block in self.blocks:
-            combined = block(combined, tvec, freqs, None, transformer_options=transformer_options)
+        transformer_options["total_blocks"] = len(self.blocks)
+        transformer_options["block_type"] = "single"
+        transformer_options["img_slice"] = [txtlen, combined.shape[1]]
+        for i, block in enumerate(self.blocks):
+            transformer_options["block_index"] = i
+            combined = block(combined, tvec, freqs, None, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options)
 
         final = self.last(combined, t)
-        out = final[:, txtlen:txtlen + imglen, :]
+        del combined
+        out = final[:, txtlen:txtlen + img_tokens, :]
         out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
                         h=h_, w=w_, ph=patch, pw=patch, c=self.channels)
-        out = out[:, :, :H_orig, :W_orig]  # crop padding back off
+        out = out[:, :, :h_orig, :w_orig]  # crop padding back off
         if temporal:
-            out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2)
+            out = out.reshape(b5, t5, self.channels, h_orig, w_orig).movedim(1, 2)
         return out
 
     def _unpack_context(self, context):

comfy/ldm/lightricks/embeddings_connector.py

@@ -6,9 +6,8 @@ import torch
 from comfy.ldm.lightricks.model import (
     CrossAttention,
     FeedForward,
+    freqs_cis_matrix,
     generate_freq_grid_np,
-    interleaved_freqs_cis,
-    split_freqs_cis,
 )
 from torch import nn
 
@@ -244,12 +243,15 @@ class Embeddings1DConnector(nn.Module):
             expected_freqs = dim // 2
             current_freqs = freqs.shape[-1]
             pad_size = expected_freqs - current_freqs
-            cos_freq, sin_freq = split_freqs_cis(
-                freqs, pad_size, self.num_attention_heads
-            )
         else:
-            cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
-        return cos_freq.to(dtype=out_dtype), sin_freq.to(dtype=out_dtype), self.split_rope
+            pad_size = dim % n_elem
+        return freqs_cis_matrix(
+            freqs,
+            pad_size,
+            self.split_rope,
+            self.num_attention_heads,
+            out_dtype,
+        )
 
     def forward(
         self,

comfy/ldm/lightricks/latent_upsampler.py

@@ -97,11 +97,11 @@ class SpatialRationalResampler(nn.Module):
     For dims==3, work per-frame for spatial scaling (temporal axis untouched).
     """
 
-    def __init__(self, mid_channels: int, scale: float):
+    def __init__(self, mid_channels: int, scale: float, operations):
         super().__init__()
         self.scale = float(scale)
         self.num, self.den = _rational_for_scale(self.scale)
-        self.conv = nn.Conv2d(
+        self.conv = operations.Conv2d(
             mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1
         )
         self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
@@ -119,18 +119,18 @@ class SpatialRationalResampler(nn.Module):
 
 class ResBlock(nn.Module):
     def __init__(
-        self, channels: int, mid_channels: Optional[int] = None, dims: int = 3
+        self, channels: int, operations, mid_channels: Optional[int] = None, dims: int = 3
     ):
         super().__init__()
         if mid_channels is None:
             mid_channels = channels
 
-        Conv = nn.Conv2d if dims == 2 else nn.Conv3d
+        Conv = operations.Conv2d if dims == 2 else operations.Conv3d
 
         self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1)
-        self.norm1 = nn.GroupNorm(32, mid_channels)
+        self.norm1 = operations.GroupNorm(32, mid_channels)
         self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1)
-        self.norm2 = nn.GroupNorm(32, channels)
+        self.norm2 = operations.GroupNorm(32, channels)
         self.activation = nn.SiLU()
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -159,6 +159,7 @@ class LatentUpsampler(nn.Module):
 
     def __init__(
         self,
+        operations,
         in_channels: int = 128,
         mid_channels: int = 512,
         num_blocks_per_stage: int = 4,
@@ -179,34 +180,34 @@ class LatentUpsampler(nn.Module):
         self.spatial_scale = float(spatial_scale)
         self.rational_resampler = rational_resampler
 
-        Conv = nn.Conv2d if dims == 2 else nn.Conv3d
+        Conv = operations.Conv2d if dims == 2 else operations.Conv3d
 
         self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1)
-        self.initial_norm = nn.GroupNorm(32, mid_channels)
+        self.initial_norm = operations.GroupNorm(32, mid_channels)
         self.initial_activation = nn.SiLU()
 
         self.res_blocks = nn.ModuleList(
-            [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+            [ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)]
         )
 
         if spatial_upsample and temporal_upsample:
             self.upsampler = nn.Sequential(
-                nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
+                operations.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
                 PixelShuffleND(3),
             )
         elif spatial_upsample:
             if rational_resampler:
                 self.upsampler = SpatialRationalResampler(
-                    mid_channels=mid_channels, scale=self.spatial_scale
+                    mid_channels=mid_channels, scale=self.spatial_scale, operations=operations
                 )
             else:
                 self.upsampler = nn.Sequential(
-                    nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
+                    operations.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
                     PixelShuffleND(2),
                 )
         elif temporal_upsample:
             self.upsampler = nn.Sequential(
-                nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
+                operations.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
                 PixelShuffleND(1),
             )
         else:
@@ -215,11 +216,14 @@ class LatentUpsampler(nn.Module):
             )
 
         self.post_upsample_res_blocks = nn.ModuleList(
-            [ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
+            [ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)]
         )
 
         self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1)
 
+    def get_dtype(self):
+        return getattr(self.initial_conv, "weight_comfy_model_dtype", self.initial_conv.weight.dtype)
+
     def forward(self, latent: torch.Tensor) -> torch.Tensor:
         b, c, f, h, w = latent.shape
 
@@ -266,7 +270,7 @@ class LatentUpsampler(nn.Module):
         return x
 
     @classmethod
-    def from_config(cls, config):
+    def from_config(cls, config, operations):
         return cls(
             in_channels=config.get("in_channels", 4),
             mid_channels=config.get("mid_channels", 128),
@@ -276,6 +280,7 @@ class LatentUpsampler(nn.Module):
             temporal_upsample=config.get("temporal_upsample", False),
             spatial_scale=config.get("spatial_scale", 2.0),
             rational_resampler=config.get("rational_resampler", False),
+            operations=operations,
         )
 
     def config(self):

comfy/ldm/lightricks/model.py

@@ -12,6 +12,8 @@ from torch import nn
 import comfy.patcher_extension
 import comfy.ldm.modules.attention
 import comfy.ldm.common_dit
+import comfy.model_management
+import comfy.quant_ops
 
 from .symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords
 
@@ -322,40 +324,42 @@ class FeedForward(nn.Module):
         return self.net(x)
 
 def apply_rotary_emb(input_tensor, freqs_cis):
-    cos_freqs, sin_freqs = freqs_cis[0], freqs_cis[1]
-    split_pe = freqs_cis[2] if len(freqs_cis) > 2 else False
-    return (
-        apply_split_rotary_emb(input_tensor, cos_freqs, sin_freqs)
-        if split_pe else
-        apply_interleaved_rotary_emb(input_tensor, cos_freqs, sin_freqs)
+    rotation_matrix, split_pe = freqs_cis
+    original_shape = input_tensor.shape
+    input_tensor = input_tensor.reshape(
+        input_tensor.shape[0], input_tensor.shape[1], rotation_matrix.shape[2], -1
     )
 
-def apply_interleaved_rotary_emb(input_tensor, cos_freqs, sin_freqs):  # TODO: remove duplicate funcs and pick the best/fastest one
-    t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
-    t1, t2 = t_dup.unbind(dim=-1)
-    t_dup = torch.stack((-t2, t1), dim=-1)
-    input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
+    if comfy.model_management.in_training:
+        if split_pe:
+            t = input_tensor.reshape(*input_tensor.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
+        else:
+            t = input_tensor.reshape(*input_tensor.shape[:-1], -1, 1, 2)
+        t = t.to(rotation_matrix.dtype)
+        output = rotation_matrix[..., 0] * t[..., 0] + rotation_matrix[..., 1] * t[..., 1]
+        if split_pe:
+            output = output.movedim(-1, -2)
+        output = output.reshape(input_tensor.shape).type_as(input_tensor)
+    elif split_pe:
+        output = comfy.quant_ops.ck.apply_rope_split_half1(input_tensor, rotation_matrix)
+    else:
+        output = comfy.quant_ops.ck.apply_rope1(input_tensor, rotation_matrix)
+    return output.reshape(original_shape)
 
-    out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
+def apply_rotary_emb_qk(q, k, freqs_cis):
+    if comfy.model_management.in_training:
+        return apply_rotary_emb(q, freqs_cis), apply_rotary_emb(k, freqs_cis)
 
-    return out
-
-def apply_split_rotary_emb(input_tensor, cos, sin):
-    needs_reshape = False
-    if input_tensor.ndim != 4 and cos.ndim == 4:
-        B, H, T, _ = cos.shape
-        input_tensor = input_tensor.reshape(B, T, H, -1).swapaxes(1, 2)
-        needs_reshape = True
-    split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
-    first_half_input = split_input[..., :1, :]
-    second_half_input = split_input[..., 1:, :]
-    output = split_input * cos.unsqueeze(-2)
-    first_half_output = output[..., :1, :]
-    second_half_output = output[..., 1:, :]
-    first_half_output.addcmul_(-sin.unsqueeze(-2), second_half_input)
-    second_half_output.addcmul_(sin.unsqueeze(-2), first_half_input)
-    output = rearrange(output, "... d r -> ... (d r)")
-    return output.swapaxes(1, 2).reshape(B, T, -1) if needs_reshape else output
+    rotation_matrix, split_pe = freqs_cis
+    q_shape = q.shape
+    k_shape = k.shape
+    q = q.reshape(q.shape[0], q.shape[1], rotation_matrix.shape[2], -1)
+    k = k.reshape(k.shape[0], k.shape[1], rotation_matrix.shape[2], -1)
+    if split_pe:
+        q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rotation_matrix)
+    else:
+        q, k = comfy.quant_ops.ck.apply_rope(q, k, rotation_matrix)
+    return q.reshape(q_shape), k.reshape(k_shape)
 
 
 class GuideAttentionMask:
@@ -461,9 +465,13 @@ class CrossAttention(nn.Module):
         q = self.q_norm(q)
         k = self.k_norm(k)
 
+        # These norms span all heads, so the per-head RMS+RoPE kernel is not equivalent.
         if pe is not None:
-            q = apply_rotary_emb(q, pe)
-            k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
+            if k_pe is None and q.shape == k.shape:
+                q, k = apply_rotary_emb_qk(q, k, pe)
+            else:
+                q = apply_rotary_emb(q, pe)
+                k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
 
         if mask is None:
             out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, attn_precision=self.attn_precision, transformer_options=transformer_options)
@@ -653,36 +661,23 @@ def generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid):
     )
     return freqs
 
-def interleaved_freqs_cis(freqs, pad_size):
-    cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
-    sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
-    if pad_size != 0:
-        cos_padding = torch.ones_like(cos_freq[:, :, : pad_size])
-        sin_padding = torch.zeros_like(cos_freq[:, :, : pad_size])
-        cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
-        sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
-    return cos_freq, sin_freq
-
-def split_freqs_cis(freqs, pad_size, num_attention_heads):
-    cos_freq = freqs.cos()
-    sin_freq = freqs.sin()
-
-    if pad_size != 0:
-        cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
-        sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
-
-        cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
-        sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
-
-    # Reshape freqs to be compatible with multi-head attention
-    B , T, half_HD = cos_freq.shape
+def freqs_cis_matrix(freqs, pad_size, split_mode, num_attention_heads, out_dtype):
+    cos_freq = freqs.cos().to(out_dtype)
+    sin_freq = freqs.sin().to(out_dtype)
+    if pad_size:
+        matrix_pad_size = pad_size if split_mode else pad_size // 2
+        cos_padding = torch.ones_like(cos_freq[:, :, :matrix_pad_size])
+        sin_padding = torch.zeros_like(sin_freq[:, :, :matrix_pad_size])
+        cos_freq = torch.cat((cos_padding, cos_freq), dim=-1)
+        sin_freq = torch.cat((sin_padding, sin_freq), dim=-1)
 
+    B, T, half_HD = cos_freq.shape
     cos_freq = cos_freq.reshape(B, T, num_attention_heads, half_HD // num_attention_heads)
     sin_freq = sin_freq.reshape(B, T, num_attention_heads, half_HD // num_attention_heads)
-
-    cos_freq = torch.swapaxes(cos_freq, 1, 2)  # (B,H,T,D//2)
-    sin_freq = torch.swapaxes(sin_freq, 1, 2)  # (B,H,T,D//2)
-    return cos_freq, sin_freq
+    rotation_matrix = torch.stack(
+        (cos_freq, -sin_freq, sin_freq, cos_freq), dim=-1
+    )
+    return rotation_matrix.reshape(*rotation_matrix.shape[:-1], 2, 2), split_mode
 
 class LTXBaseModel(torch.nn.Module, ABC):
     """
@@ -885,12 +880,17 @@ class LTXBaseModel(torch.nn.Module, ABC):
             expected_freqs = dim // 2
             current_freqs = freqs.shape[-1]
             pad_size = expected_freqs - current_freqs
-            cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
         else:
             # 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only
             n_elem = 2 * indices_grid.shape[1]
-            cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
-        return cos_freq.to(out_dtype), sin_freq.to(out_dtype), split_mode
+            pad_size = dim % n_elem
+        return freqs_cis_matrix(
+            freqs,
+            pad_size,
+            split_mode,
+            num_attention_heads,
+            out_dtype,
+        )
 
     def _prepare_positional_embeddings(self, pixel_coords, frame_rate, x_dtype):
         """Prepare positional embeddings."""

comfy/ldm/lightricks/vae/audio_vae.py

@@ -185,7 +185,7 @@ class AudioVAE(torch.nn.Module):
             self.autoencoder.mel_bins,
         )
 
-    def num_of_latents_from_frames(self, frames_number: int, frame_rate: int) -> int:
+    def num_of_latents_from_frames(self, frames_number: int, frame_rate: float) -> int:
         return math.ceil((float(frames_number) / frame_rate) * self.latents_per_second)
 
     def run_vocoder(self, mel_spec: torch.Tensor) -> torch.Tensor:

comfy/ldm/lightricks/vae/causal_conv3d.py

@@ -49,6 +49,12 @@ class CausalConv3d(nn.Module):
         )
         self.temporal_cache_state={}
 
+    def _empty_output(self, x):
+        # empty (0 frame) outputs must still have the conv's output channels and spatial dims
+        h = (x.shape[3] + 2 * self.conv.padding[1] - self.conv.kernel_size[1]) // self.conv.stride[1] + 1
+        w = (x.shape[4] + 2 * self.conv.padding[2] - self.conv.kernel_size[2]) // self.conv.stride[2] + 1
+        return x.new_empty((x.shape[0], self.out_channels, 0, h, w))
+
     def forward(self, x, causal: bool = True):
         tid = threading.get_ident()
 
@@ -58,7 +64,7 @@ class CausalConv3d(nn.Module):
             if not causal:
                 padding_length = padding_length // 2
             if x.shape[2] == 0:
-                return x
+                return self._empty_output(x)
             cached = x[:, :, :1, :, :].repeat((1, 1, padding_length, 1, 1))
         pieces = [ cached, x ]
         if is_end and not causal:
@@ -83,7 +89,7 @@ class CausalConv3d(nn.Module):
         elif is_end:
             self.temporal_cache_state[tid] = (None, True)
 
-        return self.conv(x) if x.shape[2] >= self.time_kernel_size else x[:, :, :0, :, :]
+        return self.conv(x) if x.shape[2] >= self.time_kernel_size else self._empty_output(x)
 
     @property
     def weight(self):

comfy/ldm/lightricks/vae/causal_video_autoencoder.py

@@ -390,10 +390,10 @@ class Decoder(nn.Module):
 
         # Compute output channel to be product of all channel-multiplier blocks
         output_channel = base_channels
-        for block_name, block_params in list(reversed(blocks)):
+        for block_name, block_params in blocks:
             block_params = block_params if isinstance(block_params, dict) else {}
             if block_name == "res_x_y":
-                output_channel = output_channel * block_params.get("multiplier", 2)
+                output_channel = block_params.get("in_channels", output_channel * block_params.get("multiplier", 2))
             if block_name == "compress_all":
                 output_channel = output_channel * block_params.get("multiplier", 1)
             if block_name == "compress_space":
@@ -432,7 +432,7 @@ class Decoder(nn.Module):
                     spatial_padding_mode=spatial_padding_mode,
                 )
             elif block_name == "res_x_y":
-                output_channel = output_channel // block_params.get("multiplier", 2)
+                output_channel = block_params.get("out_channels", output_channel // block_params.get("multiplier", 2))
                 block = ResnetBlock3D(
                     dims=dims,
                     in_channels=input_channel,

comfy/ldm/lumina/model.py

@@ -6,6 +6,9 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import comfy.ldm.common_dit
+import comfy.model_management
+import comfy.ops
+import comfy.quant_ops
 
 from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
 from comfy.ldm.modules.attention import optimized_attention_masked
@@ -97,6 +100,7 @@ class JointAttention(nn.Module):
         self.n_local_kv_heads = self.n_kv_heads
         self.n_rep = self.n_local_heads // self.n_local_kv_heads
         self.head_dim = dim // n_heads
+        self.qk_norm = qk_norm
 
         self.qkv = operation_settings.get("operations").Linear(
             dim,
@@ -151,10 +155,21 @@ class JointAttention(nn.Module):
         xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
         xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
 
-        xq = self.q_norm(xq)
-        xk = self.k_norm(xk)
-
-        xq, xk = apply_rope(xq, xk, freqs_cis)
+        if self.qk_norm and not comfy.model_management.in_training:
+            q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, xq, offloadable=True)
+            k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, xk, offloadable=True)
+            epsilon = self.q_norm.eps if self.q_norm.eps is not None else torch.finfo(torch.float32).eps
+            if self.n_local_heads == self.n_local_kv_heads:
+                xq, xk = comfy.quant_ops.ck.rms_rope(xq, xk, freqs_cis, q_scale, k_scale, epsilon)
+            else:
+                xq = comfy.quant_ops.ck.rms_rope1(xq, freqs_cis, q_scale, epsilon)
+                xk = comfy.quant_ops.ck.rms_rope1(xk, freqs_cis, k_scale, epsilon)
+            comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
+            comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
+        else:
+            xq = self.q_norm(xq)
+            xk = self.k_norm(xk)
+            xq, xk = apply_rope(xq, xk, freqs_cis)
 
         n_rep = self.n_local_heads // self.n_local_kv_heads
         if n_rep >= 1:

comfy/ldm/mage_flow/model.py

@@ -0,0 +1,186 @@
+# Mage-Flow (https://github.com/microsoft/Mage) native-resolution MMDiT (MIT)
+# Architecture is a 12-layer variant of the Qwen-Image double-stream block with
+# patch_size=1 (no 2x2 packing), unrotated text tokens and a bf16-rounded
+# timestep frequency table.
+import math
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple
+
+from comfy.ldm.lightricks.model import TimestepEmbedding
+from comfy.ldm.flux.layers import EmbedND
+from comfy.ldm.qwen_image.model import QwenImageTransformerBlock, LastLayer
+import comfy.patcher_extension
+
+
+class MageTimestepProjEmbeddings(nn.Module):
+    def __init__(self, embedding_dim, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=256, time_embed_dim=embedding_dim,
+            dtype=dtype, device=device, operations=operations
+        )
+
+    def forward(self, timestep, hidden_states):
+        half_dim = 128
+        exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) / half_dim
+        emb = torch.exp(exponent).to(timestep.dtype)
+        emb = timestep[:, None].float() * emb[None, :]
+        emb = 1000.0 * emb
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)  # flip_sin_to_cos
+        return self.timestep_embedder(emb.to(dtype=hidden_states.dtype))
+
+
+class MageFlowTransformer2DModel(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 128,
+        out_channels: Optional[int] = 128,
+        num_layers: int = 12,
+        attention_head_dim: int = 128,
+        num_attention_heads: int = 24,
+        joint_attention_dim: int = 2560,
+        axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
+        image_model=None,
+        dtype=None,
+        device=None,
+        operations=None,
+    ):
+        super().__init__()
+        self.dtype = dtype
+        self.patch_size = 1
+        self.in_channels = in_channels
+        self.out_channels = out_channels or in_channels
+        self.inner_dim = num_attention_heads * attention_head_dim
+
+        self.pe_embedder = EmbedND(dim=attention_head_dim, theta=10000, axes_dim=list(axes_dims_rope))
+
+        self.time_text_embed = MageTimestepProjEmbeddings(embedding_dim=self.inner_dim, dtype=dtype, device=device, operations=operations)
+
+        self.txt_norm = operations.RMSNorm(joint_attention_dim, eps=1e-6, dtype=dtype, device=device)
+        self.img_in = operations.Linear(in_channels, self.inner_dim, dtype=dtype, device=device)
+        self.txt_in = operations.Linear(joint_attention_dim, self.inner_dim, dtype=dtype, device=device)
+
+        self.transformer_blocks = nn.ModuleList([
+            QwenImageTransformerBlock(
+                dim=self.inner_dim,
+                num_attention_heads=num_attention_heads,
+                attention_head_dim=attention_head_dim,
+                dtype=dtype,
+                device=device,
+                operations=operations
+            )
+            for _ in range(num_layers)
+        ])
+
+        self.norm_out = LastLayer(self.inner_dim, self.inner_dim, dtype=dtype, device=device, operations=operations)
+        self.proj_out = operations.Linear(self.inner_dim, self.out_channels, bias=True, dtype=dtype, device=device)
+
+    def process_img(self, x, index=0):
+        # patch_size=1: tokens are raw latent pixels, no 2x2 packing.
+        bs, c, h, w = x.shape
+        hidden_states = x.movedim(1, -1).reshape(bs, h * w, c)
+
+        img_ids = torch.zeros((h, w, 3), device=x.device)
+        # Frame axis: positive image index (0 = target, 1..N = reference images).
+        img_ids[:, :, 0] = index
+        # Mage scale_rope centering: positions [-ceil(n/2), floor(n/2)), i.e.
+        # offset by (n - n//2). Differs from Qwen-Image's -(n//2) for odd sizes.
+        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.arange(h, device=x.device)[:, None] - (h - h // 2)
+        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.arange(w, device=x.device)[None, :] - (w - w // 2)
+        return hidden_states, img_ids.reshape(h * w, 3).unsqueeze(0).expand(bs, -1, -1), (h, w)
+
+    def forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
+        return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+            self._forward,
+            self,
+            comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
+        ).execute(x, timestep, context, attention_mask, ref_latents, transformer_options, **kwargs)
+
+    def _forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, control=None, **kwargs):
+        if attention_mask is not None and not torch.is_floating_point(attention_mask):
+            attention_mask = (attention_mask - 1).to(x.dtype) * torch.finfo(x.dtype).max
+
+        hidden_states, img_ids, orig_shape = self.process_img(x)
+        num_embeds = hidden_states.shape[1]
+
+        if ref_latents is not None:
+            ref_num_tokens = []
+            index = 0
+            for ref in ref_latents:
+                index += 1
+                kontext, kontext_ids, _ = self.process_img(ref, index=index)
+                hidden_states = torch.cat([hidden_states, kontext], dim=1)
+                img_ids = torch.cat([img_ids, kontext_ids], dim=1)
+                ref_num_tokens.append(kontext.shape[1])
+            transformer_options = transformer_options.copy()
+            transformer_options["reference_image_num_tokens"] = ref_num_tokens
+
+        # Text tokens are not rotated in Mage-Flow: RoPE at position 0 is the
+        # identity rotation.
+        txt_ids = torch.zeros((x.shape[0], context.shape[1], 3), device=x.device)
+
+        hidden_states = self.img_in(hidden_states)
+        context = self.txt_norm(context)
+        context = self.txt_in(context)
+
+        temb = self.time_text_embed(timestep, hidden_states)
+
+        patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
+        blocks_replace = patches_replace.get("dit", {})
+
+        if "post_input" in patches:
+            for p in patches["post_input"]:
+                out = p({"img": hidden_states, "txt": context, "img_ids": img_ids, "txt_ids": txt_ids, "transformer_options": transformer_options})
+                hidden_states = out["img"]
+                context = out["txt"]
+                img_ids = out["img_ids"]
+                txt_ids = out["txt_ids"]
+
+        ids = torch.cat((txt_ids, img_ids), dim=1)
+        image_rotary_emb = self.pe_embedder(ids).contiguous()
+        del ids, txt_ids, img_ids
+
+        transformer_options["total_blocks"] = len(self.transformer_blocks)
+        transformer_options["block_type"] = "double"
+        for i, block in enumerate(self.transformer_blocks):
+            transformer_options["block_index"] = i
+            if ("double_block", i) in blocks_replace:
+                def block_wrap(args):
+                    out = {}
+                    out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=attention_mask, temb=args["vec"], image_rotary_emb=args["pe"], transformer_options=args["transformer_options"])
+                    return out
+                out = blocks_replace[("double_block", i)]({"img": hidden_states, "txt": context, "vec": temb, "pe": image_rotary_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
+                hidden_states = out["img"]
+                context = out["txt"]
+            else:
+                context, hidden_states = block(
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=context,
+                    encoder_hidden_states_mask=attention_mask,
+                    temb=temb,
+                    image_rotary_emb=image_rotary_emb,
+                    transformer_options=transformer_options,
+                )
+
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": hidden_states, "txt": context, "x": x, "block_index": i, "transformer_options": transformer_options})
+                    hidden_states = out["img"]
+                    context = out["txt"]
+
+            if control is not None:  # Controlnet
+                control_i = control.get("input")
+                if i < len(control_i):
+                    add = control_i[i]
+                    if add is not None:
+                        hidden_states[:, :add.shape[1]] += add
+
+        hidden_states = self.norm_out(hidden_states, temb)
+        hidden_states = self.proj_out(hidden_states)
+
+        hidden_states = hidden_states[:, :num_embeds]
+        h, w = orig_shape
+        return hidden_states.reshape(x.shape[0], h, w, self.out_channels).movedim(-1, 1)

comfy/ldm/mage_flow/vae.py

@@ -0,0 +1,477 @@
+# Mage-VAE (https://github.com/microsoft/Mage) (MIT)
+# Symmetric one-step diffusion codec: DConvEncoder (image -> 128ch latent) and
+# DConvDenoiser + CoD Decoder (latent -> image). 16x downsample, latents in the
+# Flux.2-VAE-anchored space (no patch packing, no BN normalization).
+# Both encode and decode are single forward passes at t=0.
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import comfy.ops
+from comfy.ldm.modules.diffusionmodules.model import vae_attention
+
+ops = comfy.ops.disable_weight_init
+
+
+def nonlinearity(x):
+    return torch.nn.functional.silu(x)
+
+
+def Normalize(in_channels):
+    return ops.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+
+
+def modulate(x, shift, scale):
+    if x.dim() == 4:
+        b, c = x.shape[:2]
+        return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1)
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+class LayerNorm2d(ops.LayerNorm):
+    def __init__(self, num_channels, eps=1e-6, affine=True):
+        super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 3, 1).contiguous()
+        x = super().forward(x)
+        return x.permute(0, 3, 1, 2).contiguous()
+
+
+class TimestepEmbedder(nn.Module):
+    """DConv-style timestep MLP (max_period=10000, freq_size=256)."""
+
+    def __init__(self, hidden_size, frequency_embedding_size=256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            ops.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            ops.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t, dim, max_period=10000):
+        half = dim // 2
+        freqs = torch.exp(
+            -math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
+        ).to(t.device)
+        args = t[:, None].float() * freqs[None]
+        emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if dim % 2:
+            emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
+        return emb
+
+    def forward(self, t, dtype):
+        emb = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(emb.to(dtype))
+
+
+class BottleneckPatchEmbed(nn.Module):
+    """Image patch embed concatenated with a per-patch conditioning vector."""
+
+    def __init__(self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True):
+        super().__init__()
+        self.proj1 = ops.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
+        self.proj2 = ops.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
+
+    def forward(self, x, cond):
+        return self.proj2(torch.cat([self.proj1(x), cond], dim=1))
+
+
+class DiCoBlock(nn.Module):
+    """DConv block with adaLN modulation."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = ops.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+        self.conv3 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
+
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            ops.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = ops.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = ops.Conv2d(ffn, hidden_size, 1, bias=True)
+
+        self.norm1 = LayerNorm2d(hidden_size, affine=False)
+        self.norm2 = LayerNorm2d(hidden_size, affine=False)
+
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            ops.Linear(hidden_size, 6 * hidden_size, bias=True),
+        )
+
+    def forward(self, inp, c):
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+        x = modulate(self.norm1(inp), shift_msa, scale_msa)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + gate_msa[..., None, None] * x
+        x = x + gate_mlp[..., None, None] * self.conv5(
+            F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp)))
+        )
+        return x
+
+
+class EncoderDiCoBlock(nn.Module):
+    """DiCoBlock without adaLN, for the encoder head pathway."""
+
+    def __init__(self, hidden_size, mlp_ratio=4.0):
+        super().__init__()
+        self.conv1 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.conv2 = ops.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
+        self.conv3 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
+        self.ca = nn.Sequential(
+            nn.AdaptiveAvgPool2d(1),
+            ops.Conv2d(hidden_size, hidden_size, 1, bias=True),
+            nn.Sigmoid(),
+        )
+        ffn = int(mlp_ratio * hidden_size)
+        self.conv4 = ops.Conv2d(hidden_size, ffn, 1, bias=True)
+        self.conv5 = ops.Conv2d(ffn, hidden_size, 1, bias=True)
+        self.norm1 = LayerNorm2d(hidden_size)
+        self.norm2 = LayerNorm2d(hidden_size)
+
+    def forward(self, inp):
+        x = self.norm1(inp)
+        x = F.gelu(self.conv2(self.conv1(x)))
+        x = x * self.ca(x)
+        x = self.conv3(x)
+        x = inp + x
+        return x + self.conv5(F.gelu(self.conv4(self.norm2(x))))
+
+
+class NerfEmbedder(nn.Module):
+    """Patch-position embedder used by the DConv decoder x-pathway."""
+
+    def __init__(self, in_channels, hidden_size_input, max_freqs=8):
+        super().__init__()
+        self.max_freqs = max_freqs
+        self.embedder = nn.Sequential(
+            ops.Linear(in_channels + max_freqs ** 2, hidden_size_input, bias=True),
+        )
+
+    def fetch_pos(self, patch_size, device, dtype):
+        pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
+        pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
+        pos_x = pos_x.reshape(-1, 1, 1)
+        pos_y = pos_y.reshape(-1, 1, 1)
+        freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
+        fx = freqs[None, :, None]
+        fy = freqs[None, None, :]
+        coeffs = (1 + fx * fy) ** -1
+        dct_x = torch.cos(pos_x * fx * torch.pi)
+        dct_y = torch.cos(pos_y * fy * torch.pi)
+        return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
+
+    def forward(self, x):
+        B, P2, _ = x.shape
+        ps = int(P2 ** 0.5)
+        dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1)
+        return self.embedder(torch.cat([x, dct], dim=-1))
+
+
+class NerfFinalLayer(nn.Module):
+    def __init__(self, hidden_size, out_channels):
+        super().__init__()
+        self.norm = ops.RMSNorm(hidden_size, eps=1e-6)
+        self.linear = ops.Linear(hidden_size, out_channels, bias=True)
+
+    def forward(self, x):
+        return self.linear(self.norm(x))
+
+
+class MLPResBlock(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.in_ln = ops.LayerNorm(channels, eps=1e-6)
+        self.mlp = nn.Sequential(
+            ops.Linear(channels, channels, bias=True),
+            nn.SiLU(),
+            ops.Linear(channels, channels, bias=True),
+        )
+        self.adaLN_modulation = nn.Sequential(
+            nn.SiLU(),
+            ops.Linear(channels, 3 * channels, bias=True),
+        )
+
+    def forward(self, x, y):
+        shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
+        h = self.in_ln(x) * (1 + scale) + shift
+        return x + gate * self.mlp(h)
+
+
+class SimpleMLPAdaLN(nn.Module):
+    """Final small MLP that maps NerfEmbedder features to per-patch RGB."""
+
+    def __init__(self, in_channels, model_channels, out_channels, z_channels, num_res_blocks, patch_size):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.out_channels = out_channels
+        self.num_res_blocks = num_res_blocks
+        self.patch_size = patch_size
+
+        self.cond_embed = ops.Linear(z_channels, patch_size ** 2 * model_channels)
+        self.input_proj = ops.Linear(in_channels, model_channels)
+
+        self.res_blocks = nn.ModuleList(MLPResBlock(model_channels) for _ in range(num_res_blocks))
+
+    def forward(self, x, c):
+        x = self.input_proj(x)
+        c = self.cond_embed(c).reshape(c.shape[0], self.patch_size ** 2, -1)
+        for block in self.res_blocks:
+            x = block(x, c)
+        return x
+
+
+class ResnetBlock(nn.Module):
+    """GroupNorm + Conv ResBlock used by the CoD Decoder."""
+
+    def __init__(self, *, in_channels, out_channels=None):
+        super().__init__()
+        out_channels = out_channels or in_channels
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.norm1 = Normalize(in_channels)
+        self.conv1 = ops.Conv2d(in_channels, out_channels, 3, padding=1)
+        self.norm2 = Normalize(out_channels)
+        self.conv2 = ops.Conv2d(out_channels, out_channels, 3, padding=1)
+        if in_channels != out_channels:
+            self.nin_shortcut = ops.Conv2d(in_channels, out_channels, 1)
+
+    def forward(self, x):
+        h = self.conv1(nonlinearity(self.norm1(x)))
+        h = self.conv2(nonlinearity(self.norm2(h)))
+        if self.in_channels != self.out_channels:
+            x = self.nin_shortcut(x)
+        return x + h
+
+
+class AttnBlock(nn.Module):
+    """Patched (windowed) self-attention used by the CoD Decoder."""
+
+    def __init__(self, in_channels, patch_size=32):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.norm = Normalize(in_channels)
+        self.q = ops.Conv2d(in_channels, in_channels, 1)
+        self.k = ops.Conv2d(in_channels, in_channels, 1)
+        self.v = ops.Conv2d(in_channels, in_channels, 1)
+        self.proj_out = ops.Conv2d(in_channels, in_channels, 1)
+        # VAE attention selection: full-precision backends only (no sage/quantized attention)
+        self.optimized_attention = vae_attention()
+
+    def forward(self, x):
+        h_ = self.norm(x)
+        Q = self.q(h_)
+        K = self.k(h_)
+        V = self.v(h_)
+
+        d = self.patch_size
+        b, c, H, W = Q.shape
+        pad_h = (d - H % d) % d
+        pad_w = (d - W % d) % d
+        if pad_h or pad_w:
+            Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate")
+            K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate")
+            V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate")
+        _, _, H_pad, W_pad = Q.shape
+        nph, npw = H_pad // d, W_pad // d
+        np_ = nph * npw
+
+        def to_patches(t):
+            return (t.reshape(b, c, nph, d, npw, d)
+                    .permute(0, 2, 4, 1, 3, 5)
+                    .reshape(b * np_, c, d * d))
+
+        # [b*np, c, d*d]: attention over the d*d spatial positions of each window
+        Q = to_patches(Q)
+        K = to_patches(K)
+        V = to_patches(V)
+
+        h_ = self.optimized_attention(Q, K, V)
+        h_ = h_.reshape(b, nph, npw, c, d, d).permute(0, 3, 1, 4, 2, 5).reshape(b, c, H_pad, W_pad)
+        if pad_h or pad_w:
+            h_ = h_[:, :, :H, :W]
+        return x + self.proj_out(h_)
+
+
+class CoDDecoder(nn.Module):
+    """CoD Decoder: latent -> conditioning features for the denoiser (ds=16, light)."""
+
+    def __init__(self, out_ch=384, z_ch=128):
+        super().__init__()
+        self.conv_in = ops.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.block = nn.Sequential(
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+            AttnBlock(out_ch, patch_size=32),
+            ResnetBlock(in_channels=out_ch, out_channels=out_ch),
+        )
+        self.norm_out = Normalize(out_ch)
+        self.conv_out = ops.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
+        self.ada = nn.Identity()
+
+    def forward(self, z):
+        h = self.block(self.conv_in(z))
+        h = self.conv_out(nonlinearity(self.norm_out(h)))
+        return self.ada(h)
+
+
+class DConvEncoder(nn.Module):
+    """DConvEncoder: image -> packed (mean, logvar) latent."""
+
+    def __init__(
+        self,
+        z_ch=128,
+        hidden_size=384,
+        num_blocks=21,
+        patch_size=16,
+        mlp_ratio=4.0,
+        head_size=768,
+        num_head_blocks=2,
+        out_ch_mult=2,
+    ):
+        super().__init__()
+        self.z_ch = z_ch
+        self.patch_size = patch_size
+        self.patch_cond_embed = ops.Conv2d(3, head_size, kernel_size=patch_size, stride=patch_size, bias=True)
+        self.head_blocks = nn.ModuleList([
+            EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) for _ in range(num_head_blocks)
+        ])
+        self.proj_down = ops.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
+        self.z_proj = ops.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True)
+        self.fuse_proj = ops.Conv2d(hidden_size * 2, hidden_size, kernel_size=1, bias=True)
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.blocks = nn.ModuleList([
+            DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)
+        ])
+        self.norm_out = LayerNorm2d(hidden_size)
+        self.proj_out = ops.Conv2d(hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True)
+
+    def forward_pred(self, z_t, t, y):
+        cond = self.patch_cond_embed(y)
+        for block in self.head_blocks:
+            cond = block(cond)
+        cond = self.proj_down(cond)
+
+        s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1))
+        c = self.t_embedder(t.view(-1), y.dtype)
+        for block in self.blocks:
+            s = block(s, c)
+        return self.proj_out(self.norm_out(s))
+
+
+class YEmbedder(nn.Module):
+    """Holds only the CoD decoder (the original Flux2-VAE encoder side is dropped at load)."""
+
+    def __init__(self, ch=384, z_ch=128):
+        super().__init__()
+        self.decoder = CoDDecoder(out_ch=ch, z_ch=z_ch)
+
+
+class DConvDenoiser(nn.Module):
+    """One-step DConv denoiser: latent (via cond) + zero noise -> reconstructed image."""
+
+    def __init__(
+        self,
+        patch_size=16,
+        in_channels=3,
+        hidden_size=384,
+        hidden_size_x=32,
+        mlp_ratio=4.0,
+        num_blocks=24,
+        num_cond_blocks=21,
+        bottleneck_dim=128,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.patch_size = patch_size
+        self.hidden_size = hidden_size
+        self.num_cond_blocks = num_cond_blocks
+
+        self.t_embedder = TimestepEmbedder(hidden_size)
+        self.y_embedder_x = ops.Conv2d(hidden_size, hidden_size_x * patch_size ** 2, 1, 1, 0)
+        self.x_embedder = NerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8)
+        self.s_embedder = BottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
+        self.blocks = nn.ModuleList([
+            DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks)
+        ])
+        self.dec_net = SimpleMLPAdaLN(
+            in_channels=hidden_size_x,
+            model_channels=hidden_size_x,
+            out_channels=in_channels,
+            z_channels=hidden_size,
+            num_res_blocks=num_blocks - num_cond_blocks,
+            patch_size=patch_size,
+        )
+        self.final_layer = NerfFinalLayer(hidden_size_x, in_channels)
+        self.y_embedder = YEmbedder(ch=hidden_size, z_ch=bottleneck_dim)
+
+    def forward(self, x, t, cond):
+        b, _, h, w = x.shape
+        c = self.t_embedder(t.view(-1), x.dtype)
+
+        s = self.s_embedder(x, cond)
+        for block in self.blocks:
+            s = block(s, c)
+
+        length = s.shape[-2] * s.shape[-1]
+        s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
+
+        x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size)
+        x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1)
+        x = x.reshape(b, -1, self.patch_size ** 2, length).permute(0, 3, 2, 1).flatten(0, 1)
+        x = self.x_embedder(x)
+
+        x = self.dec_net(x, s)
+        x = self.final_layer(x)
+        x = x.transpose(1, 2).reshape(b, length, -1)
+        return torch.nn.functional.fold(
+            x.transpose(1, 2).contiguous(), (h, w),
+            kernel_size=self.patch_size, stride=self.patch_size,
+        )
+
+
+class MageVAE(nn.Module):
+    """
+    Encode: DConvEncoder (one-step at t=0) -> posterior mean [B, 128, H/16, W/16]
+    Decode: DConvDenoiser + CoD Decoder    -> image [B, 3, H, W] in [-1, 1]
+    """
+
+    latent_channels = 128
+    downsample_factor = 16
+
+    def __init__(self):
+        super().__init__()
+        self.dconv_encoder = DConvEncoder()
+        self.decoder_model = DConvDenoiser()
+
+    def encode(self, x):
+        B, _, H, W = x.shape
+        ps = self.dconv_encoder.patch_size
+        z_t = torch.zeros(B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype)
+        t = torch.zeros(B, device=x.device, dtype=x.dtype)
+        out = self.dconv_encoder.forward_pred(z_t, t, x)
+        return out[:, : self.latent_channels]  # posterior mean (sample_posterior=False)
+
+    def decode(self, z):
+        cond = self.decoder_model.y_embedder.decoder(z)
+        B = z.shape[0]
+        H = z.shape[2] * self.downsample_factor
+        W = z.shape[3] * self.downsample_factor
+        noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype)
+        t = torch.zeros(B, device=z.device, dtype=z.dtype)
+        return self.decoder_model.forward(noise, t, cond)

comfy/ldm/modules/attention.py

@@ -709,7 +709,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
     return out
 
 try:
-    @torch.library.custom_op("flash_attention::flash_attn", mutates_args=())
+    @torch.library.custom_op("comfy::flash_attn", mutates_args=())
     def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
                     dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor:
         softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale

comfy/ldm/modules/diffusionmodules/model.py

@@ -22,7 +22,7 @@ def torch_cat_if_needed(xl, dim):
     else:
         return None
 
-def get_timestep_embedding(timesteps, embedding_dim):
+def get_timestep_embedding(timesteps, embedding_dim, flip_sin_to_cos=False, downscale_freq_shift=1):
     """
     This matches the implementation in Denoising Diffusion Probabilistic Models:
     From Fairseq.
@@ -33,11 +33,13 @@ def get_timestep_embedding(timesteps, embedding_dim):
     assert len(timesteps.shape) == 1
 
     half_dim = embedding_dim // 2
-    emb = math.log(10000) / (half_dim - 1)
+    emb = math.log(10000) / (half_dim - downscale_freq_shift)
     emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb)
     emb = emb.to(device=timesteps.device)
     emb = timesteps.float()[:, None] * emb[None, :]
     emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+    if flip_sin_to_cos:
+        emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1)
     if embedding_dim % 2 == 1:  # zero pad
         emb = torch.nn.functional.pad(emb, (0,1,0,0))
     return emb

comfy/ldm/pixeldit/model.py

@@ -197,6 +197,9 @@ class PixDiT_T2I(nn.Module):
         """Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
         return s
 
+    def _pre_pixel_blocks(self, s, **kwargs):
+        return s
+
     def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
         H_orig, W_orig = x.shape[2], x.shape[3]
         x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
@@ -226,6 +229,7 @@ class PixDiT_T2I(nn.Module):
             s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
         s = F.silu(t_emb + s)
 
+        s = self._pre_pixel_blocks(s, **kwargs)
         s_cond = s.view(B * L, self.hidden_size)
         x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
         for blk in self.pixel_blocks:

comfy/ldm/pixeldit/pid.py

@@ -13,15 +13,15 @@ from .model import PixDiT_T2I
 from .modules import precompute_freqs_cis_2d
 
 
-class SigmaAwareGatePerTokenPerDim(nn.Module):
+class SigmaAwareGate(nn.Module):
     """gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq.
 
     Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1.
     """
 
-    def __init__(self, dim: int, dtype=None, device=None, operations=None):
+    def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None):
         super().__init__()
-        self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device)
+        self.content_proj = operations.Linear(dim * 2, 1 if per_token else dim, dtype=dtype, device=device)
         self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device))
 
     def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
@@ -36,15 +36,15 @@ class SigmaAwareGatePerTokenPerDim(nn.Module):
 class ResBlock(nn.Module):
     """Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip."""
 
-    def __init__(self, channels: int, num_groups: int = 4, dtype=None, device=None, operations=None):
+    def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", dtype=None, device=None, operations=None):
         super().__init__()
         self.block = nn.Sequential(
             operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
             operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -62,9 +62,13 @@ class LQProjection2D(nn.Module):
         patch_size: int = 16,
         sr_scale: int = 4,
         latent_spatial_down_factor: int = 8,
+        latent_unpatchify_factor: int = 1,
         num_res_blocks: int = 4,
         num_outputs: int = 7,
         interval: int = 2,
+        conv_padding_mode: str = "zeros",
+        gate_per_token: bool = False,
+        pit_output: bool = False,
         dtype=None, device=None, operations=None,
     ):
         super().__init__()
@@ -74,34 +78,38 @@ class LQProjection2D(nn.Module):
         self.patch_size = patch_size
         self.sr_scale = sr_scale
         self.latent_spatial_down_factor = latent_spatial_down_factor
+        self.latent_unpatchify_factor = latent_unpatchify_factor
         self.num_outputs = num_outputs
         self.interval = interval
 
-        z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size
+        effective_latent_channels = latent_channels // (latent_unpatchify_factor * latent_unpatchify_factor)
+        effective_spatial_down_factor = latent_spatial_down_factor // latent_unpatchify_factor
+        z_to_patch_ratio = (sr_scale * effective_spatial_down_factor) / patch_size
         self.z_to_patch_ratio = z_to_patch_ratio
         if z_to_patch_ratio >= 1:
             self.latent_fold_factor = 0
-            latent_proj_in_ch = latent_channels
+            latent_proj_in_ch = effective_latent_channels
         else:
             fold_factor = int(1 / z_to_patch_ratio)
             assert fold_factor * z_to_patch_ratio == 1.0
             self.latent_fold_factor = fold_factor
-            latent_proj_in_ch = latent_channels * fold_factor * fold_factor
+            latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor
 
         layers = [
-            operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
         ]
         for _ in range(num_res_blocks):
-            layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations))
+            layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, dtype=dtype, device=device, operations=operations))
         self.latent_proj = nn.Sequential(*layers)
 
         self.output_heads = nn.ModuleList(
             [operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)]
         )
+        self.pit_head = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) if pit_output else None
         self.gate_modules = nn.ModuleList(
-            [SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations)
+            [SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations)
              for _ in range(num_outputs)]
         )
 
@@ -115,6 +123,11 @@ class LQProjection2D(nn.Module):
         return self.gate_modules[out_idx](x, lq_feature, sigma)
 
     def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor:
+        f = self.latent_unpatchify_factor
+        if f > 1:
+            B, C, H, W = lq_latent.shape
+            lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W)
+            lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f)
         B, z_dim = lq_latent.shape[:2]
         if self.z_to_patch_ratio >= 1:
             if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW:
@@ -134,7 +147,10 @@ class LQProjection2D(nn.Module):
         feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
         B, C, H, W = feat.shape
         tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
-        return [head(tokens) for head in self.output_heads]
+        outputs = [head(tokens) for head in self.output_heads]
+        if self.pit_head is not None:
+            outputs.append(self.pit_head(tokens))
+        return outputs
 
 
 class PidNet(PixDiT_T2I):
@@ -148,6 +164,10 @@ class PidNet(PixDiT_T2I):
         lq_interval: int = 2,
         sr_scale: int = 4,
         latent_spatial_down_factor: int = 8,
+        lq_latent_unpatchify_factor: int = 1,
+        lq_conv_padding_mode: str = "zeros",
+        lq_gate_per_token: bool = False,
+        pit_lq_inject: bool = False,
         rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
         rope_ref_w: int = 1024,
         image_model=None,
@@ -165,6 +185,8 @@ class PidNet(PixDiT_T2I):
         for blk in self.pixel_blocks:
             blk._rope_fn = _pit_rope_fn
 
+        self.pit_lq_inject = pit_lq_inject
+
         num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
         self.lq_proj = LQProjection2D(
             latent_channels=lq_latent_channels,
@@ -173,13 +195,20 @@ class PidNet(PixDiT_T2I):
             patch_size=self.patch_size,
             sr_scale=sr_scale,
             latent_spatial_down_factor=latent_spatial_down_factor,
+            latent_unpatchify_factor=lq_latent_unpatchify_factor,
             num_res_blocks=lq_num_res_blocks,
             num_outputs=num_lq_outputs,
             interval=lq_interval,
+            conv_padding_mode=lq_conv_padding_mode,
+            gate_per_token=lq_gate_per_token,
+            pit_output=pit_lq_inject,
             dtype=dtype,
             device=device,
             operations=operations,
         )
+        self.pit_lq_gate = SigmaAwareGate(
+            self.hidden_size, per_token=lq_gate_per_token, dtype=dtype, device=device, operations=operations
+        ) if pit_lq_inject else None
 
     def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
         return precompute_freqs_cis_2d(
@@ -197,6 +226,11 @@ class PidNet(PixDiT_T2I):
             return s
         return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx)
 
+    def _pre_pixel_blocks(self, s, pid_pit_lq_feature=None, pid_degrade_sigma=None, **kwargs):
+        if pid_pit_lq_feature is None:
+            return s
+        return self.pit_lq_gate(s, pid_pit_lq_feature, pid_degrade_sigma)
+
     def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
         if lq_latent is None:
             raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
@@ -216,12 +250,14 @@ class PidNet(PixDiT_T2I):
             degrade_sigma = degrade_sigma.expand(B).contiguous()
 
         lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws)
+        pit_lq_feature = lq_features.pop() if self.pit_lq_inject else None
 
         return super()._forward(
             x, timesteps,
             context=context, attention_mask=attention_mask,
             transformer_options=transformer_options,
             pid_lq_features=lq_features,
+            pid_pit_lq_feature=pit_lq_feature,
             pid_degrade_sigma=degrade_sigma,
             **kwargs,
         )

comfy/ldm/seedvr/attention.py

@@ -0,0 +1,51 @@
+import torch
+
+from comfy.ldm.modules import attention as _attention
+
+
+def _var_attention_qkv(q, k, v, heads, skip_reshape):
+    if skip_reshape:
+        return q, k, v, q.shape[-1]
+    total_tokens, embed_dim = q.shape
+    head_dim = embed_dim // heads
+    return (
+        q.view(total_tokens, heads, head_dim),
+        k.view(k.shape[0], heads, head_dim),
+        v.view(v.shape[0], heads, head_dim),
+        head_dim,
+    )
+
+
+def _var_attention_output(out, heads, head_dim, skip_output_reshape):
+    if skip_output_reshape:
+        return out
+    return out.reshape(-1, heads * head_dim)
+
+
+def var_attention_optimized_split(q, k, v, heads, cu_seqlens_q, cu_seqlens_k, *args, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    q, k, v, head_dim = _var_attention_qkv(q, k, v, heads, skip_reshape)
+
+    q_split_indices = cu_seqlens_q[1:-1]
+    k_split_indices = cu_seqlens_k[1:-1]
+    if k.shape[0] != v.shape[0]:
+        raise ValueError("cu_seqlens_k does not match v token count")
+
+    q_splits = torch.tensor_split(q, q_split_indices, dim=0)
+    k_splits = torch.tensor_split(k, k_split_indices, dim=0)
+    v_splits = torch.tensor_split(v, k_split_indices, dim=0)
+    if len(q_splits) != len(k_splits) or len(q_splits) != len(v_splits):
+        raise ValueError("cu_seqlens_q and cu_seqlens_k must describe the same sequence count")
+
+    out = []
+    for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits):
+        q_i = q_i.permute(1, 0, 2).unsqueeze(0)
+        k_i = k_i.permute(1, 0, 2).unsqueeze(0)
+        v_i = v_i.permute(1, 0, 2).unsqueeze(0)
+        out_i = _attention.optimized_attention(q_i, k_i, v_i, heads, skip_reshape=True, skip_output_reshape=True)
+        out.append(out_i.squeeze(0).permute(1, 0, 2))
+
+    out = torch.cat(out, dim=0)
+    return _var_attention_output(out, heads, head_dim, skip_output_reshape)
+
+
+optimized_var_attention = var_attention_optimized_split

comfy/ldm/seedvr/color_fix.py

@@ -0,0 +1,301 @@
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from comfy.ldm.seedvr.constants import (
+    CIELAB_DELTA,
+    CIELAB_KAPPA,
+    D65_WHITE_X,
+    D65_WHITE_Z,
+    WAVELET_DECOMP_LEVELS,
+)
+
+
+def wavelet_blur(image: Tensor, radius):
+    max_safe_radius = max(1, min(image.shape[-2:]) // 8)
+    if radius > max_safe_radius:
+        radius = max_safe_radius
+
+    num_channels = image.shape[1]
+
+    kernel_vals = [
+        [0.0625, 0.125, 0.0625],
+        [0.125,  0.25,  0.125],
+        [0.0625, 0.125, 0.0625],
+    ]
+    kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device)
+    kernel = kernel[None, None].repeat(num_channels, 1, 1, 1)
+
+    image = F.pad(image, (radius, radius, radius, radius), mode='replicate')
+    output = F.conv2d(image, kernel, groups=num_channels, dilation=radius)
+
+    return output
+
+def wavelet_decomposition(image: Tensor, levels: int = WAVELET_DECOMP_LEVELS):
+    high_freq = torch.zeros_like(image)
+
+    for i in range(levels):
+        radius = 2 ** i
+        low_freq = wavelet_blur(image, radius)
+        high_freq.add_(image).sub_(low_freq)
+        image = low_freq
+
+    return high_freq, low_freq
+
+def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor) -> Tensor:
+
+    if content_feat.shape != style_feat.shape:
+        if len(content_feat.shape) >= 3:
+            style_feat = F.interpolate(
+                style_feat,
+                size=content_feat.shape[-2:],
+                mode='bilinear',
+                align_corners=False
+            )
+
+    content_high_freq, content_low_freq = wavelet_decomposition(content_feat)
+    del content_low_freq
+
+    style_high_freq, style_low_freq = wavelet_decomposition(style_feat)
+    del style_high_freq
+
+    if content_high_freq.shape != style_low_freq.shape:
+        style_low_freq = F.interpolate(
+            style_low_freq,
+            size=content_high_freq.shape[-2:],
+            mode='bilinear',
+            align_corners=False
+        )
+
+    content_high_freq.add_(style_low_freq)
+
+    return content_high_freq.clamp_(-1.0, 1.0)
+
+def _histogram_matching_channel(source: Tensor, reference: Tensor) -> Tensor:
+    original_shape = source.shape
+
+    source_flat = source.flatten()
+    reference_flat = reference.flatten()
+
+    source_sorted, source_indices = torch.sort(source_flat)
+    reference_sorted, _ = torch.sort(reference_flat)
+    del reference_flat
+
+    n_source = len(source_sorted)
+    n_reference = len(reference_sorted)
+
+    if n_source == n_reference:
+        matched_sorted = reference_sorted
+    else:
+        source_quantiles = torch.linspace(0, 1, n_source, device=source.device)
+        ref_indices = (source_quantiles * (n_reference - 1)).long()
+        ref_indices.clamp_(0, n_reference - 1)
+        matched_sorted = reference_sorted[ref_indices]
+        del source_quantiles, ref_indices, reference_sorted
+
+    del source_sorted, source_flat
+
+    inverse_indices = torch.argsort(source_indices)
+    del source_indices
+    matched_flat = matched_sorted[inverse_indices]
+    del matched_sorted, inverse_indices
+
+    return matched_flat.reshape(original_shape)
+
+def _lab_to_rgb_batch(lab: Tensor, matrix_inv: Tensor, epsilon: float, kappa: float) -> Tensor:
+    L, a, b = lab[:, 0], lab[:, 1], lab[:, 2]
+
+    fy = (L + 16.0) / 116.0
+    fx = a.div(500.0).add_(fy)
+    fz = fy - b / 200.0
+    del L, a, b
+
+    x = torch.where(
+        fx > epsilon,
+        torch.pow(fx, 3.0),
+        fx.mul(116.0).sub_(16.0).div_(kappa)
+    )
+    y = torch.where(
+        fy > epsilon,
+        torch.pow(fy, 3.0),
+        fy.mul(116.0).sub_(16.0).div_(kappa)
+    )
+    z = torch.where(
+        fz > epsilon,
+        torch.pow(fz, 3.0),
+        fz.mul(116.0).sub_(16.0).div_(kappa)
+    )
+    del fx, fy, fz
+
+    x.mul_(D65_WHITE_X)
+    z.mul_(D65_WHITE_Z)
+
+    xyz = torch.stack([x, y, z], dim=1)
+    del x, y, z
+
+    B, _, H, W = xyz.shape
+    xyz_flat = xyz.permute(0, 2, 3, 1).reshape(-1, 3)
+    del xyz
+
+    xyz_flat = xyz_flat.to(dtype=matrix_inv.dtype)
+    rgb_linear_flat = torch.matmul(xyz_flat, matrix_inv.T)
+    del xyz_flat
+
+    rgb_linear = rgb_linear_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2)
+    del rgb_linear_flat
+
+    mask = rgb_linear > 0.0031308
+    rgb = torch.where(
+        mask,
+        torch.pow(torch.clamp(rgb_linear, min=0.0), 1.0 / 2.4).mul_(1.055).sub_(0.055),
+        rgb_linear * 12.92
+    )
+    del mask, rgb_linear
+
+    return torch.clamp(rgb, 0.0, 1.0)
+
+def _rgb_to_lab_batch(rgb: Tensor, matrix: Tensor, epsilon: float, kappa: float) -> Tensor:
+    mask = rgb > 0.04045
+    rgb_linear = torch.where(
+        mask,
+        torch.pow((rgb + 0.055) / 1.055, 2.4),
+        rgb / 12.92
+    )
+    del mask
+
+    B, _, H, W = rgb_linear.shape
+    rgb_flat = rgb_linear.permute(0, 2, 3, 1).reshape(-1, 3)
+    del rgb_linear
+
+    rgb_flat = rgb_flat.to(dtype=matrix.dtype)
+    xyz_flat = torch.matmul(rgb_flat, matrix.T)
+    del rgb_flat
+
+    xyz = xyz_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2)
+    del xyz_flat
+
+    xyz[:, 0].div_(D65_WHITE_X)
+    xyz[:, 2].div_(D65_WHITE_Z)
+
+    epsilon_cubed = epsilon ** 3
+    mask = xyz > epsilon_cubed
+    f_xyz = torch.where(
+        mask,
+        torch.pow(xyz, 1.0 / 3.0),
+        xyz.mul(kappa).add_(16.0).div_(116.0)
+    )
+    del xyz, mask
+
+    L = f_xyz[:, 1].mul(116.0).sub_(16.0)
+    a = (f_xyz[:, 0] - f_xyz[:, 1]).mul_(500.0)
+    b = (f_xyz[:, 1] - f_xyz[:, 2]).mul_(200.0)
+    del f_xyz
+
+    return torch.stack([L, a, b], dim=1)
+
+def lab_color_transfer(
+    content_feat: Tensor,
+    style_feat: Tensor,
+    luminance_weight: float = 0.8
+) -> Tensor:
+    content_feat = wavelet_reconstruction(content_feat, style_feat)
+
+    if content_feat.shape != style_feat.shape:
+        style_feat = F.interpolate(
+            style_feat,
+            size=content_feat.shape[-2:],
+            mode='bilinear',
+            align_corners=False
+        )
+
+    device = content_feat.device
+    original_dtype = content_feat.dtype
+    content_feat = content_feat.float()
+    style_feat = style_feat.float()
+
+    rgb_to_xyz_matrix = torch.tensor([
+        [0.4124564, 0.3575761, 0.1804375],
+        [0.2126729, 0.7151522, 0.0721750],
+        [0.0193339, 0.1191920, 0.9503041]
+    ], dtype=torch.float32, device=device)
+
+    xyz_to_rgb_matrix = torch.tensor([
+        [ 3.2404542, -1.5371385, -0.4985314],
+        [-0.9692660,  1.8760108,  0.0415560],
+        [ 0.0556434, -0.2040259,  1.0572252]
+    ], dtype=torch.float32, device=device)
+
+    epsilon = CIELAB_DELTA
+    kappa = CIELAB_KAPPA
+
+    content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
+    style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
+
+    content_lab = _rgb_to_lab_batch(content_feat, rgb_to_xyz_matrix, epsilon, kappa)
+    del content_feat
+
+    style_lab = _rgb_to_lab_batch(style_feat, rgb_to_xyz_matrix, epsilon, kappa)
+    del style_feat, rgb_to_xyz_matrix
+
+    matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1])
+    matched_b = _histogram_matching_channel(content_lab[:, 2], style_lab[:, 2])
+
+    if luminance_weight < 1.0:
+        matched_L = _histogram_matching_channel(content_lab[:, 0], style_lab[:, 0])
+        result_L = content_lab[:, 0].mul(luminance_weight).add_(matched_L.mul(1.0 - luminance_weight))
+        del matched_L
+    else:
+        result_L = content_lab[:, 0]
+
+    del content_lab, style_lab
+
+    result_lab = torch.stack([result_L, matched_a, matched_b], dim=1)
+    del result_L, matched_a, matched_b
+
+    result_rgb = _lab_to_rgb_batch(result_lab, xyz_to_rgb_matrix, epsilon, kappa)
+    del result_lab, xyz_to_rgb_matrix
+
+    result = result_rgb.mul_(2.0).sub_(1.0)
+    del result_rgb
+
+    result = result.to(original_dtype)
+
+    return result
+
+
+def wavelet_color_transfer(content_feat: Tensor, style_feat: Tensor) -> Tensor:
+    return wavelet_reconstruction(content_feat, style_feat)
+
+
+def adain_color_transfer(content_feat: Tensor, style_feat: Tensor, eps: float = 1e-5) -> Tensor:
+    if content_feat.shape != style_feat.shape:
+        style_feat = F.interpolate(
+            style_feat,
+            size=content_feat.shape[-2:],
+            mode='bilinear',
+            align_corners=False,
+        )
+
+    original_dtype = content_feat.dtype
+    content_feat = content_feat.float()
+    style_feat = style_feat.float()
+
+    b, c = content_feat.shape[:2]
+    content_flat = content_feat.reshape(b, c, -1)
+    style_flat = style_feat.reshape(b, c, -1)
+
+    content_mean = content_flat.mean(dim=2).reshape(b, c, 1, 1)
+    content_std = (content_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1)
+    style_mean = style_flat.mean(dim=2).reshape(b, c, 1, 1)
+    style_std = (style_flat.var(dim=2, correction=0) + eps).sqrt().reshape(b, c, 1, 1)
+    del content_flat, style_flat
+
+    normalized = (content_feat - content_mean) / content_std
+    del content_mean, content_std
+    result = normalized * style_std + style_mean
+    del normalized, style_mean, style_std
+
+    result = result.clamp_(-1.0, 1.0)
+    if result.dtype != original_dtype:
+        result = result.to(original_dtype)
+    return result

comfy/ldm/seedvr/constants.py

@@ -0,0 +1,48 @@
+"""SeedVR2 constants."""
+
+# Temporal chunk-size law: the sampler's activation wall is linear in
+# T_latent * pixel area (17-cell resolution sweep + T bisection, RTX 5090, 3b fp16):
+#   max_latent_frames = (free_GiB - RESERVED - K*SIGMA) / (GIB_PER_MPX_FRAME * megapixels)
+# RESERVED covers model staging plus fixed CUDA/torch overhead; SIGMA is the measured
+# run-to-run spread of the wall; K=4 trades ~10% smaller chunks for ~1e-5 OOM odds.
+SEEDVR2_CHUNK_GIB_PER_MPX_FRAME = 0.55
+SEEDVR2_CHUNK_RESERVED_GIB = 8.5
+SEEDVR2_CHUNK_SIGMA_GIB = 0.55
+SEEDVR2_CHUNK_SIGMA_K = 4
+
+SEEDVR2_7B_VID_DIM = 3072
+SEEDVR2_OOM_BACKOFF_DIVISOR = 2
+SEEDVR2_DTYPE_BYTES_FLOOR = 4
+SEEDVR2_7B_MLP_CHUNK = 8192
+SEEDVR2_ROPE_PARTIAL_CHUNK_TOKENS = 4096  # partial-RoPE application token-chunk.
+SEEDVR2_LATENT_CHANNELS = 16
+
+SEEDVR2_COLOR_MEM_HEADROOM = 0.75
+SEEDVR2_LAB_SCALE_MULTIPLIER = 13
+SEEDVR2_WAVELET_SCALE_MULTIPLIER = 10  # per-frame byte multiplier, wavelet path.
+SEEDVR2_ADAIN_SCALE_MULTIPLIER = 6
+
+BYTEDANCE_VAE_SCALING_FACTOR = 0.9152   # configs_3b/main.yaml:57.
+BYTEDANCE_VAE_SHIFTING_FACTOR = 0.0
+BYTEDANCE_VAE_CONV_MEM_GIB = 0.5
+BYTEDANCE_VAE_NORM_MEM_GIB = 0.5
+BYTEDANCE_LOGVAR_CLAMP_MIN = -30.0      # video_vae_v3/modules/types.py:28.
+BYTEDANCE_LOGVAR_CLAMP_MAX = 20.0       # video_vae_v3/modules/types.py:28.
+BYTEDANCE_GN_CHUNKS_FP16 = 4            # causal_inflation_lib.py:351 (GroupNorm chunk count, fp16).
+BYTEDANCE_GN_CHUNKS_FP32 = 2            # causal_inflation_lib.py:351 (GroupNorm chunk count, fp32).
+BYTEDANCE_BLOCK_OUT_CHANNELS = (128, 256, 512, 512)  # s8_c16_t4_inflation_sd3.yaml:7-11.
+BYTEDANCE_SLICING_SAMPLE_MIN = 4        # s8_c16_t4_inflation_sd3.yaml:22 (slicing_sample_min_size).
+BYTEDANCE_VAE_TEMPORAL_DOWNSAMPLE = 4   # infer.py:230 (temporal_downsample_factor); the 4n+1 factor.
+BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE = 8    # infer.py:231 (spatial_downsample_factor).
+BYTEDANCE_720P_REF_AREA = 45 * 80       # dit_v2/window.py:32 (720p reference area for window scaling).
+BYTEDANCE_MAX_TEMPORAL_WINDOW = 30      # dit_v2/window.py:35 (max temporal window frames).
+BYTEDANCE_ROPE_MAX_FREQ = 256           # dit_v2/rope.py:31 (pixel-RoPE max frequency).
+BYTEDANCE_SINUSOIDAL_DIM = 256          # dit_3b/nadit.py:120 (timestep sinusoidal embed dim).
+
+ROPE_THETA = 10000   # RoPE base; Su et al., "RoFormer", arXiv:2104.09864.
+
+CIELAB_DELTA = 6.0 / 29.0          # CIE 15 (delta).
+CIELAB_KAPPA = (29.0 / 3.0) ** 3   # CIE 15 (kappa).
+D65_WHITE_X = 0.95047              # CIE D65 standard illuminant Xn (Yn = 1).
+D65_WHITE_Z = 1.08883              # CIE D65 standard illuminant Zn.
+WAVELET_DECOMP_LEVELS = 5          # wavelet color-fix decomposition depth (GIMP/Krita; StableSR).

comfy/ldm/wan/model.py

@@ -552,6 +552,7 @@ class WanModel(torch.nn.Module):
                 List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
         """
         # embeddings
+        x_input = x
         x = self.patch_embedding(x.float()).to(x.dtype)
         grid_sizes = x.shape[2:]
         transformer_options["grid_sizes"] = grid_sizes
@@ -564,11 +565,13 @@ class WanModel(torch.nn.Module):
         e0 = self.time_projection(e).unflatten(2, (6, self.dim))
 
         full_ref = None
+        img_offset = 0
         if self.ref_conv is not None:
             full_ref = kwargs.get("reference_latent", None)
             if full_ref is not None:
                 full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
                 x = torch.concat((full_ref, x), dim=1)
+                img_offset = full_ref.shape[1]
 
         # In-context reference (Bernini)
         context_latents = kwargs.get("context_latents", None)
@@ -589,6 +592,7 @@ class WanModel(torch.nn.Module):
             context_img_len = clip_fea.shape[-2]
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -604,6 +608,11 @@ class WanModel(torch.nn.Module):
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
 
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
+                    x = out["img"]
+
         # head
         x = self.head(x, e)
 
@@ -777,6 +786,7 @@ class VaceWanModel(WanModel):
         **kwargs,
     ):
         # embeddings
+        x_input = x
         x = self.patch_embedding(x.float()).to(x.dtype)
         grid_sizes = x.shape[2:]
         transformer_options["grid_sizes"] = grid_sizes
@@ -807,6 +817,7 @@ class VaceWanModel(WanModel):
         x_orig = x
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -822,6 +833,11 @@ class VaceWanModel(WanModel):
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
 
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
+                    x = out["img"]
+
             ii = self.vace_layers_mapping.get(i, None)
             if ii is not None:
                 for iii in range(len(c)):
@@ -887,6 +903,7 @@ class CameraWanModel(WanModel):
         **kwargs,
     ):
         # embeddings
+        x_input = x
         x = self.patch_embedding(x.float()).to(x.dtype)
         if self.control_adapter is not None and camera_conditions is not None:
             x = x + self.control_adapter(camera_conditions).to(x.dtype)
@@ -909,6 +926,7 @@ class CameraWanModel(WanModel):
             context_img_len = clip_fea.shape[-2]
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -924,6 +942,11 @@ class CameraWanModel(WanModel):
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
 
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
+                    x = out["img"]
+
         # head
         x = self.head(x, e)
 
@@ -1335,6 +1358,7 @@ class WanModel_S2V(WanModel):
 
         # embeddings
         bs, _, time, height, width = x.shape
+        x_input = x
         x = self.patch_embedding(x.float()).to(x.dtype)
         if control_video is not None:
             x = x + self.cond_encoder(control_video)
@@ -1379,6 +1403,7 @@ class WanModel_S2V(WanModel):
         context = self.text_embedding(context)
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -1393,6 +1418,12 @@ class WanModel_S2V(WanModel):
                 x = out["img"]
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, transformer_options=transformer_options)
+
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
+                    x = out["img"]
+
             if audio_emb is not None:
                 x = self.audio_injector(x, i, audio_emb, audio_emb_global, seq_len)
         # head
@@ -1599,6 +1630,7 @@ class HumoWanModel(WanModel):
         bs, _, time, height, width = x.shape
 
         # embeddings
+        x_input = x
         x = self.patch_embedding(x.float()).to(x.dtype)
         grid_sizes = x.shape[2:]
         x = x.flatten(2).transpose(1, 2)
@@ -1630,6 +1662,7 @@ class HumoWanModel(WanModel):
             audio = None
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -1645,6 +1678,11 @@ class HumoWanModel(WanModel):
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, audio=audio, transformer_options=transformer_options)
 
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
+                    x = out["img"]
+
         # head
         x = self.head(x, e)
 
@@ -1660,8 +1698,14 @@ class SCAILWanModel(WanModel):
 
     def forward_orig(self, x, t, context, clip_fea=None, freqs=None, transformer_options={}, pose_latents=None, reference_latent=None, ref_mask_latents=None, sam_latents=None, **kwargs):
 
+        x_input = x
+
+        img_offset = 0
         if reference_latent is not None:
             x = torch.cat((reference_latent, x), dim=2)
+            img_offset = (reference_latent.shape[2] // self.patch_size[0]) * \
+                         (reference_latent.shape[3] // self.patch_size[1]) * \
+                         (reference_latent.shape[4] // self.patch_size[2])
 
         # embeddings
         x = self.patch_embedding(x.float()).to(x.dtype)
@@ -1697,6 +1741,7 @@ class SCAILWanModel(WanModel):
             context_img_len = clip_fea.shape[-2]
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -1712,6 +1757,11 @@ class SCAILWanModel(WanModel):
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
 
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
+                    x = out["img"]
+
         # head
         x = self.head(x, e)
 

comfy/ldm/wan/model_animate.py

@@ -493,6 +493,7 @@ class AnimateWanModel(WanModel):
         **kwargs,
     ):
         # embeddings
+        x_input = x
         x = self.patch_embedding(x.float()).to(x.dtype)
         x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
         grid_sizes = x.shape[2:]
@@ -505,11 +506,13 @@ class AnimateWanModel(WanModel):
         e0 = self.time_projection(e).unflatten(2, (6, self.dim))
 
         full_ref = None
+        img_offset = 0
         if self.ref_conv is not None:
             full_ref = kwargs.get("reference_latent", None)
             if full_ref is not None:
                 full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
                 x = torch.concat((full_ref, x), dim=1)
+                img_offset = full_ref.shape[1]
 
         # context
         context = self.text_embedding(context)
@@ -522,6 +525,7 @@ class AnimateWanModel(WanModel):
             context_img_len = clip_fea.shape[-2]
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -537,6 +541,11 @@ class AnimateWanModel(WanModel):
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
 
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
+                    x = out["img"]
+
             if i % 5 == 0 and motion_vec is not None:
                 x = x + self.face_adapter.fuser_blocks[i // 5](x, motion_vec)
 

comfy/ldm/wan/model_wandancer.py

@@ -111,6 +111,7 @@ class WanDancerModel(WanModel):
 
     def forward_orig(self, x, t, context, clip_fea=None, clip_fea_ref=None, freqs=None, audio_embed=None, fps=30, audio_inject_scale=1.0, transformer_options={}, **kwargs):
         # embeddings
+        x_input = x
         if int(fps + 0.5) != 30:
             x = self.patch_embedding_global(x.float()).to(x.dtype)
         else:
@@ -128,11 +129,13 @@ class WanDancerModel(WanModel):
         e0 = self.time_projection(e).unflatten(2, (6, self.dim))
 
         full_ref = None
+        img_offset = 0
         if self.ref_conv is not None: # model has the weight, but this wasn't used in the original pipeline
             full_ref = kwargs.get("reference_latent", None)
             if full_ref is not None:
                 full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
                 x = torch.concat((full_ref, x), dim=1)
+                img_offset = full_ref.shape[1]
 
         # context
         context = self.text_embedding(context)
@@ -163,6 +166,7 @@ class WanDancerModel(WanModel):
             context_img_len += clip_fea_ref.shape[-2]
 
         patches_replace = transformer_options.get("patches_replace", {})
+        patches = transformer_options.get("patches", {})
         blocks_replace = patches_replace.get("dit", {})
         transformer_options["total_blocks"] = len(self.blocks)
         transformer_options["block_type"] = "double"
@@ -177,6 +181,12 @@ class WanDancerModel(WanModel):
                 x = out["img"]
             else:
                 x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
+
+            if "double_block" in patches:
+                for p in patches["double_block"]:
+                    out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
+                    x = out["img"]
+
             if audio_emb is not None:
                 x = self.music_injector(x, i, audio_emb, audio_emb_global=None, seq_len=seq_len, scale=audio_inject_scale)
 

comfy/ldm/wan/uni3c.py

@@ -0,0 +1,149 @@
+# Uni3C controlnet for Wan 2.1: https://github.com/ewrfcas/Uni3C
+# Converted from the original diffusers based implementation.
+import torch
+import torch.nn as nn
+
+from comfy.ldm.flux.layers import EmbedND
+from .model import WanSelfAttention
+
+
+class Uni3CLayerNormZero(nn.Module):
+    def __init__(
+        self,
+        conditioning_dim,
+        embedding_dim,
+        eps=1e-5,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = operations.Linear(conditioning_dim, 3 * embedding_dim, device=device, dtype=dtype)
+        self.norm = operations.LayerNorm(embedding_dim, eps=eps, elementwise_affine=True, device=device, dtype=dtype)
+
+    def forward(self, x, temb):
+        shift, scale, gate = self.linear(self.silu(temb)).chunk(3, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x, gate[:, None, :]
+
+
+class Uni3CAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        dim,
+        ffn_dim,
+        num_heads,
+        time_embed_dim=5120,
+        eps=1e-6,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        operation_settings = {"operations": operations, "device": device, "dtype": dtype}
+        self.norm1 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
+        self.self_attn = WanSelfAttention(dim, num_heads, qk_norm=True, eps=eps, operation_settings=operation_settings)
+        self.norm2 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
+        self.ffn = nn.Sequential(
+            operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
+            operations.Linear(ffn_dim, dim, device=device, dtype=dtype))
+
+    def forward(self, x, temb, freqs):
+        norm_x, gate_msa = self.norm1(x, temb)
+        x = x + gate_msa * self.self_attn(norm_x, freqs)
+        norm_x, gate_ff = self.norm2(x, temb)
+        x = x + gate_ff * self.ffn(norm_x)
+        return x
+
+
+class MaskCamEmbed(nn.Module):
+    def __init__(
+        self,
+        add_channels=7,
+        mid_channels=256,
+        conv_out_dim=5120,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        self.mask_padding = [0, 0, 0, 0, 3, 0]  # first frame conditioning
+        self.mask_proj = nn.Sequential(
+            operations.Conv3d(add_channels, mid_channels, kernel_size=(4, 8, 8), stride=(4, 8, 8), device=device, dtype=dtype),
+            operations.GroupNorm(mid_channels // 8, mid_channels, device=device, dtype=dtype),
+            nn.SiLU())
+        self.mask_zero_proj = operations.Conv3d(mid_channels, conv_out_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2), device=device, dtype=dtype)
+
+    def forward(self, add_inputs):
+        add_padded = torch.nn.functional.pad(add_inputs, self.mask_padding, mode="constant", value=0)
+        add_embeds = self.mask_proj(add_padded)
+        add_embeds = self.mask_zero_proj(add_embeds)
+        add_embeds = add_embeds.flatten(2).transpose(1, 2)
+        return add_embeds
+
+
+class WanUni3CControlnet(nn.Module):
+    def __init__(
+        self,
+        in_channels=36,
+        conv_out_dim=5120,
+        dim=1024,
+        ffn_dim=8192,
+        num_heads=16,
+        num_layers=20,
+        time_embed_dim=5120,
+        out_proj_dim=5120,
+        add_channels=7,
+        mid_channels=256,
+        device=None, dtype=None, operations=None
+    ):
+        super().__init__()
+        patch_size = (1, 2, 2)
+        self.num_layers = num_layers
+
+        self.controlnet_patch_embedding = operations.Conv3d(
+            in_channels, conv_out_dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=torch.float32)
+        self.controlnet_mask_embedding = MaskCamEmbed(add_channels, mid_channels, conv_out_dim, device=device, dtype=dtype, operations=operations)
+
+        if conv_out_dim != dim:
+            self.proj_in = operations.Linear(conv_out_dim, dim, device=device, dtype=dtype)
+        else:
+            self.proj_in = nn.Identity()
+
+        self.controlnet_blocks = nn.ModuleList([
+            Uni3CAttentionBlock(dim, ffn_dim, num_heads, time_embed_dim, device=device, dtype=dtype, operations=operations)
+            for _ in range(num_layers)])
+        self.proj_out = nn.ModuleList([
+            operations.Linear(dim, out_proj_dim, device=device, dtype=dtype)
+            for _ in range(num_layers)])
+
+        head_dim = dim // num_heads
+        self.rope_embedder = EmbedND(dim=head_dim, theta=10000.0, axes_dim=[head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)])
+
+    def rope_encode(self, t_len, h_len, w_len, device=None, dtype=None):
+        img_ids = torch.zeros((t_len, h_len, w_len, 3), device=device, dtype=dtype)
+        img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.arange(t_len, device=device, dtype=dtype).reshape(-1, 1, 1)
+        img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.arange(h_len, device=device, dtype=dtype).reshape(1, -1, 1)
+        img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.arange(w_len, device=device, dtype=dtype).reshape(1, 1, -1)
+        img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
+        freqs = self.rope_embedder(img_ids).movedim(1, 2)
+        return freqs
+
+    def process_input(self, control_input, render_mask=None, camera_embedding=None):
+        # render_mask/camera_embedding are the checkpoint's extra conditioning path, not wired up yet
+        hidden = self.controlnet_patch_embedding(control_input.float()).to(control_input.dtype)
+        t_len, h_len, w_len = hidden.shape[2:]
+        freqs = self.rope_encode(t_len, h_len, w_len, device=hidden.device, dtype=hidden.dtype)
+        hidden = hidden.flatten(2).transpose(1, 2)
+
+        add_inputs = None
+        if camera_embedding is not None and render_mask is not None:
+            add_inputs = torch.cat([render_mask, camera_embedding], dim=1)
+        elif render_mask is not None:
+            add_inputs = render_mask
+
+        if add_inputs is not None:
+            hidden = hidden + self.controlnet_mask_embedding(add_inputs.to(hidden.dtype))
+
+        hidden = self.proj_in(hidden)
+        return hidden, freqs
+
+    def forward_block(self, block_index, hidden, temb, freqs):
+        hidden = self.controlnet_blocks[block_index](hidden, temb, freqs)
+        residual = self.proj_out[block_index](hidden)
+        return hidden, residual

comfy/logging.py

@@ -0,0 +1,10 @@
+import logging
+
+
+DETAIL = 15
+logging.addLevelName(DETAIL, "DETAIL")
+
+
+def detail(message, *args, **kwargs):
+    kwargs.setdefault("stacklevel", 2)
+    logging.log(DETAIL, message, *args, **kwargs)

comfy/model_base.py

@@ -55,8 +55,11 @@ import comfy.ldm.pixeldit.model
 import comfy.ldm.pixeldit.pid
 import comfy.ldm.ace.model
 import comfy.ldm.omnigen.omnigen2
+import comfy.ldm.seedvr.model
 import comfy.ldm.boogu.model
 import comfy.ldm.qwen_image.model
+import comfy.ldm.mage_flow.model
+import comfy.ldm.joyimage.model
 import comfy.ldm.ideogram4.model
 import comfy.ldm.krea2.model
 import comfy.ldm.kandinsky5.model
@@ -932,6 +935,17 @@ class HunyuanDiT(BaseModel):
         out['image_meta_size'] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width, target_height, target_width, 0, 0]]))
         return out
 
+class SeedVR2(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.seedvr.model.NaDiT)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        condition = kwargs.get("condition", None)
+        if condition is not None:
+            out["condition"] = comfy.conds.CONDRegular(condition)
+        return out
+
 class PixArt(BaseModel):
     def __init__(self, model_config, model_type=ModelType.EPS, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.pixart.pixartms.PixArtMS)
@@ -2011,11 +2025,11 @@ class WAN22_WanDancer(WAN21):
 
         fps = kwargs.get("fps", None)
         if fps is not None:
-            out['fps'] = comfy.conds.CONDRegular(torch.FloatTensor([fps]))
+            out['fps'] = comfy.conds.CONDConstant(fps)
 
         audio_inject_scale = kwargs.get("audio_inject_scale", None)
         if audio_inject_scale is not None:
-            out['audio_inject_scale'] = comfy.conds.CONDRegular(torch.FloatTensor([audio_inject_scale]))
+            out['audio_inject_scale'] = comfy.conds.CONDConstant(audio_inject_scale)
         return out
 
 class Hunyuan3Dv2(BaseModel):
@@ -2214,10 +2228,7 @@ class Omnigen2(BaseModel):
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
         ref_latents = kwargs.get("reference_latents", None)
         if ref_latents is not None:
-            latents = []
-            for lat in ref_latents:
-                latents.append(self.process_latent_in(lat))
-            out['ref_latents'] = comfy.conds.CONDList(latents)
+            out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
         return out
 
     def extra_conds_shapes(self, **kwargs):
@@ -2233,8 +2244,8 @@ class Boogu(Omnigen2):
         self.memory_usage_factor_conds = ("ref_latents",)
 
 class QwenImage(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)
+    def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel):
+        super().__init__(model_config, model_type, device=device, unet_model=unet_model)
         self.memory_usage_factor_conds = ("ref_latents",)
 
     def extra_conds(self, **kwargs):
@@ -2264,6 +2275,43 @@ class QwenImage(BaseModel):
             out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
         return out
 
+class MageFlow(QwenImage):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.mage_flow.model.MageFlowTransformer2DModel)
+
+    def process_timestep(self, timestep, **kwargs):
+        # Mage runs in bf16 and rounds its timestep frequency table to the timestep dtype, keep that on fp32 devices.
+        return timestep.to(torch.bfloat16)
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 128, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 128])
+        return out
+
+class JoyImage(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel)
+        self.memory_usage_factor_conds = ("ref_latents",)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        cross_attn = kwargs.get("cross_attn", None)
+        if cross_attn is not None:
+            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
+        return out
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
+        return out
+
 class Ideogram4(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
@@ -2282,12 +2330,30 @@ class Ideogram4(BaseModel):
 class Krea2(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.krea2.model.SingleStreamDiT)
+        self.memory_usage_factor_conds = ("ref_latents",)
 
     def extra_conds(self, **kwargs):
         out = super().extra_conds(**kwargs)
         cross_attn = kwargs.get("cross_attn", None)
         if cross_attn is not None:
             out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            latents = []
+            for lat in ref_latents:
+                latents.append(self.process_latent_in(lat))
+            out['ref_latents'] = comfy.conds.CONDList(latents)
+
+            ref_latents_method = kwargs.get("reference_latents_method", None)
+            if ref_latents_method is not None:
+                out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
+        return out
+
+    def extra_conds_shapes(self, **kwargs):
+        out = {}
+        ref_latents = kwargs.get("reference_latents", None)
+        if ref_latents is not None:
+            out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
         return out
 
 class HunyuanImage21(BaseModel):

comfy/model_detection.py

@@ -470,15 +470,46 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
     # PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
     _lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
     if _lq_w_key in state_dict_keys:
-        in_ch = int(state_dict[_lq_w_key].shape[1])
+        latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])
+        hidden_dim = int(state_dict[_lq_w_key].shape[0])
         _gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
         num_gates = len({k[len(_gate_prefix):].split('.')[0]
                          for k in state_dict_keys if k.startswith(_gate_prefix)})
+        pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys
         dit_config = {"image_model": "pid",
-                      "lq_latent_channels": in_ch,
-                      "latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
+                      "lq_hidden_dim": hidden_dim}
         if num_gates > 0:
             dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
+        if pid_v1_5:
+            pid_v1_5_variants = {
+                16: {  # Flux and QwenImage
+                    "lq_latent_channels": 16,
+                    "latent_spatial_down_factor": 8,
+                    "lq_latent_unpatchify_factor": 1,
+                },
+                32: {  # Flux2 after 2x latent unpatchify
+                    "lq_latent_channels": 128,
+                    "latent_spatial_down_factor": 16,
+                    "lq_latent_unpatchify_factor": 2,
+                },
+            }
+            variant = pid_v1_5_variants.get(latent_proj_in_channels)
+            if variant is None:
+                raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels")
+            gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)]
+            dit_config.update(variant)
+            dit_config.update({
+                "lq_conv_padding_mode": "replicate",
+                "lq_gate_per_token": gate_weight.shape[0] == 1,
+                "pit_lq_inject": True,
+                "rope_ref_h": 2048,
+                "rope_ref_w": 2048,
+            })
+        else:
+            dit_config.update({
+                "lq_latent_channels": latent_proj_in_channels,
+                "latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8,
+            })
         return dit_config
 
     if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys:  # PixelDiT T2I
@@ -598,6 +629,44 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
 
         return dit_config
 
+    seedvr2_7b_separate_key = "{}blocks.35.mlp.vid.proj_out.weight".format(key_prefix)
+    if seedvr2_7b_separate_key in state_dict_keys and state_dict[seedvr2_7b_separate_key].shape[0] == 3072: # seedvr2 7b
+        dit_config = {}
+        dit_config["image_model"] = "seedvr2"
+        dit_config["vid_dim"] = 3072
+        dit_config["heads"] = 24
+        dit_config["num_layers"] = 36
+        # This checkpoint uses separate vid/txt MMModule keys in every block.
+        dit_config["mm_layers"] = 36
+        dit_config["norm_eps"] = 1e-5
+        dit_config["rope_type"] = "rope3d"
+        dit_config["rope_dim"] = 64
+        dit_config["mlp_type"] = "normal"
+        return dit_config
+    if "{}blocks.35.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 7b
+        dit_config = {}
+        dit_config["image_model"] = "seedvr2"
+        dit_config["vid_dim"] = 3072
+        dit_config["heads"] = 24
+        dit_config["num_layers"] = 36
+        # This checkpoint uses shared all.* MMModule keys after the initial blocks.
+        dit_config["mm_layers"] = 10
+        dit_config["norm_eps"] = 1e-5
+        dit_config["rope_type"] = "rope3d"
+        dit_config["rope_dim"] = 64
+        dit_config["mlp_type"] = "swiglu"
+        return dit_config
+    if "{}blocks.31.mlp.all.proj_in_gate.weight".format(key_prefix) in state_dict_keys: # seedvr2 3b
+        dit_config = {}
+        dit_config["image_model"] = "seedvr2"
+        dit_config["vid_dim"] = 2560
+        dit_config["heads"] = 20
+        dit_config["num_layers"] = 32
+        dit_config["norm_eps"] = 1.0e-05
+        dit_config["mlp_type"] = "swiglu"
+        dit_config["vid_out_norm"] = True
+        return dit_config
+
     if '{}head.modulation'.format(key_prefix) in state_dict_keys:  # Wan 2.1
         dit_config = {}
         dit_config["image_model"] = "wan2.1"
@@ -815,6 +884,13 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
             "selected_layer_index": selected_layer_index,
         }
 
+    if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys and '{}proj_out.weight'.format(key_prefix) in state_dict_keys and state_dict['{}txt_norm.weight'.format(key_prefix)].shape[0] == 2560 and state_dict['{}proj_out.weight'.format(key_prefix)].shape[0] == 128:  # Mage-Flow (Qwen Image txt_norm/proj_out are 3584/64)
+        dit_config = {}
+        dit_config["image_model"] = "mage_flow"
+        dit_config["in_channels"] = 128
+        dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
+        return dit_config
+
     if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys:  # Qwen Image
         dit_config = {}
         dit_config["image_model"] = "qwen_image"
@@ -989,6 +1065,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 dit_config["image_model"] = "SAM31"
             return dit_config
 
+    if (
+        '{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys
+        and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys
+        and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys
+        and '{}img_in.weight'.format(key_prefix) in state_dict_keys
+        and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5
+    ):
+        img_in = state_dict['{}img_in.weight'.format(key_prefix)]
+        head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
+        return {
+            "image_model": "joyimage",
+            "in_channels": img_in.shape[1],
+            "hidden_size": img_in.shape[0],
+            "patch_size": list(img_in.shape[2:]),
+            "num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'),
+            "num_attention_heads": img_in.shape[0] // head_dim,
+            "text_dim": 4096,
+        }
+
     if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
         return None
 
@@ -1119,9 +1214,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
 
     return unet_config
 
-def model_config_from_unet_config(unet_config, state_dict=None):
+
+def model_config_from_unet_config(unet_config, state_dict=None, unet_key_prefix=""):
     for model_config in comfy.supported_models.models:
-        if model_config.matches(unet_config, state_dict):
+        if model_config.matches(unet_config, state_dict, unet_key_prefix=unet_key_prefix):
             return model_config(unet_config)
 
     logging.error("no match {}".format(unet_config))
@@ -1131,7 +1227,7 @@ def model_config_from_unet(state_dict, unet_key_prefix, use_base_if_no_match=Fal
     unet_config = detect_unet_config(state_dict, unet_key_prefix, metadata=metadata)
     if unet_config is None:
         return None
-    model_config = model_config_from_unet_config(unet_config, state_dict)
+    model_config = model_config_from_unet_config(unet_config, state_dict, unet_key_prefix)
     if model_config is None and use_base_if_no_match:
         model_config = comfy.supported_models_base.BASE(unet_config)
 

comfy/model_management.py

@@ -34,6 +34,7 @@ import comfy.utils
 import comfy.quant_ops
 import comfy_aimdo.host_buffer
 import comfy_aimdo.vram_buffer
+from comfy.logging import detail
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -473,7 +474,7 @@ except:
 
 SUPPORT_FP8_OPS = args.supports_fp8_compute
 
-AMD_RDNA2_AND_OLDER_ARCH = ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]
+AMD_RDNA2_AND_OLDER_ARCH = ["gfx1030", "gfx1031", "gfx1035", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]
 AMD_ENABLE_MIOPEN_ENV = 'COMFYUI_ENABLE_MIOPEN'
 
 try:
@@ -616,6 +617,8 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
 #Freeing registerables on pressure does imply a GPU sync, so go big on
 #the hysteresis so each expensive sync gives us back a good chunk.
 REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
+WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0
+WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2
 
 def module_size(module):
     module_mem = 0
@@ -630,19 +633,60 @@ def mark_mmap_dirty(storage):
     if mmap_refs is not None:
         DIRTY_MMAPS.add(mmap_refs[0])
 
-def free_pins(size, evict_active=False):
+PIN_SUBSETS = [ "weights", "patches" ]
+LOADED_PIN_SUBSETS = [ "weights-loaded", "patches-loaded" ]
+
+def models_for_pin_eviction(active, current_prompt=None):
+    for loaded_model in current_loaded_models:
+        model = loaded_model.model
+        if model is None or not model.is_dynamic():
+            continue
+        pin_state = model.model.dynamic_pins[model.load_device]
+        if ((active is None or pin_state["active"] == active) and
+            (current_prompt is None or pin_state["current_prompt"] == current_prompt)):
+            yield model
+
+def free_model_pins(size, subsets, current_prompt, active, registrations=False):
     freed_total = 0
-    for loaded_model in reversed(current_loaded_models):
+    for model in models_for_pin_eviction(active, current_prompt=current_prompt):
         if size <= 0:
             return freed_total
-        model = loaded_model.model
-        if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
-            freed = model.partially_unload_ram(size)
-            freed_total += freed
-            size -= freed
+        if registrations:
+            freed = model.unregister_inactive_pins(size, subsets=subsets)
+        else:
+            freed = model.partially_unload_ram(size, subsets=subsets)
+        freed_total += freed
+        size -= freed
     return freed_total
 
-def ensure_pin_budget(size, evict_active=False):
+def pin_eviction_tiers(loaded, evict_active):
+    tiers = [
+        (PIN_SUBSETS, False, None),
+        (LOADED_PIN_SUBSETS, False, None),
+        (LOADED_PIN_SUBSETS, True, None),
+    ]
+    if not loaded:
+        tiers.append((PIN_SUBSETS, True, False))
+        if evict_active:
+            tiers.append((PIN_SUBSETS, True, True))
+    return tiers
+
+def free_pins(size, evict_active=False, loaded=False):
+    freed = 0
+    for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
+        freed += free_model_pins(size - freed, subsets, current_prompt, active)
+    return freed
+
+def should_free_pins_for_ram_pressure(shortfall):
+    if shortfall <= 0:
+        return False
+    if not WINDOWS:
+        return True
+    if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
+        return True
+    return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
+
+def ensure_pin_budget(size, evict_active=False, loaded=False):
     if args.high_ram:
         return True
     if args.fast_disk:
@@ -653,32 +697,21 @@ def ensure_pin_budget(size, evict_active=False):
         return True
 
     to_free = shortfall + PIN_PRESSURE_HYSTERESIS
-    return free_pins(to_free, evict_active=evict_active) >= shortfall
+    return free_pins(to_free, evict_active=evict_active, loaded=loaded) >= shortfall
 
-def free_registrations(shortfall, evict_active=True):
+def free_registrations(shortfall, evict_active=True, loaded=False):
     if MAX_PINNED_MEMORY <= 0:
         return False
     if shortfall <= 0:
         return True
 
     shortfall += REGISTERABLE_PIN_HYSTERESIS
-    for loaded_model in reversed(current_loaded_models):
-        model = loaded_model.model
-        if model is not None and model.is_dynamic() and not model.model.dynamic_pins[model.load_device]["active"]:
-            shortfall -= model.unregister_inactive_pins(shortfall)
-            if shortfall <= 0:
-                return True
-    if evict_active:
-        for loaded_model in current_loaded_models:
-            model = loaded_model.model
-            if model is not None and model.is_dynamic() and model.model.dynamic_pins[model.load_device]["active"]:
-                shortfall -= model.unregister_inactive_pins(shortfall)
-                if shortfall <= 0:
-                    return True
+    for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
+        shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=True)
     return shortfall <= REGISTERABLE_PIN_HYSTERESIS
 
-def ensure_pin_registerable(size, evict_active=True):
-    return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active)
+def ensure_pin_registerable(size, evict_active=True, loaded=False):
+    return free_registrations(TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY, evict_active=evict_active, loaded=loaded)
 
 class LoadedModel:
     def __init__(self, model: ModelPatcher):
@@ -804,6 +837,8 @@ def minimum_inference_memory():
 
 def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins_required=0, ram_required=0):
     cleanup_models_gc()
+    if not for_dynamic:
+        detail("Non dynamic memory free called! memory_required=%s pins_required=%s ram_required=%s", memory_required, pins_required, ram_required)
     unloaded_model = []
     can_unload = []
     unloaded_models = []
@@ -942,6 +977,9 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
             lowvram_model_memory = 0.1
 
         loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
+        vram_used = 0 if is_device_cpu(torch_dev) else loaded_model.model_loaded_memory()
+        ram_used = model.loaded_ram_size() if model.is_dynamic() else loaded_model.model_memory() - vram_used
+        detail("Model loaded: patcher=%s model=%s ram_mb=%.1f vram_mb=%.1f", model.__class__.__name__, model.model.__class__.__name__, ram_used / (1024 ** 2), vram_used / (1024 ** 2))
         current_loaded_models.insert(0, loaded_model)
     return
 
@@ -1368,15 +1406,17 @@ def reset_cast_buffers():
             pin_state = model.model.dynamic_pins[model.load_device]
 
             if pin_state["active"]:
-                *_, buckets = pin_state["weights"]
-                for size, bucket in list(buckets.items()):
-                    bucket[:] = [ entry for entry in bucket if entry[-1] is not None ]
-                    if not bucket:
-                        del buckets[size]
+                for subset in ("weights", "weights-loaded"):
+                    *_, buckets = pin_state[subset]
+                    for size, bucket in list(buckets.items()):
+                        bucket[:] = [ entry for entry in bucket if entry[-1] is not None ]
+                        if not bucket:
+                            del buckets[size]
 
             pin_state["active"] = False
-            model.partially_unload_ram(1e30, subsets=[ "patches" ])
-            model.model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {})
+            model.partially_unload_ram(1e30, subsets=[ "patches", "patches-loaded" ])
+            for subset in ("patches", "patches-loaded"):
+                pin_state[subset] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {})
 
     STREAM_CAST_BUFFERS.clear()
     STREAM_AIMDO_CAST_BUFFERS.clear()

comfy/model_patcher.py

@@ -22,6 +22,7 @@ import collections
 import inspect
 import logging
 import math
+import time
 import uuid
 from typing import Callable, Optional
 
@@ -37,11 +38,58 @@ import comfy.patcher_extension
 import comfy.utils
 import comfy_aimdo.host_buffer
 from comfy.comfy_types import UnetWrapperFunction
+from comfy.logging import detail
 from comfy.quant_ops import QuantizedTensor
 from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
 
 import comfy_aimdo.model_vbar
 
+def is_model_patcher_output(output):
+    return isinstance(output, ModelPatcher) or isinstance(getattr(output, "patcher", None), ModelPatcher)
+
+class PromptModelTracker:
+    def __init__(self):
+        self.models = {}
+
+    def start(self):
+        self.end()
+
+    def add(self, outputs):
+        if isinstance(outputs, collections.abc.Mapping):
+            outputs = outputs.values()
+        elif not isinstance(outputs, (list, tuple)):
+            outputs = (outputs,)
+
+        for output in outputs:
+            if isinstance(output, (collections.abc.Mapping, list, tuple)):
+                self.add(output)
+                continue
+
+            models = []
+            if isinstance(output, ModelPatcher):
+                models.append(output)
+                models.extend(output.model_patches_models())
+                models.extend(output.get_nested_additional_models())
+            else:
+                patcher = getattr(output, "patcher", None)
+                if isinstance(patcher, ModelPatcher):
+                    models.append(patcher)
+                get_models = getattr(output, "get_models", None)
+                if callable(get_models):
+                    models.extend(get_models())
+
+            for model in models:
+                if not isinstance(model, ModelPatcher) or not model.is_dynamic():
+                    continue
+                key = (id(model.model), model.load_device)
+                self.models[key] = model
+                model.set_in_use_by_current_prompt(True)
+
+    def end(self):
+        for model in self.models.values():
+            model.set_in_use_by_current_prompt(False)
+        self.models.clear()
+
 def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
     to = model_options["transformer_options"].copy()
 
@@ -1724,14 +1772,20 @@ class ModelPatcherDynamic(ModelPatcher):
             self.model.dynamic_pins[device] = {
                 "weights": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
                 "patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
+                "weights-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
+                "patches-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
                 "hostbufs_initialized": False,
                 "failed": False,
                 "active": False,
+                "current_prompt": False,
             }
 
     def is_dynamic(self):
         return True
 
+    def set_in_use_by_current_prompt(self, in_use):
+        self.model.dynamic_pins[self.load_device]["current_prompt"] = in_use
+
     def _vbar_get(self, create=False):
         if self.load_device == torch.device("cpu"):
             return None
@@ -1802,6 +1856,8 @@ class ModelPatcherDynamic(ModelPatcher):
                 hostbuf_size = comfy.model_management.pinned_hostbuf_size(self.model_size())
                 pin_state["weights"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
                 pin_state["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
+                pin_state["weights-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
+                pin_state["patches-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
                 pin_state["hostbufs_initialized"] = True
             pin_state["failed"] = False
             pin_state["active"] = True
@@ -1935,20 +1991,37 @@ class ModelPatcherDynamic(ModelPatcher):
         assert self.load_device != torch.device("cpu")
 
         vbar = self._vbar_get()
-        freed = 0 if vbar is None else vbar.free_memory(memory_to_free)
+        vbar_freed = 0 if vbar is None else vbar.free_memory(memory_to_free)
+        freed = vbar_freed
 
+        backup_freed = 0
         if freed < memory_to_free:
-            freed += self.restore_loaded_backups()
+            backup_freed = self.restore_loaded_backups()
+            freed += backup_freed
+
+        method = "vbar+backups" if vbar_freed and backup_freed else "vbar" if vbar_freed else "backups" if backup_freed else "none"
+        free_methods = getattr(self, "_free_methods", {})
+        free_methods[method] = free_methods.get(method, 0) + 1
+        self._free_methods = free_methods
+        now = time.monotonic()
+        if now - getattr(self, "_last_free_log_time", 0) >= 5:
+            requested = "all" if memory_to_free >= 1e30 else f"{memory_to_free / (1024 ** 2):.1f}MB"
+            prevailing_method = max(free_methods, key=free_methods.get)
+            detail("AIMDO free: model=%s device=%s prevailing_method=%s methods=%s requested=%s vbar_mb=%.1f backups_mb=%.1f", self.model.__class__.__name__, self.load_device, prevailing_method, free_methods, requested, vbar_freed / (1024 ** 2), backup_freed / (1024 ** 2))
+            self._free_methods = {}
+            self._last_free_log_time = now
 
         return freed
 
     def loaded_ram_size(self):
-        return (self.model.dynamic_pins[self.load_device]["weights"][0].size)
+        pin_state = self.model.dynamic_pins[self.load_device]
+        return pin_state["weights"][0].size + pin_state["weights-loaded"][0].size
 
     def pinned_memory_size(self):
-        return (self.model.dynamic_pins[self.load_device]["weights"][3][0])
+        pin_state = self.model.dynamic_pins[self.load_device]
+        return pin_state["weights"][3][0] + pin_state["weights-loaded"][3][0]
 
-    def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights", "patches" ]):
+    def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]):
         freed = 0
         pin_state = self.model.dynamic_pins[self.load_device]
         for subset in subsets:
@@ -1956,15 +2029,17 @@ class ModelPatcherDynamic(ModelPatcher):
             split = stack_split[0]
             while split >= 0:
                 module, offset = stack[split]
+                module_pin = module._pins[subset]
                 split -= 1
                 stack_split[0] = split
-                if not module._pin_registered:
+                if not module_pin["registered"]:
                     continue
-                size = module._pin.numel() * module._pin.element_size()
-                if torch.cuda.cudart().cudaHostUnregister(module._pin.data_ptr()) != 0:
+                pin = module_pin["pin"]
+                size = pin.numel() * pin.element_size()
+                if torch.cuda.cudart().cudaHostUnregister(pin.data_ptr()) != 0:
                     comfy.model_management.discard_cuda_async_error()
                     continue
-                module._pin_registered = False
+                module_pin["registered"] = False
                 comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
                 pinned_size[0] = max(0, pinned_size[0] - size)
                 freed += size
@@ -1973,20 +2048,23 @@ class ModelPatcherDynamic(ModelPatcher):
                     return freed
         return freed
 
-    def partially_unload_ram(self, ram_to_unload, subsets=[ "weights", "patches" ]):
+    def partially_unload_ram(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]):
         freed = 0
         pin_state = self.model.dynamic_pins[self.load_device]
         for subset in subsets:
             hostbuf, stack, stack_split, pinned_size, *_ = pin_state[subset]
             while len(stack) > 0:
                 module, offset = stack.pop()
-                size = module._pin.numel() * module._pin.element_size()
-                module._pin_balancer_entry[-1] = None
-                del module._pin_balancer_entry
-                del module._pin
-                hostbuf.truncate(offset, do_unregister=module._pin_registered)
+                module_pin = module._pins[subset]
+                pin = module_pin["pin"]
+                size = pin.numel() * pin.element_size()
+                module_pin["balancer_entry"][-1] = None
+                del module_pin["balancer_entry"]
+                del module_pin["pin"]
+                registered = module_pin["registered"]
+                hostbuf.truncate(offset, do_unregister=registered)
                 stack_split[0] = min(stack_split[0], len(stack) - 1)
-                if module._pin_registered:
+                if registered:
                     comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
                     pinned_size[0] = max(0, pinned_size[0] - size)
                 freed += size

comfy/ops.py

@@ -144,8 +144,13 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
         needs_cast = False
 
         xfer_source = [ s.weight, s.bias ]
-
-        pin = comfy.pinned_memory.get_pin(s)
+        subset = "weights"
+        pin = comfy.pinned_memory.get_pin(s, subset=subset)
+        if pin is None and not args.fast_disk:
+            loaded_pin = comfy.pinned_memory.get_pin(s, subset="weights-loaded")
+            if loaded_pin is not None or signature is not None:
+                subset = "weights-loaded"
+                pin = loaded_pin
         if pin is not None:
             xfer_source = [ pin ]
 
@@ -182,12 +187,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
             if pin is not None:
                 cast_maybe_lowvram_patch([pin], dest, offload_stream)
                 return
-            if signature is None or args.high_ram:
+            if signature is None or not args.fast_disk or args.high_ram:
                 comfy.pinned_memory.pin_memory(m, subset=subset, size=size)
                 pin = comfy.pinned_memory.get_pin(m, subset=subset)
             cast_maybe_lowvram_patch(source, pin, offload_stream, xfer_dest2=dest)
 
-        handle_pin(s, pin, xfer_source, xfer_dest, size=dest_size)
+        handle_pin(s, pin, xfer_source, xfer_dest, subset=subset, size=dest_size)
 
         for param_key in ("weight", "bias"):
             lowvram_source = getattr(s, param_key + "_lowvram_function", None)
@@ -197,8 +202,16 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
                 lowvram_dest = get_cast_buffer(lowvram_size)
                 lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)
 
-                pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches")
-                handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset="patches", size=lowvram_size)
+                subset = "patches"
+                pin = comfy.pinned_memory.get_pin(lowvram_source, subset=subset)
+                if pin is None:
+                    loaded_pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches-loaded")
+                    if loaded_pin is not None:
+                        subset = "patches-loaded"
+                        pin = loaded_pin
+                    elif signature is not None and not args.fast_disk:
+                        subset = "patches-loaded"
+                handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset=subset, size=lowvram_size)
 
 
         prefetch["xfer_dest"] = xfer_dest
@@ -1104,6 +1117,21 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
                 scales["convrot_groupsize"] = int(
                     layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
                 )
+        elif module.quant_format == "convrot_w4a4":
+            scale = pop_scale("weight_scale")
+            if scale is None:
+                raise ValueError(f"Missing ConvRot W4A4 weight scale for layer {layer_name}")
+            params_conf = layer_conf.get("params", {})
+            if not isinstance(params_conf, dict):
+                params_conf = {}
+            scales = {
+                "scale": scale,
+                "convrot_groupsize": int(
+                    layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
+                ),
+                "quant_group_size": 64,
+                "linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
+            }
         else:
             raise ValueError(f"Unsupported quantization format: {module.quant_format}")
 
@@ -1150,6 +1178,11 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
         if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False):
             quant_conf["convrot"] = True
             quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
+        elif module.quant_format == "convrot_w4a4":
+            quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
+            linear_dtype = getattr(params, "linear_dtype", "int4")
+            if linear_dtype != "int4":
+                quant_conf["linear_dtype"] = linear_dtype
         if extra_quant_conf:
             quant_conf.update(extra_quant_conf)
         sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
@@ -1237,7 +1270,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 run_every_op()
 
                 input_shape = input.shape
-                reshaped_3d = False
+                reshaped_nd = False
                 #If cast needs to apply lora, it should be done in the compute dtype
                 compute_dtype = input.dtype
 
@@ -1274,12 +1307,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 # Inference path (unchanged)
                 if _use_quantized and quantize_input:
 
-                    # Reshape 3D tensors to 2D for quantization (needed for NVFP4 and others)
-                    input_reshaped = input.reshape(-1, input_shape[2]) if input.ndim == 3 else input
+                    # Reshape >=3D tensors to 2D for quantization (needed for NVFP4 and others)
+                    input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input
 
                     # Fall back to non-quantized for non-2D tensors
                     if input_reshaped.ndim == 2:
-                        reshaped_3d = input.ndim == 3
+                        reshaped_nd = input.ndim >= 3
                         # dtype is now implicit in the layout class
                         scale = getattr(self, 'input_scale', None)
                         if scale is not None:
@@ -1294,9 +1327,9 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                     weight_only_quant=weight_only_quant,
                 )
 
-                # Reshape output back to 3D if input was 3D
-                if reshaped_3d:
-                    output = output.reshape((input_shape[0], input_shape[1], self.weight.shape[0]))
+                # Reshape output back to original rank if input was >2D
+                if reshaped_nd:
+                    output = output.reshape((*input_shape[:-1], self.weight.shape[0]))
 
                 return output
 
@@ -1430,6 +1463,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 }
                 if hasattr(params, "block_scale"): # NVFP4
                     kwargs["block_scale"] = params.block_scale[i]
+                if hasattr(params, "quant_group_size"):
+                    kwargs["quant_group_size"] = params.quant_group_size
+                if hasattr(params, "convrot_groupsize"):
+                    kwargs["convrot_groupsize"] = params.convrot_groupsize
+                if hasattr(params, "linear_dtype"):
+                    kwargs["linear_dtype"] = params.linear_dtype
                 return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs))
 
             def state_dict(self, *args, destination=None, prefix="", **kwargs):
@@ -1443,12 +1482,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 if layer_conf is not None:
                     layer_conf = json.loads(layer_conf.numpy().tobytes())
 
-                # Only fp8 makes sense for embeddings (per-row dequant via index select).
+                # Only fp8 and int8_tensorwise support per-row dequant via index select.
                 # Block-scaled formats (NVFP4, MXFP8) can't do per-row lookup efficiently.
                 quant_format = layer_conf.get("format") if layer_conf is not None else None
                 manually_loaded_keys = []
 
-                if quant_format in ("float8_e4m3fn", "float8_e5m2") and weight_key in state_dict:
+                if quant_format in ("float8_e4m3fn", "float8_e5m2", "int8_tensorwise") and weight_key in state_dict:
                     self.quant_format = quant_format
                     qconfig = QUANT_ALGOS[quant_format]
                     self.layout_type = qconfig["comfy_tensor_layout"]
@@ -1462,10 +1501,16 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                         scale = scale.float()
                         manually_loaded_keys.append(scale_key)
 
+                    extra = {}
+                    if quant_format == "int8_tensorwise" and layer_conf.get("convrot", False):
+                        # rotated embedding table: record it so the forward un-rotates after lookup
+                        extra["convrot"] = True
+                        extra["convrot_groupsize"] = int(layer_conf.get("convrot_groupsize", 256))
                     params = layout_cls.Params(
                         scale=scale if scale is not None else torch.ones((), dtype=torch.float32),
                         orig_dtype=MixedPrecisionOps._compute_dtype,
                         orig_shape=(self.num_embeddings, self.embedding_dim),
+                        **extra,
                     )
                     self.weight = torch.nn.Parameter(
                         QuantizedTensor(weight.to(dtype=qconfig["storage_t"]), qconfig["comfy_tensor_layout"], params),
@@ -1487,15 +1532,23 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
             def forward_comfy_cast_weights(self, input, out_dtype=None):
                 weight = self.weight
 
-                # Optimized path: lookup in fp8, dequantize only the selected rows.
+                # Optimized path: lookup in fp8/int8, dequantize only the selected rows.
                 if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
                     qdata, _, offload_stream = cast_bias_weight(self, device=input.device, dtype=weight.dtype, offloadable=True)
                     if isinstance(qdata, QuantizedTensor):
-                        scale = qdata._params.scale
+                        params = qdata._params
+                        scale = params.scale
                         qdata = qdata._qdata
                     else:
+                        params = weight._params
                         scale = None
 
+                    # int8: per-row scale possible ConvRot, so let the layout do the gather
+                    if self.quant_format == "int8_tensorwise":
+                        x = get_layout_class(self.layout_type).dequantize_embedding(qdata, params, input)
+                        uncast_bias_weight(self, qdata, None, offload_stream)
+                        return x if out_dtype is None else x.to(dtype=out_dtype)
+
                     x = torch.nn.functional.embedding(
                         input, qdata, self.padding_idx, self.max_norm,
                         self.norm_type, self.scale_grad_by_freq, self.sparse)

comfy/pinned_memory.py

@@ -9,14 +9,14 @@ import torch
 
 from comfy.cli_args import args
 
-def _add_to_bucket(module, buckets, size, priority):
+def _add_to_bucket(module, module_pin, buckets, size, priority):
     bucket = buckets.setdefault(size, [])
     entry = [-priority, 0, module]
     entry[1] = id(entry)
     bisect.insort(bucket, entry)
-    module._pin_balancer_entry = entry
+    module_pin["balancer_entry"] = entry
 
-def _steal_pin(module, stack, buckets, size, priority):
+def _steal_pin(module, stack, buckets, size, priority, subset):
     bucket = buckets.get(size)
     if bucket is None:
         return False
@@ -31,34 +31,39 @@ def _steal_pin(module, stack, buckets, size, priority):
         return False
 
     *_, victim = bucket.pop()
-    module._pin = victim._pin
-    module._pin_registered = victim._pin_registered
-    module._pin_stack_index = victim._pin_stack_index
-    stack[module._pin_stack_index] = (module, stack[module._pin_stack_index][1])
+    module_pin = module._pins[subset]
+    victim_pin = victim._pins[subset]
+    module_pin["pin"] = victim_pin["pin"]
+    module_pin["registered"] = victim_pin["registered"]
+    module_pin["stack_index"] = victim_pin["stack_index"]
+    stack_index = module_pin["stack_index"]
+    stack[stack_index] = (module, stack[stack_index][1])
 
-    victim._pin_registered = False
-    del victim._pin
-    del victim._pin_stack_index
-    del victim._pin_balancer_entry
+    victim_pin["registered"] = False
+    del victim_pin["pin"]
+    del victim_pin["stack_index"]
+    del victim_pin["balancer_entry"]
 
-    _add_to_bucket(module, buckets, size, priority)
+    _add_to_bucket(module, module_pin, buckets, size, priority)
     return True
 
 def get_pin(module, subset="weights"):
-    pin = getattr(module, "_pin", None)
-    if pin is None or module._pin_registered or args.disable_pinned_memory:
+    pins = module.__dict__.get("_pins")
+    module_pin = None if pins is None else pins.get(subset)
+    pin = None if module_pin is None else module_pin.get("pin")
+    if pin is None or module_pin["registered"] or args.disable_pinned_memory:
         return pin
 
     _, _, stack_split, pinned_size, *_ = module._pin_state[subset]
     size = pin.nbytes
-    comfy.model_management.ensure_pin_registerable(size)
+    comfy.model_management.ensure_pin_registerable(size, loaded=subset.endswith("-loaded"))
 
     if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
         comfy.model_management.discard_cuda_async_error()
         return pin
 
-    module._pin_registered = True
-    stack_split[0] = max(stack_split[0], module._pin_stack_index)
+    module_pin["registered"] = True
+    stack_split[0] = max(stack_split[0], module_pin["stack_index"])
     comfy.model_management.TOTAL_PINNED_MEMORY += size
     pinned_size[0] += size
     return pin
@@ -72,23 +77,26 @@ def pin_memory(module, subset="weights", size=None):
     if pin is not None:
         return
 
+    pins = module.__dict__.setdefault("_pins", {})
+    module_pin = pins.setdefault(subset, {})
     hostbuf, stack, stack_split, pinned_size, counter, buckets = pin_state[subset]
     if size is None:
         size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
-    offset = hostbuf.size
     registerable_size = size
-    priority = getattr(module, "_pin_balancer_priority", None)
+    loaded = subset.endswith("-loaded")
+    priority = module_pin.get("balancer_priority")
 
     if priority is None:
         priority = comfy.utils.bit_reverse_range(counter[0], 16)
         counter[0] += 1
-        module._pin_balancer_priority = priority
+        module_pin["balancer_priority"] = priority
 
     comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
-    if (not comfy.model_management.ensure_pin_budget(size) or
-        not comfy.model_management.ensure_pin_registerable(registerable_size)):
-        return _steal_pin(module, stack, buckets, size, priority)
+    if (not comfy.model_management.ensure_pin_budget(size, loaded=loaded) or
+        not comfy.model_management.ensure_pin_registerable(registerable_size, loaded=loaded)):
+        return _steal_pin(module, stack, buckets, size, priority, subset)
 
+    offset = hostbuf.size
     extended = False
     try:
         hostbuf.extend(size=size, register=False)
@@ -97,23 +105,23 @@ def pin_memory(module, subset="weights", size=None):
         pin.untyped_storage()._comfy_hostbuf = hostbuf
         if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
             comfy.model_management.discard_cuda_async_error()
-            comfy.model_management.free_registrations(size)
+            comfy.model_management.free_registrations(size, loaded=loaded)
             if torch.cuda.cudart().cudaHostRegister(pin.data_ptr(), size, 1) != 0:
                 comfy.model_management.discard_cuda_async_error()
                 del pin
                 hostbuf.truncate(offset, do_unregister=False)
-                return _steal_pin(module, stack, buckets, size, priority)
+                return _steal_pin(module, stack, buckets, size, priority, subset)
     except RuntimeError:
         if extended:
             hostbuf.truncate(offset, do_unregister=False)
-        return _steal_pin(module, stack, buckets, size, priority)
+        return _steal_pin(module, stack, buckets, size, priority, subset)
 
-    module._pin = pin
+    module_pin["pin"] = pin
     stack.append((module, offset))
-    module._pin_registered = True
-    module._pin_stack_index = len(stack) - 1
-    stack_split[0] = max(stack_split[0], module._pin_stack_index)
+    module_pin["registered"] = True
+    module_pin["stack_index"] = len(stack) - 1
+    stack_split[0] = max(stack_split[0], module_pin["stack_index"])
     comfy.model_management.TOTAL_PINNED_MEMORY += size
     pinned_size[0] += size
-    _add_to_bucket(module, buckets, size, priority)
+    _add_to_bucket(module, module_pin, buckets, size, priority)
     return True

comfy/quant_ops.py

@@ -3,6 +3,22 @@ import logging
 
 from comfy.cli_args import args
 
+
+def _rocm_kitchen_arch_supported():
+    """comfy-kitchen's INT8 Triton kernels compile tl.dot to matrix-core instructions.
+    RDNA3/3.5/4 (gfx11xx/gfx12xx) have WMMA and CDNA (gfx9xx) has MFMA; RDNA1/RDNA2
+    (gfx10xx) have neither, so the INT8 path hangs the GPU there. Gates the automatic
+    ROCm default so those cards stay on the eager fallback (an explicit
+    --enable-triton-backend still forces it on any arch)."""
+    try:
+        arch = torch.cuda.get_device_properties(torch.cuda.current_device()).gcnArchName.split(":")[0]
+    except Exception:
+        return False
+    if arch.startswith(("gfx11", "gfx12")):
+        return True
+    return arch in ("gfx908", "gfx90a", "gfx940", "gfx941", "gfx942", "gfx950")
+
+
 try:
     import comfy_kitchen as ck
     from comfy_kitchen.tensor import (
@@ -10,6 +26,7 @@ try:
         QuantizedLayout,
         TensorCoreFP8Layout as _CKFp8Layout,
         TensorCoreNVFP4Layout as _CKNvfp4Layout,
+        TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout,
         TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
         register_layout_op,
         register_layout_class,
@@ -24,10 +41,22 @@ try:
             ck.registry.disable("cuda")
             logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
 
-    if args.enable_triton_backend:
+    # On ROCm/AMD the CUDA backend is unavailable, so Triton is the only accelerated
+    # comfy-kitchen backend. Enable it by default there, but only on Triton >= 3.7 AND a
+    # matrix-core GPU (RDNA3+ WMMA gfx11xx/gfx12xx, CDNA MFMA gfx9xx). RDNA1/RDNA2
+    # (gfx10xx) have no WMMA -> the INT8 tl.dot path hangs the GPU, so they stay eager.
+    # older Triton lacks libdevice.rint on the HIP backend and hard-crashes the INT8 path.
+    if args.disable_triton_backend:
+        ck.registry.disable("triton")
+    elif args.enable_triton_backend: # or (torch.version.hip is not None and _rocm_kitchen_arch_supported()):
         try:
             import triton
-            logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__)
+            triton_version = tuple(int(v) for v in triton.__version__.split(".")[:2])
+            if args.enable_triton_backend or triton_version >= (3, 7):
+                logging.info("Found triton %s. Enabling comfy-kitchen triton backend.", triton.__version__)
+            else:
+                logging.info("Triton %s is too old for the ROCm INT8 path (needs >= 3.7); comfy-kitchen triton backend disabled.", triton.__version__)
+                ck.registry.disable("triton")
         except ImportError as e:
             logging.error(f"Failed to import triton, Error: {e}, the comfy-kitchen triton backend will not be available.")
             ck.registry.disable("triton")
@@ -51,6 +80,9 @@ except ImportError as e:
     class _CKTensorWiseINT8Layout:
         pass
 
+    class _CKTensorCoreConvRotW4A4Layout:
+        pass
+
     def register_layout_class(name, cls):
         pass
 
@@ -179,6 +211,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
 # Backward compatibility alias - default to E4M3
 TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
 TensorWiseINT8Layout = _CKTensorWiseINT8Layout
+TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout
 
 
 # ==============================================================================
@@ -190,6 +223,7 @@ register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout)
 register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout)
 register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout)
 register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
+register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout)
 if _CK_MXFP8_AVAILABLE:
     register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
 
@@ -227,6 +261,13 @@ QUANT_ALGOS["int8_tensorwise"] = {
     "quantize_input": False,
 }
 
+QUANT_ALGOS["convrot_w4a4"] = {
+    "storage_t": torch.int8,
+    "parameters": {"weight_scale"},
+    "comfy_tensor_layout": "TensorCoreConvRotW4A4Layout",
+    "quantize_input": False,
+}
+
 
 # ==============================================================================
 # Re-exports for backward compatibility
@@ -239,6 +280,7 @@ __all__ = [
     "TensorCoreFP8E4M3Layout",
     "TensorCoreFP8E5M2Layout",
     "TensorCoreNVFP4Layout",
+    "TensorCoreConvRotW4A4Layout",
     "TensorWiseINT8Layout",
     "QUANT_ALGOS",
     "register_layout_op",

comfy/samplers.py

@@ -20,6 +20,7 @@ import comfy.hooks
 import comfy.context_windows
 import comfy.multigpu
 import comfy.utils
+from comfy.logging import detail
 import scipy.stats
 import numpy
 
@@ -991,10 +992,15 @@ class KSAMPLER(Sampler):
 
         noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
 
-        k_callback = None
         total_steps = len(sigmas) - 1
-        if callback is not None:
-            k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
+        first_step = True
+        def k_callback(x):
+            nonlocal first_step
+            if first_step:
+                detail("First sampler step: model=%s sampler=%s step=%s total_steps=%s cfg=%s seed=%s sigma=%s sigma_hat=%s latent_shape=%s denoised_shape=%s", model_wrap.model_patcher.model.__class__.__name__, self.sampler_function.__name__, x["i"], total_steps, model_wrap.cfg, extra_args.get("seed"), x.get("sigma"), x.get("sigma_hat"), tuple(x["x"].shape), tuple(x["denoised"].shape))
+                first_step = False
+            if callback is not None:
+                callback(x["i"], x["denoised"], x["x"], total_steps)
 
         samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
         samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
@@ -1270,10 +1276,13 @@ class CFGGuider:
             return latent_image
 
         if latent_image.is_nested:
+            sampler_shapes = [tuple(x.shape) for x in latent_image.unbind()]
             latent_image, latent_shapes = comfy.utils.pack_latents(latent_image.unbind())
             noise, _ = comfy.utils.pack_latents(noise.unbind())
         else:
             latent_shapes = [latent_image.shape]
+            sampler_shapes = [tuple(latent_image.shape)]
+        detail("Sampler: model=%s latent_shapes=%s", self.model_patcher.model.__class__.__name__, sampler_shapes)
 
         if denoise_mask is not None:
             if denoise_mask.is_nested:

comfy/sd.py

@@ -16,6 +16,8 @@ import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
 import comfy.ldm.wan.vae2_2
 import comfy.ldm.hunyuan3d.vae
+import comfy.ldm.seedvr.vae
+import comfy.ldm.mage_flow.vae
 import comfy.ldm.triposplat.vae
 import comfy.ldm.ace.vae.music_dcae_pipeline
 import comfy.ldm.cogvideo.vae
@@ -59,6 +61,7 @@ import comfy.text_encoders.qwen_image
 import comfy.text_encoders.hunyuan_image
 import comfy.text_encoders.z_image
 import comfy.text_encoders.krea2
+import comfy.text_encoders.mage_flow
 import comfy.text_encoders.ideogram4
 import comfy.text_encoders.ovis
 import comfy.text_encoders.kandinsky5
@@ -75,6 +78,7 @@ import comfy.text_encoders.gemma4
 import comfy.text_encoders.cogvideo
 import comfy.text_encoders.sa3
 import comfy.text_encoders.gpt_oss
+import comfy.text_encoders.joyimage
 
 import comfy.model_patcher
 import comfy.lora
@@ -473,7 +477,8 @@ class CLIP:
 
 class VAE:
     def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
-        if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
+        is_seedvr2_vae = "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd
+        if not is_seedvr2_vae and 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
             sd = diffusers_convert.convert_vae_state_dict(sd)
 
         if model_management.is_amd():
@@ -500,6 +505,8 @@ class VAE:
         self.upscale_index_formula = None
         self.extra_1d_channel = None
         self.crop_input = True
+        self.handles_tiling = False
+        self.format_encoded = None
 
         self.audio_sample_rate = 44100
 
@@ -546,6 +553,33 @@ class VAE:
                 self.first_stage_model = StageC_coder()
                 self.downscale_ratio = 32
                 self.latent_channels = 16
+            elif "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd: # seedvr2
+                self.first_stage_model = comfy.ldm.seedvr.vae.VideoAutoencoderKLWrapper()
+                self.latent_channels = comfy.ldm.seedvr.vae.SEEDVR2_LATENT_CHANNELS
+                self.latent_dim = 3
+                self.disable_offload = True
+                self.memory_used_decode = lambda shape, dtype: self.first_stage_model.comfy_memory_used_decode(shape)
+                self.memory_used_encode = lambda shape, dtype: (max(shape[2], 5) * shape[3] * shape[4] * 64) * model_management.dtype_size(dtype)
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                self.handles_tiling = True
+                self.format_encoded = self.first_stage_model.comfy_format_encoded
+                self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
+                self.downscale_index_formula = (4, 8, 8)
+                self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
+                self.upscale_index_formula = (4, 8, 8)
+                self.process_input = lambda image: image * 2.0 - 1.0
+                self.crop_input = False
+            elif "student.dconv_encoder.proj_out.weight" in sd:  # Mage-VAE (one-step diffusion codec, Flux2-anchored 128ch/16x latents)
+                sd = comfy.utils.state_dict_prefix_replace(sd, {"student.dconv_encoder.": "dconv_encoder.", "pipeline.": "decoder_model."})
+                # Drop the unused Flux2-VAE anchor encoder carried in the checkpoint.
+                sd = {k: v for k, v in sd.items() if not k.startswith("decoder_model.y_embedder.encoder.") and not k.startswith("decoder_model.y_embedder.bottleneck.")}
+                self.first_stage_model = comfy.ldm.mage_flow.vae.MageVAE()
+                self.latent_channels = 128
+                self.downscale_ratio = 16
+                self.upscale_ratio = 16
+                self.working_dtypes = [torch.bfloat16, torch.float32]
+                self.memory_used_encode = lambda shape, dtype: (400 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
+                self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * 16 * 16) * model_management.dtype_size(dtype)
             elif "decoder.conv_in.weight" in sd:
                 if sd['decoder.conv_in.weight'].shape[1] == 64:
                     ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
@@ -1012,6 +1046,10 @@ class VAE:
         decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
         return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, index_formulas=self.upscale_index_formula, output_device=self.output_device))
 
+    def _decode_tiled_owned(self, samples, **kwargs):
+        out = self.first_stage_model.decode_tiled(samples.to(self.vae_dtype).to(self.device), **kwargs)
+        return self.process_output(out.to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True))
+
     def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
         steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
         steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
@@ -1048,6 +1086,25 @@ class VAE:
         encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
         return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.downscale_ratio, out_channels=self.latent_channels, downscale=True, index_formulas=self.downscale_index_formula, output_device=self.output_device)
 
+    def _encode_tiled_owned(self, pixel_samples, **kwargs):
+        x = self.process_input(pixel_samples).to(self.vae_dtype).to(self.device)
+        out = self.first_stage_model.encode_tiled(x, **kwargs)
+        return out.to(device=self.output_device, dtype=self.vae_output_dtype())
+
+    def _owned_tiled_args(self, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
+        args = {}
+        if tile_x is not None:
+            args["tile_x"] = tile_x
+        if tile_y is not None:
+            args["tile_y"] = tile_y
+        if overlap is not None:
+            args["overlap"] = overlap
+        if tile_t is not None:
+            args["tile_t"] = tile_t
+        if overlap_t is not None:
+            args["overlap_t"] = overlap_t
+        return args
+
     def decode(self, samples_in, vae_options={}):
         self.throw_exception_if_invalid()
         pixel_samples = None
@@ -1095,11 +1152,19 @@ class VAE:
                 if dims == 1 or self.extra_1d_channel is not None:
                     pixel_samples = self.decode_tiled_1d(samples_in)
                 elif dims == 2:
-                    pixel_samples = self.decode_tiled_(samples_in)
+                    if self.handles_tiling:
+                        tile = 256 // self.spacial_compression_decode()
+                        overlap = tile // 4
+                        pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
+                    else:
+                        pixel_samples = self.decode_tiled_(samples_in)
                 elif dims == 3:
                     tile = 256 // self.spacial_compression_decode()
                     overlap = tile // 4
-                    pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+                    if self.handles_tiling:
+                        pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
+                    else:
+                        pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
 
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
         return pixel_samples
@@ -1118,7 +1183,9 @@ class VAE:
             args["overlap"] = overlap
 
         with model_management.cuda_device_context(self.device):
-            if dims == 1 or self.extra_1d_channel is not None:
+            if self.handles_tiling and dims in (2, 3):
+                output = self._decode_tiled_owned(samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
+            elif dims == 1 or self.extra_1d_channel is not None:
                 args.pop("tile_y")
                 output = self.decode_tiled_1d(samples, **args)
             elif dims == 2:
@@ -1179,12 +1246,17 @@ class VAE:
                 if self.latent_dim == 3:
                     tile = 256
                     overlap = tile // 4
-                    samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+                    if self.handles_tiling:
+                        samples = self._encode_tiled_owned(pixel_samples, tile_x=tile, tile_y=tile, overlap=overlap)
+                    else:
+                        samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
                 elif self.latent_dim == 1 or self.extra_1d_channel is not None:
                     samples = self.encode_tiled_1d(pixel_samples)
                 else:
                     samples = self.encode_tiled_(pixel_samples)
 
+        if self.format_encoded is not None:
+            samples = self.format_encoded(samples)
         return samples
 
     def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
@@ -1192,7 +1264,7 @@ class VAE:
         pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
         dims = self.latent_dim
         pixel_samples = pixel_samples.movedim(-1, 1)
-        if dims == 3:
+        if dims == 3 and pixel_samples.ndim < 5:
             if not self.not_video:
                 pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
             else:
@@ -1216,21 +1288,27 @@ class VAE:
             elif dims == 2:
                 samples = self.encode_tiled_(pixel_samples, **args)
             elif dims == 3:
-                if tile_t is not None:
-                    tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
+                if self.handles_tiling:
+                    samples = self._encode_tiled_owned(pixel_samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
                 else:
-                    tile_t_latent = 9999
-                args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
+                    if tile_t is not None:
+                        tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
+                    else:
+                        tile_t_latent = 9999
+                    args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
 
-                if overlap_t is None:
-                    args["overlap"] = (1, overlap, overlap)
-                else:
-                    args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap)
-                maximum = pixel_samples.shape[2]
-                maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
+                    spatial_overlap = overlap if overlap is not None else 64
+                    if overlap_t is None:
+                        args["overlap"] = (1, spatial_overlap, spatial_overlap)
+                    else:
+                        args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), spatial_overlap, spatial_overlap)
+                    maximum = pixel_samples.shape[2]
+                    maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
 
-                samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
+                    samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
 
+        if self.format_encoded is not None:
+            samples = self.format_encoded(samples)
         return samples
 
     def get_sd(self):
@@ -1313,6 +1391,8 @@ class CLIPType(Enum):
     IDEOGRAM4 = 30
     BOOGU = 31
     KREA2 = 32
+    JOYIMAGE = 33
+    MAGE = 34
 
 
 
@@ -1368,6 +1448,7 @@ class TEModel(Enum):
     GPT_OSS_20B = 33
     QWEN3VL_4B = 34
     QWEN3VL_8B = 35
+    GEMMA_4_12B = 36
 
 
 def detect_te_model(sd):
@@ -1397,6 +1478,9 @@ def detect_te_model(sd):
     if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
         if 'model.layers.59.self_attn.q_norm.weight' in sd:
             return TEModel.GEMMA_4_31B
+        # Gemma4 12B Unified: 48 layers, encoder-free; global layers drop v_proj (attention_k_eq_v).
+        if 'model.layers.47.self_attn.q_norm.weight' in sd and 'model.layers.5.self_attn.v_proj.weight' not in sd:
+            return TEModel.GEMMA_4_12B
         if 'model.layers.41.self_attn.q_norm.weight' in sd and 'model.layers.47.self_attn.q_norm.weight' not in sd:
             return TEModel.GEMMA_4_E4B
         if 'model.layers.34.self_attn.q_norm.weight' in sd and 'model.layers.41.self_attn.q_norm.weight' not in sd:
@@ -1552,10 +1636,11 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.clip = comfy.text_encoders.sa3.SAT5GemmaModel
             clip_target.tokenizer = comfy.text_encoders.sa3.SAT5GemmaTokenizer
             tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
-        elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B):
+        elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B, TEModel.GEMMA_4_12B):
             variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
                        TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
-                       TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B}[te_model]
+                       TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B,
+                       TEModel.GEMMA_4_12B: comfy.text_encoders.gemma4.Gemma4_12B}[te_model]
             clip_target.clip = comfy.text_encoders.gemma4.gemma4_te(**llama_detect(clip_data), model_class=variant)
             clip_target.tokenizer = variant.tokenizer
             tokenizer_data["tokenizer_json"] = clip_data[0].get("tokenizer_json", None)
@@ -1642,6 +1727,14 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
                 clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
                 clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
                 clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
+            elif clip_type == CLIPType.MAGE and te_model == TEModel.QWEN3VL_4B:  # Mage-Flow: full Qwen3-VL-4B, last hidden state, Qwen-Image-style templates.
+                clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
+                clip_target.clip = comfy.text_encoders.mage_flow.te(**llama_detect(clip_data))
+                clip_target.tokenizer = comfy.text_encoders.mage_flow.MageFlowTokenizer
+            elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B:  # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx.
+                clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
+                clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data))
+                clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer
             elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2):  # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
                 klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
                 clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
@@ -1898,7 +1991,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
         manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
     else:
         manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if model_config.clip_vision_prefix is not None:
         if output_clipvision:
@@ -2039,7 +2132,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable
         manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
     else:
         manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if custom_operations is not None:
         model_config.custom_operations = custom_operations

comfy/supported_models.py

@@ -27,6 +27,8 @@ import comfy.text_encoders.z_image
 import comfy.text_encoders.ideogram4
 import comfy.text_encoders.boogu
 import comfy.text_encoders.krea2
+import comfy.text_encoders.mage_flow
+import comfy.text_encoders.joyimage
 import comfy.text_encoders.anima
 import comfy.text_encoders.ace15
 import comfy.text_encoders.longcat_image
@@ -1685,6 +1687,40 @@ class Chroma(supported_models_base.BASE):
         t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.pixart_te(**t5_detect))
 
+class SeedVR2(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "seedvr2"
+    }
+    unet_extra_config = {}
+    required_keys = {
+        "{}positive_conditioning",
+        "{}negative_conditioning",
+    }
+    latent_format = comfy.latent_formats.SeedVR2
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+    sampling_settings = {
+        "shift": 1.0,
+    }
+
+    def set_inference_dtype(self, dtype, manual_cast_dtype, device=None):
+        if (
+            dtype == torch.float16
+            and manual_cast_dtype is None
+            and comfy.model_management.should_use_bf16(device)
+        ):
+            manual_cast_dtype = torch.bfloat16
+        super().set_inference_dtype(dtype, manual_cast_dtype, device=device)
+
+    def get_model(self, state_dict, prefix="", device=None):
+        out = model_base.SeedVR2(self, device=device)
+        return out
+
+    def clip_target(self, state_dict={}):
+        return None
+
 class ChromaRadiance(Chroma):
     unet_config = {
         "image_model": "chroma_radiance",
@@ -1848,6 +1884,35 @@ class Krea2(supported_models_base.BASE):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.krea2.Krea2Tokenizer, comfy.text_encoders.krea2.te(**hunyuan_detect))
 
+class MageFlow(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "mage_flow",
+    }
+
+    sampling_settings = {
+        "multiplier": 1.0,
+        "shift": 6.0,
+    }
+
+    memory_usage_factor = 6.5
+
+    unet_extra_config = {}
+    latent_format = latent_formats.Flux2
+
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+
+    def get_model(self, state_dict, prefix="", device=None):
+        out = model_base.MageFlow(self, device=device)
+        return out
+
+    def clip_target(self, state_dict={}):
+        pref = self.text_encoder_key_prefix[0]
+        hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref))
+        return supported_models_base.ClipTarget(comfy.text_encoders.mage_flow.MageFlowTokenizer, comfy.text_encoders.mage_flow.te(**hunyuan_detect))
+
 class QwenImage(supported_models_base.BASE):
     unet_config = {
         "image_model": "qwen_image",
@@ -1877,6 +1942,38 @@ class QwenImage(supported_models_base.BASE):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
 
+class JoyImage(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "joyimage",
+    }
+
+    sampling_settings = {
+        "multiplier": 1000,
+        "shift": 1.5,
+    }
+
+    memory_usage_factor = 1.8
+
+    unet_extra_config = {
+        "theta": 10000,
+        "rope_dim_list": [16, 56, 56],
+    }
+
+    latent_format = latent_formats.Wan21
+
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+
+    def get_model(self, state_dict, prefix="", device=None):
+        return model_base.JoyImage(self, device=device)
+
+    def clip_target(self, state_dict={}):
+        pref = self.text_encoder_key_prefix[0]
+        qwen3vl_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl.transformer.".format(pref))
+        return supported_models_base.ClipTarget(comfy.text_encoders.joyimage.JoyImageTokenizer, comfy.text_encoders.joyimage.te(**qwen3vl_detect))
+
 class HunyuanImage21(HunyuanVideo):
     unet_config = {
         "image_model": "hunyuan_video",
@@ -2348,12 +2445,15 @@ models = [
     HiDream,
     HiDreamO1,
     Chroma,
+    SeedVR2,
     ChromaRadiance,
     ACEStep,
     ACEStep15,
     Omnigen2,
     Boogu,
+    MageFlow,
     QwenImage,
+    JoyImage,
     Ideogram4,
     Krea2,
     Flux2,

comfy/supported_models_base.py

@@ -54,13 +54,13 @@ class BASE:
     optimizations = {"fp8": False}
 
     @classmethod
-    def matches(s, unet_config, state_dict=None):
+    def matches(s, unet_config, state_dict=None, unet_key_prefix=""):
         for k in s.unet_config:
             if k not in unet_config or s.unet_config[k] != unet_config[k]:
                 return False
         if state_dict is not None:
             for k in s.required_keys:
-                if k not in state_dict:
+                if k.format(unet_key_prefix) not in state_dict:
                     return False
         return True
 
@@ -115,7 +115,7 @@ class BASE:
         replace_prefix = {"": self.vae_key_prefix[0]}
         return utils.state_dict_prefix_replace(state_dict, replace_prefix)
 
-    def set_inference_dtype(self, dtype, manual_cast_dtype):
+    def set_inference_dtype(self, dtype, manual_cast_dtype, device=None):
         self.unet_config['dtype'] = dtype
         self.manual_cast_dtype = manual_cast_dtype
 

comfy/text_encoders/gemma4.py

@@ -1,11 +1,15 @@
 import torch
 import torch.nn as nn
+import torchaudio.functional as AF
+import torchvision.transforms.functional as TVF
 import numpy as np
+from tokenizers import Tokenizer
 from dataclasses import dataclass
 import math
 
 from comfy import sd1_clip
 import comfy.model_management
+import comfy.ops
 from comfy.ldm.modules.attention import optimized_attention_for_device
 from comfy.rmsnorm import rms_norm
 from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, _make_scaled_embedding
@@ -21,6 +25,10 @@ GEMMA4_VISION_CONFIG = {"hidden_size": 768, "image_size": 896, "intermediate_siz
 GEMMA4_VISION_31B_CONFIG = {"hidden_size": 1152, "image_size": 896, "intermediate_size": 4304, "num_attention_heads": 16, "num_hidden_layers": 27, "patch_size": 16, "head_dim": 72, "rms_norm_eps": 1e-6, "position_embedding_size": 10240, "pooling_kernel_size": 3}
 GEMMA4_AUDIO_CONFIG = {"hidden_size": 1024, "num_hidden_layers": 12, "num_attention_heads": 8, "intermediate_size": 4096, "conv_kernel_size": 5, "attention_chunk_size": 12, "attention_context_left": 13, "attention_context_right": 0, "attention_logit_cap": 50.0, "output_proj_dims": 1536, "rms_norm_eps": 1e-6, "residual_weight": 0.5}
 
+# Encoder-free (gemma4_unified) multimodal embedders: raw patches/waveform projected directly into LM space.
+GEMMA4_UNIFIED_VISION_CONFIG = {"model_patch_size": 48, "patch_size": 16, "pooling_kernel_size": 3, "mm_embed_dim": 3840, "mm_posemb_size": 1120, "output_proj_dims": 3840, "rms_norm_eps": 1e-6}
+GEMMA4_UNIFIED_AUDIO_CONFIG = {"audio_samples_per_token": 640, "output_proj_dims": 640, "rms_norm_eps": 1e-6}
+
 @dataclass
 class Gemma4Config:
     vocab_size: int = 262144
@@ -35,6 +43,9 @@ class Gemma4Config:
     transformer_type: str = "gemma4"
     head_dim = 256
     global_head_dim = 512
+    num_global_key_value_heads = None
+    attention_k_eq_v = False
+    vision_bidirectional = False
     rms_norm_add = False
     mlp_activation = "gelu_pytorch_tanh"
     qkv_bias = False
@@ -51,6 +62,7 @@ class Gemma4Config:
     num_kv_shared_layers: int = 18
     use_double_wide_mlp: bool = False
     stop_tokens = [1, 50, 106]
+    suppress_tokens = []
     vision_config = GEMMA4_VISION_CONFIG
     audio_config = GEMMA4_AUDIO_CONFIG
     mm_tokens_per_image = 280
@@ -72,12 +84,30 @@ class Gemma4_31B_Config(Gemma4Config):
     num_hidden_layers: int = 60
     num_attention_heads: int = 32
     num_key_value_heads: int = 16
+    vision_bidirectional = True
     sliding_attention = [1024, 1024, 1024, 1024, 1024, False]
     hidden_size_per_layer_input: int = 0
     num_kv_shared_layers: int = 0
     audio_config = None
     vision_config = GEMMA4_VISION_31B_CONFIG
 
+@dataclass
+class Gemma4_12B_Config(Gemma4Config):
+    hidden_size: int = 3840
+    intermediate_size: int = 15360
+    num_hidden_layers: int = 48
+    num_attention_heads: int = 16
+    num_key_value_heads: int = 8
+    num_global_key_value_heads = 1
+    attention_k_eq_v = True
+    vision_bidirectional = True
+    sliding_attention = [1024, 1024, 1024, 1024, 1024, False]
+    hidden_size_per_layer_input: int = 0
+    num_kv_shared_layers: int = 0
+    audio_config = GEMMA4_UNIFIED_AUDIO_CONFIG
+    vision_config = GEMMA4_UNIFIED_VISION_CONFIG
+    suppress_tokens = [258883, 258882]
+
 
 # unfused RoPE as addcmul_ RoPE diverges from reference code
 def _apply_rotary_pos_emb(x, freqs_cis):
@@ -89,17 +119,18 @@ def _apply_rotary_pos_emb(x, freqs_cis):
     return out
 
 class Gemma4Attention(nn.Module):
-    def __init__(self, config, head_dim, device=None, dtype=None, ops=None):
+    def __init__(self, config, head_dim, num_kv_heads=None, k_eq_v=False, device=None, dtype=None, ops=None):
         super().__init__()
         self.num_heads = config.num_attention_heads
-        self.num_kv_heads = config.num_key_value_heads
+        self.num_kv_heads = num_kv_heads if num_kv_heads is not None else config.num_key_value_heads
         self.hidden_size = config.hidden_size
         self.head_dim = head_dim
         self.inner_size = self.num_heads * head_dim
 
         self.q_proj = ops.Linear(config.hidden_size, self.inner_size, bias=config.qkv_bias, device=device, dtype=dtype)
         self.k_proj = ops.Linear(config.hidden_size, self.num_kv_heads * head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
-        self.v_proj = ops.Linear(config.hidden_size, self.num_kv_heads * head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
+        # k_eq_v: V reuses the K projection (no separate v_proj weight)
+        self.v_proj = None if k_eq_v else ops.Linear(config.hidden_size, self.num_kv_heads * head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
         self.o_proj = ops.Linear(self.inner_size, config.hidden_size, bias=False, device=device, dtype=dtype)
 
         self.q_norm = None
@@ -133,7 +164,10 @@ class Gemma4Attention(nn.Module):
             shareable_kv = None
         else:
             xk = self.k_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
-            xv = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
+            if self.v_proj is not None:
+                xv = self.v_proj(hidden_states).view(batch_size, seq_length, self.num_kv_heads, self.head_dim)
+            else:
+                xv = xk  # k_eq_v: V is the raw K projection (before k_norm/RoPE)
             if self.k_norm is not None:
                 xk = self.k_norm(xk)
             xv = rms_norm(xv)
@@ -186,7 +220,10 @@ class TransformerBlockGemma4(nn.Module):
 
         head_dim = config.head_dim if self.sliding_attention else config.global_head_dim
 
-        self.self_attn = Gemma4Attention(config, head_dim=head_dim, device=device, dtype=dtype, ops=ops)
+        # k_eq_v only on global layers, which then use num_global_key_value_heads
+        k_eq_v = config.attention_k_eq_v and not self.sliding_attention
+        num_kv_heads = config.num_global_key_value_heads if k_eq_v else config.num_key_value_heads
+        self.self_attn = Gemma4Attention(config, head_dim=head_dim, num_kv_heads=num_kv_heads, k_eq_v=k_eq_v, device=device, dtype=dtype, ops=ops)
 
         num_kv_shared = config.num_kv_shared_layers
         first_kv_shared = config.num_hidden_layers - num_kv_shared
@@ -203,9 +240,9 @@ class TransformerBlockGemma4(nn.Module):
             self.per_layer_input_gate = ops.Linear(config.hidden_size, self.hidden_size_per_layer_input, bias=False, device=device, dtype=dtype)
             self.per_layer_projection = ops.Linear(self.hidden_size_per_layer_input, config.hidden_size, bias=False, device=device, dtype=dtype)
             self.post_per_layer_input_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, device=device, dtype=dtype)
-            self.register_buffer("layer_scalar", torch.ones(1, device=device, dtype=dtype))
-        else:
-            self.layer_scalar = None
+
+        # layer_scalar exists on every gemma4 variant, independent of per-layer input
+        self.register_buffer("layer_scalar", torch.empty(1, device=device, dtype=dtype))
 
     def forward(self, x, attention_mask=None, freqs_cis=None, past_key_value=None, per_layer_input=None, shared_kv=None):
         sliding_window = None
@@ -244,8 +281,7 @@ class TransformerBlockGemma4(nn.Module):
             x = self.post_per_layer_input_norm(x)
             x = residual + x
 
-        if self.layer_scalar is not None:
-            x = x * self.layer_scalar
+        x = x * comfy.ops.cast_to_input(self.layer_scalar, x)
 
         return x, present_key_value, shareable_kv
 
@@ -334,6 +370,19 @@ class Gemma4Transformer(nn.Module):
             causal_mask.masked_fill_(torch.ones_like(causal_mask, dtype=torch.bool).triu_(1), min_val)
             mask = mask + causal_mask if mask is not None else causal_mask
 
+            # Bidirectional attention within each image soft-token block (prefill only; text/audio stay causal).
+            if self.config.vision_bidirectional and past_len == 0 and embeds_info:
+                block_ids = torch.full((seq_len,), -1, dtype=torch.long, device=x.device)
+                group = 0
+                for info in embeds_info:
+                    if info.get("type") == "image":
+                        start = info["index"]
+                        block_ids[start:start + info["size"]] = group
+                        group += 1
+                if group > 0:
+                    same_block = (block_ids[:, None] == block_ids[None, :]) & (block_ids[:, None] >= 0)
+                    mask = mask.masked_fill(same_block, 0.0)
+
         # Per-layer inputs
         per_layer_inputs = None
         if self.hidden_size_per_layer_input:
@@ -354,8 +403,24 @@ class Gemma4Transformer(nn.Module):
         shared_global_kv = None   # KV from last non-shared global layer
 
         intermediate = None
+        all_intermediate = None
+        only_layers = None
+        if intermediate_output is not None:
+            if isinstance(intermediate_output, list):
+                all_intermediate = []
+                only_layers = {len(self.layers) + layer if layer < 0 else layer for layer in intermediate_output}
+            elif intermediate_output == "all":
+                all_intermediate = []
+                intermediate_output = None
+            elif intermediate_output < 0:
+                intermediate_output = len(self.layers) + intermediate_output
+
         next_key_values = []
         for i, layer in enumerate(self.layers):
+            if all_intermediate is not None:
+                if only_layers is None or (i in only_layers):
+                    all_intermediate.append(x.unsqueeze(1).clone())
+
             past_kv = past_key_values[i] if past_key_values is not None and len(past_key_values) > 0 else None
 
             layer_kwargs = {}
@@ -385,7 +450,18 @@ class Gemma4Transformer(nn.Module):
         if self.norm is not None:
             x = self.norm(x)
 
-        if len(next_key_values) > 0:
+        if all_intermediate is not None:
+            if only_layers is None or (len(self.layers) in only_layers):
+                all_intermediate.append(x.unsqueeze(1).clone())
+            if len(all_intermediate) > 0:
+                intermediate = torch.cat(all_intermediate, dim=1)
+
+        if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
+            intermediate = self.norm(intermediate)
+
+        # Only hand back the KV cache when caching was actually requested; SDClipModel reads
+        # outputs[2] as the pooled output.
+        if past_key_values is not None and len(next_key_values) > 0:
             return x, intermediate, next_key_values
         return x, intermediate
 
@@ -404,6 +480,8 @@ class Gemma4Base(BaseLlama, BaseGenerate, torch.nn.Module):
         cap = self.model.config.final_logit_softcapping
         if cap:
             logits = cap * torch.tanh(logits / cap)
+        if self.model.config.suppress_tokens:
+            logits[..., self.model.config.suppress_tokens] = torch.finfo(logits.dtype).min
         return logits
 
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
@@ -441,6 +519,28 @@ class Gemma4AudioMixin:
         return None, None
 
 
+class Gemma4UnifiedBase(Gemma4Base):
+    """Encoder-free multimodal Gemma4 (gemma4_unified, e.g. 12B): raw image patches and audio frames projected directly into LM space."""
+    def _init_model(self, config, dtype, device, operations):
+        self.num_layers = config.num_hidden_layers
+        self.model = Gemma4Transformer(config, device=device, dtype=dtype, ops=operations)
+        self.dtype = dtype
+        self.vision_model = Gemma4UnifiedVisionEmbedder(config.vision_config, device=device, dtype=dtype, ops=operations)
+        self.multi_modal_projector = Gemma4RMSNormProjector(config.vision_config["output_proj_dims"], config.hidden_size, dtype=dtype, device=device, ops=operations)
+        self.audio_projector = Gemma4RMSNormProjector(config.audio_config["output_proj_dims"], config.hidden_size, dtype=dtype, device=device, ops=operations)
+
+    def preprocess_embed(self, embed, device):
+        if embed["type"] == "image":
+            pixels = embed.pop("data").movedim(-1, 1).to(device, dtype=self.dtype)  # [B, H, W, C] -> [B, C, H, W], [0,1]
+            patches, positions = self.vision_model.patchify(pixels)
+            vision_out = self.vision_model(patches, positions)
+            return self.multi_modal_projector(vision_out), None
+        if embed["type"] == "audio":
+            audio = embed.pop("data").to(device, dtype=self.dtype)  # [1, T, audio_samples_per_token]
+            return self.audio_projector(audio), None
+        return None, None
+
+
 # Vision Encoder
 
 def _compute_vision_2d_rope(head_dim, pixel_position_ids, theta=100.0, device=None):
@@ -713,6 +813,73 @@ class Gemma4MultiModalProjector(Gemma4RMSNormProjector):
         super().__init__(config.vision_config["hidden_size"], config.hidden_size, dtype=dtype, device=device, ops=ops)
 
 
+# Encoder-free vision (gemma4_unified): raw merged pixel patches projected directly into LM space.
+
+def _patches_merge(patches, positions_xy, length):
+    patch_size = math.isqrt(patches.shape[-1] // 3)
+    k = math.isqrt(patches.shape[-2] // length)
+    batch = patches.shape[:-2]
+
+    max_x = positions_xy[..., 0].max(dim=-1, keepdim=True)[0] + 1
+    kidx = torch.div(positions_xy, k, rounding_mode="floor")
+    rem = torch.remainder(positions_xy, k)
+    order = rem[..., 0] + rem[..., 1] * k + k * k * kidx[..., 0] + k * max_x * kidx[..., 1]
+    perm = order.long().argsort(dim=-1)
+
+    merged = patches.gather(-2, perm.unsqueeze(-1).expand_as(patches))
+    merged = merged.reshape(*batch, length, k, k, patch_size, patch_size, 3)
+    merged = merged.permute(*range(len(batch)), -6, -5, -3, -4, -2, -1).reshape(*batch, length, (k * patch_size) ** 2 * 3)
+
+    pos = positions_xy.gather(-2, perm.unsqueeze(-1).expand_as(positions_xy))
+    pad = (positions_xy == -1).all(dim=-1, keepdim=True)
+    pos = torch.where(pad, positions_xy, pos).reshape(*batch, length, k * k, 2)
+    pos = torch.div(pos, k, rounding_mode="floor").min(dim=-2)[0]
+    return merged, pos
+
+
+class Gemma4UnifiedVisionEmbedder(nn.Module):
+    """Encoder-free patch embedder (LN -> Dense -> LN -> +2D posemb -> LN); projection to text space is the separate multi_modal_projector."""
+    def __init__(self, config, device=None, dtype=None, ops=None):
+        super().__init__()
+        self.patch_size = config["patch_size"]
+        self.pooling_kernel_size = config["pooling_kernel_size"]
+        patch_dim = config["model_patch_size"] ** 2 * 3
+        mm_embed_dim = config["mm_embed_dim"]
+        self.patch_ln1 = ops.LayerNorm(patch_dim, device=device, dtype=dtype)
+        self.patch_dense = ops.Linear(patch_dim, mm_embed_dim, device=device, dtype=dtype)
+        self.patch_ln2 = ops.LayerNorm(mm_embed_dim, device=device, dtype=dtype)
+        self.pos_embedding = nn.Parameter(torch.empty(config["mm_posemb_size"], 2, mm_embed_dim, device=device, dtype=dtype))
+        self.pos_norm = ops.LayerNorm(mm_embed_dim, device=device, dtype=dtype)
+
+    def patchify(self, pixels):
+        """pixels: [B, C, H, W] in [0,1] -> merged patches [B, N, 6912], positions [B, N, 2]."""
+        ps, k = self.patch_size, self.pooling_kernel_size
+        out_patches, out_positions = [], []
+        for img in pixels:
+            ph, pw = img.shape[-2] // ps, img.shape[-1] // ps
+            teacher = img.reshape(img.shape[0], ph, ps, pw, ps).permute(1, 3, 2, 4, 0).reshape(ph * pw, -1)
+            grid = torch.meshgrid(torch.arange(pw, device=img.device), torch.arange(ph, device=img.device), indexing="xy")
+            tpos = torch.stack(grid, dim=-1).reshape(teacher.shape[0], 2)
+            n_model = teacher.shape[0] // (k * k)
+            mp, mpos = _patches_merge(teacher.unsqueeze(0), tpos.unsqueeze(0), n_model)
+            out_patches.append(mp.squeeze(0))
+            out_positions.append(mpos.squeeze(0))
+        return torch.stack(out_patches), torch.stack(out_positions)
+
+    def forward(self, pixel_values, image_position_ids):
+        x = self.patch_ln1(pixel_values)
+        x = self.patch_dense(x)
+        x = self.patch_ln2(x)
+
+        clamped = image_position_ids.clamp(min=0).long()
+        valid = (image_position_ids != -1).to(x.dtype).unsqueeze(-1)
+        axes = torch.arange(2, device=image_position_ids.device)
+        pos = comfy.model_management.cast_to_device(self.pos_embedding, x.device, x.dtype)
+        pos_embs = (pos[clamped, axes] * valid).sum(-2)
+        x = x + pos_embs
+        return self.pos_norm(x)
+
+
 # Audio Encoder
 
 class Gemma4AudioConvSubsampler(nn.Module):
@@ -990,6 +1157,30 @@ class Gemma4AudioProjector(Gemma4RMSNormProjector):
 
 # Tokenizer and Wrappers
 
+def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, pooling_kernel_size):
+    target_px = max_patches * patch_size ** 2
+    factor = math.sqrt(target_px / (height * width))
+    side_mult = pooling_kernel_size * patch_size
+    target_height = math.floor(factor * height / side_mult) * side_mult
+    target_width = math.floor(factor * width / side_mult) * side_mult
+
+    if target_height == 0 and target_width == 0:
+        raise ValueError(f"Attempting to resize to a 0 x 0 image. Resized height should be divisible by {side_mult}.")
+
+    max_side_length = (max_patches // pooling_kernel_size ** 2) * side_mult
+    if target_height == 0:
+        target_height = side_mult
+        target_width = min(math.floor(width / height) * side_mult, max_side_length)
+    elif target_width == 0:
+        target_width = side_mult
+        target_height = min(math.floor(height / width) * side_mult, max_side_length)
+
+    if target_height * target_width > target_px:
+        raise ValueError(f"Resizing [{height}x{width}] to [{target_height}x{target_width}] exceeds the patch budget.")
+
+    return target_height, target_width
+
+
 class Gemma4_Tokenizer():
     tokenizer_json_data = None
 
@@ -998,25 +1189,35 @@ class Gemma4_Tokenizer():
             return {"tokenizer_json": self.tokenizer_json_data}
         return {}
 
-    def _extract_mel_spectrogram(self, waveform, sample_rate):
-        """Extract 128-bin log mel spectrogram.
-        Uses numpy for FFT/matmul/log to produce bit-identical results with reference code.
-        """
-        # Mix to mono first, then resample to 16kHz
+    def _audio_token_count(self, num_samples):
+        # Default (E2B/E4B): mel frames after two stride-2 conv subsamples.
+        _fl = 320  # int(round(16000 * 20.0 / 1000.0))
+        _hl = 160  # int(round(16000 * 10.0 / 1000.0))
+        _nmel = (num_samples + _fl // 2 - (_fl + 1)) // _hl + 1
+        _t = _nmel
+        for _ in range(2):
+            _t = (_t + 2 - 3) // 2 + 1
+        return min(_t, 750)
+
+    @staticmethod
+    def _resample_16k(waveform, sample_rate):
+        """Mix to mono and resample to 16kHz. Kaiser params reproduce the reference (transformers
+        load_audio -> librosa/soxr_hq) to ~1e-12 MSE using only torchaudio."""
         if waveform.dim() > 1 and waveform.shape[0] > 1:
             waveform = waveform.mean(dim=0, keepdim=True)
         if waveform.dim() == 1:
             waveform = waveform.unsqueeze(0)
-        audio = waveform.squeeze(0).float().numpy()
+        audio = waveform.float()
         if sample_rate != 16000:
-            # Use scipy's resample_poly with a high-quality FIR filter to get as close as possible to librosa's resampling (while still not full match)
-            from scipy.signal import resample_poly, firwin
-            from math import gcd
-            g = gcd(sample_rate, 16000)
-            up, down = 16000 // g, sample_rate // g
-            L = max(up, down)
-            h = firwin(160 * L + 1, 0.96 / L, window=('kaiser', 6.5))
-            audio = resample_poly(audio, up, down, window=h).astype(np.float32)
+            audio = AF.resample(audio, sample_rate, 16000, resampling_method="sinc_interp_kaiser",
+                                lowpass_filter_width=121, rolloff=0.9568384289091556, beta=21.01531462440614)
+        return audio.squeeze(0).contiguous()
+
+    def _extract_audio_features(self, waveform, sample_rate):
+        """Default (E2B/E4B): 128-bin log mel spectrogram for the conformer audio encoder.
+        Uses numpy for FFT/matmul/log to produce bit-identical results with reference code.
+        """
+        audio = self._resample_16k(waveform, sample_rate).numpy()
         n = len(audio)
 
         # Pad to multiple of 128, build sample-level mask
@@ -1064,8 +1265,8 @@ class Gemma4_Tokenizer():
         if audio is not None:
             waveform = audio["waveform"].squeeze(0) if hasattr(audio, "__getitem__") else audio
             sample_rate = audio.get("sample_rate", 16000) if hasattr(audio, "get") else 16000
-            mel, mel_mask = self._extract_mel_spectrogram(waveform, sample_rate)
-            audio_features = [(mel.unsqueeze(0), mel_mask.unsqueeze(0))]  # ([1, T, 128], [1, T])
+            feat, feat_mask = self._extract_audio_features(waveform, sample_rate)
+            audio_features = [(feat.unsqueeze(0), feat_mask.unsqueeze(0))]  # ([1, T, D], [1, T])
 
         # Process image/video frames
         is_video = video is not None
@@ -1088,15 +1289,10 @@ class Gemma4_Tokenizer():
             h, w = samples.shape[2], samples.shape[3]
             patch_size = 16
             pooling_k = 3
-            max_soft_tokens = 70 if is_video else 280  # video uses smaller token budget per frame
+            max_soft_tokens = kwargs.get("max_soft_tokens", 70 if is_video else 280)
             max_patches = max_soft_tokens * pooling_k * pooling_k
-            target_px = max_patches * patch_size * patch_size
-            factor = (target_px / (h * w)) ** 0.5
-            side_mult = pooling_k * patch_size
-            target_h = max(int(factor * h // side_mult) * side_mult, side_mult)
-            target_w = max(int(factor * w // side_mult) * side_mult, side_mult)
+            target_h, target_w = _get_aspect_ratio_preserving_size(h, w, patch_size, max_patches, pooling_k)
 
-            import torchvision.transforms.functional as TVF
             for i in range(num_frames):
                 # rescaling to match reference code
                 s = (samples[i].clamp(0, 1) * 255).to(torch.uint8)  # [C, H, W] uint8
@@ -1115,7 +1311,7 @@ class Gemma4_Tokenizer():
                 llama_text = llama_template.format(text)
             else:
                 # Build template from modalities present
-                system = "<|turn>system\n<|think|><turn|>\n" if thinking else ""
+                system = "<|turn>system\n<|think|>\n<turn|>\n" if thinking else ""
                 media = ""
                 if len(images) > 0:
                     if is_video:
@@ -1135,15 +1331,11 @@ class Gemma4_Tokenizer():
                 if len(audio_features) > 0:
                     # Compute audio token count (always at 16kHz)
                     num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
-                    _fl = 320  # int(round(16000 * 20.0 / 1000.0))
-                    _hl = 160  # int(round(16000 * 10.0 / 1000.0))
-                    _nmel = (num_samples + _fl // 2 - (_fl + 1)) // _hl + 1
-                    _t = _nmel
-                    for _ in range(2):
-                        _t = (_t + 2 - 3) // 2 + 1
-                    n_audio_tokens = min(_t, 750)
+                    n_audio_tokens = self._audio_token_count(num_samples)
                     media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
-                llama_text = f"{system}<|turn>user\n{media}{text}<turn|>\n<|turn>model\n"
+                # Non-thinking mode primes an empty thought channel so the model answers directly.
+                model_open = "" if thinking else "<|channel>thought\n<channel|>"
+                llama_text = f"{system}<|turn>user\n{text}{media}<turn|>\n<|turn>model\n{model_open}"
 
         text_tokens = super().tokenize_with_weights(llama_text, return_word_ids)
 
@@ -1178,7 +1370,6 @@ class Gemma4_Tokenizer():
 class _Gemma4Tokenizer:
     """Tokenizer using the tokenizers (Gemma4 doesn't come with sentencepiece model)"""
     def __init__(self, tokenizer_json_bytes=None, **kwargs):
-        from tokenizers import Tokenizer
         if isinstance(tokenizer_json_bytes, torch.Tensor):
             tokenizer_json_bytes = bytes(tokenizer_json_bytes.tolist())
         self.tokenizer = Tokenizer.from_str(tokenizer_json_bytes.decode("utf-8"))
@@ -1224,6 +1415,30 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma4", tokenizer=self.tokenizer_class)
 
 
+class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
+    """Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
+    embedding_size = 3840
+
+    def _extract_audio_features(self, waveform, sample_rate):
+        audio = self._resample_16k(waveform, sample_rate)
+        spt = 640  # audio_samples_per_token (40ms at 16kHz)
+        pad = (-audio.shape[0]) % spt
+        if pad:
+            audio = torch.nn.functional.pad(audio, (0, pad))
+        num_tokens = audio.shape[0] // spt
+        feats = audio[:num_tokens * spt].reshape(num_tokens, spt)
+        feats = feats[:750]  # audio_seq_length cap (matches reference truncation, ~30s)
+        mask = torch.ones(feats.shape[0], dtype=torch.bool)
+        return feats, mask
+
+    def _audio_token_count(self, num_samples):
+        return min((num_samples + 639) // 640, 750)
+
+
+class Gemma4UnifiedTokenizer(Gemma4Tokenizer):
+    tokenizer_class = Gemma4UnifiedSDTokenizer
+
+
 # Model wrappers
 class Gemma4Model(sd1_clip.SDClipModel):
     model_class = None
@@ -1256,7 +1471,7 @@ class Gemma4Model(sd1_clip.SDClipModel):
                 expanded_idx += 1
         initial_token_ids = [ids]
         input_ids = torch.tensor(initial_token_ids, device=self.execution_device)
-        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, initial_tokens=initial_token_ids[0], presence_penalty=presence_penalty, initial_input_ids=input_ids)
+        return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, initial_tokens=initial_token_ids[0], presence_penalty=presence_penalty, initial_input_ids=input_ids, embeds_info=embeds_info)
 
 
 def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=None):
@@ -1296,3 +1511,11 @@ def _make_variant(config_cls):
 Gemma4_E4B = _make_variant(Gemma4Config)
 Gemma4_E2B = _make_variant(Gemma4_E2B_Config)
 Gemma4_31B = _make_variant(Gemma4_31B_Config)
+
+
+# Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).
+class Gemma4_12B(Gemma4UnifiedBase):
+    def __init__(self, config_dict, dtype, device, operations):
+        super().__init__()
+        self._init_model(Gemma4_12B_Config(**config_dict), dtype, device, operations)
+Gemma4_12B.tokenizer = Gemma4UnifiedTokenizer

comfy/text_encoders/joyimage.py

@@ -0,0 +1,97 @@
+import torch
+
+from comfy import sd1_clip
+import comfy.text_encoders.qwen_vl
+from comfy.text_encoders.qwen3vl import Qwen3VL, Qwen3VLTokenizer
+
+JOYIMAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
+JOYIMAGE_TEMPLATE_TEXT = (
+    "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
+    "quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
+    "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+)
+JOYIMAGE_TEMPLATE_IMAGE = (
+    "<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
+    "quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
+    f"<|im_start|>user\n{JOYIMAGE_VISION_BLOCK}{{}}<|im_end|>\n<|im_start|>assistant\n"
+)
+# The DiT was trained without the leading system-prompt tokens.
+JOYIMAGE_DROP_IDX = 34
+PAD_TOKEN = 151643
+
+
+class Qwen3VL8B_JoyImage(Qwen3VL):
+    model_type = "qwen3vl_8b"
+
+    def preprocess_embed(self, embed, device):
+        if embed["type"] == "image":
+            image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(
+                embed["data"], min_pixels=65536, max_pixels=16777216, patch_size=16,
+                image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5],
+                interpolation="bicubic",
+            )
+            merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)
+            return merged, {"grid": grid, "deepstack": deepstack}
+        return None, None
+
+
+class JoyImageTokenizer(Qwen3VLTokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        super().__init__(
+            embedding_directory=embedding_directory, tokenizer_data=tokenizer_data,
+            model_type="qwen3vl_8b",
+        )
+        self.llama_template = JOYIMAGE_TEMPLATE_TEXT
+        self.llama_template_images = JOYIMAGE_TEMPLATE_IMAGE
+
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=None, **kwargs):
+        kwargs.pop("thinking", None)
+        return super().tokenize_with_weights(
+            text, return_word_ids=return_word_ids, llama_template=llama_template,
+            images=images or [], thinking=True, **kwargs,
+        )
+
+
+class _JoyImageClipModel(sd1_clip.SDClipModel):
+    def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None,
+                 attention_mask=True, model_options={}):
+        super().__init__(
+            device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={},
+            # JoyImage conditions on the pre-final-norm output of the last decoder layer.
+            dtype=dtype, special_tokens={"pad": PAD_TOKEN}, layer_norm_hidden_state=False,
+            model_class=Qwen3VL8B_JoyImage, enable_attention_masks=attention_mask,
+            return_attention_masks=attention_mask, model_options=model_options,
+        )
+
+
+class JoyImageTEModel(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}):
+        super().__init__(
+            device=device, dtype=dtype, name="qwen3vl_8b",
+            clip_model=_JoyImageClipModel, model_options=model_options,
+        )
+
+    def encode_token_weights(self, token_weight_pairs):
+        out, pooled, extra = super().encode_token_weights(token_weight_pairs)
+        if out.shape[1] <= JOYIMAGE_DROP_IDX:
+            raise ValueError(
+                f"JoyImageTEModel: encoded sequence length {out.shape[1]} is shorter "
+                f"than drop_idx={JOYIMAGE_DROP_IDX}; the prompt did not include the "
+                f"template prefix."
+            )
+        out = out[:, JOYIMAGE_DROP_IDX:]
+        if "attention_mask" in extra:
+            extra["attention_mask"] = extra["attention_mask"][:, JOYIMAGE_DROP_IDX:]
+        return out, pooled, extra
+
+
+def te(dtype_llama=None, llama_quantization_metadata=None):
+    class JoyImageTEModel_(JoyImageTEModel):
+        def __init__(self, device="cpu", dtype=None, model_options={}):
+            if llama_quantization_metadata is not None:
+                model_options = model_options.copy()
+                model_options["quantization_metadata"] = llama_quantization_metadata
+            if dtype_llama is not None:
+                dtype = dtype_llama
+            super().__init__(device=device, dtype=dtype, model_options=model_options)
+    return JoyImageTEModel_

comfy/text_encoders/llama.py

@@ -876,7 +876,7 @@ class BaseGenerate:
                                     torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype), 0))
         return past_key_values
 
-    def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None, position_ids=None, deepstack_embeds=None, visual_pos_masks=None):
+    def generate(self, embeds=None, do_sample=True, max_length=256, temperature=1.0, top_k=50, top_p=0.9, min_p=0.0, repetition_penalty=1.0, seed=42, stop_tokens=None, initial_tokens=[], execution_dtype=None, min_tokens=0, presence_penalty=0.0, initial_input_ids=None, position_ids=None, deepstack_embeds=None, visual_pos_masks=None, embeds_info=None):
         device = embeds.device
 
         if stop_tokens is None:
@@ -911,7 +911,7 @@ class BaseGenerate:
             if step == 0 and deepstack_embeds is not None:
                 extra["deepstack_embeds"] = deepstack_embeds
                 extra["visual_pos_masks"] = visual_pos_masks
-            x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids, position_ids=position_ids, **extra)
+            x, _, past_key_values = self.model.forward(None, embeds=embeds, attention_mask=None, past_key_values=past_key_values, input_ids=current_input_ids, position_ids=position_ids, **extra, embeds_info=(embeds_info if step == 0 else None))
             logits = self.logits(x)[:, -1]
             next_token = self.sample_token(logits, temperature, top_k, top_p, min_p, repetition_penalty, initial_tokens + generated_token_ids, generator, do_sample=do_sample, presence_penalty=presence_penalty)
             token_id = next_token[0].item()

comfy/text_encoders/mage_flow.py

@@ -0,0 +1,94 @@
+"""Mage-Flow text encoder: Qwen3-VL-4B, last hidden state (2560-dim).
+
+Mage-Flow conditions on the final hidden state of Qwen3-VL-4B with the leading
+system + user-opening template tokens stripped (reference start_idx 34 for t2i,
+64 for edit). The t2i template is identical to Qwen-Image's; the edit template
+uses the same system prompt as Qwen-Image-Edit with "Image N: " reference
+prefixes and no <think> block.
+"""
+
+import numbers
+
+import torch
+
+import comfy.text_encoders.qwen3vl
+from comfy import sd1_clip
+
+MAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
+
+MAGE_T2I_TEMPLATE = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+MAGE_EDIT_TEMPLATE = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
+
+
+class MageFlowTokenizer(comfy.text_encoders.qwen3vl.Qwen3VLTokenizer):
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type="qwen3vl_4b")
+        self.llama_template = MAGE_T2I_TEMPLATE
+        self.llama_template_images = MAGE_EDIT_TEMPLATE
+
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=True, **kwargs):
+        image = kwargs.get("image", None)
+        if image is not None and len(images) == 0:
+            images = [image[i:i + 1] for i in range(image.shape[0])]
+        if llama_template is None:
+            if len(images) > 0:
+                # Training-time multi-reference body: "Image 1: <ph>Image 2: <ph>...{instruction}"
+                prefix = "".join("Image {}: {}".format(j + 1, MAGE_VISION_BLOCK) for j in range(len(images)))
+                llama_template = self.llama_template_images.replace("{}", prefix + "{}", 1)
+            else:
+                llama_template = self.llama_template
+        # thinking=True: Mage templates end at "<|im_start|>assistant\n" with no <think> block.
+        return super().tokenize_with_weights(text, return_word_ids=return_word_ids, llama_template=llama_template, images=images, prevent_empty_text=prevent_empty_text, thinking=thinking, **kwargs)
+
+
+class MageFlowQwen3VLClipModel(comfy.text_encoders.qwen3vl.Qwen3VLClipModel):
+    def __init__(self, device="cpu", dtype=None, attention_mask=True, model_options={}, model_type="qwen3vl_4b"):
+        super().__init__(device=device, dtype=dtype, attention_mask=attention_mask, model_options=model_options, model_type=model_type)
+        # apply the final RMSNorm to the tapped last layer (HF last_hidden_state)
+        self.layer_norm_hidden_state = True
+
+
+class MageFlowTEModel(sd1_clip.SD1ClipModel):
+    def __init__(self, device="cpu", dtype=None, model_options={}):
+        clip_model = lambda **kw: MageFlowQwen3VLClipModel(**kw, model_type="qwen3vl_4b")  # noqa: E731
+        super().__init__(device=device, dtype=dtype, name="qwen3vl_4b", clip_model=clip_model, model_options=model_options)
+
+    def encode_token_weights(self, token_weight_pairs, template_end=-1):
+        # Strip the system + user-opening prefix (reference drop_idx: 34 t2i / 64 edit).
+        out, pooled, extra = super().encode_token_weights(token_weight_pairs)
+        tok_pairs = token_weight_pairs["qwen3vl_4b"][0]
+        count_im_start = 0
+        if template_end == -1:
+            for i, v in enumerate(tok_pairs):
+                elem = v[0]
+                if not torch.is_tensor(elem):
+                    if isinstance(elem, numbers.Integral):
+                        if elem == 151644 and count_im_start < 2:  # <|im_start|>
+                            template_end = i
+                            count_im_start += 1
+
+            if out.shape[1] > (template_end + 3):
+                if tok_pairs[template_end + 1][0] == 872:      # "user"
+                    if tok_pairs[template_end + 2][0] == 198:  # "\n"
+                        template_end += 3
+
+        out = out[:, template_end:]
+
+        if "attention_mask" in extra:
+            extra["attention_mask"] = extra["attention_mask"][:, template_end:]
+            if extra["attention_mask"].sum() == torch.numel(extra["attention_mask"]):
+                extra.pop("attention_mask")  # attention mask is useless if no masked elements
+
+        return out, pooled, extra
+
+
+def te(dtype_llama=None, llama_quantization_metadata=None):
+    class MageFlowTEModel_(MageFlowTEModel):
+        def __init__(self, device="cpu", dtype=None, model_options={}):
+            if dtype_llama is not None:
+                dtype = dtype_llama
+            if llama_quantization_metadata is not None:
+                model_options = model_options.copy()
+                model_options["quantization_metadata"] = llama_quantization_metadata
+            super().__init__(device=device, dtype=dtype, model_options=model_options)
+    return MageFlowTEModel_

comfy/text_encoders/qwen3vl.py

@@ -90,6 +90,27 @@ class Qwen3VL(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
                 deepstack = [torch.cat([deepstack[i], ds[i]], dim=0) for i in range(len(ds))]
         return position_ids, visual_pos_masks, deepstack
 
+    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[], **kwargs):
+        position_ids = kwargs.pop("position_ids", None)
+        visual_pos_masks = kwargs.pop("visual_pos_masks", None)
+        deepstack_embeds = kwargs.pop("deepstack_embeds", None)
+        if embeds is not None and position_ids is None:
+            position_ids, visual_pos_masks, deepstack_embeds = self.build_image_inputs(embeds, embeds_info)
+        return self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            embeds=embeds,
+            num_tokens=num_tokens,
+            intermediate_output=intermediate_output,
+            final_layer_norm_intermediate=final_layer_norm_intermediate,
+            dtype=dtype,
+            position_ids=position_ids,
+            embeds_info=embeds_info,
+            visual_pos_masks=visual_pos_masks,
+            deepstack_embeds=deepstack_embeds,
+            **kwargs,
+        )
+
 
 def _make_qwen3vl_model(model_type):
     class Qwen3VL_(Qwen3VL):
@@ -137,12 +158,12 @@ class Qwen3VLTokenizer(sd1_clip.SD1Tokenizer):
         self.llama_template = "<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
         self.llama_template_images = "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
 
-    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=False, **kwargs):
+    def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=False, skip_template=False, **kwargs):
         image = kwargs.get("image", None)
         if image is not None and len(images) == 0:
             images = [image[i:i + 1] for i in range(image.shape[0])]
 
-        skip_template = text.startswith('<|im_start|>')
+        skip_template = skip_template or text.startswith('<|im_start|>')
         if prevent_empty_text and text == '':
             text = ' '
 

comfy/text_encoders/qwen_vl.py

@@ -15,6 +15,7 @@ def process_qwen2vl_images(
     merge_size: int = 2,
     image_mean: list = None,
     image_std: list = None,
+    interpolation: str = "bilinear",
 ):
     if image_mean is None:
         image_mean = [0.48145466, 0.4578275, 0.40821073]
@@ -47,10 +48,9 @@ def process_qwen2vl_images(
     img_resized = F.interpolate(
         img.unsqueeze(0),
         size=(h_bar, w_bar),
-        mode='bilinear',
+        mode=interpolation,
         align_corners=False
     ).squeeze(0)
-
     normalized = img_resized.clone()
     for c in range(3):
         normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]

comfy_api/latest/_input_impl/video_types.py

@@ -1,5 +1,6 @@
 from av.container import InputContainer
 from av.subtitles.stream import SubtitleStream
+from av.video.reformatter import ColorRange
 from fractions import Fraction
 from typing import Optional
 from .._input import AudioInput, VideoInput
@@ -9,6 +10,7 @@ import itertools
 import json
 import numpy as np
 import math
+import os
 import torch
 from .._util import VideoContainer, VideoCodec, VideoComponents
 import logging
@@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
     return max(component.bits for component in stream.format.components)
 
 
+def last_decodable_audio_stream(container: InputContainer):
+    """Streams FFmpeg has no decoder for have no codec context, and decoding their
+    packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
+    stream = next(
+        (s for s in reversed(container.streams.audio) if s.codec_context is not None),
+        None,
+    )
+    if stream is None and len(container.streams.audio):
+        logging.warning("No decodable audio stream found in video; ignoring audio.")
+    return stream
+
+
+def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
+    """Containers probed only up to a window (mpegts) leave audio codec parameters unset when
+    audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
+    Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
+    for i, packet in enumerate(container.demux(audio_stream)):
+        try:
+            frames = packet.decode()
+        except av.error.FFmpegError:
+            frames = ()
+        if frames:
+            return frames[0].sample_rate, frames[0].layout.nb_channels
+        if i >= max_packets:
+            break
+    return 0, 0
+
+
+def write_output_metadata(container: InputContainer, output, metadata: dict | None):
+    """Copy the source container's metadata, then overlay the caller's tags."""
+    for key, value in container.metadata.items():
+        if metadata is None or key not in metadata:
+            output.metadata[key] = value
+    if metadata is not None:
+        for key, value in metadata.items():
+            output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
+
+
+def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
+    if format != VideoContainer.AUTO and format != VideoContainer.MP4:
+        raise ValueError("Only MP4 format is supported for now")
+    if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
+        raise ValueError("Only H264 codec is supported for now")
+    open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
+    if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
+        open_kwargs["format"] = format.value
+    elif isinstance(path, io.BytesIO):
+        open_kwargs["format"] = "mp4"  # no file extension to infer the format from
+    return open_kwargs
+
+
 class VideoFromFile(VideoInput):
     """
     Class representing video input from a file.
@@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
                     return estimated_frames
 
             # 3. Last resort: decode frames and count them (streaming)
-            if self.__start_time < 0:
-                start_time = max(self._get_raw_duration() + self.__start_time, 0)
-            else:
-                start_time = self.__start_time
+            start_time, duration = self.get_active_trim_window()
             frame_count = 1
             start_pts = int(start_time / video_stream.time_base)
-            end_pts = int((start_time + self.__duration) / video_stream.time_base)
+            end_pts = int((start_time + duration) / video_stream.time_base)
             container.seek(start_pts, stream=video_stream)
             frame_iterator = (
                 container.decode(video_stream)
@@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
 
     def get_components_internal(self, container: InputContainer) -> VideoComponents:
         video_stream = self._get_first_video_stream(container)
-        if self.__start_time < 0:
-            start_time = max(self._get_raw_duration() + self.__start_time, 0)
-        else:
-            start_time = self.__start_time
+        start_time, duration = self.get_active_trim_window()
 
         # Get video frames
         frames = []
         audio_frames = []
         alphas = None
         start_pts = int(start_time / video_stream.time_base)
-        end_pts = int((start_time + self.__duration) / video_stream.time_base)
+        end_pts = int((start_time + duration) / video_stream.time_base)
 
         if start_pts != 0:
             container.seek(start_pts, stream=video_stream)
@@ -281,18 +328,11 @@ class VideoFromFile(VideoInput):
         video_done = False
         audio_done = True
 
-        # Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
-        # and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
-        audio_stream = next(
-            (s for s in reversed(container.streams.audio) if s.codec_context is not None),
-            None,
-        )
+        audio_stream = last_decodable_audio_stream(container)
         if audio_stream is not None:
             streams += [audio_stream]
             resampler = av.audio.resampler.AudioResampler(format='fltp')
             audio_done = False
-        elif len(container.streams.audio):
-            logging.warning("No decodable audio stream found in video; ignoring audio.")
 
         for packet in container.demux(*streams):
             if video_done and audio_done:
@@ -305,7 +345,7 @@ class VideoFromFile(VideoInput):
                     for frame in packet.decode():
                         if frame.pts < start_pts:
                             continue
-                        if self.__duration and frame.pts >= end_pts:
+                        if duration and frame.pts >= end_pts:
                             video_done = True
                             break
 
@@ -372,7 +412,7 @@ class VideoFromFile(VideoInput):
                     map(resampler.resample, packet.decode())
                 )
                 for frame in aframes:
-                    if self.__duration and frame.time > start_time + self.__duration:
+                    if duration and frame.time > start_time + duration:
                         audio_done = True
                         break
 
@@ -394,8 +434,8 @@ class VideoFromFile(VideoInput):
 
         if len(audio_frames) > 0:
             audio_data = np.concatenate(audio_frames, axis=1)  # shape: (channels, total_samples)
-            if self.__duration:
-                audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
+            if duration:
+                audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
 
             audio_tensor = torch.from_numpy(audio_data).unsqueeze(0)  # shape: (1, channels, total_samples)
             audio = AudioInput({
@@ -441,28 +481,14 @@ class VideoFromFile(VideoInput):
             if not reuse_streams:
                 if bit_depth is None:
                     bit_depth = source_bit_depth
-                components = self.get_components_internal(container)
-                video = VideoFromComponents(components)
-                return video.save_to(
-                    path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
-                )
+                return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
 
             streams = container.streams
 
             open_kwargs = get_open_write_kwargs(path, container_format, format)
             with av.open(path, **open_kwargs) as output_container:
-                # Copy over the original metadata
-                for key, value in container.metadata.items():
-                    if metadata is None or key not in metadata:
-                        output_container.metadata[key] = value
-
-                # Add our new metadata
-                if metadata is not None:
-                    for key, value in metadata.items():
-                        if isinstance(value, str):
-                            output_container.metadata[key] = value
-                        else:
-                            output_container.metadata[key] = json.dumps(value)
+                # Add metadata before writing any streams
+                write_output_metadata(container, output_container, metadata)
 
                 # Add streams to the new container. Streams with no codec context cannot be used as an output template.
                 stream_map = {}
@@ -480,6 +506,282 @@ class VideoFromFile(VideoInput):
                         packet.stream = stream_map[packet.stream]
                         output_container.mux(packet)
 
+    def _save_transcoded(
+        self,
+        container: InputContainer,
+        path: str | io.BytesIO,
+        format: VideoContainer,
+        codec: VideoCodec,
+        metadata: dict | None,
+        bit_depth: int,
+    ):
+        """Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
+        open_kwargs = mp4_output_open_kwargs(path, format, codec)
+        video_stream = self._get_first_video_stream(container)
+        start_time, duration = self.get_active_trim_window()
+        start_pts = int(start_time / video_stream.time_base)
+        end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
+        stream_end_pts = None
+        if video_stream.duration is not None:
+            stream_end_pts = (video_stream.start_time or 0) + video_stream.duration
+        output_end_pts = end_pts
+        if stream_end_pts is not None and (output_end_pts is None or stream_end_pts < output_end_pts):
+            output_end_pts = stream_end_pts
+        if start_pts != 0:
+            container.seek(start_pts, stream=video_stream)
+
+        audio_stream = last_decodable_audio_stream(container)
+        pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p"
+        rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
+
+        resampler = None
+        sample_rate = 0
+        audio_time_base = None
+        duration_cap = None
+        if audio_stream is not None:
+            sample_rate = audio_stream.codec_context.sample_rate
+            channels = audio_stream.codec_context.channels
+            if not sample_rate:
+                sample_rate, channels = probe_audio_params(container, audio_stream)
+                container.seek(start_pts, stream=video_stream)
+                if sample_rate:
+                    audio_stream.codec_context.flush_buffers()
+                else:
+                    logging.warning("Audio stream parameters could not be determined; ignoring audio.")
+                    audio_stream = None
+        if audio_stream is not None:
+            audio_time_base = Fraction(1, sample_rate)
+            layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
+            resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
+            if duration:
+                duration_cap = math.ceil(duration * sample_rate)
+
+        streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
+        pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
+        video_done = False
+        audio_done = audio_stream is None
+        video_pts_offset = None
+        last_video_pts = None
+        last_video_end = None
+        # rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise
+        video_frame_durations = {}
+        source_size = None
+        rotation_k = 0
+        rotation_filter = None
+        audio_started = False
+        samples_written = 0
+        pending_audio = []
+        # The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims),
+        # and never seeking back keeps webm/mkv leading audio intact.
+        output = None
+        out_video = None
+        out_audio = None
+
+        def audio_frame_from_ndarray(nd_planar):
+            frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout)
+            frame.sample_rate = sample_rate
+            return frame
+
+        def drain_audio(final=False):
+            # Audio may cover the pts span of the video written so far, capped by the requested duration
+            nonlocal samples_written, audio_done
+            if last_video_end is None:
+                cap = 0
+            else:
+                cap = math.ceil(last_video_end * video_stream.time_base * sample_rate)
+            if duration_cap is not None:
+                cap = min(cap, duration_cap)
+            while pending_audio and not audio_done:
+                frame = pending_audio[0]
+                if samples_written + frame.samples <= cap:
+                    frame.pts = samples_written
+                    frame.time_base = audio_time_base
+                    output.mux(out_audio.encode(frame))
+                    samples_written += frame.samples
+                    pending_audio.pop(0)
+                    continue
+                if final:
+                    keep = frame.to_ndarray()[..., :cap - samples_written]
+                    if keep.shape[-1] > 0:
+                        tail = audio_frame_from_ndarray(keep)
+                        tail.pts = samples_written
+                        tail.time_base = audio_time_base
+                        output.mux(out_audio.encode(tail))
+                        samples_written += keep.shape[-1]
+                    pending_audio.clear()
+                break
+            if duration_cap is not None and samples_written >= duration_cap:
+                audio_done = True
+            return cap
+
+        try:
+            for packet in container.demux(*streams):
+                if video_done and audio_done:
+                    break
+
+                if packet.stream == video_stream and not video_done:
+                    try:
+                        frames = packet.decode()
+                    except av.error.InvalidDataError:
+                        logging.info("pyav decode error")
+                        continue
+                    for frame in frames:
+                        if frame.pts is not None and frame.pts < start_pts:
+                            continue
+                        if end_pts is not None and frame.pts is not None and frame.pts >= end_pts:
+                            video_done = True
+                            if last_video_pts is not None:
+                                # the source continues past the window: hold the last kept frame to the window end
+                                end_offset = video_pts_offset if video_pts_offset is not None else start_pts
+                                last_video_end = max(last_video_end, end_pts - end_offset)
+                            break
+                        # the source's true display duration of this frame; average_rate is not a
+                        # frame duration (sparse/VFR sources), so it is only the fallback
+                        frame_duration = frame.duration if frame.duration else pts_step
+                        if end_pts is not None and frame.pts is not None:
+                            frame_duration = min(frame_duration, end_pts - frame.pts)
+                        if output is None:
+                            rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
+                            if rotation_k % 2:
+                                out_width, out_height = frame.height, frame.width
+                            else:
+                                out_width, out_height = frame.width, frame.height
+                            if out_width % 2 or out_height % 2:
+                                raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
+                            source_size = (frame.width, frame.height)
+                            output = av.open(path, **open_kwargs)
+                            # Add metadata before writing any streams
+                            write_output_metadata(container, output, metadata)
+                            out_video = output.add_stream("h264", rate=rate)
+                            # no B-frames: reordering makes mp4 sample durations follow decode order,
+                            # so irregular-VFR spans and trim windows land wrong
+                            out_video.codec_context.max_b_frames = 0
+                            out_video.width = out_width
+                            out_video.height = out_height
+                            out_video.pix_fmt = pix_fmt
+                            # source pts pass through (rebased to 0), so variable frame rate survives
+                            out_video.codec_context.time_base = video_stream.time_base
+                            if audio_stream is not None:
+                                out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
+                        if (frame.width, frame.height) != source_size:
+                            # encoding would silently rescale the new geometry into the old one
+                            raise ValueError(
+                                f"Video resolution changes mid-stream "
+                                f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode"
+                            )
+                        if rotation_k:
+                            if rotation_filter is None:
+                                g = av.filter.Graph()
+                                g_src = g.add_buffer(width=frame.width, height=frame.height,
+                                                     format=frame.format.name, time_base=video_stream.time_base)
+                                tail = g_src
+                                for filter_name, filter_args in {1: [("transpose", "cclock")],
+                                                                 2: [("hflip", None), ("vflip", None)],
+                                                                 3: [("transpose", "clock")]}[rotation_k]:
+                                    step = g.add(filter_name, filter_args)
+                                    tail.link_to(step)
+                                    tail = step
+                                g_sink = g.add("buffersink")
+                                tail.link_to(g_sink)
+                                g.configure()
+                                rotation_filter = (g_src, g_sink)
+                            rotation_filter[0].push(frame)
+                            frame = rotation_filter[1].pull()
+                        if frame.color_range == ColorRange.JPEG:
+                            # compress full-range sources (yuvj/MJPEG) to limited range
+                            frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
+                        else:
+                            frame = frame.reformat(format=pix_fmt)
+                        frame_output_end = None
+                        if frame.pts is not None:
+                            if video_pts_offset is None:
+                                video_pts_offset = frame.pts
+                            frame.pts -= video_pts_offset
+                            if output_end_pts is not None:
+                                frame_output_end = output_end_pts - video_pts_offset
+                                if frame.pts + frame_duration > frame_output_end:
+                                    clamped_pts = frame_output_end - frame_duration
+                                    if clamped_pts >= 0 and (last_video_pts is None or clamped_pts > last_video_pts):
+                                        frame.pts = min(frame.pts, clamped_pts)
+                                    elif frame.pts < frame_output_end:
+                                        frame_duration = frame_output_end - frame.pts
+                                    else:
+                                        continue
+                        if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts):
+                            # broken sources emit missing/backward timestamps mid-stream, which the
+                            # muxer rejects; nudge them forward by one nominal frame interval
+                            frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step
+                            if frame_output_end is not None and frame.pts + frame_duration > frame_output_end:
+                                if frame.pts >= frame_output_end:
+                                    continue
+                                frame_duration = frame_output_end - frame.pts
+                        last_video_pts = frame.pts
+                        last_video_end = frame.pts + frame_duration
+                        video_frame_durations[frame.pts] = frame_duration
+                        # the decoded pict_type would force x264's frame types (intra-only
+                        # sources like MJPEG/ProRes would come out all-keyframe)
+                        frame.pict_type = 0
+                        for out_packet in out_video.encode(frame):
+                            out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
+                            output.mux(out_packet)
+                        drain_audio()
+
+                elif packet.stream == audio_stream and not audio_done:
+                    for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
+                        frame_start = None
+                        if resampled.pts is not None:
+                            # passthrough frames keep the source stream's time base
+                            tb = resampled.time_base if resampled.time_base else audio_time_base
+                            frame_start = float(resampled.pts * tb)
+                            if duration and not audio_started and frame_start >= start_time + duration:
+                                audio_done = True
+                                break
+                        if not audio_started:
+                            if frame_start is None:
+                                frame_start = 0.0
+                            to_skip = max(0, int((start_time - frame_start) * sample_rate))
+                            if to_skip >= resampled.samples:
+                                continue
+                            audio_started = True
+                            if duration and frame_start > start_time:
+                                duration_cap = min(duration_cap, math.ceil((start_time + duration - frame_start) * sample_rate))
+                            if to_skip:
+                                pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
+                                continue
+                        pending_audio.append(resampled)
+                        if video_done:
+                            # the video window is complete so the cap is final, but containers
+                            # that interleave audio behind video (fragmented mp4) still owe most
+                            # of it: stop only once the demuxed audio covers the cap
+                            cap = drain_audio()
+                            if pending_audio or samples_written >= cap:
+                                drain_audio(final=True)
+                                audio_done = True
+                                break
+
+            if output is None:
+                raise ValueError(f"No decodable video frames found in file '{self.__file}'")
+            if out_audio is not None and not audio_done:
+                drain_audio(final=True)
+            window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
+            for out_packet in out_video.encode(None):
+                duration = video_frame_durations.pop(out_packet.pts, 0)
+                if out_packet.pts == last_video_pts:
+                    duration = max(duration, window_fill)
+                out_packet.duration = duration
+                output.mux(out_packet)
+            if out_audio is not None:
+                output.mux(out_audio.encode(None))
+        except BaseException:
+            if output is not None:
+                output.close()
+                if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
+                    os.remove(path)
+            raise
+        else:
+            if output is not None:
+                output.close()
+
     def _get_first_video_stream(self, container: InputContainer):
         if len(container.streams.video):
             return container.streams.video[0]
@@ -527,22 +829,12 @@ class VideoFromComponents(VideoInput):
         bit_depth: int | None = None,
     ):
         """Save the video to a file path or BytesIO buffer."""
-        if format != VideoContainer.AUTO and format != VideoContainer.MP4:
-            raise ValueError("Only MP4 format is supported for now")
-        if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
-            raise ValueError("Only H264 codec is supported for now")
+        open_kwargs = mp4_output_open_kwargs(path, format, codec)
         # None means "use the depth this video was created with" (CreateVideo's choice).
         if bit_depth is None:
             bit_depth = self.__bit_depth
         is_10bit = bit_depth >= 10
-        extra_kwargs = {}
-        if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
-            extra_kwargs["format"] = format.value
-        elif isinstance(path, io.BytesIO):
-            # BytesIO has no file extension, so av.open can't infer the format.
-            # Default to mp4 since that's the only supported format anyway.
-            extra_kwargs["format"] = "mp4"
-        with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
+        with av.open(path, **open_kwargs) as output:
             # Add metadata before writing any streams
             if metadata is not None:
                 for key, value in metadata.items():

comfy_api_nodes/apis/bytedance.py

@@ -17,6 +17,10 @@ class Seedream4Options(BaseModel):
     max_images: int = Field(15)
 
 
+class Seedream5OptimizePromptOptions(BaseModel):
+    thinking: Literal["auto", "enabled", "disabled"] = Field(...)
+
+
 class Seedream4TaskCreationRequest(BaseModel):
     model: str = Field(...)
     prompt: str = Field(...)
@@ -28,6 +32,7 @@ class Seedream4TaskCreationRequest(BaseModel):
     sequential_image_generation_options: Seedream4Options | None = Field(Seedream4Options(max_images=15))
     watermark: bool = Field(False)
     output_format: str | None = None
+    optimize_prompt_options: Seedream5OptimizePromptOptions | None = None
 
 
 class ImageTaskCreationResponse(BaseModel):

comfy_api_nodes/apis/gemini.py

@@ -1,6 +1,6 @@
 from datetime import date
 from enum import Enum
-from typing import Any
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field
 
@@ -242,3 +242,60 @@ class GeminiGenerateContentResponse(BaseModel):
     promptFeedback: GeminiPromptFeedback | None = Field(None)
     usageMetadata: GeminiUsageMetadata | None = Field(None)
     modelVersion: str | None = Field(None)
+
+
+class GeminiInteractionTextPart(BaseModel):
+    type: Literal["text"] = "text"
+    text: str = Field(...)
+
+
+class GeminiInteractionMediaPart(BaseModel):
+    type: str = Field(..., description="One of: image, video, audio, document.")
+    data: str | None = Field(None, description="Base64-encoded media bytes.")
+    uri: str | None = Field(None, description="URI of the media, as an alternative to inline data.")
+    mime_type: str | None = Field(None)
+
+
+class GeminiInteractionGenerationConfig(BaseModel):
+    temperature: float | None = Field(None, ge=0.0, le=2.0)
+    top_p: float | None = Field(None, ge=0.0, le=1.0)
+
+
+class GeminiInteractionRequest(BaseModel):
+    model: str = Field(...)
+    input: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = Field(...)
+    generation_config: GeminiInteractionGenerationConfig | None = Field(None)
+
+
+class GeminiInteractionModalityTokens(BaseModel):
+    modality: str | None = Field(None, description="One of: text, image, audio, video, document.")
+    tokens: int | None = Field(None)
+
+
+class GeminiInteractionUsage(BaseModel):
+    input_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
+    output_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
+    total_thought_tokens: int | None = Field(None)
+
+
+class GeminiInteractionContent(BaseModel):
+    type: str | None = Field(None)
+    text: str | None = Field(None)
+    data: str | None = Field(None)
+    uri: str | None = Field(None)
+    mime_type: str | None = Field(None)
+
+
+class GeminiInteractionStep(BaseModel):
+    type: str | None = Field(None)
+    content: list[GeminiInteractionContent] | None = Field(None)
+
+
+class GeminiInteraction(BaseModel):
+    id: str | None = Field(None)
+    status: str | None = Field(
+        None,
+        description="One of: in_progress, requires_action, completed, failed, cancelled, incomplete.",
+    )
+    steps: list[GeminiInteractionStep] | None = Field(None)
+    usage: GeminiInteractionUsage | None = Field(None)

comfy_api_nodes/apis/heygen.py

@@ -0,0 +1,452 @@
+# (label, avatar_id, avatar_type, supported engines)
+HEYGEN_AVATAR_LOOKS: list[tuple[str, str, str, tuple[str, ...]]] = [
+    (
+        "Annie Lounge Standing Side",
+        "Annie_Lounge_Standing_Side_public",
+        "studio_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Yara Modern Lecture Hall",
+        "fd6814ecc5e143cd899e615a80eaa2dc",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Brandon Business Sitting Front",
+        "Brandon_Business_Sitting_Front_public",
+        "studio_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Caroline Business Sitting Side",
+        "Caroline_Business_Sitting_Side_public",
+        "studio_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Ursula Lawyer Angle 4",
+        "f7173d2bb8584c00bfec6905c5e9a492",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Sofia Corporate Presenter 01 Angle 3",
+        "fe563971fd2d438e957372dac9e2be8c",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Seoyeon Health Nutrition Coach Angle 3",
+        "fe3c5d5028d941398d064b8fc64a2dea",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Sanne Fitness Coach Angle 4",
+        "d967f935a8bf4a0c8f0bccfd66c501d2",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    ("Sander", "f5cd7b94056f495ca0610602d64a9aa3", "photo_avatar", ("avatar_v", "avatar_iv", "avatar_iii")),
+    (
+        "Rupert Personal Development Coach Angle 4",
+        "f57b3e626adb4bc997b38f64884adce4",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Olivier Professor Angle 2",
+        "f6659bbb094b459c87c967edbb9ee481",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Obi Health Nutrition Coach Angle 5",
+        "f3dc2c38201d414382f506d2d8e8d029",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Matilda Modern Office Setting",
+        "fda889ac354a440da8dbecc410981273",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Mateo Traditional Law Office",
+        "ff172d6c499c4e47ba6fcc5de631e9fc",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Marlon Inviting Armchair Setting",
+        "f5a57db099ab462daa3e7c604a05dacc",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Margaret Professor Angle 1",
+        "fb472bc29ab04bcca576e3703978fecb",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Marek Therapy Coach Angle 3",
+        "e197768703f1463a93dc25ada1f421fb",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Maeve Warm, Professional Setting",
+        "faf66681d8cc48dc82c4283200b3e782",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    (
+        "Lorenzo Professor Angle 5",
+        "fc268dc244bb40d7a554663ce723dcf0",
+        "photo_avatar",
+        ("avatar_v", "avatar_iv", "avatar_iii"),
+    ),
+    ("Luca", "Luca_public", "studio_avatar", ("avatar_iii",)),
+    ("Bruce", "Bruce_public", "studio_avatar", ("avatar_iii",)),
+    ("Nico", "Nico_public", "studio_avatar", ("avatar_iii",)),
+    ("Lisa", "Lisa_public", "studio_avatar", ("avatar_iii",)),
+    ("Sophie", "Sophie_public", "studio_avatar", ("avatar_iii",)),
+    ("Aiko", "Aiko_public", "studio_avatar", ("avatar_iii",)),
+    ("Rebecca (portrait)", "Rebecca_public", "studio_avatar", ("avatar_iii",)),
+    ("Daphne in Grey blazer (portrait)", "Daphne_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Bryce in Black t-shirt", "Bryce_public_5", "studio_avatar", ("avatar_iii",)),
+    ("Diora in White shirt", "Diora_public_3", "studio_avatar", ("avatar_iii",)),
+    ("Freja in White blazer", "Freja_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Albert in Blue blazer", "Albert_public_2", "studio_avatar", ("avatar_iii",)),
+    ("Emery in Red blazer", "Emery_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Minho in Blue shirt", "Minho_public_6", "studio_avatar", ("avatar_iii",)),
+    ("Aditya in Brown blazer", "Aditya_public_4", "studio_avatar", ("avatar_iii",)),
+    ("Nadim in Blue blazer", "Nadim_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Iker in Black blazer", "Iker_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Nour in Black blazer", "Nour_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Saskia in Blue blazer", "Saskia_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Lucien in Blue blazer", "Lucien_public_1", "studio_avatar", ("avatar_iii",)),
+    ("Esmond in Blue suit", "Esmond_public_3", "studio_avatar", ("avatar_iii",)),
+    ("Jinwoo in Blue suit", "Jinwoo_public_5", "studio_avatar", ("avatar_iii",)),
+    ("Annelore in Red sweater (portrait)", "Annelore_public_3", "studio_avatar", ("avatar_iii",)),
+    ("Bastien in Blue shirt", "Bastien_public_4", "studio_avatar", ("avatar_iii",)),
+    ("Zosia in Khaki blazer", "Zosia_public_3", "studio_avatar", ("avatar_iii",)),
+    ("Tahlia in Dark blue suit", "Tahlia_public_4", "studio_avatar", ("avatar_iii",)),
+]
+HEYGEN_AVATAR_OPTIONS = [x[0] for x in HEYGEN_AVATAR_LOOKS]
+HEYGEN_AVATAR_MAP = {x[0]: (x[1], x[2], x[3]) for x in HEYGEN_AVATAR_LOOKS}
+
+# (label, voice_id) — Starfish-compatible voices for the TTS endpoint
+HEYGEN_VOICE_TTS: list[tuple[str, str]] = [
+    ("Chill Brian (English, male)", "d2f4f24783d04e22ab49ee8fdc3715e0"),
+    ("Zain (English, female)", "0047732240584155b1588455313e78ec"),
+    ("Narrator Mateo - Excited 🤩 (Spanish, male)", "0077225a877e457db4572ccaf245910b"),
+    ("Aria (English, female)", "007e1378fc454a9f976db570ba6164a7"),
+    ("Caryns (English, female)", "0082e70326864107823605db0d77f5e0"),
+    ("Klara (English, female)", "01209fdcd1c24a109c86dc24ee0f71c0"),
+    ("Bold Kasia - Excited 🤩 (Polish, female)", "015482a78b9a46ebae74bd0beb17765b"),
+    ("Shaun (English, male)", "01c42cddcfdc4665a57b8d89cba8ffc1"),
+    ("Senthil (English, male)", "01d674cfd32b4728a3fddd21b7e7d543"),
+    ("Cody (English, male)", "01f98ed43e6140349f47dbd37a416827"),
+    ("Saffron (English, female)", "0258bbc2cd8648cfa357adfb833f6d7b"),
+    ("Blanka - Lifelike (English, female)", "02880d1c6fd94b7799d91135581ed810"),
+    ("Rami (English, male)", "02d5366a90af4c7a87157808ff352e33"),
+    ("Rhodes (English, female)", "02dce0a169b3460084b6c914d18fb2c8"),
+    ("Michelle - Voice 1 (English, female)", "02X8sHnuxFpsq1caYWN0"),
+    ("Autumn - UGC 3 (English, female)", "03dca9ebfca441dba55fb14afa0791b7"),
+    ("Reassuring Rupert (English, male)", "03fcf8ecb0a94b6b94e9007edb7c35f8"),
+    ("Rose - UGC -2  (English, female)", "0495e14c2bd74eb3aeeef03583e0bce5"),
+    ("Derya - Lifelike - Broadcaster 🎙️ (English, female)", "04d0ae1d0af2489ca7d3bb402a39a890"),
+    ("Dynamic Derek (English, male)", "0516c2d857eb425c94e90b068241914e"),
+    ("Lotte (English, female)", "052fcfb83d1a4c2f8d0368c226fea4b9"),
+    ("Thanos - Broadcaster 🎙️ (English, male)", "054af44a167344d0af2722fdfef08d17"),
+    ("Marcia (English, female)", "05f19352e8f74b0392a8f411eba40de1"),
+    ("Camden (English, male)", "06468055edd4458aa131a1dfd813c1e9"),
+    ("Rumi (English, female)", "06672207805f41a9ad0af6797f8aa14b"),
+    ("Pippa (English, female)", "06b68c4dbb544935b9af984e80efa4fb"),
+    ("William Prescott - Broadcaster 🎙️ (English, male)", "06c816b952f14fa9b3a6c42aa151f731"),
+    ("Sammy (English, female)", "06e6facd99654b9dbb9308f67bf3a31c"),
+    ("Breezy Bagus (Indonesian, male)", "06e81a5d7c8b41818d3f0b38f7cf15a1"),
+    ("Ben (English, male)", "07ca39b243184dbcb82e7e0f0e524b21"),
+    ("Smooth Dev (English, male)", "07d2ba65847541feb97abc9b60181555"),
+    ("Daran inside booth (English, male)", "080d9383c0314056aef392892e009806"),
+    ("Peppy Stella (English, female)", "084760b4922a44599575c770070ec2d7"),
+    ("Silas (English, male)", "08f561403ec846dbbd8c691cc448f45a"),
+    ("Aditya (English, male)", "09c3d65e44e247dd8b78a97a903feb58"),
+    ("Christy (English, female)", "09d88c036bf449fa905900c08b235a37"),
+    ("Elio (English, male)", "0a0b38624ac64ec6afcd5842a977ca10"),
+    ("Luminous Laksh (Hindi, male)", "0adc547b76a5401c856274c379904eb7"),
+    ("Jeff (English, male)", "0add542e349f4ccaba6ecb3b7ced6034"),
+    ("Tahlia Brooks - Excited 🤩 (English, female)", "0b440d1ac2454d69a73302fc806522b1"),
+    ("Riya Mehta (Hindi, female)", "0b464b2f4e2249a4b5a05e60eaf41e7e"),
+    ("Ben Hart (English, male)", "0b47b5a637e944f9bfd49913999b344b"),
+    ("Skylar (English, female)", "0bbfbda5aa924a68a9d1da7b8496052a"),
+    ("Relaxed Reece (English, male)", "0c2151d538844c70a8b096de533f2828"),
+    ("Daniel (English, male)", "0c23804af39a4946ac6fda42bfff2738"),
+    ("Melani (English, female)", "0c54c6399ad64551a304e1a346677723"),
+    ("Clover (English, female)", "0ccb0bea067d4449ad367baeed7ea2e9"),
+    ("Pedro Lima - Serious 😐 (Portuguese, male)", "0d0e23e8170446e38b18a7380b2d30a8"),
+    ("Ana Carvalho (Portuguese, female)", "0d23c5b2f6004e909802a2e8bfcd52c2"),
+    ("Confident Connor - Excited 🤩 (English, male)", "0dd34c3eb79247238219eea35aeb58cd"),
+    ("Vibrant Victor (Spanish, male)", "1062976ea8bf42f4adc27c7e868b8fde"),
+    ("Young Olivier (French, male)", "1c5dc9a8f8cf4de0932f91d75f43a15d"),
+    ("Émile Noir (French, male)", "25a6a67280574d3da78e97b1935ebfc7"),
+    ("Steadfast Stefan (German, male)", "0eb85e6e8710473b82f7e88609ba3053"),
+    ("Deep Dieter (German, male)", "118949676b0a46629d1ad52981c3ef84"),
+    ("Serene Marco (Italian, male)", "72e922488a614041b5ab5f6ee07e3deb"),
+    ("Murmuring Matteo (Italian, male)", "755902b751654f30a6ef49e8bbcacfec"),
+    ("Gail in car (Multilingual, female)", "0214ac51f93e420f8711d568dcfbc50e"),
+    ("Daran outside walking (Multilingual, male)", "0ac81e725f4948dfa9638ceca216bcfa"),
+    ("BOB - Voice 1 (Chinese, unknown)", "dMkR1XwIkarpNqWUJLnX"),
+    ("Hakeem Hassan (Arabic, male)", "61a4359785664d01a59664ceb87ce6d4"),
+    ("Rami Idris (Arabic, male)", "a0bd2e5d41a74643be47ac75ca9171a2"),
+    ("Bold Kasia - Friendly 😊 (Polish, female)", "331624aec8b24a6c9287b8e16bdf54e8"),
+    ("Tranquil Tulin (Turkish, female)", "61646c861eb64e2d9036d8db51385356"),
+    ("Dynamic Derya (Turkish, female)", "664b73058b784aa89ddb2924c141d441"),
+    ("Quiet Dewa (Indonesian, male)", "1fa1193cf1d74f27ba58531c07ef9862"),
+    ("Cuong (Vietnamese, male)", "8af68d7ea38f4e7ca05cf46c3f7a590b"),
+]
+HEYGEN_VOICE_TTS_OPTIONS = [x[0] for x in HEYGEN_VOICE_TTS]
+HEYGEN_VOICE_TTS_MAP = dict(HEYGEN_VOICE_TTS)
+
+# (label, voice_id) — top-ranked voices for video narration (any engine)
+HEYGEN_VOICE_GENERAL: list[tuple[str, str]] = [
+    ("Cassidy (English, female)", "16a09e4706f74997ba4ed05ea11470f6"),
+    ("Hope (English, female)", "42d00d4aac5441279d8536cd6b52c53c"),
+    ("Archer (English, male)", "453c20e1525a429080e2ad9e4b26f2cd"),
+    ("Brittney (English, female)", "4754e1ec667544b0bd18cdf4bec7d6a7"),
+    ("Mark (English, male)", "5d8c378ba8c3434586081a52ac368738"),
+    ("Andrew (English, male)", "6be73833ef9a4eb0aeee399b8fe9d62b"),
+    ("Spuds Oxley (English, male)", "76940a9adcd0490a9ce2cfe9a64a2664"),
+    ("Patrick (English, male)", "7e157ec62c9c45f1adca12faae72c86f"),
+    ("David Castlemore (English, male)", "828b59f834fd4c7188da322b6d9b6c75"),
+    ("Michael C (English, male)", "8661cd40d6c44c709e2d0031c0186ada"),
+    ("Adam Stone (English, male)", "88bb9ee1c81b466eb2a08fdde86d3619"),
+    ("Alex (English, male)", "897d6a9b2c844f56aa077238768fe10a"),
+    ("Monika Sogam (English, female)", "97dd67ab8ce242b6a9e7689cb00c6414"),
+    ("Jessica Anne Bogart (English, female)", "b966c31caf124c2a99f19ff1479c964f"),
+    ("John Doe (English, male)", "c4a8ceb7a2954500bc047fb092bcff3f"),
+    ("Ivy (English, female)", "cef3bc4e0a84424cafcde6f2cf466c97"),
+    ("Chill Brian (English, male)", "d2f4f24783d04e22ab49ee8fdc3715e0"),
+    ("Allison (English, female)", "f8c69e517f424cafaecde32dde57096b"),
+    ("Mia Starset (Norwegian, female)", "000466f8ac6d47a49f5743d50b3778de"),
+    ("William Shanks (Spanish, male)", "001248bb63f847888d37b766ee8b3a47"),
+    ("Zain (English, female)", "0047732240584155b1588455313e78ec"),
+    ("Jora Slobod (Romanian, male)", "00631519159a402ab5d8f719e51532bb"),
+    ("Narrator Mateo - Excited 🤩 (Spanish, male)", "0077225a877e457db4572ccaf245910b"),
+    ("Aria (English, female)", "007e1378fc454a9f976db570ba6164a7"),
+    ("Caryns (English, female)", "0082e70326864107823605db0d77f5e0"),
+    ("Klara (English, female)", "01209fdcd1c24a109c86dc24ee0f71c0"),
+    ("Son Tran (Vietnamese, male)", "0132f85950a94d11ba180f885101bf84"),
+    ("Bold Kasia - Excited 🤩 (Polish, female)", "015482a78b9a46ebae74bd0beb17765b"),
+    ("Marc Aurèle (French, male)", "018a94cf15574491a0bab7f6799ac15b"),
+    ("Shaun (English, male)", "01c42cddcfdc4665a57b8d89cba8ffc1"),
+    ("Senthil (English, male)", "01d674cfd32b4728a3fddd21b7e7d543"),
+    ("Cody (English, male)", "01f98ed43e6140349f47dbd37a416827"),
+    ("Saffron (English, female)", "0258bbc2cd8648cfa357adfb833f6d7b"),
+    ("Blanka - Lifelike (English, female)", "02880d1c6fd94b7799d91135581ed810"),
+    ("Rami (English, male)", "02d5366a90af4c7a87157808ff352e33"),
+    ("Rhodes (English, female)", "02dce0a169b3460084b6c914d18fb2c8"),
+    ("Michelle - Voice 1 (English, female)", "02X8sHnuxFpsq1caYWN0"),
+    ("Tuba (, female)", "034ca0c32b6542028748d6d365d90d6a"),
+    ("Autumn - UGC 3 (English, female)", "03dca9ebfca441dba55fb14afa0791b7"),
+    ("Reassuring Rupert (English, male)", "03fcf8ecb0a94b6b94e9007edb7c35f8"),
+]
+HEYGEN_VOICE_GENERAL_OPTIONS = [x[0] for x in HEYGEN_VOICE_GENERAL]
+HEYGEN_VOICE_GENERAL_MAP = dict(HEYGEN_VOICE_GENERAL)
+
+HEYGEN_TRANSLATE_LANGUAGES = [
+    "English",
+    "Spanish",
+    "Spanish (Spain)",
+    "Spanish (Mexico)",
+    "French",
+    "French (France)",
+    "German",
+    "German (Germany)",
+    "Portuguese",
+    "Portuguese (Brazil)",
+    "Italian",
+    "Italian (Italy)",
+    "Japanese",
+    "Japanese (Japan)",
+    "Korean",
+    "Chinese (Mandarin, Simplified)",
+    "Arabic",
+    "Hindi",
+    "Hindi (India)",
+    "Russian",
+    "Russian (Russia)",
+    "Dutch",
+    "Polish",
+    "Turkish",
+    "Indonesian",
+    "Vietnamese",
+    "Ukrainian",
+    "Afrikaans (South Africa)",
+    "Albanian (Albania)",
+    "Amharic (Ethiopia)",
+    "Arabic (Algeria)",
+    "Arabic (Bahrain)",
+    "Arabic (Egypt)",
+    "Arabic (Iraq)",
+    "Arabic (Jordan)",
+    "Arabic (Kuwait)",
+    "Arabic (Lebanon)",
+    "Arabic (Libya)",
+    "Arabic (Morocco)",
+    "Arabic (Oman)",
+    "Arabic (Qatar)",
+    "Arabic (Saudi Arabia)",
+    "Arabic (Syria)",
+    "Arabic (Tunisia)",
+    "Arabic (United Arab Emirates)",
+    "Arabic (World)",
+    "Arabic (Yemen)",
+    "Armenian (Armenia)",
+    "Azerbaijani (Latin, Azerbaijan)",
+    "Bangla (Bangladesh)",
+    "Basque",
+    "Belarusian (Belarus)",
+    "Bengali (India)",
+    "Bosnian (Bosnia and Herzegovina)",
+    "Bulgarian",
+    "Bulgarian (Bulgaria)",
+    "Burmese (Myanmar)",
+    "Catalan",
+    "Chinese (Cantonese, Traditional)",
+    "Chinese (Jilu Mandarin, Simplified)",
+    "Chinese (Northeastern Mandarin, Simplified)",
+    "Chinese (Southwestern Mandarin, Simplified)",
+    "Chinese (Taiwanese Mandarin, Traditional)",
+    "Chinese (Wu, Simplified)",
+    "Chinese (Zhongyuan Mandarin Henan, Simplified)",
+    "Chinese (Zhongyuan Mandarin Shaanxi, Simplified)",
+    "Croatian",
+    "Croatian (Croatia)",
+    "Czech",
+    "Czech (Czechia)",
+    "Danish",
+    "Danish (Denmark)",
+    "Dutch (Belgium)",
+    "Dutch (Netherlands)",
+    "English (Australia)",
+    "English (Canada)",
+    "English (Hong Kong SAR)",
+    "English (India)",
+    "English (Ireland)",
+    "English (Kenya)",
+    "English (New Zealand)",
+    "English (Nigeria)",
+    "English (Philippines)",
+    "English (Singapore)",
+    "English (South Africa)",
+    "English (Tanzania)",
+    "English (UK)",
+    "English (United States)",
+    "Estonian (Estonia)",
+    "Filipino",
+    "Filipino (Cebuano)",
+    "Filipino (Philippines)",
+    "Finnish",
+    "Finnish (Finland)",
+    "French (Belgium)",
+    "French (Canada)",
+    "French (Switzerland)",
+    "Galician",
+    "Georgian (Georgia)",
+    "German (Austria)",
+    "German (Switzerland)",
+    "Greek",
+    "Greek (Greece)",
+    "Gujarati (India)",
+    "Haitian Creole (Haiti)",
+    "Hebrew (Israel)",
+    "Hungarian (Hungary)",
+    "Icelandic (Iceland)",
+    "Indonesian (Indonesia)",
+    "Irish (Ireland)",
+    "Javanese (Latin, Indonesia)",
+    "Kannada (India)",
+    "Kazakh (Kazakhstan)",
+    "Khmer (Cambodia)",
+    "Konkani (India)",
+    "Korean (Korea)",
+    "Lao (Laos)",
+    "Latin (Vatican City)",
+    "Latvian (Latvia)",
+    "Lithuanian (Lithuania)",
+    "Luxembourgish (Luxembourg)",
+    "Macedonian (North Macedonia)",
+    "Maithili (India)",
+    "Malagasy (Madagascar)",
+    "Malay",
+    "Malay (Malaysia)",
+    "Malayalam (India)",
+    "Maltese (Malta)",
+    "Mandarin",
+    "Marathi (India)",
+    "Mongolian (Mongolia)",
+    "Nepali (Nepal)",
+    "Norwegian Bokmål (Norway)",
+    "Norwegian Nynorsk (Norway)",
+    "Odia (India)",
+    "Pashto (Afghanistan)",
+    "Persian (Iran)",
+    "Polish (Poland)",
+    "Portuguese (Portugal)",
+    "Punjabi (India)",
+    "Romanian",
+    "Romanian (Romania)",
+    "Serbian (Latin, Serbia)",
+    "Sindhi (India)",
+    "Sinhala (Sri Lanka)",
+    "Slovak",
+    "Slovak (Slovakia)",
+    "Slovenian (Slovenia)",
+    "Somali (Somalia)",
+    "Spanish (Argentina)",
+    "Spanish (Bolivia)",
+    "Spanish (Chile)",
+    "Spanish (Colombia)",
+    "Spanish (Costa Rica)",
+    "Spanish (Cuba)",
+    "Spanish (Dominican Republic)",
+    "Spanish (Ecuador)",
+    "Spanish (El Salvador)",
+    "Spanish (Equatorial Guinea)",
+    "Spanish (Guatemala)",
+    "Spanish (Honduras)",
+    "Spanish (Latin America)",
+    "Spanish (Nicaragua)",
+    "Spanish (Panama)",
+    "Spanish (Paraguay)",
+    "Spanish (Peru)",
+    "Spanish (Puerto Rico)",
+    "Spanish (United States)",
+    "Spanish (Uruguay)",
+    "Spanish (Venezuela)",
+    "Sundanese (Indonesia)",
+    "Swahili (Kenya)",
+    "Swahili (Tanzania)",
+    "Swedish",
+    "Swedish (Sweden)",
+    "Tamil",
+    "Tamil (India)",
+    "Tamil (Malaysia)",
+    "Tamil (Singapore)",
+    "Tamil (Sri Lanka)",
+    "Telugu (India)",
+    "Thai (Thailand)",
+    "Turkish (Türkiye)",
+    "Ukrainian (Ukraine)",
+    "Urdu (India)",
+    "Urdu (Pakistan)",
+    "Uzbek (Latin, Uzbekistan)",
+    "Vietnamese (Vietnam)",
+    "Welsh (United Kingdom)",
+    "Zulu (South Africa)",
+]

comfy_api_nodes/apis/hunyuan3d.py

@@ -77,6 +77,7 @@ class To3DUVTaskRequest(BaseModel):
 
 class To3DPartTaskRequest(BaseModel):
     File: TaskFile3DInput = Field(...)
+    EnableStagedGeneration: bool | None = Field(None)
 
 
 class TextureEditImageInfo(BaseModel):

comfy_api_nodes/apis/openai.py

@@ -128,7 +128,7 @@ class OpenAIResponse(ModelResponseProperties, ResponseProperties):
     parallel_tool_calls: bool | None = Field(True)
     status: str | None = Field(
         None,
-        description="One of `completed`, `failed`, `in_progress`, or `incomplete`.",
+        description="One of `completed`, `failed`, `in_progress`, `incomplete`, `queued`, or `cancelled`.",
     )
     usage: ResponseUsage | None = Field(None)
 

comfy_api_nodes/apis/recraft.py

@@ -244,10 +244,10 @@ RECRAFT_V4_PRO_SIZES = [
     "2304x1792",
     "1792x2304",
     "1664x2688",
-    "1434x1024",
-    "1024x1434",
     "2560x1792",
     "1792x2560",
+    "2688x1536",
+    "1536x2688",
 ]
 
 

comfy_api_nodes/apis/sync_so.py

@@ -0,0 +1,49 @@
+from pydantic import BaseModel, Field
+
+
+class SyncInputItem(BaseModel):
+    type: str = Field(..., description="Input kind: 'video', 'image' or 'audio'.")
+    url: str = Field(...)
+
+
+class SyncActiveSpeakerDetection(BaseModel):
+    auto_detect: bool | None = Field(
+        None, description="Detect the active speaker automatically. Video input only; rejected for images."
+    )
+    frame_number: int | None = Field(
+        None, description="Frame used for manual speaker selection. Must be 0 for image inputs."
+    )
+    coordinates: list[int] | None = Field(
+        None, description="Pixel [x, y] of the speaker's face in the frame selected by frame_number."
+    )
+
+
+class SyncGenerationOptions(BaseModel):
+    sync_mode: str | None = Field(
+        None,
+        description="How to resolve an audio/video duration mismatch: "
+        "cut_off, bounce, loop, silence or remap. Ignored for image inputs.",
+    )
+    i2v_prompt: str | None = Field(
+        None, description="Motion prompt for image-to-video generation. Image input only."
+    )
+    active_speaker_detection: SyncActiveSpeakerDetection | None = Field(None)
+
+
+class SyncGenerationRequest(BaseModel):
+    model: str = Field(..., description="Generation model, e.g. 'sync-3'.")
+    input: list[SyncInputItem] = Field(
+        ..., description="Exactly one visual input (video or image) plus one audio input."
+    )
+    options: SyncGenerationOptions | None = Field(None)
+
+
+class SyncGeneration(BaseModel):
+    """Subset of the Generation object returned by POST /v2/generate and GET /v2/generate/{id}."""
+
+    id: str = Field(...)
+    status: str = Field(..., description="PENDING | PROCESSING | COMPLETED | FAILED | REJECTED")
+    outputUrl: str | None = Field(None)
+    outputDuration: float | None = Field(None)
+    error: str | None = Field(None, description="Human-readable failure message.")
+    errorCode: str | None = Field(None, description="Stable machine-readable code from the GET /v2/errors catalog.")

comfy_api_nodes/nodes_anthropic.py

@@ -28,6 +28,10 @@ ANTHROPIC_IMAGE_MAX_PIXELS = 1568 * 1568
 CLAUDE_MAX_IMAGES = 20
 
 CLAUDE_MODELS: dict[str, str] = {
+    "Opus 5": "claude-opus-5",
+    "Opus 4.8": "claude-opus-4-8",
+    "Fable 5": "claude-fable-5",
+    "Sonnet 5": "claude-sonnet-5",
     "Opus 4.7": "claude-opus-4-7",
     "Opus 4.6": "claude-opus-4-6",
     "Sonnet 4.6": "claude-sonnet-4-6",
@@ -36,9 +40,12 @@ CLAUDE_MODELS: dict[str, str] = {
 }
 
 _THINKING_UNSUPPORTED = {"Haiku 4.5"}
-# Models that use the newer "adaptive" thinking mode (Opus 4.7 requires it; older models keep the explicit budget API).
+# Models that use the newer "adaptive" thinking mode (Opus 4.7+ require it; older models keep the explicit budget API).
 # Anthropic decides the actual budget when adaptive is used, based on the `output_config.effort` hint.
-_ADAPTIVE_THINKING_MODELS = {"Opus 4.7", "Opus 4.6", "Sonnet 4.6"}
+_ADAPTIVE_THINKING_MODELS = {"Opus 4.8", "Sonnet 5", "Opus 4.7", "Opus 4.6", "Sonnet 4.6"}
+_ALWAYS_THINKING_MODELS = {"Opus 5", "Fable 5"}
+_EXPLICIT_THINKING_OFF_MODELS = {"Sonnet 5"}
+_NO_TEMPERATURE_MODELS = {"Opus 5", "Opus 4.8", "Fable 5", "Sonnet 5"}
 
 # Budget mode (Sonnet 4.5): effort -> reasoning budget in tokens. Must be < max_tokens.
 # Sized so even the "high" budget fits comfortably under the default max_tokens=32768.
@@ -60,20 +67,33 @@ def _claude_model_inputs(model_label: str):
             tooltip="Maximum number of tokens to generate (includes reasoning tokens when enabled).",
             advanced=True,
         ),
-        IO.Float.Input(
-            "temperature",
-            default=1.0,
-            min=0.0,
-            max=1.0,
-            step=0.01,
-            tooltip=(
-                "Controls randomness. 0.0 is deterministic, 1.0 is most random. "
-                "Ignored for Opus 4.7 and any model when reasoning_effort is set."
-            ),
-            advanced=True,
-        ),
     ]
-    if model_label not in _THINKING_UNSUPPORTED:
+    if model_label not in _NO_TEMPERATURE_MODELS:
+        inputs.append(
+            IO.Float.Input(
+                "temperature",
+                default=1.0,
+                min=0.0,
+                max=1.0,
+                step=0.01,
+                tooltip=(
+                    "Controls randomness. 0.0 is deterministic, 1.0 is most random. "
+                    "Ignored for Opus 4.7 and any model when reasoning_effort is set."
+                ),
+                advanced=True,
+            )
+        )
+    if model_label in _ALWAYS_THINKING_MODELS:
+        inputs.append(
+            IO.Combo.Input(
+                "reasoning_effort",
+                options=[e for e in _REASONING_EFFORTS if e != "off"],
+                default="high",
+                tooltip="Extended thinking effort. Reasoning is always enabled for this model.",
+                advanced=True,
+            )
+        )
+    elif model_label not in _THINKING_UNSUPPORTED:
         inputs.append(
             IO.Combo.Input(
                 "reasoning_effort",
@@ -86,43 +106,6 @@ def _claude_model_inputs(model_label: str):
     return inputs
 
 
-def _model_price_per_million(model: str) -> tuple[float, float] | None:
-    """Return (input_per_1M, output_per_1M) USD for a Claude model, or None if unknown."""
-    if "opus-4-7" in model or "opus-4-6" in model or "opus-4-5" in model:
-        return 5.0, 25.0
-    if "sonnet-4" in model:
-        return 3.0, 15.0
-    if "haiku-4-5" in model:
-        return 1.0, 5.0
-    return None
-
-
-def calculate_tokens_price(response: AnthropicMessagesResponse) -> float | None:
-    """Compute approximate USD price from response usage. Server-side billing is authoritative."""
-    if not response.usage or not response.model:
-        return None
-    rates = _model_price_per_million(response.model)
-    if rates is None:
-        return None
-    input_rate, output_rate = rates
-    input_tokens = response.usage.input_tokens or 0
-    output_tokens = response.usage.output_tokens or 0
-    cache_read = response.usage.cache_read_input_tokens or 0
-    cache_5m = 0
-    cache_1h = 0
-    if response.usage.cache_creation:
-        cache_5m = response.usage.cache_creation.ephemeral_5m_input_tokens or 0
-        cache_1h = response.usage.cache_creation.ephemeral_1h_input_tokens or 0
-    total = (
-        input_tokens * input_rate
-        + output_tokens * output_rate
-        + cache_read * input_rate * 0.1
-        + cache_5m * input_rate * 1.25
-        + cache_1h * input_rate * 2.0
-    )
-    return total / 1_000_000.0
-
-
 def _get_text_from_response(response: AnthropicMessagesResponse) -> str:
     if not response.content:
         return ""
@@ -213,7 +196,27 @@ class ClaudeNode(IO.ComfyNode):
                 expr="""
                 (
                   $m := widgets.model;
-                  $contains($m, "opus") ? {
+                  $contains($m, "fable") ? {
+                    "type": "list_usd",
+                    "usd": [0.0143, 0.0715],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : $contains($m, "opus 4.8") ? {
+                    "type": "list_usd",
+                    "usd": [0.00715, 0.03575],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : $contains($m, "sonnet 5") ? {
+                    "type": "list_usd",
+                    "usd": [0.00286, 0.0143],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : $contains($m, "opus 5") ? {
+                    "type": "list_usd",
+                    "usd": [0.00715, 0.03575],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : $contains($m, "opus") ? {
                     "type": "list_usd",
                     "usd": [0.005, 0.025],
                     "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
@@ -247,18 +250,23 @@ class ClaudeNode(IO.ComfyNode):
         model_label = model["model"]
         max_tokens = model.get("max_tokens", 32768)
         reasoning_effort = model.get("reasoning_effort", "off")
-        thinking_enabled = reasoning_effort not in ("off", None) and model_label not in _THINKING_UNSUPPORTED
+        always_thinking = model_label in _ALWAYS_THINKING_MODELS
+        thinking_enabled = always_thinking or (
+            reasoning_effort not in ("off", None) and model_label not in _THINKING_UNSUPPORTED
+        )
 
         # Anthropic requires temperature to be unset (defaults to 1.0) when thinking is enabled.
         # Opus 4.7 also rejects user-supplied temperature.
-        if thinking_enabled or model_label == "Opus 4.7":
+        if model_label in _NO_TEMPERATURE_MODELS or thinking_enabled or model_label == "Opus 4.7":
             temperature = None
         else:
             temperature = model.get("temperature", 1.0)
 
         thinking_cfg: AnthropicThinkingConfig | None = None
         output_cfg: AnthropicOutputConfig | None = None
-        if thinking_enabled:
+        if always_thinking:
+            output_cfg = AnthropicOutputConfig(effort=reasoning_effort)
+        elif thinking_enabled:
             if model_label in _ADAPTIVE_THINKING_MODELS:
                 # Adaptive mode - Anthropic chooses the budget based on effort hint
                 thinking_cfg = AnthropicThinkingConfig(type="adaptive")
@@ -268,6 +276,8 @@ class ClaudeNode(IO.ComfyNode):
                 budget = _REASONING_BUDGET[reasoning_effort]
                 budget = min(budget, max(1024, max_tokens - 1024))
                 thinking_cfg = AnthropicThinkingConfig(type="enabled", budget_tokens=budget)
+        elif model_label in _EXPLICIT_THINKING_OFF_MODELS:
+            thinking_cfg = AnthropicThinkingConfig(type="disabled")
 
         image_tensors: list[Input.Image] = [t for t in (images or {}).values() if t is not None]
         if sum(get_number_of_images(t) for t in image_tensors) > CLAUDE_MAX_IMAGES:
@@ -291,8 +301,12 @@ class ClaudeNode(IO.ComfyNode):
                 thinking=thinking_cfg,
                 output_config=output_cfg,
             ),
-            price_extractor=calculate_tokens_price,
         )
+        if response.stop_reason == "refusal":
+            raise ValueError(
+                "Claude declined to answer this request for safety reasons. "
+                "Rephrase the prompt or try a different model."
+            )
         return IO.NodeOutput(_get_text_from_response(response) or "Empty response from Claude model.")
 
 

comfy_api_nodes/nodes_bytedance.py

@@ -34,6 +34,7 @@ from comfy_api_nodes.apis.bytedance import (
     SeedanceVirtualLibraryCreateAssetRequest,
     Seedream4Options,
     Seedream4TaskCreationRequest,
+    Seedream5OptimizePromptOptions,
     TaskAudioContent,
     TaskAudioContentUrl,
     TaskCreationResponse,
@@ -875,6 +876,17 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                     tooltip='Whether to add an "AI generated" watermark to the image.',
                     advanced=True,
                 ),
+                IO.Boolean.Input(
+                    "thinking",
+                    default=True,
+                    tooltip=(
+                        "Enable the model's prompt-optimization reasoning ('thinking') for better adherence. "
+                        "Can substantially increase generation time — notably on Seedream 5.0 Pro. "
+                        "Can only be disabled for text-to-image (not when reference images are provided)."
+                    ),
+                    optional=True,
+                    advanced=True,
+                ),
             ],
             outputs=[
                 IO.Image.Output(),
@@ -920,6 +932,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
         model: dict,
         seed: int = 0,
         watermark: bool = False,
+        thinking: bool = True,
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         model_id = SEEDREAM_MODELS[model["model"]]
@@ -979,6 +992,10 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
             raise ValueError(
                 "The maximum number of generated images plus the number of reference images cannot exceed 15."
             )
+        if not thinking and n_input_images > 0:
+            raise ValueError(
+                "'thinking' can only be disabled for text-to-image; enable it when using reference images."
+            )
 
         reference_images_urls: list[str] = []
         if image_tensors:
@@ -992,6 +1009,9 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                 wait_label="Uploading reference images",
             )
 
+        optimize_prompt_options = None
+        if n_input_images == 0:
+            optimize_prompt_options = Seedream5OptimizePromptOptions(thinking="enabled" if thinking else "disabled")
         response = await sync_op(
             cls,
             ApiEndpoint(path=BYTEPLUS_IMAGE_ENDPOINT, method="POST"),
@@ -1005,6 +1025,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                 sequential_image_generation=None if is_pro else sequential_image_generation,
                 sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images),
                 watermark=watermark,
+                optimize_prompt_options=optimize_prompt_options,
             ),
         )
         if len(response.data) == 1:
@@ -2669,7 +2690,8 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
                 "with ByteDance Seed Audio 1.0. Describe the voice(s), emotion, ambience, background music "
                 "and sound effects in the prompt, and include the lines to speak. Optionally pick a built-in "
                 "preset voice, clone voices from up to 3 reference clips (tagged @Audio1-3 in the prompt), "
-                "or derive a voice from a character image. Up to 2 minutes of audio per run."
+                "or derive a voice from a character image. Up to 2 minutes of audio per run. "
+                "The multilingual model supports 20 languages and timestamp-based timing control."
             ),
             inputs=[
                 IO.String.Input(
@@ -2680,7 +2702,9 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
                         "Describe the voice(s), emotion, pacing, ambience, background music and sound "
                         "effects, and include the lines to speak (name characters inline for dialogue). "
                         "In 'audio reference' mode, refer to connected clips by order as @Audio1, @Audio2, "
-                        "@Audio3. Maximum 3000 characters."
+                        "@Audio3. With the multilingual model, a quoted line can start with a timestamp "
+                        'range that controls when and how long it is spoken, e.g. "[5.5s:8.0s] Wait for me!". '
+                        "Write the prompt in the same language as the lines to speak. Maximum 3000 characters."
                     ),
                 ),
                 IO.DynamicCombo.Input(
@@ -2775,6 +2799,19 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
                     tooltip="Seed controls whether the node should re-run; "
                     "results are non-deterministic regardless of seed.",
                 ),
+                IO.Combo.Input(
+                    "model",
+                    options=["seed-audio-1.0-multilingual", "seed-audio-1.0"],
+                    default="seed-audio-1.0-multilingual",
+                    optional=True,
+                    tooltip=(
+                        "seed-audio-1.0-multilingual: 20 languages (English, Chinese, Japanese, Korean, "
+                        "Mexican & Castilian Spanish, Indonesian, German, Brazilian Portuguese, French, "
+                        "Thai, Vietnamese, Malay, Filipino, Italian, Russian, Dutch, Polish, Turkish, "
+                        'Swedish) plus per-sentence timing control via "[5.5s:8.0s] ..." timestamps. '
+                        "seed-audio-1.0: English and Chinese only, no timing control."
+                    ),
+                ),
             ],
             outputs=[IO.Audio.Output()],
             hidden=[
@@ -2798,6 +2835,7 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
         loudness_rate: int,
         pitch_rate: int,
         seed: int,
+        model: str = "seed-audio-1.0-multilingual",
     ) -> IO.NodeOutput:
         mode = reference_mode["reference_mode"]
         audio_indices = connected_audio_indices(reference_mode)
@@ -2824,6 +2862,7 @@ class ByteDanceSeedAudioNode(IO.ComfyNode):
             ApiEndpoint(path="/proxy/byteplus/api/v3/tts/create", method="POST"),
             response_model=SeedAudioResponse,
             data=SeedAudioRequest(
+                model=model,
                 text_prompt=text_prompt,
                 references=references,
                 audio_config=SeedAudioConfig(

comfy_api_nodes/nodes_bytedance_llm.py

@@ -34,13 +34,6 @@ SEED_MODELS: dict[str, str] = {
     "Seed 2.0 Mini": "seed-2-0-mini-260215",
 }
 
-# USD per 1M tokens: (input, cache_hit_input, output)
-_SEED_PRICES_PER_MILLION: dict[str, tuple[float, float, float]] = {
-    "seed-2-0-pro-260328": (0.50, 0.10, 3.00),
-    "seed-2-0-lite-260228": (0.25, 0.05, 2.00),
-    "seed-2-0-mini-260215": (0.10, 0.02, 0.40),
-}
-
 
 def _seed_model_inputs(max_images: int = SEED_MAX_IMAGES, max_videos: int = SEED_MAX_VIDEOS):
     return [
@@ -74,24 +67,6 @@ def _seed_model_inputs(max_images: int = SEED_MAX_IMAGES, max_videos: int = SEED
     ]
 
 
-def _calculate_price(model_id: str, response: BytePlusResponseObject) -> float | None:
-    """Compute approximate USD price from response usage."""
-    if not response.usage:
-        return None
-    rates = _SEED_PRICES_PER_MILLION.get(model_id)
-    if rates is None:
-        return None
-    input_rate, cache_hit_rate, output_rate = rates
-    input_tokens = response.usage.input_tokens or 0
-    output_tokens = response.usage.output_tokens or 0
-    cached = 0
-    if response.usage.input_tokens_details:
-        cached = response.usage.input_tokens_details.cached_tokens or 0
-    fresh_input = max(0, input_tokens - cached)
-    total = fresh_input * input_rate + cached * cache_hit_rate + output_tokens * output_rate
-    return total / 1_000_000.0
-
-
 def _get_text_from_response(response: BytePlusResponseObject) -> str:
     """Extract concatenated text from all assistant message output_text blocks."""
     if not response.output:
@@ -251,7 +226,6 @@ class ByteDanceSeedNode(IO.ComfyNode):
                 store=False,
                 stream=False,
             ),
-            price_extractor=lambda r: _calculate_price(model_id, r),
         )
         if response.error:
             raise ValueError(f"Seed API error ({response.error.code}): {response.error.message}")

comfy_api_nodes/nodes_gemini.py

@@ -24,13 +24,17 @@ from comfy_api_nodes.apis.gemini import (
     GeminiImageGenerateContentRequest,
     GeminiImageGenerationConfig,
     GeminiInlineData,
+    GeminiInteraction,
+    GeminiInteractionGenerationConfig,
+    GeminiInteractionMediaPart,
+    GeminiInteractionRequest,
+    GeminiInteractionTextPart,
     GeminiMimeType,
     GeminiPart,
     GeminiRole,
     GeminiSystemInstructionContent,
     GeminiTextPart,
     GeminiThinkingConfig,
-    Modality,
 )
 from comfy_api_nodes.util import (
     ApiEndpoint,
@@ -51,9 +55,11 @@ from comfy_api_nodes.util import (
 )
 
 GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini"
+GEMINI_INTERACTIONS_ENDPOINT = "/proxy/gemini-interactions"
 GEMINI_MAX_INPUT_FILE_SIZE = 20 * 1024 * 1024  # 20 MB
 GEMINI_URL_INPUT_BUDGET = 10
 GEMINI_MAX_INLINE_BYTES = 18 * 1024 * 1024
+GEMINI_INTERACTIONS_MAX_INLINE_BYTES = 90 * 1024 * 1024  # the Interactions API rejects requests over ~100MiB
 GEMINI_IMAGE_SYS_PROMPT = (
     "You are an expert image-generation engine. You must ALWAYS produce an image.\n"
     "Interpret all user input—regardless of "
@@ -231,16 +237,32 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
     return torch.cat(image_tensors, dim=0)
 
 
-async def get_video_from_response(
-    response: GeminiGenerateContentResponse, cls: type[IO.ComfyNode] | None = None
+def get_text_from_interaction(interaction: GeminiInteraction) -> str:
+    """Extract and concatenate all model output text from an Interactions API response."""
+    texts = []
+    for step in interaction.steps or []:
+        if step.type != "model_output":
+            continue
+        for content in step.content or []:
+            if content.type == "text" and content.text:
+                texts.append(content.text)
+    return "\n".join(texts)
+
+
+async def get_video_from_interaction(
+    interaction: GeminiInteraction, cls: type[IO.ComfyNode] | None = None
 ) -> InputImpl.VideoFromFile:
-    parts = get_parts_by_type(response, "video/*")
-    for part in parts:
-        if part.inlineData and part.inlineData.data:
-            return InputImpl.VideoFromFile(BytesIO(base64.b64decode(part.inlineData.data)))
-        if part.fileData and part.fileData.fileUri:
-            return await download_url_to_video_output(part.fileData.fileUri, cls=cls)
-    model_message = get_text_from_response(response).strip()
+    for step in interaction.steps or []:
+        if step.type != "model_output":
+            continue
+        for content in step.content or []:
+            if content.type != "video":
+                continue
+            if content.data:
+                return InputImpl.VideoFromFile(BytesIO(base64.b64decode(content.data)))
+            if content.uri:
+                return await download_url_to_video_output(content.uri, cls=cls)
+    model_message = get_text_from_interaction(interaction).strip()
     if model_message:
         raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
     raise ValueError(
@@ -249,64 +271,6 @@ async def get_video_from_response(
     )
 
 
-def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | None:
-    if not response.modelVersion:
-        return None
-    # Define prices (Cost per 1,000,000 tokens), see https://cloud.google.com/vertex-ai/generative-ai/pricing
-    output_video_tokens_price = 0.0
-    if response.modelVersion == "gemini-2.5-pro":
-        input_tokens_price = 1.25
-        output_text_tokens_price = 10.0
-        output_image_tokens_price = 0.0
-    elif response.modelVersion == "gemini-2.5-flash":
-        input_tokens_price = 0.30
-        output_text_tokens_price = 2.50
-        output_image_tokens_price = 0.0
-    elif response.modelVersion == "gemini-2.5-flash-image":
-        input_tokens_price = 0.30
-        output_text_tokens_price = 2.50
-        output_image_tokens_price = 30.0
-    elif response.modelVersion in ("gemini-3-pro-preview", "gemini-3.1-pro-preview"):
-        input_tokens_price = 2
-        output_text_tokens_price = 12.0
-        output_image_tokens_price = 0.0
-    elif response.modelVersion in ("gemini-3.1-flash-lite-preview", "gemini-3.1-flash-lite"):
-        input_tokens_price = 0.25
-        output_text_tokens_price = 1.50
-        output_image_tokens_price = 0.0
-    elif response.modelVersion in ("gemini-3-pro-image-preview", "gemini-3-pro-image"):
-        input_tokens_price = 2
-        output_text_tokens_price = 12.0
-        output_image_tokens_price = 120.0
-    elif response.modelVersion in ("gemini-3.1-flash-image-preview", "gemini-3.1-flash-image"):
-        input_tokens_price = 0.5
-        output_text_tokens_price = 3.0
-        output_image_tokens_price = 60.0
-    elif response.modelVersion == "gemini-3.1-flash-lite-image":
-        input_tokens_price = 0.25
-        output_text_tokens_price = 1.50
-        output_image_tokens_price = 30.0
-    elif response.modelVersion == "gemini-omni-flash-preview":
-        input_tokens_price = 2.145
-        output_text_tokens_price = 12.87
-        output_image_tokens_price = 0.0
-        output_video_tokens_price = 25.025
-    else:
-        return None
-    final_price = response.usageMetadata.promptTokenCount * input_tokens_price
-    if response.usageMetadata.candidatesTokensDetails:
-        for i in response.usageMetadata.candidatesTokensDetails:
-            if i.modality == Modality.IMAGE:
-                final_price += output_image_tokens_price * i.tokenCount  # for Nano Banana models
-            elif i.modality == Modality.VIDEO:
-                final_price += output_video_tokens_price * i.tokenCount  # for Omni Flash
-            else:
-                final_price += output_text_tokens_price * i.tokenCount
-    if response.usageMetadata.thoughtsTokenCount:
-        final_price += output_text_tokens_price * response.usageMetadata.thoughtsTokenCount
-    return final_price / 1_000_000.0
-
-
 def create_video_parts(video_input: Input.Video) -> list[GeminiPart]:
     """Convert a single video input to Gemini API compatible parts (inline MP4/H.264)."""
     base_64_string = video_to_base64_string(
@@ -433,14 +397,24 @@ async def build_gemini_media_parts(
         part, nbytes = _media_inline_part(kind, payload)
         inline_bytes += nbytes
         if inline_bytes > max_inline_bytes:
+            detail = f" after the first {url_budget} inputs are uploaded as URLs" if url_budget else ""
             raise ValueError(
-                f"Too much media to send inline (over {max_inline_bytes // (1024 * 1024)}MB after the first "
-                f"{url_budget} inputs are uploaded as URLs). Reduce the number or size of attached media."
+                f"Too much media to send inline (over {max_inline_bytes // (1024 * 1024)}MB{detail}). "
+                "Reduce the number or size of attached media."
             )
         parts.append(part)
     return parts
 
 
+def to_interaction_media_part(part: GeminiPart) -> GeminiInteractionMediaPart:
+    """Convert a fileData/inlineData GeminiPart into an Interactions API media part."""
+    if part.fileData:
+        mime = part.fileData.mimeType.value
+        return GeminiInteractionMediaPart(type=mime.split("/")[0], uri=part.fileData.fileUri, mime_type=mime)
+    mime = part.inlineData.mimeType.value
+    return GeminiInteractionMediaPart(type=mime.split("/")[0], data=part.inlineData.data, mime_type=mime)
+
+
 class GeminiNode(IO.ComfyNode):
     """
     Node to generate text responses from a Gemini model.
@@ -610,7 +584,6 @@ class GeminiNode(IO.ComfyNode):
                 systemInstruction=gemini_system_prompt,
             ),
             response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
         )
 
         output_text = get_text_from_response(response)
@@ -619,11 +592,12 @@ class GeminiNode(IO.ComfyNode):
 
 GEMINI_V2_MODELS: dict[str, str] = {
     "Gemini 3.1 Pro": "gemini-3.1-pro-preview",
+    "Gemini 3.5 Flash": "gemini-3.5-flash",
     "Gemini 3.1 Flash-Lite": "gemini-3.1-flash-lite-preview",
 }
 
 
-def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
+def _gemini_text_model_inputs(thinking_default: str, thinking_options: list[str] | None = None) -> list[Input]:
     """Per-model inputs revealed by the model DynamicCombo (shared media + sampling controls)."""
     return [
         IO.Autogrow.Input(
@@ -661,7 +635,7 @@ def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
         ),
         IO.Combo.Input(
             "thinking_level",
-            options=["LOW", "HIGH"],
+            options=thinking_options or ["LOW", "HIGH"],
             default=thinking_default,
             tooltip="How hard the model reasons internally before answering. "
             "HIGH improves quality on difficult tasks but costs more (thinking) tokens and is slower.",
@@ -719,6 +693,10 @@ class GeminiNodeV2(IO.ComfyNode):
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "Gemini 3.5 Flash",
+                            _gemini_text_model_inputs("MEDIUM", ["MINIMAL", "LOW", "MEDIUM", "HIGH"]),
+                        ),
                         IO.DynamicCombo.Option("Gemini 3.1 Pro", _gemini_text_model_inputs("HIGH")),
                         IO.DynamicCombo.Option("Gemini 3.1 Flash-Lite", _gemini_text_model_inputs("LOW")),
                     ],
@@ -759,7 +737,13 @@ class GeminiNodeV2(IO.ComfyNode):
                     "type": "list_usd",
                     "usd": [0.00025, 0.0015],
                     "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
-                  } : {
+                  }
+                  : $contains($m, "3.5 flash") ? {
+                    "type": "list_usd",
+                    "usd": [0.0015, 0.009],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : {
                     "type": "list_usd",
                     "usd": [0.002, 0.012],
                     "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
@@ -814,7 +798,6 @@ class GeminiNodeV2(IO.ComfyNode):
                 systemInstruction=gemini_system_prompt,
             ),
             response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
         )
 
         output_text = get_text_from_response(response)
@@ -1027,7 +1010,6 @@ class GeminiImage(IO.ComfyNode):
                 systemInstruction=gemini_system_prompt,
             ),
             response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
         )
         return IO.NodeOutput(await get_image_from_response(response), get_text_from_response(response))
 
@@ -1133,7 +1115,9 @@ class GeminiImage2(IO.ComfyNode):
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         if model == "Nano Banana 2 (Gemini 3.1 Flash Image)":
-            model = "gemini-3.1-flash-image-preview"
+            model = "gemini-3.1-flash-image"
+        elif model == "gemini-3-pro-image-preview":
+            model = "gemini-3-pro-image"
 
         parts: list[GeminiPart] = [GeminiPart(text=prompt)]
         if images is not None:
@@ -1165,7 +1149,6 @@ class GeminiImage2(IO.ComfyNode):
                 systemInstruction=gemini_system_prompt,
             ),
             response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
         )
         return IO.NodeOutput(await get_image_from_response(response), get_text_from_response(response))
 
@@ -1325,7 +1308,6 @@ class GeminiNanoBanana2(IO.ComfyNode):
                 systemInstruction=gemini_system_prompt,
             ),
             response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
         )
         return IO.NodeOutput(
             await get_image_from_response(response),
@@ -1507,7 +1489,7 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
         validate_string(prompt, strip_whitespace=True, min_length=1)
         model_choice = model["model"]
         if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)":
-            model_id = "gemini-3.1-flash-image-preview"
+            model_id = "gemini-3.1-flash-image"
         elif model_choice == "Nano Banana 2 Lite":
             model_id = "gemini-3.1-flash-lite-image"
         else:
@@ -1550,7 +1532,6 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
                 systemInstruction=gemini_system_prompt,
             ),
             response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
         )
         return IO.NodeOutput(
             await get_image_from_response(response),
@@ -1659,7 +1640,7 @@ class GeminiVideoOmni(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                expr='{"type":"usd","usd":0.146,"format":{"suffix":"/second","approximate":true}}'
+                expr='{"type":"usd","usd":0.101,"format":{"suffix":"/second","approximate":true}}'
             ),
         )
 
@@ -1678,27 +1659,40 @@ class GeminiVideoOmni(IO.ComfyNode):
         for video in videos:
             validate_video_duration(video, max_duration=10)
 
-        parts: list[GeminiPart] = []
+        parts: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = []
         if images or videos:
-            parts.extend(await build_gemini_media_parts(cls, images, [], videos))
-        parts.append(GeminiPart(text=prompt))
-        response = await sync_op(
+            # The Interactions API accepts video only inline or as a Files API URI, not as an HTTP URL.
+            media_parts = await build_gemini_media_parts(
+                cls, [], [], videos, url_budget=0, max_inline_bytes=GEMINI_INTERACTIONS_MAX_INLINE_BYTES
+            )
+            video_inline_bytes = sum(len(p.inlineData.data) for p in media_parts)
+            media_parts += await build_gemini_media_parts(
+                cls, images, [], [], max_inline_bytes=GEMINI_INTERACTIONS_MAX_INLINE_BYTES - video_inline_bytes
+            )
+            parts.extend(to_interaction_media_part(p) for p in media_parts)
+        parts.append(GeminiInteractionTextPart(text=prompt))
+        interaction = await sync_op(
             cls,
-            ApiEndpoint(path=f"{GEMINI_BASE_ENDPOINT}/{model_id}", method="POST"),
-            data=GeminiGenerateContentRequest(
-                contents=[GeminiContent(role=GeminiRole.user, parts=parts)],
-                generationConfig=GeminiGenerationConfig(
-                    responseModalities=["TEXT", "VIDEO"],
+            ApiEndpoint(path=GEMINI_INTERACTIONS_ENDPOINT, method="POST"),
+            data=GeminiInteractionRequest(
+                model=model_id,
+                input=parts,
+                generation_config=GeminiInteractionGenerationConfig(
                     temperature=model.get("temperature", 1.0),
-                    topP=model.get("top_p", 0.95),
+                    top_p=model.get("top_p", 0.95),
                 ),
             ),
-            response_model=GeminiGenerateContentResponse,
-            price_extractor=calculate_tokens_price,
+            response_model=GeminiInteraction,
         )
+        if interaction.status != "completed":
+            model_message = get_text_from_interaction(interaction).strip()
+            raise ValueError(
+                f"Gemini interaction did not complete (status: {interaction.status})."
+                + (f" Model response: {model_message}" if model_message else "")
+            )
         return IO.NodeOutput(
-            await get_video_from_response(response, cls=cls),
-            get_text_from_response(response),
+            await get_video_from_interaction(interaction, cls=cls),
+            get_text_from_interaction(interaction),
         )
 
 

comfy_api_nodes/nodes_grok.py

@@ -155,7 +155,6 @@ class GrokImageNode(IO.ComfyNode):
                 resolution=resolution.lower(),
             ),
             response_model=ImageGenerationResponse,
-            price_extractor=_extract_grok_price,
         )
         if len(response.data) == 1:
             return IO.NodeOutput(await download_url_to_image_tensor(response.data[0].url))
@@ -351,7 +350,6 @@ class GrokImageEditNode(IO.ComfyNode):
                 aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio,
             ),
             response_model=ImageGenerationResponse,
-            price_extractor=_extract_grok_price,
         )
         if len(response.data) == 1:
             return IO.NodeOutput(await download_url_to_image_tensor(response.data[0].url))
@@ -488,7 +486,6 @@ class GrokImageEditNodeV2(IO.ComfyNode):
                 aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio,
             ),
             response_model=ImageGenerationResponse,
-            price_extractor=_extract_grok_price,
         )
         if len(response.data) == 1:
             return IO.NodeOutput(await download_url_to_image_tensor(response.data[0].url))

comfy_api_nodes/nodes_heygen.py

@@ -0,0 +1,799 @@
+import uuid
+
+import torch
+from typing_extensions import override
+
+from comfy_api.latest import IO, ComfyExtension, Input
+from comfy_api_nodes.apis.heygen import (
+    HEYGEN_AVATAR_MAP,
+    HEYGEN_AVATAR_OPTIONS,
+    HEYGEN_TRANSLATE_LANGUAGES,
+    HEYGEN_VOICE_GENERAL_MAP,
+    HEYGEN_VOICE_GENERAL_OPTIONS,
+    HEYGEN_VOICE_TTS_MAP,
+    HEYGEN_VOICE_TTS_OPTIONS,
+)
+from comfy_api_nodes.util import (
+    ApiEndpoint,
+    audio_bytes_to_audio_input,
+    download_url_as_bytesio,
+    download_url_to_image_tensor,
+    download_url_to_video_output,
+    downscale_image_tensor_by_max_side,
+    get_number_of_images,
+    poll_op_raw,
+    sync_op_raw,
+    upload_audio_to_comfyapi,
+    upload_image_to_comfyapi,
+    upload_images_to_comfyapi,
+    upload_video_to_comfyapi,
+    validate_string,
+)
+from server import PromptServer
+
+_VIDEOS_PATH = "/proxy/heygen/v3/videos"
+_TRANSLATIONS_PATH = "/proxy/heygen/v3/video-translations"
+_SPEECH_PATH = "/proxy/heygen/v3/voices/speech"
+_AVATARS_PATH = "/proxy/heygen/v3/avatars"
+_LOOKS_PATH = "/proxy/heygen/v3/avatars/looks"
+
+_DEFAULT_VOICE_OPTION = "(avatar's default voice)"
+
+_AVATARS_BY_ENGINE = {
+    e: [label for label, (_aid, _atype, engines) in HEYGEN_AVATAR_MAP.items() if e in engines]
+    for e in ("avatar_iv", "avatar_iii", "avatar_v")
+}
+
+
+async def _apply_speech_source(cls: type[IO.ComfyNode], payload: dict, speech: dict, require_voice: bool) -> None:
+    """Fill script/audio speech fields of a /v3/videos payload from the DynamicCombo dict."""
+    if speech["speech"] == "audio":
+        payload["audio_url"] = await upload_audio_to_comfyapi(
+            cls, speech["audio"], container_format="mp3", codec_name="libmp3lame", mime_type="audio/mpeg"
+        )
+    elif speech["speech"] == "script":
+        validate_string(speech["text"], strip_whitespace=True, min_length=1, max_length=5000)
+        payload["script"] = speech["text"]
+        voice_id = speech.get("custom_voice_id", "").strip()
+        if not voice_id and speech["voice"] != _DEFAULT_VOICE_OPTION:
+            voice_id = HEYGEN_VOICE_GENERAL_MAP[speech["voice"]]
+        if voice_id:
+            payload["voice_id"] = voice_id
+        elif require_voice:
+            raise ValueError("A voice is required when driving the video with a text script.")
+        speed = speech.get("voice_speed", 1.0)
+        if speed != 1.0:
+            payload["voice_settings"] = {"speed": round(speed, 2)}
+
+
+async def _create_and_poll_video(cls: type[IO.ComfyNode], payload: dict) -> dict:
+    """POST a /v3/videos payload, poll until terminal, and return the final video data."""
+    created = await sync_op_raw(
+        cls,
+        ApiEndpoint(path=_VIDEOS_PATH, method="POST", headers={"Idempotency-Key": uuid.uuid4().hex}),
+        data=payload,
+    )
+    video_id = (created.get("data") or {}).get("video_id")
+    if not video_id:
+        raise ValueError(f"HeyGen did not return a video_id: {created}")
+    final = await poll_op_raw(
+        cls,
+        ApiEndpoint(path=f"{_VIDEOS_PATH}/{video_id}"),
+        status_extractor=lambda r: (r.get("data") or {}).get("status"),
+        queued_statuses=["pending", "waiting"],
+        poll_interval=5.0,
+    )
+    data = final["data"]
+    if not data.get("video_url"):
+        raise ValueError(f"HeyGen returned no video_url for video {video_id}.")
+    return data
+
+
+async def _resolve_avatar(
+    cls: type[IO.ComfyNode], avatar_label: str, custom_avatar_id: str, engine_choice: str
+) -> tuple[str, str | None]:
+    """Resolve (avatar_id, engine_type) from the combo/override + engine widgets."""
+    custom_avatar_id = custom_avatar_id.strip()
+    if custom_avatar_id:
+        look = (
+            await sync_op_raw(
+                cls,
+                ApiEndpoint(path=f"{_LOOKS_PATH}/{custom_avatar_id}"),
+                final_label_on_success=None,
+            )
+        ).get("data") or {}
+        avatar_id = custom_avatar_id
+        avatar_label = look.get("name") or custom_avatar_id
+        supported = look.get("supported_api_engines") or []
+    else:
+        avatar_id, avatar_type, supported = HEYGEN_AVATAR_MAP[avatar_label]
+
+    if engine_choice == "auto":
+        engine = next((e for e in ("avatar_iv", "avatar_iii", "avatar_v") if e in supported), None)
+    else:
+        engine = engine_choice
+        if supported and engine not in supported:
+            raise ValueError(
+                f"Avatar '{avatar_label}' does not support the {engine} engine "
+                f"(supported: {', '.join(supported)}). Set engine to 'auto' to pick "
+                "a compatible engine automatically."
+            )
+    return avatar_id, engine
+
+
+class HeyGenTalkingPhotoNode(IO.ComfyNode):
+    """Animate a still image of a person into a lip-synced talking video."""
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="HeyGenTalkingPhotoNode",
+            display_name="HeyGen Talking Photo",
+            category="partner/video/HeyGen",
+            description="Animate any image of a person into a lip-synced talking video "
+            "(HeyGen Avatar IV). Drive it with a text script or your own audio.",
+            inputs=[
+                IO.Image.Input(
+                    "image",
+                    tooltip="Image of a person to animate. Downscaled automatically if larger than 2K.",
+                ),
+                IO.DynamicCombo.Input(
+                    "speech",
+                    display_name="speech source",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "script",
+                            [
+                                IO.String.Input(
+                                    "text",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Text for the avatar to speak (up to 5000 characters). "
+                                    "The generated speech must be at least 1 second long.",
+                                ),
+                                IO.Combo.Input(
+                                    "voice",
+                                    options=HEYGEN_VOICE_GENERAL_OPTIONS,
+                                    tooltip="Voice for the script (HeyGen's most popular voices).",
+                                ),
+                                IO.String.Input(
+                                    "custom_voice_id",
+                                    default="",
+                                    optional=True,
+                                    tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
+                                    "Any voice from HeyGen's library (2000+) can be used.",
+                                ),
+                                IO.Float.Input(
+                                    "voice_speed",
+                                    default=1.0,
+                                    min=0.5,
+                                    max=1.5,
+                                    step=0.05,
+                                    optional=True,
+                                    tooltip="Speech speed multiplier.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "audio",
+                            [
+                                IO.Audio.Input(
+                                    "audio",
+                                    tooltip="Audio for the avatar to lip-sync, up to 10 minutes.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="Drive the avatar with a text script (HeyGen text-to-speech) or your own audio.",
+                ),
+                IO.Combo.Input(
+                    "resolution",
+                    options=["720p", "1080p"],
+                    default="1080p",
+                    optional=True,
+                    tooltip="Output video resolution.",
+                ),
+                IO.Combo.Input(
+                    "aspect_ratio",
+                    options=["auto", "16:9", "9:16", "1:1", "4:5", "5:4"],
+                    default="auto",
+                    optional=True,
+                    tooltip="Output aspect ratio. 'auto' follows the input image.",
+                ),
+                IO.Combo.Input(
+                    "expressiveness",
+                    options=["low", "medium", "high"],
+                    default="low",
+                    optional=True,
+                    tooltip="How expressive the animated face and gestures are.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    optional=True,
+                    tooltip="Not sent to HeyGen; change it to force a re-run.",
+                ),
+            ],
+            outputs=[IO.Video.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.0715,"format":{"suffix":"/second"}}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        speech: dict,
+        resolution: str = "1080p",
+        aspect_ratio: str = "auto",
+        expressiveness: str = "low",
+        seed: int = 0,
+    ) -> IO.NodeOutput:
+        image = downscale_image_tensor_by_max_side(image, max_side=2000)
+        image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
+        payload = {
+            "type": "image",
+            "image": {"type": "url", "url": image_url},
+            "resolution": resolution,
+            "aspect_ratio": aspect_ratio,
+            "expressiveness": expressiveness,
+            "title": "ComfyUI Talking Photo",
+        }
+        await _apply_speech_source(cls, payload, speech, require_voice=True)
+        video = await _create_and_poll_video(cls, payload)
+        return IO.NodeOutput(await download_url_to_video_output(video["video_url"]))
+
+
+class HeyGenAvatarVideoNode(IO.ComfyNode):
+    """Generate a presenter video from a HeyGen avatar look."""
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="HeyGenAvatarVideoNode",
+            display_name="HeyGen Avatar Video",
+            category="partner/video/HeyGen",
+            description="Generate a talking-presenter video from a HeyGen avatar. "
+            "Includes HeyGen's most popular public avatars; any look ID can be supplied as an override.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "engine",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "auto",
+                            [
+                                IO.Combo.Input(
+                                    "avatar",
+                                    options=HEYGEN_AVATAR_OPTIONS,
+                                    tooltip="Avatar look to present the video (curated from HeyGen's "
+                                    "public library). The best engine the look supports is chosen "
+                                    "automatically.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "avatar_iv",
+                            [
+                                IO.Combo.Input(
+                                    "avatar",
+                                    options=_AVATARS_BY_ENGINE["avatar_iv"],
+                                    tooltip="Avatar looks that support the Avatar IV engine.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "avatar_iii",
+                            [
+                                IO.Combo.Input(
+                                    "avatar",
+                                    options=_AVATARS_BY_ENGINE["avatar_iii"],
+                                    tooltip="Avatar looks that support the Avatar III engine.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "avatar_v",
+                            [
+                                IO.Combo.Input(
+                                    "avatar",
+                                    options=_AVATARS_BY_ENGINE["avatar_v"],
+                                    tooltip="Avatar looks that support the Avatar V engine.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="Rendering engine; each choice lists only the avatars that support it. "
+                    "'auto' offers every avatar and picks its best engine (Avatar IV preferred). "
+                    "Avatar V is highest fidelity, Avatar III is the most affordable.",
+                ),
+                IO.String.Input(
+                    "custom_avatar_id",
+                    default="",
+                    optional=True,
+                    tooltip="Optional HeyGen avatar look ID. When set, overrides the avatar selected above. "
+                    "Any of HeyGen's 3000+ public looks (or your private avatars) can be used.",
+                ),
+                IO.DynamicCombo.Input(
+                    "speech",
+                    display_name="speech source",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "script",
+                            [
+                                IO.String.Input(
+                                    "text",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Text for the avatar to speak (up to 5000 characters). "
+                                    "The generated speech must be at least 1 second long.",
+                                ),
+                                IO.Combo.Input(
+                                    "voice",
+                                    options=[_DEFAULT_VOICE_OPTION] + HEYGEN_VOICE_GENERAL_OPTIONS,
+                                    tooltip="Voice for the script. The default option uses the voice HeyGen assigned to the avatar.",
+                                ),
+                                IO.String.Input(
+                                    "custom_voice_id",
+                                    default="",
+                                    optional=True,
+                                    tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
+                                    "Any voice from HeyGen's library (2000+) can be used.",
+                                ),
+                                IO.Float.Input(
+                                    "voice_speed",
+                                    default=1.0,
+                                    min=0.5,
+                                    max=1.5,
+                                    step=0.05,
+                                    optional=True,
+                                    tooltip="Speech speed multiplier.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "audio",
+                            [
+                                IO.Audio.Input(
+                                    "audio",
+                                    tooltip="Audio for the avatar to lip-sync, up to 10 minutes.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="Drive the avatar with a text script (HeyGen text-to-speech) or your own audio.",
+                ),
+                IO.Combo.Input(
+                    "resolution",
+                    options=["720p", "1080p"],
+                    default="1080p",
+                    optional=True,
+                    tooltip="Output video resolution.",
+                ),
+                IO.Combo.Input(
+                    "aspect_ratio",
+                    options=["auto", "16:9", "9:16", "1:1", "4:5", "5:4"],
+                    default="auto",
+                    optional=True,
+                    tooltip="Output aspect ratio. 'auto' follows the avatar's source footage.",
+                ),
+                IO.String.Input(
+                    "background_color",
+                    default="",
+                    optional=True,
+                    tooltip="Optional solid background color as a hex code (e.g. '#00ff00'). "
+                    "Leave empty for the avatar's own background.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    optional=True,
+                    tooltip="Not sent to HeyGen; change it to force a re-run.",
+                ),
+            ],
+            outputs=[IO.Video.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                depends_on=IO.PriceBadgeDepends(widgets=["engine"]),
+                expr="""
+                widgets.engine = "avatar_iii"
+                  ? {"type":"range_usd","min_usd":0.023881,"max_usd":0.061919,"format":{"suffix":"/second"}}
+                  : widgets.engine = "avatar_v"
+                  ? {"type":"usd","usd":0.095381,"format":{"suffix":"/second"}}
+                  : widgets.engine = "avatar_iv"
+                  ? {"type":"range_usd","min_usd":0.0715,"max_usd":0.095381,"format":{"suffix":"/second"}}
+                  : {"type":"range_usd","min_usd":0.023881,"max_usd":0.095381,"format":{"suffix":"/second"}}
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        engine: dict,
+        speech: dict,
+        custom_avatar_id: str = "",
+        resolution: str = "1080p",
+        aspect_ratio: str = "auto",
+        background_color: str = "",
+        seed: int = 0,
+    ) -> IO.NodeOutput:
+        avatar_id, engine_type = await _resolve_avatar(cls, engine["avatar"], custom_avatar_id, engine["engine"])
+        payload = {
+            "type": "avatar",
+            "avatar_id": avatar_id,
+            "resolution": resolution,
+            "aspect_ratio": aspect_ratio,
+            "title": "ComfyUI Avatar Video",
+        }
+        if engine_type:
+            payload["engine"] = {"type": engine_type}
+        background_color = background_color.strip()
+        if background_color:
+            if not background_color.startswith("#"):
+                raise ValueError("background_color must be a hex color code like '#00ff00'.")
+            payload["background"] = {"type": "color", "value": background_color}
+        await _apply_speech_source(cls, payload, speech, require_voice=False)
+        video = await _create_and_poll_video(cls, payload)
+        return IO.NodeOutput(await download_url_to_video_output(video["video_url"]))
+
+
+class HeyGenCreateAvatarNode(IO.ComfyNode):
+    """Create a reusable HeyGen avatar from a photo or a text prompt."""
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="HeyGenCreateAvatarNode",
+            display_name="HeyGen Create Avatar",
+            category="partner/video/HeyGen",
+            description="Create your own reusable HeyGen avatar from a photo of a person or "
+            "from a text prompt (a generated character). Feed the resulting avatar_id into "
+            "HeyGen Avatar Video's custom_avatar_id — and save the ID somewhere to reuse the "
+            "avatar in future workflows.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "source",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "prompt",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Description of the avatar to generate (up to 1000 characters).",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_images",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Image.Input("ref_image"),
+                                        names=[f"ref_image_{i}" for i in range(1, 4)],
+                                        min=0,
+                                    ),
+                                    tooltip="Up to 3 reference images guiding the generated look.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "photo",
+                            [
+                                IO.Image.Input(
+                                    "identity_photo",
+                                    tooltip="Photo of the person to turn into an avatar. "
+                                    "Downscaled automatically if larger than 2K.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="Generate a new character from a text prompt, or create the avatar "
+                    "from a connected photo of a person.",
+                ),
+            ],
+            outputs=[
+                IO.String.Output(
+                    display_name="avatar_id",
+                    tooltip="Avatar look ID. Pass it to HeyGen Avatar Video's custom_avatar_id; "
+                    "save it to reuse the avatar later.",
+                ),
+                IO.Image.Output(display_name="preview"),
+            ],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":1.43}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        source: dict,
+    ) -> IO.NodeOutput:
+        payload: dict = {"name": "ComfyUI Avatar"}
+        if source["source"] == "photo":
+            image = downscale_image_tensor_by_max_side(source["identity_photo"], max_side=2000)
+            image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
+            payload["type"] = "photo"
+            payload["file"] = {"type": "url", "url": image_url}
+        else:
+            validate_string(source["prompt"], strip_whitespace=True, min_length=1, max_length=1000)
+            payload["type"] = "prompt"
+            payload["prompt"] = source["prompt"]
+            ref_tensors = [t for t in (source.get("reference_images") or {}).values() if t is not None]
+            if ref_tensors:
+                n_images = sum(get_number_of_images(t) for t in ref_tensors)
+                if n_images > 3:
+                    raise ValueError(f"HeyGen accepts at most 3 reference images; got {n_images}.")
+                scaled = [downscale_image_tensor_by_max_side(t, max_side=2000) for t in ref_tensors]
+                ref_urls = await upload_images_to_comfyapi(
+                    cls, scaled, max_images=3, mime_type="image/png", total_pixels=None
+                )
+                payload["reference_images"] = [{"type": "url", "url": u} for u in ref_urls]
+        created = await sync_op_raw(
+            cls,
+            ApiEndpoint(path=_AVATARS_PATH, method="POST"),
+            data=payload,
+        )
+        look_id = ((created.get("data") or {}).get("avatar_item") or {}).get("id")
+        if not look_id:
+            raise ValueError(f"HeyGen did not return an avatar: {created}")
+        final = await poll_op_raw(
+            cls,
+            ApiEndpoint(path=f"{_LOOKS_PATH}/{look_id}"),
+            # A missing status means the look needed no training and is ready.
+            status_extractor=lambda r: (r.get("data") or {}).get("status") or "completed",
+            failed_statuses=["failed", "pending_consent"],
+            poll_interval=5.0,
+        )
+        data = final["data"]
+        if data.get("preview_image_url"):
+            preview = await download_url_to_image_tensor(data["preview_image_url"])
+        else:
+            preview = torch.zeros(1, 64, 64, 3)
+        PromptServer.instance.send_progress_text(
+            f"Please save the avatar_id for reuse.\n\navatar_id: {look_id}",
+            cls.hidden.unique_id,
+        )
+        return IO.NodeOutput(look_id, preview)
+
+
+class HeyGenVideoTranslateNode(IO.ComfyNode):
+    """Translate a spoken video into another language with voice cloning and lip sync."""
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="HeyGenVideoTranslateNode",
+            display_name="HeyGen Video Translate",
+            category="partner/video/HeyGen",
+            description="Translate a spoken video into another language. Clones the original "
+            "speaker's voice and re-animates the mouth to match the translated speech.",
+            inputs=[
+                IO.Video.Input(
+                    "video",
+                    tooltip="Video with speech to translate.",
+                ),
+                IO.Combo.Input(
+                    "output_language",
+                    options=HEYGEN_TRANSLATE_LANGUAGES,
+                    tooltip="Target language for the translated video.",
+                ),
+                IO.Combo.Input(
+                    "mode",
+                    options=["speed", "precision"],
+                    default="speed",
+                    tooltip="'speed' is faster; 'precision' produces higher-quality lip sync at twice the price.",
+                ),
+                IO.Boolean.Input(
+                    "translate_audio_only",
+                    default=False,
+                    optional=True,
+                    tooltip="Only swap the audio track, keeping the original mouth movements (no lip sync).",
+                ),
+                IO.Int.Input(
+                    "speaker_count",
+                    default=0,
+                    min=0,
+                    max=10,
+                    optional=True,
+                    tooltip="Number of speakers in the video. 0 = detect automatically.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    optional=True,
+                    tooltip="Not sent to HeyGen; change it to force a re-run.",
+                ),
+            ],
+            outputs=[IO.Video.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                depends_on=IO.PriceBadgeDepends(widgets=["mode"]),
+                expr="""{"type":"usd","usd": widgets.mode = "precision" ? 0.095381 : 0.047619,"""
+                """"format":{"suffix":"/second"}}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        video: Input.Video,
+        output_language: str,
+        mode: str,
+        translate_audio_only: bool = False,
+        speaker_count: int = 0,
+        seed: int = 0,
+    ) -> IO.NodeOutput:
+        video_url = await upload_video_to_comfyapi(cls, video)
+        payload = {
+            "video": {"type": "url", "url": video_url},
+            "output_languages": [output_language],
+            "mode": mode,
+            "translate_audio_only": translate_audio_only,
+            "title": "ComfyUI Video Translate",
+        }
+        if speaker_count > 0:
+            payload["speaker_num"] = speaker_count
+        created = await sync_op_raw(
+            cls,
+            ApiEndpoint(path=_TRANSLATIONS_PATH, method="POST"),
+            data=payload,
+        )
+        translation_ids = (created.get("data") or {}).get("video_translation_ids") or []
+        if not translation_ids:
+            raise ValueError(f"HeyGen did not return a translation ID: {created}")
+        final = await poll_op_raw(
+            cls,
+            ApiEndpoint(path=f"{_TRANSLATIONS_PATH}/{translation_ids[0]}"),
+            status_extractor=lambda r: (r.get("data") or {}).get("status"),
+            queued_statuses=["pending"],
+            poll_interval=5.0,
+        )
+        data = final["data"]
+        if not data.get("video_url"):
+            raise ValueError(f"HeyGen returned no video_url for translation {translation_ids[0]}.")
+        return IO.NodeOutput(await download_url_to_video_output(data["video_url"]))
+
+
+class HeyGenTextToSpeechNode(IO.ComfyNode):
+    """Synthesize speech audio from text with HeyGen's Starfish TTS engine."""
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="HeyGenTextToSpeechNode",
+            display_name="HeyGen Text to Speech",
+            category="partner/audio/HeyGen",
+            description="Generate speech audio from text using HeyGen's Starfish TTS engine. "
+            "Includes HeyGen's most popular voices across 17 languages.",
+            inputs=[
+                IO.String.Input(
+                    "text",
+                    multiline=True,
+                    default="",
+                    tooltip="Text to synthesize (up to 5000 characters). The generated speech "
+                    "must be at least 1 second long.",
+                ),
+                IO.Combo.Input(
+                    "voice",
+                    options=HEYGEN_VOICE_TTS_OPTIONS,
+                    tooltip="Voice to use (curated from HeyGen's most popular Starfish-compatible voices).",
+                ),
+                IO.String.Input(
+                    "custom_voice_id",
+                    default="",
+                    optional=True,
+                    tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
+                    "The voice must support the Starfish engine.",
+                ),
+                IO.Float.Input(
+                    "speed",
+                    default=1.0,
+                    min=0.5,
+                    max=2.0,
+                    step=0.05,
+                    optional=True,
+                    tooltip="Speech speed multiplier.",
+                ),
+                IO.Boolean.Input(
+                    "ssml",
+                    default=False,
+                    optional=True,
+                    tooltip="Treat the text as SSML markup (for pauses, emphasis, and pronunciation control).",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    optional=True,
+                    tooltip="Not sent to HeyGen; change it to force a re-run.",
+                ),
+            ],
+            outputs=[IO.Audio.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.00095381,"format":{"approximate":true,"suffix":"/second"}}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        text: str,
+        voice: str,
+        custom_voice_id: str = "",
+        speed: float = 1.0,
+        ssml: bool = False,
+        seed: int = 0,
+    ) -> IO.NodeOutput:
+        validate_string(text, strip_whitespace=True, min_length=1, max_length=5000)
+        payload = {
+            "text": text,
+            "voice_id": custom_voice_id.strip() or HEYGEN_VOICE_TTS_MAP[voice],
+            "speed": round(speed, 2),
+        }
+        if ssml:
+            payload["input_type"] = "ssml"
+        response = await sync_op_raw(
+            cls,
+            ApiEndpoint(path=_SPEECH_PATH, method="POST"),
+            data=payload,
+        )
+        audio_url = (response.get("data") or {}).get("audio_url")
+        if not audio_url:
+            raise ValueError(f"HeyGen did not return an audio_url: {response}")
+        audio_bytes = await download_url_as_bytesio(audio_url)
+        return IO.NodeOutput(audio_bytes_to_audio_input(audio_bytes.getvalue()))
+
+
+class HeyGenExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            HeyGenTalkingPhotoNode,
+            HeyGenAvatarVideoNode,
+            HeyGenCreateAvatarNode,
+            HeyGenVideoTranslateNode,
+            HeyGenTextToSpeechNode,
+        ]
+
+
+async def comfy_entrypoint() -> HeyGenExtension:
+    return HeyGenExtension()

comfy_api_nodes/nodes_hunyuan3d.py

@@ -642,6 +642,7 @@ class Tencent3DPartNode(IO.ComfyNode):
             response_model=To3DProTaskCreateResponse,
             data=To3DPartTaskRequest(
                 File=TaskFile3DInput(Type=file_format.upper(), Url=model_url),
+                EnableStagedGeneration=True,
             ),
             is_rate_limited=_is_tencent_rate_limited,
         )

comfy_api_nodes/nodes_openai.py

@@ -41,6 +41,9 @@ STARTING_POINT_ID_PATTERN = r"<starting_point_id:(.*)>"
 
 
 class SupportedOpenAIModel(str, Enum):
+    gpt_5_6_sol = "gpt-5.6-sol"
+    gpt_5_6_terra = "gpt-5.6-terra"
+    gpt_5_6_luna = "gpt-5.6-luna"
     gpt_5_5_pro = "gpt-5.5-pro"
     gpt_5_5 = "gpt-5.5"
     gpt_5 = "gpt-5"
@@ -361,19 +364,6 @@ class OpenAIDalle3(IO.ComfyNode):
         return IO.NodeOutput(await validate_and_cast_response(response))
 
 
-def calculate_tokens_price_image_1(response: OpenAIImageGenerationResponse) -> float | None:
-    # https://platform.openai.com/docs/pricing
-    return ((response.usage.input_tokens * 10.0) + (response.usage.output_tokens * 40.0)) / 1_000_000.0
-
-
-def calculate_tokens_price_image_1_5(response: OpenAIImageGenerationResponse) -> float | None:
-    return ((response.usage.input_tokens * 8.0) + (response.usage.output_tokens * 32.0)) / 1_000_000.0
-
-
-def calculate_tokens_price_image_2_0(response: OpenAIImageGenerationResponse) -> float | None:
-    return ((response.usage.input_tokens * 8.0) + (response.usage.output_tokens * 30.0)) / 1_000_000.0
-
-
 class OpenAIGPTImage1(IO.ComfyNode):
 
     @classmethod
@@ -567,15 +557,10 @@ class OpenAIGPTImage1(IO.ComfyNode):
             if size not in ("auto", "1024x1024", "1024x1536", "1536x1024"):
                 raise ValueError(f"Resolution {size} is only supported by GPT Image 2 model")
 
-        if model == "gpt-image-1":
-            price_extractor = calculate_tokens_price_image_1
-        elif model == "gpt-image-1.5":
-            price_extractor = calculate_tokens_price_image_1_5
-        elif model == "gpt-image-2":
-            price_extractor = calculate_tokens_price_image_2_0
+        if model == "gpt-image-2":
             if background == "transparent":
                 raise ValueError("Transparent background is not supported for GPT Image 2 model")
-        else:
+        elif model not in ("gpt-image-1", "gpt-image-1.5"):
             raise ValueError(f"Unknown model: {model}")
 
         if image is not None:
@@ -630,7 +615,6 @@ class OpenAIGPTImage1(IO.ComfyNode):
                 ),
                 content_type="multipart/form-data",
                 files=files,
-                price_extractor=price_extractor,
             )
         else:
             response = await sync_op(
@@ -647,7 +631,6 @@ class OpenAIGPTImage1(IO.ComfyNode):
                     size=size,
                     moderation="low",
                 ),
-                price_extractor=price_extractor,
             )
         return IO.NodeOutput(await validate_and_cast_response(response))
 
@@ -876,13 +859,7 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                 )
             size = f"{custom_width}x{custom_height}"
 
-        if model_id == "gpt-image-1":
-            price_extractor = calculate_tokens_price_image_1
-        elif model_id == "gpt-image-1.5":
-            price_extractor = calculate_tokens_price_image_1_5
-        elif model_id == "gpt-image-2":
-            price_extractor = calculate_tokens_price_image_2_0
-        else:
+        if model_id not in ("gpt-image-1", "gpt-image-1.5", "gpt-image-2"):
             raise ValueError(f"Unknown model: {model_id}")
 
         if image_tensors:
@@ -941,7 +918,6 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                 ),
                 content_type="multipart/form-data",
                 files=files,
-                price_extractor=price_extractor,
             )
         else:
             response = await sync_op(
@@ -957,7 +933,6 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                     size=size,
                     moderation="low",
                 ),
-                price_extractor=price_extractor,
             )
         return IO.NodeOutput(await validate_and_cast_response(response))
 
@@ -1063,6 +1038,21 @@ class OpenAIChatNode(IO.ComfyNode):
                     "usd": [0.002, 0.008],
                     "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
                   }
+                  : $contains($m, "gpt-5.6-terra") ? {
+                    "type": "list_usd",
+                    "usd": [0.0025, 0.015],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : $contains($m, "gpt-5.6-luna") ? {
+                    "type": "list_usd",
+                    "usd": [0.001, 0.006],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
+                  : $contains($m, "gpt-5.6") ? {
+                    "type": "list_usd",
+                    "usd": [0.005, 0.03],
+                    "format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
+                  }
                   : $contains($m, "gpt-5.5-pro") ? {
                     "type": "list_usd",
                     "usd": [0.03, 0.18],

comfy_api_nodes/nodes_openrouter.py

@@ -45,27 +45,40 @@ class _ModelSpec:
 
 
 MODELS: list[_ModelSpec] = [
-    _ModelSpec("anthropic/claude-opus-4.7", "frontier_reasoning", 0.000005, 0.000025, max_images=20),
-    _ModelSpec("openai/gpt-5.5-pro", "frontier_reasoning", 0.00003, 0.00018, max_images=20),
-    _ModelSpec("openai/gpt-5.5", "frontier_reasoning", 0.000005, 0.00003, max_images=20),
-    _ModelSpec("google/gemini-3.5-flash", "reasoning", 0.0000015, 0.000009, max_images=20, max_videos=4),
-    _ModelSpec("x-ai/grok-4.20", "reasoning", 0.00000125, 0.0000025, max_images=20),
-    _ModelSpec("x-ai/grok-4.3", "reasoning", 0.00000125, 0.0000025, max_images=20),
-    _ModelSpec("deepseek/deepseek-v4-pro", "reasoning", 0.000000435, 0.00000087),
-    _ModelSpec("deepseek/deepseek-v4-flash", "reasoning", 0.000000112, 0.000000224),
-    _ModelSpec("deepseek/deepseek-v3.2", "reasoning", 0.000000252, 0.000000378),
-    _ModelSpec("qwen/qwen3.6-max-preview", "reasoning", 0.00000104, 0.00000624),
-    _ModelSpec("qwen/qwen3.6-plus", "reasoning", 0.000000325, 0.00000195, max_images=10, max_videos=4),
-    _ModelSpec("qwen/qwen3.6-flash", "reasoning", 0.0000001875, 0.000001125, max_images=10, max_videos=4),
-    _ModelSpec("mistralai/mistral-large-2512", "standard", 0.0000005, 0.0000015, max_images=8),
-    _ModelSpec("mistralai/mistral-medium-3-5", "reasoning", 0.0000015, 0.0000075, max_images=8),
-    _ModelSpec("z-ai/glm-4.6", "reasoning", 0.00000043, 0.00000174),
-    _ModelSpec("z-ai/glm-5", "reasoning", 0.0000006, 0.00000192),
-    _ModelSpec("moonshotai/kimi-k2.6", "reasoning", 0.00000073, 0.00000349, max_images=10),
-    _ModelSpec("moonshotai/kimi-k2-thinking", "reasoning", 0.0000006, 0.0000025),
-    _ModelSpec("perplexity/sonar-pro", "perplexity", 0.000003, 0.000015),
-    _ModelSpec("perplexity/sonar-reasoning-pro", "perplexity_reasoning", 0.000002, 0.000008),
-    _ModelSpec("perplexity/sonar-deep-research", "perplexity_reasoning", 0.000002, 0.000008),
+    _ModelSpec("anthropic/claude-opus-5", "frontier_reasoning", 0.00000715, 0.00003575, max_images=20),
+    _ModelSpec("anthropic/claude-opus-4.8", "frontier_reasoning", 0.00000715, 0.00003575, max_images=20),
+    _ModelSpec("anthropic/claude-opus-4.7", "frontier_reasoning", 0.00000715, 0.00003575, max_images=20),
+    _ModelSpec("anthropic/claude-fable-5", "frontier_reasoning", 0.0000143, 0.0000715, max_images=20),
+    _ModelSpec("anthropic/claude-sonnet-5", "frontier_reasoning", 0.00000286, 0.0000143, max_images=20),
+    _ModelSpec("anthropic/claude-haiku-4.5", "frontier_reasoning", 0.00000143, 0.00000715, max_images=20),
+    _ModelSpec("openai/gpt-5.6-sol-pro", "frontier_reasoning", 0.00000715, 0.0000429, max_images=20),
+    _ModelSpec("openai/gpt-5.6-sol", "frontier_reasoning", 0.00000715, 0.0000429, max_images=20),
+    _ModelSpec("openai/gpt-5.6-terra-pro", "frontier_reasoning", 0.000003575, 0.00002145, max_images=20),
+    _ModelSpec("openai/gpt-5.6-terra", "frontier_reasoning", 0.000003575, 0.00002145, max_images=20),
+    _ModelSpec("openai/gpt-5.6-luna-pro", "frontier_reasoning", 0.00000143, 0.00000858, max_images=20),
+    _ModelSpec("openai/gpt-5.6-luna", "frontier_reasoning", 0.00000143, 0.00000858, max_images=20),
+    _ModelSpec("openai/gpt-5.5-pro", "frontier_reasoning", 0.0000429, 0.0002574, max_images=20),
+    _ModelSpec("openai/gpt-5.5", "frontier_reasoning", 0.00000715, 0.0000429, max_images=20),
+    _ModelSpec("google/gemini-3.5-flash", "reasoning", 0.000002145, 0.00001287, max_images=20, max_videos=4),
+    _ModelSpec("x-ai/grok-4.5", "reasoning", 0.00000286, 0.00000858, max_images=20),
+    _ModelSpec("x-ai/grok-4.20", "reasoning", 0.0000017875, 0.000003575, max_images=20),
+    _ModelSpec("x-ai/grok-4.3", "reasoning", 0.0000017875, 0.000003575, max_images=20),
+    _ModelSpec("deepseek/deepseek-v4-pro", "reasoning", 0.00000062205, 0.0000012441),
+    _ModelSpec("deepseek/deepseek-v4-flash", "reasoning", 0.00000016016, 0.00000032032),
+    _ModelSpec("deepseek/deepseek-v3.2", "reasoning", 0.00000036036, 0.00000054054),
+    _ModelSpec("qwen/qwen3.6-max-preview", "reasoning", 0.0000014872, 0.0000089232),
+    _ModelSpec("qwen/qwen3.6-plus", "reasoning", 0.00000046475, 0.0000027885, max_images=10, max_videos=4),
+    _ModelSpec("qwen/qwen3.6-flash", "reasoning", 0.000000268125, 0.00000160875, max_images=10, max_videos=4),
+    _ModelSpec("mistralai/mistral-large-2512", "standard", 0.000000715, 0.000002145, max_images=8),
+    _ModelSpec("mistralai/mistral-medium-3-5", "reasoning", 0.000002145, 0.000010725, max_images=8),
+    _ModelSpec("z-ai/glm-4.6", "reasoning", 0.0000006149, 0.0000024882),
+    _ModelSpec("z-ai/glm-5", "reasoning", 0.000000858, 0.0000027456),
+    _ModelSpec("moonshotai/kimi-k3", "reasoning", 0.00000429, 0.00002145, max_images=10),
+    _ModelSpec("moonshotai/kimi-k2.6", "reasoning", 0.0000010439, 0.0000049907, max_images=10),
+    _ModelSpec("moonshotai/kimi-k2-thinking", "reasoning", 0.000000858, 0.000003575),
+    _ModelSpec("perplexity/sonar-pro", "perplexity", 0.00000429, 0.00002145),
+    _ModelSpec("perplexity/sonar-reasoning-pro", "perplexity_reasoning", 0.00000286, 0.00001144),
+    _ModelSpec("perplexity/sonar-deep-research", "perplexity_reasoning", 0.00000286, 0.00001144),
 ]
 
 _MODELS_BY_SLUG: dict[str, _ModelSpec] = {m.slug: m for m in MODELS}
@@ -146,12 +159,6 @@ def _build_model_options() -> list[IO.DynamicCombo.Option]:
     return [IO.DynamicCombo.Option(spec.slug, _inputs_for_model(spec)) for spec in MODELS]
 
 
-def _calculate_price(response: OpenRouterChatResponse) -> float | None:
-    if response.usage and response.usage.cost is not None:
-        return float(response.usage.cost)
-    return None
-
-
 def _price_badge_jsonata() -> str:
     rates_pairs = []
     for spec in MODELS:
@@ -269,8 +276,8 @@ class OpenRouterLLMNode(IO.ComfyNode):
             essentials_category="Text Generation",
             description=(
                 "Generate text responses through OpenRouter. Routes to a curated set of popular "
-                "models from xAI, DeepSeek, Qwen, Mistral, Z.AI (GLM), Moonshot (Kimi), and "
-                "Perplexity Sonar."
+                "models from Anthropic (Claude), OpenAI (GPT), Google (Gemini), xAI (Grok), "
+                "DeepSeek, Qwen, Mistral, Z.AI (GLM), Moonshot (Kimi), and Perplexity Sonar."
             ),
             inputs=[
                 IO.String.Input(
@@ -359,7 +366,6 @@ class OpenRouterLLMNode(IO.ComfyNode):
             ApiEndpoint(path=OPENROUTER_CHAT_ENDPOINT, method="POST"),
             response_model=OpenRouterChatResponse,
             data=request,
-            price_extractor=_calculate_price,
         )
         return IO.NodeOutput(_extract_text(response))
 

comfy_api_nodes/nodes_recraft.py

@@ -399,7 +399,7 @@ class RecraftTextToImageNode(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="RecraftTextToImageNode",
-            display_name="Recraft Text to Image",
+            display_name="Recraft V3 Text to Image",
             category="partner/image/Recraft",
             description="Generates images synchronously based on prompt and resolution.",
             inputs=[
@@ -511,7 +511,7 @@ class RecraftImageToImageNode(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="RecraftImageToImageNode",
-            display_name="Recraft Image to Image",
+            display_name="Recraft V3 Image to Image",
             category="partner/image/Recraft",
             description="Modify image based on prompt and strength.",
             inputs=[
@@ -731,7 +731,7 @@ class RecraftTextToVectorNode(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="RecraftTextToVectorNode",
-            display_name="Recraft Text to Vector",
+            display_name="Recraft V3 Text to Vector",
             category="partner/image/Recraft",
             description="Generates SVG synchronously based on prompt and resolution.",
             inputs=[
@@ -1087,7 +1087,7 @@ class RecraftV4TextToImageNode(IO.ComfyNode):
             node_id="RecraftV4TextToImageNode",
             display_name="Recraft V4 Text to Image",
             category="partner/image/Recraft",
-            description="Generates images using Recraft V4 or V4 Pro models.",
+            description="Generates images using Recraft V4 and V4.1 models.",
             inputs=[
                 IO.String.Input(
                     "prompt",
@@ -1097,11 +1097,56 @@ class RecraftV4TextToImageNode(IO.ComfyNode):
                 IO.String.Input(
                     "negative_prompt",
                     multiline=True,
-                    tooltip="An optional text description of undesired elements on an image.",
+                    tooltip="This input is ignored: negative prompt is not supported by "
+                    "Recraft V4 and V4.1 models.",
                 ),
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_SIZES,
+                                    default="1024x1024",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_utility",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_SIZES,
+                                    default="1024x1024",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_pro",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_PRO_SIZES,
+                                    default="2048x2048",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_utility_pro",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_PRO_SIZES,
+                                    default="2048x2048",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
                         IO.DynamicCombo.Option(
                             "recraftv4",
                             [
@@ -1162,7 +1207,14 @@ class RecraftV4TextToImageNode(IO.ComfyNode):
                 depends_on=IO.PriceBadgeDepends(widgets=["model", "n"]),
                 expr="""
                 (
-                    $prices := {"recraftv4": 0.04, "recraftv4_pro": 0.25};
+                    $prices := {
+                        "recraftv4_1": 0.035,
+                        "recraftv4_1_utility": 0.035,
+                        "recraftv4_1_pro": 0.21,
+                        "recraftv4_1_utility_pro": 0.21,
+                        "recraftv4": 0.04,
+                        "recraftv4_pro": 0.25
+                    };
                     {"type":"usd","usd": $lookup($prices, widgets.model) * widgets.n}
                 )
                 """,
@@ -1179,14 +1231,13 @@ class RecraftV4TextToImageNode(IO.ComfyNode):
         seed: int,
         recraft_controls: RecraftControls | None = None,
     ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=False, min_length=1, max_length=10000)
+        validate_string(prompt, strip_whitespace=True, min_length=1, max_length=10000)
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/recraft/image_generation", method="POST"),
             response_model=RecraftImageGenerationResponse,
             data=RecraftImageGenerationRequest(
                 prompt=prompt,
-                negative_prompt=negative_prompt if negative_prompt else None,
                 model=model["model"],
                 size=model["size"],
                 n=n,
@@ -1211,7 +1262,7 @@ class RecraftV4TextToVectorNode(IO.ComfyNode):
             node_id="RecraftV4TextToVectorNode",
             display_name="Recraft V4 Text to Vector",
             category="partner/image/Recraft",
-            description="Generates SVG using Recraft V4 or V4 Pro models.",
+            description="Generates SVG using Recraft V4 and V4.1 models.",
             inputs=[
                 IO.String.Input(
                     "prompt",
@@ -1221,11 +1272,56 @@ class RecraftV4TextToVectorNode(IO.ComfyNode):
                 IO.String.Input(
                     "negative_prompt",
                     multiline=True,
-                    tooltip="An optional text description of undesired elements on an image.",
+                    tooltip="This input is ignored: negative prompt is not supported by "
+                    "Recraft V4 and V4.1 models.",
                 ),
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_vector",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_SIZES,
+                                    default="1024x1024",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_utility_vector",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_SIZES,
+                                    default="1024x1024",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_pro_vector",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_PRO_SIZES,
+                                    default="2048x2048",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "recraftv4_1_utility_pro_vector",
+                            [
+                                IO.Combo.Input(
+                                    "size",
+                                    options=RECRAFT_V4_PRO_SIZES,
+                                    default="2048x2048",
+                                    tooltip="The size of the generated image.",
+                                ),
+                            ],
+                        ),
                         IO.DynamicCombo.Option(
                             "recraftv4",
                             [
@@ -1286,7 +1382,14 @@ class RecraftV4TextToVectorNode(IO.ComfyNode):
                 depends_on=IO.PriceBadgeDepends(widgets=["model", "n"]),
                 expr="""
                 (
-                    $prices := {"recraftv4": 0.08, "recraftv4_pro": 0.30};
+                    $prices := {
+                        "recraftv4_1_vector": 0.08,
+                        "recraftv4_1_utility_vector": 0.08,
+                        "recraftv4_1_pro_vector": 0.30,
+                        "recraftv4_1_utility_pro_vector": 0.30,
+                        "recraftv4": 0.08,
+                        "recraftv4_pro": 0.30
+                    };
                     {"type":"usd","usd": $lookup($prices, widgets.model) * widgets.n}
                 )
                 """,
@@ -1303,18 +1406,17 @@ class RecraftV4TextToVectorNode(IO.ComfyNode):
         seed: int,
         recraft_controls: RecraftControls | None = None,
     ) -> IO.NodeOutput:
-        validate_string(prompt, strip_whitespace=False, min_length=1, max_length=10000)
+        validate_string(prompt, strip_whitespace=True, min_length=1, max_length=10000)
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/recraft/image_generation", method="POST"),
             response_model=RecraftImageGenerationResponse,
             data=RecraftImageGenerationRequest(
                 prompt=prompt,
-                negative_prompt=negative_prompt if negative_prompt else None,
                 model=model["model"],
                 size=model["size"],
                 n=n,
-                style="vector_illustration",
+                style=None if model["model"].endswith("_vector") else "vector_illustration",
                 substyle=None,
                 controls=recraft_controls.create_api_model() if recraft_controls else None,
             ),

comfy_api_nodes/nodes_reve.py

@@ -62,13 +62,6 @@ def _postprocessing_inputs():
     ]
 
 
-def _reve_price_extractor(headers: dict) -> float | None:
-    credits_used = headers.get("x-reve-credits-used")
-    if credits_used is not None:
-        return float(credits_used) / 524.48
-    return None
-
-
 def _reve_response_header_validator(headers: dict) -> None:
     error_code = headers.get("x-reve-error-code")
     if error_code:
@@ -180,7 +173,6 @@ class ReveImageCreateNode(IO.ComfyNode):
                 headers={"Accept": "image/webp"},
             ),
             as_binary=True,
-            price_extractor=_reve_price_extractor,
             response_header_validator=_reve_response_header_validator,
             data=ReveImageCreateRequest(
                 prompt=prompt,
@@ -279,7 +271,6 @@ class ReveImageEditNode(IO.ComfyNode):
                 headers={"Accept": "image/webp"},
             ),
             as_binary=True,
-            price_extractor=_reve_price_extractor,
             response_header_validator=_reve_response_header_validator,
             data=ReveImageEditRequest(
                 edit_instruction=edit_instruction,
@@ -396,7 +387,6 @@ class ReveImageRemixNode(IO.ComfyNode):
                 headers={"Accept": "image/webp"},
             ),
             as_binary=True,
-            price_extractor=_reve_price_extractor,
             response_header_validator=_reve_response_header_validator,
             data=ReveImageRemixRequest(
                 prompt=prompt,

comfy_api_nodes/nodes_runway.py

@@ -194,6 +194,7 @@ class RunwayImageToVideoNodeGen3a(IO.ComfyNode):
                 depends_on=IO.PriceBadgeDepends(widgets=["duration"]),
                 expr="""{"type":"usd","usd": 0.0715 * widgets.duration}""",
             ),
+            is_deprecated=True,
         )
 
     @classmethod
@@ -390,6 +391,7 @@ class RunwayFirstLastFrameNode(IO.ComfyNode):
                 depends_on=IO.PriceBadgeDepends(widgets=["duration"]),
                 expr="""{"type":"usd","usd": 0.0715 * widgets.duration}""",
             ),
+            is_deprecated=True,
         )
 
     @classmethod

comfy_api_nodes/nodes_sync_so.py

@@ -0,0 +1,391 @@
+from typing_extensions import override
+
+from comfy_api.latest import IO, ComfyExtension, Input
+from comfy_api_nodes.apis.sync_so import (
+    SyncActiveSpeakerDetection,
+    SyncGeneration,
+    SyncGenerationOptions,
+    SyncGenerationRequest,
+    SyncInputItem,
+)
+from comfy_api_nodes.util import (
+    ApiEndpoint,
+    download_url_to_video_output,
+    downscale_image_tensor,
+    downscale_image_tensor_by_max_side,
+    get_image_dimensions,
+    get_number_of_images,
+    poll_op,
+    sync_op,
+    upload_audio_to_comfyapi,
+    upload_image_to_comfyapi,
+    upload_video_to_comfyapi,
+    validate_audio_duration,
+)
+
+
+class SyncLipSyncNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="SyncLipSyncNode",
+            display_name="sync.so Lip Sync",
+            category="partner/video/sync.so",
+            description=(
+                "Re-sync mouth movement in a video to new speech audio using sync.so. "
+                "Handles close-ups, profiles and obstructions automatically while preserving "
+                "the speaker's expression. Cost scales with output duration."
+            ),
+            inputs=[
+                IO.Video.Input(
+                    "video",
+                    tooltip="Footage of the speaker to re-sync. Up to 4K (4096x2160); "
+                    "a constant frame rate of 24/25/30 fps works best.",
+                ),
+                IO.Audio.Input(
+                    "audio",
+                    tooltip="Speech audio to sync the mouth to.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    tooltip="Seed controls whether the node should re-run; "
+                    "results are non-deterministic regardless of seed.",
+                ),
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "sync-3",
+                            [
+                                IO.Combo.Input(
+                                    "sync_mode",
+                                    options=["bounce", "cut_off", "loop", "silence", "remap"],
+                                    default="bounce",
+                                    tooltip=(
+                                        "How to handle a duration mismatch between video and audio; "
+                                        "this also sets the output length. "
+                                        "bounce: video plays forward then backward until the audio ends "
+                                        "(output = audio length). "
+                                        "loop: video restarts until the audio ends (output = audio length). "
+                                        "remap: video is time-stretched to match the audio (output = audio length). "
+                                        "cut_off: the longer track is trimmed (output = shorter length). "
+                                        "silence: nothing is trimmed; the shorter track is padded "
+                                        "(output = longer length)."
+                                    ),
+                                ),
+                                IO.Combo.Input(
+                                    "speaker_selection",
+                                    options=["default", "auto-detect", "coordinates"],
+                                    default="default",
+                                    tooltip=(
+                                        "Which face to lipsync when several people are visible. "
+                                        "default: let the model decide. "
+                                        "auto-detect: detect and follow the active speaker. "
+                                        "coordinates: target the face at pixel (speaker_x, speaker_y) "
+                                        "in the frame chosen by speaker_frame."
+                                    ),
+                                ),
+                                IO.Int.Input(
+                                    "speaker_frame",
+                                    default=0,
+                                    min=0,
+                                    max=1_000_000,
+                                    advanced=True,
+                                    tooltip="Video frame used to locate the speaker. "
+                                    "Only used when speaker_selection is 'coordinates'.",
+                                ),
+                                IO.Int.Input(
+                                    "speaker_x",
+                                    default=0,
+                                    min=0,
+                                    max=4096,
+                                    advanced=True,
+                                    tooltip="X pixel coordinate of the speaker's face. "
+                                    "Only used when speaker_selection is 'coordinates'.",
+                                ),
+                                IO.Int.Input(
+                                    "speaker_y",
+                                    default=0,
+                                    min=0,
+                                    max=4096,
+                                    advanced=True,
+                                    tooltip="Y pixel coordinate of the speaker's face. "
+                                    "Only used when speaker_selection is 'coordinates'.",
+                                ),
+                            ],
+                        )
+                    ],
+                    tooltip="sync.so generation model.",
+                ),
+            ],
+            outputs=[IO.Video.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        video: Input.Video,
+        audio: Input.Audio,
+        seed: int,
+        model: dict,
+    ) -> IO.NodeOutput:
+        try:
+            width, height = video.get_dimensions()
+        except Exception:
+            width = height = None
+        if width and height and (max(width, height) > 4096 or width * height > 4096 * 2160):
+            raise ValueError(
+                f"sync.so rejects videos above 4K (4096x2160); got {width}x{height}. Downscale the video first."
+            )
+        validate_audio_duration(audio, max_duration=600)
+
+        if model["speaker_selection"] == "auto-detect":
+            speaker_detection = SyncActiveSpeakerDetection(auto_detect=True)
+        elif model["speaker_selection"] == "coordinates":
+            speaker_detection = SyncActiveSpeakerDetection(
+                frame_number=model["speaker_frame"],
+                coordinates=[model["speaker_x"], model["speaker_y"]],
+            )
+        else:
+            speaker_detection = None
+
+        video_url = await upload_video_to_comfyapi(cls, video, max_duration=600)
+        audio_url = await upload_audio_to_comfyapi(cls, audio)
+
+        generation = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
+            response_model=SyncGeneration,
+            data=SyncGenerationRequest(
+                model=model["model"],
+                input=[
+                    SyncInputItem(type="video", url=video_url),
+                    SyncInputItem(type="audio", url=audio_url),
+                ],
+                options=SyncGenerationOptions(
+                    sync_mode=model["sync_mode"],
+                    active_speaker_detection=speaker_detection,
+                ),
+            ),
+        )
+        generation = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
+            response_model=SyncGeneration,
+            status_extractor=lambda g: g.status,
+            completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
+            failed_statuses=[],
+            queued_statuses=["PENDING"],
+            poll_interval=10.0,
+        )
+        if generation.status != "COMPLETED":
+            code = f" [{generation.errorCode}]" if generation.errorCode else ""
+            raise ValueError(
+                f"sync.so generation {generation.status.lower()}{code}: "
+                f"{generation.error or 'no error details provided'}"
+            )
+        if not generation.outputUrl:
+            raise ValueError("sync.so generation completed but no output URL was returned.")
+        return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
+
+
+class SyncTalkingImageNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="SyncTalkingImageNode",
+            display_name="sync.so Talking Image",
+            category="partner/video/sync.so",
+            description=(
+                "Animate a still portrait into a talking video driven by speech audio, "
+                "using sync.so's sync-3 model. The output duration matches the audio. "
+                "Cost scales with output duration."
+            ),
+            inputs=[
+                IO.Image.Input(
+                    "image",
+                    tooltip="A single image with a clearly visible face, up to 4K (4096x2160).",
+                ),
+                IO.Audio.Input(
+                    "audio",
+                    tooltip="Speech audio driving the talking video; the output duration matches it. "
+                    "Chain any TTS node here to drive the animation from text.",
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Optional guidance for how the portrait comes to life, e.g. "
+                    "'make the subject smile and look at the camera'. "
+                    "Leave empty for natural talking motion.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=2147483647,
+                    control_after_generate=True,
+                    tooltip="Seed controls whether the node should re-run; "
+                    "results are non-deterministic regardless of seed.",
+                ),
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "sync-3",
+                            [
+                                IO.Combo.Input(
+                                    "speaker_selection",
+                                    options=["default", "coordinates"],
+                                    default="default",
+                                    tooltip=(
+                                        "Which face to animate when several people are visible. "
+                                        "default: let the model decide. "
+                                        "coordinates: target the face at pixel (speaker_x, speaker_y) "
+                                        "in the image. Auto-detection is not supported for images."
+                                    ),
+                                ),
+                                IO.Int.Input(
+                                    "speaker_x",
+                                    default=0,
+                                    min=0,
+                                    max=4096,
+                                    advanced=True,
+                                    tooltip="X pixel coordinate of the speaker's face. "
+                                    "Only used when speaker_selection is 'coordinates'.",
+                                ),
+                                IO.Int.Input(
+                                    "speaker_y",
+                                    default=0,
+                                    min=0,
+                                    max=4096,
+                                    advanced=True,
+                                    tooltip="Y pixel coordinate of the speaker's face. "
+                                    "Only used when speaker_selection is 'coordinates'.",
+                                ),
+                                IO.Boolean.Input(
+                                    "auto_downscale",
+                                    default=True,
+                                    advanced=True,
+                                    tooltip="Automatically downscale the image if it exceeds the 4K "
+                                    "(4096x2160) input limit; speaker coordinates are scaled to match. "
+                                    "When disabled, an oversized image raises an error instead.",
+                                ),
+                            ],
+                        )
+                    ],
+                    tooltip="sync.so generation model. Image input is exclusive to sync-3.",
+                ),
+            ],
+            outputs=[IO.Video.Output()],
+            hidden=[
+                IO.Hidden.auth_token_comfy_org,
+                IO.Hidden.api_key_comfy_org,
+                IO.Hidden.unique_id,
+            ],
+            is_api_node=True,
+            price_badge=IO.PriceBadge(
+                expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        audio: Input.Audio,
+        prompt: str,
+        seed: int,
+        model: dict,
+    ) -> IO.NodeOutput:
+        if get_number_of_images(image) != 1:
+            raise ValueError("Exactly one image is required; got a batch. Pick one frame first.")
+        validate_audio_duration(audio, max_duration=600)
+
+        height, width = get_image_dimensions(image)
+        speaker_x, speaker_y = model["speaker_x"], model["speaker_y"]
+        if max(width, height) > 4096 or width * height > 4096 * 2160:
+            if not model["auto_downscale"]:
+                raise ValueError(
+                    f"sync.so rejects images above 4K (4096x2160); got {width}x{height}. "
+                    "Downscale the image first or enable auto_downscale."
+                )
+            image = downscale_image_tensor(image, total_pixels=4096 * 2160)
+            image = downscale_image_tensor_by_max_side(image, max_side=4096)
+            new_height, new_width = get_image_dimensions(image)
+            # speaker coordinates are given in the original image's pixel space
+            speaker_x = min(new_width - 1, round(speaker_x * new_width / width))
+            speaker_y = min(new_height - 1, round(speaker_y * new_height / height))
+
+        if model["speaker_selection"] == "coordinates":
+            speaker_detection = SyncActiveSpeakerDetection(
+                frame_number=0,  # images have a single frame; auto_detect is rejected by the API
+                coordinates=[speaker_x, speaker_y],
+            )
+        else:
+            speaker_detection = None
+
+        image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
+        audio_url = await upload_audio_to_comfyapi(cls, audio)
+
+        generation = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
+            response_model=SyncGeneration,
+            data=SyncGenerationRequest(
+                model=model["model"],
+                input=[
+                    SyncInputItem(type="image", url=image_url),
+                    SyncInputItem(type="audio", url=audio_url),
+                ],
+                options=SyncGenerationOptions(
+                    i2v_prompt=prompt.strip() or None,
+                    active_speaker_detection=speaker_detection,
+                ),
+            ),
+        )
+        generation = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
+            response_model=SyncGeneration,
+            status_extractor=lambda g: g.status,
+            completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
+            failed_statuses=[],
+            queued_statuses=["PENDING"],
+            poll_interval=10.0,
+        )
+        if generation.status != "COMPLETED":
+            code = f" [{generation.errorCode}]" if generation.errorCode else ""
+            raise ValueError(
+                f"sync.so generation {generation.status.lower()}{code}: "
+                f"{generation.error or 'no error details provided'}"
+            )
+        if not generation.outputUrl:
+            raise ValueError("sync.so generation completed but no output URL was returned.")
+        return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
+
+
+class SyncExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            SyncLipSyncNode,
+            SyncTalkingImageNode,
+        ]
+
+
+async def comfy_entrypoint() -> SyncExtension:
+    return SyncExtension()

comfy_api_nodes/util/_helpers.py

@@ -11,9 +11,12 @@ from io import BytesIO
 from yarl import URL
 
 from comfy.cli_args import args
+from comfy.comfy_api_env import normalize_comfy_api_base
 from comfy.deploy_environment import get_deploy_environment
 from comfy.model_management import processing_interrupted
 from comfy_api.latest import IO
+from comfy_execution.utils import get_executing_context
+from comfyui_version import __version__ as comfyui_version
 
 from .common_exceptions import ProcessingInterrupted
 
@@ -55,15 +58,20 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
     relative/cloud URLs resolved against ``default_base_url()``; because the result
     includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
     """
-    return {
+    headers = {
         **get_auth_header(node_cls),
         "Comfy-Env": get_deploy_environment(),
         "Comfy-Usage-Source": get_usage_source(node_cls),
+        "Comfy-Core-Version": comfyui_version,
     }
+    ctx = get_executing_context()
+    if ctx is not None:
+        headers["Comfy-Job-Id"] = ctx.prompt_id
+    return headers
 
 
 def default_base_url() -> str:
-    return getattr(args, "comfy_api_base", "https://api.comfy.org")
+    return normalize_comfy_api_base(getattr(args, "comfy_api_base", "https://api.comfy.org"))
 
 
 async def sleep_with_interrupt(

comfy_api_nodes/util/client.py

@@ -2,8 +2,10 @@ import asyncio
 import contextlib
 import json
 import logging
+import math
 import time
 import uuid
+import weakref
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from enum import Enum
@@ -84,11 +86,37 @@ class _PollUIState:
 
 _RETRY_STATUS = {408, 500, 502, 503, 504}  # status 429 is handled separately
 _MAX_RETRY_AFTER_WAIT = 150.0  # Cap a server Retry-After at this many seconds so a large hint can't block execution
+
+PRICE_CREDITS_HEADER = "X-Comfy-Credits-Used"
+"""Proxy response header with the actual cost in Comfy credits. When present on any successful proxied response,
+it takes precedence over ``price_extractor``."""
+
+_credits_used_by_execution: "weakref.WeakKeyDictionary[type, float]" = weakref.WeakKeyDictionary()
+"""Last PRICE_CREDITS_HEADER value per node execution, keyed by the node's per-execution class clone."""
 COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed", "finished", "done", "complete"]
 FAILED_STATUSES = ["cancelled", "canceled", "canceling", "fail", "failed", "error"]
 QUEUED_STATUSES = ["created", "queued", "queueing", "submitted", "initializing", "wait", "in_queue"]
 
 
+def _maybe_remember_credits_used(node_cls: type[IO.ComfyNode], header_value: str | None) -> None:
+    """Remember a PRICE_CREDITS_HEADER value from a successful proxied response."""
+    if not header_value:
+        return
+    try:
+        credits_used = float(header_value)
+    except (TypeError, ValueError):
+        logging.debug("Ignoring malformed %s header: %r", PRICE_CREDITS_HEADER, header_value)
+        return
+    if not math.isfinite(credits_used) or credits_used < 0:
+        logging.debug("Ignoring out-of-range %s header: %r", PRICE_CREDITS_HEADER, header_value)
+        return
+    _credits_used_by_execution[node_cls] = credits_used + 0.0  # normalize -0.0
+
+
+def _get_remembered_credits_used(node_cls: type[IO.ComfyNode]) -> float | None:
+    return _credits_used_by_execution.get(node_cls)
+
+
 async def sync_op(
     cls: type[IO.ComfyNode],
     endpoint: ApiEndpoint,
@@ -450,10 +478,15 @@ def _display_text(
     display_lines: list[str] = []
     if status:
         display_lines.append(f"Status: {status.capitalize() if isinstance(status, str) else status}")
-    if price is not None:
+    server_credits = _get_remembered_credits_used(node_cls)
+    if server_credits is not None:
+        p = f"{server_credits:,.2f}".rstrip("0").rstrip(".")
+    elif price is not None:
         p = f"{float(price) * 211:,.1f}".rstrip("0").rstrip(".")
-        if p != "0":
-            display_lines.append(f"Price: {p} credits")
+    else:
+        p = None
+    if p is not None and p != "0":
+        display_lines.append(f"Price: {p} credits")
     if text is not None:
         display_lines.append(text)
     if display_lines:
@@ -606,7 +639,8 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
     """Core request with retries, per-second interruption monitoring, true cancellation, and friendly errors."""
     url = cfg.endpoint.path
     parsed_url = urlparse(url)
-    if not parsed_url.scheme and not parsed_url.netloc:  # is URL relative?
+    is_comfy_api_request = not parsed_url.scheme and not parsed_url.netloc  # is URL relative?
+    if is_comfy_api_request:
         url = urljoin(default_base_url().rstrip("/") + "/", url.lstrip("/"))
 
     method = cfg.endpoint.method
@@ -644,7 +678,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
         logging.debug("[DEBUG] HTTP %s %s (attempt %d)", method, url, attempt)
 
         payload_headers = {"Accept": "*/*"} if expect_binary else {"Accept": "application/json"}
-        if not parsed_url.scheme and not parsed_url.netloc:  # is URL relative?
+        if is_comfy_api_request:
             payload_headers.update(get_comfy_api_headers(cfg.node_cls))
         if cfg.endpoint.headers:
             payload_headers.update(cfg.endpoint.headers)
@@ -804,6 +838,8 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                                 )
                     bytes_payload = bytes(buff)
                     resp_headers = {k.lower(): v for k, v in resp.headers.items()}
+                    if is_comfy_api_request:
+                        _maybe_remember_credits_used(cfg.node_cls, resp.headers.get(PRICE_CREDITS_HEADER))
                     if cfg.price_extractor:
                         with contextlib.suppress(Exception):
                             extracted_price = cfg.price_extractor(resp_headers)
@@ -831,6 +867,8 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
                         except json.JSONDecodeError:
                             payload = {"_raw": text}
                         response_content_to_log = payload if isinstance(payload, dict) else text
+                    if is_comfy_api_request:
+                        _maybe_remember_credits_used(cfg.node_cls, resp.headers.get(PRICE_CREDITS_HEADER))
                     with contextlib.suppress(Exception):
                         extracted_price = cfg.price_extractor(payload) if cfg.price_extractor else None
                     operation_succeeded = True

comfy_execution/caching.py

@@ -5,7 +5,7 @@ import psutil
 import time
 import torch
 from typing import Sequence, Mapping, Dict
-from comfy.model_patcher import ModelPatcher
+from comfy.model_patcher import is_model_patcher_output
 from comfy_execution.graph import DynamicPrompt
 from abc import ABC, abstractmethod
 
@@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
 
 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
 
+RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
+
 
 def all_outputs_dynamic(outputs):
     if outputs is None:
@@ -517,12 +519,18 @@ def all_outputs_dynamic(outputs):
 
     return True
 
-
 class RAMPressureCache(LRUCache):
 
     def __init__(self, key_class, enable_providers=False):
         super().__init__(key_class, 0, enable_providers=enable_providers)
         self.timestamps = {}
+        self.active_evictions = False
+        self.full_evictions = False
+
+    async def set_prompt(self, dynprompt, node_ids, is_changed_cache):
+        self.active_evictions = False
+        self.full_evictions = False
+        await super().set_prompt(dynprompt, node_ids, is_changed_cache)
 
     def clean_unused(self):
         self._clean_subcaches()
@@ -539,9 +547,9 @@ class RAMPressureCache(LRUCache):
         self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
         super().set_local(node_id, value)
 
-    def ram_release(self, target, free_active=False):
+    def ram_release(self, target, free_active=False, min_entry_size=0):
         if psutil.virtual_memory().available >= target:
-            return
+            return 0
 
         clean_list = []
 
@@ -555,8 +563,9 @@ class RAMPressureCache(LRUCache):
             oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
 
             ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
+            oom_ram_usage = ram_usage
             def scan_list_for_ram_usage(outputs):
-                nonlocal ram_usage
+                nonlocal ram_usage, oom_ram_usage
                 if outputs is None:
                     return
                 for output in outputs:
@@ -564,19 +573,30 @@ class RAMPressureCache(LRUCache):
                         scan_list_for_ram_usage(output)
                     elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
                         ram_usage += output.numel() * output.element_size()
-                    elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
+                        oom_ram_usage += output.numel() * output.element_size()
+                    elif is_model_patcher_output(output) and self.used_generation[key] != self.generation:
                         #old ModelPatchers are the first to go
-                        ram_usage = 1e30
+                        oom_ram_usage = 1e30
             scan_list_for_ram_usage(cache_entry.outputs)
 
-            oom_score *= ram_usage
+            if ram_usage < min_entry_size:
+                continue
+
+            oom_score *= oom_ram_usage
             #In the case where we have no information on the node ram usage at all,
             #break OOM score ties on the last touch timestamp (pure LRU)
-            bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
+            bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage))
 
+        freed = 0
         while psutil.virtual_memory().available < target and clean_list:
-            _, _, key = clean_list.pop()
+            _, _, key, ram_usage = clean_list.pop()
             del self.cache[key]
             self.used_generation.pop(key, None)
             self.timestamps.pop(key, None)
             self.children.pop(key, None)
+            freed += ram_usage
+        if freed and free_active:
+            self.active_evictions = True
+            if min_entry_size == 0:
+                self.full_evictions = True
+        return freed

comfy_execution/graph.py

@@ -195,9 +195,10 @@ class ExecutionList(TopologicalSort):
     ExecutionList implements a topological dissolve of the graph. After a node is staged for execution,
     it can still be returned to the graph after having further dependencies added.
     """
-    def __init__(self, dynprompt, output_cache):
+    def __init__(self, dynprompt, output_cache, output_link_callback=None):
         super().__init__(dynprompt)
         self.output_cache = output_cache
+        self.output_link_callback = output_link_callback
         self.staged_node_id = None
         self.execution_cache = {}
         self.execution_cache_listeners = {}
@@ -205,13 +206,16 @@ class ExecutionList(TopologicalSort):
     def is_cached(self, node_id):
         return self.output_cache.get_local(node_id) is not None
 
-    def cache_link(self, from_node_id, to_node_id):
+    def cache_link(self, from_node_id, to_node_id, from_socket=None):
         if to_node_id not in self.execution_cache:
             self.execution_cache[to_node_id] = {}
-        self.execution_cache[to_node_id][from_node_id] = self.output_cache.get_local(from_node_id)
+        value = self.output_cache.get_local(from_node_id)
+        self.execution_cache[to_node_id][from_node_id] = value
         if from_node_id not in self.execution_cache_listeners:
             self.execution_cache_listeners[from_node_id] = set()
-        self.execution_cache_listeners[from_node_id].add(to_node_id)
+        self.execution_cache_listeners[from_node_id].add((to_node_id, from_socket))
+        if value is not None and from_socket is not None and self.output_link_callback is not None:
+            self.output_link_callback(value.outputs[from_socket])
 
     def get_cache(self, from_node_id, to_node_id):
         if to_node_id not in self.execution_cache:
@@ -225,13 +229,15 @@ class ExecutionList(TopologicalSort):
 
     def cache_update(self, node_id, value):
         if node_id in self.execution_cache_listeners:
-            for to_node_id in self.execution_cache_listeners[node_id]:
+            for to_node_id, from_socket in self.execution_cache_listeners[node_id]:
                 if to_node_id in self.execution_cache:
                     self.execution_cache[to_node_id][node_id] = value
+                if from_socket is not None and self.output_link_callback is not None:
+                    self.output_link_callback(value.outputs[from_socket])
 
     def add_strong_link(self, from_node_id, from_socket, to_node_id):
         super().add_strong_link(from_node_id, from_socket, to_node_id)
-        self.cache_link(from_node_id, to_node_id)
+        self.cache_link(from_node_id, to_node_id, from_socket)
 
     async def stage_node_execution(self):
         assert self.staged_node_id is None

comfy_execution/jobs.py

@@ -56,6 +56,9 @@ PREVIEWABLE_MEDIA_TYPES = frozenset({'images', 'video', 'audio', '3d', 'text'})
 # 3D file extensions for preview fallback (no dedicated media_type exists)
 THREE_D_EXTENSIONS = frozenset({'.obj', '.fbx', '.gltf', '.glb', '.usdz'})
 
+# Text file extensions for preview fallback (the formats SaveText can produce)
+TEXT_EXTENSIONS = frozenset({'.txt', '.md', '.json'})
+
 
 def has_3d_extension(filename: str) -> bool:
     lower = filename.lower()
@@ -143,9 +146,10 @@ def is_previewable(media_type: str, item: dict) -> bool:
     Maintains backwards compatibility with existing logic.
 
     Priority:
-    1. media_type is 'images', 'video', 'audio', or '3d'
+    1. media_type is 'images', 'video', 'audio', '3d', or 'text'
     2. format field starts with 'video/' or 'audio/'
     3. filename has a 3D extension (.obj, .fbx, .gltf, .glb, .usdz)
+    4. filename has a text extension (.txt, .md, .json, ...)
     """
     if media_type in PREVIEWABLE_MEDIA_TYPES:
         return True
@@ -156,10 +160,12 @@ def is_previewable(media_type: str, item: dict) -> bool:
     if fmt and (fmt.startswith('video/') or fmt.startswith('audio/')):
         return True
 
-    # Check for 3D files by extension
+    # Check for 3D and text files by extension
     filename = item.get('filename', '').lower()
     if any(filename.endswith(ext) for ext in THREE_D_EXTENSIONS):
         return True
+    if any(filename.endswith(ext) for ext in TEXT_EXTENSIONS):
+        return True
 
     return False
 
@@ -255,6 +261,10 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
     Preview priority (matching frontend):
     1. type="output" with previewable media
     2. Any previewable media
+
+    Text content entries (strings under 'text') are preview-only metadata,
+    matching the frontend's METADATA_KEYS: they can serve as the fallback
+    preview but are not counted as outputs.
     """
     count = 0
     preview_output = None
@@ -275,7 +285,6 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
                     if normalized is None:
                         # Not a 3D file string — check for text preview
                         if media_type == 'text':
-                            count += 1
                             if preview_output is None:
                                 if isinstance(item, tuple):
                                     text_value = item[0] if item else ''

comfy_extras/nodes_audio.py

@@ -298,6 +298,7 @@ class PreviewAudio(IO.ComfyNode):
             search_aliases=["play audio"],
             display_name="Preview Audio",
             category="audio",
+            description="Preview the audio without saving it to the ComfyUI output directory.",
             inputs=[
                 IO.Audio.Input("audio"),
             ],

comfy_extras/nodes_bounding_boxes.py

@@ -1,3 +1,5 @@
+import json
+
 import numpy as np
 import torch
 from PIL import Image, ImageDraw, ImageEnhance, ImageFont
@@ -166,6 +168,111 @@ def boxes_to_regions(boxes, width: int, height: int) -> list:
     return regions
 
 
+def normalize_incoming_boxes(bboxes) -> list:
+    if isinstance(bboxes, dict):
+        frame = [bboxes]
+    elif not isinstance(bboxes, list) or not bboxes:
+        frame = []
+    elif isinstance(bboxes[0], dict):
+        frame = bboxes
+    else:
+        frame = bboxes[0] if isinstance(bboxes[0], list) else []
+    boxes = []
+    for box in frame:
+        if not isinstance(box, dict):
+            continue
+        norm = {
+            "x": box.get("x", 0),
+            "y": box.get("y", 0),
+            "width": box.get("width", 0),
+            "height": box.get("height", 0),
+        }
+        meta = box.get("metadata")
+        if isinstance(meta, dict):
+            norm["metadata"] = meta
+        boxes.append(norm)
+    return boxes
+
+
+def _looks_like_element(box: dict) -> bool:
+    bbox = box.get("bbox")
+    return isinstance(bbox, (list, tuple)) and len(bbox) == 4
+
+
+def _looks_like_bbox(box: dict) -> bool:
+    return all(key in box for key in ("x", "y", "width", "height"))
+
+
+def elements_to_boxes(elements: list, width: int, height: int) -> list:
+    boxes = []
+    for element in elements:
+        if not isinstance(element, dict):
+            continue
+        bbox = element.get("bbox")
+        if not (isinstance(bbox, (list, tuple)) and len(bbox) == 4):
+            raise ValueError("bboxes element is missing a valid 'bbox' [ymin, xmin, ymax, xmax]")
+        try:
+            ymin, xmin, ymax, xmax = (float(v) / 1000.0 for v in bbox)
+        except (TypeError, ValueError):
+            raise ValueError("bboxes element 'bbox' must contain four numbers")
+        etype = "text" if element.get("type") == "text" else "obj"
+        boxes.append({
+            "x": round(min(xmin, xmax) * width),
+            "y": round(min(ymin, ymax) * height),
+            "width": round(abs(xmax - xmin) * width),
+            "height": round(abs(ymax - ymin) * height),
+            "metadata": {
+                "type": etype,
+                "text": element.get("text", "") if etype == "text" else "",
+                "desc": element.get("desc", ""),
+                "palette": element.get("color_palette", []) or [],
+            },
+        })
+    return boxes
+
+
+def boxes_from_input(data, width: int, height: int) -> list:
+    if data is None:
+        return []
+    if isinstance(data, str):
+        text = data.strip()
+        if not text:
+            return []
+        try:
+            data = json.loads(text)
+        except (ValueError, TypeError) as exc:
+            raise ValueError(f"bboxes string input is not valid JSON: {exc}") from exc
+    if isinstance(data, dict):
+        if _looks_like_element(data):
+            return elements_to_boxes([data], width, height)
+        if _looks_like_bbox(data):
+            return normalize_incoming_boxes(data)
+        raise ValueError(
+            "bboxes dict must be a bounding box (x, y, width, height) or an element (with a 'bbox')"
+        )
+    if not isinstance(data, list):
+        raise ValueError(
+            "bboxes input must be bounding boxes, elements, or a JSON string, "
+            f"got {type(data).__name__}"
+        )
+    if not data:
+        return []
+    first = data[0]
+    if isinstance(first, list):
+        return normalize_incoming_boxes(data)
+    if isinstance(first, dict):
+        if _looks_like_element(first):
+            return elements_to_boxes(data, width, height)
+        if _looks_like_bbox(first):
+            return normalize_incoming_boxes(data)
+        raise ValueError(
+            "bboxes items must be bounding boxes (x, y, width, height) or elements (with a 'bbox')"
+        )
+    raise ValueError(
+        f"bboxes list must contain bounding boxes or elements, got {type(first).__name__}"
+    )
+
+
 def _norm_bbox(region: dict) -> list[int]:
     def grid(value: float) -> int:
         return max(0, min(1000, round(value * 1000)))
@@ -217,29 +324,48 @@ class CreateBoundingBoxes(io.ComfyNode):
                     optional=True,
                     tooltip="Optional image used as background in the canvas and preview.",
                 ),
+                io.MultiType.Input(
+                    "bboxes",
+                    [io.BoundingBox, io.Array, io.String],
+                    optional=True,
+                    tooltip="Bounding boxes, elements, or a JSON string to initialize the canvas. A new upstream value initializes the canvas; edits made on the canvas take priority and are kept until the upstream value changes again.",
+                ),
                 io.Int.Input("width", default=1024, min=64, max=16384, step=16,
                              tooltip="Width of the canvas and the pixel grid for the bounding boxes."),
                 io.Int.Input("height", default=1024, min=64, max=16384, step=16,
                              tooltip="Height of the canvas and the pixel grid for the bounding boxes."),
                 editor_state,
+                io.BoundingBoxes.Input(
+                    "last_incoming",
+                    optional=True,
+                    tooltip="Internal state managed by the canvas: the upstream bboxes value that last initialized it. Leave empty to re-initialize the canvas from the bboxes input on the next run.",
+                ),
             ],
             outputs=[
                 io.Image.Output(display_name="preview"),
                 io.BoundingBox.Output(display_name="bboxes"),
                 io.Array.Output(display_name="elements"),
             ],
+            is_output_node=True,
             is_experimental=True,
         )
 
     @classmethod
-    def execute(cls, width, height, editor_state=None, background=None) -> io.NodeOutput:
-        regions = boxes_to_regions(editor_state, width, height)
+    def execute(cls, width, height, editor_state=None, last_incoming=None, background=None, bboxes=None) -> io.NodeOutput:
+        incoming = boxes_from_input(bboxes, width, height)
+        applied = last_incoming if isinstance(last_incoming, list) else []
+        upstream_changed = bool(incoming) and incoming != applied
+        source = incoming if upstream_changed else (editor_state or [])
+        regions = boxes_to_regions(source, width, height)
         preview = render_preview(regions, width, height, _bg_from_image(background))
+        ui = {"dims": [width, height]}
+        if incoming:
+            ui["input_bboxes"] = incoming
         return io.NodeOutput(
             preview,
             fractions_to_bbox_frame(regions, width, height),
             build_elements(regions),
-            ui={"dims": [width, height]},
+            ui=ui,
         )
 
 

comfy_extras/nodes_dataset.py

@@ -2,6 +2,7 @@ import logging
 import os
 import json
 
+import av
 import numpy as np
 import torch
 from PIL import Image
@@ -9,7 +10,7 @@ from typing_extensions import override
 
 import folder_paths
 import node_helpers
-from comfy_api.latest import ComfyExtension, io
+from comfy_api.latest import ComfyExtension, io, Input, InputImpl, Types
 
 
 def load_and_process_images(image_files, input_dir):
@@ -42,6 +43,130 @@ def load_and_process_images(image_files, input_dir):
     return output_images
 
 
+def secure_subfolder_path(base_dir, folder_name):
+    """Resolve folder_name inside base_dir, rejecting anything that escapes it.
+
+    Blocks '..', absolute paths, drive letters and symlink escapes using the
+    same realpath containment check as the core file endpoints.
+    """
+    target = os.path.abspath(os.path.join(base_dir, folder_name))
+    if not folder_paths.is_within_directory(base_dir, target):
+        raise ValueError(f"Invalid folder name {folder_name!r}: resolves outside of {base_dir}")
+    return target
+
+
+def list_dataset_folders():
+    """Relative paths of dataset folders found under all dataset roots.
+
+    Any subfolder containing a metadata.json or *.safetensors shard counts as
+    a dataset; the walk doesn't descend into matched folders.
+
+    Symlinked directories are followed, but symlink loops are avoided.
+    """
+    found = set()
+
+    for root in folder_paths.get_folder_paths("datasets"):
+        if not os.path.isdir(root):
+            continue
+
+        root = os.path.abspath(root)
+        seen_dirs = set()
+
+        for dirpath, subdirs, filenames in os.walk(root, followlinks=True):
+            try:
+                st = os.stat(dirpath)  # follows symlinks
+            except OSError:
+                subdirs[:] = []
+                continue
+
+            dir_key = (st.st_dev, st.st_ino)
+            if dir_key in seen_dirs:
+                subdirs[:] = []
+                continue
+
+            seen_dirs.add(dir_key)
+
+            if dirpath != root and (
+                "metadata.json" in filenames
+                or any(f.endswith(".safetensors") for f in filenames)
+            ):
+                found.add(os.path.relpath(dirpath, root).replace(os.sep, "/"))
+                subdirs[:] = []
+                continue
+
+            kept_subdirs = []
+            for name in subdirs:
+                child = os.path.join(dirpath, name)
+                try:
+                    child_st = os.stat(child)  # follows symlinks
+                except OSError:
+                    continue
+
+                child_key = (child_st.st_dev, child_st.st_ino)
+                if child_key not in seen_dirs:
+                    kept_subdirs.append(name)
+
+            subdirs[:] = kept_subdirs
+
+    return sorted(found)
+
+
+def get_dataset_save_dir(folder_name):
+    """Resolve the folder to save a new dataset into, inside the default root.
+
+    The folder is not created here; callers makedirs after validation.
+    """
+    root = folder_paths.get_folder_paths("datasets")[0]
+    target = secure_subfolder_path(root, folder_name)
+    if os.path.realpath(target) == os.path.realpath(root):
+        raise ValueError("folder_name must name a subfolder of the datasets directory, e.g. 'my_dataset'.")
+    return target
+
+
+def get_dataset_dir(folder_name):
+    """Find an existing dataset folder by relative name across all dataset roots."""
+    roots = folder_paths.get_folder_paths("datasets")
+    for root in roots:
+        target = secure_subfolder_path(root, folder_name)
+        if os.path.realpath(target) == os.path.realpath(root):
+            raise ValueError("folder_name must name a subfolder of the datasets directory, e.g. 'my_dataset'.")
+        if os.path.isdir(target):
+            return target
+    raise ValueError(f"Dataset folder {folder_name!r} not found in: {', '.join(roots)}")
+
+
+VALID_VIDEO_EXTENSIONS = [".mp4", ".avi", ".mov", ".webm", ".mkv", ".flv"]
+
+
+def _decode_selected_frames(video: Input.Video, indices: list[int]) -> Input.Video:
+    """Decode only the requested frame indices from a video.
+
+    Opens the underlying container once, decodes frames in presentation order,
+    keeps only the ones whose index is in ``indices``, and returns the result
+    wrapped in a VideoFromComponents so it still satisfies the VideoInput
+    contract for downstream nodes.
+    """
+    indices_sorted = sorted(set(indices))
+    max_idx = indices_sorted[-1]
+    source = video.get_stream_source()
+
+    frames_by_idx: dict[int, torch.Tensor] = {}
+    with av.open(source, mode="r") as container:
+        stream = container.streams.video[0]
+        wanted = set(indices_sorted)
+        for frame_idx, frame in enumerate(container.decode(stream)):
+            if frame_idx in wanted:
+                img = frame.to_ndarray(format="rgb24")
+                frames_by_idx[frame_idx] = torch.from_numpy(img.copy()).float() / 255.0
+            if frame_idx >= max_idx:
+                break
+
+    stacked = torch.stack([frames_by_idx[i] for i in indices])
+    return InputImpl.VideoFromComponents(
+        Types.VideoComponents(images=stacked, frame_rate=video.get_frame_rate())
+    )
+
+
 class LoadImageDataSetFromFolderNode(io.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -157,6 +282,116 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
         return io.NodeOutput(output_tensor, captions)
 
 
+class LoadVideoDataSetFromFolderNode(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LoadVideoDataSetFromFolder",
+            search_aliases=["load folder", "load from folder", "load dataset", "load videos", "import dataset"],
+            display_name="Load Video (from Folder)",
+            category="video",
+            description="Load a dataset of videos from a specified folder and return a list of videos. Supported formats: MP4, AVI, MOV, WEBM, MKV, FLV.",
+            is_experimental=True,
+            inputs=[
+                io.Combo.Input(
+                    "folder",
+                    options=folder_paths.get_input_subfolders(),
+                    tooltip="The folder containing video files.",
+                ),
+            ],
+            outputs=[
+                io.Video.Output(
+                    display_name="videos",
+                    is_output_list=True,
+                    tooltip="Lazy video references; frames are decoded only when needed downstream.",
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, folder):
+        sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
+        video_files = sorted([
+            f for f in os.listdir(sub_input_dir)
+            if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS)
+        ])
+
+        if not video_files:
+            raise ValueError(f"No video files found in {sub_input_dir}")
+
+        videos = [InputImpl.VideoFromFile(os.path.join(sub_input_dir, f)) for f in video_files]
+        logging.info(f"Loaded {len(videos)} lazy video references from {sub_input_dir}")
+        return io.NodeOutput(videos)
+
+
+class LoadVideoTextDataSetFromFolderNode(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LoadVideoTextDataSetFromFolder",
+            search_aliases=["load folder", "load from folder", "load dataset", "load videos", "import dataset"],
+            display_name="Load Video-Text (from Folder)",
+            category="video",
+            description="Load a dataset of pairs of videos and text captions from a specified folder and return them as a list. Supported formats: MP4, AVI, MOV, WEBM, MKV, FLV.",
+            is_experimental=True,
+            inputs=[
+                io.Combo.Input(
+                    "folder",
+                    options=folder_paths.get_input_subfolders(),
+                    tooltip="The folder containing video files and .txt captions.",
+                ),
+            ],
+            outputs=[
+                io.Video.Output(
+                    display_name="videos",
+                    is_output_list=True,
+                    tooltip="Lazy video references; frames are decoded only when needed downstream.",
+                ),
+                io.String.Output(
+                    display_name="texts",
+                    is_output_list=True,
+                    tooltip="List of text captions.",
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, folder):
+        sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
+
+        video_files = []
+        for item in sorted(os.listdir(sub_input_dir)):
+            path = os.path.join(sub_input_dir, item)
+            if any(item.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS):
+                video_files.append(path)
+            elif os.path.isdir(path):
+                # Support kohya-ss/sd-scripts folder structure: {repeat}_{desc}/
+                repeat = 1
+                if item.split("_")[0].isdigit():
+                    repeat = int(item.split("_")[0])
+                video_files.extend([
+                    os.path.join(path, f)
+                    for f in sorted(os.listdir(path))
+                    if any(f.lower().endswith(ext) for ext in VALID_VIDEO_EXTENSIONS)
+                ] * repeat)
+
+        if not video_files:
+            raise ValueError(f"No video files found in {sub_input_dir}")
+
+        captions = []
+        for vf in video_files:
+            caption_path = os.path.splitext(vf)[0] + ".txt"
+            if os.path.exists(caption_path):
+                with open(caption_path, "r", encoding="utf-8") as f:
+                    captions.append(f.read().strip())
+            else:
+                captions.append("")
+
+        videos = [InputImpl.VideoFromFile(vf) for vf in video_files]
+        logging.info(f"Loaded {len(videos)} lazy video references with captions from {sub_input_dir}")
+        return io.NodeOutput(videos, captions)
+
+
 def save_images_to_folder(image_list, output_dir, prefix="image", overwrite=True):
     """Utility function to save a list of image tensors to disk.
 
@@ -252,7 +487,7 @@ class SaveImageDataSetToFolderNode(io.ComfyNode):
         filename_prefix = filename_prefix[0]
         mode = mode[0]
 
-        output_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
+        output_dir = secure_subfolder_path(folder_paths.get_output_directory(), folder_name)
         saved_files = save_images_to_folder(images, output_dir, filename_prefix, mode=='overwrite')
 
         logging.info(f"Saved {len(saved_files)} images to {output_dir}.")
@@ -306,7 +541,7 @@ class SaveImageTextDataSetToFolderNode(io.ComfyNode):
         filename_prefix = filename_prefix[0]
         mode = mode[0]
 
-        output_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
+        output_dir = secure_subfolder_path(folder_paths.get_output_directory(), folder_name)
         saved_files = save_images_to_folder(images, output_dir, filename_prefix, mode=='overwrite')
 
         # Save captions
@@ -470,7 +705,15 @@ class ImageProcessingNode(io.ComfyNode):
 
     @classmethod
     def execute(cls, images, **kwargs):
-        """Execute the node. Routes to _process or _group_process based on mode."""
+        """Execute the node. Routes to _process or _group_process based on mode.
+
+        For individual processing (_process), automatically handles multi-frame
+        inputs (video tensors [T, H, W, C]) by applying _process per-frame and
+        concatenating the results. This allows all spatial transform nodes to
+        work with video without modification. Nodes that natively handle batched
+        tensors (e.g. pure tensor math) can set per_frame_process = False to
+        skip the per-frame loop.
+        """
         is_group = cls._detect_processing_mode()
 
         if is_group:
@@ -489,7 +732,16 @@ class ImageProcessingNode(io.ComfyNode):
             result = cls._group_process(images, **params)
         else:
             # Individual processing: images is single item, call _process
-            result = cls._process(images, **params)
+            # Auto-loop over frames for multi-frame inputs (video [T, H, W, C])
+            # so that PIL-based spatial transforms work per-frame automatically.
+            if images.shape[0] > 1 and getattr(cls, 'per_frame_process', True):
+                results = []
+                for i in range(images.shape[0]):
+                    frame_result = cls._process(images[i:i + 1], **params)
+                    results.append(frame_result)
+                result = torch.cat(results, dim=0)
+            else:
+                result = cls._process(images, **params)
 
         return io.NodeOutput(result)
 
@@ -803,6 +1055,7 @@ class NormalizeImagesNode(ImageProcessingNode):
     display_name = "Normalize Image Colors"
     category = "image/color"
     description = "Normalize images using mean and standard deviation."
+    per_frame_process = False  # Pure tensor math, handles any batch size
     extra_inputs = [
         io.Float.Input(
             "mean",
@@ -833,6 +1086,7 @@ class AdjustBrightnessNode(ImageProcessingNode):
     display_name = "Adjust Brightness"
     category="image/adjustments"
     description = "Adjust the brightness of an image."
+    per_frame_process = False  # Pure tensor math, handles any batch size
     extra_inputs = [
         io.Float.Input(
             "factor",
@@ -854,6 +1108,7 @@ class AdjustContrastNode(ImageProcessingNode):
     display_name = "Adjust Contrast"
     category="image/adjustments"
     description = "Adjust the contrast of an image."
+    per_frame_process = False  # Pure tensor math, handles any batch size
     extra_inputs = [
         io.Float.Input(
             "factor",
@@ -935,6 +1190,261 @@ class ShuffleImageTextDatasetNode(io.ComfyNode):
         return io.NodeOutput(shuffled_images, shuffled_texts)
 
 
+# ========== Video Processing Nodes ==========
+
+
+class VideoFrameSampleNode(io.ComfyNode):
+    """Sample a fixed number of frames from a video using various strategies.
+
+    For contiguous strategies ("head"/"tail") the result is a fully lazy
+    VideoInput (no frames decoded). For non-contiguous strategies
+    ("uniform"/"random") only the selected indices are decoded.
+    """
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="VideoFrameSample",
+            search_aliases=["sample frames", "extract frames"],
+            display_name="Sample Video Frame",
+            category="video",
+            description="Sample a fixed number of frames from a video using various strategies.",
+            is_experimental=True,
+            inputs=[
+                io.Video.Input("video", tooltip="Input video."),
+                io.Int.Input(
+                    "num_frames",
+                    default=16,
+                    min=1,
+                    max=9999,
+                    tooltip="Number of frames to sample.",
+                ),
+                io.Combo.Input(
+                    "strategy",
+                    options=["uniform", "head", "tail", "random"],
+                    default="uniform",
+                    tooltip="uniform: evenly spaced, head: first N, tail: last N, random: random sorted.",
+                ),
+                io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=0xFFFFFFFFFFFFFFFF,
+                    tooltip="Random seed (only used with 'random' strategy).",
+                ),
+            ],
+            outputs=[
+                io.Video.Output(display_name="video", tooltip="Sampled video."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, video, num_frames, strategy, seed):
+        total_frames = video.get_frame_count()
+        num_frames = min(num_frames, total_frames)
+        fps = float(video.get_frame_rate())
+
+        if strategy == "head":
+            return io.NodeOutput(
+                video.as_trimmed(0.0, num_frames / fps, strict_duration=False)
+            )
+        if strategy == "tail":
+            start_t = (total_frames - num_frames) / fps
+            return io.NodeOutput(
+                video.as_trimmed(start_t, num_frames / fps, strict_duration=False)
+            )
+
+        if strategy == "uniform":
+            if num_frames == 1:
+                indices = [total_frames // 2]
+            else:
+                indices = [round(i * (total_frames - 1) / (num_frames - 1)) for i in range(num_frames)]
+        elif strategy == "random":
+            rng = np.random.RandomState(seed % (2**32 - 1))
+            indices = sorted(rng.choice(total_frames, size=num_frames, replace=False).tolist())
+        else:
+            raise ValueError(f"Unknown strategy: {strategy}")
+
+        return io.NodeOutput(_decode_selected_frames(video, indices))
+
+
+class VideoTemporalCropNode(io.ComfyNode):
+    """Crop a continuous range of frames from a video (fully lazy)."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="VideoTemporalCrop",
+            search_aliases=["crop", "crop video", "temporal crop", "truncate video"],
+            display_name="Crop Video (Temporal)",
+            category="video/transform",
+            description="Crop a continuous range of frames from a video.",
+            is_experimental=True,
+            inputs=[
+                io.Video.Input("video", tooltip="Input video."),
+                io.Int.Input(
+                    "start_frame",
+                    default=0,
+                    min=0,
+                    max=99999,
+                    tooltip="Starting frame index.",
+                ),
+                io.Int.Input(
+                    "length",
+                    default=16,
+                    min=1,
+                    max=99999,
+                    tooltip="Number of frames to keep.",
+                ),
+            ],
+            outputs=[
+                io.Video.Output(display_name="video", tooltip="Cropped video (lazy)."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, video, start_frame, length):
+        total_frames = video.get_frame_count()
+        fps = float(video.get_frame_rate())
+        start_frame = min(start_frame, max(total_frames - 1, 0))
+        length = min(length, total_frames - start_frame)
+        return io.NodeOutput(
+            video.as_trimmed(start_frame / fps, length / fps, strict_duration=False)
+        )
+
+
+class VideoRandomTemporalCropNode(io.ComfyNode):
+    """Randomly crop a continuous range of frames from a video (fully lazy)."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="VideoRandomTemporalCrop",
+            search_aliases=["crop", "crop video", "temporal crop", "truncate video", "random crop"],
+            display_name="Crop Video (Temporal Random)",
+            category="video/transform",
+            description="Randomly crop a continuous range of frames from a video.",
+            is_experimental=True,
+            inputs=[
+                io.Video.Input("video", tooltip="Input video."),
+                io.Int.Input(
+                    "length",
+                    default=16,
+                    min=1,
+                    max=99999,
+                    tooltip="Number of frames to keep.",
+                ),
+                io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=0xFFFFFFFFFFFFFFFF,
+                    tooltip="Random seed.",
+                ),
+            ],
+            outputs=[
+                io.Video.Output(display_name="video", tooltip="Cropped video (lazy)."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, video, length, seed):
+        total_frames = video.get_frame_count()
+        fps = float(video.get_frame_rate())
+        length = min(length, total_frames)
+        max_start = total_frames - length
+        rng = np.random.RandomState(seed % (2**32 - 1))
+        start = rng.randint(0, max_start + 1) if max_start > 0 else 0
+        return io.NodeOutput(
+            video.as_trimmed(start / fps, length / fps, strict_duration=False)
+        )
+
+
+class ShuffleVideoDatasetNode(io.ComfyNode):
+    """Randomly shuffle the order of videos in the dataset."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="ShuffleVideoDataset",
+            search_aliases=["shuffle", "randomize", "mix"],
+            display_name="Shuffle Videos List",
+            category="video/batch",
+            description="Randomly shuffle the order of videos in a list.",
+            is_experimental=True,
+            is_input_list=True,
+            inputs=[
+                io.Video.Input("videos", tooltip="List of videos to shuffle."),
+                io.Int.Input(
+                    "seed", default=0, min=0, max=0xFFFFFFFFFFFFFFFF, tooltip="Random seed."
+                ),
+            ],
+            outputs=[
+                io.Video.Output(
+                    display_name="videos",
+                    is_output_list=True,
+                    tooltip="Shuffled videos",
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, videos, seed):
+        seed = seed[0] if isinstance(seed, list) else seed
+        np.random.seed(seed % (2**32 - 1))
+        indices = np.random.permutation(len(videos))
+        return io.NodeOutput([videos[i] for i in indices])
+
+
+class ShuffleVideoTextDatasetNode(io.ComfyNode):
+    """Shuffle videos and their captions together, preserving pairs."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="ShuffleVideoTextDataset",
+            search_aliases=["shuffle", "randomize", "mix"],
+            display_name="Shuffle Pairs of Video-Text",
+            category="dataset/video",
+            description="Randomly shuffle the order of pairs of video-text in a list.",
+            is_experimental=True,
+            is_input_list=True,
+            inputs=[
+                io.Video.Input("videos", tooltip="List of videos to shuffle."),
+                io.String.Input("texts", tooltip="List of texts to shuffle."),
+                io.Int.Input(
+                    "seed",
+                    default=0,
+                    min=0,
+                    max=0xFFFFFFFFFFFFFFFF,
+                    tooltip="Random seed.",
+                ),
+            ],
+            outputs=[
+                io.Video.Output(
+                    display_name="videos",
+                    is_output_list=True,
+                    tooltip="Shuffled videos",
+                ),
+                io.String.Output(
+                    display_name="texts",
+                    is_output_list=True,
+                    tooltip="Shuffled texts",
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, videos, texts, seed):
+        seed = seed[0] if isinstance(seed, list) else seed
+        np.random.seed(seed % (2**32 - 1))
+        indices = np.random.permutation(len(videos))
+        return io.NodeOutput(
+            [videos[i] for i in indices],
+            [texts[i] for i in indices],
+        )
+
+
 # ========== Text Transform Nodes ==========
 
 
@@ -1443,7 +1953,7 @@ class SaveTrainingDataset(io.ComfyNode):
                 io.String.Input(
                     "folder_name",
                     default="training_dataset",
-                    tooltip="Name of folder to save dataset (inside output directory).",
+                    tooltip="Name of folder to save the dataset into, inside the datasets directory. Subfolders like 'project/run1' are allowed.",
                 ),
                 io.Int.Input(
                     "shard_size",
@@ -1473,8 +1983,8 @@ class SaveTrainingDataset(io.ComfyNode):
                 f"Something went wrong in dataset preparation."
             )
 
-        # Create output directory
-        output_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
+        # Create output directory (inside the datasets root, traversal-safe)
+        output_dir = get_dataset_save_dir(folder_name)
         os.makedirs(output_dir, exist_ok=True)
 
         # Prepare data pairs
@@ -1533,10 +2043,10 @@ class LoadTrainingDataset(io.ComfyNode):
             description="Load encoded training dataset (latents + conditioning) from disk for use in training.",
             is_experimental=True,
             inputs=[
-                io.String.Input(
+                io.Combo.Input(
                     "folder_name",
-                    default="training_dataset",
-                    tooltip="Name of folder containing the saved dataset (inside output directory).",
+                    options=list_dataset_folders(),
+                    tooltip="Saved dataset to load, from the datasets directory.",
                 ),
             ],
             outputs=[
@@ -1555,11 +2065,8 @@ class LoadTrainingDataset(io.ComfyNode):
 
     @classmethod
     def execute(cls, folder_name):
-        # Get dataset directory
-        dataset_dir = os.path.join(folder_paths.get_output_directory(), folder_name)
-
-        if not os.path.exists(dataset_dir):
-            raise ValueError(f"Dataset directory not found: {dataset_dir}")
+        # Get dataset directory (searched across all dataset roots, traversal-safe)
+        dataset_dir = get_dataset_dir(folder_name)
 
         # Find all shard files
         shard_files = sorted(
@@ -1608,7 +2115,10 @@ class DatasetExtension(ComfyExtension):
             LoadImageTextDataSetFromFolderNode,
             SaveImageDataSetToFolderNode,
             SaveImageTextDataSetToFolderNode,
-            # Image transform nodes
+            # Video data loading nodes
+            LoadVideoDataSetFromFolderNode,
+            LoadVideoTextDataSetFromFolderNode,
+            # Image transform nodes (auto-handle video via per-frame processing)
             ResizeImagesByShorterEdgeNode,
             ResizeImagesByLongerEdgeNode,
             CenterCropImagesNode,
@@ -1618,6 +2128,12 @@ class DatasetExtension(ComfyExtension):
             AdjustContrastNode,
             ShuffleDatasetNode,
             ShuffleImageTextDatasetNode,
+            # Video processing nodes (lazy VideoInput in/out)
+            VideoFrameSampleNode,
+            VideoTemporalCropNode,
+            VideoRandomTemporalCropNode,
+            ShuffleVideoDatasetNode,
+            ShuffleVideoTextDatasetNode,
             # Text transform nodes
             TextToLowercaseNode,
             TextToUppercaseNode,

comfy_extras/nodes_fresca.py

@@ -10,7 +10,7 @@ def Fourier_filter(x, scale_low=1.0, scale_high=1.5, freq_cutoff=20):
     Apply frequency-dependent scaling to an image tensor using Fourier transforms.
 
     Parameters:
-        x:           Input tensor of shape (B, C, H, W)
+        x:           Input tensor of shape (..., H, W)
         scale_low:   Scaling factor for low-frequency components (default: 1.0)
         scale_high:  Scaling factor for high-frequency components (default: 1.5)
         freq_cutoff: Number of frequency indices around center to consider as low-frequency (default: 20)
@@ -31,8 +31,8 @@ def Fourier_filter(x, scale_low=1.0, scale_high=1.5, freq_cutoff=20):
     # Initialize mask with high-frequency scaling factor
     mask = torch.ones(x_freq.shape, device=device) * scale_high
     m = mask
-    for d in range(len(x_freq.shape) - 2):
-        dim = d + 2
+    for d in range(2):
+        dim = len(x_freq.shape) - 2 + d
         cc = x_freq.shape[dim] // 2
         f_c = min(freq_cutoff, cc)
         m = m.narrow(dim, cc - f_c, f_c * 2)

comfy_extras/nodes_hunyuan.py

@@ -2,6 +2,8 @@ import nodes
 import node_helpers
 import torch
 import comfy.model_management
+import comfy.model_patcher
+import comfy.ops
 from typing_extensions import override
 from comfy_api.latest import ComfyExtension, io
 from comfy.ldm.hunyuan_video.upsampler import HunyuanVideo15SRModel
@@ -217,8 +219,11 @@ class LatentUpscaleModelLoader(io.ComfyNode):
             model.load_sd(sd)
         elif "post_upsample_res_blocks.0.conv2.bias" in sd:
             config = json.loads(metadata["config"])
-            model = LatentUpsampler.from_config(config).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32]))
-            model.load_state_dict(sd)
+            model = LatentUpsampler.from_config(config, operations=comfy.ops.disable_weight_init).to(dtype=comfy.model_management.vae_dtype(allowed_dtypes=[torch.bfloat16, torch.float32]))
+            comfy.model_management.archive_model_dtypes(model)
+            model_patcher = comfy.model_patcher.CoreModelPatcher(model, load_device=comfy.model_management.get_torch_device(), offload_device=comfy.model_management.unet_offload_device())
+            model.load_state_dict(sd, assign=model_patcher.is_dynamic())
+            model = model_patcher
 
         return io.NodeOutput(model)
 

comfy_extras/nodes_images.py

@@ -844,15 +844,18 @@ class ImageMergeTileList(IO.ComfyNode):
 # Format specifications
 # ---------------------------------------------------------------------------
 
-# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format,
-# stream pix_fmt). Keeps the encode path declarative instead of branchy.
+# Maps (file_format, bit_depth, num_channels) -> (quantization scale, numpy dtype,
+# av frame pix_fmt, stream pix_fmt). Keeps the encode path declarative instead of branchy.
 _FORMAT_SPECS = {
-    ("png", "8-bit", False):  {"scale": 255.0,   "dtype": np.uint8,   "frame_fmt": "rgb24",     "stream_fmt": "rgb24"},
-    ("png", "8-bit", True):   {"scale": 255.0,   "dtype": np.uint8,   "frame_fmt": "rgba",      "stream_fmt": "rgba"},
-    ("png", "16-bit", False): {"scale": 65535.0, "dtype": np.uint16,  "frame_fmt": "rgb48le",   "stream_fmt": "rgb48be"},
-    ("png", "16-bit", True):  {"scale": 65535.0, "dtype": np.uint16,  "frame_fmt": "rgba64le",  "stream_fmt": "rgba64be"},
-    ("exr", "32-bit float", False): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le",  "stream_fmt": "gbrpf32le"},
-    ("exr", "32-bit float", True):  {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
+    ("png", "8-bit", 1):  {"scale": 255.0,   "dtype": np.uint8,   "frame_fmt": "gray",      "stream_fmt": "gray"},
+    ("png", "8-bit", 3):  {"scale": 255.0,   "dtype": np.uint8,   "frame_fmt": "rgb24",     "stream_fmt": "rgb24"},
+    ("png", "8-bit", 4):  {"scale": 255.0,   "dtype": np.uint8,   "frame_fmt": "rgba",      "stream_fmt": "rgba"},
+    ("png", "16-bit", 1): {"scale": 65535.0, "dtype": np.uint16,  "frame_fmt": "gray16le",  "stream_fmt": "gray16be"},
+    ("png", "16-bit", 3): {"scale": 65535.0, "dtype": np.uint16,  "frame_fmt": "rgb48le",   "stream_fmt": "rgb48be"},
+    ("png", "16-bit", 4): {"scale": 65535.0, "dtype": np.uint16,  "frame_fmt": "rgba64le",  "stream_fmt": "rgba64be"},
+    ("exr", "32-bit float", 1): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "grayf32le",  "stream_fmt": "grayf32le"},
+    ("exr", "32-bit float", 3): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le",  "stream_fmt": "gbrpf32le"},
+    ("exr", "32-bit float", 4): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
 }
 
 
@@ -891,10 +894,11 @@ def hlg_to_linear(t: torch.Tensor) -> torch.Tensor:
         return torch.cat([hlg_to_linear(rgb), alpha], dim=-1)
 
     # Piecewise: sqrt branch below 0.5, log branch above.
-    # Clamp inside the log branch so negative / out-of-range values don't blow up;
+    # Clamp the log branch at the 0.5 branch point (not above it) so the
+    # unselected lane stays finite in exp() without altering selected values;
     # values above 1.0 are allowed and extrapolate naturally.
     low = (t ** 2) / 3.0
-    high = (torch.exp((t.clamp(min=_HLG_C) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
+    high = (torch.exp((t.clamp(min=0.5) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
     return torch.where(t <= 0.5, low, high)
 
 
@@ -1087,7 +1091,8 @@ def _encode_image(
     bit_depth: str,
     colorspace: str,
 ) -> bytes:
-    """Encode a single HxWxC tensor to PNG or EXR bytes in memory.
+    """Encode a single HxWxC (or channel-less HxW grayscale) tensor to PNG or
+    EXR bytes in memory. Grayscale is written as single-channel PNG / Y-only EXR.
 
     For EXR the input is interpreted according to `colorspace` and converted
     to scene-linear (EXR's convention) before writing:
@@ -1101,10 +1106,16 @@ def _encode_image(
     For PNG, colorspace selection does not modify pixels — PNG is delivered
     sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
     """
+    if img_tensor.ndim == 2:
+        img_tensor = img_tensor.unsqueeze(-1)  # Some nodes emit grayscale as (H, W) with no channel dim, mask-style.
     height, width, num_channels = img_tensor.shape
-    has_alpha = num_channels == 4
 
-    spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)]
+    spec = _FORMAT_SPECS.get((file_format, bit_depth, num_channels))
+    if spec is None:
+        raise ValueError(
+            f"No {file_format}/{bit_depth} encoder for {num_channels}-channel images: "
+            "supported channel counts are 1 (grayscale), 3 (RGB) and 4 (RGBA)."
+        )
 
     if spec["dtype"] == np.float32:
         # EXR path: preserve full range, no clamp.

comfy_extras/nodes_joyimage.py

@@ -0,0 +1,102 @@
+from typing_extensions import override
+
+import comfy.utils
+import node_helpers
+from comfy_api.latest import ComfyExtension, io
+
+
+# fmt: off
+BUCKETS_1024 = [
+    (512, 1792), (512, 1856), (512, 1920), (512, 1984), (512, 2048),
+    (576, 1600), (576, 1664), (576, 1728), (576, 1792),
+    (640, 1472), (640, 1536), (640, 1600),
+    (704, 1344), (704, 1408), (704, 1472),
+    (768, 1216), (768, 1280), (768, 1344),
+    (832, 1152), (832, 1216),
+    (896, 1088), (896, 1152),
+    (960, 1024), (960, 1088),
+    (1024, 960), (1024, 1024),
+    (1088, 896), (1088, 960),
+    (1152, 832), (1152, 896),
+    (1216, 768), (1216, 832),
+    (1280, 768),
+    (1344, 704), (1344, 768),
+    (1408, 704),
+    (1472, 640), (1472, 704),
+    (1536, 640),
+    (1600, 576), (1600, 640),
+    (1664, 576),
+    (1728, 576),
+    (1792, 512), (1792, 576),
+    (1856, 512),
+    (1920, 512),
+    (1984, 512),
+    (2048, 512),
+]
+# fmt: on
+
+
+def _find_best_bucket(height: int, width: int) -> tuple[int, int]:
+    target_ratio = height / width
+    return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio))
+
+
+def _resize_reference(image):
+    if image.shape[0] != 1:
+        raise ValueError("JoyImage reference inputs must contain one image each")
+    samples = image.movedim(-1, 1)
+    bucket_h, bucket_w = _find_best_bucket(samples.shape[2], samples.shape[3])
+    resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center")
+    return resized.movedim(1, -1)[:, :, :, :3]
+
+
+def _encode(clip, prompt, vae, images):
+    resized_images = [_resize_reference(image) for image in images]
+    conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=resized_images))
+    if vae is not None and resized_images:
+        ref_latents = [vae.encode(image) for image in resized_images]
+        conditioning = node_helpers.conditioning_set_values(
+            conditioning, {"reference_latents": ref_latents}, append=True,
+        )
+    return conditioning
+
+
+class TextEncodeJoyImageEdit(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        image_template = io.Autogrow.TemplatePrefix(
+            io.Image.Input("image"),
+            prefix="image",
+            min=0,
+            max=6,
+        )
+        return io.Schema(
+            node_id="TextEncodeJoyImageEdit",
+            category="model/conditioning/joyimage",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
+                io.Vae.Input("vae", optional=True),
+                io.Autogrow.Input("images", template=image_template, optional=True),
+            ],
+            outputs=[
+                io.Conditioning.Output(),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput:
+        images = images or {}
+        return io.NodeOutput(_encode(clip, prompt, vae, list(images.values())))
+
+
+class JoyImageExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            TextEncodeJoyImageEdit,
+        ]
+
+
+async def comfy_entrypoint() -> JoyImageExtension:
+    return JoyImageExtension()

comfy_extras/nodes_load_3d.py

@@ -61,14 +61,10 @@ class Load3D(IO.ComfyNode):
 
     @classmethod
     def execute(cls, model_file, image, **kwargs) -> IO.NodeOutput:
-        image_path = folder_paths.get_annotated_filepath(image['image'])
-        mask_path = folder_paths.get_annotated_filepath(image['mask'])
-        normal_path = folder_paths.get_annotated_filepath(image['normal'])
-
         load_image_node = nodes.LoadImage()
-        output_image, ignore_mask = load_image_node.load_image(image=image_path)
-        ignore_image, output_mask = load_image_node.load_image(image=mask_path)
-        normal_image, ignore_mask2 = load_image_node.load_image(image=normal_path)
+        output_image, ignore_mask = load_image_node.load_image(image=image['image'])
+        ignore_image, output_mask = load_image_node.load_image(image=image['mask'])
+        normal_image, ignore_mask2 = load_image_node.load_image(image=image['normal'])
 
         video = None
 
@@ -96,6 +92,7 @@ class Preview3D(IO.ComfyNode):
             search_aliases=["view mesh", "3d viewer"],
             display_name="Preview 3D & Animation",
             category="3d",
+            description="Preview a 3D model file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             inputs=[
@@ -140,6 +137,7 @@ class Preview3DAdvanced(IO.ComfyNode):
             display_name="Preview 3D (Advanced)",
             search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"],
             category="3d",
+            description="Preview a 3D model file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             inputs=[
@@ -176,8 +174,9 @@ class Preview3DAdvanced(IO.ComfyNode):
         filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
         model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
 
+        viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
         camera_info_input = kwargs.get("camera_info", None)
-        camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
+        camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
         model_3d_info_input = kwargs.get("model_3d_info", None)
         model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
         return IO.NodeOutput(
@@ -197,6 +196,7 @@ class PreviewGaussianSplat(IO.ComfyNode):
             node_id="PreviewGaussianSplat",
             display_name="Preview Splat",
             category="3d",
+            description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             search_aliases=[
@@ -244,8 +244,9 @@ class PreviewGaussianSplat(IO.ComfyNode):
         filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
         model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
 
+        viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
         camera_info_input = kwargs.get("camera_info", None)
-        camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
+        camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
         model_3d_info_input = kwargs.get("model_3d_info", None)
         model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
         return IO.NodeOutput(
@@ -265,6 +266,7 @@ class PreviewPointCloud(IO.ComfyNode):
             node_id="PreviewPointCloud",
             display_name="Preview Point Cloud",
             category="3d",
+            description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             search_aliases=[
@@ -303,8 +305,9 @@ class PreviewPointCloud(IO.ComfyNode):
         filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
         model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
 
+        viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
         camera_info_input = kwargs.get("camera_info", None)
-        camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
+        camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
         model_3d_info_input = kwargs.get("model_3d_info", None)
         model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
         return IO.NodeOutput(
@@ -375,8 +378,9 @@ class Load3DAdvanced(IO.ComfyNode):
         file_3d = None
         if model_file and model_file != "none":
             file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
+        viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
         model_3d_info = viewport_state.get('model_3d_info', [])
-        return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height)
+        return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height)
 
 
 class Load3DExtension(ComfyExtension):

comfy_extras/nodes_lt.py

@@ -50,8 +50,8 @@ class GetICLoRAParameters(io.ComfyNode):
         factor = 1
         if metadata:
             try:
-                factor = max(1, round(float(metadata.get("reference_downscale_factor", 1))))
-            except (TypeError, ValueError):
+                factor = max(1, round(float(next(v for k, v in metadata.items() if k.endswith("reference_downscale_factor")))))
+            except (StopIteration, TypeError, ValueError):
                 factor = 1
         parameters = {"reference_downscale_factor": factor}
         return io.NodeOutput(parameters)

comfy_extras/nodes_lt_audio.py

@@ -107,14 +107,17 @@ class LTXVEmptyLatentAudio(io.ComfyNode):
                     display_mode=io.NumberDisplay.number,
                     tooltip="Number of frames.",
                 ),
-                io.Int.Input(
-                    "frame_rate",
-                    default=25,
-                    min=1,
-                    max=1000,
-                    step=1,
-                    display_mode=io.NumberDisplay.number,
-                    tooltip="Number of frames per second.",
+                io.MultiType.Input(
+                    io.Float.Input(
+                        "frame_rate",
+                        default=25.0,
+                        min=1.0,
+                        max=1000.0,
+                        step=0.01,
+                        display_mode=io.NumberDisplay.number,
+                        tooltip="Number of frames per second.",
+                    ),
+                    [io.Int],
                 ),
                 io.Int.Input(
                     "batch_size",
@@ -137,7 +140,7 @@ class LTXVEmptyLatentAudio(io.ComfyNode):
     def execute(
         cls,
         frames_number: int,
-        frame_rate: int,
+        frame_rate: float,
         batch_size: int,
         audio_vae,
     ) -> io.NodeOutput:

comfy_extras/nodes_lt_upsampler.py

@@ -38,26 +38,20 @@ class LTXVLatentUpsampler(IO.ComfyNode):
         Returns:
             tuple: Tuple containing the upsampled latent
         """
-        device = model_management.get_torch_device()
-        memory_required = model_management.module_size(upscale_model)
-
-        model_dtype = next(upscale_model.parameters()).dtype
+        device = upscale_model.load_device
+        model = upscale_model.model
+        model_dtype = upscale_model.model_dtype()
         latents = samples["samples"]
         input_dtype = latents.dtype
 
-        memory_required += math.prod(latents.shape) * 3000.0  # TODO: more accurate
-        model_management.free_memory(memory_required, device)
+        memory_required = math.prod(latents.shape) * 3000.0  # TODO: more accurate
+        model_management.load_models_gpu([upscale_model], memory_required=memory_required)
 
-        try:
-            upscale_model.to(device)  # TODO: use the comfy model management system.
+        latents = latents.to(dtype=model_dtype, device=device)
 
-            latents = latents.to(dtype=model_dtype, device=device)
-
-            """Upsample latents without tiling."""
-            latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents)
-            upsampled_latents = upscale_model(latents)
-        finally:
-            upscale_model.cpu()
+        """Upsample latents without tiling."""
+        latents = vae.first_stage_model.per_channel_statistics.un_normalize(latents)
+        upsampled_latents = model(latents)
 
         upsampled_latents = vae.first_stage_model.per_channel_statistics.normalize(
             upsampled_latents

comfy_extras/nodes_mage.py

@@ -0,0 +1,103 @@
+from typing_extensions import override
+
+import comfy.utils
+import node_helpers
+import torch
+import comfy.model_management
+from comfy_api.latest import ComfyExtension, io
+
+
+class TextEncodeMageFlowEdit(io.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="TextEncodeMageFlowEdit",
+            category="model/conditioning/mage",
+            description="Encode an edit instruction with one or more reference images for Mage-Flow-Edit. Reference latents are resized to the output resolution (width/height, or the first image's size when 0). Use the latent output for sampling so the sizes always match.",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.String.Input("prompt", multiline=True, dynamic_prompts=True),
+                io.String.Input("negative_prompt", multiline=True, dynamic_prompts=True, advanced=True),
+                io.Vae.Input("vae", optional=True),
+                io.Autogrow.Input(
+                    "images",
+                    template=io.Autogrow.TemplateNames(
+                        io.Image.Input("image"),
+                        names=[f"image_{i}" for i in range(1, 17)],
+                        min=0,
+                    ),
+                    tooltip="Reference image(s) to edit. All references are resized to the output resolution before encoding.",
+                ),
+                io.Int.Input("width", default=0, min=0, max=8192, step=16, tooltip="Output width. 0 = use the first reference image's size."),
+                io.Int.Input("height", default=0, min=0, max=8192, step=16, tooltip="Output height. 0 = use the first reference image's size."),
+                io.Int.Input("batch_size", default=1, min=1, max=4096),
+            ],
+            outputs=[
+                io.Conditioning.Output(display_name="positive"),
+                io.Conditioning.Output(display_name="negative"),
+                io.Latent.Output(display_name="latent"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip, prompt, negative_prompt="", vae=None, images: io.Autogrow.Type = None, width=0, height=0, batch_size=1) -> io.NodeOutput:
+        ref_latents = []
+        images = images or {}
+        images = [images[name] for name in sorted(images, key=lambda n: int(n.rsplit("_", 1)[-1])) if images[name] is not None]
+        images_vl = []
+
+        # Output resolution: explicit width/height, else the primary reference's own size, floored to /16.
+        # Each dimension falls back independently so a 0 on one axis keeps an explicit value on the other.
+        if width == 0 or height == 0:
+            if len(images) > 0:
+                ref_h, ref_w = images[0].shape[1], images[0].shape[2]
+            else:
+                ref_h, ref_w = 1024, 1024
+            height = height or ref_h
+            width = width or ref_w
+        width = max(16, (width // 16) * 16)
+        height = max(16, (height // 16) * 16)
+
+        for image in images:
+            samples = image.movedim(-1, 1)
+
+            # VL conditioning copy: cap the long edge at 384 (training preprocessing).
+            long_edge = max(samples.shape[3], samples.shape[2])
+            if long_edge > 384:
+                scale_by = 384 / long_edge
+                s = comfy.utils.common_upscale(samples, max(1, round(samples.shape[3] * scale_by)), max(1, round(samples.shape[2] * scale_by)), "bicubic", "disabled")
+                images_vl.append(s.movedim(1, -1))
+            else:
+                images_vl.append(image)
+
+            if vae is not None:
+                # All references are resized to the output resolution before encoding, because Mage's RoPE aligns reference and target content by position
+                if samples.shape[3] != width or samples.shape[2] != height:
+                    s = comfy.utils.common_upscale(samples, width, height, "bicubic", "disabled")
+                else:
+                    s = samples
+                ref_latents.append(vae.encode(s.movedim(1, -1)[:, :, :, :3]))
+
+        # Negative branch keeps the same reference images (VL tokens + ref latents), only the instruction differs.
+        positive = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=images_vl))
+        negative = clip.encode_from_tokens_scheduled(clip.tokenize(negative_prompt if negative_prompt else " ", images=images_vl))
+
+        if len(ref_latents) > 0:
+            positive = node_helpers.conditioning_set_values(positive, {"reference_latents": ref_latents}, append=True)
+            negative = node_helpers.conditioning_set_values(negative, {"reference_latents": ref_latents}, append=True)
+
+        latent = torch.zeros([batch_size, 128, height // 16, width // 16], device=comfy.model_management.intermediate_device())
+        return io.NodeOutput(positive, negative, {"samples": latent})
+
+
+class MageExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            TextEncodeMageFlowEdit,
+        ]
+
+
+async def comfy_entrypoint() -> MageExtension:
+    return MageExtension()

comfy_extras/nodes_mask.py

@@ -419,17 +419,18 @@ class MaskPreview(IO.ComfyNode):
             search_aliases=["show mask", "view mask", "inspect mask", "debug mask"],
             display_name="Preview Mask",
             category="image/mask",
-            description="Saves the input images to your ComfyUI output directory.",
+            description="Preview the masks without saving them to the ComfyUI output directory.",
             inputs=[
                 IO.Mask.Input("mask"),
             ],
             hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
             is_output_node=True,
+            outputs=[IO.Mask.Output(display_name="mask")]
         )
 
     @classmethod
     def execute(cls, mask, filename_prefix="ComfyUI") -> IO.NodeOutput:
-        return IO.NodeOutput(ui=UI.PreviewMask(mask))
+        return IO.NodeOutput(mask, ui=UI.PreviewMask(mask))
 
 
 class MaskExtension(ComfyExtension):

comfy_extras/nodes_model_patch.py

@@ -8,6 +8,8 @@ import comfy.ldm.common_dit
 import comfy.latent_formats
 import comfy.ldm.lumina.controlnet
 import comfy.ldm.supir.supir_modules
+import comfy.ldm.anima.lllite
+import comfy.ldm.wan.uni3c
 from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
 from comfy_api.latest import io
 from comfy.ldm.supir.supir_patch import SUPIRPatch
@@ -236,10 +238,12 @@ class ModelPatchLoader:
 
     def load_model_patch(self, name):
         model_patch_path = folder_paths.get_full_path_or_raise("model_patches", name)
-        sd = comfy.utils.load_torch_file(model_patch_path, safe_load=True)
+        sd, metadata = comfy.utils.load_torch_file(model_patch_path, safe_load=True, return_metadata=True)
         dtype = comfy.utils.weight_dtype(sd)
 
-        if 'controlnet_blocks.0.y_rms.weight' in sd:
+        if 'lllite_conditioning1.conv1.weight' in sd:
+            model = comfy.ldm.anima.lllite.AnimaLLLite(sd, metadata, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
+        elif 'controlnet_blocks.0.y_rms.weight' in sd:
             additional_in_dim = sd["img_in.weight"].shape[1] - 64
             model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
         elif 'feature_embedder.mid_layer_norm.bias' in sd:
@@ -261,6 +265,37 @@ class ModelPatchLoader:
                     if torch.count_nonzero(ref_weight) == 0:
                         config['broken'] = True
             model = comfy.ldm.lumina.controlnet.ZImage_Control(device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast, **config)
+        elif 'controlnet_patch_embedding.weight' in sd:  # Uni3C controlnet for Wan
+            attn_key_replace = {".self_attn.to_q.": ".self_attn.q.",
+                                ".self_attn.to_k.": ".self_attn.k.",
+                                ".self_attn.to_v.": ".self_attn.v.",
+                                ".self_attn.to_out.0.": ".self_attn.o."}
+            converted_sd = {}
+            for k, w in sd.items():
+                for r, rr in attn_key_replace.items():
+                    k = k.replace(r, rr)
+                converted_sd[k] = w
+            sd = converted_sd
+
+            num_layers = sum(1 for k in sd if k.startswith("proj_out.") and k.endswith(".weight"))
+            conv_out_dim = sd["controlnet_patch_embedding.weight"].shape[0]
+            if "proj_in.weight" in sd:
+                dim = sd["proj_in.weight"].shape[0]
+            else:
+                dim = conv_out_dim
+            model = comfy.ldm.wan.uni3c.WanUni3CControlnet(
+                    in_channels=sd["controlnet_patch_embedding.weight"].shape[1],
+                    conv_out_dim=conv_out_dim,
+                    dim=dim,
+                    ffn_dim=sd["controlnet_blocks.0.ffn.0.bias"].shape[0],
+                    num_layers=num_layers,
+                    time_embed_dim=sd["controlnet_blocks.0.norm1.linear.weight"].shape[1],
+                    out_proj_dim=sd["proj_out.0.weight"].shape[0],
+                    add_channels=sd["controlnet_mask_embedding.mask_proj.0.weight"].shape[1],
+                    mid_channels=sd["controlnet_mask_embedding.mask_proj.0.weight"].shape[0],
+                    device=comfy.model_management.unet_offload_device(),
+                    dtype=dtype,
+                    operations=comfy.ops.manual_cast)
         elif "audio_proj.proj1.weight" in sd:
             model = MultiTalkModelPatch(
                     audio_window=5, context_tokens=32, vae_scale=4,
@@ -296,6 +331,50 @@ class ModelPatchLoader:
         return (model_patcher,)
 
 
+class AnimaLLLiteApply:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {"model": ("MODEL",),
+                             "model_patch": ("MODEL_PATCH",),
+                             "image": ("IMAGE",),
+                             "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                             "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                             },
+                "optional": {"mask": ("MASK",),
+                             }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply_patch"
+    EXPERIMENTAL = True
+
+    CATEGORY = "model_patches/anima"
+
+    def apply_patch(self, model, model_patch, image, strength, start_percent, end_percent, mask=None):
+        image = image[..., :3]
+
+        if model_patch.model.cond_in_channels == 4 and mask is None:
+            mask = torch.zeros_like(image[..., 0])
+        elif model_patch.model.cond_in_channels != 4:
+            mask = None
+
+        model_sampling = model.get_model_object("model_sampling")
+        sigma_start = float(model_sampling.percent_to_sigma(start_percent))
+        sigma_end = float(model_sampling.percent_to_sigma(end_percent))
+        patch = comfy.ldm.anima.lllite.AnimaLLLitePatch(model_patch, image, mask, strength, sigma_start, sigma_end)
+        model_patched = model.clone()
+        model_patched.set_model_post_input_patch(patch)
+        model_patched.set_model_attn1_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(
+            patch,
+            {"q": "self_attn_q_proj", "k": "self_attn_k_proj", "v": "self_attn_v_proj"},
+        ))
+        model_patched.set_model_attn2_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(
+            patch,
+            {"q": "cross_attn_q_proj"},
+        ))
+        model_patched.set_model_patch(comfy.ldm.anima.lllite.AnimaLLLiteMLPPatch(patch), "mlp_patch")
+        return (model_patched,)
+
+
 class DiffSynthCnetPatch:
     def __init__(self, model_patch, vae, image, strength, mask=None):
         self.model_patch = model_patch
@@ -514,6 +593,150 @@ class ZImageFunControlnet(QwenImageDiffsynthControlnet):
 
     CATEGORY = "model/patch/z-image"
 
+class WanUni3CCnetPatch:
+    def __init__(self, model_patch, render_video, vae, latent_format, strength, sigma_start, sigma_end):
+        self.model_patch = model_patch
+        self.render_video = render_video
+        self.vae = vae
+        self.latent_format = latent_format
+        self.strength = strength
+        self.sigma_start = sigma_start
+        self.sigma_end = sigma_end
+        self.prepared_render = None
+        self.temp_data = None
+
+    def encode_render_video(self, target_latent_shape):
+        t_len, h_len, w_len = target_latent_shape
+        temporal_compression = self.vae.temporal_compression_decode() or 1
+        spatial_compression = self.vae.spacial_compression_encode()
+        target_frames = (t_len - 1) * temporal_compression + 1
+        target_height = h_len * spatial_compression
+        target_width = w_len * spatial_compression
+
+        frames = self.render_video
+        if frames.shape[0] > target_frames:
+            frames = frames[:target_frames]
+        elif frames.shape[0] < target_frames:
+            last_frame = frames[-1:].expand(target_frames - frames.shape[0], -1, -1, -1)
+            frames = torch.cat([frames, last_frame], dim=0)
+
+        if frames.shape[1] != target_height or frames.shape[2] != target_width:
+            frames = comfy.utils.common_upscale(frames.movedim(-1, 1), target_width, target_height, "bilinear", "center").movedim(1, -1)
+
+        loaded_models = comfy.model_management.loaded_models(only_currently_used=True)
+        render_latent = self.vae.encode(frames)
+        comfy.model_management.load_models_gpu(loaded_models)
+        return self.latent_format.process_in(render_latent)
+
+    def build_controlnet_input(self, x, dtype, samples_per_cond):
+        # first 20 channels of the model input: noise latent + I2V mask (zero padded for T2V)
+        hidden = x[:samples_per_cond, :20].to(dtype)
+        if hidden.shape[1] < 20:
+            pad_shape = list(hidden.shape)
+            pad_shape[1] = 20 - hidden.shape[1]
+            hidden = torch.cat([hidden, torch.zeros(pad_shape, dtype=hidden.dtype, device=hidden.device)], dim=1)
+
+        render = self.prepared_render
+        if render is None or render.shape[2:] != hidden.shape[2:]:
+            render = self.encode_render_video(hidden.shape[2:])
+        render = render.to(device=hidden.device, dtype=dtype)
+        self.prepared_render = render
+        if render.shape[0] != hidden.shape[0]:
+            render = render.expand(hidden.shape[0], -1, -1, -1, -1)
+        return torch.cat([hidden, render], dim=1)
+
+    def __call__(self, kwargs):
+        img = kwargs.get("img")
+        block_index = kwargs.get("block_index")
+        transformer_options = kwargs.get("transformer_options", {})
+
+        if block_index == 0:
+            self.temp_data = None
+            active = True
+            sigmas = transformer_options.get("sigmas", None)
+            if sigmas is not None:
+                sigma = sigmas[0].item()
+                if sigma > self.sigma_start or sigma < self.sigma_end:
+                    active = False
+            if active:
+                x = kwargs.get("x")
+                # cond and uncond chunks share latents, so we can reuse residuals
+                num_conds = len(transformer_options.get("cond_or_uncond", [0]))
+                samples_per_cond = x.shape[0]
+                if num_conds > 0 and x.shape[0] % num_conds == 0:
+                    samples_per_cond = x.shape[0] // num_conds
+                temb = kwargs.get("vec")[:samples_per_cond]
+                if temb.ndim == 3:
+                    temb = temb[:, 0]
+                model = self.model_patch.model
+                controlnet_input = self.build_controlnet_input(x, img.dtype, samples_per_cond)
+                hidden, freqs = model.process_input(controlnet_input)
+                self.temp_data = (hidden, temb.to(img.dtype), freqs)
+
+        num_layers = self.model_patch.model.num_layers
+        if self.temp_data is not None and block_index < num_layers:
+            hidden, temb, freqs = self.temp_data
+            hidden, residual = self.model_patch.model.forward_block(block_index, hidden, temb, freqs)
+            residual = residual.to(img.dtype) * self.strength
+            if residual.shape[0] != img.shape[0]:
+                residual = residual.repeat(img.shape[0] // residual.shape[0], 1, 1)
+            img_offset = kwargs.get("img_offset", 0)
+            img[:, img_offset:img_offset + residual.shape[1]] += residual
+            if block_index >= num_layers - 1:
+                self.temp_data = None
+            else:
+                self.temp_data = (hidden, temb, freqs)
+
+        return kwargs
+
+    def to(self, device_or_dtype):
+        if isinstance(device_or_dtype, torch.device):
+            if self.prepared_render is not None:
+                self.prepared_render = self.prepared_render.to(device_or_dtype)
+            self.temp_data = None
+        return self
+
+    def models(self):
+        return [self.model_patch]
+
+
+class WanUni3CControlnetApply:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": { "model": ("MODEL",),
+                              "model_patch": ("MODEL_PATCH",),
+                              "vae": ("VAE",),
+                              "render_video": ("IMAGE", {"tooltip": "The guidance video rendered from the camera trajectory, most commonly warped point cloud renders of the input image."}),
+                              "strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
+                              "start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                              "end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
+                              }}
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "apply_patch"
+    EXPERIMENTAL = True
+
+    CATEGORY = "model/patch/wan"
+
+    def apply_patch(self, model, model_patch, vae, render_video, strength, start_percent, end_percent):
+        if not isinstance(model_patch.model, comfy.ldm.wan.uni3c.WanUni3CControlnet):
+            raise ValueError("The connected model patch is not a Uni3C ControlNet.")
+        cnet_dim = model_patch.model.controlnet_blocks[0].norm1.linear.in_features
+        model_dim = getattr(model.get_model_object("diffusion_model"), "dim", None)
+        if model_dim is None:
+            raise ValueError("The Uni3C ControlNet only works with Wan models.")
+        if model_dim != cnet_dim:
+            raise ValueError("This Uni3C ControlNet expects a Wan model with dim {}, the loaded model has dim {}.".format(cnet_dim, model_dim))
+
+        model_patched = model.clone()
+        model_sampling = model.get_model_object("model_sampling")
+        sigma_start = model_sampling.percent_to_sigma(start_percent)
+        sigma_end = model_sampling.percent_to_sigma(end_percent)
+        latent_format = model.get_model_object("latent_format")
+        patch = WanUni3CCnetPatch(model_patch, render_video[:, :, :, :3], vae, latent_format, strength, sigma_start, sigma_end)
+        model_patched.set_model_double_block_patch(patch)
+        return (model_patched,)
+
+
 class UsoStyleProjectorPatch:
     def __init__(self, model_patch, encoded_image):
         self.model_patch = model_patch
@@ -672,14 +895,18 @@ NODE_CLASS_MAPPINGS = {
     "ModelPatchLoader": ModelPatchLoader,
     "QwenImageDiffsynthControlnet": QwenImageDiffsynthControlnet,
     "ZImageFunControlnet": ZImageFunControlnet,
+    "WanUni3CControlnetApply": WanUni3CControlnetApply,
     "USOStyleReference": USOStyleReference,
     "SUPIRApply": SUPIRApply,
+    "AnimaLLLiteApply": AnimaLLLiteApply,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
     "ModelPatchLoader": "Load Model Patch",
     "QwenImageDiffsynthControlnet": "Apply Qwen Image DiffSynth ControlNet",
     "ZImageFunControlnet": "Apply Z-Image Fun ControlNet",
+    "WanUni3CControlnetApply": "Apply Wan Uni3C ControlNet",
     "USOStyleReference": "Apply USO Style Reference",
     "SUPIRApply": "Apply SUPIR Patch",
+    "AnimaLLLiteApply": "Apply Anima LLLite",
 }

comfy_extras/nodes_preview_any.py

@@ -18,6 +18,7 @@ class PreviewAny():
 
     CATEGORY = "utilities"
     SEARCH_ALIASES = ["show output", "inspect", "debug", "print value", "show text"]
+    DESCRIPTION = "Preview any input value as text."
 
     def main(self, source=None):
         torch.set_printoptions(edgeitems=6)

comfy_extras/nodes_primitive.py

@@ -10,11 +10,10 @@ class String(io.ComfyNode):
         return io.Schema(
             node_id="PrimitiveString",
             search_aliases=["text", "string", "text box", "prompt"],
-            display_name="Text String (DEPRECATED)",
+            display_name="Text",
             category="utilities/primitive",
             inputs=[io.String.Input("value")],
-            outputs=[io.String.Output()],
-            is_deprecated=True
+            outputs=[io.String.Output()]
         )
 
     @classmethod
@@ -28,7 +27,7 @@ class StringMultiline(io.ComfyNode):
         return io.Schema(
             node_id="PrimitiveStringMultiline",
             search_aliases=["text", "string", "text multiline", "string multiline", "text box", "prompt"],
-            display_name="Input Text",
+            display_name="Text (Multiline)",
             category="utilities/primitive",
             essentials_category="Basics",
             inputs=[io.String.Input("value", multiline=True)],

comfy_extras/nodes_save_3d.py

@@ -13,7 +13,7 @@ from typing_extensions import override
 
 import folder_paths
 from comfy.cli_args import args
-from comfy_api.latest import ComfyExtension, IO, Types
+from comfy_api.latest import ComfyExtension, IO, Types, UI
 
 
 def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False):
@@ -406,10 +406,165 @@ class SaveGLB(IO.ComfyNode):
         return IO.NodeOutput(ui={"3d": results})
 
 
+def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str:
+    full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
+        filename_prefix, folder_paths.get_output_directory()
+    )
+    ext = model_3d.format or "glb"
+    saved_filename = f"{filename}_{counter:05}.{ext}"
+    model_3d.save_to(os.path.join(full_output_folder, saved_filename))
+    return f"{subfolder}/{saved_filename}" if subfolder else saved_filename
+
+
+def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
+    model_file = _save_file3d_to_output(model_3d, filename_prefix)
+    viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
+    camera_info_input = kwargs.get("camera_info", None)
+    camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
+    model_3d_info_input = kwargs.get("model_3d_info", None)
+    model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
+    return IO.NodeOutput(
+        model_3d,
+        model_3d_info,
+        camera_info,
+        width,
+        height,
+        ui=UI.PreviewUI3DAdvanced(model_file, camera_info, model_3d_info),
+    )
+
+
+class Save3DAdvanced(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="Save3DAdvanced",
+            display_name="Save 3D (Advanced)",
+            search_aliases=["save 3d", "export 3d model", "save mesh advanced"],
+            category="3d",
+            is_experimental=True,
+            is_output_node=True,
+            inputs=[
+                IO.MultiType.Input(
+                    "model_3d",
+                    types=[
+                        IO.File3DGLB,
+                        IO.File3DGLTF,
+                        IO.File3DFBX,
+                        IO.File3DOBJ,
+                        IO.File3DSTL,
+                        IO.File3DUSDZ,
+                        IO.File3DAny,
+                    ],
+                    tooltip="3D model file from an upstream 3D node.",
+                ),
+                IO.String.Input("filename_prefix", default="3d/ComfyUI"),
+                IO.Load3D.Input("viewport_state"),
+                IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
+                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
+                IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
+                IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
+            ],
+            outputs=[
+                IO.File3DAny.Output(display_name="model_3d"),
+                IO.Load3DModelInfo.Output(display_name="model_3d_info"),
+                IO.Load3DCamera.Output(display_name="camera_info"),
+                IO.Int.Output(display_name="width"),
+                IO.Int.Output(display_name="height"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
+        return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
+
+
+class SaveGaussianSplat(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="SaveGaussianSplat",
+            display_name="Save Splat",
+            search_aliases=["save splat", "save gaussian splat", "export gaussian", "export splat"],
+            category="3d",
+            is_experimental=True,
+            is_output_node=True,
+            inputs=[
+                IO.MultiType.Input(
+                    "model_3d",
+                    types=[
+                        IO.File3DSplatAny,
+                        IO.File3DPLY,
+                        IO.File3DSPLAT,
+                        IO.File3DSPZ,
+                        IO.File3DKSPLAT,
+                    ],
+                    tooltip="A gaussian splat 3D file.",
+                ),
+                IO.String.Input("filename_prefix", default="3d/ComfyUI"),
+                IO.Load3D.Input("viewport_state"),
+                IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
+                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
+                IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
+                IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
+            ],
+            outputs=[
+                IO.File3DSplatAny.Output(display_name="model_3d"),
+                IO.Load3DModelInfo.Output(display_name="model_3d_info"),
+                IO.Load3DCamera.Output(display_name="camera_info"),
+                IO.Int.Output(display_name="width"),
+                IO.Int.Output(display_name="height"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
+        return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
+
+
+class SavePointCloud(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="SavePointCloud",
+            display_name="Save Point Cloud",
+            search_aliases=["save point cloud", "save pointcloud", "export point cloud"],
+            category="3d",
+            is_experimental=True,
+            is_output_node=True,
+            inputs=[
+                IO.MultiType.Input(
+                    "model_3d",
+                    types=[
+                        IO.File3DPointCloudAny,
+                        IO.File3DPLY,
+                    ],
+                    tooltip="Point cloud file (.ply)",
+                ),
+                IO.String.Input("filename_prefix", default="3d/ComfyUI"),
+                IO.Load3D.Input("viewport_state"),
+                IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
+                IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
+                IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
+                IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
+            ],
+            outputs=[
+                IO.File3DPointCloudAny.Output(display_name="model_3d"),
+                IO.Load3DModelInfo.Output(display_name="model_3d_info"),
+                IO.Load3DCamera.Output(display_name="camera_info"),
+                IO.Int.Output(display_name="width"),
+                IO.Int.Output(display_name="height"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
+        return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
+
+
 class Save3DExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
-        return [SaveGLB]
+        return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud]
 
 
 async def comfy_entrypoint() -> Save3DExtension: