Add MiniMax Music 3 to readme. (#15656)

Signed-off-by: Alexander Piskun <bigcat88@icloud.com>

.coderabbit.yaml

@@ -11,6 +11,10 @@ reviews:
   review_status: false
   review_details: true
   commit_status: true
+  # Without this, a review that never happened (rate limit, internal error)
+  # still posts a green "CodeRabbit" commit status, so a throttled review is
+  # indistinguishable from a clean one.
+  fail_commit_status: true
   collapse_walkthrough: true
   changed_files_summary: false
   sequence_diagrams: false

AGENTS.md

@@ -305,6 +305,12 @@
 
 - Follow existing node conventions: `INPUT_TYPES`, `RETURN_TYPES`, `FUNCTION`,
   `CATEGORY`, and registration through the local mapping used by that file.
+- Treat legacy combo inputs, `io.Combo`, and `io.DynamicCombo` values as
+  untrusted when they affect filesystem access. Any value used as a file or
+  folder name, path component, format, or extension must be validated again at
+  the load/save boundary using an existing `folder_paths` resolver or
+  containment helper, or a fixed allowlist/mapping. Do not rely only on the
+  advertised combo options or prompt validation.
 - Keep node changes backward compatible by default. Add inputs with sensible
   defaults and avoid changing output types unless the request requires it.
 - Model implementations should add the minimal number of ComfyUI nodes required

README.md

@@ -37,7 +37,7 @@
 
 ComfyUI is the AI creation engine for visual professionals who demand control over every model, every parameter, and every output. Its powerful and modular node graph interface empowers creatives to generate images, videos, 3D models, audio, and more...
 - ComfyUI natively supports the latest open-source state of the art models.
-- API nodes provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
+- [Partner nodes](https://docs.comfy.org/tutorials/partner-nodes/overview#partner-nodes) provide access to the best closed source models such as Nano Banana, Seedance, Hunyuan3D, etc.
 - It is available on Windows, Linux, and macOS, locally with our [desktop application](https://www.comfy.org/download), our [portable install](#installing) or on our [cloud](https://www.comfy.org/cloud).
 - The most sophisticated workflows can be exposed through a simple UI thanks to App Mode.
 - It integrates seamlessly into production pipelines with our API endpoints.
@@ -66,67 +66,24 @@ Supports all operating systems and GPU types (NVIDIA, AMD, Intel, Apple Silicon,
 See what ComfyUI can do with the [newer template workflows](https://comfy.org/workflows) or old [example workflows](https://comfyanonymous.github.io/ComfyUI_examples/).
 
 ## Features
-- Nodes/graph/flowchart interface to experiment and create complex Stable Diffusion workflows without needing to code anything.
-- NOTE: There are many more models supported than the list below, if you want to see what is supported see our templates list inside ComfyUI.
-- Image Models
-   - SD1.x, SD2.x ([unCLIP](https://comfyanonymous.github.io/ComfyUI_examples/unclip/))
-   - [SDXL](https://comfyanonymous.github.io/ComfyUI_examples/sdxl/), [SDXL Turbo](https://comfyanonymous.github.io/ComfyUI_examples/sdturbo/)
-   - [Stable Cascade](https://comfyanonymous.github.io/ComfyUI_examples/stable_cascade/)
-   - [SD3 and SD3.5](https://comfyanonymous.github.io/ComfyUI_examples/sd3/)
-   - Pixart Alpha and Sigma
-   - [AuraFlow](https://comfyanonymous.github.io/ComfyUI_examples/aura_flow/)
-   - [HunyuanDiT](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_dit/)
-   - [Flux](https://comfyanonymous.github.io/ComfyUI_examples/flux/)
-   - [Lumina Image 2.0](https://comfyanonymous.github.io/ComfyUI_examples/lumina2/)
-   - [HiDream](https://comfyanonymous.github.io/ComfyUI_examples/hidream/)
-   - [Qwen Image](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/)
-   - [Hunyuan Image 2.1](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_image/)
-   - [Flux 2](https://comfyanonymous.github.io/ComfyUI_examples/flux2/)
-   - [Z Image](https://comfyanonymous.github.io/ComfyUI_examples/z_image/)
-   - Ernie Image
-- Image Editing Models
-   - [Omnigen 2](https://comfyanonymous.github.io/ComfyUI_examples/omnigen/)
-   - [Flux Kontext](https://comfyanonymous.github.io/ComfyUI_examples/flux/#flux-kontext-image-editing-model)
-   - [HiDream E1.1](https://comfyanonymous.github.io/ComfyUI_examples/hidream/#hidream-e11)
-   - [Qwen Image Edit](https://comfyanonymous.github.io/ComfyUI_examples/qwen_image/#edit-model)
-- Video Models
-   - [Stable Video Diffusion](https://comfyanonymous.github.io/ComfyUI_examples/video/)
-   - [Mochi](https://comfyanonymous.github.io/ComfyUI_examples/mochi/)
-   - [LTX-Video](https://comfyanonymous.github.io/ComfyUI_examples/ltxv/)
-   - [Hunyuan Video](https://comfyanonymous.github.io/ComfyUI_examples/hunyuan_video/)
-   - [Wan 2.1](https://comfyanonymous.github.io/ComfyUI_examples/wan/)
-   - [Wan 2.2](https://comfyanonymous.github.io/ComfyUI_examples/wan22/)
-   - [Hunyuan Video 1.5](https://docs.comfy.org/tutorials/video/hunyuan/hunyuan-video-1-5)
-- Audio Models
-   - [Stable Audio](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
-   - [ACE Step](https://comfyanonymous.github.io/ComfyUI_examples/audio/)
-- 3D Models
-   - [Hunyuan3D 2.0](https://docs.comfy.org/tutorials/3d/hunyuan3D-2)
-- Asynchronous Queue system
-- Many optimizations: Only re-executes the parts of the workflow that changes between executions.
-- Smart memory management: can automatically run large models on GPUs with as low as 1GB vram with smart offloading.
-- Works even if you don't have a GPU with: ```--cpu``` (slow)
-- Can load ckpt and safetensors: All in one checkpoints or standalone diffusion models, VAEs and CLIP models.
-- Safe loading of ckpt, pt, pth, etc.. files.
-- Embeddings/Textual inversion
-- [Loras (regular, locon and loha)](https://comfyanonymous.github.io/ComfyUI_examples/lora/)
-- [Hypernetworks](https://comfyanonymous.github.io/ComfyUI_examples/hypernetworks/)
-- Loading full workflows (with seeds) from generated PNG, WebP and FLAC files.
-- Saving/Loading workflows as Json files.
-- Nodes interface can be used to create complex workflows like one for [Hires fix](https://comfyanonymous.github.io/ComfyUI_examples/2_pass_txt2img/) or much more advanced ones.
-- [Area Composition](https://comfyanonymous.github.io/ComfyUI_examples/area_composition/)
-- [Inpainting](https://comfyanonymous.github.io/ComfyUI_examples/inpaint/) with both regular and inpainting models.
-- [ControlNet and T2I-Adapter](https://comfyanonymous.github.io/ComfyUI_examples/controlnet/)
-- [Upscale Models (ESRGAN, ESRGAN variants, SwinIR, Swin2SR, etc...)](https://comfyanonymous.github.io/ComfyUI_examples/upscale_models/)
-- [GLIGEN](https://comfyanonymous.github.io/ComfyUI_examples/gligen/)
-- [Model Merging](https://comfyanonymous.github.io/ComfyUI_examples/model_merging/)
-- [LCM models and Loras](https://comfyanonymous.github.io/ComfyUI_examples/lcm/)
-- Latent previews with [TAESD](#how-to-show-high-quality-previews)
-- Works fully offline: core will never download anything unless you want to.
-- Optional API nodes to use paid models from external providers through the online [Comfy API](https://docs.comfy.org/tutorials/api-nodes/overview) disable with: `--disable-api-nodes`
-- [Config file](extra_model_paths.yaml.example) to set the search paths for models.
+- A visual node graph for building and reusing image, video, audio, 3D, and text workflows without code.
+- Reusable subgraphs, workflow templates, App Mode, and a local API for integrating workflows into applications.
+- Efficient local execution with asynchronous queueing, partial graph re-execution, smart VRAM and RAM management, model offloading, and support for quantized models.
+- Broad native model support. This is a representative list; browse the [workflow library](https://comfy.org/workflows/) for maintained, ready-to-run templates.
+  - [Image generation](https://comfy.org/workflows/tag/text-to-image/): Stable Diffusion 1.5, SDXL, SD3.5, Flux.1, Flux.2, Qwen Image, Z-Image, Hunyuan Image 2.1, HiDream, Lumina Image 2.0, Chroma, Anima, LongCat Image, Ideogram 4, Krea 2, MageFlow, Microsoft Lens, PixelDiT, Kandinsky 5, and Ernie Image.
+  - [Image editing](https://comfy.org/workflows/tag/image-edit/): Flux Kontext, Flux.2 Klein, Qwen Image Edit, HiDream E1.1 and O1, OmniGen2, Boogu, JoyImage Edit, MageFlow Edit, and LongCat Image Edit.
+  - [Video generation](https://comfy.org/workflows/tag/video-generation/): Wan 2.1 and 2.2, LTX-Video 2 and 2.3, HunyuanVideo 1.5, Kandinsky 5 Video, CogVideoX, Cosmos Predict2, Bernini-R, SCAIL 2, and Mochi.
+  - [Audio and video generation](https://comfy.org/workflows/): MiniMax H3 and LTX-AV.
+  - [Audio generation](https://comfy.org/workflows/tag/text-to-audio/): ACE-Step 1.5, Stable Audio 3 and MiniMax Music 3
+  - [3D and vision](https://comfy.org/workflows/): Hunyuan3D 2.1, TripoSplat, SeedVR2, SUPIR, Depth Anything 3, MoGe, SAM 3 and 3.1, RT-DETRv4, and BiRefNet.
+  - [Text generation](https://comfy.org/workflows/tag/text-generation/): Gemma 3 and 4, Qwen3, Qwen3.5, and Qwen3-VL, including multimodal inputs.
+- Load complete checkpoints or separate diffusion models, VAEs, text encoders, LoRAs, ControlNets, adapters, and upscalers from supported model formats.
+- Built-in tools for inpainting, outpainting, reference conditioning, masks and compositing, model merging, upscaling, frame interpolation, segmentation, depth estimation, and media processing.
+- Save and load workflows as JSON, or recover complete workflows and seeds from supported generated media.
+- Runs fully offline: core does not download anything unless you request it. Use `--disable-api-nodes` to disable the optional paid [Comfy API nodes](https://docs.comfy.org/tutorials/api-nodes/overview) and force all built-in functionality to stay offline.
+- Extend ComfyUI with custom nodes
+- Configure additional model locations with [`extra_model_paths.yaml`](extra_model_paths.yaml.example).
 
-Workflow examples can be found on the [Examples page](https://comfyanonymous.github.io/ComfyUI_examples/)
 
 ## Release Process
 
@@ -188,11 +145,19 @@ ComfyUI follows a weekly release cycle targeting Monday but this regularly chang
 
 # Installing
 
+## Windows and Mac
+
+We highly recommend using the [desktop app](https://comfy.org/download):
+
+### [Link to Download](https://comfy.org/download)
+
+The desktop app is the easiest and best way to use ComfyUI for new users.
+
 ## Windows Portable
 
-There is a portable standalone build for Windows that should work for running on Nvidia GPUs or for running on your CPU only on the [releases page](https://github.com/comfyanonymous/ComfyUI/releases).
+There is a portable standalone build for Windows that should work for running on Nvidia GPUs or for running on your CPU only. It is not recommended for regular users. Regular users should use the desktop app above.
 
-### [Direct link to download](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z)
+[Direct link to download (nvidia)](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z)
 
 Simply download, extract with [7-Zip](https://7-zip.org) or with the windows explorer on recent windows versions and run. For smaller models you normally only need to put the checkpoints (the huge ckpt/safetensors files) in: ComfyUI\models\checkpoints but many of the larger models have multiple files. Make sure to follow the instructions to know which subfolder to put them in ComfyUI\models\
 
@@ -208,7 +173,7 @@ The portable above currently comes with python 3.13 and pytorch cuda 13.0. Updat
 
 [Portable for Nvidia GPUs](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia.7z) (supports 20 series and above).
 
-[Portable for Nvidia GPUs with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs).
+[Portable for Nvidia GPUs with pytorch cuda 12.6 and python 3.12](https://github.com/comfyanonymous/ComfyUI/releases/latest/download/ComfyUI_windows_portable_nvidia_cu126.7z) (Supports Nvidia 10 series and older GPUs, DO NOT USE THIS ON NEWER 20 SERIES AND ABOVE GPUS).
 
 #### How do I share models between another UI and ComfyUI?
 
@@ -229,7 +194,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia
 
 Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12
 
-torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
+torch 2.7 is minimally supported but using a newer version is extremely recommended. Using a cu130 or above version of pytorch is required on Nvidia 20 series and above. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it.
 
 ### Instructions:
 

app/assets/api/routes.py

@@ -18,7 +18,7 @@ from app.assets.api.schemas_in import (
     AssetValidationError,
     UploadError,
 )
-from app.assets.helpers import validate_blake3_hash
+from app.assets.helpers import normalize_tags, validate_blake3_hash
 from app.assets.api.upload import (
     delete_temp_file_if_exists,
     parse_multipart_upload,
@@ -117,6 +117,87 @@ def _build_validation_error_response(code: str, ve: ValidationError) -> web.Resp
     return _build_error_response(400, code, "Validation failed.", {"errors": errors})
 
 
+class InvalidTagFilterError(Exception):
+    """Invalid combination of tag-filter query parameters."""
+
+    def __init__(self, message: str, details: dict):
+        super().__init__(message)
+        self.details = details
+
+
+# Caps the per-tag EXISTS fan-out; deliberately covers the legacy spellings too.
+MAX_TAG_FILTER_TAGS = 100
+
+
+def _resolve_tag_filters(
+    q: schemas_in.ListAssetsQuery | schemas_in.TagsRefineQuery,
+) -> tuple[list[str], list[str], list[str]]:
+    """Resolve legacy (include/exclude) and new (all/any/none) tag-filter
+    spellings into effective (all, any, none) lists.
+
+    Combination validation applies only when the request uses at least one
+    new-name parameter (non-empty after normalisation); requests using only
+    the legacy names keep their historical behaviour, including degenerate
+    combinations like include_tags=a&exclude_tags=a.
+    """
+    # model_dump, not attribute access: deprecated fields warn on every attribute read.
+    legacy = q.model_dump(include={"include_tags", "exclude_tags"})
+    include_tags = normalize_tags(legacy["include_tags"])
+    exclude_tags = normalize_tags(legacy["exclude_tags"])
+    tags_all = normalize_tags(q.tags_all)
+    tags_any = normalize_tags(q.tags_any)
+    tags_none = normalize_tags(q.tags_none)
+
+    for param_name, values in (
+        ("include_tags", include_tags),
+        ("exclude_tags", exclude_tags),
+        ("tags_all", tags_all),
+        ("tags_any", tags_any),
+        ("tags_none", tags_none),
+    ):
+        if len(values) > MAX_TAG_FILTER_TAGS:
+            raise InvalidTagFilterError(
+                f"'{param_name}' lists {len(values)} tags; the maximum is "
+                f"{MAX_TAG_FILTER_TAGS}.",
+                {
+                    "parameter": param_name,
+                    "count": len(values),
+                    "max": MAX_TAG_FILTER_TAGS,
+                },
+            )
+
+    if not (tags_all or tags_any or tags_none):
+        return include_tags, [], exclude_tags
+
+    if include_tags and tags_all:
+        raise InvalidTagFilterError(
+            "Cannot combine 'include_tags' with 'tags_all'; use 'tags_all'.",
+            {"parameters": ["include_tags", "tags_all"]},
+        )
+    if exclude_tags and tags_none:
+        raise InvalidTagFilterError(
+            "Cannot combine 'exclude_tags' with 'tags_none'; use 'tags_none'.",
+            {"parameters": ["exclude_tags", "tags_none"]},
+        )
+
+    all_param, all_list = (
+        ("tags_all", tags_all) if tags_all else ("include_tags", include_tags)
+    )
+    none_param, none_list = (
+        ("tags_none", tags_none) if tags_none else ("exclude_tags", exclude_tags)
+    )
+
+    conflicting = sorted(set(all_list) & set(none_list))
+    if conflicting:
+        raise InvalidTagFilterError(
+            f"Query can never match: {', '.join(repr(t) for t in conflicting)} "
+            f"required by '{all_param}' but rejected by '{none_param}'.",
+            {"conflicting_tags": conflicting, "parameters": [all_param, none_param]},
+        )
+
+    return all_list, tags_any, none_list
+
+
 def _validate_sort_field(requested: str | None) -> str:
     if not requested:
         return "created_at"
@@ -217,6 +298,11 @@ async def list_assets_route(request: web.Request) -> web.Response:
     except ValidationError as ve:
         return _build_validation_error_response("INVALID_QUERY", ve)
 
+    try:
+        tags_all, tags_any, tags_none = _resolve_tag_filters(q)
+    except InvalidTagFilterError as e:
+        return _build_error_response(400, "INVALID_TAG_FILTER", str(e), e.details)
+
     sort = _validate_sort_field(q.sort)
     order_candidate = (q.order or "desc").lower()
     order = order_candidate if order_candidate in {"asc", "desc"} else "desc"
@@ -224,8 +310,9 @@ async def list_assets_route(request: web.Request) -> web.Response:
     try:
         result = list_assets_page(
             owner_id=USER_MANAGER.get_request_user_id(request),
-            include_tags=q.include_tags,
-            exclude_tags=q.exclude_tags,
+            include_tags=tags_all,
+            exclude_tags=tags_none,
+            any_tags=tags_any,
             name_contains=q.name_contains,
             metadata_filter=q.metadata_filter,
             limit=q.limit,
@@ -715,10 +802,16 @@ async def get_tags_refine(request: web.Request) -> web.Response:
     except ValidationError as ve:
         return _build_validation_error_response("INVALID_QUERY", ve)
 
+    try:
+        tags_all, tags_any, tags_none = _resolve_tag_filters(q)
+    except InvalidTagFilterError as e:
+        return _build_error_response(400, "INVALID_TAG_FILTER", str(e), e.details)
+
     tag_counts = list_tag_histogram(
         owner_id=USER_MANAGER.get_request_user_id(request),
-        include_tags=q.include_tags,
-        exclude_tags=q.exclude_tags,
+        include_tags=tags_all,
+        exclude_tags=tags_none,
+        any_tags=tags_any,
         name_contains=q.name_contains,
         metadata_filter=q.metadata_filter,
         limit=q.limit,

app/assets/api/schemas_in.py

@@ -50,8 +50,12 @@ class ParsedUpload:
 
 
 class ListAssetsQuery(BaseModel):
-    include_tags: list[str] = Field(default_factory=list)
-    exclude_tags: list[str] = Field(default_factory=list)
+    # Deprecated spellings: include_tags ≡ tags_all, exclude_tags ≡ tags_none.
+    include_tags: list[str] = Field(default_factory=list, deprecated=True)
+    exclude_tags: list[str] = Field(default_factory=list, deprecated=True)
+    tags_all: list[str] = Field(default_factory=list)
+    tags_any: list[str] = Field(default_factory=list)
+    tags_none: list[str] = Field(default_factory=list)
     name_contains: str | None = None
 
     # Accept either a JSON string (query param) or a dict
@@ -70,7 +74,10 @@ class ListAssetsQuery(BaseModel):
     )
     order: Literal["asc", "desc"] = "desc"
 
-    @field_validator("include_tags", "exclude_tags", mode="before")
+    @field_validator(
+        "include_tags", "exclude_tags", "tags_all", "tags_any", "tags_none",
+        mode="before",
+    )
     @classmethod
     def _split_csv_tags(cls, v):
         # Accept "a,b,c" or ["a","b"] (we are liberal in what we accept)
@@ -154,13 +161,20 @@ class CreateFromHashBody(BaseModel):
 
 
 class TagsRefineQuery(BaseModel):
-    include_tags: list[str] = Field(default_factory=list)
-    exclude_tags: list[str] = Field(default_factory=list)
+    # Deprecated spellings: include_tags ≡ tags_all, exclude_tags ≡ tags_none.
+    include_tags: list[str] = Field(default_factory=list, deprecated=True)
+    exclude_tags: list[str] = Field(default_factory=list, deprecated=True)
+    tags_all: list[str] = Field(default_factory=list)
+    tags_any: list[str] = Field(default_factory=list)
+    tags_none: list[str] = Field(default_factory=list)
     name_contains: str | None = None
     metadata_filter: dict[str, Any] | None = None
     limit: conint(ge=1, le=1000) = 100
 
-    @field_validator("include_tags", "exclude_tags", mode="before")
+    @field_validator(
+        "include_tags", "exclude_tags", "tags_all", "tags_any", "tags_none",
+        mode="before",
+    )
     @classmethod
     def _split_csv_tags(cls, v):
         if v is None:

app/assets/database/queries/asset_reference.py

@@ -268,6 +268,8 @@ def list_references_page(
     order: str | None = None,
     after_cursor_value: object | None = None,
     after_cursor_id: str | None = None,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> tuple[list[AssetReference], dict[str, list[str]], int]:
     """List references with pagination, filtering, and sorting.
 
@@ -293,7 +295,7 @@ def list_references_page(
         escaped, esc = escape_sql_like_string(name_contains)
         base = base.where(AssetReference.name.ilike(f"%{escaped}%", escape=esc))
 
-    base = apply_tag_filters(base, include_tags, exclude_tags)
+    base = apply_tag_filters(base, include_tags, exclude_tags, any_tags)
     base = apply_metadata_filter(base, metadata_filter)
 
     sort = (sort or "created_at").lower()
@@ -345,7 +347,7 @@ def list_references_page(
         count_stmt = count_stmt.where(
             AssetReference.name.ilike(f"%{escaped}%", escape=esc)
         )
-    count_stmt = apply_tag_filters(count_stmt, include_tags, exclude_tags)
+    count_stmt = apply_tag_filters(count_stmt, include_tags, exclude_tags, any_tags)
     count_stmt = apply_metadata_filter(count_stmt, metadata_filter)
 
     total = int(session.execute(count_stmt).scalar_one() or 0)

app/assets/database/queries/common.py

@@ -60,10 +60,13 @@ def apply_tag_filters(
     stmt: sa.sql.Select,
     include_tags: Sequence[str] | None = None,
     exclude_tags: Sequence[str] | None = None,
+    any_tags: Sequence[str] | None = None,
 ) -> sa.sql.Select:
-    """include_tags: every tag must be present; exclude_tags: none may be present."""
+    """include_tags: every tag must be present; any_tags: at least one must be
+    present; exclude_tags: none may be present."""
     include_tags = normalize_tags(include_tags)
     exclude_tags = normalize_tags(exclude_tags)
+    any_tags = normalize_tags(any_tags)
 
     if include_tags:
         for tag_name in include_tags:
@@ -74,6 +77,14 @@ def apply_tag_filters(
                 )
             )
 
+    if any_tags:
+        stmt = stmt.where(
+            exists().where(
+                (AssetReferenceTag.asset_reference_id == AssetReference.id)
+                & (AssetReferenceTag.tag_name.in_(any_tags))
+            )
+        )
+
     if exclude_tags:
         stmt = stmt.where(
             ~exists().where(

app/assets/database/queries/tags.py

@@ -340,6 +340,8 @@ def list_tag_counts_for_filtered_assets(
     name_contains: str | None = None,
     metadata_filter: dict | None = None,
     limit: int = 100,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> dict[str, int]:
     """Return tag counts for assets matching the given filters.
 
@@ -359,7 +361,7 @@ def list_tag_counts_for_filtered_assets(
         escaped, esc = escape_sql_like_string(name_contains)
         ref_sq = ref_sq.where(AssetReference.name.ilike(f"%{escaped}%", escape=esc))
 
-    ref_sq = apply_tag_filters(ref_sq, include_tags, exclude_tags)
+    ref_sq = apply_tag_filters(ref_sq, include_tags, exclude_tags, any_tags)
     ref_sq = apply_metadata_filter(ref_sq, metadata_filter)
     ref_sq = ref_sq.subquery()
 

app/assets/services/asset_management.py

@@ -279,6 +279,8 @@ def list_assets_page(
     sort: str = "created_at",
     order: str = "desc",
     after: str | None = None,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> ListAssetsResult:
     """List assets with optional cursor pagination.
 
@@ -317,6 +319,7 @@ def list_assets_page(
             owner_id=owner_id,
             include_tags=include_tags,
             exclude_tags=exclude_tags,
+            any_tags=any_tags,
             name_contains=name_contains,
             metadata_filter=metadata_filter,
             limit=fetch_limit,

app/assets/services/tagging.py

@@ -85,6 +85,8 @@ def list_tag_histogram(
     name_contains: str | None = None,
     metadata_filter: dict | None = None,
     limit: int = 100,
+    # Appended last so pre-existing positional callers keep binding correctly.
+    any_tags: Sequence[str] | None = None,
 ) -> dict[str, int]:
     with create_session() as session:
         return list_tag_counts_for_filtered_assets(
@@ -92,6 +94,7 @@ def list_tag_histogram(
             owner_id=owner_id,
             include_tags=include_tags,
             exclude_tags=exclude_tags,
+            any_tags=any_tags,
             name_contains=name_contains,
             metadata_filter=metadata_filter,
             limit=limit,

comfy/background_removal/birefnet.py

@@ -433,19 +433,16 @@ class DeformableConv2d(nn.Module):
     def forward(self, x):
         offset = self.offset_conv(x)
         modulator = 2. * torch.sigmoid(self.modulator_conv(x))
-        weight, bias, offload_info = comfy.ops.cast_bias_weight(self.regular_conv, x, offloadable=True)
-
-        x = deform_conv2d(
-            input=x,
-            offset=offset,
-            weight=weight,
-            bias=None,
-            padding=self.padding,
-            mask=modulator,
-            stride=self.stride,
-        )
-        comfy.ops.uncast_bias_weight(self.regular_conv, weight, bias, offload_info)
-        return x
+        with comfy.ops.CastBiasWeightContext(self.regular_conv, x, offloadable=True) as (weight, _bias):
+            return deform_conv2d(
+                input=x,
+                offset=offset,
+                weight=weight,
+                bias=None,
+                padding=self.padding,
+                mask=modulator,
+                stride=self.stride,
+            )
 
 class BasicDecBlk(nn.Module):
     def __init__(self, in_channels=64, out_channels=64, inter_channels=64, device=None, dtype=None, operations=None):

comfy/cli_args.py

@@ -149,6 +149,7 @@ attn_group.add_argument("--use-quad-cross-attention", action="store_true", help=
 attn_group.add_argument("--use-pytorch-cross-attention", action="store_true", help="Use the new pytorch 2.0 cross attention function.")
 attn_group.add_argument("--use-sage-attention", action="store_true", help="Use sage attention.")
 attn_group.add_argument("--use-flash-attention", action="store_true", help="Use FlashAttention.")
+attn_group.add_argument("--use-ck-attention", action="store_true", help="Use Comfy Kitchen attention.")
 
 parser.add_argument("--disable-xformers", action="store_true", help="Disable xformers.")
 
@@ -179,6 +180,7 @@ parser.add_argument("--disable-async-offload", action="store_true", help="Disabl
 parser.add_argument("--disable-dynamic-vram", action="store_true", help="Disable dynamic VRAM and use estimate based model loading.")
 parser.add_argument("--enable-dynamic-vram", action="store_true", help="Enable dynamic VRAM on systems where it's not enabled by default.")
 parser.add_argument("--fast-disk", action="store_true", help="Prefer disk-backed dynamic loading and offload over unpinned RAM. Can be faster for users with fast NVME disks.")
+parser.add_argument("--disable-cuda-graphs", action="store_true", help="Disable CUDA graphs.")
 
 parser.add_argument("--force-non-blocking", action="store_true", help="Force ComfyUI to use non-blocking operations for all applicable tensors. This may improve performance on some non-Nvidia systems but can cause issues with some workflows.")
 

comfy/clip_model.py

@@ -314,13 +314,18 @@ class CLIPVisionModelProjection(torch.nn.Module):
         if "projection_dim" in config_dict:
             self.visual_projection = operations.Linear(config_dict["hidden_size"], config_dict["projection_dim"], bias=False)
         else:
-            self.visual_projection = lambda a: a
+            self.visual_projection = torch.nn.Identity()
 
         if "llava3" == config_dict.get("projector_type", None):
             self.multi_modal_projector = LlavaProjector(config_dict["hidden_size"], 4096, dtype, device, operations)
         else:
             self.multi_modal_projector = None
 
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        if "{}visual_projection.weight".format(prefix) not in state_dict:
+            self.visual_projection = torch.nn.Identity()
+        super()._load_from_state_dict(state_dict, prefix, *args, **kwargs)
+
     def forward(self, *args, **kwargs):
         x = self.vision_model(*args, **kwargs)
         out = self.visual_projection(x[2])

comfy/controlnet.py

@@ -381,13 +381,10 @@ class ControlLoraOps:
             self.bias = None
 
         def forward(self, input):
-            weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
-            if self.up is not None:
-                x = torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
-            else:
-                x = torch.nn.functional.linear(input, weight, bias)
-            comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with comfy.ops.CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                if self.up is None:
+                    return torch.nn.functional.linear(input, weight, bias)
+                return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
 
     class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
         def __init__(

comfy/k_diffusion/sampling.py

@@ -1445,7 +1445,7 @@ def res_multistep(model, x, sigmas, extra_args=None, callback=None, disable=None
                 x = sigma_fn(h) * x + h * (b1 * denoised + b2 * old_denoised)
 
         # Noise addition
-        if sigmas[i + 1] > 0:
+        if sigma_up > 0:
             x = x + noise_sampler(sigmas[i], sigmas[i + 1]) * s_noise * sigma_up
 
         if cfg_pp:

comfy/latent_formats.py

@@ -957,6 +957,11 @@ class ACEAudio15(LatentFormat):
     latent_dimensions = 1
     temporal_downscale_ratio = 1764
 
+class MiniMaxMusic3(LatentFormat):
+    latent_channels = 128
+    latent_dimensions = 1
+    temporal_downscale_ratio = 512
+
 class ChromaRadiance(LatentFormat):
     latent_channels = 3
     spacial_downscale_ratio = 1

comfy/ldm/lightricks/av_model.py

@@ -16,7 +16,9 @@ from comfy.ldm.lightricks.model import (
 from comfy.ldm.lightricks.symmetric_patchifier import AudioPatchifier
 from comfy.ldm.lightricks.embeddings_connector import Embeddings1DConnector
 import comfy.ldm.common_dit
+import comfy.model_management
 import comfy.model_prefetch
+import comfy.quant_ops
 
 class CompressedTimestep:
     """Store video timestep embeddings in compressed form using per-frame indexing."""
@@ -94,6 +96,8 @@ class BasicAVTransformerBlock(nn.Module):
         attn_precision=None,
         apply_gated_attention=False,
         cross_attention_adaln=False,
+        ff_bias=True,
+        audio_ff_bias=True,
         dtype=None,
         device=None,
         operations=None,
@@ -176,10 +180,10 @@ class BasicAVTransformerBlock(nn.Module):
         )
 
         self.ff = FeedForward(
-            v_dim, dim_out=v_dim, glu=True, dtype=dtype, device=device, operations=operations
+            v_dim, dim_out=v_dim, glu=True, ff_bias=ff_bias, dtype=dtype, device=device, operations=operations
         )
         self.audio_ff = FeedForward(
-            a_dim, dim_out=a_dim, glu=True, dtype=dtype, device=device, operations=operations
+            a_dim, dim_out=a_dim, glu=True, ff_bias=audio_ff_bias, dtype=dtype, device=device, operations=operations
         )
 
         num_ada_params = ADALN_CROSS_ATTN_PARAMS_COUNT if cross_attention_adaln else ADALN_BASE_PARAMS_COUNT
@@ -271,7 +275,10 @@ class BasicAVTransformerBlock(nn.Module):
         if run_vx:
             # video self-attention
             vshift_msa, vscale_msa = (self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(0, 2)))
-            norm_vx = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_msa) + vshift_msa
+            if comfy.model_management.in_training:
+                norm_vx = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_msa) + vshift_msa
+            else:
+                norm_vx = comfy.quant_ops.ck.rms_adaln(vx, vscale_msa, vshift_msa)
             del vshift_msa, vscale_msa
             attn1_out = self.attn1(norm_vx, pe=v_pe, mask=self_attention_mask, transformer_options=transformer_options)
             del norm_vx
@@ -305,7 +312,6 @@ class BasicAVTransformerBlock(nn.Module):
 
         # video - audio cross attention.
         if run_a2v or run_v2a:
-            vx_norm3 = comfy.ldm.common_dit.rms_norm(vx)
             ax_norm3 = comfy.ldm.common_dit.rms_norm(ax)
 
             # audio to video cross attention
@@ -315,7 +321,10 @@ class BasicAVTransformerBlock(nn.Module):
                 scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v = self.get_ada_values(
                     self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[:2]
 
-                vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_a2v_v) + shift_ca_video_hidden_states_a2v_v
+                if comfy.model_management.in_training:
+                    vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + scale_ca_video_hidden_states_a2v_v) + shift_ca_video_hidden_states_a2v_v
+                else:
+                    vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v)
                 ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_a2v) + shift_ca_audio_hidden_states_a2v
                 del scale_ca_video_hidden_states_a2v_v, shift_ca_video_hidden_states_a2v_v, scale_ca_audio_hidden_states_a2v, shift_ca_audio_hidden_states_a2v
 
@@ -334,7 +343,10 @@ class BasicAVTransformerBlock(nn.Module):
                     self.scale_shift_table_a2v_ca_video[:4, :], vx.shape[0], v_cross_scale_shift_timestep)[2:4]
 
                 ax_scaled = ax_norm3 * (1 + scale_ca_audio_hidden_states_v2a) + shift_ca_audio_hidden_states_v2a
-                vx_scaled = vx_norm3 * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
+                if comfy.model_management.in_training:
+                    vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + scale_ca_video_hidden_states_v2a) + shift_ca_video_hidden_states_v2a
+                else:
+                    vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a)
                 del scale_ca_video_hidden_states_v2a, shift_ca_video_hidden_states_v2a, scale_ca_audio_hidden_states_v2a, shift_ca_audio_hidden_states_v2a
 
                 v2a_out = self.video_to_audio_attn(ax_scaled, context=vx_scaled, pe=a_cross_pe, k_pe=v_cross_pe, transformer_options=transformer_options)
@@ -344,12 +356,14 @@ class BasicAVTransformerBlock(nn.Module):
                 ax.addcmul_(v2a_out, gate_out_v2a)
                 del gate_out_v2a, v2a_out
 
-            del vx_norm3, ax_norm3
 
         # video feedforward
         if run_vx:
             vshift_mlp, vscale_mlp = self.get_ada_values(self.scale_shift_table, vx.shape[0], v_timestep, slice(3, 5))
-            vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_mlp) + vshift_mlp
+            if comfy.model_management.in_training:
+                vx_scaled = comfy.ldm.common_dit.rms_norm(vx) * (1 + vscale_mlp) + vshift_mlp
+            else:
+                vx_scaled = comfy.quant_ops.ck.rms_adaln(vx, vscale_mlp, vshift_mlp)
             del vshift_mlp, vscale_mlp
 
             ff_out = self.ff(vx_scaled)
@@ -401,12 +415,16 @@ class LTXAVModel(LTXVModel):
         apply_gated_attention=False,
         caption_proj_before_connector=False,
         cross_attention_adaln=False,
+        ff_bias=True,
+        audio_ff_bias=True,
+        use_prompt_adaln_single=True,
         dtype=None,
         device=None,
         operations=None,
         **kwargs,
     ):
         # Store audio-specific parameters
+        self.audio_ff_bias = audio_ff_bias
         self.audio_in_channels = audio_in_channels
         self.audio_cross_attention_dim = audio_cross_attention_dim
         self.audio_attention_head_dim = audio_attention_head_dim
@@ -439,6 +457,8 @@ class LTXAVModel(LTXVModel):
             timestep_scale_multiplier=timestep_scale_multiplier,
             caption_proj_before_connector=caption_proj_before_connector,
             cross_attention_adaln=cross_attention_adaln,
+            ff_bias=ff_bias,
+            use_prompt_adaln_single=use_prompt_adaln_single,
             dtype=dtype,
             device=device,
             operations=operations,
@@ -463,7 +483,7 @@ class LTXAVModel(LTXVModel):
             operations=self.operations,
         )
 
-        if self.cross_attention_adaln:
+        if self.cross_attention_adaln and self.use_prompt_adaln_single:
             self.audio_prompt_adaln_single = AdaLayerNormSingle(
                 self.audio_inner_dim,
                 embedding_coefficient=2,
@@ -594,6 +614,8 @@ class LTXAVModel(LTXVModel):
                     a_context_dim=self.audio_cross_attention_dim,
                     apply_gated_attention=self.apply_gated_attention,
                     cross_attention_adaln=self.cross_attention_adaln,
+                    ff_bias=self.ff_bias,
+                    audio_ff_bias=self.audio_ff_bias,
                     dtype=dtype,
                     device=device,
                     operations=self.operations,
@@ -912,9 +934,15 @@ class LTXAVModel(LTXVModel):
         blocks_replace = patches_replace.get("dit", {})
         prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.transformer_blocks), vx.device, transformer_options)
 
+        # Blocks whose self-attention should be perturbed to a value-passthrough (STG).
+        stg_self_attn_blocks = transformer_options.get("stg_self_attn_blocks", ())
+
         # Process transformer blocks
         for i, block in enumerate(self.transformer_blocks):
             comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, vx.device, block)
+            block_transformer_options = transformer_options
+            if i in stg_self_attn_blocks:
+                block_transformer_options = {**transformer_options, "stg_skip_self_attn": True}
             if ("double_block", i) in blocks_replace:
 
                 def block_wrap(args):
@@ -957,7 +985,7 @@ class LTXAVModel(LTXVModel):
                         "a_cross_scale_shift_timestep": av_ca_audio_scale_shift_timestep,
                         "v_cross_gate_timestep": av_ca_a2v_gate_noise_timestep,
                         "a_cross_gate_timestep": av_ca_v2a_gate_noise_timestep,
-                        "transformer_options": transformer_options,
+                        "transformer_options": block_transformer_options,
                         "self_attention_mask": self_attention_mask,
                         "v_prompt_timestep": v_prompt_timestep,
                         "a_prompt_timestep": a_prompt_timestep,
@@ -981,7 +1009,7 @@ class LTXAVModel(LTXVModel):
                     a_cross_scale_shift_timestep=av_ca_audio_scale_shift_timestep,
                     v_cross_gate_timestep=av_ca_a2v_gate_noise_timestep,
                     a_cross_gate_timestep=av_ca_v2a_gate_noise_timestep,
-                    transformer_options=transformer_options,
+                    transformer_options=block_transformer_options,
                     self_attention_mask=self_attention_mask,
                     v_prompt_timestep=v_prompt_timestep,
                     a_prompt_timestep=a_prompt_timestep,

comfy/ldm/lightricks/duration_head.py

@@ -0,0 +1,81 @@
+"""LTX 2.4 DurationHead: predicts the natural shot duration (in seconds) from
+the caption connector token outputs, without running the diffusion pipeline.
+"""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class AttentionPooler(nn.Module):
+    """Cross-attend ``num_queries`` learnable tokens against ``tokens``."""
+
+    def __init__(self, hidden_dim=256, num_queries=1, num_heads=4):
+        super().__init__()
+        self.num_queries = num_queries
+        self.query_tokens = nn.Parameter(torch.empty(num_queries, hidden_dim))
+        self.cross_attn = nn.MultiheadAttention(embed_dim=hidden_dim, num_heads=num_heads, batch_first=True)
+
+    def forward(self, tokens):
+        queries = self.query_tokens.unsqueeze(0).expand(tokens.shape[0], -1, -1)
+        pooled, _ = self.cross_attn(queries, tokens, tokens, need_weights=False)
+        return pooled
+
+
+class DurationHead(nn.Module):
+    """Predict duration in seconds from one or both connector outputs."""
+
+    def __init__(
+        self,
+        video_cross_attention_dim=4096,
+        audio_cross_attention_dim=2048,
+        pooler_hidden_dim=256,
+        num_queries=1,
+        num_pooler_heads=4,
+        mlp_hidden=256,
+    ):
+        super().__init__()
+        self.video_input_proj = nn.Linear(video_cross_attention_dim, pooler_hidden_dim)
+        self.video_modality_emb = nn.Parameter(torch.empty(pooler_hidden_dim))
+        self.audio_input_proj = nn.Linear(audio_cross_attention_dim, pooler_hidden_dim)
+        self.audio_modality_emb = nn.Parameter(torch.empty(pooler_hidden_dim))
+        self.attention_pooler = AttentionPooler(
+            hidden_dim=pooler_hidden_dim, num_queries=num_queries, num_heads=num_pooler_heads)
+        self.mlp_hidden = nn.Linear(pooler_hidden_dim * num_queries, mlp_hidden)
+        self.mlp_out = nn.Linear(mlp_hidden, 1)
+
+    def forward(self, video_tokens=None, audio_tokens=None):
+        """``video_tokens``: (B, T_v, 4096), ``audio_tokens``: (B, T_a, 2048);
+        at least one required. Returns duration in seconds, shape (B,)."""
+        token_groups = []
+        if video_tokens is not None:
+            token_groups.append(self.video_input_proj(video_tokens) + self.video_modality_emb)
+        if audio_tokens is not None:
+            token_groups.append(self.audio_input_proj(audio_tokens) + self.audio_modality_emb)
+        if not token_groups:
+            raise ValueError("DurationHead requires at least one of video_tokens / audio_tokens")
+        pooled = self.attention_pooler(torch.cat(token_groups, dim=1))
+        pooled = pooled.reshape(pooled.shape[0], -1)
+        hidden = F.gelu(self.mlp_hidden(pooled), approximate="tanh")
+        return self.mlp_out(hidden).squeeze(-1).exp()
+
+
+def normalize_state_dict(sd):
+    for prefix in ("model.diffusion_model.duration_head.", "duration_head."):
+        stripped = {k[len(prefix):]: v for k, v in sd.items() if k.startswith(prefix)}
+        if stripped:
+            return stripped
+    return sd
+
+
+def seconds_to_num_frames(seconds, frame_rate, min_seconds, max_seconds, time_scale=8):
+    """Convert seconds to a frame count clamped to ``[min_seconds, max_seconds]``
+    and snapped (floor) to the VAE's ``8k + 1`` causal temporal grid; snapping
+    that undershoots the minimum bumps up to the next grid point instead."""
+    min_frames = max(1, round(min_seconds * frame_rate))
+    max_frames = round(max_seconds * frame_rate)
+    raw_frames = max(min_frames, min(round(seconds * frame_rate), max_frames))
+    frames = (raw_frames - 1) // time_scale * time_scale + 1
+    if frames < min_frames:
+        frames = min(-(-(min_frames - 1) // time_scale) * time_scale + 1, max_frames)
+    return frames

comfy/ldm/lightricks/embeddings_connector.py

@@ -50,6 +50,7 @@ class BasicTransformerBlock1D(nn.Module):
         context_dim=None,
         attn_precision=None,
         apply_gated_attention=False,
+        ff_bias=True,
         dtype=None,
         device=None,
         operations=None,
@@ -74,6 +75,7 @@ class BasicTransformerBlock1D(nn.Module):
             dim,
             dim_out=dim,
             glu=True,
+            ff_bias=ff_bias,
             dtype=dtype,
             device=device,
             operations=operations,
@@ -123,6 +125,7 @@ class Embeddings1DConnector(nn.Module):
         causal_temporal_positioning=False,
         num_learnable_registers: Optional[int] = 128,
         apply_gated_attention=False,
+        connector_ff_bias=True,
         dtype=None,
         device=None,
         operations=None,
@@ -148,6 +151,7 @@ class Embeddings1DConnector(nn.Module):
                     attention_head_dim,
                     context_dim=cross_attention_dim,
                     apply_gated_attention=apply_gated_attention,
+                    ff_bias=connector_ff_bias,
                     dtype=dtype,
                     device=device,
                     operations=operations,

comfy/ldm/lightricks/model.py

@@ -13,6 +13,7 @@ import comfy.patcher_extension
 import comfy.ldm.modules.attention
 import comfy.ldm.common_dit
 import comfy.model_management
+import comfy.ops
 import comfy.quant_ops
 
 from .symmetric_patchifier import SymmetricPatchifier, latent_to_pixel_coords
@@ -302,26 +303,30 @@ class NormSingleLinearTextProjection(nn.Module):
 
 
 class GELU_approx(nn.Module):
-    def __init__(self, dim_in, dim_out, dtype=None, device=None, operations=None):
+    def __init__(self, dim_in, dim_out, bias=True, dtype=None, device=None, operations=None):
         super().__init__()
-        self.proj = operations.Linear(dim_in, dim_out, dtype=dtype, device=device)
+        self.proj = operations.Linear(dim_in, dim_out, bias=bias, dtype=dtype, device=device)
 
     def forward(self, x):
         return torch.nn.functional.gelu(self.proj(x), approximate="tanh")
 
 
 class FeedForward(nn.Module):
-    def __init__(self, dim, dim_out, mult=4, glu=False, dropout=0.0, dtype=None, device=None, operations=None):
+    def __init__(self, dim, dim_out, mult=4, glu=False, dropout=0.0, ff_bias=True, dtype=None, device=None, operations=None):
         super().__init__()
         inner_dim = int(dim * mult)
-        project_in = GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations)
+        project_in = GELU_approx(dim, inner_dim, bias=ff_bias, dtype=dtype, device=device, operations=operations)
 
         self.net = nn.Sequential(
-            project_in, nn.Dropout(dropout), operations.Linear(inner_dim, dim_out, dtype=dtype, device=device)
+            project_in, nn.Dropout(dropout), operations.Linear(inner_dim, dim_out, bias=ff_bias, dtype=dtype, device=device)
         )
 
     def forward(self, x):
-        return self.net(x)
+        # net = [GELU_approx(proj), Dropout, Linear]; the fused path skips the
+        # Dropout, so leave it to the stock path whenever it could be active.
+        if comfy.model_management.in_training:
+            return self.net(x)
+        return comfy.ops.linear_input_act(self.net[2], self.net[0].proj(x), "gelu_tanh")
 
 def apply_rotary_emb(input_tensor, freqs_cis):
     rotation_matrix, split_pe = freqs_cis
@@ -457,28 +462,34 @@ class CrossAttention(nn.Module):
         )
 
     def forward(self, x, context=None, mask=None, pe=None, k_pe=None, transformer_options={}):
+        self_attn = context is None
         q = self.to_q(x)
         context = x if context is None else context
         k = self.to_k(context)
         v = self.to_v(context)
 
-        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:
-            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)
-        elif isinstance(mask, GuideAttentionMask):
-            out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+        # Spatio-Temporal Guidance (STG) perturbation: for the flagged self-attention
+        # layers, the attention degrades to a passthrough of the value projection (out = V).
+        if self_attn and transformer_options.get("stg_skip_self_attn", False):
+            out = v
         else:
-            out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            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:
+                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)
+            elif isinstance(mask, GuideAttentionMask):
+                out = _attention_with_guide_mask(q, k, v, self.heads, mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
+            else:
+                out = comfy.ldm.modules.attention.optimized_attention(q, k, v, self.heads, mask=mask, attn_precision=self.attn_precision, transformer_options=transformer_options)
 
         # Apply per-head gating if enabled
         if self.to_gate_logits is not None:
@@ -497,7 +508,7 @@ ADALN_CROSS_ATTN_PARAMS_COUNT = 9
 
 class BasicTransformerBlock(nn.Module):
     def __init__(
-        self, dim, n_heads, d_head, context_dim=None, attn_precision=None, cross_attention_adaln=False, dtype=None, device=None, operations=None
+        self, dim, n_heads, d_head, context_dim=None, attn_precision=None, cross_attention_adaln=False, ff_bias=True, dtype=None, device=None, operations=None
     ):
         super().__init__()
 
@@ -513,7 +524,7 @@ class BasicTransformerBlock(nn.Module):
             device=device,
             operations=operations,
         )
-        self.ff = FeedForward(dim, dim_out=dim, glu=True, dtype=dtype, device=device, operations=operations)
+        self.ff = FeedForward(dim, dim_out=dim, glu=True, ff_bias=ff_bias, dtype=dtype, device=device, operations=operations)
 
         self.attn2 = CrossAttention(
             query_dim=dim,
@@ -535,7 +546,12 @@ class BasicTransformerBlock(nn.Module):
     def forward(self, x, context=None, attention_mask=None, timestep=None, pe=None, transformer_options={}, self_attention_mask=None, prompt_timestep=None):
         shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (self.scale_shift_table[None, None, :6].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)[:, :, :6, :]).unbind(dim=2)
 
-        x += self.attn1(comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa, pe=pe, mask=self_attention_mask, transformer_options=transformer_options) * gate_msa
+        if comfy.model_management.in_training:
+            norm_x = comfy.ldm.common_dit.rms_norm(x) * (1 + scale_msa) + shift_msa
+        else:
+            norm_x = comfy.quant_ops.ck.rms_adaln(x, scale_msa, shift_msa)
+
+        x += self.attn1(norm_x, pe=pe, mask=self_attention_mask, transformer_options=transformer_options) * gate_msa
 
         if self.cross_attention_adaln:
             shift_q_mca, scale_q_mca, gate_mca = (self.scale_shift_table[None, None, 6:9].to(device=x.device, dtype=x.dtype) + timestep.reshape(x.shape[0], timestep.shape[1], self.scale_shift_table.shape[0], -1)[:, :, 6:9, :]).unbind(dim=2)
@@ -589,7 +605,10 @@ def apply_cross_attention_adaln(
         prompt_scale_shift_table[None, None].to(device=x.device, dtype=x.dtype)
         + prompt_timestep.reshape(batch_size, prompt_timestep.shape[1], 2, -1)
     ).unbind(dim=2)
-    attn_input = comfy.ldm.common_dit.rms_norm(x) * (1 + q_scale) + q_shift
+    if comfy.model_management.in_training:
+        attn_input = comfy.ldm.common_dit.rms_norm(x) * (1 + q_scale) + q_shift
+    else:
+        attn_input = comfy.quant_ops.ck.rms_adaln(x, q_scale, q_shift)
     encoder_hidden_states = context * (1 + scale_kv) + shift_kv
     return attn(attn_input, context=encoder_hidden_states, mask=attention_mask, transformer_options=transformer_options) * q_gate
 
@@ -704,6 +723,9 @@ class LTXBaseModel(torch.nn.Module, ABC):
         caption_proj_before_connector=False,
         cross_attention_adaln=False,
         caption_projection_first_linear=True,
+        ff_bias=True,
+        use_prompt_adaln_single=True,
+        use_keyframes_abs_pos_embedding=False,
         dtype=None,
         device=None,
         operations=None,
@@ -733,6 +755,9 @@ class LTXBaseModel(torch.nn.Module, ABC):
         self.caption_proj_before_connector = caption_proj_before_connector
         self.cross_attention_adaln = cross_attention_adaln
         self.caption_projection_first_linear = caption_projection_first_linear
+        self.ff_bias = ff_bias
+        self.use_prompt_adaln_single = use_prompt_adaln_single
+        self.use_keyframes_abs_pos_embedding = use_keyframes_abs_pos_embedding
 
         # Common dimensions
         self.inner_dim = num_attention_heads * attention_head_dim
@@ -760,12 +785,17 @@ class LTXBaseModel(torch.nn.Module, ABC):
             self.in_channels, self.inner_dim, bias=True, dtype=dtype, device=device
         )
 
+        if self.use_keyframes_abs_pos_embedding:
+            self.keyframes_abs_pos_embedding = nn.Parameter(torch.zeros(1, self.inner_dim, dtype=dtype, device=device))
+        else:
+            self.keyframes_abs_pos_embedding = None
+
         embedding_coefficient = ADALN_CROSS_ATTN_PARAMS_COUNT if self.cross_attention_adaln else ADALN_BASE_PARAMS_COUNT
         self.adaln_single = AdaLayerNormSingle(
             self.inner_dim, embedding_coefficient=embedding_coefficient, use_additional_conditions=False, dtype=dtype, device=device, operations=self.operations
         )
 
-        if self.cross_attention_adaln:
+        if self.cross_attention_adaln and self.use_prompt_adaln_single:
             self.prompt_adaln_single = AdaLayerNormSingle(
                 self.inner_dim, embedding_coefficient=2, use_additional_conditions=False, dtype=dtype, device=device, operations=self.operations
             )
@@ -1057,6 +1087,7 @@ class LTXVModel(LTXBaseModel):
                     self.attention_head_dim,
                     context_dim=self.cross_attention_dim,
                     cross_attention_adaln=self.cross_attention_adaln,
+                    ff_bias=self.ff_bias,
                     dtype=dtype,
                     device=device,
                     operations=self.operations,
@@ -1086,6 +1117,15 @@ class LTXVModel(LTXBaseModel):
 
         grid_mask = None
         if keyframe_idxs is not None and keyframe_idxs.shape[2] > 0:
+            tokens_per_frame = self.tokens_per_latent_frame(additional_args["orig_shape"])
+            if keyframe_idxs.shape[2] % tokens_per_frame != 0:
+                raise ValueError(
+                    f"keyframe_idxs holds {keyframe_idxs.shape[2]} tokens, which is not a whole number of "
+                    f"{tokens_per_frame}-token latent frames. The appended frames were recorded against a "
+                    "different spatial resolution than the latent being sampled, so their positions would land "
+                    "on the wrong tokens. Crop the guides and separate the generated keyframes before "
+                    "upscaling the latent."
+                )
             additional_args.update({ "orig_patchified_shape": list(x.shape)})
             denoise_mask = self.patchifier.patchify(denoise_mask)[0]
             grid_mask = ~torch.any(denoise_mask < 0, dim=-1)[0]
@@ -1128,8 +1168,64 @@ class LTXVModel(LTXBaseModel):
             additional_args["num_guide_tokens"] = keyframe_idxs.shape[2]
 
         x = self.patchify_proj(x)
+        x = self.apply_keyframes_abs_pos_embedding(
+            x,
+            pixel_coords,
+            orig_shape=additional_args["orig_shape"],
+            grid_mask=grid_mask,
+            num_guide_tokens=additional_args.get("num_guide_tokens", 0),
+            generated_keyframes=kwargs.get("generated_keyframes", None),
+        )
         return x, pixel_coords, additional_args
 
+    def tokens_per_latent_frame(self, orig_shape):
+        """Token count of a single latent frame at the given latent shape."""
+        patch_size = self.patchifier.patch_size
+        return (orig_shape[3] // patch_size[1]) * (orig_shape[4] // patch_size[2])
+
+    def keyframes_abs_pos_mask(self, pixel_coords, orig_shape, grid_mask, num_guide_tokens, generated_keyframes):
+        """Per-token mask selecting the latents that encode a single standalone pixel frame.
+
+        Returns a (batch, tokens) boolean mask over the already grid-filtered token sequence.
+        """
+        temporal_start = pixel_coords[:, 0]
+        if temporal_start.ndim == 3:  # (batch, tokens, [start, end])
+            temporal_start = temporal_start[..., 0]
+        mask = temporal_start == 0
+        if num_guide_tokens > 0:
+            mask[:, -num_guide_tokens:] = False
+
+        if generated_keyframes is not None:
+            # The temporal patch size is always 1, so one latent frame is one row of tokens.
+            tokens_per_frame = self.tokens_per_latent_frame(orig_shape)
+            if generated_keyframes["tokens_per_frame"] != tokens_per_frame:
+                raise ValueError(
+                    f"The generated keyframes were recorded at {generated_keyframes['tokens_per_frame']} tokens "
+                    f"per latent frame but this latent has {tokens_per_frame}. Separate the generated keyframes "
+                    "before upscaling the latent."
+                )
+            first_token = generated_keyframes["first_latent_frame"] * tokens_per_frame
+            num_slot_tokens = generated_keyframes["num_keyframes"] * tokens_per_frame
+            slots = torch.zeros(orig_shape[2] * tokens_per_frame, dtype=torch.bool, device=mask.device)
+            slots[first_token:first_token + num_slot_tokens] = True
+            if grid_mask is not None:
+                slots = slots[grid_mask]
+            mask = mask | slots
+
+        return mask
+
+    def apply_keyframes_abs_pos_embedding(self, x, pixel_coords, orig_shape, grid_mask, num_guide_tokens, generated_keyframes):
+        """Add the learned keyframe marker to the single-pixel-frame tokens.
+
+        A no-op for every checkpoint built without the parameter.
+        """
+        if self.keyframes_abs_pos_embedding is None:
+            return x
+
+        mask = self.keyframes_abs_pos_mask(pixel_coords, orig_shape, grid_mask, num_guide_tokens, generated_keyframes)
+        embedding = self.keyframes_abs_pos_embedding.to(device=x.device, dtype=x.dtype)
+        return x + mask.unsqueeze(-1).to(x.dtype) * embedding
+
     def _build_guide_self_attention_mask(self, x, transformer_options, merged_args):
         """Build self-attention mask for per-guide attention attenuation.
 

comfy/ldm/lightricks/vae/audio_vae.py

@@ -1,6 +1,5 @@
 import json
 from dataclasses import dataclass
-import math
 import torch
 import torchaudio
 
@@ -186,7 +185,7 @@ class AudioVAE(torch.nn.Module):
         )
 
     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)
+        return round((float(frames_number) / frame_rate) * self.latents_per_second)
 
     def run_vocoder(self, mel_spec: torch.Tensor) -> torch.Tensor:
         audio_channels = self.autoencoder.decoder.out_ch

comfy/ldm/lightricks/vae/na_diffusion_decoder.py

@@ -0,0 +1,520 @@
+"""LTX 2.4 diffusion video VAE decoder (NADiffusionDecoder).
+
+Port of the reference ``DiffusionVideoDecoder`` without the NATTEN dependency:
+``natten.na3d`` is replaced by ``comfy_kitchen.na3d``, which reproduces
+NATTEN's semantics (window of exactly ``kernel_size`` per query, shifted
+inward at grid boundaries, dilation 1) and dispatches cuda/triton/eager per
+device and dtype (the eager backend covers CPU and fp32).
+
+Stages 1-4 deterministically upsample the latent into a context volume via
+NA transformer blocks + linear pixel-shuffle upsamples. Stage 5 runs
+``DiffusionNABlock``s that denoise patchified noised pixels ``x_t`` guided by
+that context through AdaLN-Zero scale/shift. The 2.4 checkpoint is single-step
+``x0``: one forward pass yields the pixels directly, no Euler loop.
+
+State dict keys match the shipped checkpoints directly (fused ``attn.qkv``,
+``t_embedder.mlp.{0,2}``, ``shared_adaln.proj``); no rename pass is needed.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+from einops import rearrange
+from torch import nn
+import comfy.model_management
+
+from comfy.ldm.lightricks.model import get_timestep_embedding
+from .causal_video_autoencoder import Encoder, processor
+
+import comfy_kitchen
+
+# Token chunk for the SwiGLU MLP (bounds the [chunk, hidden] workspace).
+MLP_TOKEN_CHUNK = 65536
+
+
+def rms_norm(x, weight, eps=1e-6):
+    if hasattr(F, "rms_norm"):
+        return F.rms_norm(x, (x.shape[-1],), weight=weight.to(x.dtype), eps=eps)
+    x_f = x.float()
+    x_f = x_f * torch.rsqrt(x_f.pow(2).mean(-1, keepdim=True) + eps)
+    return (x_f * weight.float()).to(x.dtype)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim, eps=1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x):
+        return rms_norm(x, self.weight, self.eps)
+
+
+def patchify(x, patch_size_hw, patch_size_t=1):
+    if patch_size_hw == 1 and patch_size_t == 1:
+        return x
+    return rearrange(x, "b c (f p) (h q) (w r) -> b (c p r q) f h w", p=patch_size_t, q=patch_size_hw, r=patch_size_hw)
+
+
+def unpatchify(x, patch_size_hw, patch_size_t=1):
+    if patch_size_hw == 1 and patch_size_t == 1:
+        return x
+    return rearrange(x, "b (c p r q) f h w -> b c (f p) (h q) (w r)", p=patch_size_t, q=patch_size_hw, r=patch_size_hw)
+
+
+# --- Absolute per-axis RoPE (matches ltx-core rope.py numerics) ---
+
+def default_rope_dim_split(head_dim):
+    d_t = (head_dim // 4) // 2 * 2
+    d_hw = (head_dim - d_t) // 2
+    if d_hw % 2 != 0:
+        d_t -= 2
+        d_hw = (head_dim - d_t) // 2
+    return (d_t, d_hw, d_hw)
+
+
+def rope_inv_freqs(dim, base=10000.0, device=None):
+    out_device = device
+    if not comfy.model_management.supports_fp64(device):
+        device = torch.device("cpu")
+
+    exponents = torch.arange(0, dim, 2, dtype=torch.float64, device=device) / dim
+    return (1.0 / torch.pow(torch.tensor(float(base), dtype=torch.float64, device=device), exponents)).to(dtype=torch.float32, device=out_device)
+
+
+def _rope_tables(lengths, inv_freqs, device):
+    """Precompute per-axis fp32 cos/sin tables for global 0-based positions."""
+    tables = []
+    for length, inv in zip(lengths, inv_freqs):
+        pos = torch.arange(length, dtype=torch.float32, device=device)
+        ang = pos[:, None] * inv[None, :]
+        tables.append((ang.cos(), ang.sin()))
+    return tables
+
+
+def _rope_matrices_slice(tables, t0, t1, h, w):
+    """Per-token rotation matrices ``(1, ts*h*w, 1, hd/2, 2, 2)`` fp32 for
+    ``comfy_kitchen.rms_rope_`` (interleaved-pair convention), covering global
+    frames ``[t0, t1)`` of the axis-factorized tables."""
+    parts = []
+    for (c, s), sl in zip(tables, (slice(t0, t1), slice(None), slice(None))):
+        c, s = c[sl], s[sl]
+        parts.append(torch.stack([c, -s, s, c], dim=-1).reshape(c.shape[0], 1, 1, c.shape[1], 2, 2))
+    ts = t1 - t0
+    freqs = torch.cat([
+        parts[0].expand(ts, h, w, -1, 2, 2),
+        parts[1].transpose(0, 1).expand(ts, h, w, -1, 2, 2),
+        parts[2].movedim(0, 2).expand(ts, h, w, -1, 2, 2),
+    ], dim=3)
+    return freqs.reshape(1, ts * h * w, 1, -1, 2, 2)
+
+
+class NeighborhoodAttention3D(nn.Module):
+    """QKV (fused, matching checkpoint keys) + q/k RMSNorm + abs RoPE + NA."""
+
+    def __init__(self, dim, kernel_size, head_dim=64, rope_base=10000.0):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = dim // head_dim
+        self.head_dim = head_dim
+        self.kernel_size = tuple(kernel_size)
+        self.scale = head_dim ** -0.5
+        self.rope_split = default_rope_dim_split(head_dim)
+        self.rope_base = rope_base
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.proj = nn.Linear(dim, dim, bias=True)
+        self.q_norm = RMSNorm(head_dim, eps=1e-6)
+        self.k_norm = RMSNorm(head_dim, eps=1e-6)
+
+    def forward(self, x, pre=None, add_to=None):
+        """``pre`` (per-token norm/modulate) is applied slice-wise so the full
+        pre-attention tensor is never materialized; ``add_to`` streams the
+        output projection into it in place (residual add) and returns it.
+        Both bound peak memory without changing results."""
+        batch, t, h, w, _ = x.shape
+        inv_freqs = tuple(rope_inv_freqs(d, self.rope_base, device=x.device) for d in self.rope_split)
+        tables = _rope_tables((t, h, w), inv_freqs, x.device)
+        shape = (batch, t, h, w, self.num_heads, self.head_dim)
+        q = torch.empty(shape, dtype=x.dtype, device=x.device)
+        k = torch.empty(shape, dtype=x.dtype, device=x.device)
+        v = torch.empty(shape, dtype=x.dtype, device=x.device)
+        q_weight = (self.q_norm.weight.detach() * self.scale).to(x.dtype)  # scale commutes with the rotation
+        k_weight = self.k_norm.weight.detach().to(x.dtype)
+        chunk = max(1, (2 ** 25) // max(h * w * self.dim, 1))
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            sl = x[:, t0:t1] if pre is None else pre(x[:, t0:t1])
+            qc, kc, vc = self.qkv(sl).chunk(3, dim=-1)
+            cshape = (batch, t1 - t0, h, w, self.num_heads, self.head_dim)
+            q[:, t0:t1] = qc.reshape(cshape)
+            k[:, t0:t1] = kc.reshape(cshape)
+            v[:, t0:t1] = vc.reshape(cshape)
+            freqs = _rope_matrices_slice(tables, t0, t1, h, w)
+            nt = (t1 - t0) * h * w
+            for b in range(batch):
+                comfy_kitchen.rms_rope_(
+                    q[b, t0:t1].view(1, nt, self.num_heads, self.head_dim),
+                    k[b, t0:t1].view(1, nt, self.num_heads, self.head_dim),
+                    freqs, q_weight, k_weight)
+        out = comfy_kitchen.na3d(q, k, v, list(self.kernel_size), None, 1.0)
+        del q, k, v
+        out = out.reshape(batch, t, h, w, self.dim)
+        res = add_to if add_to is not None else torch.empty_like(out)
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            if add_to is not None:
+                res[:, t0:t1] += self.proj(out[:, t0:t1])
+            else:
+                res[:, t0:t1] = self.proj(out[:, t0:t1])
+        return res
+
+
+class SwiGLU(nn.Module):
+    """``w_down(silu(w_gate(x)) * w_up(x))``, chunked over tokens to bound the
+    ``[chunk, hidden]`` workspace."""
+
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.w_up = nn.Linear(dim, hidden_dim, bias=False)
+        self.w_gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.w_down = nn.Linear(hidden_dim, dim, bias=False)
+
+    def forward(self, x, pre=None, add_to=None):
+        """``pre``/``add_to`` as in ``NeighborhoodAttention3D.forward``."""
+        _, t, h, w, _ = x.shape
+        chunk = max(1, MLP_TOKEN_CHUNK // max(h * w, 1))
+        out = add_to if add_to is not None else torch.empty_like(x)
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            sl = x[:, t0:t1] if pre is None else pre(x[:, t0:t1])
+            y = self.w_down(F.silu(self.w_gate(sl)) * self.w_up(sl))
+            if add_to is not None:
+                out[:, t0:t1] += y
+            else:
+                out[:, t0:t1] = y
+        return out
+
+
+class NABlock(nn.Module):
+    """Pre-norm transformer block: NA -> SwiGLU MLP with residual adds."""
+
+    def __init__(self, dim, kernel_size, head_dim=64, mlp_ratio=4.0):
+        super().__init__()
+        self.norm1 = RMSNorm(dim, eps=1e-6)
+        self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim)
+        self.norm2 = RMSNorm(dim, eps=1e-6)
+        hidden = (int(dim * mlp_ratio) + 15) // 16 * 16
+        self.mlp = SwiGLU(dim, hidden)
+
+    def forward(self, x):
+        x = self.attn(x, pre=self.norm1, add_to=x)
+        return self.mlp(x, pre=self.norm2, add_to=x)
+
+
+def modulate(x, scale, shift):
+    return x * (1.0 + scale) + shift
+
+
+class AdaLNZero(nn.Module):
+    """``t_emb`` -> 7 (scale/shift/gate) chunks; gate slots unused (folded at export)."""
+
+    NUM_CHUNKS = 7
+
+    def __init__(self, dim, t_emb_dim):
+        super().__init__()
+        self.proj = nn.Linear(t_emb_dim, self.NUM_CHUNKS * dim, bias=True)
+
+    def forward(self, t_emb):
+        h = self.proj(F.silu(t_emb))
+        return tuple(c[:, None, None, None, :] for c in h.chunk(self.NUM_CHUNKS, dim=-1))
+
+
+class DiffusionNABlock(nn.Module):
+    """NA + SwiGLU with shared AdaLN-Zero scale/shift (ungated residuals)."""
+
+    def __init__(self, dim, kernel_size, context_channels, head_dim=64, mlp_ratio=4.0):
+        super().__init__()
+        self.context_proj = nn.Linear(context_channels, dim, bias=True)
+        self.scale_shift_table = nn.Parameter(torch.zeros(AdaLNZero.NUM_CHUNKS, dim))
+        self.norm1 = RMSNorm(dim, eps=1e-6)
+        self.attn = NeighborhoodAttention3D(dim, kernel_size, head_dim=head_dim)
+        self.norm2 = RMSNorm(dim, eps=1e-6)
+        hidden = (int(dim * mlp_ratio) + 15) // 16 * 16
+        self.mlp = SwiGLU(dim, hidden)
+
+    def forward(self, x, latent_context, modulation):
+        scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [
+            modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(AdaLNZero.NUM_CHUNKS)
+        ]
+        chunk = max(1, MLP_TOKEN_CHUNK // max(x.shape[2] * x.shape[3], 1))
+        for t0 in range(0, x.shape[1], chunk):
+            x[:, t0:t0 + chunk] += self.context_proj(latent_context[:, t0:t0 + chunk])
+        x = self.attn(x, pre=lambda s: modulate(self.norm1(s), scale_msa, shift_msa), add_to=x)
+        return self.mlp(x, pre=lambda s: modulate(self.norm2(s), scale_mlp, shift_mlp), add_to=x)
+
+
+class LinearPixelShuffleUpsample(nn.Module):
+    """Linear channel-expand, then channels-last pixel shuffle."""
+
+    def __init__(self, in_channels, stride, out_channels_reduction_factor=1):
+        super().__init__()
+        self.stride = tuple(stride)
+        proj_out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
+        self.out_channels = proj_out_channels // math.prod(stride)
+        self.proj = nn.Linear(in_channels, proj_out_channels, bias=True)
+
+    def forward(self, x, drop_leading_frame=True):
+        batch, t, h, w, _ = x.shape
+        p1, p2, p3 = self.stride
+        out = torch.empty((batch, t * p1, h * p2, w * p3, self.out_channels), dtype=x.dtype, device=x.device)
+        chunk = max(1, MLP_TOKEN_CHUNK // max(h * w, 1))
+        for t0 in range(0, t, chunk):
+            t1 = min(t0 + chunk, t)
+            out[:, t0 * p1:t1 * p1] = rearrange(
+                self.proj(x[:, t0:t1]), "b t h w (c p1 p2 p3) -> b (t p1) (h p2) (w p3) c",
+                p1=p1, p2=p2, p3=p3,
+            )
+        if p1 == 2 and drop_leading_frame:
+            # The causal temporal pixel-shuffle duplicates the leading frame.
+            out = out[:, 1:]
+        return out
+
+
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal(256) -> MLP. ``mlp.{0,2}`` naming matches the checkpoint."""
+
+    def __init__(self, t_emb_dim=384, freq_dim=256):
+        super().__init__()
+        self.freq_dim = freq_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(freq_dim, t_emb_dim, bias=True),
+            nn.SiLU(),
+            nn.Linear(t_emb_dim, t_emb_dim, bias=True),
+        )
+
+    def forward(self, timestep, dtype):
+        emb = get_timestep_embedding(timestep.flatten(), self.freq_dim, flip_sin_to_cos=True,
+                                     downscale_freq_shift=0, scale=1)
+        return self.mlp(emb.to(dtype))
+
+
+class NADiffusionDecoder(nn.Module):
+    """Stages 1-4 (deterministic NA upsample) + stage-5 diffusion blocks.
+
+    Input latent must already be un-normalized (the wrapper applies
+    ``per_channel_statistics.un_normalize``, same as the conv VAE path).
+    """
+
+    def __init__(
+        self,
+        in_channels=128,
+        out_channels=3,
+        patch_size=4,
+        head_dim=64,
+        stage_channels=(2048, 1024, 512, 512, 256),
+        stage_depths=(4, 6, 4, 2, 8),
+        stage_kernels=((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5), (11, 11, 11)),
+        upsamples=(((1, 2, 2), 2), ((2, 1, 1), 2), ((2, 2, 2), 1), ((2, 2, 2), 2)),
+        stage5_kernel=(11, 11, 11),
+        t_emb_dim=384,
+        default_num_inference_steps=1,
+        timestep_scale_multiplier=1000.0,
+        model_output_type="x0",
+    ):
+        super().__init__()
+        self.patch_size = patch_size
+        self.out_channels = out_channels
+        self.timestep_scale_multiplier = timestep_scale_multiplier
+        self.model_output_type = model_output_type
+        self.register_buffer(
+            "default_inference_timesteps",
+            torch.linspace(1.0, 1.0 / default_num_inference_steps, default_num_inference_steps),
+            persistent=False,
+        )
+        self.temporal_upscale = math.prod(s[0] for s, _ in upsamples)
+        self.spatial_upscale = math.prod(s[1] for s, _ in upsamples) * patch_size
+        # NATTEN-style last-frame border mitigation: replicate the last latent
+        # frame through stages 1-4, crop the appendix off the context after.
+        self.trailing_pad_latent_frames = (stage_kernels[0][0] // 2) * 2
+
+        self.conv_in = nn.Linear(in_channels, stage_channels[0], bias=True)
+
+        self.det_stages = nn.ModuleList()
+        self.upsamples = nn.ModuleList()
+        for stage_i in range(len(stage_channels) - 1):
+            c = stage_channels[stage_i]
+            self.det_stages.append(nn.ModuleList(
+                [NABlock(c, stage_kernels[stage_i], head_dim=head_dim) for _ in range(stage_depths[stage_i])]
+            ))
+            stride, reduction = upsamples[stage_i]
+            self.upsamples.append(LinearPixelShuffleUpsample(c, stride, out_channels_reduction_factor=reduction))
+
+        self.t_embedder = TimestepEmbedder(t_emb_dim=t_emb_dim)
+
+        c5 = stage_channels[-1]
+        self.context_channels = c5
+        noised_pixel_channels = out_channels * (patch_size ** 2)
+        self.conv_in_x_t = nn.Linear(noised_pixel_channels, c5, bias=True)
+        self.shared_adaln = AdaLNZero(c5, t_emb_dim)
+        self.diff_blocks = nn.ModuleList([
+            DiffusionNABlock(c5, stage5_kernel, context_channels=c5, head_dim=head_dim)
+            for _ in range(stage_depths[-1])
+        ])
+        self.norm_out = RMSNorm(c5, eps=1e-6)
+        self.conv_out = nn.Linear(c5, noised_pixel_channels, bias=True)
+
+    def forward_pre_diffusion(self, z, drop_leading_frame=True, pad_trailing=True):
+        """Stages 1-4: latent -> stage-5 context, channels-last.
+
+        ``drop_leading_frame`` must be True only when ``z`` contains the
+        latent's true temporal origin (t=0); tiled callers decoding a later
+        temporal chunk pass False (the duplicate leading frame belongs solely
+        to the origin chunk). ``pad_trailing`` only for chunks containing the
+        latent's last frame."""
+        n = self.trailing_pad_latent_frames if pad_trailing else 0
+        if n > 0:
+            z = torch.cat([z, z[:, :, -1:].expand(-1, -1, n, -1, -1)], dim=2)
+        x = z.permute(0, 2, 3, 4, 1)
+        x = self.conv_in(x)
+        for stage_i, blocks in enumerate(self.det_stages):
+            for block in blocks:
+                x = block(x)
+            x = self.upsamples[stage_i](x, drop_leading_frame=drop_leading_frame)
+        if n > 0:
+            x = x[:, :-(n * self.temporal_upscale)]
+        return x
+
+    def forward_diff_step(self, context, x_t, t):
+        x = patchify(x_t, patch_size_hw=self.patch_size, patch_size_t=1)
+        x = self.conv_in_x_t(x.permute(0, 2, 3, 4, 1))
+        t_emb = self.t_embedder(self.timestep_scale_multiplier * t, dtype=x.dtype)
+        modulation = self.shared_adaln(t_emb)
+        for block in self.diff_blocks:
+            x = block(x, context, modulation)
+        x = self.norm_out(x)
+        x = self.conv_out(x)
+        x = x.permute(0, 4, 1, 2, 3)
+        return unpatchify(x, patch_size_hw=self.patch_size, patch_size_t=1)
+
+    def forward(self, z, generator=None, drop_leading_frame=True, pad_trailing=True):
+        context = self.forward_pre_diffusion(z, drop_leading_frame=drop_leading_frame, pad_trailing=pad_trailing)
+        batch, t5, h5, w5, _ = context.shape
+        pixel_shape = (batch, self.out_channels, t5, h5 * self.patch_size, w5 * self.patch_size)
+        x_t = torch.randn(pixel_shape, dtype=z.dtype, device=z.device, generator=generator)
+
+        timesteps = self.default_inference_timesteps.to(z.device)
+        num_steps = timesteps.shape[0]
+        for i in range(num_steps):
+            t_now = timesteps[i].expand(batch)
+            model_out = self.forward_diff_step(context, x_t, t_now)
+            if self.model_output_type == "x0":
+                x0 = model_out
+                if i == num_steps - 1:
+                    return x0
+                velocity = (x_t.float() - x0.float()) / timesteps[i]
+            else:  # "v"
+                velocity = model_out.float()
+                if i == num_steps - 1:
+                    return (x_t.float() - timesteps[i] * velocity).to(z.dtype)
+            t_next = timesteps[i + 1] if i + 1 < num_steps else torch.zeros_like(timesteps[i])
+            x_t = (x_t.float() - (timesteps[i] - t_next) * velocity).to(z.dtype)
+        return x_t
+
+
+LTX_24_VAE_CONFIG = {
+    "_class_name": "CausalDiffusionVAE",
+    "dims": 3,
+    "model_output_type": "x0",
+    "encoder": {
+        "dims": 3,
+        "in_channels": 3,
+        "out_channels": 128,
+        "blocks": [
+            ["res_x", {"num_layers": 4}],
+            ["compress_space_res", {"multiplier": 2}],
+            ["res_x", {"num_layers": 6}],
+            ["compress_time_res", {"multiplier": 2}],
+            ["res_x", {"num_layers": 4}],
+            ["compress_all_res", {"multiplier": 2}],
+            ["res_x", {"num_layers": 2}],
+            ["compress_all_res", {"multiplier": 1}],
+            ["res_x", {"num_layers": 2}],
+        ],
+        "patch_size": 4,
+        "latent_log_var": "constant",
+        "norm_layer": "pixel_norm",
+        "base_channels": 128,
+        "spatial_padding_mode": "zeros",
+    },
+    "decoder": {
+        "in_channels": 128,
+        "out_channels": 3,
+        "patch_size": 4,
+        "head_dim": 64,
+        "stage_channels": [2048, 1024, 512, 512, 256],
+        "stage_depths": [4, 6, 4, 2, 8],
+        "stage_kernels": [[3, 7, 7], [3, 7, 7], [3, 5, 5], [3, 5, 5], [11, 11, 11]],
+        "upsamples": [[[1, 2, 2], 2], [[2, 1, 1], 2], [[2, 2, 2], 1], [[2, 2, 2], 2]],
+        "stage5_kernel": [11, 11, 11],
+        "timestep_scale_multiplier": 1000.0,
+        "default_num_inference_steps": 1,
+    },
+}
+
+
+class CausalDiffusionVAE(nn.Module):
+    """LTX 2.4 video VAE: conv encoder (shared with the 2.0 arch) + NA
+    diffusion decoder. Interface mirrors ``causal_video_autoencoder.VideoVAE``.
+    """
+
+    def __init__(self, config=None):
+        super().__init__()
+        if config is None:
+            config = LTX_24_VAE_CONFIG
+        self.config = config
+        enc = config.get("encoder", LTX_24_VAE_CONFIG["encoder"])
+        dec = config.get("decoder", LTX_24_VAE_CONFIG["decoder"])
+        dec_defaults = LTX_24_VAE_CONFIG["decoder"]
+
+        self.encoder = Encoder(
+            dims=enc.get("dims", 3),
+            in_channels=enc.get("in_channels", 3),
+            out_channels=enc.get("out_channels", 128),
+            blocks=enc.get("blocks", LTX_24_VAE_CONFIG["encoder"]["blocks"]),
+            patch_size=enc.get("patch_size", 4),
+            latent_log_var=enc.get("latent_log_var", "constant"),
+            norm_layer=enc.get("norm_layer", "pixel_norm"),
+            spatial_padding_mode=enc.get("spatial_padding_mode", "zeros"),
+            base_channels=enc.get("base_channels", 128),
+        )
+
+        self.decoder = NADiffusionDecoder(
+            in_channels=dec.get("in_channels", 128),
+            out_channels=dec.get("out_channels", 3),
+            patch_size=dec.get("patch_size", 4),
+            head_dim=dec.get("head_dim", 64),
+            stage_channels=tuple(dec.get("stage_channels", dec_defaults["stage_channels"])),
+            stage_depths=tuple(dec.get("stage_depths", dec_defaults["stage_depths"])),
+            stage_kernels=tuple(tuple(k) for k in dec.get("stage_kernels", dec_defaults["stage_kernels"])),
+            upsamples=tuple((tuple(s), r) for s, r in dec.get("upsamples", dec_defaults["upsamples"])),
+            stage5_kernel=tuple(dec.get("stage5_kernel", dec_defaults["stage5_kernel"])),
+            t_emb_dim=dec.get("t_emb_dim", 384),
+            default_num_inference_steps=dec.get("default_num_inference_steps", 1),
+            timestep_scale_multiplier=dec.get("timestep_scale_multiplier", 1000.0),
+            model_output_type=config.get("model_output_type", "x0"),
+        )
+
+        self.per_channel_statistics = processor()
+
+    def encode(self, x, device=None):
+        x = x[:, :, :max(1, 1 + ((x.shape[2] - 1) // 8) * 8), :, :]
+        means, logvar = torch.chunk(self.encoder(x, device=device), 2, dim=1)
+        return self.per_channel_statistics.normalize(means)
+
+    def decode(self, x):
+        # Fixed-seed noise so decodes are reproducible TODO: expose?
+        generator = torch.Generator(device=x.device)
+        generator.manual_seed(0)
+        return self.decoder(self.per_channel_statistics.un_normalize(x), generator=generator)

comfy/ldm/minimax/audio_vae.py

@@ -35,7 +35,8 @@ class Snake1d(nn.Module):
         self.alpha = nn.Parameter(torch.empty(1, channels, 1))
 
     def forward(self, x):
-        return snake(x, self.alpha, self.alpha)
+        alpha = comfy.ops.cast_to_input(self.alpha, x)
+        return snake(x, alpha, alpha)
 
 
 class SnakeBeta(nn.Module):
@@ -47,8 +48,8 @@ class SnakeBeta(nn.Module):
         self.beta = nn.Parameter(torch.empty(in_features))
 
     def forward(self, x):
-        alpha = torch.exp(self.alpha).view(1, -1, 1)
-        beta = torch.exp(self.beta).view(1, -1, 1)
+        alpha = torch.exp(comfy.ops.cast_to_input(self.alpha, x)).view(1, -1, 1)
+        beta = torch.exp(comfy.ops.cast_to_input(self.beta, x)).view(1, -1, 1)
         return snake(x, alpha, beta)
 
 
@@ -98,7 +99,7 @@ class UpSample1d(nn.Module):
     def forward(self, x):
         _, C, _ = x.shape
         x = F.pad(x, (self.pad, self.pad), mode="replicate")
-        x = F.conv_transpose1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C).mul_(self.ratio)
+        x = F.conv_transpose1d(x, comfy.ops.cast_to_input(self.filter.expand(C, -1, -1), x), stride=self.stride, groups=C).mul_(self.ratio)
         x = x[..., self.pad_left:-self.pad_right]
         return x
 
@@ -114,7 +115,7 @@ class LowPassFilter1d(nn.Module):
     def forward(self, x):
         _, C, _ = x.shape
         x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
-        return F.conv1d(x, self.filter.expand(C, -1, -1).to(x.dtype), stride=self.stride, groups=C)
+        return F.conv1d(x, comfy.ops.cast_to_input(self.filter.expand(C, -1, -1), x), stride=self.stride, groups=C)
 
 
 class DownSample1d(nn.Module):
@@ -239,7 +240,7 @@ class CausalAttention(nn.Module):
     def forward(self, x):
         B, N, C = x.shape
         weight, _, offload_stream = comfy.ops.cast_bias_weight(self.qkv, x, offloadable=True)
-        qkv = F.linear(x, weight=weight, bias=torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)))
+        qkv = F.linear(x, weight=weight, bias=comfy.ops.cast_to_input(torch.cat((self.q_bias, self.zero_k_bias, self.v_bias)), x))
         comfy.ops.uncast_bias_weight(self.qkv, weight, None, offload_stream)
         q, k, v = qkv.reshape(B, N, 3, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4).unbind(0)
 

comfy/ldm/minimax/model.py

@@ -9,8 +9,9 @@ The packed sequence is:
 Timestep domain: the model receives the *video* sigma from the sampler and
 derives per-token timesteps t = 1 - sigma internally; the audio stream runs on
 its own shifted schedule (sigma_shift video 12.0 / audio 3.0), mapped from the
-video sigma in closed form. The audio velocity is returned scaled by the
-schedule map's derivative d(sigma_a)/d(sigma_v).
+video sigma in closed form. The sampler carries the audio latent scaled onto the
+video schedule (ModelSamplingAV); forward() undoes that scale and converts the
+velocity back, so _forward only ever sees the stream's own latent.
 """
 
 import math
@@ -24,7 +25,7 @@ import comfy.model_prefetch
 import comfy.ops
 import comfy.patcher_extension
 import comfy.quant_ops
-from comfy.ldm.modules.attention import optimized_attention
+from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention
 
 FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
 FRAME_RESCALE = 5.0 / 3.0
@@ -38,17 +39,6 @@ def time_shift_sigma(sigma, from_shift, to_shift):
     return to_shift * base / (1.0 + (to_shift - 1.0) * base)
 
 
-def time_shift_slope(sigma, from_shift, to_shift):
-    """d(sigma_to)/d(sigma_from) at the same base-grid point.
-
-    Scaling a stream's returned velocity by this slope makes the flat ODE that
-    any sampler integrates on the from-schedule equal to that stream's true ODE
-    on its own schedule.
-    """
-    base = sigma / (from_shift + sigma * (1.0 - from_shift))
-    return (to_shift * (1.0 + (from_shift - 1.0) * base) ** 2) / (from_shift * (1.0 + (to_shift - 1.0) * base) ** 2)
-
-
 def patchify_video(latent, patch_size=(1, 2, 2)):
     # [B, C, T, H, W] -> [B*t*h*w, C*pt*ph*pw]
     b, c, t_full, h_full, w_full = latent.shape
@@ -101,6 +91,18 @@ def _video_t_grid(n, origin):
     return float(origin) + torch.cat([torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)])
 
 
+def _ref_t_span(blk):
+    # time-axis span a reference block occupies ahead of the target streams
+    kind = blk["kind"]
+    if kind == "image":
+        return 1.0
+    if kind == "audio":
+        return float(blk["ref_audio_t"])
+    if kind in ("video", "video_audio"):
+        return max(float(blk["ref_audio_t"]), sum(_video_t_spans(blk["latent_t"])))
+    return 0.0
+
+
 def _audio_grid(cursor, t, w_low, w_high):
     # channel-major stereo rows: t advances per latent frame, w pinned to the grid extremes per stereo channel, h stays 0
     g = torch.zeros(t * 2, 3, dtype=torch.float64)
@@ -175,9 +177,10 @@ class Attention(nn.Module):
         else:
             q = self.q_norm(q.view(s, self.heads, self.head_dim))
             k = self.k_norm(k.view(s, self.heads, self.head_dim))
-        q = q.transpose(0, 1).unsqueeze(0)
-        k = k.transpose(0, 1).unsqueeze(0)
-        v = v.transpose(0, 1).unsqueeze(0)
+        v = v.clone()
+        q = AttentionTensorContainer(q.transpose(0, 1).unsqueeze(0))
+        k = AttentionTensorContainer(k.transpose(0, 1).unsqueeze(0))
+        v = AttentionTensorContainer(v.transpose(0, 1).unsqueeze(0))
         out = optimized_attention(q, k, v, self.heads, mask=None, skip_reshape=True, transformer_options=transformer_options)
         return self.out_proj(out.squeeze(0))
 
@@ -297,7 +300,7 @@ class FinalLayer(nn.Module):
 class PackedLayout:
     """Static packed-sequence structure for one shape/conditioning signature."""
 
-    def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None, frame_count=None):
+    def __init__(self, text_len, latent_t, latent_h, latent_w, audio_t, keyframes=None, refs=None):
         frame, w_grid = _frame_grid(latent_h, latent_w)
         frame_rows = frame.shape[0]
 
@@ -308,29 +311,37 @@ class PackedLayout:
 
         img_pos, img_update = [], []
         audio_pos, audio_update = [], []
-        cursor = text_len
         row = text_len
 
-        if keyframes:
-            # fl2va: keyframe cond rows right after text, sharing the target spatial grid
-            for kf in keyframes:
-                pixel_index = kf["resolved_frame_index"]
-                if pixel_index == 0:
-                    cond_t = float(text_len)
-                elif frame_count is not None and pixel_index == frame_count - 1:
-                    cond_t = float(text_len) + sum(_video_t_spans(latent_t)) - FRAME_RESCALE
-                else:
-                    raise ValueError("only first/last keyframe anchors are supported")
-                g = torch.empty(frame_rows, 3, dtype=torch.float64)
-                g[:, 0] = cond_t
-                g[:, 1:] = frame
-                segments.append(("cond", frame_rows))
-                pos.append(g)
-                img_pos.append(torch.arange(row, row + frame_rows))
-                img_update.append(torch.zeros(frame_rows, dtype=torch.bool))
-                row += frame_rows
-
         target_audio_w = (float(w_grid[0]), float(w_grid[-1]))
+        # refs pack between text and the targets, so the target timeline starts after their spans
+        cursor = float(text_len)
+        for blk in refs or ():
+            cursor += _ref_t_span(blk)
+
+        if keyframes:
+            # fl2va: keyframe cond rows right after text, sharing the target spatial grid;
+            # anchors count from the target timeline origin, FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
+            for kf in keyframes:
+                cond_t = cursor + FRAME_RESCALE * kf["resolved_frame_index"]
+                video_latent = kf.get("latent")
+                if video_latent is not None:
+                    vt = video_latent.shape[2]
+                    n = vt * frame_rows
+                    segments.append(("cond", n))
+                    pos.append(_video_grid(vt, frame, cond_t))
+                    img_pos.append(torch.arange(row, row + n))
+                    img_update.append(torch.zeros(n, dtype=torch.bool))
+                    row += n
+                audio_latent = kf.get("audio_latent")
+                if audio_latent is not None:
+                    rt = audio_latent.shape[-1]
+                    segments.append(("cond_audio", rt * 2))
+                    pos.append(_audio_grid(cond_t, rt, *target_audio_w))
+                    audio_pos.append(torch.arange(row, row + rt * 2))
+                    audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
+                    row += rt * 2
+
         if refs:
             cursor = float(text_len)
             for blk in refs:
@@ -398,7 +409,7 @@ class PackedLayout:
         self.audio_update = torch.cat(audio_update)
         self.signature = (text_len, latent_t, latent_h, latent_w, audio_t)
         # contiguous segment table (start, stop, kind)
-        # kinds: text / cond / ref_img / ref_audio / audio / video
+        # kinds: text / cond / cond_audio / ref_img / ref_audio / audio / video
         # the packed sequence is uniform per segment in (modality tag, timestep class),
         # except the text span (tag runs resolved at forward time from the presentation tags)
         seg_abs = []
@@ -496,12 +507,30 @@ class MiniMaxH3Model(nn.Module):
         return torch.cat(rows, dim=0) if rows else None
 
     def forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs):
-        return comfy.patcher_extension.WrapperExecutor.new_class_executor(
+        # the sampler carries the audio as (sigma_v / sigma_a) * x_audio; undo it outside
+        # the wrappers so they and the network see the stream's own latent and velocity
+        scale = float((minimax_payload or {}).get("audio_scale", 1.0))
+        audio_src = x[1]
+        if scale != 1.0:
+            shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
+            shift_a = float(transformer_options.get("minimax_h3_sigma_shift_audio", self.sigma_shift_audio))
+            sigma_v = (timestep.flatten()[0] / 1000.0).float().clamp(min=1e-6)
+            sigma_a = time_shift_sigma(sigma_v, shift_v, shift_a)
+            carry = (sigma_a / sigma_v).to(audio_src.dtype)
+            x = [x[0], audio_src * carry]
+
+        out = 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, transformer_options, minimax_payload=minimax_payload, **kwargs)
 
+        if scale != 1.0:
+            # d/d(sigma_v) of the carried variable
+            out[1] = ((1.0 - scale) * (audio_src * carry)
+                      + (1.0 + (scale - 1.0) * sigma_a).to(out[1].dtype) * out[1])
+        return out
+
     def _forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, **kwargs):
         video_x, audio_x = x[0], x[1]
         orig_t, orig_h, orig_w = video_x.shape[2], video_x.shape[3], video_x.shape[4]
@@ -520,8 +549,7 @@ class MiniMaxH3Model(nn.Module):
         if layout is None or layout.signature != (text_len, latent_t, lat_h, lat_w, audio_t):
             layout = PackedLayout(text_len, latent_t, lat_h, lat_w, audio_t,
                                   keyframes=payload.get("keyframes"),
-                                  refs=payload.get("refs"),
-                                  frame_count=payload.get("frame_count"))
+                                  refs=payload.get("refs"))
 
         # model_base passes model_sampling.timestep(sigma) = sigma * 1000
         shift_v = float(transformer_options.get("minimax_h3_sigma_shift_video", self.sigma_shift_video))
@@ -534,14 +562,14 @@ class MiniMaxH3Model(nn.Module):
         vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
         aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP))
         has_vis_cond = any(k in ("cond", "ref_img") for _, _, k in layout.segments)
-        has_aud_cond = any(k == "ref_audio" for _, _, k in layout.segments)
+        has_aud_cond = any(k in ("cond_audio", "ref_audio") for _, _, k in layout.segments)
         seg_t = {"text": t_v, "video": t_v, "audio": t_a,
                  "cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
-                 "ref_audio": max(t_a, aud_aug)}
+                 "cond_audio": max(t_a, aud_aug), "ref_audio": max(t_a, aud_aug)}
         unique_t = sorted({t_v, t_a} | ({seg_t["cond"]} if has_vis_cond else set())
                           | ({seg_t["ref_audio"]} if has_aud_cond else set()))
         t_row = {t: i for i, t in enumerate(unique_t)}
-        seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "ref_audio": 2}
+        seg_tag = {"text": 1, "video": 0, "audio": 2, "cond": 0, "ref_img": 0, "cond_audio": 2, "ref_audio": 2}
 
         text_tags = payload.get("text_token_tags")
         mod_segments = []
@@ -639,8 +667,4 @@ class MiniMaxH3Model(nn.Module):
         video_out = video_out[:, :, :orig_t, :orig_h, :orig_w]
         audio_out = unpack_audio(a)
 
-        # The sampler integrates the flat ODE dX/dsigma_v = (X - denoised)/sigma_v.
-        # Scaling the audio velocity by d(sigma_a)/d(sigma_v) makes that ODE equal
-        # to the audio stream's true ODE on its own shifted schedule.
-        slope_a = time_shift_slope(sigma_v, shift_v, shift_a).to(audio_out.dtype)
-        return [-video_out.to(video_x.dtype), (-slope_a) * audio_out.to(audio_x.dtype)]
+        return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]

comfy/ldm/minimax/vae.py

@@ -6,6 +6,7 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
+import comfy.model_management
 import comfy.ops
 import comfy.quant_ops
 import comfy.rmsnorm
@@ -199,11 +200,11 @@ class RotaryEmbeddingND(nn.Module):
 
 class FeedForward(nn.Module):
     # Gated SiLU FFN.
-    def __init__(self, dim, mult=4, bias=True):
+    def __init__(self, dim, mult=4, bias=True, operations=ops):
         super().__init__()
         inner_dim = dim * mult
-        self.w1 = ops.Linear(dim, inner_dim * 2, bias=bias)
-        self.w2 = ops.Linear(inner_dim, dim, bias=bias)
+        self.w1 = operations.Linear(dim, inner_dim * 2, bias=bias)
+        self.w2 = operations.Linear(inner_dim, dim, bias=bias)
 
     def forward(self, x):
         gate, x = self.w1(x).chunk(2, dim=-1)
@@ -211,15 +212,15 @@ class FeedForward(nn.Module):
 
 
 class Attention(nn.Module):
-    def __init__(self, heads, dim_head, bias=True, eps=1e-5):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
         super().__init__()
         self.dim_head = dim_head
         self.heads = heads
         inner_dim = dim_head * heads
         self.norm_q = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
         self.norm_k = ops.RMSNorm(dim_head, eps=eps, elementwise_affine=False)
-        self.to_qkv = ops.Linear(inner_dim, inner_dim * 3, bias=bias)
-        self.to_out = ops.Linear(inner_dim, inner_dim, bias=bias)
+        self.to_qkv = operations.Linear(inner_dim, inner_dim * 3, bias=bias)
+        self.to_out = operations.Linear(inner_dim, inner_dim, bias=bias)
 
     def forward(self, x, rotary_pos_emb=None):
         batch_size, seq_len, _ = x.shape
@@ -242,24 +243,24 @@ class Attention(nn.Module):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, heads, dim_head, bias=True, eps=1e-5):
+    def __init__(self, heads, dim_head, bias=True, eps=1e-5, operations=ops):
         super().__init__()
         dim = heads * dim_head
         self.norm1 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
-        self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
+        self.attn = Attention(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
         self.scale1 = nn.Parameter(torch.empty(dim))
         self.norm2 = ops.RMSNorm(dim, elementwise_affine=True, eps=eps)
-        self.ff = FeedForward(dim=dim, bias=bias)
+        self.ff = FeedForward(dim=dim, bias=bias, operations=operations)
         self.scale2 = nn.Parameter(torch.empty(dim))
 
     def forward(self, x, rotary_pos_emb=None):
-        x = x.addcmul_(self.attn(comfy.rmsnorm.rms_norm(x, self.norm1.weight, self.norm1.eps), rotary_pos_emb), self.scale1)
-        return x.addcmul_(self.ff(comfy.rmsnorm.rms_norm(x, self.norm2.weight, self.norm2.eps)), self.scale2)
+        x = x.addcmul_(self.attn(comfy.rmsnorm.rms_norm(x, self.norm1.weight, self.norm1.eps), rotary_pos_emb), comfy.ops.cast_to_input(self.scale1, x))
+        return x.addcmul_(self.ff(comfy.rmsnorm.rms_norm(x, self.norm2.weight, self.norm2.eps)), comfy.ops.cast_to_input(self.scale2, x))
 
 
 class ViT3DDecoder(nn.Module):
     def __init__(self, patch_size=16, patch_size_t=4, in_channels=24, out_channels=3, num_layers=36, heads=32, dim_head=64, rope_theta=100.0,
-                 rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4):
+                 rope_dim_ratio=0.75, bias=True, eps=1e-5, num_register_tokens=4, operations=ops):
         super().__init__()
         dim = heads * dim_head
         self.patch_size = patch_size
@@ -274,7 +275,7 @@ class ViT3DDecoder(nn.Module):
         self.register_buffer("mask_token", torch.empty(1, 1, dim))
 
         self.transformer_blocks = nn.ModuleList(
-            [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps)
+            [TransformerBlock(heads=heads, dim_head=dim_head, bias=bias, eps=eps, operations=operations)
              for _ in range(num_layers)]
         )
 
@@ -289,7 +290,7 @@ class ViT3DDecoder(nn.Module):
         num_patches = h.shape[1]
         num_suffix = 1 + self.num_register_tokens
 
-        h = torch.cat([h, self.register_tokens.expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1)
+        h = torch.cat([h, comfy.ops.cast_to_input(self.register_tokens, h).expand(B, -1, -1), torch.zeros_like(h[:, 0:1, :])], dim=1)
 
         img_ids = create_token_ids((latent_T, latent_H, latent_W), x.device, x.dtype).expand(B, -1, -1)
         suffix_ids = torch.zeros((B, num_suffix, 3), device=x.device, dtype=img_ids.dtype)
@@ -321,6 +322,8 @@ class ViT3DDecoder(nn.Module):
 # Full VAE
 
 class MiniMaxH3VideoVAE(nn.Module):
+    comfy_has_chunked_io = True
+
     def __init__(
         self,
         in_channels=3,
@@ -337,6 +340,7 @@ class MiniMaxH3VideoVAE(nn.Module):
         tile_size=256,
         tile_overlap_min=64,
         tiling=True,
+        operations=ops,
     ):
         super().__init__()
         self.vae_ratio = int(math.prod(space_down))
@@ -372,6 +376,7 @@ class MiniMaxH3VideoVAE(nn.Module):
             patch_size_t=self.vae_ratio_t,
             in_channels=z_channels,
             out_channels=out_ch,
+            operations=operations,
         )
 
         self.register_buffer("latents_mean", torch.tensor(LATENTS_MEAN))
@@ -387,6 +392,23 @@ class MiniMaxH3VideoVAE(nn.Module):
     def _decode_pixels(self, z):
         return self.decoder(self.post_quant_conv(z))
 
+    def _normalize_pixels(self, x):
+        return x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
+
+    def _finalize_pixels(self, part):
+        # raw decoder output -> float32 pixels in [0, 1] (the VAE wrapper's process_output is identity)
+        part = part * self.pixel_std.to(device=part.device, dtype=torch.float32)
+        return part.add_(self.pixel_mean.to(device=part.device, dtype=torch.float32)).clamp_(0.0, 1.0)
+
+    def decode_output_shape(self, input_shape):
+        b, c, t, h, w = input_shape
+        if t == 1:
+            frames = 1
+        else:
+            pad_tokens, num_chunks = self._decode_temporal_chunks(t)
+            frames = self._decode_temporal_frame_plan(t + pad_tokens, num_chunks, pad_tokens)
+        return (b, self.decoder.out_channels, frames, h * self.vae_ratio, w * self.vae_ratio)
+
     def _adaptive_encode(self, x):
         if self.tiling:
             return self.tiled_encode(x)
@@ -519,18 +541,15 @@ class MiniMaxH3VideoVAE(nn.Module):
 
     # temporal chunking
 
-    def encode_temporal(self, x):
-        if x.shape[2] % self.clip_length != 0:
-            pad_size = (-x.shape[2]) % self.clip_length
-            pad_frames = x[:, :, -1:].repeat(1, 1, pad_size, 1, 1)
-            x = torch.cat([x, pad_frames], dim=2)
-
-        num_chunks = x.shape[2] // self.clip_length
-
+    def encode_temporal(self, x, device):
+        # chunked input io: x may live on the CPU, clips move to the device as they encode
         z_list = []
-        for i in range(num_chunks):
-            clip_x = x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :]
-            z_list.append(self._adaptive_encode(clip_x))
+        for i in range(math.ceil(x.shape[2] / self.clip_length)):
+            clip_x = x[:, :, i * self.clip_length:(i + 1) * self.clip_length, :, :].to(device)
+            if clip_x.shape[2] < self.clip_length:
+                pad_frames = clip_x[:, :, -1:].repeat(1, 1, self.clip_length - clip_x.shape[2], 1, 1)
+                clip_x = torch.cat([clip_x, pad_frames], dim=2)
+            z_list.append(self._adaptive_encode(self._normalize_pixels(clip_x)))
 
         z = torch.cat(z_list, dim=2)
         if self.token_drop > 0:
@@ -575,43 +594,42 @@ class MiniMaxH3VideoVAE(nn.Module):
         total_frames += final_overlap_frames
         return total_frames - self._decode_temporal_pad_frames(z_len, pad_tokens)
 
-    def decode_temporal(self, z):
-        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
-        split_count = int(self.token_drop > 0) + 1
-
-        pseudo_total_tokens = z.shape[2] + self.token_drop
-
-        pad_tokens = 0
-        remainder = pseudo_total_tokens % self.tokens_chunk_size
-        if remainder != 0:
-            pad_tokens = self.tokens_chunk_size - remainder
-            pseudo_total_tokens += pad_tokens
+    def _decode_temporal_chunks(self, z_len):
+        pseudo_total_tokens = z_len + self.token_drop
+        pad_tokens = (-pseudo_total_tokens) % self.tokens_chunk_size
+        pseudo_total_tokens += pad_tokens
 
         num_chunks = pseudo_total_tokens // self.tokens_chunk_size - int(self.token_drop > 0)
         if num_chunks < 1:
             # too few tokens for one chunk (e.g. T_lat == 2): pad one extra chunk
             pad_tokens += self.tokens_chunk_size
             num_chunks += 1
+        return pad_tokens, num_chunks
 
+    def decode_temporal(self, z, output_buffer=None):
+        chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
+        split_count = int(self.token_drop > 0) + 1
+
+        if output_buffer is None:
+            # finalized chunks stream out of VRAM so the full video never sits on the GPU
+            output_buffer = torch.empty(self.decode_output_shape(z.shape), dtype=torch.float32,
+                                        device=comfy.model_management.intermediate_device())
+
+        pad_tokens, num_chunks = self._decode_temporal_chunks(z.shape[2])
         if pad_tokens > 0:
             pad_z = z[:, :, -1:, :, :].repeat(1, 1, pad_tokens, 1, 1)
             z = torch.cat([z, pad_z], dim=2)
 
-        output_frames = self._decode_temporal_frame_plan(z.shape[2], num_chunks, pad_tokens)
-
-        dec = None
+        dec = output_buffer
         dec_overlap = None
         write_pos = 0
 
         def write_part(part):
-            nonlocal dec, write_pos
+            nonlocal write_pos
             part_frames = part.shape[2]
             if part_frames <= 0:
                 return
-            if dec is None:
-                out_shape = list(part.shape)
-                out_shape[2] = output_frames
-                dec = torch.empty(out_shape, dtype=part.dtype, device=part.device)
+            part = self._finalize_pixels(part)
             copy_frames = min(part_frames, max(0, dec.shape[2] - write_pos))
             if copy_frames > 0:
                 dec[:, :, write_pos:write_pos + copy_frames, :, :].copy_(
@@ -651,18 +669,18 @@ class MiniMaxH3VideoVAE(nn.Module):
         return dec
 
 
-    def encode(self, x):
+    def encode(self, x, device=None):
         # x: [B, 3, T, H, W] in [-1, 1] -> normalized latents [B, 24, T_lat, H/16, W/16]
         if x.ndim == 4:
             x = x.unsqueeze(2)
-
-        x = x.add(1.0).mul_(0.5).sub_(self.pixel_mean.to(x)).div_(self.pixel_std.to(x))
+        if device is None:
+            device = x.device
 
         if x.shape[2] == 1:
-            moments = self._adaptive_encode(x)
+            moments = self._adaptive_encode(self._normalize_pixels(x.to(device)))
             moments = moments[:, :, -1:, :, :]
         else:
-            moments = self.encode_temporal(x)
+            moments = self.encode_temporal(x, device)
 
         mean = torch.chunk(moments.float(), 2, dim=1)[0]
 
@@ -677,18 +695,16 @@ class MiniMaxH3VideoVAE(nn.Module):
     def decode_tiled(self, z, **kwargs):
         return self.decode(z)
 
-    def decode(self, z):
-        # z: [B, 24, T_lat, H_lat, W_lat] normalized latents -> pixels [B, 3, T, H, W] in [-1, 1]
+    def decode(self, z, output_buffer=None):
+        # z: [B, 24, T_lat, H_lat, W_lat] normalized latents -> float32 pixels [B, 3, T, H, W] in [0, 1]
         latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(z)
         latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(z)
         z = z * latents_std + latents_mean
 
         if z.shape[2] == 1:
-            dec = self._adaptive_decode(z)
-            dec = dec[:, :, -1:, :, :]
-        else:
-            dec = self.decode_temporal(z)
-
-        dec = dec.float()
-        dec.mul_(self.pixel_std.to(dec)).add_(self.pixel_mean.to(dec)).clamp_(0.0, 1.0).mul_(2.0).sub_(1.0)
-        return dec
+            dec = self._finalize_pixels(self._adaptive_decode(z)[:, :, -1:, :, :])
+            if output_buffer is None:
+                return dec
+            output_buffer.copy_(dec)
+            return output_buffer
+        return self.decode_temporal(z, output_buffer)

comfy/ldm/minimax_music/ar.py

@@ -0,0 +1,343 @@
+import dataclasses
+import hashlib
+
+import torch
+from torch import nn
+
+import comfy.model_management
+import comfy.model_prefetch
+import comfy.ops
+import comfy.utils
+from comfy.ldm.modules.attention import optimized_attention_for_device
+from comfy.text_encoders.llama import Llama2_, Qwen3_8BConfig
+
+from .prompt import AUDIO_CODE_OFFSET, SPECIAL_TOKEN_IDS
+
+
+CFG_SCALE = 1.5
+CFG_TOP_K = 50
+C0_VOCAB_SIZE = 16384
+MAX_PROMPT_TOKENS = 5000
+MAX_AUDIO_FRAMES = 9000
+AUDIO_FRAMES_PER_SECOND = 25
+
+
+def derive_seed(seed, *parts):
+    digest = hashlib.blake2b(digest_size=8, person=b"minimax-ttm")
+    digest.update(int(seed).to_bytes(8, "little", signed=False))
+    for part in parts:
+        value = str(part).encode("utf-8")
+        digest.update(len(value).to_bytes(4, "little"))
+        digest.update(value)
+    return int.from_bytes(digest.digest(), "little") & ((1 << 63) - 1)
+
+
+def sample_topk(logits, top_k, generator):
+    values = torch.nan_to_num(logits.float(), nan=-1e9, posinf=1e9, neginf=-1e9)
+    top_k = min(top_k, values.shape[-1])
+    threshold = torch.topk(values, top_k, dim=-1).values[..., -1, None]
+    values = values.masked_fill(values < threshold, -float("inf"))
+    probabilities = torch.nan_to_num(torch.softmax(values, dim=-1), nan=0.0)
+    probabilities = probabilities / probabilities.sum(dim=-1, keepdim=True).clamp_min(1e-12)
+    return torch.multinomial(probabilities, 1, generator=generator).squeeze(-1)
+
+
+class RVQAttention(nn.Module):
+    def __init__(self, hidden_size, num_heads, merged_qkv, dtype, device, operations):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = hidden_size // num_heads
+        self.merged_qkv = merged_qkv
+        if merged_qkv:
+            self.qkv_proj = operations.Linear(hidden_size, hidden_size * 3, bias=False, dtype=dtype, device=device)
+        else:
+            self.q_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+            self.k_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+            self.v_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+        self.o_proj = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x):
+        batch, length, hidden_size = x.shape
+        if self.merged_qkv:
+            q, k, v = self.qkv_proj(x).chunk(3, dim=-1)
+        else:
+            q = self.q_proj(x)
+            k = self.k_proj(x)
+            v = self.v_proj(x)
+        q = q.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
+        k = k.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
+        v = v.reshape(batch, length, self.num_heads, self.head_dim).transpose(1, 2)
+        mask = torch.full((length, length), torch.finfo(q.dtype).min, device=q.device, dtype=q.dtype).triu_(1)
+        attention = optimized_attention_for_device(q.device, mask=True, small_input=True)
+        out = attention(q, k, v, self.num_heads, mask=mask, skip_reshape=True)
+        return self.o_proj(out)
+
+
+class RVQRMSNorm(nn.Module):
+    def __init__(self, hidden_size, dtype, device):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(hidden_size, dtype=dtype, device=device))
+
+    def forward(self, x):
+        return torch.nn.functional.rms_norm(x, (x.shape[-1],), comfy.ops.cast_to_input(self.weight, x), 1e-6)
+
+
+class RVQMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, merged_mlp, dtype, device, operations):
+        super().__init__()
+        self.merged_mlp = merged_mlp
+        if merged_mlp:
+            self.gate_up_proj = operations.Linear(hidden_size, intermediate_size * 2, bias=False, dtype=dtype, device=device)
+        else:
+            self.gate_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
+            self.up_proj = operations.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype, device=device)
+        self.down_proj = operations.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x):
+        if self.merged_mlp:
+            return comfy.ops.linear_input_act(self.down_proj, self.gate_up_proj(x), "swiglu")
+        return self.down_proj(torch.nn.functional.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class RVQDecoderBlock(nn.Module):
+    def __init__(self, hidden_size, num_heads, intermediate_size, merged_qkv, merged_mlp, dtype, device, operations):
+        super().__init__()
+        self.input_layernorm = RVQRMSNorm(hidden_size, dtype, device)
+        self.self_attn = RVQAttention(hidden_size, num_heads, merged_qkv, dtype, device, operations)
+        self.post_attention_layernorm = RVQRMSNorm(hidden_size, dtype, device)
+        self.mlp = RVQMLP(hidden_size, intermediate_size, merged_mlp, dtype, device, operations)
+
+    def forward(self, x):
+        x = x + self.self_attn(self.input_layernorm(x))
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+
+class RVQDepthDecoder(nn.Module):
+    def __init__(self, config, dtype, device, operations):
+        super().__init__()
+        hidden_size = int(config["hidden_size"])
+        audio_vocab_size = int(config["audio_vocab_size"])
+        merged_qkv = config.get("decoder_merged_qkv", False)
+        merged_mlp = config.get("decoder_merged_mlp", False)
+        num_codebooks = int(config["audio_num_codebooks"])
+        self.projection = operations.Linear(hidden_size, hidden_size, bias=False, dtype=dtype, device=device)
+        self.pos_embedding = operations.Embedding(16, hidden_size, dtype=dtype, device=device)
+        self.audio_heads = nn.ModuleList([
+            operations.Linear(hidden_size, audio_vocab_size, bias=False, dtype=dtype, device=device)
+            for _ in range(num_codebooks - 1)
+        ])
+        self.layers = nn.ModuleList([
+            RVQDecoderBlock(
+                hidden_size,
+                int(config["decoder_num_heads"]),
+                int(config["decoder_intermediate_size"]),
+                merged_qkv,
+                merged_mlp,
+                dtype,
+                device,
+                operations,
+            )
+            for _ in range(int(config["decoder_num_layers"]))
+        ])
+        self.norm = RVQRMSNorm(hidden_size, dtype, device)
+
+    def forward(self, sequence):
+        positions = torch.arange(sequence.shape[1], device=sequence.device)
+        x = sequence + self.pos_embedding(positions, out_dtype=sequence.dtype).unsqueeze(0)
+        for layer in self.layers:
+            x = layer(x)
+        return self.norm(x)
+
+
+class MiniMaxMusic3AR(nn.Module):
+    def __init__(self, config, dtype, device, operations):
+        super().__init__()
+        config_fields = {field.name for field in dataclasses.fields(Qwen3_8BConfig)}
+        qwen_config = Qwen3_8BConfig(**{key: value for key, value in config.items() if key in config_fields})
+        qwen_config.lm_head = False
+        qwen_config.fixed_kv = True
+        self.model = Llama2_(qwen_config, device=device, dtype=dtype, ops=operations)
+        self.model.prefetch_dynamic_vbars = True
+        self.model.graph_dynamic_vbar_blocks = True
+        self.model.lm_head = operations.Linear(qwen_config.hidden_size, qwen_config.vocab_size, bias=False, dtype=dtype, device=device)
+        self.model.lm_head_pruned = operations.Linear(qwen_config.hidden_size, C0_VOCAB_SIZE + 1, bias=False, dtype=dtype, device=device)
+        self.model.embed_tokens_prefill = operations.Embedding(AUDIO_CODE_OFFSET, qwen_config.hidden_size, dtype=dtype, device=device)
+        self.model.embed_tokens_audio = operations.Embedding(C0_VOCAB_SIZE, qwen_config.hidden_size, dtype=dtype, device=device)
+        self.model.pruned_lm_head = None
+        self.model.pruned_embedding = None
+        self.model.audio_extra_embedding = operations.Embedding(
+            int(config["audio_vocab_size"]) * (int(config["audio_num_codebooks"]) - 1),
+            qwen_config.hidden_size,
+            dtype=dtype,
+            device=device,
+        )
+        self.model.audio_decoder = RVQDepthDecoder(config, dtype, device, operations)
+        self.audio_vocab_size = int(config["audio_vocab_size"])
+        self.num_codebooks = int(config["audio_num_codebooks"])
+        self.embedding_scale = self.num_codebooks ** -0.5
+
+    def _guided_c0(self, logits, cfg_scale, top_k):
+        conditioned = logits[0:1].float()
+        unconditioned = logits[1:2].float()
+        guided = unconditioned + (conditioned - unconditioned) * cfg_scale
+        threshold = torch.topk(conditioned, top_k, dim=-1).values[..., -1, None]
+        return guided.masked_fill(conditioned < threshold, -float("inf"))
+
+    def _depth_codes(self, hidden, c0, c0_embed, generator, execution_dtype, cfg_scale, top_k):
+        decoder = self.model.audio_decoder
+        sequence = [decoder.projection(hidden).unsqueeze(1)]
+        sequence.append(decoder.projection(c0_embed).unsqueeze(1))
+        codes = [c0]
+        hidden_parts = []
+        for index in range(1, self.num_codebooks):
+            out = decoder(torch.cat(sequence, dim=1))[:, -1]
+            hidden_parts.append(out[:1].detach())
+            logits = decoder.audio_heads[index - 1](out)
+            conditioned = logits[:1].float()
+            unconditioned = logits[1:2].float()
+            code = sample_topk(unconditioned + (conditioned - unconditioned) * cfg_scale, top_k, generator).repeat(2)
+            codes.append(code)
+            if index < self.num_codebooks - 1:
+                embedding = self.model.audio_extra_embedding(
+                    code + (index - 1) * self.audio_vocab_size,
+                    out_dtype=execution_dtype,
+                )
+                sequence.append(decoder.projection(embedding).unsqueeze(1))
+        return torch.stack(codes, dim=1), torch.cat(hidden_parts, dim=-1)
+
+    def _embed_c0(self, codes, execution_dtype):
+        if self.model.pruned_embedding:
+            return self.model.embed_tokens_audio(codes, out_dtype=execution_dtype)
+        return self.model.embed_tokens(codes + AUDIO_CODE_OFFSET, out_dtype=execution_dtype)
+
+    def _embed_audio_frame(self, codes, execution_dtype):
+        c0 = self._embed_c0(codes[:, 0], execution_dtype)
+        offsets = torch.arange(self.num_codebooks - 1, device=codes.device) * self.audio_vocab_size
+        extra = self.model.audio_extra_embedding(codes[:, 1:] + offsets.unsqueeze(0), out_dtype=execution_dtype).sum(dim=1)
+        return ((c0 + extra) * self.embedding_scale).unsqueeze(1)
+
+    def _sample_c0(self, hidden, cfg_scale, top_k, generator, vocab_mask):
+        if self.model.pruned_lm_head:
+            guided = self._guided_c0(self.model.lm_head_pruned(hidden).float(), cfg_scale, top_k)
+            code = sample_topk(guided, top_k, generator)
+            stop_token = 0
+            offset = 1
+        else:
+            logits = self.model.lm_head(hidden).float()
+            stop_token = SPECIAL_TOKEN_IDS["<|audio_end|>"]
+            logits = logits.masked_fill(vocab_mask, -float("inf"))
+            guided = self._guided_c0(logits, cfg_scale, top_k).masked_fill(vocab_mask, -float("inf"))
+            code = sample_topk(guided, top_k, generator)
+            offset = AUDIO_CODE_OFFSET
+        return torch.where(code == stop_token, 0, code - offset), code, stop_token
+
+    def generate(self, input_ids, seed, max_audio_frames, device, cfg_scale=CFG_SCALE, top_k=CFG_TOP_K):
+        prompt_tokens = int(input_ids.shape[1])
+        if prompt_tokens > MAX_PROMPT_TOKENS:
+            raise ValueError(f"MiniMax Music3 prompt has {prompt_tokens} tokens; maximum is {MAX_PROMPT_TOKENS}")
+
+        input_ids = input_ids.to(device)
+        if comfy.model_management.should_use_bf16(device):
+            execution_dtype = torch.bfloat16
+        else:
+            execution_dtype = torch.float32
+        unconditioned = input_ids.clone()
+        unconditioned[:, 1:-2] = SPECIAL_TOKEN_IDS["<|audio_cfg|>"]
+        text_ids = torch.cat((input_ids, unconditioned), dim=0)
+        if self.model.pruned_embedding:
+            text_embeds = self.model.embed_tokens_prefill(text_ids, out_dtype=execution_dtype)
+        else:
+            text_embeds = self.model.embed_tokens(text_ids, out_dtype=execution_dtype)
+        decode_limit = min(int(max_audio_frames), MAX_AUDIO_FRAMES)
+        past = self.model.init_kv_cache(2, prompt_tokens + decode_limit + 1, device, execution_dtype)
+        output = self.model(None, embeds=text_embeds, past_key_values=past, dtype=execution_dtype)
+        last_hidden = output[0][:, -1]
+        past = output[2]
+
+        generator = torch.Generator(device=device).manual_seed(derive_seed(seed, "ar"))
+        decoder = self.model.audio_decoder
+        depth_io = {
+            "hidden": torch.empty_like(last_hidden),
+            "c0": torch.empty((last_hidden.shape[0],), dtype=torch.long, device=device),
+            "c0_embed": torch.empty_like(last_hidden),
+            "codes": torch.empty((last_hidden.shape[0], self.num_codebooks), dtype=torch.long, device=device),
+            "depth_hidden": torch.empty((1, last_hidden.shape[-1] * (self.num_codebooks - 1)), dtype=execution_dtype, device=device),
+        }
+        decoder._comfy_cross_step_state = depth_io
+        comfy.model_management._register_cross_step(decoder)
+        hidden_frames = []
+        pending_code = None
+        stop_token = None
+        pending_event = None
+        pending_hidden = None
+        progress = comfy.utils.ProgressBar(decode_limit)
+        cuda_device = torch.device(device).type == "cuda"
+        vocab_mask = None
+        if not self.model.pruned_lm_head:
+            vocab_mask = torch.ones(self.model.vocab_size, dtype=torch.bool, device=device)
+            vocab_mask[AUDIO_CODE_OFFSET:AUDIO_CODE_OFFSET + C0_VOCAB_SIZE] = False
+            vocab_mask[SPECIAL_TOKEN_IDS["<|audio_end|>"]] = False
+
+        for frame_index in comfy.utils.model_trange(decode_limit + 1, desc="AR sampling"):
+            comfy.model_management.throw_exception_if_processing_interrupted()
+            if pending_code is not None:
+                if pending_event is not None:
+                    pending_event.synchronize()
+                if int(pending_code.item()) == stop_token:
+                    pending_hidden = None
+                    break
+                if pending_hidden is not None:
+                    hidden_frames.append(pending_hidden)
+                    progress.update_absolute(len(hidden_frames))
+                    if len(hidden_frames) >= decode_limit:
+                        break
+
+            c0, code_or_stop, stop_token = self._sample_c0(last_hidden, cfg_scale, top_k, generator, vocab_mask)
+            if pending_code is None:
+                pending_code = torch.empty_like(code_or_stop, device="cpu", pin_memory=cuda_device)
+                if cuda_device:
+                    pending_event = torch.cuda.Event()
+            pending_code.copy_(code_or_stop, non_blocking=cuda_device)
+            if pending_event is not None:
+                pending_event.record()
+
+            c0 = c0.repeat(2)
+            c0_embed = self._embed_c0(c0, execution_dtype)
+            depth_io["hidden"].copy_(last_hidden)
+            depth_io["c0"].copy_(c0)
+            depth_io["c0_embed"].copy_(c0_embed)
+
+            def depth_core():
+                codes, depth_hidden = self._depth_codes(
+                    depth_io["hidden"], depth_io["c0"], depth_io["c0_embed"], generator, execution_dtype, cfg_scale, top_k
+                )
+                depth_io["codes"].copy_(codes)
+                depth_io["depth_hidden"].copy_(depth_hidden)
+
+            depth_queue = comfy.model_prefetch.make_prefetch_queue(
+                [[decoder, self.model.audio_extra_embedding]], device, {"prefetch_dynamic_vbars": True}
+            )
+            comfy.model_prefetch.prefetch_queue_pop(
+                depth_queue, device, decoder, execution_dtype, core=depth_core, enable_graph=True, generator=generator
+            )
+            comfy.model_prefetch.prefetch_queue_pop(depth_queue, device, None)
+            feedback_codes = depth_io["codes"]
+            depth_hidden = depth_io["depth_hidden"]
+            frame_hidden = torch.cat((last_hidden[:1].detach(), depth_hidden), dim=-1)
+            if frame_index > 0:
+                pending_hidden = frame_hidden[0].clone()
+
+            feedback = self._embed_audio_frame(feedback_codes, execution_dtype)
+            output = self.model(None, embeds=feedback, past_key_values=past, dtype=execution_dtype)
+            last_hidden = output[0][:, -1]
+            past = output[2]
+
+        if pending_hidden is not None and len(hidden_frames) < decode_limit:
+            if pending_event is not None:
+                pending_event.synchronize()
+            if int(pending_code.item()) != stop_token:
+                hidden_frames.append(pending_hidden)
+
+        if not hidden_frames:
+            raise ValueError("MiniMax Music3 generated zero audio frames")
+        return torch.stack(hidden_frames).to(device="cpu")

comfy/ldm/minimax_music/dav.py

@@ -0,0 +1,137 @@
+import math
+
+import torch
+from torch import nn
+
+import comfy.ops
+
+
+def snake(x, alpha):
+    shape = x.shape
+    flat = x.reshape(shape[0], shape[1], -1)
+    alpha = comfy.ops.cast_to_input(alpha, flat)
+    flat = flat + (alpha + 1e-9).reciprocal() * torch.sin(alpha * flat).pow(2)
+    return flat.reshape(shape)
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels, dtype, device):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.empty(1, channels, 1, dtype=dtype, device=device))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+def _weight_norm_conv(operations, *args, **kwargs):
+    return nn.utils.parametrizations.weight_norm(operations.Conv1d(*args, **kwargs))
+
+
+def _weight_norm_conv_transpose(operations, *args, **kwargs):
+    return nn.utils.parametrizations.weight_norm(operations.ConvTranspose1d(*args, **kwargs))
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, dim, dilation, dtype, device, operations):
+        super().__init__()
+        padding = 3 * dilation
+        self.block = nn.Sequential(
+            Snake1d(dim, dtype, device),
+            _weight_norm_conv(
+                operations,
+                dim,
+                dim,
+                kernel_size=7,
+                dilation=dilation,
+                padding=padding,
+                dtype=dtype,
+                device=device,
+            ),
+            Snake1d(dim, dtype, device),
+            _weight_norm_conv(operations, dim, dim, kernel_size=1, dtype=dtype, device=device),
+        )
+
+    def forward(self, x):
+        residual = self.block(x)
+        if residual.shape[-1] != x.shape[-1]:
+            padding = (x.shape[-1] - residual.shape[-1]) // 2
+            x = x[..., padding:x.shape[-1] - padding]
+        return x + residual
+
+
+class DecoderBlock(nn.Module):
+    def __init__(self, input_dim, output_dim, stride, dtype, device, operations):
+        super().__init__()
+        self.block = nn.Sequential(
+            Snake1d(input_dim, dtype, device),
+            _weight_norm_conv_transpose(
+                operations,
+                input_dim,
+                output_dim,
+                kernel_size=2 * stride,
+                stride=stride,
+                padding=math.ceil(stride / 2),
+                dtype=dtype,
+                device=device,
+            ),
+            ResidualUnit(output_dim, 1, dtype, device, operations),
+            ResidualUnit(output_dim, 3, dtype, device, operations),
+            ResidualUnit(output_dim, 9, dtype, device, operations),
+        )
+
+    def forward(self, x):
+        return self.block(x)
+
+
+class Decoder(nn.Module):
+    def __init__(self, dtype, device, operations):
+        super().__init__()
+        layers = [
+            _weight_norm_conv(
+                operations,
+                1024,
+                1536,
+                kernel_size=7,
+                padding=3,
+                dtype=dtype,
+                device=device,
+            )
+        ]
+        channels = 1536
+        output_dim = channels
+        for index, stride in enumerate((8, 8, 4, 2)):
+            input_dim = channels // (2 ** index)
+            output_dim = channels // (2 ** (index + 1))
+            layers.append(DecoderBlock(input_dim, output_dim, stride, dtype, device, operations))
+        layers.extend((
+            Snake1d(output_dim, dtype, device),
+            _weight_norm_conv(
+                operations,
+                output_dim,
+                1,
+                kernel_size=7,
+                padding=3,
+                dtype=dtype,
+                device=device,
+            ),
+            nn.Tanh(),
+        ))
+        self.model = nn.Sequential(*layers)
+
+    def forward(self, x):
+        return self.model(x)
+
+
+class MiniMaxMusic3DAV(nn.Module):
+    def __init__(self, dtype=None, device=None, operations=None):
+        super().__init__()
+        self.dec_in_proj = operations.Conv1d(64, 1024, kernel_size=1, dtype=dtype, device=device)
+        self.decoder = Decoder(dtype, device, operations)
+
+    def decode(self, latent):
+        batch, _, frames = latent.shape
+        folded = latent.reshape(batch * 2, 64, frames)
+        waveform = self.decoder(self.dec_in_proj(folded))
+        return waveform.reshape(batch, 2, -1)
+
+    forward = decode

comfy/ldm/minimax_music/dit.py

@@ -0,0 +1,213 @@
+import math
+
+import torch
+from torch import nn
+
+import comfy.model_management
+import comfy.ops
+import comfy.quant_ops
+from comfy.ldm.modules.attention import optimized_attention_for_device
+
+
+MAX_CONDITION_FRAMES = 200
+CONDITION_HOP_FRAMES = 100
+
+
+def latent_length(audio_frames):
+    return max(1, int(audio_frames * 44100 / 24000 * 960 / 512))
+
+
+class FourierFeatures(nn.Module):
+    def __init__(self, in_features, out_features, dtype, device):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty(out_features // 2, in_features, dtype=dtype, device=device))
+
+    def forward(self, value):
+        weight = comfy.ops.cast_to_input(self.weight, value)
+        features = 2.0 * math.pi * value @ weight.T
+        return torch.cat((features.cos(), features.sin()), dim=-1)
+
+
+class LayerNorm(nn.Module):
+    def __init__(self, dim, dtype, device):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.empty(dim, dtype=dtype, device=device))
+        self.register_buffer("beta", torch.empty(dim, dtype=dtype, device=device))
+
+    def forward(self, x):
+        return torch.nn.functional.layer_norm(
+            x,
+            (x.shape[-1],),
+            comfy.ops.cast_to_input(self.gamma, x),
+            comfy.ops.cast_to_input(self.beta, x),
+        )
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim, dtype, device):
+        super().__init__()
+        self.register_buffer("inv_freq", torch.empty(dim // 2, dtype=dtype, device=device))
+
+    def forward_from_seq_len(self, length, device, dtype):
+        positions = torch.arange(length, device=device, dtype=torch.float32)
+        frequencies = torch.outer(positions, comfy.ops.cast_to_input(self.inv_freq, positions))
+        frequencies = frequencies.to(dtype)
+        cos, sin = frequencies.cos(), frequencies.sin()
+        return torch.stack((cos, -sin, sin, cos), dim=-1).reshape(1, 1, length, frequencies.shape[-1], 2, 2)
+
+
+def _apply_rope(x, rotation_matrix):
+    x_dtype = x.dtype
+    x = x.reshape(*x.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).to(rotation_matrix.dtype)
+    x = rotation_matrix[..., 0] * x[..., 0] + rotation_matrix[..., 1] * x[..., 1]
+    return x.movedim(-1, -2).flatten(-2).to(x_dtype)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, dim_heads, dtype, device, operations):
+        super().__init__()
+        self.num_heads = dim // dim_heads
+        self.dim_heads = dim_heads
+        self.to_qkv = operations.Linear(dim, dim * 3, bias=False, dtype=dtype, device=device)
+        self.to_out = operations.Linear(dim, dim, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x, rotation_matrix):
+        batch, length, dim = x.shape
+        q, k, v = self.to_qkv(x).chunk(3, dim=-1)
+        q = q.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2)
+        k = k.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2)
+        v = v.reshape(batch, length, self.num_heads, self.dim_heads).transpose(1, 2)
+        rotary_dims = rotation_matrix.shape[-3] * 2
+        if comfy.model_management.in_training:
+            q = torch.cat((_apply_rope(q[..., :rotary_dims], rotation_matrix), q[..., rotary_dims:]), dim=-1)
+            k = torch.cat((_apply_rope(k[..., :rotary_dims], rotation_matrix), k[..., rotary_dims:]), dim=-1)
+        else:
+            rotated_q, rotated_k = comfy.quant_ops.ck.apply_rope_split_half(q[..., :rotary_dims], k[..., :rotary_dims], rotation_matrix)
+            q = torch.cat((rotated_q, q[..., rotary_dims:]), dim=-1)
+            k = torch.cat((rotated_k, k[..., rotary_dims:]), dim=-1)
+        attention = optimized_attention_for_device(q.device)
+        out = attention(q, k, v, self.num_heads, skip_reshape=True)
+        return self.to_out(out)
+
+
+class GLU(nn.Module):
+    def __init__(self, dim, inner_dim, dtype, device, operations):
+        super().__init__()
+        self.proj = operations.Linear(dim, inner_dim * 2, dtype=dtype, device=device)
+
+    def forward(self, x):
+        value, gate = self.proj(x).chunk(2, dim=-1)
+        return value * torch.nn.functional.silu(gate)
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, inner_dim, dtype, device, operations):
+        super().__init__()
+        self.ff = nn.Sequential(
+            GLU(dim, inner_dim, dtype, device, operations),
+            nn.Identity(),
+            operations.Linear(inner_dim, dim, dtype=dtype, device=device),
+        )
+
+    def forward(self, x):
+        return self.ff(x)
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, dim, dim_heads, inner_dim, dtype, device, operations):
+        super().__init__()
+        self.pre_norm = LayerNorm(dim, dtype, device)
+        self.self_attn = Attention(dim, dim_heads, dtype, device, operations)
+        self.ff_norm = LayerNorm(dim, dtype, device)
+        self.ff = FeedForward(dim, inner_dim, dtype, device, operations)
+
+    def forward(self, x, rotation_matrix):
+        x = x + self.self_attn(self.pre_norm(x), rotation_matrix)
+        return x + self.ff(self.ff_norm(x))
+
+
+class ContinuousTransformer(nn.Module):
+    def __init__(self, dtype, device, operations):
+        super().__init__()
+        self.project_in = operations.Linear(2304, 2048, bias=False, dtype=dtype, device=device)
+        self.project_out = operations.Linear(2048, 128, bias=False, dtype=dtype, device=device)
+        self.rotary_pos_emb = RotaryEmbedding(32, dtype, device)
+        self.layers = nn.ModuleList([
+            TransformerBlock(2048, 64, 8192, dtype, device, operations)
+            for _ in range(36)
+        ])
+
+    def forward(self, x, timestep_embedding):
+        x = self.project_in(x)
+        x = torch.cat((timestep_embedding.unsqueeze(1), x), dim=1)
+        rotation_matrix = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], x.device, x.dtype)
+        for layer in self.layers:
+            x = layer(x, rotation_matrix)
+        return self.project_out(x[:, 1:])
+
+
+class DiffusionTransformer(nn.Module):
+    def __init__(self, dtype, device, operations):
+        super().__init__()
+        self.transformer = ContinuousTransformer(dtype, device, operations)
+        self.timestep_features = FourierFeatures(1, 256, dtype, device)
+        self.to_timestep_embed = nn.Sequential(
+            operations.Linear(256, 2048, dtype=dtype, device=device),
+            nn.SiLU(),
+            operations.Linear(2048, 2048, dtype=dtype, device=device),
+        )
+        self.preprocess_conv = operations.Conv1d(2304, 2304, 1, bias=False, dtype=dtype, device=device)
+        self.postprocess_conv = operations.Conv1d(128, 128, 1, bias=False, dtype=dtype, device=device)
+
+    def forward(self, x, timestep, condition):
+        full = torch.cat((x, torch.zeros_like(x), condition), dim=1)
+        full = self.preprocess_conv(full) + full
+        timestep_features = self.timestep_features(timestep[:, None]).to(dtype=x.dtype)
+        timestep_embedding = self.to_timestep_embed(timestep_features)
+        out = self.transformer(full.transpose(1, 2), timestep_embedding).transpose(1, 2)
+        return self.postprocess_conv(out) + out
+
+
+class MiniMaxMusic3DiT(nn.Module):
+    def __init__(self, dtype=None, device=None, operations=None, **kwargs):
+        super().__init__()
+        self.dtype = dtype
+        self.latent_conditioners = nn.Sequential(
+            operations.Conv1d(4096, 2048, kernel_size=3, padding=1, dtype=dtype, device=device)
+        )
+        self.diffusion_transformer = DiffusionTransformer(dtype, device, operations)
+        self.cond_layer_logits = nn.Parameter(torch.empty(8, dtype=dtype, device=device))
+        self.cond_layer_scale = nn.Parameter(torch.empty(1, dtype=dtype, device=device))
+
+    def aligned_condition(self, hidden):
+        frames = hidden.shape[1]
+        hidden = hidden.transpose(1, 2).reshape(hidden.shape[0], 8, 4096, frames)
+        weights = torch.softmax(comfy.ops.cast_to_input(self.cond_layer_logits, hidden), dim=0)
+        hidden = torch.einsum("blht,l->bht", hidden, weights)
+        hidden = comfy.ops.cast_to_input(self.cond_layer_scale, hidden) * hidden
+        condition = self.latent_conditioners(hidden)
+        return torch.nn.functional.interpolate(condition, size=latent_length(frames), mode="nearest")
+
+    def forward(self, x, timestep, context, conditioning_scale, **kwargs):
+        condition = self.aligned_condition(context)
+        condition = condition * conditioning_scale[:, :1, :1]
+        if condition.shape[-1] < x.shape[-1]:
+            condition = torch.nn.functional.pad(condition, (0, x.shape[-1] - condition.shape[-1]))
+        else:
+            condition = condition[..., :x.shape[-1]]
+        window = latent_length(MAX_CONDITION_FRAMES)
+        if x.shape[-1] <= window:
+            return -self.diffusion_transformer(x, timestep, condition)
+
+        output = torch.zeros_like(x)
+        count = torch.zeros((1, 1, x.shape[-1]), device=x.device, dtype=x.dtype)
+        hop = latent_length(CONDITION_HOP_FRAMES)
+        start = 0
+        while start < x.shape[-1]:
+            end = min(start + window, x.shape[-1])
+            output[..., start:end] -= self.diffusion_transformer(x[..., start:end], timestep, condition[..., start:end])
+            count[..., start:end] += 1
+            if end == x.shape[-1]:
+                break
+            start += hop
+        return output / count

comfy/ldm/minimax_music/prompt.py

@@ -0,0 +1,70 @@
+import re
+
+
+SPECIAL_TOKEN_IDS = {
+    "<|im_start|>": 151644,
+    "<|im_end|>": 151645,
+    "<|audio_cfg|>": 151654,
+    "<|audio_start|>": 151669,
+    "<|audio_end|>": 151670,
+    "<|caption_start|>": 151671,
+    "<|caption_end|>": 151672,
+    "<|lyrics_start|>": 151673,
+    "<|lyrics_end|>": 151674,
+}
+AUDIO_CODE_OFFSET = 151675
+
+_SPECIAL_TAG_RE = re.compile(r"<\|([^|]*)\|>")
+_LYRIC_TAG_RE = re.compile(r"\s*(\[[^\]]+\])\s*")
+
+
+def _remove_markdown_format(text):
+    lines = []
+    for raw_line in text.splitlines():
+        line = re.sub(r"^\s{0,3}#{1,6}\s+", "", raw_line)
+        line = re.sub(r"^\s*[*+-]\s+", "", line)
+        while "**" in line:
+            updated = re.sub(r"\*\*([^*]+)\*\*", r"\1", line)
+            if updated == line:
+                break
+            line = updated
+        line = re.sub(r"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", line)
+        lines.append(line.rstrip())
+    text = "\n".join(lines)
+    text = re.sub(r"^\s*[-*_]{3,}\s*$", "", text, flags=re.MULTILINE)
+    return text.replace("• ", "").replace("    ", "")
+
+
+def clean_caption(caption):
+    def replace_special(match):
+        inner = match.group(1).strip()
+        parts = inner.split(None, 1)
+        return f"{parts[0]} is {parts[1]}" if len(parts) == 2 else inner
+
+    text = _SPECIAL_TAG_RE.sub(replace_special, caption)
+    text = _remove_markdown_format(text)
+    return re.sub(r"\n{2,}", "\n", text)
+
+
+def normalize_lyrics(lyrics):
+    parts = _LYRIC_TAG_RE.split(lyrics)
+    text = "\n".join(part.lower() if part.startswith("[") else part for part in parts if part)
+    text = text.replace(" ^ ", "\n")
+    return f"[start]\n{text}"
+
+
+def build_prompt(caption, lyrics):
+    return (
+        "<|im_start|><|caption_start|>"
+        f"{clean_caption(caption)}"
+        "<|caption_end|><|lyrics_start|>"
+        f"{normalize_lyrics(lyrics)}"
+        "<|lyrics_end|><|im_end|><|audio_start|>"
+    )
+
+
+def validate_tokenizer(tokenizer):
+    for token, expected in SPECIAL_TOKEN_IDS.items():
+        token_id = tokenizer.convert_tokens_to_ids(token)
+        if token_id != expected:
+            raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}: expected {expected}, got {token_id}")

comfy/ldm/modules/attention.py

@@ -10,6 +10,8 @@ from typing import Optional, Any, Callable, Union
 import logging
 import functools
 
+import comfy_kitchen
+
 from .diffusionmodules.util import AlphaBlender, timestep_embedding
 from .sub_quadratic_attention import efficient_dot_product_attention
 
@@ -49,6 +51,8 @@ except ImportError:
         logging.error(f"\n\nTo use the `--use-flash-attention` feature, the `flash-attn` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install flash-attn")
         exit(-1)
 
+COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE = comfy_kitchen.int8_attention_is_available()
+
 REGISTERED_ATTENTION_FUNCTIONS = {}
 def register_attention_function(name: str, func: Callable):
     # avoid replacing existing functions
@@ -145,9 +149,34 @@ def Normalize(in_channels, dtype=None, device=None):
     return torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True, dtype=dtype, device=device)
 
 
+class AttentionTensorContainer:
+    """Single-owner tensor input consumed by an optimized attention backend."""
+
+    __slots__ = ("tensor",)
+
+    def __init__(self, tensor: torch.Tensor):
+        self.tensor: torch.Tensor | None = tensor
+
+    def peek(self) -> torch.Tensor:
+        if self.tensor is None:
+            raise RuntimeError("attention tensor container has already been consumed")
+        return self.tensor
+
+    def take(self) -> torch.Tensor:
+        tensor = self.peek()
+        self.tensor = None
+        return tensor
+
+
 def wrap_attn(func):
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
+        containers = None
+        if len(args) >= 3 and isinstance(args[0], AttentionTensorContainer):
+            if not isinstance(args[1], AttentionTensorContainer) or not isinstance(args[2], AttentionTensorContainer):
+                raise TypeError("q, k, and v must all be attention tensor containers")
+            containers = args[:3]
+
         remove_attn_wrapper_key = False
         try:
             if "_inside_attn_wrapper" not in kwargs:
@@ -156,11 +185,22 @@ def wrap_attn(func):
                 kwargs["_inside_attn_wrapper"] = True
                 if transformer_options is not None:
                     if "optimized_attention_override" in transformer_options:
-                        return transformer_options["optimized_attention_override"](func, *args, **kwargs)
+                        optimized_attention_override = transformer_options["optimized_attention_override"]
+                        if containers is not None:
+                            if hasattr(optimized_attention_override, "container_function"):
+                                return optimized_attention_override.container_function(*args, **kwargs)
+                            args = tuple(container.take() for container in containers) + args[3:]
+                        return optimized_attention_override(func, *args, **kwargs)
+
+            if containers is not None:
+                if wrapper.container_function is not None:
+                    return wrapper.container_function(*args, **kwargs)
+                args = tuple(container.take() for container in containers) + args[3:]
             return func(*args, **kwargs)
         finally:
             if remove_attn_wrapper_key:
                 del kwargs["_inside_attn_wrapper"]
+    wrapper.container_function = None
     return wrapper
 
 @wrap_attn
@@ -545,6 +585,63 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha
             ).transpose(1, 2).reshape(-1, q.shape[2], heads * dim_head)
     return out
 
+def _comfy_kitchen_int8_inputs(q, k, v, heads, mask, skip_reshape, enable_gqa):
+    dim_head = q.shape[-1] if skip_reshape else q.shape[-1] // heads
+    b = q.shape[0]
+    if not skip_reshape:
+        q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa, expand_kv=False)
+        q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v))
+
+    if mask is not None:
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+        if mask.ndim == 3:
+            mask = mask.unsqueeze(1)
+
+    return q, k, v, mask, b, dim_head
+
+
+@wrap_attn
+def attention_comfy_kitchen_int8(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    q, k, v, mask, b, dim_head = _comfy_kitchen_int8_inputs(
+        q, k, v, heads, mask, skip_reshape, kwargs.get("enable_gqa", False)
+    )
+    out = comfy_kitchen.int8_attention(
+        q,
+        k,
+        v,
+        scale=kwargs.get("scale", None),
+        attn_mask=mask,
+    )
+    if not skip_output_reshape:
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+    return out
+
+
+def _attention_comfy_kitchen_int8_containers(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
+    q = q.take()
+    k = k.take()
+    v = v.take()
+    q, k, v, mask, b, dim_head = _comfy_kitchen_int8_inputs(
+        q, k, v, heads, mask, skip_reshape, kwargs.get("enable_gqa", False)
+    )
+    quantized = comfy_kitchen.prequantize_int8_attention(
+        q,
+        k,
+        v,
+        scale=kwargs.get("scale", None),
+        attn_mask=mask,
+    )
+    del q, k, v
+    out = comfy_kitchen.int8_attention_from_prequantized(quantized)
+    if not skip_output_reshape:
+        out = out.transpose(1, 2).reshape(b, -1, heads * dim_head)
+    return out
+
+
+attention_comfy_kitchen_int8.container_function = _attention_comfy_kitchen_int8_containers
+
+
 @wrap_attn
 def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs):
     if kwargs.get("low_precision_attention", True) is False or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK):
@@ -775,10 +872,20 @@ else:
         logging.info("Using sub quadratic optimization for attention, if you have memory or speed issues try using: --use-split-cross-attention")
         optimized_attention = attention_sub_quad
 
+if model_management.comfy_kitchen_attention_enabled():
+    if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE:
+        logging.info("Using Comfy Kitchen attention")
+        optimized_attention = attention_comfy_kitchen_int8
+    else:
+        logging.error("Comfy Kitchen attention is unavailable. Install a Comfy Kitchen build with attention support to use --use-ck-attention.")
+        exit(-1)
+
 optimized_attention_masked = optimized_attention
 
 
 # register core-supported attention functions
+if COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE:
+    register_attention_function("comfy_kitchen_int8", attention_comfy_kitchen_int8)
 if SAGE_ATTENTION_IS_AVAILABLE:
     register_attention_function("sage", attention_sage)
 if SAGE_ATTENTION3_IS_AVAILABLE:

comfy/ldm/wan/model.py

@@ -11,6 +11,7 @@ from comfy.ldm.flux.layers import EmbedND
 from comfy.ldm.flux.math import apply_rope1, rope
 import comfy.ldm.common_dit
 import comfy.model_management
+import comfy.ops
 import comfy.patcher_extension
 
 
@@ -146,7 +147,8 @@ class WanI2VCrossAttention(WanSelfAttention):
         v = self.v(context)
         k_img = self.norm_k_img(self.k_img(context_img))
         v_img = self.v_img(context_img)
-        img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options)
+        # Sageattn can cause Nans here, don't allow it as there is no speed difference anyway as img attention is tiny.
+        img_x = optimized_attention(q, k_img, v_img, heads=self.num_heads, transformer_options=transformer_options, low_precision_attention=False)
         # compute attention
         x = optimized_attention(q, k, v, heads=self.num_heads, transformer_options=transformer_options)
 
@@ -174,6 +176,13 @@ def repeat_e(e, x):
         return torch.repeat_interleave(e, repeats + 1, dim=1)[:, :x.size(1)]
 
 
+class WanFeedForward(nn.Sequential):
+    """[Linear, GELU(tanh), Linear], with the GELU folded into the down-projection."""
+
+    def forward(self, x):
+        return comfy.ops.linear_input_act(self[2], self[0](x), "gelu_tanh")
+
+
 class WanAttentionBlock(nn.Module):
 
     def __init__(self,
@@ -207,7 +216,7 @@ class WanAttentionBlock(nn.Module):
                                                                       qk_norm,
                                                                       eps, operation_settings=operation_settings)
         self.norm2 = operation_settings.get("operations").LayerNorm(dim, eps, elementwise_affine=False, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.ffn = nn.Sequential(
+        self.ffn = WanFeedForward(
             operation_settings.get("operations").Linear(dim, ffn_dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")), nn.GELU(approximate='tanh'),
             operation_settings.get("operations").Linear(ffn_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")))
 

comfy/ldm/wan/model_animate2.py

@@ -0,0 +1,387 @@
+# Wan-Animate-2: https://github.com/Wan-Video/Wan-Animate-2
+"""Wan2.1-I2V-14B weights, driven by a video instead of a motion extractor.
+
+A pose branch over the pose video's latents runs in lockstep with the generation
+branch, feeding it K/V per block. The reference image is one extra latent frame at the
+front of the generation branch, trimmed off by the caller. Upstream calls the pose video
+the driving video and its branch forward_ref, not to be confused with the reference image.
+"""
+
+import torch
+
+import comfy.ldm.common_dit
+import comfy.model_management
+import comfy.quant_ops
+import comfy.utils
+from comfy.ldm.flux.math import apply_rope1
+from comfy.ldm.modules.attention import optimized_attention
+
+from .model import WanAttentionBlock, WanModel, WanSelfAttention, repeat_e, sinusoidal_embedding_1d
+
+
+class WanAnimate2SelfAttention(WanSelfAttention):
+
+    def qkv(self, x, freqs):
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        q = apply_rope1(self.norm_q(self.q(x)).view(b, s, n, d), freqs)
+        k = apply_rope1(self.norm_k(self.k(x)).view(b, s, n, d), freqs)
+        return q, k, self.v(x).view(b, s, n, d)
+
+    def _attn1_patch(self, x, q, k, transformer_options):
+        for p in transformer_options.get("patches", {}).get("attn1_patch", []):
+            x = p({"x": x, "q": q, "k": k, "transformer_options": transformer_options})
+        return x
+
+    def kv(self, x, freqs):
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        return apply_rope1(self.norm_k(self.k(x)).view(b, s, n, d), freqs), self.v(x).view(b, s, n, d)
+
+    def forward_pose(self, x, freqs, transformer_options={}):
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        q, k, v = self.qkv(x, freqs)
+        out = optimized_attention(q.reshape(b, s, n * d), k.reshape(b, s, n * d), v.reshape(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options)
+        return self.o(self._attn1_patch(out, q, k, transformer_options)), k, v
+
+    def forward_gen(self, x, freqs, k_pose, v_pose, f_gen, hw, buffers, ref_strength=1.0, transformer_options={}):
+        # frame j attends every gen token plus pose frame j-1 (frame 0 is the reference slot and has none)
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+        q, k, v = self.qkv(x, freqs)
+        if ref_strength != 1.0:
+            v[:, :hw] *= ref_strength  # frame 0 is the reference image's slot
+
+        if k_pose is None:  # pose influence windowed out: plain self-attention, no per-frame loop
+            out = optimized_attention(q.reshape(b, s, n * d), k.reshape(b, s, n * d), v.reshape(b, s, n * d), heads=self.num_heads, transformer_options=transformer_options)
+            return self.o(self._attn1_patch(out, q, k, transformer_options))
+
+        # gen half is the same every frame; only the hw-token pose tail is rewritten
+        kbuf, vbuf, out = buffers
+        kbuf[:, :s] = k
+        vbuf[:, :s] = v
+
+        for j in range(f_gen):
+            q_j = q[:, j * hw:(j + 1) * hw].reshape(b, hw, n * d)
+            if j == 0:
+                kk, vv = k, v
+            else:
+                kbuf[:, s:] = k_pose[:, (j - 1) * hw:j * hw]
+                vbuf[:, s:] = v_pose[:, (j - 1) * hw:j * hw]
+                kk, vv = kbuf, vbuf
+            out[:, j * hw:(j + 1) * hw] = optimized_attention(q_j, kk.reshape(b, kk.shape[1], n * d), vv.reshape(b, kk.shape[1], n * d), heads=self.num_heads, transformer_options=transformer_options)
+        return self.o(self._attn1_patch(out, q, k, transformer_options))
+
+
+class WanAnimate2Block(WanAttentionBlock):
+
+    def __init__(self, cross_attn_type, dim, ffn_dim, num_heads, window_size=(-1, -1), qk_norm=True, cross_attn_norm=False, eps=1e-6, operation_settings={}):
+        super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
+        self.self_attn = WanAnimate2SelfAttention(dim, num_heads, window_size, qk_norm, eps, operation_settings=operation_settings)
+
+    def _modulation(self, e, x):
+        if e.ndim < 4:
+            return (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
+        return (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e).unbind(2)
+
+    def _cross_attn_ffn(self, x, e, context, context_img_len, transformer_options):
+        x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
+        for p in transformer_options.get("patches", {}).get("attn2_patch", []):
+            x = p({"x": x, "transformer_options": transformer_options})
+        y = self.ffn(torch.addcmul(repeat_e(e[3], x), self.norm2(x), 1 + repeat_e(e[4], x)))
+        return torch.addcmul(x, y, repeat_e(e[5], x))
+
+    def forward_pose(self, x, e, freqs, context, context_img_len=257, transformer_options={}):
+        e = self._modulation(e, x)
+        x = x.contiguous()
+        y, k, v = self.self_attn.forward_pose(torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)), freqs, transformer_options=transformer_options)
+        x = torch.addcmul(x, y, repeat_e(e[2], x))
+        del y
+        return self._cross_attn_ffn(x, e, context, context_img_len, transformer_options), k, v
+
+    def kv_from_input(self, x_pose, e, freqs, transformer_options={}):
+        e = self._modulation(e, x_pose)
+        x_pose = x_pose.contiguous()
+        return self.self_attn.kv(torch.addcmul(repeat_e(e[0], x_pose), self.norm1(x_pose), 1 + repeat_e(e[1], x_pose)), freqs)
+
+    def forward_gen(self, x, e, freqs, context, k_pose, v_pose, f_gen, hw, buffers, ref_strength=1.0, context_img_len=257, transformer_options={}):
+        e = self._modulation(e, x)
+        x = x.contiguous()
+        y = self.self_attn.forward_gen(torch.addcmul(repeat_e(e[0], x), self.norm1(x), 1 + repeat_e(e[1], x)), freqs, k_pose, v_pose, f_gen, hw, buffers, ref_strength=ref_strength, transformer_options=transformer_options)
+        x = torch.addcmul(x, y, repeat_e(e[2], x))
+        del y
+        return self._cross_attn_ffn(x, e, context, context_img_len, transformer_options)
+
+
+class PoseBranchCache:
+    """Pose-branch block inputs, reused across the sampling steps of one execution.
+
+    Caching the block input rather than its K/V halves the memory; reprojecting K/V on read
+    costs ~4% of re-running the block. One slot per distinct pose sequence, so under
+    context windows each window keeps its own; least recently used slots are evicted when
+    the store device runs low on memory. Created and freed by WanAnimate2Cache.
+    """
+
+    CONVROT_GROUPSIZE = 256
+
+    def __init__(self, store_device=None, dtype="default"):
+        self.store_device = torch.device(store_device) if store_device is not None else torch.device("cpu")
+        self.dtype = dtype
+        self.slots = []  # most recently used last
+        self.slot = None
+        self._pending = {}
+        self._staging = {}
+
+    def select(self, pose_latents):
+        # select runs at a forward boundary: an interrupted forward can leave copies in flight that a different slot's forward would then mistake for its own
+        if self._pending:
+            for t, stream in self._pending.values():
+                if stream is not None:
+                    stream.synchronize()
+            self._pending = {}
+        # keyed on batch element 0, so a cond batch size change mid-run stays valid
+        k = pose_latents[:1]
+        for s in self.slots:
+            if s["key"].shape == k.shape and torch.equal(s["key"], k.to(s["key"].device)):
+                self.slots.remove(s)
+                self.slots.append(s)
+                self.slot = s
+                return
+        # cache what fits: a filled slot is the size estimate for the next one, and least recently used slots make room when the store device runs low
+        est = max((self._slot_bytes(s) for s in self.slots), default=0) * 1.5
+        while self.slots and comfy.model_management.get_free_memory(self.store_device) < est:
+            self._free_slot(self.slots.pop(0))
+        self.slot = {"key": k.clone().to(self.store_device), "blocks": {}, "params": {}, "shape": None, "pinned": []}
+        self.slots.append(self.slot)
+
+    def _free_slot(self, s):
+        for t, stream in self._pending.values():
+            if stream is not None:
+                stream.synchronize()  # an aborted forward can leave a copy in flight, still reading memory we are about to unpin
+        self._pending = {}
+        for t in s["pinned"]:
+            comfy.model_management.unpin_memory(t)
+
+    def free(self):
+        for s in self.slots:
+            self._free_slot(s)
+        self.slots = []
+        self.slot = None
+        self._staging = {}
+
+    def filled(self, num_blocks):
+        return self.slot is not None and len(self.slot["blocks"]) == num_blocks
+
+    def put(self, i, x_pose):
+        t = x_pose[:1]
+        params = None
+        if self.dtype in ("int8", "int4"):
+            # convrot is what lets low-bit survive the ~125x per-channel outliers here, and over a [tokens, dim] view per-row scale means per-token. The kernels want 2D and a power-of-4 group that divides dim.
+            self.slot["shape"] = t.shape
+            g = self.CONVROT_GROUPSIZE
+            while g > 4 and t.shape[-1] % g:
+                g //= 4
+            if self.dtype == "int4":
+                t, params = comfy.quant_ops.TensorCoreConvRotW4A4Layout.quantize(t.reshape(-1, t.shape[-1]), convrot_groupsize=g)
+            else:
+                t, params = comfy.quant_ops.TensorWiseINT8Layout.quantize(t.reshape(-1, t.shape[-1]), is_weight=True, per_channel=True, convrot=True, convrot_groupsize=g)
+
+        t = t.to(self.store_device, copy=True)
+        if comfy.model_management.pin_memory(t):
+            self.slot["pinned"].append(t)
+        self.slot["blocks"][i] = t
+        # the scales follow the blocks off the GPU: per-window slots would otherwise pile them up in VRAM (~200 MB per window at 480p int4)
+        self.slot["params"][i] = params if params is None else params.to_device(self.store_device)
+
+    def prefetch(self, i, device, dtype):
+        # call before the compute this should overlap, so the stream waits only on work already enqueued
+        if i not in self.slot["blocks"] or i in self._pending:
+            return
+        t = self.slot["blocks"][i]
+        cast_dtype = None if self.slot["params"][i] is not None else dtype  # int8 entries move in their stored dtype and widen in take()
+        stream = None
+        r = None
+        if t.device != device:
+            stream = comfy.model_management.get_offload_stream(device)
+            cs = comfy.model_management.current_stream(device)
+            if stream is not None and cs is not None:
+                # the handed-out stream last waited on the main stream a full rotation ago, which does not cover the previous consumer's reads of this slot; wait now so the copy cannot overwrite a slot still being read
+                stream.wait_stream(cs)
+            # two persistent staging buffers per tensor shape instead of a fresh allocation per block (~29 GB of churn per pass at 720p); windows of different lengths get their own pair
+            buf_key = (tuple(t.shape), cast_dtype if cast_dtype is not None else t.dtype)
+            if buf_key not in self._staging:
+                self._staging[buf_key] = [torch.empty(t.shape, dtype=buf_key[1], device=device) for _ in range(2)]
+            r = self._staging[buf_key][i % 2]
+        self._pending[i] = (comfy.model_management.cast_to(t, cast_dtype, device, non_blocking=True, stream=stream, r=r), stream)
+
+    def take(self, i, device, dtype, batch_size):
+        if i not in self._pending:
+            self.prefetch(i, device, dtype)
+        t, stream = self._pending.pop(i)
+        comfy.model_management.sync_stream(device, stream)
+        params = self.slot["params"][i]
+        if params is not None:
+            layout = comfy.quant_ops.TensorCoreConvRotW4A4Layout if self.dtype == "int4" else comfy.quant_ops.TensorWiseINT8Layout
+            t = layout.dequantize(t, params.to_device(t.device)).reshape(self.slot["shape"]).to(dtype)
+        return comfy.utils.repeat_to_batch_size(t, batch_size)
+
+    def _slot_bytes(self, s):
+        return sum(t.numel() * t.element_size() for t in s["blocks"].values())
+
+    def memory_bytes(self):
+        return sum(self._slot_bytes(s) for s in self.slots)
+
+
+class WanAnimate2Model(WanModel):
+
+    def __init__(self,
+                 model_type='animate2',
+                 patch_size=(1, 2, 2),
+                 text_len=512,
+                 in_dim=36,
+                 dim=5120,
+                 ffn_dim=13824,
+                 freq_dim=256,
+                 text_dim=4096,
+                 out_dim=16,
+                 num_heads=40,
+                 num_layers=40,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6,
+                 flf_pos_embed_token_number=None,
+                 in_dim_ref_conv=None,
+                 image_model=None,
+                 device=None, dtype=None, operations=None,
+                 ):
+        # model_type is 'animate2' in unet_config, but the checkpoint is i2v-shaped
+        super().__init__(model_type='i2v', patch_size=patch_size, text_len=text_len, in_dim=in_dim, dim=dim, ffn_dim=ffn_dim, freq_dim=freq_dim,
+                         text_dim=text_dim, out_dim=out_dim, num_heads=num_heads, num_layers=num_layers, window_size=window_size, qk_norm=qk_norm,
+                         cross_attn_norm=cross_attn_norm, eps=eps, flf_pos_embed_token_number=flf_pos_embed_token_number, in_dim_ref_conv=in_dim_ref_conv,
+                         wan_attn_block_class=WanAnimate2Block, image_model=image_model, device=device, dtype=dtype, operations=operations)
+
+    def rope_encode_pose(self, t, h, w, w_patches, device=None, dtype=None):
+        # t_start=1 lines pose frame j up with gen frame j+1, past the reference slot; shift_x parks it in its own strip of rope space.
+        # The caller's rope_options are a user scaling knob and deliberately not forwarded.
+        return super().rope_encode(t, h, w, t_start=1, device=device, dtype=dtype, transformer_options={"rope_options": {"shift_x": float(w_patches)}})
+
+    def _forward(self, x, timestep, context, clip_fea=None, time_dim_concat=None, transformer_options={}, pose_latents=None, clip_fea_pose=None, context_pose=None, **kwargs):
+        bs, c, t, h, w = x.shape
+        x = comfy.ldm.common_dit.pad_to_patch_size(x, self.patch_size)
+
+        # h/w pre-pad: rope_encode's rounding reproduces the post-pad grid
+        freqs = self.rope_encode(t, h, w, device=x.device, dtype=x.dtype, transformer_options=transformer_options)
+        freqs_pose = None
+        if pose_latents is not None:  # absent when the node's timestep window excludes this step
+            pose_latents = comfy.ldm.common_dit.pad_to_patch_size(pose_latents.to(x.dtype), self.patch_size)
+            w_patches = (w + (self.patch_size[2] // 2)) // self.patch_size[2]
+            freqs_pose = self.rope_encode_pose(pose_latents.shape[2], h, w, w_patches, device=x.device, dtype=x.dtype)
+
+        return self.forward_orig(x, timestep, context, clip_fea=clip_fea, freqs=freqs, freqs_pose=freqs_pose, pose_latents=pose_latents,
+                                 clip_fea_pose=clip_fea_pose, context_pose=context_pose, transformer_options=transformer_options, **kwargs)[:, :, :t, :h, :w]
+
+    def forward_orig(self, x, t, context, clip_fea=None, freqs=None, freqs_pose=None, pose_latents=None, clip_fea_pose=None, context_pose=None, pose_strength=1.0, reference_strength=1.0, transformer_options={}, **kwargs):
+        x_input = x[:, :, 1:]  # video-only: frame 0 is the reference slot, offset past it below
+        x = self.patch_embedding(x.float()).to(x.dtype)
+        grid_sizes = x.shape[2:]
+        transformer_options["grid_sizes"] = grid_sizes
+        f_gen, gh, gw = grid_sizes
+        hw = gh * gw
+        x = x.flatten(2).transpose(1, 2)
+
+        # the node windows the pose influence via cond timestep ranges: outside the window the cond carries no pose latents, and the branch, its cache traffic and the per-frame attention loop are all skipped
+        apply_pose = pose_latents is not None
+        if apply_pose and pose_latents.shape[2] != f_gen - 1:  # before cache.select, which would otherwise keep an empty slot keyed to the rejected latents
+            raise ValueError("pose branch has {} latent frames, expected {} (generation frames minus the reference-image slot)".format(pose_latents.shape[2], f_gen - 1))
+
+        cache = transformer_options.get("animate2_cache", None) if apply_pose else None
+        if cache is not None:
+            cache.select(pose_latents)
+        cached = cache is not None and cache.filled(len(self.blocks))
+
+        x_pose = None
+        if not cached and apply_pose:
+            # 36ch = [latents(16) | mask(4) | latents(16)]; latents twice, and the mask is all ones since every pose frame is known
+            x_pose = self.patch_embedding(torch.cat([pose_latents, torch.ones_like(pose_latents[:, :4]), pose_latents], dim=1).float()).to(x.dtype)
+            x_pose = x_pose.flatten(2).transpose(1, 2)
+
+        # time embeddings
+        e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x.dtype))
+        e = e.reshape(t.shape[0], -1, e.shape[-1])
+        e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+
+        e0_pose = None
+        if apply_pose:
+            t_pose = torch.ones_like(t.flatten())
+            e_pose = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t_pose).to(dtype=x.dtype))
+            e_pose = e_pose.reshape(t.shape[0], -1, e_pose.shape[-1])
+            e0_pose = self.time_projection(e_pose).unflatten(2, (6, self.dim))
+
+        context_gen = self.text_embedding(context)
+
+        context_img_len = None
+        if clip_fea is not None:
+            if self.img_emb is not None:
+                context_gen = torch.cat([self.img_emb(clip_fea), context_gen], dim=1)
+            context_img_len = clip_fea.shape[-2]
+
+        context_img_len_pose = None
+        if not cached and apply_pose:
+            context_pose = self.text_embedding(context if context_pose is None else context_pose)
+            clip_fea_pose = clip_fea if clip_fea_pose is None else clip_fea_pose
+            if clip_fea_pose is not None:
+                if self.img_emb is not None:
+                    context_pose = torch.cat([self.img_emb(clip_fea_pose), context_pose], dim=1)
+                context_img_len_pose = clip_fea_pose.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"
+
+        if cache is not None and not cached and apply_pose and "context_window" in transformer_options:
+            # pose-only prepass, to avoid inflating dynamic VRAM calibration when using multiple context windows
+            for i, block in enumerate(self.blocks):
+                transformer_options["block_index"] = i
+                cache.put(i, x_pose)
+                x_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)[0]
+            x_pose = None
+            cached = True
+
+        buffers = None
+        if apply_pose:
+            # allocated once and reused by every block
+            n, d = self.num_heads, self.dim // self.num_heads
+            buffers = (x.new_empty(x.shape[0], x.shape[1] + hw, n, d), x.new_empty(x.shape[0], x.shape[1] + hw, n, d), x.new_empty(x.shape[0], x.shape[1], self.dim))
+
+        for i, block in enumerate(self.blocks):
+            transformer_options["block_index"] = i
+
+            if not apply_pose:
+                k_pose = v_pose = None
+            elif cached:
+                x_pose_in = cache.take(i, x.device, x.dtype, x.shape[0])
+                cache.prefetch(i + 1, x.device, x.dtype)  # queue the next block before the gen compute it should overlap
+                k_pose, v_pose = block.kv_from_input(x_pose_in, e0_pose, freqs_pose, transformer_options=transformer_options)
+                del x_pose_in
+            else:
+                if cache is not None:
+                    cache.put(i, x_pose)
+                # runs even under a block replace: its state has to reach block i+1
+                x_pose, k_pose, v_pose = block.forward_pose(x_pose, e0_pose, freqs_pose, context_pose, context_img_len=context_img_len_pose, transformer_options=transformer_options)
+            if v_pose is not None and pose_strength != 1.0:
+                v_pose = v_pose * pose_strength
+
+            if ("double_block", i) in blocks_replace:
+                def block_wrap(args, block=block, k_pose=k_pose, v_pose=v_pose):
+                    return {"img": block.forward_gen(args["img"], args["vec"], args["pe"], args["txt"], k_pose, v_pose, f_gen, hw, buffers, ref_strength=reference_strength, context_img_len=context_img_len, transformer_options=args["transformer_options"])}
+                out = blocks_replace[("double_block", i)]({"img": x, "txt": context_gen, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
+                x = out["img"]
+            else:
+                x = block.forward_gen(x, e0, freqs, context_gen, k_pose, v_pose, f_gen, hw, buffers, ref_strength=reference_strength, 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": hw, "transformer_options": transformer_options})
+                    x = out["img"]
+
+        return self.unpatchify(self.head(x, e), grid_sizes)

comfy/ldm/wan/uni3c.py

@@ -4,7 +4,7 @@ import torch
 import torch.nn as nn
 
 from comfy.ldm.flux.layers import EmbedND
-from .model import WanSelfAttention
+from .model import WanFeedForward, WanSelfAttention
 
 
 class Uni3CLayerNormZero(nn.Module):
@@ -41,7 +41,7 @@ class Uni3CAttentionBlock(nn.Module):
         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(
+        self.ffn = WanFeedForward(
             operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
             operations.Linear(ffn_dim, dim, device=device, dtype=dtype))
 

comfy/model_base.py

@@ -22,6 +22,8 @@ import torch
 import logging
 import comfy.ldm.lightricks.av_model
 import comfy.ldm.minimax.model
+import comfy.ldm.minimax_music.dit
+import comfy.nested_tensor
 import comfy.ldm.lightricks.symmetric_patchifier
 import comfy.context_windows
 from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
@@ -45,6 +47,7 @@ import comfy.ldm.cosmos.predict2
 import comfy.ldm.lumina.model
 import comfy.ldm.wan.model
 import comfy.ldm.wan.model_animate
+import comfy.ldm.wan.model_animate2
 import comfy.ldm.wan.ar_model
 import comfy.ldm.wan.model_wandancer
 import comfy.ldm.hunyuan3d.model
@@ -100,6 +103,7 @@ class ModelType(Enum):
     FLOW_COSMOS = 10
     IMG_TO_IMG_FLOW = 11
     V_PREDICTION_DDPM = 12
+    FLOW_AV = 13
 
 
 def model_sampling(model_config, model_type):
@@ -136,6 +140,9 @@ def model_sampling(model_config, model_type):
         c = comfy.model_sampling.IMG_TO_IMG_FLOW
     elif model_type == ModelType.V_PREDICTION_DDPM:
         c = comfy.model_sampling.V_PREDICTION_DDPM
+    elif model_type == ModelType.FLOW_AV:
+        c = comfy.model_sampling.CONST
+        s = comfy.model_sampling.ModelSamplingAV
 
     class ModelSampling(s, c):
         pass
@@ -180,6 +187,7 @@ class BaseModel(torch.nn.Module):
 
         self.model_type = model_type
         self.model_sampling = model_sampling(model_config, model_type)
+        self.latent_shapes = None  # set by the sampler for models that pack several streams into one latent
 
         self.adm_channels = unet_config.get("adm_in_channels", None)
         if self.adm_channels is None:
@@ -1146,6 +1154,10 @@ class LTXV(BaseModel):
         if guide_attention_entries is not None:
             out['guide_attention_entries'] = comfy.conds.CONDConstant(guide_attention_entries)
 
+        generated_keyframes = kwargs.get("generated_keyframes", None)
+        if generated_keyframes is not None:
+            out['generated_keyframes'] = comfy.conds.CONDConstant(generated_keyframes)
+
         return out
 
     def process_timestep(self, timestep, x, denoise_mask=None, **kwargs):
@@ -1206,6 +1218,10 @@ class LTXAV(BaseModel):
         if ref_audio is not None:
             out['ref_audio'] = comfy.conds.CONDConstant(ref_audio)
 
+        generated_keyframes = kwargs.get("generated_keyframes", None)
+        if generated_keyframes is not None:
+            out['generated_keyframes'] = comfy.conds.CONDConstant(generated_keyframes)
+
         return out
 
     def process_timestep(self, timestep, x, denoise_mask=None, audio_denoise_mask=None, **kwargs):
@@ -1807,6 +1823,41 @@ class WAN22_Animate(WAN21):
             return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1)
         return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
 
+class WAN_Animate2(WAN21):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model_animate2.WanAnimate2Model)
+        self.image_to_video = True
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+
+        pose_video_latent = kwargs.get("pose_video_latent", None)
+        if pose_video_latent is not None:
+            out['pose_latents'] = comfy.conds.CONDRegular(self.process_latent_in(pose_video_latent))
+
+        clip_vision_output_pose = kwargs.get("clip_vision_output_pose", None)
+        if clip_vision_output_pose is not None:
+            out['clip_fea_pose'] = comfy.conds.CONDRegular(clip_vision_output_pose.penultimate_hidden_states)
+
+        cross_attn_pose = kwargs.get("cross_attn_pose", None)
+        if cross_attn_pose is not None:
+            out['context_pose'] = comfy.conds.CONDRegular(cross_attn_pose)
+
+        pose_strength = kwargs.get("pose_strength", 1.0)
+        if pose_strength != 1.0:
+            out['pose_strength'] = comfy.conds.CONDConstant(pose_strength)
+
+        reference_strength = kwargs.get("reference_strength", 1.0)
+        if reference_strength != 1.0:
+            out['reference_strength'] = comfy.conds.CONDConstant(reference_strength)
+
+        return out
+
+    def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]):
+        if cond_key == "pose_latents":
+            return comfy.context_windows.slice_cond(cond_value, window, x_in, device, temporal_dim=2, temporal_offset=1)
+        return super().resize_cond_for_context_window(cond_key, cond_value, window, x_in, device, retain_index_list=retain_index_list)
+
 class WAN22_S2V(WAN21):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.WanModel_S2V)
@@ -2065,9 +2116,33 @@ class Hunyuan3Dv2_1(BaseModel):
         return out
 
 class MiniMaxH3(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+    def __init__(self, model_config, model_type=ModelType.FLOW_AV, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax.model.MiniMaxH3Model)
 
+    def audio_scale(self):
+        """Scale the sampler carries the audio stream at, 1.0 when not sampling the packed latent."""
+        if self.latent_shapes is None or len(self.latent_shapes) < 2:
+            return 1.0
+        return self.model_sampling.audio_scale
+
+    def _scale_audio_slice(self, latent, scale):
+        # the sampler carries the audio stream scaled onto the video schedule
+        if scale == 1.0:
+            return latent
+        if latent.is_nested:  # the x0 output hands back the unpacked view
+            streams = latent.unbind()
+            return comfy.nested_tensor.NestedTensor([streams[0], streams[1] * scale] + list(streams[2:]))
+        n = math.prod(self.latent_shapes[0][1:])
+        latent = latent.clone()
+        latent[..., n:] *= scale
+        return latent
+
+    def process_latent_in(self, latent):
+        return self._scale_audio_slice(super().process_latent_in(latent), self.audio_scale())
+
+    def process_latent_out(self, latent):
+        return super().process_latent_out(self._scale_audio_slice(latent, 1.0 / self.audio_scale()))
+
     def extra_conds(self, **kwargs):
         out = super().extra_conds(**kwargs)
         cross_attn = kwargs.get("cross_attn", None)
@@ -2090,31 +2165,30 @@ class MiniMaxH3(BaseModel):
         keyframes = kwargs.get("minimax_keyframes", None)
         if keyframes is not None:
             payload["keyframes"] = keyframes
-            payload["frame_count"] = kwargs.get("minimax_frame_count", None)
-            payload["cond_video_latents"] = [kf["latent"] for kf in keyframes]
+            payload["cond_video_latents"] = [kf["latent"] for kf in keyframes if kf.get("latent") is not None]
+            payload["cond_audio_latents"] = [kf["audio_latent"] for kf in keyframes if kf.get("audio_latent") is not None]
         refs = kwargs.get("minimax_refs", None)
         if refs is not None:
             payload["refs"] = refs
-            payload["cond_video_latents"] = [r["latent"] for r in refs if "latent" in r]
-            payload["cond_audio_latents"] = [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
+            payload["cond_video_latents"] = payload.get("cond_video_latents", []) + [r["latent"] for r in refs if "latent" in r]
+            payload["cond_audio_latents"] = payload.get("cond_audio_latents", []) + [r["audio_latent"] for r in refs if r.get("audio_latent") is not None]
         if kwargs.get("minimax_visual_cond_noise_aug", None) is not None:
             payload["visual_cond_noise_aug"] = kwargs["minimax_visual_cond_noise_aug"]
         if kwargs.get("minimax_audio_cond_noise_aug", None) is not None:
             payload["audio_cond_noise_aug"] = kwargs["minimax_audio_cond_noise_aug"]
         payload["seed"] = kwargs.get("seed", 0)
+        # same value process_latent_in/out used, so the model never undoes a scale that was not applied
+        payload["audio_scale"] = self.audio_scale()
         if cross_attn is not None and latent_shapes is not None and len(latent_shapes) > 1:
             # packed layout built once per sampling run, h/w rounded up to the DiT's 2x2 patch
             vs = latent_shapes[0]
             payload["layout"] = comfy.ldm.minimax.model.PackedLayout(
                 cross_attn.shape[1], vs[2], (vs[3] + 1) // 2 * 2, (vs[4] + 1) // 2 * 2,
                 latent_shapes[1][-1], keyframes=payload.get("keyframes"),
-                refs=payload.get("refs"), frame_count=payload.get("frame_count"))
+                refs=payload.get("refs"))
         out['minimax_payload'] = comfy.conds.CONDConstant(payload)
         return out
 
-    def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
-        return latent_image
-
 class TripoSplat(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.triposplat.model.LatentSeqMMFlowModel)
@@ -2264,6 +2338,18 @@ class ACEStep15(BaseModel):
         out['refer_audio'] = comfy.conds.CONDRegular(refer_audio)
         return out
 
+class MiniMaxMusic3(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.minimax_music.dit.MiniMaxMusic3DiT)
+
+    def process_timestep(self, timestep, **kwargs):
+        return 1.0 - timestep
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        out["conditioning_scale"] = comfy.conds.CONDRegular(kwargs["conditioning_scale"])
+        return out
+
 class Omnigen2(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.omnigen.omnigen2.OmniGen2Transformer2DModel)

comfy/model_detection.py

@@ -44,6 +44,13 @@ def calculate_transformer_depth(prefix, state_dict_keys, state_dict):
 def detect_unet_config(state_dict, key_prefix, metadata=None):
     state_dict_keys = list(state_dict.keys())
 
+    if (
+        '{}cond_layer_logits'.format(key_prefix) in state_dict_keys
+        and '{}latent_conditioners.0.weight'.format(key_prefix) in state_dict_keys
+        and '{}diffusion_transformer.transformer.layers.0.self_attn.to_qkv.weight'.format(key_prefix) in state_dict_keys
+    ):
+        return {"audio_model": "minimax_music3"}
+
     if '{}joint_blocks.0.context_block.attn.qkv.weight'.format(key_prefix) in state_dict_keys: #mmdit model
         unet_config = {}
         unet_config["in_channels"] = state_dict['{}x_embedder.proj.weight'.format(key_prefix)].shape[1]
@@ -397,6 +404,7 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["cross_attention_dim"] = shape[1]
         if metadata is not None and "config" in metadata:
             dit_config.update(json.loads(metadata["config"]).get("transformer", {}))
+        dit_config["use_keyframes_abs_pos_embedding"] = '{}keyframes_abs_pos_embedding'.format(key_prefix) in state_dict_keys
         return dit_config
 
     if '{}genre_embedder.weight'.format(key_prefix) in state_dict_keys: #ACE-Step model
@@ -829,11 +837,10 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
 
         dit_config["use_adaln_lora"] = True
         dit_config["adaln_lora_dim"] = 256
+        dit_config["num_blocks"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
         if dit_config["model_channels"] == 2048:
-            dit_config["num_blocks"] = 28
             dit_config["num_heads"] = 16
         elif dit_config["model_channels"] == 5120:
-            dit_config["num_blocks"] = 36
             dit_config["num_heads"] = 40
 
         if dit_config["in_channels"] == 16:

comfy/model_management.py

@@ -490,28 +490,36 @@ try:
         except:
             rocm_version = (6, -1)
 
-        def aotriton_supported(gpu_arch):
-            path = torch.__path__[0]
-            path = os.path.join(os.path.join(path, "lib"), "aotriton.images")
-            gfx = set(map(lambda a: a[4:], filter(lambda a: a.startswith("amd-gfx"), os.listdir(path))))
-            if gpu_arch in gfx:
-                return True
-            if "{}x".format(gpu_arch[:-1]) in gfx:
-                return True
-            if "{}xx".format(gpu_arch[:-2]) in gfx:
-                return True
-            return False
+        def aotriton_supported():
+            """Whether pytorch reports flash attention as usable on this gpu.
+
+            can_use_flash_attention() evaluates runtime eligibility for the given
+            parameters; on a ROCm build that includes checking the gpu arch against the
+            kernel images AOTriton was compiled for. Querying it avoids assuming where
+            those images live inside the torch install. The probe tensor is shaped and
+            typed to pass the unrelated SDPA checks, so False means no hardware support
+            rather than a rejected shape.
+            """
+            try:
+                if not torch.backends.cuda.is_flash_attention_available():  # not built with flash attention
+                    return False
+                q = torch.empty((1, 1, 8, 64), dtype=torch.float16, device=get_torch_device())
+                params = torch.backends.cuda.SDPAParams(q, q, q, None, 0.0, False, False)
+                return torch.backends.cuda.can_use_flash_attention(params, False)
+            except (AttributeError, RuntimeError, TypeError) as e:
+                logging.warning("Could not query aotriton support: {}".format(e))
+                return False
 
         logging.info("AMD arch: {}".format(arch))
         logging.info("ROCm version: {}".format(rocm_version))
         if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
-            if aotriton_supported(arch):  # AMD efficient attention implementation depends on aotriton.
+            if aotriton_supported():  # AMD efficient attention implementation depends on aotriton.
                 if torch_version_numeric >= (2, 7):  # works on 2.6 but doesn't actually seem to improve much
                     if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151"]):  # TODO: more arches, TODO: gfx950
                         ENABLE_PYTORCH_ATTENTION = True
                 if rocm_version >= (7, 0):
-                   if any((a in arch) for a in ["gfx1200", "gfx1201"]):
-                       ENABLE_PYTORCH_ATTENTION = True
+                    if any((a in arch) for a in ["gfx1200", "gfx1201"]):
+                        ENABLE_PYTORCH_ATTENTION = True
         if torch_version_numeric >= (2, 7) and rocm_version >= (6, 4):
             if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950"]):  # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
                 SUPPORT_FP8_OPS = True
@@ -1360,9 +1368,14 @@ STREAM_CAST_BUFFERS = {}
 LARGEST_CASTED_WEIGHT = (None, 0)
 STREAM_AIMDO_CAST_BUFFERS = {}
 LARGEST_AIMDO_CASTED_WEIGHT = (None, 0)
+CROSS_STEP_STATE = weakref.WeakSet()
 
 DEFAULT_AIMDO_CAST_BUFFER_RESERVATION_SIZE = 16 * 1024 ** 3
 
+# NOTE: devs/agents: this is temporary and will be removed in a future comfy. Not supported for custom node use.
+def _register_cross_step(module):
+    CROSS_STEP_STATE.add(module)
+
 def get_cast_buffer(offload_stream, device, size, ref):
     global LARGEST_CASTED_WEIGHT
 
@@ -1417,6 +1430,10 @@ def reset_cast_buffers():
         mmap_obj.bounce()
     DIRTY_MMAPS.clear()
 
+    for module in CROSS_STEP_STATE:
+        del module._comfy_cross_step_state
+    CROSS_STEP_STATE.clear()
+
     for loaded_model in current_loaded_models:
         model = loaded_model.model
         if model is not None and model.is_dynamic():
@@ -1543,13 +1560,31 @@ def cast_to_device(tensor, device, dtype, copy=False):
 PINNED_MEMORY = {}
 TOTAL_PINNED_MEMORY = 0
 MAX_PINNED_MEMORY = -1
+
+def get_disk_swap_total():
+    if not os.path.exists("/proc/swaps"):
+        return 0
+
+    total = 0
+    try:
+        with open("/proc/swaps", encoding="utf-8") as swaps:
+            next(swaps, None)
+            for line in swaps:
+                filename, _, size, _, _ = line.rsplit(maxsplit=4)
+                if os.path.basename(os.path.realpath(filename)).startswith("zram"):
+                    continue
+                total += int(size) * 1024
+    except:
+        logging.warning("Could not get amount of swap memory on system.")
+    return total
+
 if not args.disable_pinned_memory:
     if is_nvidia() or is_amd():
         ram = get_total_memory(torch.device("cpu"))
         if WINDOWS:
             MAX_PINNED_MEMORY = ram * 0.40  # Windows limit is apparently 50%
         else:
-            MAX_PINNED_MEMORY = ram * 0.90
+            MAX_PINNED_MEMORY = max(ram * 0.40, min(ram * 0.90, ram - 4 * 1024 ** 3, ram + get_disk_swap_total() - 16 * 1024 ** 3))
         logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
 
 PINNING_ALLOWED_TYPES = set(["Tensor", "Parameter", "QuantizedTensor"])
@@ -1640,6 +1675,9 @@ def unpin_memory(tensor):
 def sage_attention_enabled():
     return args.use_sage_attention
 
+def comfy_kitchen_attention_enabled():
+    return args.use_ck_attention
+
 def flash_attention_enabled():
     return args.use_flash_attention
 

comfy/model_patcher.py

@@ -685,6 +685,14 @@ class ModelPatcher:
     def set_model_attn2_output_patch(self, patch):
         self.set_model_patch(patch, "attn2_output_patch")
 
+    def set_model_optimized_attention(self, optimized_attention):
+        def optimized_attention_override(_, *args, **kwargs):
+            return optimized_attention(*args, **kwargs)
+
+        if hasattr(optimized_attention, "container_function") and optimized_attention.container_function is not None:
+            optimized_attention_override.container_function = optimized_attention.container_function
+        self.model_options["transformer_options"]["optimized_attention_override"] = optimized_attention_override
+
     def set_model_input_block_patch(self, patch):
         self.set_model_patch(patch, "input_block_patch")
 
@@ -1879,8 +1887,29 @@ class ModelPatcherDynamic(ModelPatcher):
             loading = self._load_list(for_dynamic=True, default_device=device_to)
             loading.sort()
 
+            get_units = getattr(self.model, "get_dynamic_vram__units", None)
+            dynamic_units, last_dynamic_units = get_units() if get_units is not None else ([], [])
+            dynamic_units = list(dynamic_units)
+            last_dynamic_units = list(last_dynamic_units)
+            loading_by_module = {entry[-2]: entry for entry in loading}
+            loading = []
+            for unit in dynamic_units:
+                unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,)
+                modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module]
+                for index, module in enumerate(modules):
+                    loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None))
+            last_loading = []
+            for unit in last_dynamic_units:
+                unit_modules = unit if isinstance(unit, (list, tuple)) else (unit,)
+                modules = [module for root in unit_modules for module in root.modules() if module in loading_by_module]
+                for index, module in enumerate(modules):
+                    last_loading.append((*loading_by_module.pop(module), unit if index == len(modules) - 1 else None))
+            loading.extend((*entry, None) for entry in loading_by_module.values())
+            loading.extend(last_loading)
+            v_block = None
+
             for x in loading:
-                *_, module_mem, n, m, params = x
+                *_, module_mem, n, m, params, end_of_block = x
 
                 def set_dirty(item, dirty):
                     if dirty or not hasattr(item, "_v_signature"):
@@ -1973,6 +2002,13 @@ class ModelPatcherDynamic(ModelPatcher):
 
                 move_weight_functions(m, device_to)
 
+                if hasattr(m, "_v"):
+                    v_block = m._v if v_block is None else (v_block[0], v_block[1], max(v_block[2], m._v[1] + m._v[2] - v_block[1]))
+                if end_of_block is not None:
+                    unit = end_of_block
+                    (unit[0] if isinstance(unit, (list, tuple)) else unit)._v_block = v_block
+                    v_block = None
+
             for key, buf in self.model.named_buffers(recurse=True):
                 if key not in self.backup_buffers:
                     self.backup_buffers[key] = buf

comfy/model_prefetch.py

@@ -1,11 +1,19 @@
+import torch
+import warnings
+import weakref
+
 import comfy_aimdo.model_vbar
+from comfy.cli_args import args
 import comfy.memory_management
 import comfy.model_management
 import comfy.ops
 
 PREFETCH_QUEUES = []
+GRAPH_MODULES = weakref.WeakSet()
+GRAPH_WARMED_MODULES = weakref.WeakSet()
+GRAPH_CAPTURE_STREAMS = {}
 
-def cleanup_prefetched_modules(comfy_modules):
+def cleanup_prefetched_modules(module, comfy_modules):
     for s in comfy_modules:
         prefetch = getattr(s, "_prefetch", None)
         if prefetch is None:
@@ -17,39 +25,86 @@ def cleanup_prefetched_modules(comfy_modules):
         if prefetch["signature"] is not None:
             comfy_aimdo.model_vbar.vbar_unpin(s._v)
         delattr(s, "_prefetch")
+    if getattr(module, "_v_block_faulted", False):
+        comfy_aimdo.model_vbar.vbar_unpin(module._v_block)
+        del module._v_block_faulted
+
+def _drop_graph(module):
+    graph = getattr(module, "_comfy_graph", None)
+    if graph is None:
+        return
+    # reset() through the bound method surfaces the allocator's benign
+    # "uncaptured free of a captured allocation" as catchable Python warnings;
+    # a plain del frees from the C++ dealloc path and spams stderr instead
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        graph["graph"].reset()
+    del module._comfy_graph
 
 def cleanup_prefetch_queues():
-    global PREFETCH_QUEUES
+    global PREFETCH_QUEUES, GRAPH_CAPTURE_STREAMS
 
     for queue in PREFETCH_QUEUES:
         for entry in queue:
             if entry is None or not isinstance(entry, tuple):
                 continue
             _, prefetch_state = entry
-            comfy_modules = prefetch_state[1]
+            prefetched_module, comfy_modules = prefetch_state
             if comfy_modules is not None:
-                cleanup_prefetched_modules(comfy_modules)
+                cleanup_prefetched_modules(prefetched_module, comfy_modules)
     PREFETCH_QUEUES = []
+    for module in GRAPH_MODULES:
+        _drop_graph(module)
+    GRAPH_MODULES.clear()
+    GRAPH_WARMED_MODULES.clear()
+    GRAPH_CAPTURE_STREAMS = {}
 
-def prefetch_queue_pop(queue, device, module):
+def prefetch_queue_pop(queue, device, module, dtype=None, core=None, enable_graph=False, generator=None):
+    enable_graph = enable_graph and not args.disable_cuda_graphs and comfy.model_management.is_device_cuda(device) and getattr(module, "_v_block", None) is not None
     if queue is None:
+        if core is not None:
+            core()
         return
 
+    capture_stream = None
+    if enable_graph:
+        capture_stream = GRAPH_CAPTURE_STREAMS.get(device)
+        if capture_stream is None:
+            capture_stream = torch.cuda.Stream(device=device)
+            GRAPH_CAPTURE_STREAMS[device] = capture_stream
+
+    signature = None
+    graph_hit = False
+    graph = getattr(module, "_comfy_graph", None) if enable_graph else None
+    if graph is not None:
+        signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block)
+        if signature is not None:
+            module._v_block_faulted = True
+            graph_hit = comfy_aimdo.model_vbar.vbar_signature_compare(signature, graph["signature"])
+
     consumed = queue.pop(0)
     if consumed is not None:
         offload_stream, prefetch_state = consumed
         if offload_stream is not None:
             offload_stream.wait_stream(comfy.model_management.current_stream(device))
-        _, comfy_modules = prefetch_state
+        prefetched_module, comfy_modules = prefetch_state
         if comfy_modules is not None:
-            cleanup_prefetched_modules(comfy_modules)
+            cleanup_prefetched_modules(prefetched_module, comfy_modules)
 
+    if graph_hit:
+        queue[0] = (None, (module, []))
+        graph["graph"].replay()
+        return
+
+    fully_faulted = False
     prefetch = queue[0]
     if prefetch is not None:
         comfy_modules = []
-        for s in prefetch.modules():
-            if hasattr(s, "_v"):
-                comfy_modules.append(s)
+        prefetch_modules = prefetch if isinstance(prefetch, (list, tuple)) else (prefetch,)
+        for root in prefetch_modules:
+            for s in root.modules():
+                if hasattr(s, "_v"):
+                    comfy_modules.append(s)
 
         registerable_size = 0
         for s in comfy_modules:
@@ -59,11 +114,42 @@ def prefetch_queue_pop(queue, device, module):
                 if lowvram_fn is not None:
                     registerable_size += lowvram_fn.memory_required()
 
-        offload_stream = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True)
+        offload_stream, fully_faulted = comfy.ops.cast_modules_with_vbar(comfy_modules, None, device, None, True, return_faulted=True)
         if not comfy.model_management.args.fast_disk:
             comfy.model_management.ensure_pin_registerable(registerable_size)
         comfy.model_management.sync_stream(device, offload_stream)
-        queue[0] = (offload_stream, (prefetch, comfy_modules))
+        if fully_faulted and dtype is not None:
+            for comfy_module in comfy_modules:
+                comfy.ops.resolve_cast_module_with_vbar(comfy_module, dtype, device, dtype, None, False, return_weights=False)
+        queue[0] = (offload_stream, (module, comfy_modules))
+
+    if core is not None:
+        if enable_graph and fully_faulted and module in GRAPH_WARMED_MODULES:
+            if signature is None:
+                signature = comfy_aimdo.model_vbar.vbar_fault(module._v_block)
+                if signature is not None:
+                    module._v_block_faulted = True
+            if signature is not None:
+                _drop_graph(module)
+                graph = torch.cuda.CUDAGraph()
+                if generator is not None:
+                    graph.register_generator_state(generator)
+                capture_stream.wait_stream(comfy.model_management.current_stream(device))
+                with torch.cuda.graph(graph, stream=capture_stream, capture_error_mode="thread_local"):
+                    core()
+                comfy.model_management.current_stream(device).wait_stream(capture_stream)
+                graph.replay()
+                module._comfy_graph = {"graph": graph, "signature": signature}
+                GRAPH_MODULES.add(module)
+                return
+        if capture_stream is None:
+            core()
+        else:
+            capture_stream.wait_stream(comfy.model_management.current_stream(device))
+            with torch.cuda.stream(capture_stream):
+                core()
+            comfy.model_management.current_stream(device).wait_stream(capture_stream)
+            GRAPH_WARMED_MODULES.add(module)
 
 def make_prefetch_queue(queue, device, transformer_options):
     if (not transformer_options.get("prefetch_dynamic_vbars", False)

comfy/model_sampling.py

@@ -325,6 +325,27 @@ class ModelSamplingDiscreteFlow(torch.nn.Module):
             return 0.0
         return time_snr_shift(self.shift, 1.0 - percent)
 
+class ModelSamplingAV(ModelSamplingDiscreteFlow):
+    """Flow sampling for packed audio-video latents whose audio stream has its own flow shift.
+
+    Carrying the audio latent scaled onto the video schedule makes the pack an ordinary
+    single-schedule flow latent whose audio target is scaled by audio_scale.
+    """
+    def __init__(self, model_config=None):
+        super().__init__(model_config)
+        sampling_settings = model_config.sampling_settings if model_config is not None else {}
+        self.audio_shift = sampling_settings.get("audio_shift", None)
+
+    def set_parameters(self, shift=1.0, audio_shift=None, timesteps=1000, multiplier=1000):
+        self.audio_shift = audio_shift
+        super().set_parameters(shift=shift, timesteps=timesteps, multiplier=multiplier)
+
+    @property
+    def audio_scale(self):
+        if self.audio_shift is None:
+            return 1.0
+        return self.shift / self.audio_shift
+
 class StableCascadeSampling(ModelSamplingDiscrete):
     def __init__(self, model_config=None):
         super().__init__()

comfy/nested_tensor.py

@@ -83,6 +83,9 @@ class NestedTensor:
     def layout(self):
         return self.tensors[0].layout
 
+    def __repr__(self):
+        return f"{type(self).__name__}({self.tensors!r})"
+
 
 def cat_nested(tensors, *args, **kwargs):
     cated_tensors = []

comfy/ops.py

@@ -73,6 +73,8 @@ try:
             ]
 
             def scaled_dot_product_attention(q, k, v, *args, **kwargs):
+                if q.nelement() < 1024 * 128:  # arbitrary number, for small inputs cudnn attention seems slower
+                    return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
                 attn_mask = args[0] if len(args) > 0 else kwargs.get("attn_mask")
                 if kwargs.get("enable_gqa", False) and attn_mask is not None and not comfy.model_management.is_nvidia():
                     k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
@@ -121,10 +123,12 @@ def materialize_meta_param(s, param_keys):
 
 
 # FIXME: add n=1 cache hit fast path
-def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking):
+def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blocking, return_faulted=False):
     offload_stream = None
     cast_buffer = None
     cast_buffer_offset = 0
+    if return_faulted:
+        fully_faulted = all(not getattr(s, param_key + "_function", []) for s in comfy_modules for param_key in ("weight", "bias"))
 
     def ensure_offload_stream(module, required_size, check_largest):
         nonlocal offload_stream
@@ -161,6 +165,8 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
     for s in comfy_modules:
         signature = comfy_aimdo.model_vbar.vbar_fault(s._v)
         resident = comfy_aimdo.model_vbar.vbar_signature_compare(signature, s._v_signature)
+        if return_faulted and (signature is None or not resident):
+            fully_faulted = False
         prefetch = {
             "signature": signature,
             "resident": resident,
@@ -253,10 +259,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
         prefetch["needs_cast"] = needs_cast
         s._prefetch = prefetch
 
+    if return_faulted:
+        return offload_stream, fully_faulted
     return offload_stream
 
 
-def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant):
+def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, want_requant, return_weights=True):
 
     prefetch = getattr(s, "_prefetch", None)
 
@@ -296,7 +304,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
                 tensor = tensor.dequantize()
             return tensor
 
-        if orig.dtype != dtype or len(fns) > 0:
+        if (return_weights and orig.dtype != dtype) or len(fns) > 0:
             x = to_dequant(x, dtype)
         if not resident and lowvram_fn is not None:
             x = to_dequant(x, dtype if compute_dtype is None else compute_dtype)
@@ -323,7 +331,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
     if prefetch["signature"] is not None:
         prefetch["resident"] = True
 
-    return weight, bias
+    return (weight, bias) if return_weights else None
 
 
 def cast_bias_weight(s, input=None, dtype=None, device=None, bias_dtype=None, offloadable=False, compute_dtype=None, want_requant=False):
@@ -450,6 +458,26 @@ def uncast_bias_weight(s, weight, bias, offload_stream):
             device = bias_a.device
     os.wait_stream(comfy.model_management.current_stream(device))
 
+class CastBiasWeightContext:
+    # When initialized with no arguments or the first is None, the context
+    # will return the tuple (None, None).
+    def __init__(self, *args, **kwargs):
+        self.slf = args[0] if len(args) else None
+        self.state = (None, None) if self.slf is None else cast_bias_weight(*args, **kwargs)
+
+    def __enter__(self):
+        result = self.state
+        if len(result) < 3 or result[2] is None:
+            # Not offloaded, immediately drop references.
+            self.state = self.slf = None
+        return result[:2]
+
+    def __exit__(self, *_args) -> None:
+        if self.slf is None:
+            return
+        slf, state = self.slf, self.state
+        self.state = self.slf = None
+        uncast_bias_weight(slf, *state)
 
 class CastWeightBiasOp:
     comfy_cast_weights = False
@@ -536,10 +564,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.linear(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.linear(input, weight, bias)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -553,10 +579,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = self._conv_forward(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return self._conv_forward(input, weight, bias)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -570,10 +594,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = self._conv_forward(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return self._conv_forward(input, weight, bias)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -598,10 +620,8 @@ class disable_weight_init:
                 return super()._conv_forward(input, weight, bias, *args, **kwargs)
 
         def forward_comfy_cast_weights(self, input, autopad=None):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = self._conv_forward(input, weight, bias, autopad=autopad)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return self._conv_forward(input, weight, bias, autopad=autopad)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -615,10 +635,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -632,12 +650,10 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
-            running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
-            x = torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
+                running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
+                return torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -651,15 +667,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            if self.weight is not None:
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            else:
-                weight = None
-                bias = None
-                offload_stream = None
-            x = torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -674,15 +683,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            if self.weight is not None:
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            else:
-                weight = None
-                bias = None
-                offload_stream = None
-            x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -701,12 +703,10 @@ class disable_weight_init:
                 input, output_size, self.stride, self.padding, self.kernel_size,
                 num_spatial_dims, self.dilation)
 
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.conv_transpose2d(
-                input, weight, bias, self.stride, self.padding,
-                output_padding, self.groups, self.dilation)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.conv_transpose2d(
+                       input, weight, bias, self.stride, self.padding,
+                       output_padding, self.groups, self.dilation)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -725,12 +725,10 @@ class disable_weight_init:
                 input, output_size, self.stride, self.padding, self.kernel_size,
                 num_spatial_dims, self.dilation)
 
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.conv_transpose1d(
-                input, weight, bias, self.stride, self.padding,
-                output_padding, self.groups, self.dilation)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.conv_transpose1d(
+                       input, weight, bias, self.stride, self.padding,
+                       output_padding, self.groups, self.dilation)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -793,10 +791,8 @@ class disable_weight_init:
             output_dtype = out_dtype
             if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
                 out_dtype = None
-            weight, bias, offload_stream = cast_bias_weight(self, device=input.device, dtype=out_dtype, offloadable=True)
-            x = torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, device=input.device, dtype=out_dtype, offloadable=True) as (weight, bias):
+                return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
 
 
         def forward(self, *args, **kwargs):
@@ -872,7 +868,6 @@ def fp8_linear(self, input):
     if input.ndim != 2:
         return None
     lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
-    w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True)
     scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
 
     scale_input = torch.ones((), device=input.device, dtype=torch.float32)
@@ -881,15 +876,16 @@ def fp8_linear(self, input):
     layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
     quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
 
-    # Wrap weight in QuantizedTensor - this enables unified dispatch
-    # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
-    layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
-    quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
-    o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
+    with CastBiasWeightContext(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True) as (w, bias):
+        # Wrap weight in QuantizedTensor - this enables unified dispatch
+        # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
+        w_shape = tuple(w.shape)
+        layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=w_shape)
+        quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
+        o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
 
-    uncast_bias_weight(self, w, bias, offload_stream)
     if tensor_3d:
-        o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))
+        o = o.reshape((input_shape[0], input_shape[1], w_shape[0]))
 
     return o
 
@@ -909,10 +905,8 @@ class fp8_ops(manual_cast):
                 except Exception as e:
                     logging.info("Exception during fp8 op: {}".format(e))
 
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.linear(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.linear(input, weight, bias)
 
 CUBLAS_IS_AVAILABLE = False
 try:
@@ -928,10 +922,8 @@ if CUBLAS_IS_AVAILABLE:
                 return None
 
             def forward_comfy_cast_weights(self, input):
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                x = cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
-                uncast_bias_weight(self, weight, bias, offload_stream)
-                return x
+                with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                    return cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
 
             def forward(self, *args, **kwargs):
                 run_every_op()
@@ -1209,6 +1201,26 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
                 "quant_group_size": 64,
                 "linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
             }
+        elif module.quant_format == "asym_w4a8_int8":
+            # int4 weight (packed int8 [N,K/2]) + fp8 per-group scale (weight_s_rel),
+            # fp32 per-channel scale (weight_s_channel) + optional Lloyd-Max codebook.
+            scale = pop_scale("weight_s_rel")
+            if scale is None:
+                raise ValueError(f"Missing W4A8 group scale (weight_s_rel) for layer {layer_name}")
+            if scale.dtype == torch.uint8:
+                scale = scale.view(torch.float8_e4m3fn)
+            params_conf = layer_conf.get("params", {})
+            if not isinstance(params_conf, dict):
+                params_conf = {}
+            scales = {
+                "scale": scale,
+                "s_channel": pop_scale("weight_s_channel"),
+                "codebook": pop_scale("weight_codebook"),
+                "group_size": int(layer_conf.get("group_size", params_conf.get("group_size", 16))),
+                "convrot_groupsize": int(
+                    layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
+                ),
+            }
         else:
             raise ValueError(f"Unsupported quantization format: {module.quant_format}")
 
@@ -1260,6 +1272,9 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
             linear_dtype = getattr(params, "linear_dtype", "int4")
             if linear_dtype != "int4":
                 quant_conf["linear_dtype"] = linear_dtype
+        elif module.quant_format == "asym_w4a8_int8":
+            quant_conf["group_size"] = getattr(params, "group_size", 16)
+            quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
         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)
@@ -1319,29 +1334,28 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 want_requant=False,
                 weight_only_quant=False,
             ):
-                if weight_only_quant:
-                    weight, bias, offload_stream = cast_bias_weight(
-                        self,
-                        input=None,
-                        dtype=self.weight.dtype,
-                        device=input.device,
-                        bias_dtype=input.dtype,
-                        offloadable=True,
-                        compute_dtype=compute_dtype,
-                        want_requant=True,
-                    )
-                    weight = weight.to(dtype=input.dtype)
-                else:
-                    weight, bias, offload_stream = cast_bias_weight(
+                if not weight_only_quant:
+                    with CastBiasWeightContext(
                         self,
                         input,
                         offloadable=True,
                         compute_dtype=compute_dtype,
                         want_requant=want_requant,
-                    )
-                x = self._forward(input, weight, bias)
-                uncast_bias_weight(self, weight, bias, offload_stream)
-                return x
+                    ) as (weight, bias):
+                        return self._forward(input, weight, bias)
+
+                with CastBiasWeightContext(
+                    self,
+                    input=None,
+                    dtype=self.weight.dtype,
+                    device=input.device,
+                    bias_dtype=input.dtype,
+                    offloadable=True,
+                    compute_dtype=compute_dtype,
+                    want_requant=True,
+                ) as (weight, bias):
+                    weight = weight.to(dtype=input.dtype)
+                    return self._forward(input, weight, bias)
 
             def forward(self, input, *args, **kwargs):
                 run_every_op()
@@ -1366,25 +1380,20 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
 
                 # Training path: quantized forward with compute_dtype backward via autograd function
                 if (input.requires_grad and _use_quantized and quantize_input):
-
-                    weight, bias, offload_stream = cast_bias_weight(
+                    with CastBiasWeightContext(
                         self,
                         input,
                         offloadable=True,
                         compute_dtype=compute_dtype,
                         want_requant=True
-                    )
+                    ) as (weight, bias):
+                        scale = getattr(self, 'input_scale', None)
+                        if scale is not None:
+                            scale = comfy.model_management.cast_to_device(scale, input.device, None)
 
-                    scale = getattr(self, 'input_scale', None)
-                    if scale is not None:
-                        scale = comfy.model_management.cast_to_device(scale, input.device, None)
-
-                    output = QuantLinearFunc.apply(
-                        input, weight, bias, self.layout_type, scale, compute_dtype
-                    )
-
-                    uncast_bias_weight(self, weight, bias, offload_stream)
-                    return output
+                        return QuantLinearFunc.apply(
+                            input, weight, bias, self.layout_type, scale, compute_dtype
+                        )
 
                 # Inference path (unchanged)
                 if _use_quantized and quantize_input:
@@ -1495,13 +1504,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 """Cast the whole bank once; expert_linear inside reuses the cast.
                 Not re-entrant — do not nest calls on the same instance.
                 """
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                self._resident_bank = (weight, bias)
-                try:
-                    yield self
-                finally:
-                    self._resident_bank = None
-                    uncast_bias_weight(self, weight, bias, offload_stream)
+                with CastBiasWeightContext(self, input, offloadable=True) as self._resident_bank:
+                    try:
+                        yield self
+                    finally:
+                        self._resident_bank = None
 
             def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor:
                 """Linear against expert i's weight (with optional bias)."""
@@ -1509,11 +1516,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 if resident is not None:
                     weight, bias = resident
                     return self._expert_linear_impl(input, weight, bias, i)
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                try:
+                with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
                     return self._expert_linear_impl(input, weight, bias, i)
-                finally:
-                    uncast_bias_weight(self, weight, bias, offload_stream)
 
             def _expert_linear_impl(self, input, weight, bias, i):
                 if isinstance(weight, QuantizedTensor):
@@ -1616,28 +1620,26 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
 
                 # 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):
-                        params = qdata._params
-                        scale = params.scale
-                        qdata = qdata._qdata
-                    else:
-                        params = weight._params
-                        scale = None
+                    with CastBiasWeightContext(self, device=input.device, dtype=weight.dtype, offloadable=True) as (qdata, _bias):
+                        if isinstance(qdata, QuantizedTensor):
+                            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)
+                        # 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)
+                            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)
-                    uncast_bias_weight(self, qdata, None, offload_stream)
+                        x = torch.nn.functional.embedding(
+                            input, qdata, self.padding_idx, self.max_norm,
+                            self.norm_type, self.scale_grad_by_freq, self.sparse)
                     target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
                     x = x.to(dtype=target_dtype)
-                    if scale is not None and scale != 1.0:
+                    if scale is not None:
                         x = x * scale.to(dtype=target_dtype)
                     return x
 

comfy/quant_ops.py

@@ -28,6 +28,7 @@ try:
         TensorCoreNVFP4Layout as _CKNvfp4Layout,
         TensorCoreConvRotW4A4Layout as _CKTensorCoreConvRotW4A4Layout,
         TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
+        AsymW4A8Int8Layout as _CKAsymW4A8Int8Layout,
         register_layout_op,
         register_layout_class,
         get_layout_class,
@@ -39,7 +40,7 @@ try:
         cuda_version = tuple(map(int, str(torch.version.cuda).split('.')))
         if cuda_version < (13,):
             ck.registry.disable("cuda")
-            logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.")
+            logging.warning("WARNING: You need pytorch with cu130 or higher to use optimized CUDA operations.\nWARNING WARNING WARNING\nIf you are on nvidia 20 series and above it is required that you update your pytorch to cu130 or higher.\n")
 
     # 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
@@ -83,6 +84,9 @@ except ImportError as e:
     class _CKTensorCoreConvRotW4A4Layout:
         pass
 
+    class _CKAsymW4A8Int8Layout:
+        pass
+
     def register_layout_class(name, cls):
         pass
 
@@ -212,7 +216,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
 TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
 TensorWiseINT8Layout = _CKTensorWiseINT8Layout
 TensorCoreConvRotW4A4Layout = _CKTensorCoreConvRotW4A4Layout
-
+AsymW4A8Int8Layout = _CKAsymW4A8Int8Layout
 
 # ==============================================================================
 # Registry
@@ -226,6 +230,7 @@ register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
 register_layout_class("TensorCoreConvRotW4A4Layout", _CKTensorCoreConvRotW4A4Layout)
 if _CK_MXFP8_AVAILABLE:
     register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
+register_layout_class("AsymW4A8Int8Layout", _CKAsymW4A8Int8Layout)
 
 QUANT_ALGOS = {
     "float8_e4m3fn": {
@@ -268,6 +273,13 @@ QUANT_ALGOS["convrot_w4a4"] = {
     "quantize_input": False,
 }
 
+QUANT_ALGOS["asym_w4a8_int8"] = {
+    "storage_t": torch.int8,
+    "parameters": {"weight_scale"},
+    "comfy_tensor_layout": "AsymW4A8Int8Layout",
+    "quantize_input": False,
+}
+
 
 # ==============================================================================
 # Re-exports for backward compatibility
@@ -282,6 +294,7 @@ __all__ = [
     "TensorCoreNVFP4Layout",
     "TensorCoreConvRotW4A4Layout",
     "TensorWiseINT8Layout",
+    "AsymW4A8Int8Layout",
     "QUANT_ALGOS",
     "register_layout_op",
 ]

comfy/sample.py

@@ -37,6 +37,11 @@ def prepare_noise(latent_image, seed, noise_inds=None):
 
     return noises
 
+def prepare_empty_noise(latent_image):
+    if latent_image.is_nested:
+        return comfy.nested_tensor.NestedTensor([torch.zeros_like(t, device="cpu") for t in latent_image.unbind()])
+    return torch.zeros_like(latent_image, device="cpu")
+
 def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None, downscale_ratio_temporal=None):
     if latent_image.is_nested:
         return latent_image

comfy/samplers.py

@@ -1218,6 +1218,8 @@ class CFGGuider:
         return sampling_function(self.inner_model, x, timestep, self.conds.get("negative", None), self.conds.get("positive", None), self.cfg, model_options=model_options, seed=seed)
 
     def inner_sample(self, noise, latent_image, device, sampler, sigmas, denoise_mask, callback, disable_pbar, seed, latent_shapes=None):
+        self.inner_model.latent_shapes = latent_shapes
+
         if latent_image is not None and torch.count_nonzero(latent_image) > 0: #Don't shift the empty latent image.
             latent_image = self.inner_model.process_latent_in(latent_image)
 

comfy/sd.py

@@ -11,6 +11,7 @@ from .ldm.cascade.stage_c_coder import StageC_coder
 from .ldm.audio.autoencoder import AudioOobleckVAE
 import comfy.ldm.genmo.vae.model
 import comfy.ldm.lightricks.vae.causal_video_autoencoder
+import comfy.ldm.lightricks.vae.na_diffusion_decoder
 import comfy.ldm.lightricks.vae.audio_vae
 import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
@@ -24,6 +25,7 @@ import comfy.ldm.cogvideo.vae
 import comfy.ldm.hunyuan_video.vae
 import comfy.ldm.mmaudio.vae.autoencoder
 import comfy.ldm.audio.vae_sa3
+import comfy.ldm.minimax_music.dav
 import comfy.pixel_space_convert
 import comfy.weight_adapter
 import yaml
@@ -31,6 +33,7 @@ import math
 import os
 
 import comfy.utils
+import comfy.ops
 
 from . import clip_vision
 from . import gligen
@@ -73,6 +76,7 @@ import comfy.text_encoders.longcat_image
 import comfy.text_encoders.qwen35
 import comfy.text_encoders.qwen3vl
 import comfy.text_encoders.minimax
+import comfy.text_encoders.minimax_music
 import comfy.ldm.minimax.vae
 import comfy.ldm.minimax.audio_vae
 import comfy.text_encoders.boogu
@@ -514,7 +518,22 @@ class VAE:
         self.audio_sample_rate = 44100
 
         if config is None:
-            if "decoder.mid.block_1.mix_factor" in sd:
+            if "dec_in_proj.weight" in sd and "decoder.model.0.weight_g" in sd:  # MiniMax Music3 DAV
+                self.first_stage_model = comfy.ldm.minimax_music.dav.MiniMaxMusic3DAV(operations=comfy.ops.disable_weight_init)
+                self.latent_channels = 128
+                self.output_channels = 2
+                self.upscale_ratio = 512
+                self.downscale_ratio = 512
+                self.latent_dim = 1
+                self.process_output = lambda audio: audio
+                self.process_input = lambda audio: audio
+                self.working_dtypes = [torch.float32]
+                self.disable_offload = True
+                self.memory_used_decode = lambda shape, dtype: (shape[-1] * 512 * 1400 + 800_000_000) * model_management.dtype_size(dtype)
+                def _no_encode(*args, **kwargs):
+                    raise RuntimeError("MiniMax Music3 DAV cannot encode audio")
+                self.memory_used_encode = _no_encode
+            elif "decoder.mid.block_1.mix_factor" in sd:
                 encoder_config = {'double_z': True, 'z_channels': 4, 'resolution': 256, 'in_channels': 3, 'out_ch': 3, 'ch': 128, 'ch_mult': [1, 2, 4, 4], 'num_res_blocks': 2, 'attn_resolutions': [], 'dropout': 0.0}
                 decoder_config = encoder_config.copy()
                 decoder_config["video_kernel_size"] = [3, 1, 1]
@@ -583,6 +602,22 @@ class VAE:
                 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_x_t.weight" in sd:  # lightricks LTX 2.4 diffusion VAE decoder
+                vae_config = None
+                if metadata is not None and "config" in metadata:
+                    vae_config = json.loads(metadata["config"]).get("vae", None)
+                self.first_stage_model = comfy.ldm.lightricks.vae.na_diffusion_decoder.CausalDiffusionVAE(config=vae_config)
+                self.latent_channels = sd["decoder.conv_in.weight"].shape[1]
+                self.latent_dim = 3
+                self.disable_offload = True
+                self.crop_input = False  # generic crop would narrow the frame axis by the 32x spatial ratio
+                self.memory_used_decode = lambda shape, dtype: (1700 * shape[2] * shape[3] * shape[4] * (8 * 8 * 8)) * model_management.dtype_size(dtype)
+                self.memory_used_encode = lambda shape, dtype: (80 * max(shape[2], 7) * shape[3] * shape[4]) * model_management.dtype_size(dtype)
+                self.upscale_ratio = (lambda a: max(0, a * 8 - 7), 32, 32)
+                self.upscale_index_formula = (8, 32, 32)
+                self.downscale_ratio = (lambda a: max(0, math.floor((a + 7) / 8)), 32, 32)
+                self.downscale_index_formula = (8, 32, 32)
+                self.working_dtypes = [torch.bfloat16, torch.float32]
             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}
@@ -940,7 +975,11 @@ class VAE:
                 if not comfy.memory_management.aimdo_enabled:
                     self.disable_offload = True
             elif "decoder.transformer_blocks.0.scale1" in sd and "encoder.down.5.block.0.conv1.weight" in sd:  # MiniMax H3 video VAE
-                self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE()
+                minimax_ops = comfy.ops.disable_weight_init
+                minimax_quant = comfy.utils.detect_layer_quantization(sd, "")
+                if minimax_quant is not None:  # int8+convrot quantized decoder
+                    minimax_ops = comfy.ops.mixed_precision_ops(minimax_quant, dtype if dtype is not None else torch.float16)
+                self.first_stage_model = comfy.ldm.minimax.vae.MiniMaxH3VideoVAE(operations=minimax_ops)
                 self.latent_channels = 24
                 self.latent_dim = 3
                 # frames 17k+5 <-> latents 5k+2, 16x spatial
@@ -951,13 +990,21 @@ class VAE:
                 self.working_dtypes = [torch.float16, torch.float32]
                 # the model tiles internally (256px spatial, 17-frame temporal chunks)
                 self.handles_tiling = True
+                # decode finalizes straight to [0, 1] while streaming chunks out
+                self.process_output = lambda image: image
+                # one decoded temporal chunk (with overlap) is all that ever sits in VRAM
+                chunk_frames = (self.first_stage_model.tokens_chunk_size + self.first_stage_model.token_overlap) * self.first_stage_model.vae_ratio_t
+
                 def estimate_encode_memory(frames, height, width, dtype):
                     fixed = 110_000_000 if frames == 1 else 1_300_000_000
                     elements_per_pixel = 7 if frames == 1 else 9.5
+                    # only one clip of the input video is ever resident on the GPU
+                    frames = min(frames, self.first_stage_model.clip_length)
                     return (elements_per_pixel * frames * height * width + fixed) * model_management.dtype_size(dtype) * 1.03
 
                 def estimate_decode_memory(frames, height, width, dtype):
                     fixed = 110_000_000 if frames <= 22 else 270_000_000
+                    frames = min(frames, chunk_frames + 2)
                     return (9.5 * frames * height * width + fixed) * model_management.dtype_size(dtype) * 1.03
 
                 self.memory_used_encode = lambda shape, dtype: estimate_encode_memory(shape[2], shape[3], shape[4], dtype)
@@ -1194,6 +1241,7 @@ class VAE:
                 do_tile = True
 
             if do_tile:
+                pixel_samples = None
                 comfy.model_management.soft_empty_cache()
                 dims = samples_in.ndim - 2
                 if dims == 1 or self.extra_1d_channel is not None:
@@ -1209,16 +1257,48 @@ class VAE:
                     tile = 256 // self.spacial_compression_decode()
                     overlap = tile // 4
                     if self.handles_tiling:
+                        memory_used = self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, tile, tile, None), self.vae_dtype)
+                        model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
                         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))
+                        # Reserve as much as an untiled decode could use (capped by what the device can provide), then size the tiles to fill that reservation:
+                        # shrink the temporal tile until one tile fits, then grow the spatial tile while it still fits.
+                        budget = min(memory_used, int(model_management.get_total_memory(self.device) * 0.8))
+                        model_management.load_models_gpu([self.patcher], memory_required=budget, force_full_load=self.disable_offload)
+                        tile_t = samples_in.shape[2]
+                        est = lambda tt, txy: self.memory_used_decode(self._tile_bounded_shape(samples_in.shape, txy, txy, tt), self.vae_dtype)
+                        while tile_t > 2 and est(tile_t, tile) > budget:
+                            tile_t = -(-tile_t // 2)
+                        while tile * 2 <= max(samples_in.shape[3], samples_in.shape[4]) and est(tile_t, tile * 2) <= budget:
+                            tile *= 2
+                        overlap = tile // 4
+                        pixel_samples = self.decode_tiled_3d(samples_in, tile_t=tile_t, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
 
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
         return pixel_samples
 
+    def _tile_bounded_shape(self, shape, tile_x, tile_y, tile_t):
+        """Clamp a latent shape to one tile for memory estimates: peak memory of a tiled decode is per-tile. Only caller-provided tile dims are clamped."""
+        s = list(shape)
+        if len(s) == 5:
+            if tile_t is not None:
+                s[2] = min(s[2], tile_t)
+            if tile_y is not None:
+                s[3] = min(s[3], tile_y)
+            if tile_x is not None:
+                s[4] = min(s[4], tile_x)
+        elif len(s) == 4 and self.extra_1d_channel is None:
+            if tile_y is not None:
+                s[2] = min(s[2], tile_y)
+            if tile_x is not None:
+                s[3] = min(s[3], tile_x)
+        elif tile_x is not None:
+            s[-1] = min(s[-1], tile_x)
+        return tuple(s)
+
     def decode_tiled(self, samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
         self.throw_exception_if_invalid()
-        memory_used = self.memory_used_decode(samples.shape, self.vae_dtype) #TODO: calculate mem required for tile
+        memory_used = self.memory_used_decode(self._tile_bounded_shape(samples.shape, tile_x, tile_y, tile_t), self.vae_dtype)
         model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
         dims = samples.ndim - 2
         args = {}
@@ -1630,7 +1710,16 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
     clip_target.params = {}
     if len(clip_data) == 1:
         te_model = detect_te_model(clip_data[0])
-        if te_model == TEModel.CLIP_G:
+        if clip_type == CLIPType.MINIMAX and "model.audio_decoder.projection.weight" in clip_data[0]:
+            tokenizer_data["tokenizer_json"] = clip_data[0].pop("tokenizer_json", None)
+            quant = comfy.utils.detect_layer_quantization(clip_data[0], "")
+            if quant is not None:
+                model_options = model_options.copy()
+                model_options["quantization_metadata"] = quant
+            clip_target.params["projection_config"] = comfy.text_encoders.minimax_music.detect_merged_config(clip_data[0])
+            clip_target.clip = comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel
+            clip_target.tokenizer = comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer
+        elif te_model == TEModel.CLIP_G:
             if clip_type == CLIPType.STABLE_CASCADE:
                 clip_target.clip = sdxl_clip.StableCascadeClipModel
                 clip_target.tokenizer = sdxl_clip.StableCascadeTokenizer
@@ -1689,12 +1778,21 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             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, 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,
-                       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
+            if te_model == TEModel.GEMMA_4_12B and "text_embedding_projection.video_aggregate_embed.weight" in clip_data[0]:
+                clip_target.clip = comfy.text_encoders.lt.ltxav_te(
+                    **llama_detect(clip_data),
+                    **comfy.text_encoders.lt.sd_detect(clip_data),
+                    text_encoder_model=comfy.text_encoders.gemma4.gemma4_text_encoder_model(comfy.text_encoders.gemma4.Gemma4_12B),
+                    text_encoder_key="gemma4",
+                )
+                clip_target.tokenizer = comfy.text_encoders.lt.ltxav_gemma4_tokenizer(comfy.text_encoders.gemma4.Gemma4_12B.tokenizer)
+            else:
+                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,
+                           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)
         elif te_model == TEModel.GEMMA_2_2B:
             if clip_type == CLIPType.PIXELDIT:
@@ -1862,9 +1960,30 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.clip = comfy.text_encoders.kandinsky5.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.kandinsky5.Kandinsky5TokenizerImage
         elif clip_type == CLIPType.LTXV:
-            clip_target.clip = comfy.text_encoders.lt.ltxav_te(**llama_detect(clip_data), **comfy.text_encoders.lt.sd_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.lt.LTXAVGemmaTokenizer
-            tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
+            te_models = [detect_te_model(sd) for sd in clip_data]
+            gemma4_models = {
+                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,
+                TEModel.GEMMA_4_12B: comfy.text_encoders.gemma4.Gemma4_12B,
+            }
+            gemma4_type = next((model for model in te_models if model in gemma4_models), None)
+            if gemma4_type is None:
+                clip_target.clip = comfy.text_encoders.lt.ltxav_te(**llama_detect(clip_data), **comfy.text_encoders.lt.sd_detect(clip_data))
+                clip_target.tokenizer = comfy.text_encoders.lt.LTXAVGemmaTokenizer
+                gemma_sd = clip_data[te_models.index(TEModel.GEMMA_3_12B)] if TEModel.GEMMA_3_12B in te_models else clip_data[0]
+                tokenizer_data["spiece_model"] = gemma_sd.get("spiece_model", None)
+            else:
+                variant = gemma4_models[gemma4_type]
+                clip_target.clip = comfy.text_encoders.lt.ltxav_te(
+                    **llama_detect(clip_data),
+                    **comfy.text_encoders.lt.sd_detect(clip_data),
+                    text_encoder_model=comfy.text_encoders.gemma4.gemma4_text_encoder_model(variant),
+                    text_encoder_key="gemma4",
+                )
+                clip_target.tokenizer = comfy.text_encoders.lt.ltxav_gemma4_tokenizer(variant.tokenizer)
+                gemma_sd = clip_data[te_models.index(gemma4_type)]
+                tokenizer_data["tokenizer_json"] = gemma_sd.get("tokenizer_json", None)
         elif clip_type == CLIPType.NEWBIE:
             clip_target.clip = comfy.text_encoders.newbie.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.newbie.NewBieTokenizer

comfy/supported_models.py

@@ -16,6 +16,7 @@ import comfy.text_encoders.genmo
 import comfy.text_encoders.lt
 import comfy.text_encoders.hunyuan_video
 import comfy.text_encoders.minimax
+import comfy.text_encoders.minimax_music
 import comfy.text_encoders.cosmos
 import comfy.text_encoders.lumina2
 import comfy.text_encoders.wan
@@ -963,6 +964,7 @@ class MiniMaxH3(supported_models_base.BASE):
 
     sampling_settings = {
         "shift": 12.0,
+        "audio_shift": 3.0,
     }
 
     unet_extra_config = {}
@@ -1449,6 +1451,20 @@ class WAN22_Animate(WAN21_T2V):
         out = model_base.WAN22_Animate(self, device=device)
         return out
 
+class WAN_Animate2(WAN21_T2V):
+    unet_config = {
+        "image_model": "wan2.1",
+        "model_type": "animate2",
+    }
+
+    sampling_settings = {
+        "shift": 5.0,
+    }
+
+    def get_model(self, state_dict, prefix="", device=None):
+        out = model_base.WAN_Animate2(self, device=device)
+        return out
+
 class WAN22_T2V(WAN21_T2V):
     unet_config = {
         "image_model": "wan2.1",
@@ -2185,6 +2201,28 @@ class ACEStep15(supported_models_base.BASE):
 
         return supported_models_base.ClipTarget(comfy.text_encoders.ace15.ACE15Tokenizer, comfy.text_encoders.ace15.te(**detect))
 
+class MiniMaxMusic3(supported_models_base.BASE):
+    unet_config = {
+        "audio_model": "minimax_music3",
+    }
+
+    latent_format = comfy.latent_formats.MiniMaxMusic3
+    memory_usage_factor = 2.0
+    supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+    sampling_settings = {"multiplier": 1.0}
+
+    def get_model(self, state_dict, prefix="", device=None):
+        return model_base.MiniMaxMusic3(self, device=device)
+
+    def model_type(self, state_dict, prefix=""):
+        return model_base.ModelType.FLOW
+
+    def clip_target(self, state_dict={}):
+        detect = comfy.text_encoders.minimax_music.detect_merged_config(state_dict, self.text_encoder_key_prefix[0])
+        target = supported_models_base.ClipTarget(comfy.text_encoders.minimax_music.MiniMaxMusic3Tokenizer, comfy.text_encoders.minimax_music.MiniMaxMusic3TEModel)
+        target.params["projection_config"] = detect
+        return target
+
 
 class LongCatImage(supported_models_base.BASE):
     unet_config = {
@@ -2463,6 +2501,7 @@ models = [
     WAN22_S2V,
     WAN21_HuMo,
     WAN22_Animate,
+    WAN_Animate2,
     WAN21_FlowRVS,
     WAN21_SCAIL,
     WAN21_SCAIL2,
@@ -2478,6 +2517,7 @@ models = [
     ChromaRadiance,
     ACEStep,
     ACEStep15,
+    MiniMaxMusic3,
     Omnigen2,
     Boogu,
     MageFlow,

comfy/text_encoders/bpe_tokenizer.py

@@ -0,0 +1,333 @@
+"""
+Pure-Python byte-level BPE tokenizer.
+Supports loading from HuggingFace tokenizer.json (LLaMA-style)
+and from Mistral tekken JSON blobs.
+No dependency on the `transformers`, `tokenizers`, or `regex` packages.
+"""
+import base64
+import json
+import os
+import re
+import unicodedata
+
+
+# This is also the default pattern used by the previous MistralConverter path.
+_LLAMA_PATTERN = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+_CONTRACTIONS = ("'re", "'ve", "'ll", "'s", "'t", "'m", "'d")
+
+
+def _is_letter(c):
+    return unicodedata.category(c)[0] == "L"
+
+
+def _is_number(c):
+    return unicodedata.category(c)[0] == "N"
+
+
+def _is_whitespace(c):
+    return c in " \t\n\r\v\f\x85\u2028\u2029" or unicodedata.category(c) == "Zs"
+
+
+def _split_llama(text):
+    pieces = []
+    i = 0
+    while i < len(text):
+        contraction = None
+        if text[i] == "'":
+            for suffix in _CONTRACTIONS:
+                if text[i:i + len(suffix)].casefold() == suffix:
+                    contraction = text[i:i + len(suffix)]
+                    break
+        if contraction is not None:
+            pieces.append(contraction)
+            i += len(contraction)
+            continue
+
+        j = i
+        if text[j] not in "\r\n" and not _is_letter(text[j]) and not _is_number(text[j]):
+            j += 1
+        if j < len(text) and _is_letter(text[j]):
+            j += 1
+            while j < len(text) and _is_letter(text[j]):
+                j += 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        if _is_number(text[i]):
+            j = i + 1
+            while j < len(text) and j - i < 3 and _is_number(text[j]):
+                j += 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        j = i
+        if text[j] == " ":
+            j += 1
+        punct_start = j
+        while j < len(text) and not _is_whitespace(text[j]) and not _is_letter(text[j]) and not _is_number(text[j]):
+            j += 1
+        if j > punct_start:
+            while j < len(text) and text[j] in "\r\n":
+                j += 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        if _is_whitespace(text[i]):
+            j = i + 1
+            while j < len(text) and _is_whitespace(text[j]):
+                j += 1
+            last_newline = max(text.rfind("\r", i, j), text.rfind("\n", i, j))
+            if last_newline >= i:
+                j = last_newline + 1
+            elif j < len(text) and j - i > 1:
+                j -= 1
+            pieces.append(text[i:j])
+            i = j
+            continue
+
+        pieces.append(text[i])
+        i += 1
+    return pieces
+
+
+def _make_split_pattern(pattern_str):
+    if pattern_str != _LLAMA_PATTERN:
+        raise ValueError(f"Unsupported tokenizer split pattern: {pattern_str}")
+    return _split_llama
+
+
+def _bytes_to_unicode():
+    bs = (list(range(ord("!"), ord("~") + 1))
+          + list(range(ord("¡"), ord("¬") + 1))
+          + list(range(ord("®"), ord("ÿ") + 1)))
+    cs = bs[:]
+    n = 0
+    for b in range(2**8):
+        if b not in bs:
+            bs.append(b)
+            cs.append(2**8 + n)
+            n += 1
+    cs = [chr(n) for n in cs]
+    return dict(zip(bs, cs))
+
+
+class BPETokenizer:
+    """Byte-level BPE tokenizer with optional BOS prepending."""
+
+    def __init__(self, vocab, merges_by_pair, special_token_ids, pattern_str,
+                 byte_encoder, byte_decoder, bos_id=None):
+        self._vocab = vocab                          # str -> int
+        self._inv_vocab = {v: k for k, v in vocab.items()}
+        self._merges = merges_by_pair                # (str, str) -> priority int
+        self._special_token_ids = special_token_ids  # str -> int
+        self._special_ids = set(special_token_ids.values())
+        self._byte_encoder = byte_encoder
+        self._byte_decoder = byte_decoder
+        self._bos_id = bos_id
+
+        self._split = _make_split_pattern(pattern_str)
+        sorted_specials = sorted(special_token_ids.keys(), key=len, reverse=True)
+        if sorted_specials:
+            self._special_split = re.compile(
+                '(' + '|'.join(re.escape(s) for s in sorted_specials) + ')'
+            )
+        else:
+            self._special_split = None
+
+    def _bpe_encode_piece(self, chars):
+        if len(chars) <= 1:
+            return chars
+        while True:
+            min_rank = float('inf')
+            best_pair = None
+            for i in range(len(chars) - 1):
+                r = self._merges.get((chars[i], chars[i + 1]), float('inf'))
+                if r < min_rank:
+                    min_rank = r
+                    best_pair = (chars[i], chars[i + 1])
+            if best_pair is None:
+                break
+            merged = best_pair[0] + best_pair[1]
+            new_chars = []
+            i = 0
+            while i < len(chars):
+                if i < len(chars) - 1 and chars[i] == best_pair[0] and chars[i + 1] == best_pair[1]:
+                    new_chars.append(merged)
+                    i += 2
+                else:
+                    new_chars.append(chars[i])
+                    i += 1
+            chars = new_chars
+            if len(chars) == 1:
+                break
+        return chars
+
+    def _encode_raw(self, text):
+        ids = []
+        parts = self._special_split.split(text) if self._special_split else [text]
+        for part in parts:
+            if not part:
+                continue
+            if part in self._special_token_ids:
+                ids.append(self._special_token_ids[part])
+            else:
+                for piece in self._split(part):
+                    byte_chars = [self._byte_encoder[b] for b in piece.encode('utf-8')]
+                    for tok in self._bpe_encode_piece(byte_chars):
+                        ids.append(self._vocab[tok])
+        return ids
+
+    def __call__(self, text):
+        ids = self._encode_raw(text)
+        if self._bos_id is not None:
+            ids = [self._bos_id] + ids
+        return {"input_ids": ids}
+
+    def get_vocab(self):
+        return dict(self._vocab)
+
+    def decode(self, token_ids, skip_special_tokens=True):
+        buf = bytearray()
+        for tid in token_ids:
+            s = self._inv_vocab.get(tid, '')
+            if tid in self._special_ids:
+                if not skip_special_tokens:
+                    buf.extend(s.encode('utf-8'))
+            else:
+                for c in s:
+                    buf.append(self._byte_decoder[c])
+        return buf.decode('utf-8', errors='replace')
+
+
+def _extract_pattern(pretok):
+    if pretok.get('type') == 'Sequence':
+        for sub in pretok.get('pretokenizers', []):
+            if sub.get('type') == 'Split':
+                pat = sub.get('pattern', {})
+                if 'Regex' in pat:
+                    return pat['Regex']
+    elif pretok.get('type') == 'Split':
+        pat = pretok.get('pattern', {})
+        if 'Regex' in pat:
+            return pat['Regex']
+    return None
+
+
+def _extract_bos_id(post_processor, special_token_ids):
+    if post_processor.get('type') == 'TemplateProcessing':
+        single = post_processor.get('single', [])
+        if single and 'SpecialToken' in single[0]:
+            bos_str = single[0]['SpecialToken']['id']
+            return special_token_ids.get(bos_str)
+    return None
+
+
+def from_tokenizer_json(path):
+    """Load a BPETokenizer from a directory containing tokenizer.json."""
+    tok_file = os.path.join(path, 'tokenizer.json')
+    with open(tok_file, encoding='utf-8') as f:
+        data = json.load(f)
+
+    vocab = dict(data['model']['vocab'])  # str -> int
+
+    merges_by_pair = {}
+    for i, merge_str in enumerate(data['model'].get('merges', [])):
+        a, b = merge_str.split(' ', 1)
+        if (a, b) not in merges_by_pair:
+            merges_by_pair[(a, b)] = i
+
+    special_token_ids = {}
+    for tok in data.get('added_tokens', []):
+        special_token_ids[tok['content']] = tok['id']
+        vocab[tok['content']] = tok['id']  # include in vocab for inv_vocab decode
+
+    pattern = _extract_pattern(data.get('pre_tokenizer', {}))
+    if pattern is None:
+        raise ValueError(f"Could not extract regex pattern from {tok_file}")
+
+    bos_id = _extract_bos_id(data.get('post_processor', {}), special_token_ids)
+
+    byte_encoder = _bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+
+    return BPETokenizer(vocab, merges_by_pair, special_token_ids, pattern,
+                        byte_encoder, byte_decoder, bos_id=bos_id)
+
+
+def from_tekken_json(data):
+    """Build a BPETokenizer from a Mistral tekken JSON blob (bytes or str)."""
+    mistral_vocab = json.loads(data)
+    config = mistral_vocab["config"]
+
+    byte_encoder = _bytes_to_unicode()
+    byte_decoder = {v: k for k, v in byte_encoder.items()}
+
+    def tbts(b):
+        return "".join(byte_encoder[ord(c)] for c in b.decode("latin-1"))
+
+    special_token_offset = config["default_num_special_tokens"]
+    max_vocab = config["default_vocab_size"] - special_token_offset
+
+    raw_vocab = {}
+    for w in mistral_vocab["vocab"]:
+        r = w["rank"]
+        if r >= max_vocab:
+            continue
+        raw_vocab[base64.b64decode(w["token_bytes"])] = r + special_token_offset
+
+    special_tokens_dict = {}
+    for w in mistral_vocab["special_tokens"]:
+        if "token_bytes" in w:
+            special_tokens_dict[base64.b64decode(w["token_bytes"])] = w["rank"]
+        else:
+            special_tokens_dict[w["token_str"]] = w["rank"]
+
+    all_special = list(special_tokens_dict.keys())
+    combined = dict(special_tokens_dict)
+    combined.update(raw_vocab)
+
+    bpe_vocab = {}
+    merge_triples = []
+    for token, rank in combined.items():
+        if token not in all_special:
+            bpe_vocab[tbts(token)] = rank
+            if len(token) == 1:
+                continue
+            local = []
+            for i in range(1, len(token)):
+                pl, pr = token[:i], token[i:]
+                if pl in combined and pr in combined and (pl + pr) in combined:
+                    local.append((pl, pr, rank))
+            local.sort(key=lambda x: (combined[x[0]], combined[x[1]]))
+            merge_triples.extend(local)
+        else:
+            tok_str = token.decode("utf-8", errors="replace") if isinstance(token, bytes) else token
+            bpe_vocab[tok_str] = rank
+
+    merge_triples.sort(key=lambda v: v[2])
+
+    merges_by_pair = {}
+    for i, (pl, pr, _) in enumerate(merge_triples):
+        pair = (tbts(pl), tbts(pr))
+        if pair not in merges_by_pair:
+            merges_by_pair[pair] = i
+
+    special_str_ids = {}
+    for tok in all_special:
+        tok_str = tok.decode("utf-8", errors="replace") if isinstance(tok, bytes) else tok
+        if tok_str in bpe_vocab:
+            special_str_ids[tok_str] = bpe_vocab[tok_str]
+
+    return BPETokenizer(bpe_vocab, merges_by_pair, special_str_ids, _LLAMA_PATTERN,
+                        byte_encoder, byte_decoder, bos_id=None)
+
+
+class LlamaTokenizerFast:
+    """Drop-in replacement for transformers.LlamaTokenizerFast (read-only use)."""
+
+    @staticmethod
+    def from_pretrained(path, **kwargs):
+        return from_tokenizer_json(path)

comfy/text_encoders/flux.py

@@ -3,11 +3,10 @@ import comfy.text_encoders.t5
 import comfy.text_encoders.sd3_clip
 import comfy.text_encoders.llama
 import comfy.model_management
-from transformers import T5TokenizerFast, LlamaTokenizerFast, Qwen2Tokenizer
+from transformers import T5TokenizerFast, Qwen2Tokenizer
+from .bpe_tokenizer import from_tekken_json
 import torch
 import os
-import json
-import base64
 
 class T5XXLTokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None, tokenizer_data={}):
@@ -75,45 +74,13 @@ def flux_clip(dtype_t5=None, t5_quantization_metadata=None):
 def load_mistral_tokenizer(data):
     if torch.is_tensor(data):
         data = data.numpy().tobytes()
+    return {"tokenizer_object": from_tekken_json(data)}
 
-    try:
-        from transformers.integrations.mistral import MistralConverter
-    except ModuleNotFoundError:
-        from transformers.models.pixtral.convert_pixtral_weights_to_hf import MistralConverter
-
-    mistral_vocab = json.loads(data)
-
-    special_tokens = {}
-    vocab = {}
-
-    max_vocab = mistral_vocab["config"]["default_vocab_size"]
-    max_vocab -= len(mistral_vocab["special_tokens"])
-
-    for w in mistral_vocab["vocab"]:
-        r = w["rank"]
-        if r >= max_vocab:
-            continue
-
-        vocab[base64.b64decode(w["token_bytes"])] = r
-
-    for w in mistral_vocab["special_tokens"]:
-        if "token_bytes" in w:
-            special_tokens[base64.b64decode(w["token_bytes"])] = w["rank"]
-        else:
-            special_tokens[w["token_str"]] = w["rank"]
-
-    all_special = []
-    for v in special_tokens:
-        all_special.append(v)
-
-    special_tokens.update(vocab)
-    vocab = special_tokens
-    return {"tokenizer_object": MistralConverter(vocab=vocab, additional_special_tokens=all_special).converted(), "legacy": False}
 
 class MistralTokenizerClass:
     @staticmethod
-    def from_pretrained(path, **kwargs):
-        return LlamaTokenizerFast(**kwargs)
+    def from_pretrained(path, tokenizer_object=None, **kwargs):
+        return tokenizer_object
 
 class Mistral3Tokenizer(sd1_clip.SDTokenizer):
     def __init__(self, embedding_directory=None, embedding_size=5120, embedding_key='mistral3_24b', tokenizer_data={}):

comfy/text_encoders/gemma4.py

@@ -6,13 +6,16 @@ import numpy as np
 from tokenizers import Tokenizer
 from dataclasses import dataclass
 import math
+import re
 
 from comfy import sd1_clip
 import comfy.model_management
+import comfy.model_prefetch
 import comfy.ops
+import comfy.quant_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
+from comfy.text_encoders.llama import RMSNorm, MLP, BaseLlama, BaseGenerate, FixedKV, _make_scaled_embedding
 
 
 # Intentional minor divergences from transformers -reference implementation:
@@ -109,7 +112,28 @@ class Gemma4_12B_Config(Gemma4Config):
     suppress_tokens = [258883, 258882]
 
 
-# unfused RoPE as addcmul_ RoPE diverges from reference code
+class RingKV(FixedKV):
+    # sliding-window ring: writes wrap at capacity, validity saturates
+    def prepare(self, num_tokens):
+        capacity = self.key.shape[2]
+        self.position.fill_(self.index % capacity)
+        self.seqlen.fill_(min(self.index + num_tokens, capacity))
+
+
+def _fixed_kv_decode_mask(mask, cache, min_val):
+    capacity = cache.key.shape[2]
+    valid = min(cache.index + 1, capacity)
+    output = mask.new_full((*mask.shape[:-1], capacity), min_val)
+    if isinstance(cache, RingKV):
+        positions = torch.arange(cache.index + 1 - valid, cache.index + 1, device=mask.device) % capacity
+        output.index_copy_(-1, positions, mask[..., -valid:])
+    else:
+        output[..., :valid] = mask[..., :valid]
+    return output
+
+
+# unfused RoPE as addcmul_ RoPE diverges from reference code (vision only; text
+# layers use the kitchen split-half kernel, bitwise-equal to this with bf16 freqs)
 def _apply_rotary_pos_emb(x, freqs_cis):
     cos, sin = freqs_cis[0], freqs_cis[1]
     half = x.shape[-1] // 2
@@ -140,6 +164,23 @@ class Gemma4Attention(nn.Module):
         if config.k_norm == "gemma3":
             self.k_norm = RMSNorm(head_dim, eps=config.rms_norm_eps, device=device, dtype=dtype)
 
+    def _decode_attention(self, xq, cache, bias):
+        if bias is None:
+            # eager decode: slice the cache to the valid length (python-side index,
+            # no mask needed; a full ring is order-invariant under softmax)
+            n = min(cache.index + 1, cache.key.shape[2])
+            gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {}
+            attention = optimized_attention_for_device(xq.device, mask=False, small_input=True)
+            return attention(xq, cache.key[:, :, :n], cache.value[:, :, :n], self.num_heads, skip_reshape=True, scale=1.0, **gqa_kwargs)
+        # graph capture: fixed-length masked attention over the full capacity, explicit
+        # math (SDPA leaves its fast path on broadcast-bias + GQA and costs ~0.5ms/layer)
+        batch_size = xq.shape[0]
+        groups = self.num_heads // self.num_kv_heads
+        q = xq.reshape(batch_size, self.num_kv_heads, groups, self.head_dim)
+        scores = q @ cache.key.transpose(-1, -2) + bias
+        probs = torch.softmax(scores.float(), dim=-1).to(xq.dtype)
+        return (probs @ cache.value).reshape(batch_size, 1, self.inner_size)
+
     def forward(
         self,
         hidden_states: torch.Tensor,
@@ -156,10 +197,16 @@ class Gemma4Attention(nn.Module):
         if self.q_norm is not None:
             xq = self.q_norm(xq)
 
+        if isinstance(shared_kv, FixedKV):
+            # decode on a KV-shared layer: attend the source layer's fixed cache
+            xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
+            output = self._decode_attention(xq, shared_kv, attention_mask)
+            return self.o_proj(output), None, None
+
         if shared_kv is not None:
             xk, xv = shared_kv
             # Apply RoPE to Q only (K already has RoPE from source layer)
-            xq = _apply_rotary_pos_emb(xq, freqs_cis)
+            xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
             present_key_value = None
             shareable_kv = None
         else:
@@ -173,11 +220,39 @@ class Gemma4Attention(nn.Module):
             xv = rms_norm(xv)
             xk = xk.transpose(1, 2)
             xv = xv.transpose(1, 2)
-            xq = _apply_rotary_pos_emb(xq, freqs_cis)
-            xk = _apply_rotary_pos_emb(xk, freqs_cis)
+            xq = comfy.quant_ops.ck.apply_rope_split_half1(xq, freqs_cis)
+            xk = comfy.quant_ops.ck.apply_rope_split_half1(xk, freqs_cis)
 
             present_key_value = None
-            if past_key_value is not None:
+            fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None
+            if fixed_cache is not None:
+                if seq_length == 1:
+                    # CUDA-graphable decode: write at the device-side ring/linear position
+                    fixed_cache.key.index_copy_(2, fixed_cache.position, xk)
+                    fixed_cache.value.index_copy_(2, fixed_cache.position, xv)
+                    output = self._decode_attention(xq, fixed_cache, attention_mask)
+                    return self.o_proj(output), fixed_cache, None
+
+                # prefill: attend the local sequence, persist the tail into the cache
+                capacity = fixed_cache.key.shape[2]
+                index = fixed_cache.index
+                if index + seq_length <= capacity:
+                    fixed_cache.key[:, :, index:index + seq_length] = xk
+                    fixed_cache.value[:, :, index:index + seq_length] = xv
+                    if index > 0:
+                        xk = fixed_cache.key[:, :, :index + seq_length]
+                        xv = fixed_cache.value[:, :, :index + seq_length]
+                elif index == 0:
+                    # prefill longer than the sliding ring: attend the full local K/V
+                    # (per-query windows come from the prefill sliding mask), cache only
+                    # the last `capacity` keys at their wrapped slots (position % capacity)
+                    slots = torch.arange(seq_length - capacity, seq_length, device=xk.device) % capacity
+                    fixed_cache.key.index_copy_(2, slots, xk[:, :, -capacity:])
+                    fixed_cache.value.index_copy_(2, slots, xv[:, :, -capacity:])
+                else:
+                    raise RuntimeError("gemma4: chunked prefill past the sliding window is not supported")
+                present_key_value = fixed_cache
+            elif past_key_value is not None:
                 cumulative_len = 0
                 if len(past_key_value) > 0:
                     past_key, past_value, cumulative_len = past_key_value
@@ -245,6 +320,7 @@ class TransformerBlockGemma4(nn.Module):
         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):
+        output = x
         sliding_window = None
         if self.sliding_attention:
             sliding_window = self.sliding_attention
@@ -281,7 +357,8 @@ class TransformerBlockGemma4(nn.Module):
             x = self.post_per_layer_input_norm(x)
             x = residual + x
 
-        x = x * comfy.ops.cast_to_input(self.layer_scalar, x)
+        # in-place into the input buffer so CUDA-graph replays land in the static x
+        x = torch.mul(x, comfy.ops.cast_to_input(self.layer_scalar, x), out=output)
 
         return x, present_key_value, shareable_kv
 
@@ -290,6 +367,9 @@ class Gemma4Transformer(nn.Module):
     def __init__(self, config, device=None, dtype=None, ops=None):
         super().__init__()
         self.config = config
+        self.fixed_kv = True
+        self.prefetch_dynamic_vbars = True
+        self.graph_dynamic_vbar_blocks = True
 
         self.embed_tokens = _make_scaled_embedding(ops, config.vocab_size, config.hidden_size, config.hidden_size ** 0.5, device, dtype)
 
@@ -298,6 +378,19 @@ class Gemma4Transformer(nn.Module):
             for i in range(config.num_hidden_layers)
         ])
 
+        # KV-shared layers never run k_proj/v_proj/k_norm: their never-resolved vbar
+        # signatures would block layer graph capture, so prefetch only what executes
+        first_kv_shared = config.num_hidden_layers - config.num_kv_shared_layers if config.num_kv_shared_layers > 0 else config.num_hidden_layers
+        self._prefetch_units = []
+        for i, layer in enumerate(self.layers):
+            if i >= first_kv_shared:
+                dead = {layer.self_attn.k_proj, layer.self_attn.v_proj, layer.self_attn.k_norm}
+                self._prefetch_units.append([
+                    m for m in layer.modules() if next(m.children(), None) is None and m not in dead
+                ])
+            else:
+                self._prefetch_units.append(layer)
+
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps, device=device, dtype=dtype) if config.final_norm else None
 
         # Precompute RoPE inv_freq on CPU to match reference code's exact value
@@ -311,6 +404,9 @@ class Gemma4Transformer(nn.Module):
         sliding_inv = 1.0 / (config.rope_theta[1] ** (torch.arange(0, config.head_dim, 2).float() / config.head_dim))
         self.register_buffer("_sliding_inv_freq", sliding_inv, persistent=False)
 
+        if config.suppress_tokens:
+            self.register_buffer("_suppress_tokens", torch.tensor(config.suppress_tokens, dtype=torch.long), persistent=False)
+
         # Per-layer input mechanism
         self.hidden_size_per_layer_input = config.hidden_size_per_layer_input
         if self.hidden_size_per_layer_input:
@@ -322,19 +418,26 @@ class Gemma4Transformer(nn.Module):
                 self.hidden_size_per_layer_input, eps=config.rms_norm_eps,
                 device=device, dtype=dtype)
 
+    def get_dynamic_vram__units(self):
+        return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], [])
+
     def get_past_len(self, past_key_values):
         for kv in past_key_values:
+            if isinstance(kv, FixedKV):
+                return kv.index
             if len(kv) >= 3:
                 return kv[2]
         return 0
 
     def _freqs_from_inv(self, inv_freq, position_ids, device, dtype):
-        """Compute cos/sin from stored inv_freq"""
+        """Compute per-pair 2x2 rotation matrices [B, 1, S, d/2, 2, 2] from stored inv_freq"""
         inv_exp = inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(device)
         pos_exp = position_ids[:, None, :].float()
         freqs = (inv_exp @ pos_exp).transpose(1, 2)
-        emb = torch.cat((freqs, freqs), dim=-1)
-        return emb.cos().unsqueeze(1).to(dtype), emb.sin().unsqueeze(1).to(dtype)
+        cos, sin = freqs.cos(), freqs.sin()
+        mat = torch.stack((torch.stack((cos, -sin), dim=-1),
+                           torch.stack((sin, cos), dim=-1)), dim=-2)
+        return mat.unsqueeze(1).to(dtype)
 
     def compute_freqs_cis(self, position_ids, device, dtype=None):
         global_freqs = self._freqs_from_inv(self._global_inv_freq, position_ids, device, dtype)
@@ -401,6 +504,71 @@ class Gemma4Transformer(nn.Module):
         first_kv_shared = self.config.num_hidden_layers - num_kv_shared if num_kv_shared > 0 else self.config.num_hidden_layers
         shared_sliding_kv = None  # KV from last non-shared sliding layer
         shared_global_kv = None   # KV from last non-shared global layer
+        share_source = {}
+        if num_kv_shared > 0:
+            for i in range(first_kv_shared):
+                share_source[bool(self.layers[i].sliding_attention)] = i
+
+        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(
+            list(self._prefetch_units), x.device,
+            {"prefetch_dynamic_vbars": self.prefetch_dynamic_vbars and past_key_values is not None})
+
+        fixed_kv = (past_key_values is not None and len(past_key_values) > 0
+                    and isinstance(past_key_values[0], FixedKV))
+        decode = fixed_kv and seq_len == 1
+        # mirror the conditions under which prefetch_queue_pop can actually capture, so
+        # eager fallbacks keep the sliced decode path instead of the full-capacity one
+        enable_graph = (decode and mask is None and self.graph_dynamic_vbar_blocks
+                        and prefetch_queue is not None
+                        and hasattr(self.layers[0], "_v_block")
+                        and not comfy.model_management.args.disable_cuda_graphs
+                        and comfy.model_management.is_device_cuda(x.device))
+        decode_bias = None
+        decode_masks = None
+        if decode:
+            prepared = set()
+            for kv in past_key_values:
+                if isinstance(kv, FixedKV) and id(kv.position) not in prepared:
+                    kv.prepare(seq_len)
+                    prepared.add(id(kv.position))
+            if mask is not None:
+                decode_masks = {}
+                for kv in past_key_values:
+                    if isinstance(kv, FixedKV) and id(kv.position) not in decode_masks:
+                        decode_masks[id(kv.position)] = _fixed_kv_decode_mask(mask, kv, min_val)
+        if enable_graph:
+            # static buffers + per-capacity attention biases: layer graphs replay against
+            # stable storage, refreshed eagerly each step
+            capacities = tuple(sorted({kv.key.shape[2] for kv in past_key_values if isinstance(kv, FixedKV)}))
+            state_key = (x.shape, x.dtype, x.device, tuple(t.shape for t in freqs_cis), capacities,
+                         None if per_layer_inputs is None else per_layer_inputs.shape)
+            state = getattr(self, "_comfy_cross_step_state", None)
+            if state is None or state["key"] != state_key:
+                state = {"key": state_key,
+                         "x": torch.empty_like(x),
+                         "freqs_cis": [torch.empty_like(t) for t in freqs_cis],
+                         "bias": {c: torch.empty((1, 1, 1, c), dtype=x.dtype, device=x.device) for c in capacities},
+                         "per_layer": None if per_layer_inputs is None else torch.empty_like(per_layer_inputs),
+                         "bias_valid": -1}
+                self._comfy_cross_step_state = state
+                comfy.model_management._register_cross_step(self)
+            state["x"].copy_(x)
+            for source, target in zip(freqs_cis, state["freqs_cis"]):
+                target.copy_(source)
+            x = state["x"]
+            freqs_cis = state["freqs_cis"]
+            if per_layer_inputs is not None:
+                state["per_layer"].copy_(per_layer_inputs)
+                per_layer_inputs = state["per_layer"]
+            valid = past_len + 1
+            for capacity, bias in state["bias"].items():
+                if state["bias_valid"] != past_len:
+                    bias.fill_(min_val)
+                    bias[..., :min(valid, capacity)] = 0
+                elif past_len < capacity:
+                    bias[..., past_len:valid] = 0
+            state["bias_valid"] = valid
+            decode_bias = state["bias"]
 
         intermediate = None
         all_intermediate = None
@@ -429,12 +597,36 @@ class Gemma4Transformer(nn.Module):
 
             is_sliding = hasattr(layer, 'sliding_attention') and layer.sliding_attention
             if i >= first_kv_shared and num_kv_shared > 0:
-                shared = shared_sliding_kv if is_sliding else shared_global_kv
-                if shared is not None:
-                    layer_kwargs['shared_kv'] = shared
+                if decode:
+                    layer_kwargs['shared_kv'] = past_key_values[share_source[bool(is_sliding)]]
+                else:
+                    shared = shared_sliding_kv if is_sliding else shared_global_kv
+                    if shared is not None:
+                        layer_kwargs['shared_kv'] = shared
 
-            x, current_kv, shareable_kv = layer(x=x, attention_mask=mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
+            if enable_graph:
+                bias_cache = layer_kwargs.get('shared_kv', past_kv)
+                layer_mask = decode_bias[bias_cache.key.shape[2]]
+            elif decode:
+                bias_cache = layer_kwargs.get('shared_kv', past_kv)
+                layer_mask = None if decode_masks is None else decode_masks[id(bias_cache.position)]
+            else:
+                layer_mask = mask
 
+            result = []
+
+            def core():
+                nonlocal x
+                x, current_kv, shareable_kv = layer(x=x, attention_mask=layer_mask, freqs_cis=freqs_cis, past_key_value=past_kv, **layer_kwargs)
+                result.append((current_kv, shareable_kv))
+
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
+
+            if result:
+                current_kv, shareable_kv = result[0]
+            else:
+                # graph replay: the cache already holds this step's write
+                current_kv, shareable_kv = past_kv, None
             next_key_values.append(current_kv if current_kv is not None else ())
 
             # Only track the last sliding/global before the sharing boundary
@@ -447,6 +639,14 @@ class Gemma4Transformer(nn.Module):
             if i == intermediate_output:
                 intermediate = x.clone()
 
+        if prefetch_queue is not None:
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
+
+        if fixed_kv:
+            for kv in past_key_values:
+                if isinstance(kv, FixedKV):
+                    kv.advance(seq_len)
+
         if self.norm is not None:
             x = self.norm(x)
 
@@ -481,14 +681,37 @@ class Gemma4Base(BaseLlama, BaseGenerate, torch.nn.Module):
         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
+            logits.index_fill_(-1, self.model._suppress_tokens, torch.finfo(logits.dtype).min)
         return logits
 
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
-        past_key_values = []
-        for _ in range(self.model.config.num_hidden_layers):
-            past_key_values.append(())
-        return past_key_values
+        cfg = self.model.config
+        num_layers = cfg.num_hidden_layers
+        if not self.model.fixed_kv:
+            return [() for _ in range(num_layers)]
+        first_shared = num_layers - cfg.num_kv_shared_layers if cfg.num_kv_shared_layers > 0 else num_layers
+        # position/seqlen device tensors are shared per cache geometry and filled once per step
+        trackers = {}
+        caches = []
+        for i in range(num_layers):
+            if i >= first_shared:
+                caches.append(())
+                continue
+            sliding = cfg.sliding_attention[i % len(cfg.sliding_attention)] if cfg.sliding_attention else False
+            head_dim = cfg.head_dim if sliding else cfg.global_head_dim
+            k_eq_v = cfg.attention_k_eq_v and not sliding
+            kv_heads = cfg.num_global_key_value_heads if k_eq_v else cfg.num_key_value_heads
+            length = min(sliding, max_cache_len) if sliding else max_cache_len
+            cache_cls = RingKV if sliding else FixedKV
+            tracker = trackers.get((cache_cls, length))
+            if tracker is None:
+                tracker = (torch.empty((1,), device=device, dtype=torch.int64),
+                           torch.empty((batch,), device=device, dtype=torch.int32))
+                trackers[(cache_cls, length)] = tracker
+            # zero-init: decode attends full capacity with masked tails, 0*0 stays finite
+            key = torch.zeros((batch, kv_heads, length, head_dim), device=device, dtype=execution_dtype)
+            caches.append(cache_cls(key, torch.zeros_like(key), 0, tracker[0], tracker[1]))
+        return caches
 
     def preprocess_embed(self, embed, device):
         if embed["type"] == "image":
@@ -1183,6 +1406,7 @@ def _get_aspect_ratio_preserving_size(height, width, patch_size, max_patches, po
 
 class Gemma4_Tokenizer():
     tokenizer_json_data = None
+    prime_empty_thought = False
 
     def state_dict(self):
         if self.tokenizer_json_data is not None:
@@ -1333,8 +1557,8 @@ class Gemma4_Tokenizer():
                     num_samples = int(waveform.shape[-1] * 16000 / sample_rate) if sample_rate != 16000 else waveform.shape[-1]
                     n_audio_tokens = self._audio_token_count(num_samples)
                     media += "<|audio>" + "<|audio|>" * n_audio_tokens + "<audio|>"
-                # Non-thinking mode primes an empty thought channel so the model answers directly.
-                model_open = "" if thinking else "<|channel>thought\n<channel|>"
+                # 12B/31B prime a closed thought block for non-thinking mode, E2B/E4B must not: it cues them into reasoning inline.
+                model_open = "<|channel>thought\n<channel|>" if self.prime_empty_thought and not thinking else ""
                 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)
@@ -1401,11 +1625,13 @@ class Gemma4SDTokenizer(Gemma4_Tokenizer, sd1_clip.SDTokenizer):
 
     def decode(self, token_ids, **kwargs):
         text = super().decode(token_ids, skip_special_tokens=False)
-        # Translate thinking channel markers to standard <think>/</think> tags
+        # Only a close that ends a thought channel becomes </think>: generation primed with
+        # another channel leaves its opener in the prompt, so its close is not reasoning.
+        text = re.sub(r"<\|channel>thought\n(.*?)<channel\|>", r"<think>\n\1</think>", text, flags=re.DOTALL)
         text = text.replace("<|channel>thought\n", "<think>\n")
-        text = text.replace("<channel|>", "</think>")
         # Strip remaining special tokens
-        text = text.replace("<turn|>", "").replace("<eos>", "").strip()
+        text = re.sub(r"<\|channel>\w*\n?|<channel\|>|<\|turn>\w*\n?|<turn\|>", "", text)
+        text = text.replace("<eos>", "").strip()
         return text
 
 
@@ -1418,6 +1644,7 @@ class Gemma4Tokenizer(sd1_clip.SD1Tokenizer):
 class Gemma4UnifiedSDTokenizer(Gemma4SDTokenizer):
     """Encoder-free (gemma4_unified) audio: raw 16kHz waveform frames instead of mel spectrogram."""
     embedding_size = 3840
+    prime_empty_thought = True
 
     def _extract_audio_features(self, waveform, sample_rate):
         audio = self._resample_16k(waveform, sample_rate)
@@ -1443,14 +1670,14 @@ class Gemma4UnifiedTokenizer(Gemma4Tokenizer):
 class Gemma4Model(sd1_clip.SDClipModel):
     model_class = None
     def __init__(self, device="cpu", layer="all", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
+        llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
+        if llama_quantization_metadata is not None:
+            model_options = model_options.copy()
+            model_options["quantization_metadata"] = llama_quantization_metadata
         self.dtypes = set()
         self.dtypes.add(dtype)
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=self.model_class, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-    def process_tokens(self, tokens, device):
-        embeds, _, _, _ = super().process_tokens(tokens, device)
-        return embeds
-
     def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty=0.0):
         if isinstance(tokens, dict):
             tokens = next(iter(tokens.values()))
@@ -1474,8 +1701,19 @@ class Gemma4Model(sd1_clip.SDClipModel):
         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_clip_model(model_class):
+    return type('Gemma4Model_', (Gemma4Model,), {'model_class': model_class})
+
+
+def gemma4_text_encoder_model(model_class):
+    return type('Gemma4TextEncoderModel_', (Gemma4Model,), {
+        'model_class': model_class,
+        'process_tokens': sd1_clip.SDClipModel.process_tokens,
+    })
+
+
 def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=None):
-    clip_model = type('Gemma4Model_', (Gemma4Model,), {'model_class': model_class})
+    clip_model = gemma4_clip_model(model_class)
     class Gemma4TEModel_(sd1_clip.SD1ClipModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
             if llama_quantization_metadata is not None:
@@ -1484,12 +1722,15 @@ def gemma4_te(dtype_llama=None, llama_quantization_metadata=None, model_class=No
             if dtype_llama is not None:
                 dtype = dtype_llama
             super().__init__(device=device, dtype=dtype, name="gemma4", clip_model=clip_model, model_options=model_options)
+
+        def get_dynamic_vram__units(self):
+            return getattr(self, self.clip).transformer.model.get_dynamic_vram__units()
     return Gemma4TEModel_
 
 
 # Variants
 
-def _make_variant(config_cls):
+def _make_variant(config_cls, prime_empty_thought=False):
     audio = config_cls.audio_config is not None
     bases = (Gemma4AudioMixin, Gemma4Base) if audio else (Gemma4Base,)
     class Variant(*bases):
@@ -1499,8 +1740,8 @@ def _make_variant(config_cls):
             if audio:
                 self._init_audio(self.model.config, dtype, device, operations)
     embedding_size = config_cls.hidden_size
-    if embedding_size != Gemma4SDTokenizer.embedding_size:
-        tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size})
+    if embedding_size != Gemma4SDTokenizer.embedding_size or prime_empty_thought:
+        tok_cls = type('T', (Gemma4SDTokenizer,), {'embedding_size': embedding_size, 'prime_empty_thought': prime_empty_thought})
         class Tokenizer(Gemma4Tokenizer):
             tokenizer_class = tok_cls
         Variant.tokenizer = Tokenizer
@@ -1510,7 +1751,7 @@ 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_31B = _make_variant(Gemma4_31B_Config, prime_empty_thought=True)
 
 
 # Gemma4 12B Unified: encoder-free multimodal, distinct base/tokenizer (not via _make_variant).

comfy/text_encoders/hunyuan_video.py

@@ -2,7 +2,7 @@ from comfy import sd1_clip
 import comfy.model_management
 import comfy.text_encoders.llama
 from .hunyuan_image import HunyuanImageTokenizer
-from transformers import LlamaTokenizerFast
+from .bpe_tokenizer import LlamaTokenizerFast
 import torch
 import os
 import numbers

comfy/text_encoders/llama.py

@@ -5,15 +5,33 @@ from typing import Optional, Any, Tuple
 import math
 from tqdm import tqdm
 import comfy.utils
+import comfy_kitchen
 
 from comfy.ldm.modules.attention import optimized_attention_for_device
 import comfy.model_management
+import comfy.model_prefetch
 import comfy.ops
 import comfy.ldm.common_dit
 import comfy.clip_model
 
 from . import qwen_vl
 
+
+@dataclass
+class FixedKV:
+    key: torch.Tensor
+    value: torch.Tensor
+    index: int
+    position: torch.Tensor
+    seqlen: torch.Tensor
+
+    def prepare(self, num_tokens):
+        self.position.fill_(self.index)
+        self.seqlen.fill_(self.index + num_tokens)
+
+    def advance(self, num_tokens):
+        self.index += num_tokens
+
 @dataclass
 class Llama2Config:
     vocab_size: int = 128320
@@ -249,6 +267,9 @@ class Qwen3_8BConfig:
     rope_scale = None
     final_norm: bool = True
     lm_head: bool = True
+    fixed_kv: bool = False
+    merged_qkv: bool = False
+    merged_mlp: bool = False
     stop_tokens = [151643, 151645]
 
 @dataclass
@@ -498,9 +519,14 @@ class Attention(nn.Module):
         self.inner_size = self.num_heads * self.head_dim
 
         ops = ops or nn
-        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 * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
-        self.v_proj = ops.Linear(config.hidden_size, self.num_kv_heads * self.head_dim, bias=config.qkv_bias, device=device, dtype=dtype)
+        self.kv_size = self.num_kv_heads * self.head_dim
+        self.merged_qkv = getattr(config, "merged_qkv", False)
+        if self.merged_qkv:
+            self.qkv_proj = ops.Linear(config.hidden_size, self.inner_size + self.kv_size * 2, bias=config.qkv_bias, device=device, dtype=dtype)
+        else:
+            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.kv_size, bias=config.qkv_bias, device=device, dtype=dtype)
+            self.v_proj = ops.Linear(config.hidden_size, self.kv_size, 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
@@ -522,9 +548,12 @@ class Attention(nn.Module):
     ):
         batch_size, seq_length, _ = hidden_states.shape
 
-        xq = self.q_proj(hidden_states)
-        xk = self.k_proj(hidden_states)
-        xv = self.v_proj(hidden_states)
+        if self.merged_qkv:
+            xq, xk, xv = self.qkv_proj(hidden_states).split((self.inner_size, self.kv_size, self.kv_size), dim=-1)
+        else:
+            xq = self.q_proj(hidden_states)
+            xk = self.k_proj(hidden_states)
+            xv = self.v_proj(hidden_states)
 
         xq = xq.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
         xk = xk.view(batch_size, seq_length, self.num_kv_heads, self.head_dim).transpose(1, 2)
@@ -537,8 +566,29 @@ class Attention(nn.Module):
 
         xq, xk = apply_rope(xq, xk, freqs_cis=freqs_cis)
 
-        present_key_value = None
-        if past_key_value is not None:
+        fixed_cache = past_key_value if isinstance(past_key_value, FixedKV) else None
+        if fixed_cache is not None:
+            xq = xq.transpose(1, 2)
+            xk = xk.transpose(1, 2)
+            xv = xv.transpose(1, 2)
+            if seq_length == 1:
+                # CUDA-graphable decode path.
+                fixed_cache.key.index_copy_(1, fixed_cache.position, xk)
+                fixed_cache.value.index_copy_(1, fixed_cache.position, xv)
+                output = comfy_kitchen.flash_attention_decode(xq, fixed_cache.key, fixed_cache.value, fixed_cache.seqlen)
+                return self.o_proj(output.view(batch_size, seq_length, self.inner_size)), fixed_cache
+
+            fixed_cache.key[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xk)
+            fixed_cache.value[:, fixed_cache.index:fixed_cache.index + seq_length].copy_(xv)
+            xk = fixed_cache.key[:, :fixed_cache.index + seq_length]
+            xv = fixed_cache.value[:, :fixed_cache.index + seq_length]
+
+            xq = xq.transpose(1, 2)
+            xk = xk.transpose(1, 2)
+            xv = xv.transpose(1, 2)
+
+        present_key_value = fixed_cache
+        if fixed_cache is None and past_key_value is not None:
             index = 0
             num_tokens = xk.shape[2]
             if len(past_key_value) > 0:
@@ -569,15 +619,27 @@ class MLP(nn.Module):
     def __init__(self, config: Llama2Config, device=None, dtype=None, ops: Any = None, intermediate_size=None):
         super().__init__()
         intermediate_size = intermediate_size or config.intermediate_size
-        self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
-        self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
+        self.merged_mlp = getattr(config, "merged_mlp", False)
+        if self.merged_mlp:
+            self.gate_up_proj = ops.Linear(config.hidden_size, intermediate_size * 2, bias=False, device=device, dtype=dtype)
+        else:
+            self.gate_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
+            self.up_proj = ops.Linear(config.hidden_size, intermediate_size, bias=False, device=device, dtype=dtype)
         self.down_proj = ops.Linear(intermediate_size, config.hidden_size, bias=False, device=device, dtype=dtype)
         if config.mlp_activation == "silu":
             self.activation = torch.nn.functional.silu
+            self.merged_input_act = "swiglu"
         elif config.mlp_activation == "gelu_pytorch_tanh":
             self.activation = lambda a: torch.nn.functional.gelu(a, approximate="tanh")
+            self.merged_input_act = None
 
     def forward(self, x):
+        if self.merged_mlp:
+            x = self.gate_up_proj(x)
+            if self.merged_input_act is not None:
+                return comfy.ops.linear_input_act(self.down_proj, x, self.merged_input_act)
+            gate, up = x.chunk(2, dim=-1)
+            return self.down_proj(self.activation(gate) * up)
         return self.down_proj(self.activation(self.gate_proj(x)) * self.up_proj(x))
 
 class TransformerBlock(nn.Module):
@@ -596,6 +658,7 @@ class TransformerBlock(nn.Module):
         optimized_attention=None,
         past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ):
+        output = x
         # Self Attention
         residual = x
         x = self.input_layernorm(x)
@@ -612,7 +675,7 @@ class TransformerBlock(nn.Module):
         residual = x
         x = self.post_attention_layernorm(x)
         x = self.mlp(x)
-        x = residual + x
+        x = torch.add(residual, x, out=output)
 
         return x, present_key_value
 
@@ -641,6 +704,7 @@ class TransformerBlockGemma2(nn.Module):
         optimized_attention=None,
         past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ):
+        output = x
         sliding_window = None
         if self.transformer_type == 'gemma3':
             if self.sliding_attention:
@@ -676,7 +740,7 @@ class TransformerBlockGemma2(nn.Module):
         x = self.pre_feedforward_layernorm(x)
         x = self.mlp(x)
         x = self.post_feedforward_layernorm(x)
-        x = residual + x
+        x = torch.add(residual, x, out=output)
 
         return x, present_key_value
 
@@ -688,9 +752,14 @@ def _make_scaled_embedding(ops, vocab_size, hidden_size, scale, device, dtype):
 
 
 class Llama2_(nn.Module):
+    fixed_kv = False
+    graph_dynamic_vbar_blocks = False
+
     def __init__(self, config, device=None, dtype=None, ops=None):
         super().__init__()
         self.config = config
+        self.fixed_kv = getattr(config, "fixed_kv", False)
+        self.graph_dynamic_vbar_blocks = False
         self.vocab_size = config.vocab_size
 
         if self.config.transformer_type == "gemma2" or self.config.transformer_type == "gemma3":
@@ -713,8 +782,27 @@ class Llama2_(nn.Module):
         if config.lm_head:
             self.lm_head = ops.Linear(config.hidden_size, config.vocab_size, bias=False, device=device, dtype=dtype)
 
+    def get_dynamic_vram__units(self):
+        return (list(self.layers), []) if self.graph_dynamic_vbar_blocks else ([], [])
+
     def get_past_len(self, past_key_values):
-        return past_key_values[0][2]
+        first = past_key_values[0]
+        return first.index if isinstance(first, FixedKV) else first[2]
+
+    def init_kv_cache(self, batch, capacity, device, dtype):
+        caches = []
+        fixed_kv = self.fixed_kv and comfy_kitchen.flash_attention_decode_is_available(device)
+        for _ in range(self.config.num_hidden_layers):
+            if fixed_kv:
+                key = torch.empty((batch, capacity, self.config.num_key_value_heads, self.config.head_dim), device=device, dtype=dtype)
+                value = torch.empty_like(key)
+                position = torch.empty((1,), device=device, dtype=torch.int64)
+                seqlen = torch.empty((batch,), device=device, dtype=torch.int32)
+                caches.append(FixedKV(key, value, 0, position, seqlen))
+            else:
+                key = torch.empty((batch, self.config.num_key_value_heads, capacity, self.config.head_dim), device=device, dtype=dtype)
+                caches.append((key, torch.empty_like(key), 0))
+        return caches
 
     def compute_freqs_cis(self, position_ids, device):
         return precompute_freqs_cis(self.config.head_dim,
@@ -756,6 +844,33 @@ class Llama2_(nn.Module):
 
         optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
 
+        fixed_kv = past_key_values is not None and len(past_key_values) > 0 and isinstance(past_key_values[0], FixedKV)
+        enable_graph = self.graph_dynamic_vbar_blocks and fixed_kv and seq_len == 1 and mask is None
+        if enable_graph:
+            freqs_cis_groups = freqs_cis if isinstance(freqs_cis, list) else [freqs_cis]
+            cross_step_state_key = [(x.shape, x.stride(), x.dtype, x.device)]
+            for group in freqs_cis_groups:
+                for tensor in group:
+                    cross_step_state_key.append((tensor.shape, tensor.stride(), tensor.dtype, tensor.device))
+            cross_step_state_key = tuple(cross_step_state_key)
+            cross_step_state = getattr(self, "_comfy_cross_step_state", None)
+            if cross_step_state is None or cross_step_state["key"] != cross_step_state_key:
+                static_freqs_cis = []
+                for group in freqs_cis_groups:
+                    static_freqs_cis.append(tuple(torch.empty_like(tensor) for tensor in group))
+                if not isinstance(freqs_cis, list):
+                    static_freqs_cis = static_freqs_cis[0]
+                cross_step_state = {"key": cross_step_state_key, "x": torch.empty_like(x), "freqs_cis": static_freqs_cis}
+                self._comfy_cross_step_state = cross_step_state
+                comfy.model_management._register_cross_step(self)
+            cross_step_state["x"].copy_(x)
+            static_freqs_cis_groups = cross_step_state["freqs_cis"] if isinstance(freqs_cis, list) else [cross_step_state["freqs_cis"]]
+            for source_group, target_group in zip(freqs_cis_groups, static_freqs_cis_groups):
+                for source, target in zip(source_group, target_group):
+                    target.copy_(source)
+            x = cross_step_state["x"]
+            freqs_cis = cross_step_state["freqs_cis"]
+
         intermediate = None
         all_intermediate = None
         only_layers = None
@@ -769,7 +884,8 @@ class Llama2_(nn.Module):
             elif intermediate_output < 0:
                 intermediate_output = len(self.layers) + intermediate_output
 
-        next_key_values = []
+        prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.layers), x.device, {"prefetch_dynamic_vbars": getattr(self, "prefetch_dynamic_vbars", False)})
+        next_key_values = list(past_key_values) if past_key_values is not None else []
         for i, layer in enumerate(self.layers):
             if all_intermediate is not None:
                 if only_layers is None or (i in only_layers):
@@ -779,16 +895,24 @@ class Llama2_(nn.Module):
             if past_key_values is not None:
                 past_kv = past_key_values[i] if len(past_key_values) > 0 else []
 
-            x, current_kv = layer(
-                x=x,
-                attention_mask=mask,
-                freqs_cis=freqs_cis,
-                optimized_attention=optimized_attention,
-                past_key_value=past_kv,
-            )
+            if fixed_kv:
+                past_kv.prepare(seq_len)
 
-            if current_kv is not None:
-                next_key_values.append(current_kv)
+            def core():
+                nonlocal x
+                x, current_kv = layer(
+                    x=x,
+                    attention_mask=mask,
+                    freqs_cis=freqs_cis,
+                    optimized_attention=optimized_attention,
+                    past_key_value=past_kv,
+                )
+                if next_key_values:
+                    next_key_values[i] = current_kv
+
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, layer, x.dtype, core=core, enable_graph=enable_graph)
+            if fixed_kv:
+                next_key_values[i].advance(seq_len)
 
             # DeepStack: add per-layer visual features into the first len() decoder layers at image positions (Qwen3-VL)
             if deepstack_embeds is not None and i < len(deepstack_embeds):
@@ -797,6 +921,9 @@ class Llama2_(nn.Module):
             if i == intermediate_output:
                 intermediate = x.clone()
 
+        if prefetch_queue is not None:
+            comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, x.device, None)
+
         if self.norm is not None:
             x = self.norm(x)
 
@@ -810,7 +937,7 @@ class Llama2_(nn.Module):
         if intermediate is not None and final_layer_norm_intermediate and self.norm is not None:
             intermediate = self.norm(intermediate)
 
-        if len(next_key_values) > 0:
+        if next_key_values:
             return x, intermediate, next_key_values
         else:
             return x, intermediate
@@ -868,24 +995,13 @@ class BaseGenerate:
         else:
             module = self.model.embed_tokens
 
-        offload_stream = None
-        if module.comfy_cast_weights:
-            weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True)
-        else:
-            weight = self.model.embed_tokens.weight.to(x)
-
-        x = torch.nn.functional.linear(input, weight, None)
-
-        comfy.ops.uncast_bias_weight(module, weight, None, offload_stream)
-        return x
+        if not module.comfy_cast_weights:
+            return torch.nn.functional.linear(input, self.model.embed_tokens.weight.to(x), None)
+        with comfy.ops.CastBiasWeightContext(module, input, offloadable=True) as (weight, _bias):
+            return torch.nn.functional.linear(input, weight, None)
 
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
-        model_config = self.model.config
-        past_key_values = []
-        for x in range(model_config.num_hidden_layers):
-            past_key_values.append((torch.empty([batch, model_config.num_key_value_heads, max_cache_len, model_config.head_dim], device=device, dtype=execution_dtype),
-                                    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
+        return self.model.init_kv_cache(batch, max_cache_len, device, execution_dtype)
 
     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

comfy/text_encoders/lt.py

@@ -81,6 +81,17 @@ class LTXAVGemmaTokenizer(sd1_clip.SD1Tokenizer):
         super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, name="gemma3_12b", tokenizer=Gemma3_12BTokenizer)
 
 
+def ltxav_gemma4_tokenizer(tokenizer):
+    class LTXAVGemma4Tokenizer(tokenizer):
+        def __init__(self, embedding_directory=None, tokenizer_data={}):
+            super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
+            gemma_tokenizer = getattr(self, self.clip)
+            if gemma_tokenizer.min_length == 1:
+                gemma_tokenizer.min_length = 1024
+
+    return LTXAVGemma4Tokenizer
+
+
 class Gemma3_12BModel(sd1_clip.SDClipModel):
     def __init__(self, device="cpu", layer="all", layer_idx=None, dtype=None, attention_mask=True, model_options={}):
         llama_quantization_metadata = model_options.get("llama_quantization_metadata", None)
@@ -97,10 +108,10 @@ class Gemma3_12BModel(sd1_clip.SDClipModel):
         return self.transformer.generate(embeds, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, stop_tokens=[106], presence_penalty=presence_penalty)  # 106 is <end_of_turn>
 
 class DualLinearProjection(torch.nn.Module):
-    def __init__(self, in_dim, out_dim_video, out_dim_audio, dtype=None, device=None, operations=None):
+    def __init__(self, in_dim, out_dim_video, out_dim_audio, video_bias=True, audio_bias=True, dtype=None, device=None, operations=None):
         super().__init__()
-        self.audio_aggregate_embed = operations.Linear(in_dim, out_dim_audio, bias=True, dtype=dtype, device=device)
-        self.video_aggregate_embed = operations.Linear(in_dim, out_dim_video, bias=True, dtype=dtype, device=device)
+        self.audio_aggregate_embed = operations.Linear(in_dim, out_dim_audio, bias=audio_bias, dtype=dtype, device=device)
+        self.video_aggregate_embed = operations.Linear(in_dim, out_dim_video, bias=video_bias, dtype=dtype, device=device)
 
     def forward(self, x):
         source_dim = x.shape[-1]
@@ -112,22 +123,28 @@ class DualLinearProjection(torch.nn.Module):
         return torch.cat((video, audio), dim=-1)
 
 class LTXAVTEModel(torch.nn.Module):
-    def __init__(self, dtype_llama=None, device="cpu", dtype=None, text_projection_type="single_linear", model_options={}):
+    def __init__(self, dtype_llama=None, device="cpu", dtype=None, text_projection_type="single_linear", text_encoder_model=Gemma3_12BModel, text_encoder_key="gemma3_12b", video_projection_dim=3840, audio_projection_dim=2048, video_projection_bias=None, audio_projection_bias=True, model_options={}):
         super().__init__()
         self.dtypes = set()
         self.dtypes.add(dtype)
         self.compat_mode = False
         self.text_projection_type = text_projection_type
+        self.text_encoder_key = text_encoder_key
+        self.execution_device = None
 
-        self.gemma3_12b = Gemma3_12BModel(device=device, dtype=dtype_llama, model_options=model_options, layer="all", layer_idx=None)
+        self.gemma3_12b = text_encoder_model(device=device, dtype=dtype_llama, model_options=model_options, layer="all", layer_idx=None)
         self.dtypes.add(dtype_llama)
 
         operations = self.gemma3_12b.operations # TODO
+        text_encoder_config = self.gemma3_12b.transformer.model.config
+        projection_in_dim = text_encoder_config.hidden_size * (text_encoder_config.num_hidden_layers + 1)
+        if video_projection_bias is None:
+            video_projection_bias = self.text_projection_type == "dual_linear"
 
         if self.text_projection_type == "single_linear":
-            self.text_embedding_projection = operations.Linear(3840 * 49, 3840, bias=False, dtype=dtype, device=device)
+            self.text_embedding_projection = operations.Linear(projection_in_dim, video_projection_dim, bias=video_projection_bias, dtype=dtype, device=device)
         elif self.text_projection_type == "dual_linear":
-            self.text_embedding_projection = DualLinearProjection(3840 * 49, 4096, 2048, dtype=dtype, device=device, operations=operations)
+            self.text_embedding_projection = DualLinearProjection(projection_in_dim, video_projection_dim, audio_projection_dim, video_bias=video_projection_bias, audio_bias=audio_projection_bias, dtype=dtype, device=device, operations=operations)
 
 
     def enable_compat_mode(self):  # TODO: remove
@@ -161,7 +178,7 @@ class LTXAVTEModel(torch.nn.Module):
         self.execution_device = None
 
     def encode_token_weights(self, token_weight_pairs):
-        token_weight_pairs = token_weight_pairs["gemma3_12b"]
+        token_weight_pairs = token_weight_pairs[self.text_encoder_key]
 
         out, pooled, extra = self.gemma3_12b.encode_token_weights(token_weight_pairs)
         out = out[:, :, -torch.sum(extra["attention_mask"]).item():]
@@ -189,51 +206,54 @@ class LTXAVTEModel(torch.nn.Module):
         return out.to(device=out_device, dtype=torch.float), pooled, extra
 
     def generate(self, tokens, do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty):
-        return self.gemma3_12b.generate(tokens["gemma3_12b"], do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty)
+        return self.gemma3_12b.generate(tokens[self.text_encoder_key], do_sample, max_length, temperature, top_k, top_p, min_p, repetition_penalty, seed, presence_penalty)
 
     def load_sd(self, sd):
-        if "model.layers.47.self_attn.q_norm.weight" in sd:
-            return self.gemma3_12b.load_sd(sd)
-        else:
-            sdo = comfy.utils.state_dict_prefix_replace(sd, {"text_embedding_projection.aggregate_embed.weight": "text_embedding_projection.weight", "text_embedding_projection.": "text_embedding_projection."}, filter_keys=True)
-            if len(sdo) == 0:
-                sdo = sd
+        missing_all = []
+        unexpected_all = []
 
-            missing_all = []
-            unexpected_all = []
+        if "model.layers.0.self_attn.q_norm.weight" in sd:
+            gemma_sd = {k: v for k, v in sd.items() if not k.startswith("text_embedding_projection.")}
+            missing, unexpected = self.gemma3_12b.load_sd(gemma_sd)
+            missing_all.extend(missing)
+            unexpected_all.extend(unexpected)
 
-            for prefix, component in [("text_embedding_projection.", self.text_embedding_projection)]:
-                component_sd = {k.replace(prefix, ""): v for k, v in sdo.items() if k.startswith(prefix)}
-                if component_sd:
-                    missing, unexpected = component.load_state_dict(component_sd, strict=False, assign=getattr(self, "can_assign_sd", False))
-                    missing_all.extend([f"{prefix}{k}" for k in missing])
-                    unexpected_all.extend([f"{prefix}{k}" for k in unexpected])
+        sdo = comfy.utils.state_dict_prefix_replace(sd, {"text_embedding_projection.aggregate_embed.": "text_embedding_projection.", "text_embedding_projection.": "text_embedding_projection."}, filter_keys=True)
+        if len(sdo) == 0:
+            sdo = sd
 
-            if "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.2.attn1.to_q.bias" not in sd:  # TODO: remove
-                ww = sd.get("model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.bias", None)
-                if ww is not None:
-                    if ww.shape[0] == 3840:
-                        self.enable_compat_mode()
-                        sdv = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.video_embeddings_connector.": ""}, filter_keys=True)
-                        self.video_embeddings_connector.load_state_dict(sdv, strict=False, assign=getattr(self, "can_assign_sd", False))
-                        sda = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.audio_embeddings_connector.": ""}, filter_keys=True)
-                        self.audio_embeddings_connector.load_state_dict(sda, strict=False, assign=getattr(self, "can_assign_sd", False))
+        for prefix, component in [("text_embedding_projection.", self.text_embedding_projection)]:
+            component_sd = {k.replace(prefix, ""): v for k, v in sdo.items() if k.startswith(prefix)}
+            if component_sd:
+                missing, unexpected = component.load_state_dict(component_sd, strict=False, assign=getattr(self, "can_assign_sd", False))
+                missing_all.extend([f"{prefix}{k}" for k in missing])
+                unexpected_all.extend([f"{prefix}{k}" for k in unexpected])
 
-            return (missing_all, unexpected_all)
+        if "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.2.attn1.to_q.bias" not in sd:  # TODO: remove
+            ww = sd.get("model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.0.attn1.to_q.bias", None)
+            if ww is not None:
+                if ww.shape[0] == 3840:
+                    self.enable_compat_mode()
+                    sdv = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.video_embeddings_connector.": ""}, filter_keys=True)
+                    self.video_embeddings_connector.load_state_dict(sdv, strict=False, assign=getattr(self, "can_assign_sd", False))
+                    sda = comfy.utils.state_dict_prefix_replace(sd, {"model.diffusion_model.audio_embeddings_connector.": ""}, filter_keys=True)
+                    self.audio_embeddings_connector.load_state_dict(sda, strict=False, assign=getattr(self, "can_assign_sd", False))
+
+        return (missing_all, unexpected_all)
 
     def memory_estimation_function(self, token_weight_pairs, device=None):
         constant = 6.0
         if comfy.model_management.should_use_bf16(device):
             constant /= 2.0
 
-        token_weight_pairs = token_weight_pairs.get("gemma3_12b", [])
+        token_weight_pairs = token_weight_pairs.get(self.text_encoder_key, [])
         m = min([sum(1 for _ in itertools.takewhile(lambda x: x[0] == 0, sub)) for sub in token_weight_pairs])
 
         num_tokens = sum(map(lambda a: len(a), token_weight_pairs)) - m
         num_tokens = max(num_tokens, 642)
         return num_tokens * constant * 1024 * 1024
 
-def ltxav_te(dtype_llama=None, llama_quantization_metadata=None, text_projection_type="single_linear"):
+def ltxav_te(dtype_llama=None, llama_quantization_metadata=None, text_projection_type="single_linear", text_encoder_model=Gemma3_12BModel, text_encoder_key="gemma3_12b", video_projection_dim=3840, audio_projection_dim=2048, video_projection_bias=None, audio_projection_bias=True):
     class LTXAVTEModel_(LTXAVTEModel):
         def __init__(self, device="cpu", dtype=None, model_options={}):
             if llama_quantization_metadata is not None:
@@ -241,16 +261,29 @@ def ltxav_te(dtype_llama=None, llama_quantization_metadata=None, text_projection
                 model_options["llama_quantization_metadata"] = llama_quantization_metadata
             if dtype_llama is not None:
                 dtype = dtype_llama
-            super().__init__(dtype_llama=dtype_llama, device=device, dtype=dtype, text_projection_type=text_projection_type, model_options=model_options)
+            super().__init__(dtype_llama=dtype_llama, device=device, dtype=dtype, text_projection_type=text_projection_type, text_encoder_model=text_encoder_model, text_encoder_key=text_encoder_key, video_projection_dim=video_projection_dim, audio_projection_dim=audio_projection_dim, video_projection_bias=video_projection_bias, audio_projection_bias=audio_projection_bias, model_options=model_options)
     return LTXAVTEModel_
 
 
 def sd_detect(state_dict_list, prefix=""):
     for sd in state_dict_list:
-        if "{}text_embedding_projection.audio_aggregate_embed.bias".format(prefix) in sd:
-            return {"text_projection_type": "dual_linear"}
-        if "{}text_embedding_projection.weight".format(prefix) in sd or "{}text_embedding_projection.aggregate_embed.weight".format(prefix) in sd:
-            return {"text_projection_type": "single_linear"}
+        video_key = "{}text_embedding_projection.video_aggregate_embed.weight".format(prefix)
+        audio_key = "{}text_embedding_projection.audio_aggregate_embed.weight".format(prefix)
+        if video_key in sd and audio_key in sd:
+            return {
+                "text_projection_type": "dual_linear",
+                "video_projection_dim": sd[video_key].shape[0],
+                "audio_projection_dim": sd[audio_key].shape[0],
+                "video_projection_bias": "{}text_embedding_projection.video_aggregate_embed.bias".format(prefix) in sd,
+                "audio_projection_bias": "{}text_embedding_projection.audio_aggregate_embed.bias".format(prefix) in sd,
+            }
+        for key in ("{}text_embedding_projection.weight".format(prefix), "{}text_embedding_projection.aggregate_embed.weight".format(prefix)):
+            if key in sd:
+                return {
+                    "text_projection_type": "single_linear",
+                    "video_projection_dim": sd[key].shape[0],
+                    "video_projection_bias": key.removesuffix("weight") + "bias" in sd,
+                }
     return {}
 
 

comfy/text_encoders/lumina2.py

@@ -49,10 +49,6 @@ class Gemma3_4B_Vision_Model(sd1_clip.SDClipModel):
 
         super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"start": 2, "pad": 0}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Gemma3_4B_Vision, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
 
-    def process_tokens(self, tokens, device):
-        embeds, _, _, _ = super().process_tokens(tokens, device)
-        return embeds
-
 class LuminaModel(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, model_options={}, name="gemma2_2b", clip_model=Gemma2_2BModel):
         super().__init__(device=device, dtype=dtype, name=name, clip_model=clip_model, model_options=model_options)

comfy/text_encoders/minimax_music.py

@@ -0,0 +1,117 @@
+import torch
+from tokenizers import Tokenizer
+
+import comfy.ops
+from comfy.ldm.minimax_music.ar import CFG_SCALE, CFG_TOP_K, MAX_AUDIO_FRAMES, MiniMaxMusic3AR
+from comfy.ldm.minimax_music.prompt import SPECIAL_TOKEN_IDS, build_prompt
+
+
+MODEL_CONFIG = {
+    "vocab_size": 200000,
+    "hidden_size": 4096,
+    "intermediate_size": 12288,
+    "num_hidden_layers": 36,
+    "num_attention_heads": 32,
+    "num_key_value_heads": 8,
+    "max_position_embeddings": 10240,
+    "rms_norm_eps": 1e-6,
+    "rope_theta": 1000000.0,
+    "head_dim": 128,
+    "audio_vocab_size": 1024,
+    "audio_num_codebooks": 8,
+    "decoder_num_heads": 16,
+    "decoder_intermediate_size": 6144,
+    "decoder_num_layers": 4,
+}
+
+
+def detect_merged_config(state_dict, prefix=""):
+    return {
+        "merged_qkv": "{}model.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
+        "merged_mlp": "{}model.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
+        "decoder_merged_qkv": "{}model.audio_decoder.layers.0.self_attn.qkv_proj.weight".format(prefix) in state_dict,
+        "decoder_merged_mlp": "{}model.audio_decoder.layers.0.mlp.gate_up_proj.weight".format(prefix) in state_dict,
+    }
+
+
+class MiniMaxMusic3Tokenizer:
+    def __init__(self, embedding_directory=None, tokenizer_data={}):
+        tokenizer_json = tokenizer_data.get("tokenizer_json")
+        if tokenizer_json is None:
+            raise ValueError("MiniMax Music3 text encoder checkpoint is missing tokenizer_json")
+        if torch.is_tensor(tokenizer_json):
+            tokenizer_json = tokenizer_json.detach().cpu().numpy().tobytes()
+        self.tokenizer_json = tokenizer_json
+        self.tokenizer = Tokenizer.from_str(tokenizer_json.decode("utf-8"))
+        for token, expected in SPECIAL_TOKEN_IDS.items():
+            if self.tokenizer.token_to_id(token) != expected:
+                raise ValueError(f"MiniMax Music3 tokenizer mismatch for {token}")
+
+    def tokenize_with_weights(self, text, return_word_ids=False, **kwargs):
+        prompt = build_prompt(text, kwargs.get("lyrics", ""))
+        token_ids = self.tokenizer.encode(prompt, add_special_tokens=False).ids
+        return {
+            "minimax_music3": [[(token, 1.0) for token in token_ids]],
+            "seed": int(kwargs.get("seed", 0)),
+            "max_audio_frames": int(kwargs.get("max_audio_frames", MAX_AUDIO_FRAMES)),
+            "cfg_scale": float(kwargs.get("cfg_scale", CFG_SCALE)),
+            "top_k": int(kwargs.get("top_k", CFG_TOP_K)),
+        }
+
+    def state_dict(self):
+        return {"tokenizer_json": torch.frombuffer(bytearray(self.tokenizer_json), dtype=torch.uint8)}
+
+    def decode(self, token_ids, skip_special_tokens=True):
+        return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens)
+
+
+class MiniMaxMusic3TEModel(MiniMaxMusic3AR):
+    def __init__(self, device="cpu", dtype=None, model_options={}, projection_config=None):
+        dtype = torch.bfloat16
+        quant_config = model_options.get("quantization_metadata", None)
+        operations = model_options.get("custom_operations", None)
+        if operations is None:
+            operations = comfy.ops.mixed_precision_ops(quant_config, dtype) if quant_config is not None else comfy.ops.manual_cast
+        super().__init__({**MODEL_CONFIG, **(projection_config or {})}, dtype, device, operations)
+        self.dtypes = {dtype}
+        self.execution_device = device
+
+    def set_clip_options(self, options):
+        self.execution_device = options.get("execution_device", self.execution_device)
+
+    def reset_clip_options(self):
+        pass
+
+    def get_dynamic_vram__units(self):
+        units, last_units = self.model.get_dynamic_vram__units()
+        if self.model.pruned_embedding:
+            last_units = [*last_units, self.model.embed_tokens_prefill]
+        return [(self.model.audio_decoder, self.model.audio_extra_embedding), *units], last_units
+
+    def encode_token_weights(self, token_weight_pairs):
+        token_ids = [token for token, _ in token_weight_pairs["minimax_music3"][0]]
+        input_ids = torch.tensor([token_ids], dtype=torch.long)
+        seed = token_weight_pairs["seed"]
+        max_audio_frames = token_weight_pairs["max_audio_frames"]
+        cfg_scale = token_weight_pairs["cfg_scale"]
+        top_k = token_weight_pairs["top_k"]
+        hidden = self.generate(input_ids, seed, max_audio_frames, self.execution_device, cfg_scale, top_k)
+        return hidden.unsqueeze(0), None, {}
+
+    def load_state_dict(self, state_dict, strict=True, assign=False):
+        if self.model.pruned_embedding is None:
+            self.model.pruned_embedding = "model.embed_tokens_prefill.weight" in state_dict
+            if self.model.pruned_embedding:
+                del self.model.embed_tokens
+            else:
+                del self.model.embed_tokens_prefill, self.model.embed_tokens_audio
+        if self.model.pruned_lm_head is None:
+            self.model.pruned_lm_head = "model.lm_head_pruned.weight" in state_dict
+            if self.model.pruned_lm_head:
+                del self.model.lm_head
+            else:
+                del self.model.lm_head_pruned
+        return super().load_state_dict(state_dict, strict=strict, assign=assign)
+
+    def load_sd(self, state_dict):
+        return self.load_state_dict(state_dict, strict=False, assign=getattr(self, "can_assign_sd", False))

comfy_api/latest/_io.py

@@ -847,6 +847,61 @@ class Load3DAnimation(Load3D):
     ...
 
 
+@comfytype(io_type="LAYERS")
+class Layers(ComfyTypeIO):
+    BlendMode = Literal[
+        "normal", "multiply", "screen", "overlay", "darken", "lighten",
+        "color-dodge", "color-burn", "hard-light", "soft-light", "difference",
+        "exclusion", "linear-dodge", "linear-burn", "vivid-light", "pin-light",
+        "linear-light", "hard-mix", "subtract", "divide", "grain-extract",
+        "grain-merge", "hue", "saturation", "color", "luminosity",
+    ]
+
+    class LayerItem(TypedDict):
+        image: torch.Tensor
+        type: Literal["raster"]
+        x: NotRequired[int]
+        y: NotRequired[int]
+        mask: NotRequired[torch.Tensor]
+        z_index: int
+        name: NotRequired[str]
+        opacity: NotRequired[float]
+        blend_mode: NotRequired["Layers.BlendMode"]
+        visible: NotRequired[bool]
+        flip_h: NotRequired[bool]
+        flip_v: NotRequired[bool]
+        rotation: NotRequired[float]
+        w: NotRequired[int]
+        h: NotRequired[int]
+
+    class Document(TypedDict):
+        version: int
+        canvas: NotRequired[tuple[int, int]]
+        layers: list["Layers.LayerItem"]
+
+    Type = Document
+
+
+@comfytype(io_type="COMPOSITOR")
+class Compositor(ComfyTypeIO):
+    class LayerState(TypedDict):
+        version: NotRequired[int]
+        canvas: dict
+        background: NotRequired[dict]
+        inputs: NotRequired[list[str]]
+        order: NotRequired[list[int]]
+        layers: list[dict]
+
+    Type = LayerState
+
+    class Input(WidgetInput):
+        def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None,
+                     socketless: bool=True, default: dict=None, advanced: bool=None):
+            super().__init__(id, display_name, optional, tooltip, None, default, socketless, None, None, None, None, advanced)
+            if default is None:
+                self.default = {}
+
+
 @comfytype(io_type="PHOTOMAKER")
 class Photomaker(ComfyTypeIO):
     Type = Any
@@ -2403,6 +2458,8 @@ __all__ = [
     "Load3DModelInfo",
     "Load3D",
     "Load3DAnimation",
+    "Compositor",
+    "Layers",
     "Photomaker",
     "Point",
     "FaceAnalysis",

comfy_api/latest/_ui.py

@@ -260,6 +260,7 @@ class ImageSaveHelper:
 class AudioSaveHelper:
     """A helper class with static methods to handle audio saving and metadata."""
     _OPUS_RATES = [8000, 12000, 16000, 24000, 48000]
+    _FORMATS = {"flac", "mp3", "opus"}
 
     @staticmethod
     def save_audio(
@@ -270,6 +271,9 @@ class AudioSaveHelper:
         format: str = "flac",
         quality: str = "128k",
     ) -> list[SavedResult]:
+        if format not in AudioSaveHelper._FORMATS:
+            raise ValueError(f"Unsupported audio format: {format!r}")
+
         full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
             filename_prefix, _get_directory_by_folder_type(folder_type)
         )

comfy_api_nodes/apis/bfl.py

@@ -1,7 +1,7 @@
 from enum import Enum
 from typing import Any
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field
 
 
 class BFLFluxExpandImageRequest(BaseModel):
@@ -121,6 +121,8 @@ class BFLFluxProGenerateResponse(BaseModel):
 class BFLStatus(str, Enum):
     task_not_found = "Task not found"
     pending = "Pending"
+    reasoning = "Reasoning"
+    generating = "Generating"
     request_moderated = "Request Moderated"
     content_moderated = "Content Moderated"
     ready = "Ready"
@@ -132,3 +134,35 @@ class BFLFluxStatusResponse(BaseModel):
     status: BFLStatus = Field(...)
     result: dict[str, Any] | None = Field(None)
     progress: float | None = Field(None, ge=0.0, le=1.0)
+
+
+class Flux3VideoRequest(BaseModel):
+    """Fields shared by every generation mode of /v1/flux-3-video."""
+
+    model_config = ConfigDict(extra="forbid")
+
+    prompt: str = Field(...)
+    aspect_ratio: str = Field("auto")
+    duration: int | str = Field("auto", description="Whole seconds, or 'auto'.")
+    resolution: str = Field("hd", description="'hd' is the 720p class, 'fhd' the 1080p class.")
+    generate_audio: bool = Field(True)
+    safety_tolerance: int = Field(2, description="0 is the strictest; conditioned modes cap at 2.")
+
+
+class Flux3TextToVideoRequest(Flux3VideoRequest):
+    mode: str = Field("t2v")
+
+
+class Flux3ImageToVideoRequest(Flux3VideoRequest):
+    mode: str = Field("i2v")
+    keyframes: list[str] | list[tuple[float, str]] = Field(
+        ...,
+        description="Images (URL or base64), or [seconds, image] pairs pinning each to a time.",
+    )
+
+
+class Flux3VideoContinuationRequest(Flux3VideoRequest):
+    mode: str = Field("v2v")
+    start_video: str = Field(
+        ..., description="MP4 (URL or base64); the new clip carries on from its final frames."
+    )

comfy_api_nodes/apis/bria.py

@@ -57,6 +57,81 @@ class BriaRemoveBackgroundRequest(BaseModel):
     seed: int = Field(...)
 
 
+class BriaGenFillRequest(BaseModel):
+    image: str = Field(...)
+    mask: str = Field(
+        ...,
+        description="Binary mask defining the region to fill: white (255) pixels are generated, "
+        "black (0) pixels are preserved. Must have the same aspect ratio as the image.",
+    )
+    prompt: str = Field(...)
+    negative_prompt: str | None = Field(None)
+    refine_prompt: bool = Field(True)
+    seed: int = Field(...)
+    prompt_content_moderation: bool = Field(False, description="If true, returns 422 on prompt moderation failure.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image or mask moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
+class BriaEraseRequest(BaseModel):
+    image: str = Field(...)
+    mask: str = Field(
+        ...,
+        description="Binary mask defining the region to erase: white (255) pixels are removed, "
+        "black (0) pixels are preserved. Must have the same aspect ratio as the image.",
+    )
+    mask_type: str = Field("manual", description="'manual' for hand-drawn masks, 'automatic' for segmentation masks.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image or mask moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
+class BriaExpandRequest(BaseModel):
+    image: str = Field(...)
+    aspect_ratio: str | float | None = Field(
+        None,
+        description="Target ratio: a preset string (1:1, 2:3, 3:2, 3:4, 4:3, 4:5, 5:4, 9:16, 16:9) "
+        "or a float between 0.5 and 3.0. When set, the canvas/placement fields are ignored.",
+    )
+    canvas_size: list[int] | None = Field(None, description="Output canvas [width, height]; area up to 5000x5000.")
+    original_image_size: list[int] | None = Field(
+        None, description="Size [width, height] of the original image inside the canvas."
+    )
+    original_image_location: list[int] | None = Field(
+        None,
+        description="Top-left corner [x, y] of the original image inside the canvas; "
+        "values may fall outside the canvas, cropping the image.",
+    )
+    prompt: str | None = Field(None, description="If omitted, Bria auto-generates a prompt from the image.")
+    negative_prompt: str | None = Field(None)
+    seed: int = Field(...)
+    prompt_content_moderation: bool = Field(False, description="If true, returns 422 on prompt moderation failure.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
+class BriaIncreaseResolutionRequest(BaseModel):
+    image: str = Field(...)
+    desired_increase: int = Field(..., description="Resolution multiplier, 2 or 4.")
+    visual_input_content_moderation: bool = Field(
+        False, description="If true, returns 422 on image moderation failure."
+    )
+    visual_output_content_moderation: bool = Field(
+        False, description="If true, returns 422 on visual output moderation failure."
+    )
+
+
 class BriaStatusResponse(BaseModel):
     request_id: str = Field(...)
     status_url: str = Field(...)
@@ -72,6 +147,26 @@ class BriaRemoveBackgroundResponse(BaseModel):
     result: BriaRemoveBackgroundResult | None = Field(None)
 
 
+class BriaImageResult(BaseModel):
+    image_url: str = Field(...)
+
+
+class BriaImageResultResponse(BaseModel):
+    status: str = Field(...)
+    result: BriaImageResult | None = Field(None)
+
+
+class BriaExpandResult(BaseModel):
+    image_url: str = Field(...)
+    prompt: str | None = Field(None)
+    seed: int | None = Field(None)
+
+
+class BriaExpandResponse(BaseModel):
+    status: str = Field(...)
+    result: BriaExpandResult | None = Field(None)
+
+
 class BriaImageEditResult(BaseModel):
     structured_prompt: str = Field(...)
     image_url: str = Field(...)

comfy_api_nodes/apis/bytedance.py

@@ -35,6 +35,23 @@ class Seedream4TaskCreationRequest(BaseModel):
     optimize_prompt_options: Seedream5OptimizePromptOptions | None = None
 
 
+class Seedream5LayerOptimizePromptOptions(BaseModel):
+    mode: Literal["standard", "fast"] = Field(...)
+
+
+class Seedream5LayerSeparationRequest(BaseModel):
+    model: str = Field(...)
+    prompt: str | None = Field(None)
+    image: str = Field(..., description="Single image URL")
+    size: str = Field("auto")
+    seed: int = Field(..., ge=0, le=2147483647)
+    response_format: str = Field("url")
+    output_format: str = Field("png")
+    layer_decomposition: bool = Field(True)
+    watermark: bool = Field(False)
+    optimize_prompt_options: Seedream5LayerOptimizePromptOptions | None = Field(None)
+
+
 class ImageTaskCreationResponse(BaseModel):
     model: str = Field(...)
     created: int = Field(..., description="Unix timestamp (in seconds) indicating time when the request was created.")
@@ -95,9 +112,11 @@ class Seedance2TaskCreationRequest(BaseModel):
     generate_audio: bool | None = Field(None)
     resolution: str | None = Field(None)
     ratio: str | None = Field(None)
-    duration: int | None = Field(None, ge=4, le=15)
+    duration: int | None = Field(None)
     seed: int | None = Field(None, ge=0, le=2147483647)
     watermark: bool | None = Field(None)
+    output_format: str | None = Field(None)
+    omni_reference_task_type: str | None = Field(None, description="One of: auto, reference, edit, extend.")
 
 
 class TaskCreationResponse(BaseModel):
@@ -186,6 +205,10 @@ SEEDANCE2_PRICE_PER_1K_TOKENS = {
     ("dreamina-seedance-2-0-mini", True, "480p"): 0.0021,
     ("dreamina-seedance-2-0-mini", False, "720p"): 0.0035,
     ("dreamina-seedance-2-0-mini", True, "720p"): 0.0021,
+    ("dreamina-seedance-2-5-260628", False, "480p"): 0.0107,
+    ("dreamina-seedance-2-5-260628", True, "480p"): 0.0064,
+    ("dreamina-seedance-2-5-260628", False, "720p"): 0.0107,
+    ("dreamina-seedance-2-5-260628", True, "720p"): 0.0064,
 }
 
 
@@ -304,8 +327,31 @@ SEEDANCE2_REF_VIDEO_PIXEL_LIMITS = {
         "480p": {"min": 409_600, "max": 927_408},
         "720p": {"min": 409_600, "max": 927_408},
     },
+    "dreamina-seedance-2-5-260628": {
+        "480p": {"min": 409_600, "max": 8_295_044},
+        "720p": {"min": 409_600, "max": 8_295_044},
+    },
 }
 
+SEEDANCE2_REFERENCE_LIMITS_DEFAULT = {
+    "max_images": 9,
+    "max_videos": 3,
+    "max_audios": 3,
+    "max_total_seconds": 15.1,
+}
+SEEDANCE2_REFERENCE_LIMITS = {
+    "dreamina-seedance-2-5-260628": {
+        "max_images": 30,
+        "max_videos": 10,
+        "max_audios": 10,
+        "max_total_seconds": 30.1,
+    },
+}
+
+
+def seedance2_reference_limits(model_id: str) -> dict:
+    return SEEDANCE2_REFERENCE_LIMITS.get(model_id, SEEDANCE2_REFERENCE_LIMITS_DEFAULT)
+
 # The time in this dictionary are given for 10 seconds duration.
 VIDEO_TASKS_EXECUTION_TIME = {
     "seedance-1-0-lite-t2v-250428": {

comfy_api_nodes/apis/grok.py

@@ -9,6 +9,7 @@ class ImageGenerationRequest(BaseModel):
     seed: int = Field(...)
     response_format: str = Field("url")
     resolution: str = Field(...)
+    quality: str | None = Field(None)
 
 
 class InputUrlObject(BaseModel):
@@ -28,6 +29,7 @@ class ImageEditRequest(BaseModel):
     seed: int = Field(...)
     response_format: str = Field("url")
     aspect_ratio: str | None = Field(...)
+    quality: str | None = Field(None)
 
 
 class VideoGenerationRequest(BaseModel):

comfy_api_nodes/apis/minimax.py

@@ -161,12 +161,30 @@ class Hailuo03TaskCreationRequest(BaseModel):
         ..., min_length=1
     )
     resolution: str = Field(...)
-    duration: int = Field(..., ge=5, le=15)
+    duration: int = Field(..., ge=4, le=15)
     ratio: str | None = Field(None)
     seed: int | None = Field(None, ge=0, le=4294967295)
     aigc_watermark: bool | None = Field(None)
 
 
+class Hailuo03ContextIRRequest(BaseModel):
+    model: str = Field(...)
+    content: list[Hailuo03TextContent | Hailuo03ImageContent | Hailuo03VideoContent | Hailuo03AudioContent] = Field(
+        ..., min_length=1
+    )
+    duration: int = Field(..., ge=4, le=15)
+    ratio: str | None = Field(None)
+
+
+class Hailuo03RegenerationRequest(BaseModel):
+    model: str = Field(...)
+    content: list[Hailuo03TextContent | Hailuo03ImageContent | Hailuo03VideoContent | Hailuo03AudioContent] = Field(
+        ..., min_length=1
+    )
+    resolution: str = Field(...)
+    aigc_watermark: bool | None = Field(None)
+
+
 class Hailuo03TaskCreationResponse(BaseModel):
     task_id: str = Field(...)
 
@@ -178,6 +196,7 @@ class Hailuo03TaskError(BaseModel):
 
 class Hailuo03TaskContent(BaseModel):
     url: str | None = Field(None)
+    prompt: str | None = Field(None)
 
 
 class Hailuo03TaskUsage(BaseModel):

comfy_api_nodes/apis/qwen.py

@@ -0,0 +1,46 @@
+from pydantic import BaseModel, Field
+
+
+class QwenImageContentItem(BaseModel):
+    image: str | None = Field(None)
+    text: str | None = Field(None)
+
+
+class QwenImageMessage(BaseModel):
+    role: str = Field("user")
+    content: list[QwenImageContentItem] = Field(...)
+
+
+class QwenImageInputField(BaseModel):
+    messages: list[QwenImageMessage] = Field(...)
+
+
+class QwenImageParametersField(BaseModel):
+    size: str | None = Field(None, description="Output resolution as 'width*height'; omit for the model default.")
+    n: int = Field(1, ge=1, le=6)
+    seed: int = Field(..., ge=0, le=2147483647)
+    prompt_extend: bool = Field(True)
+    watermark: bool = Field(False)
+    negative_prompt: str | None = Field(None)
+
+
+class QwenImageGenerationRequest(BaseModel):
+    model: str = Field(...)
+    input: QwenImageInputField = Field(...)
+    parameters: QwenImageParametersField = Field(...)
+
+
+class QwenImageChoice(BaseModel):
+    finish_reason: str | None = Field(None)
+    message: QwenImageMessage | None = Field(None)
+
+
+class QwenImageOutputField(BaseModel):
+    choices: list[QwenImageChoice] = Field(default_factory=list)
+
+
+class QwenImageGenerationResponse(BaseModel):
+    output: QwenImageOutputField | None = Field(None)
+    request_id: str = Field(...)
+    code: str | None = Field(None, description="Error code for the failed request.")
+    message: str | None = Field(None, description="Details about the failed request.")

comfy_api_nodes/apis/topaz.py

@@ -20,6 +20,41 @@ class ImageEnhanceRequest(BaseModel):
     color_preservation: str = Field("true", description="To preserve the original color")
 
 
+class ImageEnhanceRequestV2(BaseModel):
+    model: str = Field(...)
+    output_format: str = Field("png")
+    source_url: str = Field(...)
+    output_width: Optional[int] = Field(None)
+    output_height: Optional[int] = Field(None)
+    crop_to_fill: Optional[bool] = Field(None, description="Available for Reimagine only")
+    prompt: Optional[str] = Field(None, description="Available for Reimagine and Bloom 2")
+    creativity: Optional[int] = Field(None, description="From 1 to 9; available for Reimagine and Bloom 2")
+    subject_detection: Optional[str] = Field(None, description="Available for Reimagine only")
+    face_enhancement: Optional[bool] = Field(None, description="Available for Reimagine only")
+    face_enhancement_creativity: Optional[float] = Field(None, description="Is ignored if face_enhancement is false")
+    face_enhancement_strength: Optional[float] = Field(None, description="Is ignored if face_enhancement is false")
+    face_preservation: Optional[str] = Field(
+        None, description='String "true" or "false"; available for Reimagine only'
+    )
+    color_preservation: Optional[str] = Field(
+        None, description='String "true" or "false"; available for Reimagine and Bloom 2'
+    )
+    autoprompt: Optional[str] = Field(
+        None, description='String "true" or "false"; auto-generate a prompt, available for Bloom 2 only'
+    )
+    seed: Optional[int] = Field(None, description="Available for Bloom 2 only")
+    enhancement_strength: Optional[str] = Field(
+        None, description="low, medium or high; available for Wonder 3.5 only"
+    )
+    grain: Optional[str] = Field(
+        None, description='String "true" or "false"; available for Bloom 2 and Wonder 3.5'
+    )
+    grain_model: Optional[str] = Field(None, description="silver, gaussian or grey")
+    grain_strength: Optional[float] = Field(None, description="From 0 to 1")
+    grain_size: Optional[float] = Field(None, description="From 1 to 5")
+    grain_density: Optional[float] = Field(None, description="From 0 to 1")
+
+
 class ImageAsyncTaskResponse(BaseModel):
     process_id: str = Field(...)
 

comfy_api_nodes/nodes_bfl.py

@@ -1,3 +1,5 @@
+import math
+
 import torch
 from pydantic import BaseModel
 from typing_extensions import override
@@ -14,16 +16,23 @@ from comfy_api_nodes.apis.bfl import (
     BFLFluxVTORequest,
     BFLStatus,
     Flux2ProGenerateRequest,
+    Flux3ImageToVideoRequest,
+    Flux3TextToVideoRequest,
+    Flux3VideoContinuationRequest,
+    Flux3VideoRequest,
 )
 from comfy_api_nodes.util import (
     ApiEndpoint,
     convert_mask_to_image,
     download_url_to_image_tensor,
+    download_url_to_video_output,
     get_number_of_images,
     poll_op,
     resize_mask_to_image,
     sync_op,
     tensor_to_base64_string,
+    upload_images_to_comfyapi,
+    upload_video_to_comfyapi,
     validate_aspect_ratio_string,
     validate_image_dimensions,
     validate_string,
@@ -1007,6 +1016,385 @@ class Flux2ImageNode(IO.ComfyNode):
         return IO.NodeOutput(await download_url_to_image_tensor(response.result["sample"]))
 
 
+_FLUX3_ASPECT_RATIOS = ["auto", "21:9", "2:1", "16:9", "4:3", "1:1", "3:4", "9:16"]
+_FLUX3_MIN_DURATION = 5
+_FLUX3_MAX_DURATION = 20
+_FLUX3_DURATIONS = ["auto"] + [str(i) for i in range(_FLUX3_MIN_DURATION, _FLUX3_MAX_DURATION + 1)]
+_FLUX3_RESOLUTIONS = {"720p": "hd", "1080p": "fhd"}
+_FLUX3_MAX_IMAGES = 10
+_FLUX3_MIN_IMAGE_SIDE = 256
+_FLUX3_MAX_IMAGE_ASPECT = 64
+
+
+def _flux3_validate_image(image: torch.Tensor) -> None:
+    validate_image_dimensions(image, min_width=_FLUX3_MIN_IMAGE_SIDE, min_height=_FLUX3_MIN_IMAGE_SIDE)
+    height, width = image.shape[-3], image.shape[-2]
+    if max(width, height) > _FLUX3_MAX_IMAGE_ASPECT * min(width, height):
+        raise ValueError(
+            f"Image aspect ratio is too extreme ({width}x{height}); "
+            f"FLUX 3 accepts at most {_FLUX3_MAX_IMAGE_ASPECT}:1."
+        )
+
+
+def _flux3_collect_images(images: dict | None, field_name: str) -> list[torch.Tensor]:
+    """Flatten Autogrow slots (each possibly batched) into single images and validate them."""
+    flat: list[torch.Tensor] = []
+    for tensor in (images or {}).values():
+        if tensor is None:
+            continue
+        if tensor.ndim == 4:
+            flat.extend(tensor[i] for i in range(tensor.shape[0]))
+        else:
+            flat.append(tensor)
+    if len(flat) > _FLUX3_MAX_IMAGES:
+        raise ValueError(f"FLUX 3 supports at most {_FLUX3_MAX_IMAGES} {field_name}, got {len(flat)}.")
+    for tensor in flat:
+        _flux3_validate_image(tensor)
+    return flat
+
+
+def _flux3_parse_times(value: str, image_count: int, duration: int | str) -> list[float]:
+    """Parse one keyframe time in seconds per image: increasing, inside the clip."""
+    parts = [part.strip() for part in value.split(",") if part.strip()]
+    if len(parts) != image_count:
+        raise ValueError(
+            f"Give one time per keyframe image: got {len(parts)} time(s) for {image_count} image(s)."
+        )
+    try:
+        times = [float(part) for part in parts]
+    except ValueError as exc:
+        raise ValueError(f"Keyframe times must be numbers in seconds, comma-separated; got '{value}'.") from exc
+    if not all(math.isfinite(time) for time in times):
+        raise ValueError(f"Keyframe times must be finite numbers in seconds; got '{value}'.")
+    if any(later <= earlier for earlier, later in zip(times, times[1:])):
+        raise ValueError(f"Keyframe times must increase; got {times}.")
+    if times[0] < 0:
+        raise ValueError(f"Keyframe times cannot be negative; got {times[0]}.")
+    cap = _FLUX3_MAX_DURATION if duration == "auto" else int(duration)
+    if times[-1] > cap:
+        raise ValueError(f"Keyframe time {times[-1]}s is past the end of a {cap}s clip.")
+    return times
+
+
+class Flux3VideoNodeBase(IO.ComfyNode):
+    """Shared widgets, request plumbing and polling for the FLUX 3 generation modes."""
+
+    RATE_HD: float
+    RATE_FHD: float
+
+    @classmethod
+    def common_inputs(cls) -> list:
+        return [
+            IO.Combo.Input(
+                "aspect_ratio",
+                options=_FLUX3_ASPECT_RATIOS,
+                default="auto",
+                tooltip="Output aspect ratio. 'auto' picks one from the prompt and inputs.",
+            ),
+            IO.Combo.Input(
+                "duration",
+                options=_FLUX3_DURATIONS,
+                default="auto",
+                tooltip="Clip length in seconds. 'auto' fits the length to the content.",
+            ),
+            IO.Combo.Input(
+                "resolution",
+                options=list(_FLUX3_RESOLUTIONS),
+                default="720p",
+                tooltip="Output resolution.",
+            ),
+            IO.Boolean.Input(
+                "generate_audio",
+                default=True,
+                tooltip="Generate synchronized audio (ambient, speech, effects). "
+                "Off produces a video with no audio track.",
+            ),
+            IO.Int.Input(
+                "safety_tolerance",
+                default=2,
+                min=0,
+                max=4,
+                advanced=True,
+                tooltip="Moderation tolerance, 0 is the strictest. Requests that send images or "
+                "video are capped at 2 whatever you set here.",
+            ),
+            IO.Int.Input(
+                "seed",
+                default=42,
+                min=0,
+                max=0xFFFFFFFF,
+                control_after_generate=True,
+                tooltip="Seed to determine if node should re-run; FLUX 3 picks its own seed, so "
+                "actual results are nondeterministic regardless of this value.",
+            ),
+        ]
+
+    @classmethod
+    def common_fields(
+        cls,
+        prompt: str,
+        aspect_ratio: str,
+        duration: str,
+        resolution: str,
+        generate_audio: bool,
+        safety_tolerance: int,
+    ) -> dict:
+        validate_string(prompt, field_name="prompt", min_length=1)
+        return {
+            "prompt": prompt,
+            "aspect_ratio": aspect_ratio,
+            "duration": duration if duration == "auto" else int(duration),
+            "resolution": _FLUX3_RESOLUTIONS[resolution],
+            "generate_audio": generate_audio,
+            "safety_tolerance": safety_tolerance,
+        }
+
+    @classmethod
+    def price_badge(cls) -> IO.PriceBadge:
+        return IO.PriceBadge(
+            depends_on=IO.PriceBadgeDepends(widgets=["resolution", "duration"]),
+            expr=f"""
+                (
+                  $rate := widgets.resolution = "1080p" ? {cls.RATE_FHD} : {cls.RATE_HD};
+                  $type(widgets.duration) = "string" and widgets.duration != "auto"
+                    ? {{"type":"usd","usd": $rate * $number(widgets.duration)}}
+                    : {{"type":"usd","usd": $rate, "format": {{"suffix": "/second"}}}}
+                )
+                """,
+        )
+
+
+async def _flux3_execute(cls: type[IO.ComfyNode], request: Flux3VideoRequest) -> IO.NodeOutput:
+    initial_response = await sync_op(
+        cls,
+        ApiEndpoint(path="/proxy/bfl/v1/flux-3-video", method="POST"),
+        response_model=BFLFluxProGenerateResponse,
+        data=request,
+    )
+
+    def price_extractor(_r: BaseModel) -> float | None:
+        return None if initial_response.cost is None else initial_response.cost / 100
+
+    response = await poll_op(
+        cls,
+        ApiEndpoint(initial_response.polling_url),
+        response_model=BFLFluxStatusResponse,
+        status_extractor=lambda r: r.status,
+        progress_extractor=lambda r: r.progress,
+        price_extractor=price_extractor,
+        completed_statuses=[BFLStatus.ready],
+        failed_statuses=[
+            BFLStatus.request_moderated,
+            BFLStatus.content_moderated,
+            BFLStatus.error,
+            BFLStatus.task_not_found,
+        ],
+        queued_statuses=[BFLStatus.pending],
+        poll_interval=8.0,
+        # a failed task answers the poll with a retryable-class HTTP 5xx (500 and 503 observed);
+        # a small retry budget surfaces real failures quickly
+        max_retries_per_poll=3,
+    )
+    return IO.NodeOutput(await download_url_to_video_output(response.result["sample"]))
+
+
+class Flux3TextToVideoNode(Flux3VideoNodeBase):
+    RATE_HD = 0.2431
+    RATE_FHD = 0.4147
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="Flux3TextToVideoNode",
+            display_name="Flux 3 Text to Video",
+            category="partner/video/BFL",
+            description="Generates a video with synchronized audio from a text prompt via FLUX 3.",
+            inputs=[
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="What you want, in plain language; the prompt is interpreted and expanded "
+                    "before generation. Describe ambient sound, music and speech separately for layered audio.",
+                ),
+                *cls.common_inputs(),
+            ],
+            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=cls.price_badge(),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        prompt: str,
+        aspect_ratio: str,
+        duration: str,
+        resolution: str,
+        generate_audio: bool,
+        safety_tolerance: int,
+        seed: int,
+    ) -> IO.NodeOutput:
+        request = Flux3TextToVideoRequest(
+            **cls.common_fields(prompt, aspect_ratio, duration, resolution, generate_audio, safety_tolerance)
+        )
+        return await _flux3_execute(cls, request)
+
+
+class Flux3ImageToVideoNode(Flux3VideoNodeBase):
+    RATE_HD = 0.2431
+    RATE_FHD = 0.4147
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="Flux3ImageToVideoNode",
+            display_name="Flux 3 Image to Video",
+            category="partner/video/BFL",
+            description="Animates 1 to 10 images with FLUX 3. Each image becomes a frame of the clip: "
+            "one image opens it, two morph from the first to the second, and more are spread across it "
+            "or pinned to times you choose.",
+            inputs=[
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="How the scene should move and sound; the prompt is interpreted and "
+                    "expanded before generation.",
+                ),
+                IO.Autogrow.Input(
+                    "keyframes",
+                    template=IO.Autogrow.TemplatePrefix(
+                        IO.Image.Input("image", tooltip="Keyframe image."),
+                        prefix="image_",
+                        min=1,
+                        max=_FLUX3_MAX_IMAGES,
+                    ),
+                    tooltip="1 to 10 images, in playback order. Minimum 256x256 pixels each.",
+                ),
+                IO.DynamicCombo.Input(
+                    "placement",
+                    options=[
+                        IO.DynamicCombo.Option("spread across the clip", []),
+                        IO.DynamicCombo.Option(
+                            "at times",
+                            [
+                                IO.String.Input(
+                                    "times",
+                                    default="0",
+                                    tooltip="One time in seconds per image, comma-separated and "
+                                    "increasing, e.g. '0, 2.5, 5'.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="'spread across the clip' lets FLUX 3 place the images (one opens the clip, "
+                    "two become its start and end); 'at times' pins every image to a second you choose.",
+                ),
+                *cls.common_inputs(),
+            ],
+            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=cls.price_badge(),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        prompt: str,
+        keyframes: IO.Autogrow.Type,
+        placement: dict,
+        aspect_ratio: str,
+        duration: str,
+        resolution: str,
+        generate_audio: bool,
+        safety_tolerance: int,
+        seed: int,
+    ) -> IO.NodeOutput:
+        fields = cls.common_fields(prompt, aspect_ratio, duration, resolution, generate_audio, safety_tolerance)
+        images = _flux3_collect_images(keyframes, "keyframes")
+        if not images:
+            raise ValueError("Connect at least one keyframe image.")
+        times = None
+        if placement["placement"] == "at times":
+            times = _flux3_parse_times(placement["times"], len(images), fields["duration"])
+        elif len(images) >= 3 and fields["duration"] == "auto":
+            # spread images land evenly between the first and last, which needs a known length
+            raise ValueError(
+                f"Spreading {len(images)} images across the clip needs an explicit duration: "
+                "set duration, or place the images yourself with 'at times'."
+            )
+        urls = await upload_images_to_comfyapi(
+            cls, images, max_images=_FLUX3_MAX_IMAGES, wait_label="Uploading keyframes"
+        )
+        request = Flux3ImageToVideoRequest(
+            keyframes=list(zip(times, urls)) if times is not None else urls,
+            **fields,
+        )
+        return await _flux3_execute(cls, request)
+
+
+class Flux3VideoContinuationNode(Flux3VideoNodeBase):
+    RATE_HD = 0.5863
+    RATE_FHD = 0.7579
+
+    @classmethod
+    def define_schema(cls) -> IO.Schema:
+        return IO.Schema(
+            node_id="Flux3VideoContinuationNode",
+            display_name="Flux 3 Video Continuation",
+            category="partner/video/BFL",
+            description="Continues a video with FLUX 3: the new clip carries on from the final frames "
+            "of the one you provide.",
+            inputs=[
+                IO.Video.Input("video", tooltip="The clip to continue."),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="What the continuation should show; the prompt is interpreted and expanded "
+                    "before generation.",
+                ),
+                *cls.common_inputs(),
+            ],
+            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=cls.price_badge(),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        video: Input.Video,
+        prompt: str,
+        aspect_ratio: str,
+        duration: str,
+        resolution: str,
+        generate_audio: bool,
+        safety_tolerance: int,
+        seed: int,
+    ) -> IO.NodeOutput:
+        fields = cls.common_fields(prompt, aspect_ratio, duration, resolution, generate_audio, safety_tolerance)
+        url = await upload_video_to_comfyapi(cls, video, wait_label="Uploading source video")
+        request = Flux3VideoContinuationRequest(start_video=url, **fields)
+        return await _flux3_execute(cls, request)
+
+
 class BFLExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@@ -1021,6 +1409,9 @@ class BFLExtension(ComfyExtension):
             Flux2ProImageNode,
             Flux2MaxImageNode,
             Flux2ImageNode,
+            Flux3TextToVideoNode,
+            Flux3ImageToVideoNode,
+            Flux3VideoContinuationNode,
         ]
 
 

comfy_api_nodes/nodes_bria.py

@@ -6,7 +6,13 @@ from typing_extensions import override
 from comfy_api.latest import IO, ComfyExtension, Input
 from comfy_api_nodes.apis.bria import (
     BriaEditImageRequest,
+    BriaEraseRequest,
+    BriaExpandRequest,
+    BriaExpandResponse,
+    BriaGenFillRequest,
     BriaImageEditResponse,
+    BriaImageResultResponse,
+    BriaIncreaseResolutionRequest,
     BriaRemoveBackgroundRequest,
     BriaRemoveBackgroundResponse,
     BriaRemoveVideoBackgroundRequest,
@@ -21,13 +27,30 @@ from comfy_api_nodes.util import (
     convert_mask_to_image,
     download_url_to_image_tensor,
     download_url_to_video_output,
+    downscale_image_tensor_by_max_side,
+    get_image_dimensions,
     poll_op,
     sync_op,
     upload_image_to_comfyapi,
     upload_video_to_comfyapi,
+    validate_string,
     validate_video_duration,
 )
 
+BRIA_MAX_OUTPUT_SIDE = 8192
+BRIA_MIN_RATIO = 0.5
+BRIA_MAX_RATIO = 3.0
+BRIA_MIN_SHORT_SIDE = 224
+
+
+def _upscaled_output_side(height: int, width: int, multiplier: int) -> int:
+    prescale = max(1.0, BRIA_MIN_SHORT_SIDE / min(height, width))
+    return round(max(height, width) * prescale * multiplier)
+
+
+def _smallest_output_side(height: int, width: int, multiplier: int) -> int:
+    return round(max(height, width) / min(height, width) * BRIA_MIN_SHORT_SIDE * multiplier)
+
 
 class BriaImageEditNode(IO.ComfyNode):
 
@@ -243,6 +266,503 @@ class BriaRemoveImageBackground(IO.ComfyNode):
         return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
 
 
+def _mask_to_binary_image(mask: Input.Image, action: str) -> torch.Tensor:
+    binary = (mask > 0.5).float()
+    if not binary.any():
+        raise ValueError(
+            f"The mask is empty, so there is nothing to {action}. Masks are binarized at 50%: "
+            f"areas painted at less than half opacity are ignored."
+        )
+    return convert_mask_to_image(binary)
+
+
+def _validate_mask_aspect_ratio(image: Input.Image, mask: Input.Image) -> None:
+    ih, iw = image.shape[1], image.shape[2]
+    mh, mw = mask.shape[-2], mask.shape[-1]
+    if abs(iw * mh - ih * mw) > 0.01 * ih * mw:
+        raise ValueError(f"Mask must have the same aspect ratio as the image: image is {iw}x{ih}, mask is {mw}x{mh}.")
+
+
+class BriaGenFill(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaGenFill",
+            display_name="Bria Generative Fill",
+            category="partner/image/Bria",
+            description="Generate objects or scenery inside a masked region of an image using Bria.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Mask.Input(
+                    "mask",
+                    tooltip="White areas are filled with generated content, black areas are preserved. "
+                    "The mask is binarized before sending, so partially painted areas count as white. "
+                    "Must have the same aspect ratio as the image.",
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Description of what to generate inside the masked region.",
+                ),
+                IO.String.Input("negative_prompt", multiline=True, default=""),
+                IO.Boolean.Input(
+                    "refine_prompt",
+                    default=True,
+                    tooltip="Automatically adjust the prompt for better results; "
+                    "disable to use the prompt exactly as written.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=1,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("prompt_content_moderation", default=False),
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[IO.Image.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.0429}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        mask: Input.Image,
+        prompt: str,
+        negative_prompt: str,
+        refine_prompt: bool,
+        seed: int,
+        moderation: InputModerationSettings,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1)
+        _validate_mask_aspect_ratio(image, mask)
+        mask_image = _mask_to_binary_image(mask, "fill")
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/gen_fill", method="POST"),
+            data=BriaGenFillRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                mask=await upload_image_to_comfyapi(
+                    cls, mask_image, total_pixels=None, wait_label="Uploading mask"
+                ),
+                prompt=prompt,
+                negative_prompt=negative_prompt if negative_prompt else None,
+                refine_prompt=refine_prompt,
+                seed=seed,
+                prompt_content_moderation=moderation.get("prompt_content_moderation", False),
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaImageResultResponse,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
+
+
+class BriaEraser(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaEraser",
+            display_name="Bria Eraser",
+            category="partner/image/Bria",
+            description="Remove objects or areas outlined by a mask from an image using Bria.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Mask.Input(
+                    "mask",
+                    tooltip="White areas are erased, black areas are preserved. "
+                    "The mask is binarized before sending, so partially painted areas count as white. "
+                    "Must have the same aspect ratio as the image.",
+                ),
+                IO.Combo.Input(
+                    "mask_type",
+                    options=["manual", "automatic"],
+                    tooltip="manual for hand-drawn or brush masks, "
+                    "automatic for masks produced by segmentation models such as SAM.",
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[IO.Image.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.0286}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        mask: Input.Image,
+        mask_type: str,
+        moderation: dict,
+    ) -> IO.NodeOutput:
+        _validate_mask_aspect_ratio(image, mask)
+        mask_image = _mask_to_binary_image(mask, "erase")
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/erase", method="POST"),
+            data=BriaEraseRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                mask=await upload_image_to_comfyapi(
+                    cls, mask_image, total_pixels=None, wait_label="Uploading mask"
+                ),
+                mask_type=mask_type,
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaImageResultResponse,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
+
+
+class BriaExpandImage(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaExpandImage",
+            display_name="Bria Expand Image",
+            category="partner/image/Bria",
+            description="Expand an image beyond its borders with generated content using Bria.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.DynamicCombo.Input(
+                    "expand_mode",
+                    options=[
+                        *[IO.DynamicCombo.Option(ratio, []) for ratio in
+                          ["1:1", "2:3", "3:2", "3:4", "4:3", "4:5", "5:4", "9:16", "16:9"]],
+                        IO.DynamicCombo.Option(
+                            "custom_ratio",
+                            [
+                                IO.Int.Input(
+                                    "ratio_width",
+                                    default=21,
+                                    min=1,
+                                    max=100,
+                                    tooltip="Width side of the target ratio: 21 and 9 give 21:9.",
+                                ),
+                                IO.Int.Input(
+                                    "ratio_height",
+                                    default=9,
+                                    min=1,
+                                    max=100,
+                                    tooltip="Height side of the target ratio: 21 and 9 give 21:9. "
+                                    f"Bria only accepts width/height between {BRIA_MIN_RATIO} and "
+                                    f"{BRIA_MAX_RATIO}, so anything taller than 1:2 needs the manual mode.",
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "manual",
+                            [
+                                IO.Int.Input("canvas_width", default=1000, min=64, max=5000),
+                                IO.Int.Input("canvas_height", default=1000, min=64, max=5000),
+                                IO.Int.Input(
+                                    "image_width",
+                                    default=500,
+                                    min=1,
+                                    max=5000,
+                                    tooltip="Width of the original image inside the canvas.",
+                                ),
+                                IO.Int.Input(
+                                    "image_height",
+                                    default=500,
+                                    min=1,
+                                    max=5000,
+                                    tooltip="Height of the original image inside the canvas.",
+                                ),
+                                IO.Int.Input(
+                                    "image_x",
+                                    default=250,
+                                    min=-5000,
+                                    max=5000,
+                                    tooltip="X position of the image's top-left corner inside the canvas; "
+                                    "may fall outside the canvas, cropping the image.",
+                                ),
+                                IO.Int.Input(
+                                    "image_y",
+                                    default=250,
+                                    min=-5000,
+                                    max=5000,
+                                    tooltip="Y position of the image's top-left corner inside the canvas; "
+                                    "may fall outside the canvas, cropping the image.",
+                                ),
+                            ],
+                        ),
+                    ],
+                    tooltip="Target shape of the expanded image: a preset aspect ratio, a custom ratio, "
+                    "or manual placement of the original image on a canvas. "
+                    "Manual is the only mode that can reach a canvas taller than 1:2.",
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                    tooltip="Optional description of the expanded scene; "
+                    "when empty, Bria generates one from the image.",
+                ),
+                IO.String.Input("negative_prompt", multiline=True, default=""),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=1,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("prompt_content_moderation", default=False),
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[
+                IO.Image.Output(),
+                IO.String.Output(display_name="prompt", tooltip="The prompt used for the expansion; "
+                                 "auto-generated by Bria when the prompt input is empty."),
+            ],
+            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.0286}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        expand_mode: dict,
+        prompt: str,
+        negative_prompt: str,
+        seed: int,
+        moderation: InputModerationSettings,
+    ) -> IO.NodeOutput:
+        mode = expand_mode["expand_mode"]
+        aspect_ratio = canvas_size = original_image_size = original_image_location = None
+        if mode == "manual":
+            canvas_size = [expand_mode["canvas_width"], expand_mode["canvas_height"]]
+            original_image_size = [expand_mode["image_width"], expand_mode["image_height"]]
+            original_image_location = [expand_mode["image_x"], expand_mode["image_y"]]
+        elif mode == "custom_ratio":
+            ratio_width, ratio_height = expand_mode["ratio_width"], expand_mode["ratio_height"]
+            aspect_ratio = ratio_width / ratio_height
+            if not BRIA_MIN_RATIO <= aspect_ratio <= BRIA_MAX_RATIO:
+                raise ValueError(
+                    f"Bria accepts a width-to-height ratio between {BRIA_MIN_RATIO} and {BRIA_MAX_RATIO}: "
+                    f"{ratio_width}:{ratio_height} is {aspect_ratio:.4f}. "
+                    f"Use the manual expand mode to reach a canvas of any shape."
+                )
+        else:
+            aspect_ratio = mode
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/expand", method="POST"),
+            data=BriaExpandRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                aspect_ratio=aspect_ratio,
+                canvas_size=canvas_size,
+                original_image_size=original_image_size,
+                original_image_location=original_image_location,
+                prompt=prompt if prompt else None,
+                negative_prompt=negative_prompt if negative_prompt else None,
+                seed=seed,
+                prompt_content_moderation=moderation.get("prompt_content_moderation", False),
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaExpandResponse,
+        )
+        return IO.NodeOutput(
+            await download_url_to_image_tensor(response.result.image_url),
+            response.result.prompt or "",
+        )
+
+
+class BriaIncreaseResolution(IO.ComfyNode):
+
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="BriaIncreaseResolution",
+            display_name="Bria Increase Resolution",
+            category="partner/image/Bria",
+            description="Upscale an image by 2x or 4x using Bria, preserving the original content.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.Combo.Input(
+                    "desired_increase",
+                    options=["2", "4"],
+                    tooltip="Resolution multiplier. The output must fit within 8192 pixels on each side.",
+                ),
+                IO.Boolean.Input(
+                    "auto_downscale",
+                    default=False,
+                    tooltip="Automatically lower the multiplier, and downscale the input image if that is "
+                    "still not enough, when the output would exceed the limit.",
+                ),
+                IO.DynamicCombo.Input(
+                    "moderation",
+                    options=[
+                        IO.DynamicCombo.Option("false", []),
+                        IO.DynamicCombo.Option(
+                            "true",
+                            [
+                                IO.Boolean.Input("visual_input_moderation", default=False),
+                                IO.Boolean.Input("visual_output_moderation", default=False),
+                            ],
+                        ),
+                    ],
+                    tooltip="Moderation settings",
+                ),
+            ],
+            outputs=[IO.Image.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.0286}""",
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        desired_increase: str,
+        auto_downscale: bool,
+        moderation: dict,
+    ) -> IO.NodeOutput:
+        multiplier = int(desired_increase)
+        height, width = get_image_dimensions(image)
+        if _upscaled_output_side(height, width, multiplier) > BRIA_MAX_OUTPUT_SIDE:
+            candidates = [c for c in (4, 2) if c <= multiplier]
+            if not auto_downscale:
+                predicted = _upscaled_output_side(height, width, multiplier)
+                raise ValueError(
+                    f"Bria can upscale up to a maximum output dimension of {BRIA_MAX_OUTPUT_SIDE} pixels: "
+                    f"input is {width}x{height}, x{multiplier} would be {predicted} pixels on the long side. "
+                    f"Enable auto_downscale, or use a smaller input image or a lower multiplier."
+                )
+            fitted = next(
+                (c for c in candidates if _upscaled_output_side(height, width, c) <= BRIA_MAX_OUTPUT_SIDE), None
+            )
+            if fitted is not None:
+                multiplier = fitted
+            else:
+                shrinkable = next((c for c in sorted(candidates) if _smallest_output_side(height, width, c)
+                                   <= BRIA_MAX_OUTPUT_SIDE), None)
+                if shrinkable is None:
+                    raise ValueError(
+                        f"This image cannot be upscaled by Bria at any multiplier: it is {width}x{height}, and "
+                        f"Bria first enlarges the short side to {BRIA_MIN_SHORT_SIDE} pixels, which pushes the "
+                        f"long side past the {BRIA_MAX_OUTPUT_SIDE} pixel limit. Crop it to a squarer shape first."
+                    )
+                multiplier = shrinkable
+                image = downscale_image_tensor_by_max_side(image, max_side=BRIA_MAX_OUTPUT_SIDE // multiplier)
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/bria/v2/image/edit/increase_resolution", method="POST"),
+            data=BriaIncreaseResolutionRequest(
+                image=await upload_image_to_comfyapi(cls, image, total_pixels=None, wait_label="Uploading image"),
+                desired_increase=multiplier,
+                visual_input_content_moderation=moderation.get("visual_input_moderation", False),
+                visual_output_content_moderation=moderation.get("visual_output_moderation", False),
+            ),
+            response_model=BriaStatusResponse,
+        )
+        response = await poll_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/bria/v2/status/{response.request_id}"),
+            status_extractor=lambda r: r.status,
+            response_model=BriaImageResultResponse,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(response.result.image_url))
+
+
 class BriaRemoveVideoBackground(IO.ComfyNode):
 
     @classmethod
@@ -572,6 +1092,10 @@ class BriaExtension(ComfyExtension):
         return [
             BriaImageEditNode,
             BriaRemoveImageBackground,
+            BriaGenFill,
+            BriaEraser,
+            BriaExpandImage,
+            BriaIncreaseResolution,
             BriaRemoveVideoBackground,
             BriaVideoGreenScreen,
             BriaVideoReplaceBackground,

comfy_api_nodes/nodes_gemini.py

@@ -43,6 +43,7 @@ from comfy_api_nodes.util import (
     download_url_to_image_tensor,
     download_url_to_video_output,
     get_number_of_images,
+    pad_images_to_common_channels,
     sync_op,
     tensor_to_base64_string,
     upload_audio_to_comfyapi,
@@ -76,8 +77,8 @@ GEMINI_IMAGE_2_PRICE_BADGE = IO.PriceBadge(
       $m := widgets.model;
       $r := widgets.resolution;
       $isFlash := $contains($m, "nano banana 2");
-      $flashPrices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
-      $proPrices := {"1k": 0.134, "2k": 0.134, "4k": 0.24};
+      $flashPrices := {"1k": 0.0835, "2k": 0.1217, "4k": 0.1848};
+      $proPrices := {"1k": 0.1608, "2k": 0.1608, "4k": 0.288};
       $prices := $isFlash ? $flashPrices : $proPrices;
       {"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
     )
@@ -233,8 +234,8 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
                 "Try rephrasing your prompt or changing the response modality to 'IMAGE+TEXT' "
                 "to see the model's reasoning."
             )
-        return torch.zeros((1, 1024, 1024, 4))
-    return torch.cat(image_tensors, dim=0)
+        return torch.zeros((1, 1024, 1024, 3))
+    return torch.cat(pad_images_to_common_channels(image_tensors), dim=0)
 
 
 def get_text_from_interaction(interaction: GeminiInteraction) -> str:
@@ -1464,10 +1465,10 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
                 expr="""
                 (
                   $contains(widgets.model, "lite")
-                    ? {"type":"usd","usd": 0.034, "format":{"suffix":"/Image","approximate":true}}
+                    ? {"type":"usd","usd": 0.0408, "format":{"suffix":"/Image","approximate":true}}
                     : (
                         $r := $lookup(widgets, "model.resolution");
-                        $prices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
+                        $prices := {"1k": 0.0835, "2k": 0.1217, "4k": 0.1848};
                         {"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
                       )
                 )

comfy_api_nodes/nodes_grok.py

@@ -36,6 +36,26 @@ _GROK_VIDEO_MODEL_API_IDS = {
     "grok-imagine-video-1.5": "grok-imagine-video-1.5",
 }
 
+_GROK_IMAGE_MODEL_API_IDS = {
+    "grok-imagine-image-2.0": "grok-imagine-image-2.0",
+}
+
+_GROK_IMAGE_QUALITY_MODELS = {"grok-imagine-image-2.0"}
+
+_GROK_IMAGE_QUALITY_OPTIONS = ["medium", "low"]
+
+_GROK_IMAGE_EDIT_MAX_IMAGES = {
+    "grok-imagine-image-2.0": 3,
+    "grok-imagine-image-pro": 1,
+    "grok-imagine-image-quality": 3,
+    "grok-imagine-image": 3,
+}
+
+_GROK_IMAGE_EDIT_ASPECT_RATIO_NEEDS_MULTIPLE = {
+    "grok-imagine-image-quality",
+    "grok-imagine-image",
+}
+
 _GROK_VOICE_OPTIONS = [
     "none",
     "ara",
@@ -132,6 +152,7 @@ class GrokImageNode(IO.ComfyNode):
                 IO.Combo.Input(
                     "model",
                     options=[
+                        "grok-imagine-image-2.0",
                         "grok-imagine-image-quality",
                         "grok-imagine-image-pro",
                         "grok-imagine-image",
@@ -181,6 +202,12 @@ class GrokImageNode(IO.ComfyNode):
                     "actual results are nondeterministic regardless of seed.",
                 ),
                 IO.Combo.Input("resolution", options=["1K", "2K"], optional=True),
+                IO.Combo.Input(
+                    "quality",
+                    options=_GROK_IMAGE_QUALITY_OPTIONS,
+                    optional=True,
+                    tooltip="Quality level, supported only by the grok-imagine-image-2.0 model.",
+                ),
             ],
             outputs=[
                 IO.Image.Output(),
@@ -192,12 +219,15 @@ class GrokImageNode(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model", "number_of_images", "resolution", "quality"]),
                 expr="""
                 (
-                  $rate := widgets.model = "grok-imagine-image-quality"
-                    ? (widgets.resolution = "1k" ? 0.05 : 0.07)
-                    : ($contains(widgets.model, "pro") ? 0.07 : 0.02);
+                  $is1k := widgets.resolution = "1k";
+                  $rate := widgets.model = "grok-imagine-image-2.0"
+                    ? (widgets.quality = "low" ? ($is1k ? 0.04 : 0.06) : ($is1k ? 0.06 : 0.08))
+                    : (widgets.model = "grok-imagine-image-quality"
+                        ? ($is1k ? 0.05 : 0.07)
+                        : ($contains(widgets.model, "pro") ? 0.07 : 0.02));
                   {"type":"usd","usd": $rate * widgets.number_of_images}
                 )
                 """,
@@ -213,18 +243,20 @@ class GrokImageNode(IO.ComfyNode):
         number_of_images: int,
         seed: int,
         resolution: str = "1K",
+        quality: str = "medium",
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         response = await sync_op(
             cls,
             ApiEndpoint(path="/proxy/xai/v1/images/generations", method="POST"),
             data=ImageGenerationRequest(
-                model=model,
+                model=_GROK_IMAGE_MODEL_API_IDS.get(model, model),
                 prompt=prompt,
                 aspect_ratio=aspect_ratio,
                 n=number_of_images,
                 seed=seed,
                 resolution=resolution.lower(),
+                quality=quality if model in _GROK_IMAGE_QUALITY_MODELS else None,
             ),
             response_model=ImageGenerationResponse,
         )
@@ -255,7 +287,9 @@ _GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS = [
 ]
 
 
-def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: bool):
+def _grok_image_edit_model_inputs(
+    *, max_ref_images: int, with_aspect_ratio: bool, with_quality: bool = False, aspect_ratio_needs_multiple: bool = True
+):
     inputs = [
         IO.Autogrow.Input(
             "images",
@@ -281,12 +315,18 @@ def _grok_image_edit_model_inputs(*, max_ref_images: int, with_aspect_ratio: boo
             display_mode=IO.NumberDisplay.number,
         ),
     ]
+    if with_quality:
+        inputs.append(IO.Combo.Input("quality", options=_GROK_IMAGE_QUALITY_OPTIONS))
     if with_aspect_ratio:
         inputs.append(
             IO.Combo.Input(
                 "aspect_ratio",
                 options=_GROK_IMAGE_EDIT_ASPECT_RATIO_OPTIONS,
-                tooltip="Only allowed when multiple images are connected.",
+                tooltip=(
+                    "Only allowed when multiple images are connected."
+                    if aspect_ratio_needs_multiple
+                    else "Aspect ratio of the edited image."
+                ),
             )
         )
     return inputs
@@ -451,6 +491,15 @@ class GrokImageEditNodeV2(IO.ComfyNode):
                 IO.DynamicCombo.Input(
                     "model",
                     options=[
+                        IO.DynamicCombo.Option(
+                            "grok-imagine-image-2.0",
+                            _grok_image_edit_model_inputs(
+                                max_ref_images=3,
+                                with_aspect_ratio=True,
+                                with_quality=True,
+                                aspect_ratio_needs_multiple=False,
+                            ),
+                        ),
                         IO.DynamicCombo.Option(
                             "grok-imagine-image-quality",
                             _grok_image_edit_model_inputs(max_ref_images=3, with_aspect_ratio=True),
@@ -488,18 +537,23 @@ class GrokImageEditNodeV2(IO.ComfyNode):
             is_api_node=True,
             price_badge=IO.PriceBadge(
                 depends_on=IO.PriceBadgeDepends(
-                    widgets=["model", "model.resolution", "model.number_of_images"],
+                    widgets=["model", "model.resolution", "model.number_of_images", "model.quality"],
                 ),
                 expr="""
                 (
-                  $isQualityModel := widgets.model = "grok-imagine-image-quality";
+                  $is20 := widgets.model = "grok-imagine-image-2.0";
                   $isPro := $contains(widgets.model, "pro");
                   $res := $lookup(widgets, "model.resolution");
                   $n := $lookup(widgets, "model.number_of_images");
-                  $rate := $isQualityModel
-                    ? ($res = "1k" ? 0.05 : 0.07)
-                    : ($isPro ? 0.07 : 0.02);
-                  $base := $isQualityModel ? 0.01 : 0.002;
+                  $is1k := $res = "1k";
+                  $rate := $is20
+                    ? ($lookup(widgets, "model.quality") = "low"
+                        ? ($is1k ? 0.04 : 0.06)
+                        : ($is1k ? 0.06 : 0.08))
+                    : (widgets.model = "grok-imagine-image-quality"
+                        ? ($is1k ? 0.05 : 0.07)
+                        : ($isPro ? 0.07 : 0.02));
+                  $base := ($is20 or widgets.model = "grok-imagine-image-quality") ? 0.01 : 0.002;
                   $output := $rate * $n;
                   $isPro
                     ? {"type":"usd","usd": $base + $output}
@@ -525,13 +579,15 @@ class GrokImageEditNodeV2(IO.ComfyNode):
 
         image_tensors: list[Input.Image] = [t for t in images_dict.values() if t is not None]
         n_images = sum(get_number_of_images(t) for t in image_tensors)
+        max_images = _GROK_IMAGE_EDIT_MAX_IMAGES.get(model_id, 3)
         if n_images < 1:
             raise ValueError("At least one image is required for editing.")
-        if model_id == "grok-imagine-image-pro" and n_images > 1:
-            raise ValueError("The pro model supports only 1 input image.")
-        if model_id != "grok-imagine-image-pro" and n_images > 3:
-            raise ValueError("A maximum of 3 input images is supported.")
-        if aspect_ratio != "auto" and n_images == 1:
+        if n_images > max_images:
+            raise ValueError(
+                f"The {model_id} model supports at most {max_images} input "
+                f"image{'s' if max_images > 1 else ''}; {n_images} are connected."
+            )
+        if aspect_ratio != "auto" and model_id in _GROK_IMAGE_EDIT_ASPECT_RATIO_NEEDS_MULTIPLE and n_images == 1:
             raise ValueError(
                 "Custom aspect ratio is only allowed when multiple images are connected to the image input."
             )
@@ -547,7 +603,7 @@ class GrokImageEditNodeV2(IO.ComfyNode):
             cls,
             ApiEndpoint(path="/proxy/xai/v1/images/edits", method="POST"),
             data=ImageEditRequest(
-                model=model_id,
+                model=_GROK_IMAGE_MODEL_API_IDS.get(model_id, model_id),
                 images=[
                     InputUrlObject(url=f"data:image/png;base64,{tensor_to_base64_string(i)}") for i in flat_tensors
                 ],
@@ -556,6 +612,7 @@ class GrokImageEditNodeV2(IO.ComfyNode):
                 n=number_of_images,
                 seed=seed,
                 aspect_ratio=None if aspect_ratio == "auto" else aspect_ratio,
+                quality=model.get("quality") if model_id in _GROK_IMAGE_QUALITY_MODELS else None,
             ),
             response_model=ImageGenerationResponse,
         )

comfy_api_nodes/nodes_kling.py

@@ -5,17 +5,13 @@ For source of truth on the allowed permutations of request fields, please refere
 """
 
 import logging
-import math
 import re
 
 import torch
 from typing_extensions import override
 
-from comfy_api.latest import IO, ComfyExtension, Input, InputImpl
+from comfy_api.latest import IO, ComfyExtension, Input
 from comfy_api_nodes.apis import (
-    KlingCameraControl,
-    KlingCameraConfig,
-    KlingCameraControlType,
     KlingVideoGenDuration,
     KlingVideoGenMode,
     KlingVideoGenAspectRatio,
@@ -30,23 +26,12 @@ from comfy_api_nodes.apis import (
     KlingLipSyncInputObject,
     KlingLipSyncRequest,
     KlingLipSyncResponse,
-    KlingVirtualTryOnModelName,
-    KlingVirtualTryOnRequest,
-    KlingVirtualTryOnResponse,
     KlingVideoResult,
     KlingImageResult,
     KlingImageGenerationsRequest,
     KlingImageGenerationsResponse,
     KlingImageGenImageReferenceType,
     KlingImageGenAspectRatio,
-    KlingVideoEffectsRequest,
-    KlingVideoEffectsResponse,
-    KlingDualCharacterEffectsScene,
-    KlingSingleImageEffectsScene,
-    KlingDualCharacterEffectInput,
-    KlingSingleImageEffectInput,
-    KlingCharacterEffectModelName,
-    KlingSingleImageEffectModelName,
 )
 from comfy_api_nodes.apis.kling import (
     ImageToVideoWithAudioRequest,
@@ -119,9 +104,6 @@ PATH_TEXT_TO_VIDEO = f"/proxy/kling/{KLING_API_VERSION}/videos/text2video"
 PATH_IMAGE_TO_VIDEO = f"/proxy/kling/{KLING_API_VERSION}/videos/image2video"
 PATH_VIDEO_EXTEND = f"/proxy/kling/{KLING_API_VERSION}/videos/video-extend"
 PATH_LIP_SYNC = f"/proxy/kling/{KLING_API_VERSION}/videos/lip-sync"
-PATH_VIDEO_EFFECTS = f"/proxy/kling/{KLING_API_VERSION}/videos/effects"
-PATH_CHARACTER_IMAGE = f"/proxy/kling/{KLING_API_VERSION}/images/generations"
-PATH_VIRTUAL_TRY_ON = f"/proxy/kling/{KLING_API_VERSION}/images/kolors-virtual-try-on"
 PATH_IMAGE_GENERATIONS = f"/proxy/kling/{KLING_API_VERSION}/images/generations"
 
 MAX_PROMPT_LENGTH_T2V = 2500
@@ -133,21 +115,11 @@ MAX_PROMPT_LENGTH_LIP_SYNC = 120
 AVERAGE_DURATION_T2V = 319
 AVERAGE_DURATION_I2V = 164
 AVERAGE_DURATION_LIP_SYNC = 455
-AVERAGE_DURATION_VIRTUAL_TRY_ON = 19
 AVERAGE_DURATION_IMAGE_GEN = 32
-AVERAGE_DURATION_VIDEO_EFFECTS = 320
 AVERAGE_DURATION_VIDEO_EXTEND = 320
 
 
 MODE_TEXT2VIDEO = {
-    "standard mode / 5s duration / kling-v1-6": ("std", "5", "kling-v1-6"),
-    "standard mode / 10s duration / kling-v1-6": ("std", "10", "kling-v1-6"),
-    "pro mode / 5s duration / kling-v2-master": ("pro", "5", "kling-v2-master"),
-    "pro mode / 10s duration / kling-v2-master": ("pro", "10", "kling-v2-master"),
-    "standard mode / 5s duration / kling-v2-master": ("std", "5", "kling-v2-master"),
-    "standard mode / 10s duration / kling-v2-master": ("std", "10", "kling-v2-master"),
-    "pro mode / 5s duration / kling-v2-1-master": ("pro", "5", "kling-v2-1-master"),
-    "pro mode / 10s duration / kling-v2-1-master": ("pro", "10", "kling-v2-1-master"),
     "pro mode / 5s duration / kling-v2-5-turbo": ("pro", "5", "kling-v2-5-turbo"),
     "pro mode / 10s duration / kling-v2-5-turbo": ("pro", "10", "kling-v2-5-turbo"),
 }
@@ -160,12 +132,6 @@ See: [Kling API Docs Capability Map](https://app.klingai.com/global/dev/document
 
 
 MODE_START_END_FRAME = {
-    "pro mode / 5s duration / kling-v1-5": ("pro", "5", "kling-v1-5"),
-    "pro mode / 10s duration / kling-v1-5": ("pro", "10", "kling-v1-5"),
-    "pro mode / 5s duration / kling-v1-6": ("pro", "5", "kling-v1-6"),
-    "pro mode / 10s duration / kling-v1-6": ("pro", "10", "kling-v1-6"),
-    "pro mode / 5s duration / kling-v2-1": ("pro", "5", "kling-v2-1"),
-    "pro mode / 10s duration / kling-v2-1": ("pro", "10", "kling-v2-1"),
     "pro mode / 5s duration / kling-v2-5-turbo": ("pro", "5", "kling-v2-5-turbo"),
     "pro mode / 10s duration / kling-v2-5-turbo": ("pro", "10", "kling-v2-5-turbo"),
 }
@@ -287,11 +253,6 @@ async def finish_omni_video_task(cls: type[IO.ComfyNode], response: TaskStatusRe
     return IO.NodeOutput(await download_url_to_video_output(final_response.data.task_result.videos[0].url))
 
 
-def is_valid_camera_control_configs(configs: list[float]) -> bool:
-    """Verifies that at least one camera control configuration is non-zero."""
-    return any(not math.isclose(value, 0.0) for value in configs)
-
-
 def is_valid_task_creation_response(response: KlingText2VideoResponse) -> bool:
     """Verifies that the initial response contains a task ID."""
     return bool(response.data.task_id)
@@ -307,7 +268,7 @@ def is_valid_video_response(response: KlingText2VideoResponse) -> bool:
     )
 
 
-def is_valid_image_response(response: KlingVirtualTryOnResponse) -> bool:
+def is_valid_image_response(response: KlingImageGenerationsResponse) -> bool:
     """Verifies that the response contains a task result with at least one image."""
     return (
         response.data is not None
@@ -430,7 +391,6 @@ async def execute_text2video(
     model_mode: str,
     duration: str,
     aspect_ratio: str,
-    camera_control: KlingCameraControl | None = None,
 ) -> IO.NodeOutput:
     validate_prompts(prompt, negative_prompt, MAX_PROMPT_LENGTH_T2V)
     task_creation_response = await sync_op(
@@ -445,7 +405,6 @@ async def execute_text2video(
             model_name=model_name,
             cfg_scale=cfg_scale,
             aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio),
-            camera_control=camera_control,
         ),
     )
 
@@ -475,19 +434,11 @@ async def execute_image2video(
     model_mode: str,
     aspect_ratio: str,
     duration: str,
-    camera_control: KlingCameraControl | None = None,
     end_frame: torch.Tensor | None = None,
 ) -> IO.NodeOutput:
     validate_prompts(prompt, negative_prompt, MAX_PROMPT_LENGTH_I2V)
     validate_input_image(start_frame)
 
-    if camera_control is not None:
-        # Camera control type for image 2 video is always `simple`
-        camera_control.type = KlingCameraControlType.simple
-
-    if model_mode == "std" and model_name == KlingVideoGenModelName.kling_v2_5_turbo.value:
-        model_mode = "pro"  # October 5: currently "std" mode is not supported for this model
-
     task_creation_response = await sync_op(
         cls,
         ApiEndpoint(path=PATH_IMAGE_TO_VIDEO, method="POST"),
@@ -505,7 +456,6 @@ async def execute_image2video(
             cfg_scale=cfg_scale,
             mode=KlingVideoGenMode(model_mode),
             duration=KlingVideoGenDuration(duration),
-            camera_control=camera_control,
         ),
     )
 
@@ -525,59 +475,6 @@ async def execute_image2video(
     return IO.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
 
 
-async def execute_video_effect(
-    cls: type[IO.ComfyNode],
-    dual_character: bool,
-    effect_scene: KlingDualCharacterEffectsScene | KlingSingleImageEffectsScene,
-    model_name: str,
-    duration: KlingVideoGenDuration,
-    image_1: torch.Tensor,
-    image_2: torch.Tensor | None = None,
-    model_mode: KlingVideoGenMode | None = None,
-) -> tuple[InputImpl.VideoFromFile, str, str]:
-    if dual_character:
-        request_input_field = KlingDualCharacterEffectInput(
-            model_name=model_name,
-            mode=model_mode,
-            images=[
-                tensor_to_base64_string(image_1),
-                tensor_to_base64_string(image_2),
-            ],
-            duration=duration,
-        )
-    else:
-        request_input_field = KlingSingleImageEffectInput(
-            model_name=model_name,
-            image=tensor_to_base64_string(image_1),
-            duration=duration,
-        )
-
-    task_creation_response = await sync_op(
-        cls,
-        endpoint=ApiEndpoint(path=PATH_VIDEO_EFFECTS, method="POST"),
-        response_model=KlingVideoEffectsResponse,
-        data=KlingVideoEffectsRequest(
-            effect_scene=effect_scene,
-            input=request_input_field,
-        ),
-    )
-
-    validate_task_creation_response(task_creation_response)
-    task_id = task_creation_response.data.task_id
-
-    final_response = await poll_op(
-        cls,
-        ApiEndpoint(path=f"{PATH_VIDEO_EFFECTS}/{task_id}"),
-        response_model=KlingVideoEffectsResponse,
-        estimated_duration=AVERAGE_DURATION_VIDEO_EFFECTS,
-        status_extractor=lambda r: (r.data.task_status.value if r.data and r.data.task_status else None),
-    )
-    validate_video_result_response(final_response)
-
-    video = get_video_from_response(final_response)
-    return await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration)
-
-
 async def execute_lipsync(
     cls: type[IO.ComfyNode],
     video: Input.Video,
@@ -640,125 +537,6 @@ async def execute_lipsync(
     return IO.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
 
 
-class KlingCameraControls(IO.ComfyNode):
-    """Kling Camera Controls Node"""
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingCameraControls",
-            display_name="Kling Camera Controls",
-            category="partner/video/Kling",
-            description="Allows specifying configuration options for Kling Camera Controls and motion control effects.",
-            inputs=[
-                IO.Combo.Input("camera_control_type", options=KlingCameraControlType),
-                IO.Float.Input(
-                    "horizontal_movement",
-                    default=0.0,
-                    min=-10.0,
-                    max=10.0,
-                    step=0.25,
-                    display_mode=IO.NumberDisplay.slider,
-                    tooltip="Controls camera's movement along horizontal axis (x-axis). Negative indicates left, positive indicates right",
-                ),
-                IO.Float.Input(
-                    "vertical_movement",
-                    default=0.0,
-                    min=-10.0,
-                    max=10.0,
-                    step=0.25,
-                    display_mode=IO.NumberDisplay.slider,
-                    tooltip="Controls camera's movement along vertical axis (y-axis). Negative indicates downward, positive indicates upward.",
-                ),
-                IO.Float.Input(
-                    "pan",
-                    default=0.5,
-                    min=-10.0,
-                    max=10.0,
-                    step=0.25,
-                    display_mode=IO.NumberDisplay.slider,
-                    tooltip="Controls camera's rotation in vertical plane (x-axis). Negative indicates downward rotation, positive indicates upward rotation.",
-                ),
-                IO.Float.Input(
-                    "tilt",
-                    default=0.0,
-                    min=-10.0,
-                    max=10.0,
-                    step=0.25,
-                    display_mode=IO.NumberDisplay.slider,
-                    tooltip="Controls camera's rotation in horizontal plane (y-axis). Negative indicates left rotation, positive indicates right rotation.",
-                ),
-                IO.Float.Input(
-                    "roll",
-                    default=0.0,
-                    min=-10.0,
-                    max=10.0,
-                    step=0.25,
-                    display_mode=IO.NumberDisplay.slider,
-                    tooltip="Controls camera's rolling amount (z-axis). Negative indicates counterclockwise, positive indicates clockwise.",
-                ),
-                IO.Float.Input(
-                    "zoom",
-                    default=0.0,
-                    min=-10.0,
-                    max=10.0,
-                    step=0.25,
-                    display_mode=IO.NumberDisplay.slider,
-                    tooltip="Controls change in camera's focal length. Negative indicates narrower field of view, positive indicates wider field of view.",
-                ),
-            ],
-            outputs=[IO.Custom("CAMERA_CONTROL").Output(display_name="camera_control")],
-        )
-
-    @classmethod
-    def validate_inputs(
-        cls,
-        horizontal_movement: float,
-        vertical_movement: float,
-        pan: float,
-        tilt: float,
-        roll: float,
-        zoom: float,
-    ) -> bool | str:
-        if not is_valid_camera_control_configs(
-            [
-                horizontal_movement,
-                vertical_movement,
-                pan,
-                tilt,
-                roll,
-                zoom,
-            ]
-        ):
-            return "Invalid camera control configs: at least one of the values must be non-zero"
-        return True
-
-    @classmethod
-    def execute(
-        cls,
-        camera_control_type: str,
-        horizontal_movement: float,
-        vertical_movement: float,
-        pan: float,
-        tilt: float,
-        roll: float,
-        zoom: float,
-    ) -> IO.NodeOutput:
-        return IO.NodeOutput(
-            KlingCameraControl(
-                type=KlingCameraControlType(camera_control_type),
-                config=KlingCameraConfig(
-                    horizontal=horizontal_movement,
-                    vertical=vertical_movement,
-                    pan=pan,
-                    roll=roll,
-                    tilt=tilt,
-                    zoom=zoom,
-                ),
-            )
-        )
-
-
 class KlingTextToVideoNode(IO.ComfyNode):
     """Kling Text to Video Node"""
 
@@ -782,7 +560,7 @@ class KlingTextToVideoNode(IO.ComfyNode):
                 IO.Combo.Input(
                     "mode",
                     options=modes,
-                    default=modes[8],
+                    default=modes[0],
                     tooltip="The configuration to use for the video generation following the format: mode / duration / model_name.",
                 ),
             ],
@@ -802,25 +580,7 @@ class KlingTextToVideoNode(IO.ComfyNode):
                 expr="""
                 (
                   $m := widgets.mode;
-                  $contains($m,"v2-5-turbo")
-                    ? ($contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35})
-                    : $contains($m,"v2-1-master")
-                      ? ($contains($m,"10s") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
-                      : $contains($m,"v2-master")
-                        ? ($contains($m,"10s") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
-                        : $contains($m,"v1-6")
-                          ? (
-                              $contains($m,"pro")
-                                ? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                                : ($contains($m,"10s") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
-                            )
-                          : $contains($m,"v1")
-                            ? (
-                                $contains($m,"pro")
-                                  ? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                                  : ($contains($m,"10s") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
-                              )
-                            : {"type":"usd","usd":0.14}
+                  $contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35}
                 )
                 """,
             ),
@@ -1716,71 +1476,6 @@ class OmniProImageNode(IO.ComfyNode):
         return IO.NodeOutput(torch.cat(tensors, dim=0))
 
 
-class KlingCameraControlT2VNode(IO.ComfyNode):
-    """
-    Kling Text to Video Camera Control Node. This node is a text to video node, but it supports controlling the camera.
-    Duration, mode, and model_name request fields are hard-coded because camera control is only supported in pro mode with the kling-v1-5 model at 5s duration as of 2025-05-02.
-    """
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingCameraControlT2VNode",
-            display_name="Kling Text to Video (Camera Control)",
-            category="partner/video/Kling",
-            description="Transform text into cinematic videos with professional camera movements that simulate real-world cinematography. Control virtual camera actions including zoom, rotation, pan, tilt, and first-person view, while maintaining focus on your original text.",
-            inputs=[
-                IO.String.Input("prompt", multiline=True, tooltip="Positive text prompt"),
-                IO.String.Input("negative_prompt", multiline=True, tooltip="Negative text prompt"),
-                IO.Float.Input("cfg_scale", default=0.75, min=0.0, max=1.0),
-                IO.Combo.Input(
-                    "aspect_ratio",
-                    options=KlingVideoGenAspectRatio,
-                    default="16:9",
-                ),
-                IO.Custom("CAMERA_CONTROL").Input(
-                    "camera_control",
-                    tooltip="Can be created using the Kling Camera Controls node. Controls the camera movement and motion during the video generation.",
-                ),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="video_id"),
-                IO.String.Output(display_name="duration"),
-            ],
-            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.14}""",
-            ),
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        prompt: str,
-        negative_prompt: str,
-        cfg_scale: float,
-        aspect_ratio: str,
-        camera_control: KlingCameraControl | None = None,
-    ) -> IO.NodeOutput:
-        return await execute_text2video(
-            cls,
-            model_name=KlingVideoGenModelName.kling_v1,
-            cfg_scale=cfg_scale,
-            model_mode=KlingVideoGenMode.std,
-            aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio),
-            duration=KlingVideoGenDuration.field_5,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            camera_control=camera_control,
-        )
-
-
 class KlingImage2VideoNode(IO.ComfyNode):
     """Kling Image to Video Node"""
 
@@ -1796,11 +1491,10 @@ class KlingImage2VideoNode(IO.ComfyNode):
                 IO.String.Input("negative_prompt", multiline=True, tooltip="Negative text prompt"),
                 IO.Combo.Input(
                     "model_name",
-                    options=KlingVideoGenModelName,
-                    default="kling-v2-master",
+                    options=["kling-v2-5-turbo"],
                 ),
                 IO.Float.Input("cfg_scale", default=0.8, min=0.0, max=1.0),
-                IO.Combo.Input("mode", options=KlingVideoGenMode, default=KlingVideoGenMode.std),
+                IO.Combo.Input("mode", options=["pro"]),
                 IO.Combo.Input(
                     "aspect_ratio",
                     options=KlingVideoGenAspectRatio,
@@ -1820,29 +1514,10 @@ class KlingImage2VideoNode(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["mode", "model_name", "duration"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["duration"]),
                 expr="""
                 (
-                  $mode := widgets.mode;
-                  $model := widgets.model_name;
-                  $dur := widgets.duration;
-                  $contains($model,"v2-5-turbo")
-                    ? ($contains($dur,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35})
-                    : ($contains($model,"v2-1-master") or $contains($model,"v2-master"))
-                      ? ($contains($dur,"10") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
-                      : ($contains($model,"v2-1") or $contains($model,"v1-6") or $contains($model,"v1-5"))
-                        ? (
-                            $contains($mode,"pro")
-                              ? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                              : ($contains($dur,"10") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
-                          )
-                        : $contains($model,"v1")
-                          ? (
-                              $contains($mode,"pro")
-                                ? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                                : ($contains($dur,"10") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
-                            )
-                          : {"type":"usd","usd":0.14}
+                  $contains(widgets.duration,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35}
                 )
                 """,
             ),
@@ -1859,8 +1534,6 @@ class KlingImage2VideoNode(IO.ComfyNode):
         mode: str,
         aspect_ratio: str,
         duration: str,
-        camera_control: KlingCameraControl | None = None,
-        end_frame: torch.Tensor | None = None,
     ) -> IO.NodeOutput:
         return await execute_image2video(
             cls,
@@ -1872,79 +1545,6 @@ class KlingImage2VideoNode(IO.ComfyNode):
             aspect_ratio=aspect_ratio,
             model_mode=mode,
             duration=duration,
-            camera_control=camera_control,
-            end_frame=end_frame,
-        )
-
-
-class KlingCameraControlI2VNode(IO.ComfyNode):
-    """
-    Kling Image to Video Camera Control Node. This node is a image to video node, but it supports controlling the camera.
-    Duration, mode, and model_name request fields are hard-coded because camera control is only supported in pro mode with the kling-v1-5 model at 5s duration as of 2025-05-02.
-    """
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingCameraControlI2VNode",
-            display_name="Kling Image to Video (Camera Control)",
-            category="partner/video/Kling",
-            description="Transform still images into cinematic videos with professional camera movements that simulate real-world cinematography. Control virtual camera actions including zoom, rotation, pan, tilt, and first-person view, while maintaining focus on your original image.",
-            inputs=[
-                IO.Image.Input(
-                    "start_frame",
-                    tooltip="Reference Image - URL or Base64 encoded string, cannot exceed 10MB, resolution not less than 300*300px, aspect ratio between 1:2.5 ~ 2.5:1. Base64 should not include data:image prefix.",
-                ),
-                IO.String.Input("prompt", multiline=True, tooltip="Positive text prompt"),
-                IO.String.Input("negative_prompt", multiline=True, tooltip="Negative text prompt"),
-                IO.Float.Input("cfg_scale", default=0.75, min=0.0, max=1.0),
-                IO.Combo.Input(
-                    "aspect_ratio",
-                    options=KlingVideoGenAspectRatio,
-                    default=KlingVideoGenAspectRatio.field_16_9,
-                ),
-                IO.Custom("CAMERA_CONTROL").Input(
-                    "camera_control",
-                    tooltip="Can be created using the Kling Camera Controls node. Controls the camera movement and motion during the video generation.",
-                ),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="video_id"),
-                IO.String.Output(display_name="duration"),
-            ],
-            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.49}""",
-            ),
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        start_frame: torch.Tensor,
-        prompt: str,
-        negative_prompt: str,
-        cfg_scale: float,
-        aspect_ratio: str,
-        camera_control: KlingCameraControl,
-    ) -> IO.NodeOutput:
-        return await execute_image2video(
-            cls,
-            model_name=KlingVideoGenModelName.kling_v1_5,
-            start_frame=start_frame,
-            cfg_scale=cfg_scale,
-            model_mode=KlingVideoGenMode.pro,
-            aspect_ratio=KlingVideoGenAspectRatio(aspect_ratio),
-            duration=KlingVideoGenDuration.field_5,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            camera_control=camera_control,
         )
 
 
@@ -1977,7 +1577,7 @@ class KlingStartEndFrameNode(IO.ComfyNode):
                 IO.Combo.Input(
                     "mode",
                     options=modes,
-                    default=modes[6],
+                    default=modes[0],
                     tooltip="The configuration to use for the video generation following the format: mode / duration / model_name.",
                 ),
             ],
@@ -1997,25 +1597,7 @@ class KlingStartEndFrameNode(IO.ComfyNode):
                 expr="""
                 (
                   $m := widgets.mode;
-                  $contains($m,"v2-5-turbo")
-                    ? ($contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35})
-                    : $contains($m,"v2-1")
-                      ? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                      : $contains($m,"v2-master")
-                        ? ($contains($m,"10s") ? {"type":"usd","usd":2.8} : {"type":"usd","usd":1.4})
-                        : $contains($m,"v1-6")
-                          ? (
-                              $contains($m,"pro")
-                                ? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                                : ($contains($m,"10s") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
-                            )
-                          : $contains($m,"v1")
-                            ? (
-                                $contains($m,"pro")
-                                  ? ($contains($m,"10s") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                                  : ($contains($m,"10s") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
-                              )
-                            : {"type":"usd","usd":0.14}
+                  $contains($m,"10") ? {"type":"usd","usd":0.7} : {"type":"usd","usd":0.35}
                 )
                 """,
             ),
@@ -2126,169 +1708,6 @@ class KlingVideoExtendNode(IO.ComfyNode):
         return IO.NodeOutput(await download_url_to_video_output(str(video.url)), str(video.id), str(video.duration))
 
 
-class KlingDualCharacterVideoEffectNode(IO.ComfyNode):
-    """Kling Dual Character Video Effect Node"""
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingDualCharacterVideoEffectNode",
-            display_name="Kling Dual Character Video Effects",
-            category="partner/video/Kling",
-            description="Achieve different special effects when generating a video based on the effect_scene. First image will be positioned on left side, second on right side of the composite.",
-            inputs=[
-                IO.Image.Input("image_left", tooltip="Left side image"),
-                IO.Image.Input("image_right", tooltip="Right side image"),
-                IO.Combo.Input(
-                    "effect_scene",
-                    options=[i.value for i in KlingDualCharacterEffectsScene],
-                ),
-                IO.Combo.Input(
-                    "model_name",
-                    options=[i.value for i in KlingCharacterEffectModelName],
-                    default="kling-v1",
-                ),
-                IO.Combo.Input(
-                    "mode",
-                    options=[i.value for i in KlingVideoGenMode],
-                    default="std",
-                ),
-                IO.Combo.Input(
-                    "duration",
-                    options=[i.value for i in KlingVideoGenDuration],
-                ),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="duration"),
-            ],
-            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", "model_name", "duration"]),
-                expr="""
-                (
-                  $mode := widgets.mode;
-                  $model := widgets.model_name;
-                  $dur := widgets.duration;
-                  ($contains($model,"v1-6") or $contains($model,"v1-5"))
-                    ? (
-                        $contains($mode,"pro")
-                          ? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                          : ($contains($dur,"10") ? {"type":"usd","usd":0.56} : {"type":"usd","usd":0.28})
-                      )
-                    : $contains($model,"v1")
-                      ? (
-                          $contains($mode,"pro")
-                            ? ($contains($dur,"10") ? {"type":"usd","usd":0.98} : {"type":"usd","usd":0.49})
-                            : ($contains($dur,"10") ? {"type":"usd","usd":0.28} : {"type":"usd","usd":0.14})
-                        )
-                      : {"type":"usd","usd":0.14}
-                )
-                """,
-            ),
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        image_left: torch.Tensor,
-        image_right: torch.Tensor,
-        effect_scene: KlingDualCharacterEffectsScene,
-        model_name: KlingCharacterEffectModelName,
-        mode: KlingVideoGenMode,
-        duration: KlingVideoGenDuration,
-    ) -> IO.NodeOutput:
-        video, _, duration = await execute_video_effect(
-            cls,
-            dual_character=True,
-            effect_scene=effect_scene,
-            model_name=model_name,
-            model_mode=mode,
-            duration=duration,
-            image_1=image_left,
-            image_2=image_right,
-        )
-        return IO.NodeOutput(video, duration)
-
-
-class KlingSingleImageVideoEffectNode(IO.ComfyNode):
-    """Kling Single Image Video Effect Node"""
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingSingleImageVideoEffectNode",
-            display_name="Kling Video Effects",
-            category="partner/video/Kling",
-            description="Achieve different special effects when generating a video based on the effect_scene.",
-            inputs=[
-                IO.Image.Input(
-                    "image",
-                    tooltip=" Reference Image. URL or Base64 encoded string (without data:image prefix). File size cannot exceed 10MB, resolution not less than 300*300px, aspect ratio between 1:2.5 ~ 2.5:1",
-                ),
-                IO.Combo.Input(
-                    "effect_scene",
-                    options=[i.value for i in KlingSingleImageEffectsScene],
-                ),
-                IO.Combo.Input(
-                    "model_name",
-                    options=[i.value for i in KlingSingleImageEffectModelName],
-                ),
-                IO.Combo.Input(
-                    "duration",
-                    options=[i.value for i in KlingVideoGenDuration],
-                ),
-            ],
-            outputs=[
-                IO.Video.Output(),
-                IO.String.Output(display_name="video_id"),
-                IO.String.Output(display_name="duration"),
-            ],
-            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=["effect_scene"]),
-                expr="""
-                (
-                  ($contains(widgets.effect_scene,"dizzydizzy") or $contains(widgets.effect_scene,"bloombloom"))
-                    ? {"type":"usd","usd":0.49}
-                    : {"type":"usd","usd":0.28}
-                )
-                """,
-            ),
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        image: torch.Tensor,
-        effect_scene: KlingSingleImageEffectsScene,
-        model_name: KlingSingleImageEffectModelName,
-        duration: KlingVideoGenDuration,
-    ) -> IO.NodeOutput:
-        return IO.NodeOutput(
-            *(
-                await execute_video_effect(
-                    cls,
-                    dual_character=False,
-                    effect_scene=effect_scene,
-                    model_name=model_name,
-                    duration=duration,
-                    image_1=image,
-                )
-            )
-        )
-
-
 class KlingLipSyncAudioToVideoNode(IO.ComfyNode):
     """Kling Lip Sync Audio to Video Node. Syncs mouth movements in a video file to the audio content of an audio file."""
 
@@ -2409,73 +1828,6 @@ class KlingLipSyncTextToVideoNode(IO.ComfyNode):
         )
 
 
-class KlingVirtualTryOnNode(IO.ComfyNode):
-    """Kling Virtual Try On Node."""
-
-    @classmethod
-    def define_schema(cls) -> IO.Schema:
-        return IO.Schema(
-            node_id="KlingVirtualTryOnNode",
-            display_name="Kling Virtual Try On",
-            category="partner/image/Kling",
-            description="Kling Virtual Try On Node. Input a human image and a cloth image to try on the cloth on the human. You can merge multiple clothing item pictures into one image with a white background.",
-            inputs=[
-                IO.Image.Input("human_image"),
-                IO.Image.Input("cloth_image"),
-                IO.Combo.Input(
-                    "model_name",
-                    options=[i.value for i in KlingVirtualTryOnModelName],
-                    default="kolors-virtual-try-on-v1",
-                ),
-            ],
-            outputs=[
-                IO.Image.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.7}""",
-            ),
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        human_image: torch.Tensor,
-        cloth_image: torch.Tensor,
-        model_name: KlingVirtualTryOnModelName,
-    ) -> IO.NodeOutput:
-        task_creation_response = await sync_op(
-            cls,
-            ApiEndpoint(path=PATH_VIRTUAL_TRY_ON, method="POST"),
-            response_model=KlingVirtualTryOnResponse,
-            data=KlingVirtualTryOnRequest(
-                human_image=tensor_to_base64_string(human_image),
-                cloth_image=tensor_to_base64_string(cloth_image),
-                model_name=model_name,
-            ),
-        )
-
-        validate_task_creation_response(task_creation_response)
-        task_id = task_creation_response.data.task_id
-
-        final_response = await poll_op(
-            cls,
-            ApiEndpoint(path=f"{PATH_VIRTUAL_TRY_ON}/{task_id}"),
-            response_model=KlingVirtualTryOnResponse,
-            estimated_duration=AVERAGE_DURATION_VIRTUAL_TRY_ON,
-            status_extractor=lambda r: (r.data.task_status.value if r.data and r.data.task_status else None),
-        )
-        validate_image_result_response(final_response)
-
-        images = get_images_from_response(final_response)
-        return IO.NodeOutput(await image_result_to_node_output(images))
-
-
 class KlingImageGenerationNode(IO.ComfyNode):
     """Kling Image Generation Node. Generate an image from a text prompt with an optional reference image."""
 
@@ -2514,7 +1866,7 @@ class KlingImageGenerationNode(IO.ComfyNode):
                     tooltip="Subject reference similarity",
                     advanced=True,
                 ),
-                IO.Combo.Input("model_name", options=["kling-v3", "kling-v2", "kling-v1-5"]),
+                IO.Combo.Input("model_name", options=["kling-v3", "kling-v2"]),
                 IO.Combo.Input(
                     "aspect_ratio",
                     options=[i.value for i in KlingImageGenAspectRatio],
@@ -2550,14 +1902,10 @@ class KlingImageGenerationNode(IO.ComfyNode):
             ],
             is_api_node=True,
             price_badge=IO.PriceBadge(
-                depends_on=IO.PriceBadgeDepends(widgets=["model_name", "n"], inputs=["image"]),
+                depends_on=IO.PriceBadgeDepends(widgets=["model_name", "n"]),
                 expr="""
                 (
-                  $m := widgets.model_name;
-                  $base :=
-                    $contains($m,"kling-v1-5")
-                      ? (inputs.image.connected ? 0.028 : 0.014)
-                      : $contains($m,"kling-v3") ? 0.028 : 0.014;
+                  $base := $contains(widgets.model_name,"kling-v3") ? 0.028 : 0.014;
                   {"type":"usd","usd": $base * widgets.n}
                 )
                 """,
@@ -3394,19 +2742,13 @@ class KlingExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
         return [
-            KlingCameraControls,
             KlingTextToVideoNode,
             KlingImage2VideoNode,
-            KlingCameraControlI2VNode,
-            KlingCameraControlT2VNode,
             KlingStartEndFrameNode,
             KlingVideoExtendNode,
             KlingLipSyncAudioToVideoNode,
             KlingLipSyncTextToVideoNode,
-            KlingVirtualTryOnNode,
             KlingImageGenerationNode,
-            KlingSingleImageVideoEffectNode,
-            KlingDualCharacterVideoEffectNode,
             OmniProTextToVideoNode,
             OmniProFirstLastFrameNode,
             OmniProImageToVideoNode,

comfy_api_nodes/nodes_ltxv.py

@@ -6,8 +6,12 @@ from typing_extensions import override
 from comfy_api.latest import IO, ComfyExtension, Input, InputImpl
 from comfy_api_nodes.util import (
     ApiEndpoint,
+    download_url_to_video_output,
     get_number_of_images,
+    poll_op,
+    sync_op,
     sync_op_raw,
+    upload_audio_to_comfyapi,
     upload_images_to_comfyapi,
     validate_string,
 )
@@ -17,6 +21,11 @@ MODELS_MAP = {
     "LTX-2 (Fast)": "ltx-2-fast",
 }
 
+V25_MODELS_MAP = {
+    "LTX-2.5 (Fast)": "ltx-2-5-fast",
+    "LTX-2.5 (Pro)": "ltx-2-5-pro",
+}
+
 
 class ExecuteTaskRequest(BaseModel):
     prompt: str = Field(...)
@@ -26,6 +35,48 @@ class ExecuteTaskRequest(BaseModel):
     fps: int | None = Field(25)
     generate_audio: bool | None = Field(True)
     image_uri: str | None = Field(None)
+    last_frame_uri: str | None = Field(None)
+
+
+class AudioToVideoRequest(BaseModel):
+    prompt: str = Field(...)
+    model: str = Field(...)
+    resolution: str = Field(...)
+    audio_uri: str = Field(...)
+    image_uri: str | None = Field(None)
+
+
+class Ltx25SubmitResponse(BaseModel):
+    id: str = Field(...)
+
+
+class Ltx25JobResult(BaseModel):
+    video_url: str | None = Field(None)
+
+
+class Ltx25JobStatusResponse(BaseModel):
+    id: str = Field(...)
+    status: str = Field(...)
+    result: Ltx25JobResult | None = Field(None)
+
+
+async def _v25_submit_and_poll(cls: type[IO.ComfyNode], route: str, data: BaseModel) -> IO.NodeOutput:
+    submit = await sync_op(
+        cls,
+        ApiEndpoint(f"/proxy/ltx/v2/{route}", "POST"),
+        response_model=Ltx25SubmitResponse,
+        data=data,
+        max_retries=1,
+    )
+    job = await poll_op(
+        cls,
+        ApiEndpoint(f"/proxy/ltx/v2/{route}/{submit.id}"),
+        response_model=Ltx25JobStatusResponse,
+        status_extractor=lambda r: r.status,
+    )
+    if not job.result or not job.result.video_url:
+        raise RuntimeError(f"LTX job {job.id} completed without a video URL.")
+    return IO.NodeOutput(await download_url_to_video_output(job.result.video_url, cls=cls))
 
 
 PRICE_BADGE = IO.PriceBadge(
@@ -43,6 +94,128 @@ PRICE_BADGE = IO.PriceBadge(
     """,
 )
 
+V25_PRICE_BADGE = IO.PriceBadge(
+    depends_on=IO.PriceBadgeDepends(widgets=["model", "model.duration", "model.resolution"]),
+    expr="""
+    (
+      $prices := {
+        "ltx-2.5 (fast)": {
+          "1280x720":0.1287,"720x1280":0.1287,
+          "1920x1080":0.1859,"1080x1920":0.1859,
+          "2560x1440":0.2717,"1440x2560":0.2717,
+          "3840x2160":0.429,"2160x3840":0.429
+        },
+        "ltx-2.5 (pro)": {
+          "1280x720":0.1716,"720x1280":0.1716,
+          "1920x1080":0.2431,"1080x1920":0.2431
+        }
+      };
+      $model := $lookup(widgets, "model");
+      $table := $type($model) = "string" ? $lookup($prices, $model) : undefined;
+      $res := $lookup(widgets, "model.resolution");
+      $pps := $type($table) = "object" and $type($res) = "string" ? $lookup($table, $res) : undefined;
+      $durRaw := $lookup(widgets, "model.duration");
+      $dur := $type($durRaw) in ["string", "number"] ? $number($durRaw) : undefined;
+      $type($pps) = "number" and $type($dur) = "number"
+        ? {"type":"usd","usd": $pps * $dur}
+        : undefined
+    )
+    """,
+)
+
+V25_A2V_PRICE_BADGE = IO.PriceBadge(
+    depends_on=IO.PriceBadgeDepends(widgets=["model"]),
+    expr="""
+    (
+      $rates := {"ltx-2.5 (fast)":0.1859, "ltx-2.5 (pro)":0.2431};
+      $model := $lookup(widgets, "model");
+      $rate := $type($model) = "string" ? $lookup($rates, $model) : undefined;
+      $type($rate) = "number"
+        ? {"type":"usd","usd": $rate, "format":{"suffix":"/second"}}
+        : undefined
+    )
+    """,
+)
+
+
+def _v25_generation_inputs(
+    durations: list[str], resolutions: list[str], fps_options: list[str], tooltip: str | None
+) -> list:
+    return [
+        IO.Combo.Input(
+            "duration",
+            options=durations,
+            default="8",
+            tooltip=tooltip,
+        ),
+        IO.Combo.Input(
+            "resolution",
+            options=resolutions,
+            default="1920x1080",
+        ),
+        IO.Combo.Input("fps", options=fps_options, default="25"),
+        IO.Boolean.Input(
+            "generate_audio",
+            default=True,
+            tooltip="When true, the generated video will include AI-generated audio matching the scene.",
+            advanced=True,
+        ),
+    ]
+
+
+def _v25_model_combo() -> IO.DynamicCombo.Input:
+    return IO.DynamicCombo.Input(
+        "model",
+        options=[
+            IO.DynamicCombo.Option(
+                "LTX-2.5 (Fast)",
+                _v25_generation_inputs(
+                    ["2", "3", "4", "5", "6", "8", "10", "12", "14", "16", "18", "20"],
+                    [
+                        "1280x720",
+                        "720x1280",
+                        "1920x1080",
+                        "1080x1920",
+                        "2560x1440",
+                        "1440x2560",
+                        "3840x2160",
+                        "2160x3840",
+                    ],
+                    ["24", "25", "48", "50"],
+                    "Video duration in seconds. Durations over 10s require a 720p/1080p resolution and 24/25 FPS.",
+                ),
+            ),
+            IO.DynamicCombo.Option(
+                "LTX-2.5 (Pro)",
+                _v25_generation_inputs(
+                    ["2", "3", "4", "5", "6", "8", "10"],
+                    ["1280x720", "720x1280", "1920x1080", "1080x1920"],
+                    ["24", "25", "50"],
+                    "Video duration in seconds.",
+                ),
+            ),
+        ],
+    )
+
+
+def _v25_seed_input() -> IO.Int.Input:
+    return IO.Int.Input(
+        "seed",
+        default=42,
+        min=0,
+        max=0xFFFFFFFF,
+        control_after_generate=True,
+        tooltip="Seed to determine if node should re-run; "
+        "actual results are nondeterministic regardless of seed.",
+    )
+
+
+def _v25_validate_settings(model: dict) -> None:
+    if int(model["duration"]) > 10 and (
+        int(model["fps"]) > 25 or model["resolution"] in ("2560x1440", "1440x2560", "3840x2160", "2160x3840")
+    ):
+        raise ValueError("Durations over 10s require a 720p or 1080p resolution and 24/25 FPS.")
+
 
 class TextToVideoNode(IO.ComfyNode):
     @classmethod
@@ -86,6 +259,7 @@ class TextToVideoNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=PRICE_BADGE,
         )
 
@@ -164,6 +338,7 @@ class ImageToVideoNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=PRICE_BADGE,
         )
 
@@ -203,12 +378,217 @@ class ImageToVideoNode(IO.ComfyNode):
         return IO.NodeOutput(InputImpl.VideoFromFile(BytesIO(response)))
 
 
+class Ltx25TextToVideoNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="LtxApi25TextToVideo",
+            display_name="LTX 2.5 Text To Video",
+            category="partner/video/LTXV",
+            description="Professional-quality videos with customizable duration and resolution.",
+            inputs=[
+                _v25_model_combo(),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                ),
+                _v25_seed_input(),
+            ],
+            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=V25_PRICE_BADGE,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        prompt: str,
+        seed: int = 42,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1, max_length=10000)
+        _v25_validate_settings(model)
+        return await _v25_submit_and_poll(
+            cls,
+            "text-to-video",
+            ExecuteTaskRequest(
+                prompt=prompt,
+                model=V25_MODELS_MAP[model["model"]],
+                duration=int(model["duration"]),
+                resolution=model["resolution"],
+                fps=int(model["fps"]),
+                generate_audio=model["generate_audio"],
+            ),
+        )
+
+
+class Ltx25ImageToVideoNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="LtxApi25ImageToVideo",
+            display_name="LTX 2.5 Image To Video",
+            category="partner/video/LTXV",
+            description="Professional-quality videos with customizable duration and resolution based on start image.",
+            inputs=[
+                IO.Image.Input("image", tooltip="First frame to be used for the video."),
+                _v25_model_combo(),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                ),
+                _v25_seed_input(),
+                IO.Image.Input(
+                    "last_frame",
+                    optional=True,
+                    tooltip="Last frame to be used for the video.",
+                ),
+            ],
+            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=V25_PRICE_BADGE,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        model: dict,
+        prompt: str,
+        seed: int = 42,
+        last_frame: Input.Image | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1, max_length=10000)
+        _v25_validate_settings(model)
+        if get_number_of_images(image) != 1:
+            raise ValueError("Currently only one input image is supported.")
+        last_frame_uri = None
+        if last_frame is not None:
+            if get_number_of_images(last_frame) != 1:
+                raise ValueError("Currently only one last frame image is supported.")
+            last_frame_uri = (await upload_images_to_comfyapi(cls, last_frame, max_images=1, mime_type="image/png"))[0]
+        return await _v25_submit_and_poll(
+            cls,
+            "image-to-video",
+            ExecuteTaskRequest(
+                image_uri=(await upload_images_to_comfyapi(cls, image, max_images=1, mime_type="image/png"))[0],
+                last_frame_uri=last_frame_uri,
+                prompt=prompt,
+                model=V25_MODELS_MAP[model["model"]],
+                duration=int(model["duration"]),
+                resolution=model["resolution"],
+                fps=int(model["fps"]),
+                generate_audio=model["generate_audio"],
+            ),
+        )
+
+
+class Ltx25AudioToVideoNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="LtxApi25AudioToVideo",
+            display_name="LTX 2.5 Audio To Video",
+            category="partner/video/LTXV",
+            description="Generate a video driven by an audio track, with an optional first frame image.",
+            inputs=[
+                IO.Audio.Input(
+                    "audio",
+                    tooltip="Audio track driving the video. Its length (2-20 seconds) sets the video duration.",
+                ),
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "LTX-2.5 (Fast)",
+                            [IO.Combo.Input("resolution", options=["1920x1080", "1080x1920"])],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "LTX-2.5 (Pro)",
+                            [IO.Combo.Input("resolution", options=["1920x1080", "1080x1920"])],
+                        ),
+                    ],
+                ),
+                IO.String.Input(
+                    "prompt",
+                    multiline=True,
+                    default="",
+                ),
+                _v25_seed_input(),
+                IO.Image.Input(
+                    "image",
+                    optional=True,
+                    tooltip="Optional first frame to be used for the video.",
+                ),
+            ],
+            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=V25_A2V_PRICE_BADGE,
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        audio: Input.Audio,
+        model: dict,
+        prompt: str,
+        seed: int = 42,
+        image: Input.Image | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(prompt, min_length=1, max_length=10000)
+        audio_duration = audio["waveform"].shape[-1] / audio["sample_rate"]
+        if not 2 <= audio_duration <= 20:
+            raise ValueError(f"Audio duration must be between 2 and 20 seconds, got {audio_duration:.1f}s.")
+        image_uri = None
+        if image is not None:
+            if get_number_of_images(image) != 1:
+                raise ValueError("Currently only one input image is supported.")
+            image_uri = (await upload_images_to_comfyapi(cls, image, max_images=1, mime_type="image/png"))[0]
+        return await _v25_submit_and_poll(
+            cls,
+            "audio-to-video",
+            AudioToVideoRequest(
+                prompt=prompt,
+                model=V25_MODELS_MAP[model["model"]],
+                resolution=model["resolution"],
+                audio_uri=await upload_audio_to_comfyapi(cls, audio),
+                image_uri=image_uri,
+            ),
+        )
+
+
 class LtxvApiExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
         return [
             TextToVideoNode,
             ImageToVideoNode,
+            Ltx25TextToVideoNode,
+            Ltx25ImageToVideoNode,
+            Ltx25AudioToVideoNode,
         ]
 
 

comfy_api_nodes/nodes_minimax.py

@@ -3,12 +3,14 @@ from typing import Optional
 import torch
 from typing_extensions import override
 
-from comfy_api.latest import IO, ComfyExtension
+from comfy_api.latest import IO, ComfyExtension, Input
 from comfy_api_nodes.apis.minimax import (
     Hailuo03AudioContent,
     Hailuo03AudioContentUrl,
+    Hailuo03ContextIRRequest,
     Hailuo03ImageContent,
     Hailuo03ImageContentUrl,
+    Hailuo03RegenerationRequest,
     Hailuo03TaskCreationRequest,
     Hailuo03TaskCreationResponse,
     Hailuo03TaskQueryResponse,
@@ -456,6 +458,9 @@ HAILUO_03_QUERY_ENDPOINT = "/proxy/minimax/v2/query/video_generation"  # + /{tas
 HAILUO_03_MODELS = {"MiniMax H3": "MiniMax-H3"}
 HAILUO_03_FAILED_STATUSES = ["failed", "cancelled", "expired"]
 
+HAILUO_03_CONTEXT_IR_ENDPOINT = "/proxy/minimax/v2/h3_context_ir"
+HAILUO_03_REGENERATION_ENDPOINT = "/proxy/minimax/v2/video_regeneration"
+
 
 def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = True):
     inputs = [
@@ -487,10 +492,10 @@ def _hailuo03_model_inputs(include_ratio: bool = True, allow_adaptive: bool = Tr
         IO.Int.Input(
             "duration",
             default=5,
-            min=5,
+            min=4,
             max=15,
             step=1,
-            tooltip="Duration of the output video in seconds (5-15).",
+            tooltip="Duration of the output video in seconds (4-15).",
             display_mode=IO.NumberDisplay.slider,
         )
     )
@@ -939,6 +944,592 @@ class MinimaxHailuo03ReferenceNode(IO.ComfyNode):
         )
 
 
+class MinimaxHailuo03ContextIRNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="MinimaxHailuo03ContextIRNode",
+            display_name="MiniMax H3 Context IR (Prompt Enhancer)",
+            category="partner/video/MiniMax",
+            description="Analyze text and media context with MiniMax H3 Context IR and produce an enhanced, "
+            "structured video prompt. Feed the output into the prompt of a MiniMax H3 video node and attach "
+            "the same media there in the same order, because the enhanced prompt refers to the attached "
+            "media by position.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "MiniMax H3",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Description of the video you intend to generate.",
+                                ),
+                                IO.Int.Input(
+                                    "duration",
+                                    default=5,
+                                    min=4,
+                                    max=15,
+                                    step=1,
+                                    tooltip="Duration of the video you intend to generate, in seconds (4-15).",
+                                    display_mode=IO.NumberDisplay.slider,
+                                ),
+                                IO.Combo.Input(
+                                    "ratio",
+                                    options=["adaptive", "16:9", "4:3", "1:1", "3:4", "9:16", "21:9"],
+                                    default="adaptive",
+                                    tooltip="Aspect ratio of the video you intend to generate. 'adaptive' "
+                                    "requires at least one image, video, or audio input.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_images",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Image.Input("reference_image"),
+                                        names=[
+                                            "image_1",
+                                            "image_2",
+                                            "image_3",
+                                            "image_4",
+                                            "image_5",
+                                            "image_6",
+                                            "image_7",
+                                            "image_8",
+                                            "image_9",
+                                        ],
+                                        min=0,
+                                    ),
+                                    tooltip="Subject or style reference images, referred to in the prompt "
+                                    "as 'Image 1'..'Image 9' in connection order. Up to 9 images.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_videos",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Video.Input("reference_video"),
+                                        names=["video_1", "video_2", "video_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Motion or scene reference videos, referred to in the prompt "
+                                    "as 'Video 1'..'Video 3' in connection order. Up to 3 videos, "
+                                    "2-15 seconds each, 15 seconds in total.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_audios",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Audio.Input("reference_audio"),
+                                        names=["audio_1", "audio_2", "audio_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Audio references, referred to in the prompt as "
+                                    "'Audio 1'..'Audio 3' in connection order. Up to 3 clips, "
+                                    "2-15 seconds each, 15 seconds in total. Cannot be used without "
+                                    "a reference image or video.",
+                                ),
+                            ],
+                        )
+                    ],
+                    tooltip="Model to use for prompt enhancement.",
+                ),
+                IO.Image.Input(
+                    "first_frame",
+                    tooltip="First frame of the video you intend to generate. Cannot be combined with "
+                    "reference media.",
+                    optional=True,
+                ),
+                IO.Image.Input(
+                    "last_frame",
+                    tooltip="Last frame of the video you intend to generate. Cannot be combined with "
+                    "reference media.",
+                    optional=True,
+                ),
+            ],
+            outputs=[
+                IO.String.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(
+                    inputs=["first_frame", "last_frame"],
+                    input_groups=["model.reference_images", "model.reference_videos", "model.reference_audios"],
+                ),
+                expr="""
+                (
+                  $imgsRaw := $lookup(inputGroups, "model.reference_images");
+                  $imgs := $imgsRaw ? $imgsRaw : 0;
+                  $vidsRaw := $lookup(inputGroups, "model.reference_videos");
+                  $vids := $vidsRaw ? $vidsRaw : 0;
+                  $audsRaw := $lookup(inputGroups, "model.reference_audios");
+                  $auds := $audsRaw ? $audsRaw : 0;
+                  $frames := (inputs.first_frame.connected ? 1 : 0) + (inputs.last_frame.connected ? 1 : 0);
+                  ($imgs + $vids + $auds) > 0
+                    ? {"type": "range_usd", "min_usd": 0.06, "max_usd": 0.11, "format": {"approximate": true}}
+                    : $frames > 0
+                      ? {"type": "usd", "usd": 0.05, "format": {"approximate": true}}
+                      : {"type": "usd", "usd": 0.02, "format": {"approximate": true}}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        first_frame: torch.Tensor | None = None,
+        last_frame: torch.Tensor | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(model["prompt"], strip_whitespace=True, min_length=1)
+
+        reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
+        reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
+        reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
+        has_frames = first_frame is not None or last_frame is not None
+        has_references = bool(reference_images) or bool(reference_videos) or bool(reference_audios)
+        if has_frames and has_references:
+            raise ValueError(
+                "First/last frame and reference media are mutually exclusive. Use frames for an "
+                "image-to-video prompt, or reference media for a reference-to-video prompt."
+            )
+        if reference_audios and not reference_images and not reference_videos:
+            raise ValueError("Reference audio cannot be used without a reference image or video.")
+        if not has_frames and not has_references and model["ratio"] == "adaptive":
+            raise ValueError(
+                "Ratio 'adaptive' is not supported for text-only requests; select an explicit aspect ratio."
+            )
+
+        for frame in (first_frame, last_frame):
+            if frame is not None:
+                validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+                validate_image_dimensions(frame, min_width=256, min_height=256)
+        for image in reference_images.values():
+            validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+            validate_image_dimensions(image, min_width=256, min_height=256)
+
+        total_video_duration = 0.0
+        for i, video in enumerate(reference_videos.values(), 1):
+            try:
+                fps = float(video.get_frame_rate())
+            except Exception:
+                fps = 0.0
+            if fps and not (23.9 <= fps <= 60.5):
+                raise ValueError(f"Reference video {i} is {fps:.2f} FPS. Supported range is 23.976-60 FPS.")
+            try:
+                dur = video.get_duration()
+            except Exception:
+                continue
+            if dur < 1.8:
+                raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_video_duration += dur
+        if total_video_duration > 15.1:
+            raise ValueError(
+                f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        total_audio_duration = 0.0
+        for i, audio in enumerate(reference_audios.values(), 1):
+            dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
+            if dur < 1.8:
+                raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_audio_duration += dur
+        if total_audio_duration > 15.1:
+            raise ValueError(
+                f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        content: list = [Hailuo03TextContent(text=model["prompt"])]
+        if first_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, first_frame, max_images=1, wait_label="Uploading first frame"
+                            )
+                        )[0],
+                    ),
+                    role="first_frame",
+                )
+            )
+        if last_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, last_frame, max_images=1, wait_label="Uploading last frame"
+                            )
+                        )[0],
+                    ),
+                    role="last_frame",
+                )
+            )
+        for i, image in enumerate(reference_images.values(), 1):
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, image, max_images=1, wait_label=f"Uploading image {i}"
+                            )
+                        )[0],
+                    ),
+                    role="reference_image",
+                )
+            )
+        for i, video in enumerate(reference_videos.values(), 1):
+            content.append(
+                Hailuo03VideoContent(
+                    video_url=Hailuo03VideoContentUrl(
+                        url=await upload_video_to_comfyapi(cls, video, wait_label=f"Uploading video {i}"),
+                    ),
+                )
+            )
+        for audio in reference_audios.values():
+            content.append(
+                Hailuo03AudioContent(
+                    audio_url=Hailuo03AudioContentUrl(
+                        url=await upload_audio_to_comfyapi(
+                            cls,
+                            audio,
+                            container_format="mp3",
+                            codec_name="libmp3lame",
+                            mime_type="audio/mpeg",
+                        ),
+                    ),
+                )
+            )
+
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=HAILUO_03_CONTEXT_IR_ENDPOINT, method="POST"),
+            response_model=Hailuo03TaskCreationResponse,
+            data=Hailuo03ContextIRRequest(
+                model=HAILUO_03_MODELS[model["model"]],
+                content=content,
+                duration=model["duration"],
+                ratio=None if model["ratio"] == "adaptive" else model["ratio"],
+            ),
+        )
+        task_result = await poll_op(
+            cls,
+            ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
+            response_model=Hailuo03TaskQueryResponse,
+            status_extractor=lambda r: r.task.status,
+            failed_statuses=HAILUO_03_FAILED_STATUSES,
+            poll_interval=5,
+        )
+        prompt = task_result.task.content.prompt if task_result.task.content else None
+        if not prompt:
+            raise Exception(f"No enhanced prompt in the response: {task_result.model_dump()}")
+        return IO.NodeOutput(prompt)
+
+
+class MinimaxHailuo03RegenerateNode(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="MinimaxHailuo03RegenerateNode",
+            display_name="MiniMax H3 Regenerate to 2K",
+            category="partner/video/MiniMax",
+            description="Re-render a MiniMax H3 768P output at 2K resolution. Connect the unmodified 768P "
+            "video and the exact prompt used to generate it; if the original generation used first/last "
+            "frames or reference media, attach the same inputs.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "MiniMax H3",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="The exact prompt used to generate the source video.",
+                                ),
+                                IO.Combo.Input(
+                                    "resolution",
+                                    options=["2K"],
+                                    tooltip="Resolution to re-render the source video at.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_images",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Image.Input("reference_image"),
+                                        names=[
+                                            "image_1",
+                                            "image_2",
+                                            "image_3",
+                                            "image_4",
+                                            "image_5",
+                                            "image_6",
+                                            "image_7",
+                                            "image_8",
+                                            "image_9",
+                                        ],
+                                        min=0,
+                                    ),
+                                    tooltip="Reference images from the original generation, in the same "
+                                    "order. Up to 9 images.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_videos",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Video.Input("reference_video"),
+                                        names=["video_1", "video_2", "video_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Reference videos from the original generation, in the same "
+                                    "order. Up to 3 videos, 2-15 seconds each, 15 seconds in total.",
+                                ),
+                                IO.Autogrow.Input(
+                                    "reference_audios",
+                                    template=IO.Autogrow.TemplateNames(
+                                        IO.Audio.Input("reference_audio"),
+                                        names=["audio_1", "audio_2", "audio_3"],
+                                        min=0,
+                                    ),
+                                    tooltip="Audio references from the original generation, in the same "
+                                    "order. Up to 3 clips, 2-15 seconds each, 15 seconds in total. "
+                                    "Cannot be used without a reference image or video.",
+                                ),
+                            ],
+                        )
+                    ],
+                    tooltip="Model to use for video regeneration.",
+                ),
+                IO.Video.Input(
+                    "video",
+                    tooltip="The MiniMax H3 768P output video to re-render. Connect the unmodified output "
+                    "of a MiniMax H3 video node (24 FPS, 4-15 seconds). 2K outputs cannot be used.",
+                ),
+                IO.Image.Input(
+                    "first_frame",
+                    tooltip="First frame image from the original generation, if one was used.",
+                    optional=True,
+                ),
+                IO.Image.Input(
+                    "last_frame",
+                    tooltip="Last frame image from the original generation, if one was used.",
+                    optional=True,
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add an AIGC watermark to the video.",
+                    advanced=True,
+                ),
+            ],
+            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,
+        model: dict,
+        video: Input.Video,
+        watermark: bool,
+        first_frame: torch.Tensor | None = None,
+        last_frame: torch.Tensor | None = None,
+    ) -> IO.NodeOutput:
+        validate_string(model["prompt"], strip_whitespace=True, min_length=1)
+
+        try:
+            fps = float(video.get_frame_rate())
+        except Exception:
+            fps = 0.0
+        if fps and not (23.9 <= fps <= 24.1):
+            raise ValueError(
+                f"The source video is {fps:.2f} FPS. Regeneration accepts unmodified MiniMax H3 768P "
+                "outputs, which are 24 FPS."
+            )
+        try:
+            width, height = video.get_dimensions()
+        except Exception:
+            width = height = 0
+        if width and height and (width % 32 or height % 32 or width * height > 1_032_192):
+            raise ValueError(
+                f"The source video is {width}x{height}. Regeneration accepts MiniMax H3 768P outputs "
+                "(width and height divisible by 32, at most 1,032,192 total pixels); 2K outputs cannot "
+                "be used as a source."
+            )
+        try:
+            frame_count = video.get_frame_count()
+        except Exception:
+            frame_count = 0
+        if frame_count and (frame_count < 107 or frame_count > 362 or (frame_count - 107) % 17):
+            raise ValueError(
+                f"The source video has {frame_count} frames. Regeneration accepts unmodified "
+                "MiniMax H3 outputs, whose length is 107 to 362 frames in steps of 17 "
+                "(4 to 15 seconds at 24 FPS)."
+            )
+
+        reference_images = {k: v for k, v in (model.get("reference_images") or {}).items() if v is not None}
+        reference_videos = {k: v for k, v in (model.get("reference_videos") or {}).items() if v is not None}
+        reference_audios = {k: v for k, v in (model.get("reference_audios") or {}).items() if v is not None}
+        if (first_frame is not None or last_frame is not None) and (
+            reference_images or reference_videos or reference_audios
+        ):
+            raise ValueError(
+                "First/last frame and reference media are mutually exclusive. Use frames for an "
+                "image-to-video prompt, or reference media for a reference-to-video prompt."
+            )
+        if reference_audios and not reference_images and not reference_videos:
+            raise ValueError("Reference audio cannot be used without a reference image or video.")
+
+        for frame in (first_frame, last_frame):
+            if frame is not None:
+                validate_image_aspect_ratio(frame, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+                validate_image_dimensions(frame, min_width=256, min_height=256)
+        for image in reference_images.values():
+            validate_image_aspect_ratio(image, (2, 5), (5, 2), strict=False)  # 0.4 to 2.5
+            validate_image_dimensions(image, min_width=256, min_height=256)
+
+        total_video_duration = 0.0
+        for i, ref_video in enumerate(reference_videos.values(), 1):
+            try:
+                ref_fps = float(ref_video.get_frame_rate())
+            except Exception:
+                ref_fps = 0.0
+            if ref_fps and not (23.9 <= ref_fps <= 60.5):
+                raise ValueError(f"Reference video {i} is {ref_fps:.2f} FPS. Supported range is 23.976-60 FPS.")
+            try:
+                dur = ref_video.get_duration()
+            except Exception:
+                continue
+            if dur < 1.8:
+                raise ValueError(f"Reference video {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_video_duration += dur
+        if total_video_duration > 15.1:
+            raise ValueError(
+                f"Total reference video duration is {total_video_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        total_audio_duration = 0.0
+        for i, audio in enumerate(reference_audios.values(), 1):
+            dur = int(audio["waveform"].shape[-1]) / int(audio["sample_rate"])
+            if dur < 1.8:
+                raise ValueError(f"Reference audio {i} is too short: {dur:.1f}s. Minimum duration is 2 seconds.")
+            total_audio_duration += dur
+        if total_audio_duration > 15.1:
+            raise ValueError(
+                f"Total reference audio duration is {total_audio_duration:.1f}s. Maximum is 15 seconds."
+            )
+
+        content: list = [
+            Hailuo03VideoContent(
+                video_url=Hailuo03VideoContentUrl(
+                    url=await upload_video_to_comfyapi(cls, video, wait_label="Uploading source video"),
+                ),
+                role="base_video",
+            ),
+            Hailuo03TextContent(text=model["prompt"]),
+        ]
+        if first_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, first_frame, max_images=1, wait_label="Uploading first frame"
+                            )
+                        )[0],
+                    ),
+                    role="first_frame",
+                )
+            )
+        if last_frame is not None:
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, last_frame, max_images=1, wait_label="Uploading last frame"
+                            )
+                        )[0],
+                    ),
+                    role="last_frame",
+                )
+            )
+        for i, image in enumerate(reference_images.values(), 1):
+            content.append(
+                Hailuo03ImageContent(
+                    image_url=Hailuo03ImageContentUrl(
+                        url=(
+                            await upload_images_to_comfyapi(
+                                cls, image, max_images=1, wait_label=f"Uploading image {i}"
+                            )
+                        )[0],
+                    ),
+                    role="reference_image",
+                )
+            )
+        for i, ref_video in enumerate(reference_videos.values(), 1):
+            content.append(
+                Hailuo03VideoContent(
+                    video_url=Hailuo03VideoContentUrl(
+                        url=await upload_video_to_comfyapi(cls, ref_video, wait_label=f"Uploading video {i}"),
+                    ),
+                )
+            )
+        for audio in reference_audios.values():
+            content.append(
+                Hailuo03AudioContent(
+                    audio_url=Hailuo03AudioContentUrl(
+                        url=await upload_audio_to_comfyapi(
+                            cls,
+                            audio,
+                            container_format="mp3",
+                            codec_name="libmp3lame",
+                            mime_type="audio/mpeg",
+                        ),
+                    ),
+                )
+            )
+
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=HAILUO_03_REGENERATION_ENDPOINT, method="POST"),
+            response_model=Hailuo03TaskCreationResponse,
+            data=Hailuo03RegenerationRequest(
+                model=HAILUO_03_MODELS[model["model"]],
+                content=content,
+                resolution=model["resolution"],
+                aigc_watermark=watermark,
+            ),
+        )
+        task_result = await poll_op(
+            cls,
+            ApiEndpoint(path=f"{HAILUO_03_QUERY_ENDPOINT}/{response.task_id}"),
+            response_model=Hailuo03TaskQueryResponse,
+            status_extractor=lambda r: r.task.status,
+            failed_statuses=HAILUO_03_FAILED_STATUSES,
+            poll_interval=10,
+        )
+        video_url = task_result.task.content.url if task_result.task.content else None
+        if not video_url:
+            raise Exception(f"No video URL in the response: {task_result.model_dump()}")
+        return IO.NodeOutput(await download_url_to_video_output(video_url))
+
+
 class MinimaxExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
@@ -950,6 +1541,8 @@ class MinimaxExtension(ComfyExtension):
             MinimaxHailuo03TextToVideoNode,
             MinimaxHailuo03FirstLastFrameNode,
             MinimaxHailuo03ReferenceNode,
+            MinimaxHailuo03ContextIRNode,
+            MinimaxHailuo03RegenerateNode,
         ]
 
 

comfy_api_nodes/nodes_openai.py

@@ -494,9 +494,9 @@ class OpenAIGPTImage1(IO.ComfyNode):
                       "high":   [0.133, 0.22]
                     },
                     "gpt-image-2": {
-                      "low":    [0.0048, 0.019],
-                      "medium": [0.041, 0.168],
-                      "high":   [0.165, 0.67]
+                      "low":    [0.0058, 0.0228],
+                      "medium": [0.0492, 0.2016],
+                      "high":   [0.198, 0.804]
                     }
                   };
                   $range := $lookup($lookup($ranges, widgets.model), widgets.quality);
@@ -792,9 +792,9 @@ class OpenAIGPTImageNodeV2(IO.ComfyNode):
                       "high":   [0.133, 0.22]
                     },
                     "gpt-image-2": {
-                      "low":    [0.0048, 0.019],
-                      "medium": [0.041, 0.168],
-                      "high":   [0.165, 0.67]
+                      "low":    [0.0058, 0.0228],
+                      "medium": [0.0492, 0.2016],
+                      "high":   [0.198, 0.804]
                     }
                   };
                   $range := $lookup($lookup($ranges, widgets.model), $lookup(widgets, "model.quality"));

comfy_api_nodes/nodes_qwen.py

@@ -0,0 +1,442 @@
+import math
+import re
+
+import torch
+from typing_extensions import override
+
+from comfy_api.latest import IO, ComfyExtension
+from comfy_api_nodes.apis.qwen import (
+    QwenImageContentItem,
+    QwenImageGenerationRequest,
+    QwenImageGenerationResponse,
+    QwenImageInputField,
+    QwenImageMessage,
+    QwenImageParametersField,
+)
+from comfy_api_nodes.util import (
+    ApiEndpoint,
+    download_url_to_image_tensor,
+    sync_op,
+    tensor_to_base64_string,
+    validate_string,
+)
+
+GENERATION_PATH = "/proxy/qwen/api/v1/services/aigc/multimodal-generation/generation"
+QWEN_IMAGE_MODELS = ["qwen-image-3.0-pro", "qwen-image-3.0"]
+MIN_AREA = 262144  # 512*512
+MAX_AREA = 6553600  # 2560*2560
+MAX_ASPECT = 8  # the API allows aspect ratios from 1:8 to 8:1
+MAX_INPUT_BYTES = 10 * 1024 * 1024  # the API rejects decoded input images over 10MB
+
+_IMAGE_REF_RE = re.compile(r"@image(?P<idx>\d*)(?!\w)", re.IGNORECASE | re.ASCII)
+
+
+def _resolve_image_refs(prompt: str, total_images: int) -> str:
+    """Rewrite @Image1-style references (shared partner-node syntax, 1-based; an unnumbered
+    @image means the first image) into the plain 'Image N' wording the model resolves
+    natively. A tag counts only at a word boundary or right after a previous tag, so
+    adjacent tags like '@Image1@Image2' all resolve while addresses like user@image1.com
+    pass through untouched."""
+    parts = []
+    pos = 0
+    prev_end = -1
+    for match in _IMAGE_REF_RE.finditer(prompt):
+        start = match.start()
+        if start > 0 and start != prev_end and (prompt[start - 1].isalnum() or prompt[start - 1] == "_"):
+            continue
+        idx = int(match.group("idx") or 1)
+        if not 1 <= idx <= total_images:
+            raise ValueError(
+                f"The prompt references @Image{idx}, but only {total_images} reference images "
+                f"are connected (a batched input counts once per image)."
+            )
+        parts.append(prompt[pos:start])
+        parts.append(f"Image {idx}")
+        pos = match.end()
+        prev_end = match.end()
+    parts.append(prompt[pos:])
+    return "".join(parts)
+
+
+def _validate_size(width: int, height: int) -> None:
+    if not MIN_AREA <= width * height <= MAX_AREA:
+        raise ValueError(
+            f"Image area must be between {MIN_AREA} (512x512) and {MAX_AREA} (2560x2560) pixels; "
+            f"got {width}x{height} = {width * height}."
+        )
+    if width > MAX_ASPECT * height or height > MAX_ASPECT * width:
+        raise ValueError(f"Aspect ratio must be between 1:8 and 8:1; got {width}x{height}.")
+
+
+def _fit_to_size(width: int, height: int) -> tuple[int, int]:
+    """Scale dimensions into the supported pixel area and 1:8..8:1 aspect range, preserving
+    the aspect ratio where possible."""
+    if width > MAX_ASPECT * height:
+        height = math.ceil(width / MAX_ASPECT)
+    elif height > MAX_ASPECT * width:
+        width = math.ceil(height / MAX_ASPECT)
+    area = width * height
+    if area < MIN_AREA:
+        scale = math.sqrt(MIN_AREA / area)
+        width, height = math.ceil(width * scale), math.ceil(height * scale)
+    elif area > MAX_AREA:
+        scale = math.sqrt(MAX_AREA / area)
+        width, height = math.floor(width * scale), math.floor(height * scale)
+    # rounding can push the ratio a hair past the limit; trimming only ever shrinks the area
+    return min(width, MAX_ASPECT * height), min(height, MAX_ASPECT * width)
+
+
+def _image_data_uri(image: torch.Tensor) -> str:
+    """PNG data URI of an RGB view of the image, downscaled to <=2048x2048; falls back to
+    JPEG when the PNG exceeds the API's decoded-size cap (e.g. noisy, incompressible images)."""
+    image = image[..., :3]
+    b64 = tensor_to_base64_string(image, total_pixels=2048 * 2048)
+    if len(b64) * 3 > MAX_INPUT_BYTES * 4:
+        return "data:image/jpeg;base64," + tensor_to_base64_string(
+            image, total_pixels=2048 * 2048, mime_type="image/jpeg"
+        )
+    return "data:image/png;base64," + b64
+
+
+async def _download_result_images(response: QwenImageGenerationResponse) -> torch.Tensor:
+    if not response.output:
+        raise Exception(f"An unknown error occurred: {response.code} - {response.message}")
+    urls = [
+        item.image
+        for choice in response.output.choices
+        if choice.message
+        for item in choice.message.content
+        if item.image
+    ]
+    if not urls:
+        raise Exception(f"The response contains no images: {response.code} - {response.message}")
+    return torch.cat([await download_url_to_image_tensor(url) for url in urls])
+
+
+def _size_inputs() -> list[IO.Int.Input]:
+    return [
+        IO.Int.Input(
+            "width",
+            default=1024,
+            min=256,
+            max=2560,
+            step=16,
+            tooltip="The total pixel area must be between 512x512 and 2560x2560; "
+            "any aspect ratio within that area works.",
+        ),
+        IO.Int.Input(
+            "height",
+            default=1024,
+            min=256,
+            max=2560,
+            step=16,
+            tooltip="The total pixel area must be between 512x512 and 2560x2560; "
+            "any aspect ratio within that area works.",
+        ),
+    ]
+
+
+def _t2i_model_option(model_id: str) -> IO.DynamicCombo.Option:
+    return IO.DynamicCombo.Option(
+        model_id,
+        [
+            IO.String.Input(
+                "prompt",
+                multiline=True,
+                default="",
+                tooltip="Prompt describing the image. Supports English and Chinese.",
+            ),
+            IO.String.Input(
+                "negative_prompt",
+                multiline=True,
+                default="",
+                tooltip="Negative prompt describing what to avoid.",
+            ),
+            *_size_inputs(),
+        ],
+    )
+
+
+def _edit_model_option(model_id: str) -> IO.DynamicCombo.Option:
+    return IO.DynamicCombo.Option(
+        model_id,
+        [
+            IO.Autogrow.Input(
+                "images",
+                template=IO.Autogrow.TemplateNames(
+                    IO.Image.Input("image"),
+                    names=["image_1", "image_2", "image_3"],
+                    min=1,
+                ),
+                tooltip="1-3 reference images. Refer to them in the prompt as @Image1, @Image2, "
+                "@Image3, numbered in input order; a batched input counts once per image.",
+            ),
+            IO.String.Input(
+                "prompt",
+                multiline=True,
+                default="",
+                tooltip="Editing instructions. Supports English and Chinese, "
+                "and @Image1-style references to the input images.",
+            ),
+            IO.String.Input(
+                "negative_prompt",
+                multiline=True,
+                default="",
+                tooltip="Negative prompt describing what to avoid.",
+            ),
+        ],
+    )
+
+
+class QwenImageTextToImageApi(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="QwenImageTextToImageApi",
+            display_name="Qwen Image 3 Text to Image",
+            category="partner/image/Qwen",
+            description="Generates images from a text prompt using the Qwen-Image 3.0 models.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[_t2i_model_option(model_id) for model_id in QWEN_IMAGE_MODELS],
+                    tooltip="Model to use.",
+                ),
+                IO.Int.Input(
+                    "n",
+                    default=1,
+                    min=1,
+                    max=6,
+                    display_mode=IO.NumberDisplay.number,
+                    tooltip="Number of images to generate, returned as a batch.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation.",
+                ),
+                IO.Boolean.Input(
+                    "prompt_extend",
+                    default=True,
+                    tooltip="Whether to enhance the prompt with AI assistance.",
+                    advanced=True,
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add an AI-generated watermark to the result.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Image.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=["model", "model.width", "model.height", "n"]),
+                expr="""
+                (
+                  $isPro := widgets.model = "qwen-image-3.0-pro";
+                  $area := $lookup(widgets, "model.width") * $lookup(widgets, "model.height");
+                  $rate := $isPro ? ($area > 2250000 ? 0.10725 : 0.0572) : 0.0429;
+                  {"type":"usd","usd": $rate * widgets.n}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        n: int = 1,
+        seed: int = 42,
+        prompt_extend: bool = True,
+        watermark: bool = False,
+    ):
+        validate_string(model["prompt"], strip_whitespace=False, min_length=1)
+        width, height = model["width"], model["height"]
+        _validate_size(width, height)
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=GENERATION_PATH, method="POST"),
+            response_model=QwenImageGenerationResponse,
+            data=QwenImageGenerationRequest(
+                model=model["model"],
+                input=QwenImageInputField(
+                    messages=[QwenImageMessage(content=[QwenImageContentItem(text=model["prompt"])])],
+                ),
+                parameters=QwenImageParametersField(
+                    size=f"{width}*{height}",
+                    n=n,
+                    seed=seed,
+                    prompt_extend=prompt_extend,
+                    watermark=watermark,
+                    negative_prompt=model["negative_prompt"] or None,
+                ),
+            ),
+        )
+        return IO.NodeOutput(await _download_result_images(response))
+
+
+class QwenImageEditApi(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="QwenImageEditApi",
+            display_name="Qwen Image 3 Edit",
+            category="partner/image/Qwen",
+            description="Edits or combines up to 3 reference images guided by a text prompt "
+            "using the Qwen-Image 3.0 models.",
+            inputs=[
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[_edit_model_option(model_id) for model_id in QWEN_IMAGE_MODELS],
+                    tooltip="Model to use.",
+                ),
+                IO.DynamicCombo.Input(
+                    "size",
+                    options=[
+                        IO.DynamicCombo.Option("match input", []),
+                        IO.DynamicCombo.Option("auto", []),
+                        IO.DynamicCombo.Option("custom", _size_inputs()),
+                    ],
+                    tooltip="Output resolution. 'match input' reuses the first reference image's size, "
+                    "'auto' lets the model pick a size with the same aspect ratio, "
+                    "'custom' sets an explicit width and height.",
+                ),
+                IO.Int.Input(
+                    "n",
+                    default=1,
+                    min=1,
+                    max=6,
+                    display_mode=IO.NumberDisplay.number,
+                    tooltip="Number of images to generate, returned as a batch.",
+                ),
+                IO.Int.Input(
+                    "seed",
+                    default=42,
+                    min=0,
+                    max=2147483647,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    control_after_generate=True,
+                    tooltip="Seed to use for generation.",
+                ),
+                IO.Boolean.Input(
+                    "prompt_extend",
+                    default=True,
+                    tooltip="Whether to enhance the prompt with AI assistance.",
+                    advanced=True,
+                ),
+                IO.Boolean.Input(
+                    "watermark",
+                    default=False,
+                    tooltip="Whether to add an AI-generated watermark to the result.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Image.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=["model", "size", "size.width", "size.height", "n"],
+                    input_groups=["model.images"],
+                ),
+                expr="""
+                (
+                  $isPro := widgets.model = "qwen-image-3.0-pro";
+                  $mode := widgets.size;
+                  $count := $max([$lookup(inputGroups, "model.images"), 1]);
+                  $inputCost := 0.00429 * $count;
+                  $area := $mode = "custom"
+                    ? $lookup(widgets, "size.width") * $lookup(widgets, "size.height") : 0;
+                  $customRate := $area > 2250000 ? 0.10725 : 0.0572;
+                  $isPro and $mode != "custom"
+                    ? {"type":"range_usd",
+                       "min_usd": 0.0572 * widgets.n + $inputCost,
+                       "max_usd": 0.10725 * widgets.n + $inputCost}
+                    : {"type":"usd",
+                       "usd": ($isPro ? $customRate : 0.0429) * widgets.n + $inputCost}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        model: dict,
+        size: dict,
+        n: int = 1,
+        seed: int = 42,
+        prompt_extend: bool = True,
+        watermark: bool = False,
+    ):
+        validate_string(model["prompt"], strip_whitespace=False, min_length=1)
+        reference_images = [image for key in model["images"] for image in model["images"][key]]
+        if len(reference_images) > 3:
+            raise ValueError(
+                f"A maximum of 3 reference images is supported; got {len(reference_images)} "
+                f"(a batched input counts once per image)."
+            )
+        prompt = _resolve_image_refs(model["prompt"], len(reference_images))
+        if size["size"] == "custom":
+            _validate_size(size["width"], size["height"])
+            size_str = f"{size['width']}*{size['height']}"
+        elif size["size"] == "match input":
+            height, width = reference_images[0].shape[0], reference_images[0].shape[1]
+            width, height = _fit_to_size(width, height)
+            size_str = f"{width}*{height}"
+        else:  # auto: the API picks a size preserving the input aspect ratio (1.9-4.2 MP)
+            size_str = None
+        content = [QwenImageContentItem(image=_image_data_uri(image)) for image in reference_images]
+        content.append(QwenImageContentItem(text=prompt))
+        response = await sync_op(
+            cls,
+            ApiEndpoint(path=GENERATION_PATH, method="POST"),
+            response_model=QwenImageGenerationResponse,
+            data=QwenImageGenerationRequest(
+                model=model["model"],
+                input=QwenImageInputField(messages=[QwenImageMessage(content=content)]),
+                parameters=QwenImageParametersField(
+                    size=size_str,
+                    n=n,
+                    seed=seed,
+                    prompt_extend=prompt_extend,
+                    watermark=watermark,
+                    negative_prompt=model["negative_prompt"] or None,
+                ),
+            ),
+        )
+        return IO.NodeOutput(await _download_result_images(response))
+
+
+class QwenApiExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[IO.ComfyNode]]:
+        return [
+            QwenImageTextToImageApi,
+            QwenImageEditApi,
+        ]
+
+
+async def comfy_entrypoint() -> QwenApiExtension:
+    return QwenApiExtension()

comfy_api_nodes/nodes_recraft.py

@@ -27,6 +27,7 @@ from comfy_api_nodes.util import (
     ApiEndpoint,
     bytesio_to_image_tensor,
     download_url_as_bytesio,
+    pad_images_to_common_channels,
     resize_mask_to_image,
     sync_op,
     tensor_to_bytesio,
@@ -621,7 +622,7 @@ class RecraftImageToImageNode(IO.ComfyNode):
                 images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftImageInpaintingNode(IO.ComfyNode):
@@ -723,7 +724,7 @@ class RecraftImageInpaintingNode(IO.ComfyNode):
                 images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftTextToVectorNode(IO.ComfyNode):
@@ -954,7 +955,7 @@ class RecraftReplaceBackgroundNode(IO.ComfyNode):
             images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftRemoveBackgroundNode(IO.ComfyNode):
@@ -995,7 +996,7 @@ class RecraftRemoveBackgroundNode(IO.ComfyNode):
                 image=image[i],
                 path="/proxy/recraft/images/removeBackground",
             )
-            images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
+            images.append(torch.cat([bytesio_to_image_tensor(x, mode="RGBA") for x in sub_bytes], dim=0))
             pbar.update(1)
 
         images_tensor = torch.cat(images, dim=0)
@@ -1047,7 +1048,7 @@ class RecraftCrispUpscaleNode(IO.ComfyNode):
             images.append(torch.cat([bytesio_to_image_tensor(x) for x in sub_bytes], dim=0))
             pbar.update(1)
 
-        return IO.NodeOutput(torch.cat(images, dim=0))
+        return IO.NodeOutput(torch.cat(pad_images_to_common_channels(images), dim=0))
 
 
 class RecraftCreativeUpscaleNode(RecraftCrispUpscaleNode):

comfy_api_nodes/nodes_reve.py

@@ -137,6 +137,7 @@ class ReveImageCreateNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=IO.PriceBadge(
                 depends_on=IO.PriceBadgeDepends(
                     widgets=["upscale", "upscale.upscale_factor"],
@@ -228,6 +229,7 @@ class ReveImageEditNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=IO.PriceBadge(
                 depends_on=IO.PriceBadgeDepends(
                     widgets=["model", "upscale", "upscale.upscale_factor"],
@@ -337,6 +339,7 @@ class ReveImageRemixNode(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
             price_badge=IO.PriceBadge(
                 depends_on=IO.PriceBadgeDepends(
                     widgets=["model", "upscale", "upscale.upscale_factor"],

comfy_api_nodes/nodes_topaz.py

@@ -12,6 +12,7 @@ from comfy_api_nodes.apis.topaz import (
     ImageAsyncTaskResponse,
     ImageDownloadResponse,
     ImageEnhanceRequest,
+    ImageEnhanceRequestV2,
     ImageStatusResponse,
     OutputInformationVideo,
     Resolution,
@@ -51,7 +52,7 @@ class TopazImageEnhance(IO.ComfyNode):
     def define_schema(cls):
         return IO.Schema(
             node_id="TopazImageEnhance",
-            display_name="Topaz Image Enhance",
+            display_name="Topaz Image Enhance (Legacy)",
             category="partner/image/Topaz",
             description="Industry-standard upscaling and image enhancement.",
             inputs=[
@@ -162,6 +163,7 @@ class TopazImageEnhance(IO.ComfyNode):
                 IO.Hidden.unique_id,
             ],
             is_api_node=True,
+            is_deprecated=True,
         )
 
     @classmethod
@@ -229,6 +231,355 @@ class TopazImageEnhance(IO.ComfyNode):
         return IO.NodeOutput(await download_url_to_image_tensor(results.download_url))
 
 
+class TopazImageEnhanceV2(IO.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return IO.Schema(
+            node_id="TopazImageEnhanceV2",
+            display_name="Topaz Image Enhance",
+            category="partner/image/Topaz",
+            description="Industry-standard upscaling and image enhancement.",
+            inputs=[
+                IO.Image.Input("image"),
+                IO.DynamicCombo.Input(
+                    "model",
+                    options=[
+                        IO.DynamicCombo.Option(
+                            "Reimagine",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Optional text prompt for creative upscaling guidance.",
+                                ),
+                                IO.Int.Input(
+                                    "creativity",
+                                    default=3,
+                                    min=1,
+                                    max=9,
+                                    step=1,
+                                    display_mode=IO.NumberDisplay.slider,
+                                ),
+                                IO.Combo.Input(
+                                    "subject_detection",
+                                    options=["All", "Foreground", "Background"],
+                                    advanced=True,
+                                ),
+                                IO.Boolean.Input(
+                                    "face_enhancement",
+                                    default=True,
+                                    tooltip="Enhance faces (if present) during processing.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "face_enhancement_creativity",
+                                    default=0.0,
+                                    min=0.0,
+                                    max=1.0,
+                                    step=0.01,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Set the creativity level for face enhancement.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "face_enhancement_strength",
+                                    default=1.0,
+                                    min=0.0,
+                                    max=1.0,
+                                    step=0.01,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Controls how sharp enhanced faces are relative to the background.",
+                                    advanced=True,
+                                ),
+                                IO.Boolean.Input(
+                                    "face_preservation",
+                                    default=True,
+                                    tooltip="Preserve subjects' facial identity.",
+                                    advanced=True,
+                                ),
+                                IO.Boolean.Input(
+                                    "color_preservation",
+                                    default=True,
+                                    tooltip="Preserve the original colors.",
+                                    advanced=True,
+                                ),
+                                IO.Boolean.Input(
+                                    "crop_to_fill",
+                                    default=False,
+                                    tooltip="By default, the image is letterboxed when the output aspect "
+                                    "ratio differs. Enable to crop the image to fill the output dimensions.",
+                                    advanced=True,
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "Bloom 2",
+                            [
+                                IO.String.Input(
+                                    "prompt",
+                                    multiline=True,
+                                    default="",
+                                    tooltip="Optional text prompt for generation. "
+                                    "Leave empty to auto-generate a prompt from the input image.",
+                                ),
+                                IO.Int.Input(
+                                    "creativity",
+                                    default=3,
+                                    min=1,
+                                    max=9,
+                                    step=1,
+                                    display_mode=IO.NumberDisplay.slider,
+                                    tooltip="1 is restrained enhancement, 9 is pronounced reinterpretation "
+                                    "with newly generated detail.",
+                                ),
+                                IO.Int.Input(
+                                    "seed",
+                                    default=2,
+                                    min=1,
+                                    max=2000,
+                                    control_after_generate=True,
+                                    tooltip="Seed for reproducible generation.",
+                                ),
+                                IO.Boolean.Input(
+                                    "color_preservation",
+                                    default=True,
+                                    tooltip="Preserve the original colors.",
+                                    advanced=True,
+                                ),
+                                IO.Boolean.Input(
+                                    "grain",
+                                    default=False,
+                                    tooltip="Add grain to the output image.",
+                                    advanced=True,
+                                ),
+                                IO.Combo.Input(
+                                    "grain_model",
+                                    options=["silver", "gaussian", "grey"],
+                                    tooltip="Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "grain_strength",
+                                    default=0.5,
+                                    min=0.0,
+                                    max=1.0,
+                                    step=0.01,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Strength of the grain effect. Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "grain_size",
+                                    default=1.0,
+                                    min=1.0,
+                                    max=5.0,
+                                    step=0.1,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Size of the grain particles. Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "grain_density",
+                                    default=0.5,
+                                    min=0.0,
+                                    max=1.0,
+                                    step=0.01,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Intensity of the grain effect. Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                            ],
+                        ),
+                        IO.DynamicCombo.Option(
+                            "Wonder 3.5",
+                            [
+                                IO.Combo.Input(
+                                    "enhancement_strength",
+                                    options=["low", "medium", "high"],
+                                    default="high",
+                                    tooltip="Enhancement level for varying input conditions.",
+                                ),
+                                IO.Boolean.Input(
+                                    "grain",
+                                    default=False,
+                                    tooltip="Add grain to the output image.",
+                                    advanced=True,
+                                ),
+                                IO.Combo.Input(
+                                    "grain_model",
+                                    options=["silver", "gaussian", "grey"],
+                                    tooltip="Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "grain_strength",
+                                    default=0.5,
+                                    min=0.0,
+                                    max=1.0,
+                                    step=0.01,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Strength of the grain effect. Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "grain_size",
+                                    default=1.0,
+                                    min=1.0,
+                                    max=5.0,
+                                    step=0.1,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Size of the grain particles. Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                                IO.Float.Input(
+                                    "grain_density",
+                                    default=0.5,
+                                    min=0.0,
+                                    max=1.0,
+                                    step=0.01,
+                                    display_mode=IO.NumberDisplay.number,
+                                    tooltip="Intensity of the grain effect. Is ignored if grain is disabled.",
+                                    advanced=True,
+                                ),
+                            ],
+                        ),
+                    ],
+                ),
+                IO.Int.Input(
+                    "output_width",
+                    default=0,
+                    min=0,
+                    max=32000,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    optional=True,
+                    tooltip="Zero value means to calculate automatically (usually it will be original size "
+                    "or scaled proportionally to output_height if specified). "
+                    "Wonder 3.5 supports upscale factors from 1x to 6x only. "
+                    "Bloom 2 and Wonder 3.5 preserve the input aspect ratio and treat the "
+                    "requested size as a target.",
+                    advanced=True,
+                ),
+                IO.Int.Input(
+                    "output_height",
+                    default=0,
+                    min=0,
+                    max=32000,
+                    step=1,
+                    display_mode=IO.NumberDisplay.number,
+                    optional=True,
+                    tooltip="Zero value means to output in the same height as original or scaled "
+                    "proportionally to output_width if specified. "
+                    "Wonder 3.5 supports upscale factors from 1x to 6x only. "
+                    "Bloom 2 and Wonder 3.5 preserve the input aspect ratio and treat the "
+                    "requested size as a target.",
+                    advanced=True,
+                ),
+            ],
+            outputs=[
+                IO.Image.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=["model"]),
+                expr="""
+                (
+                  $usdPer8Mp := $lookup(
+                    {"reimagine": 0.32, "bloom 2": 0.4576, "wonder 3.5": 0.1144},
+                    $lookup(widgets, "model")
+                  );
+                  {"type":"usd","usd": $usdPer8Mp, "format": {"suffix": "/8MP", "approximate": true}}
+                )
+                """,
+            ),
+        )
+
+    @classmethod
+    async def execute(
+        cls,
+        image: Input.Image,
+        model: dict,
+        output_width: int = 0,
+        output_height: int = 0,
+    ) -> IO.NodeOutput:
+        if get_number_of_images(image) != 1:
+            raise ValueError("Only one input image is supported.")
+        model_choice = model["model"]
+        download_url = await upload_images_to_comfyapi(
+            cls, image, max_images=1, mime_type="image/png", total_pixels=4096 * 4096
+        )
+        request = ImageEnhanceRequestV2(
+            model=model_choice,
+            source_url=download_url[0],
+            output_width=output_width if output_width else None,
+            output_height=output_height if output_height else None,
+        )
+        if model_choice == "Reimagine":
+            request.prompt = model["prompt"]
+            request.creativity = model["creativity"]
+            request.subject_detection = model["subject_detection"]
+            request.face_enhancement = model["face_enhancement"]
+            request.face_enhancement_creativity = model["face_enhancement_creativity"]
+            request.face_enhancement_strength = model["face_enhancement_strength"]
+            request.face_preservation = str(model["face_preservation"]).lower()
+            request.color_preservation = str(model["color_preservation"]).lower()
+            request.crop_to_fill = model["crop_to_fill"]
+        elif model_choice == "Bloom 2":
+            prompt = model["prompt"].strip()
+            if prompt:
+                request.prompt = prompt
+                request.autoprompt = "false"
+            else:
+                request.autoprompt = "true"
+            request.creativity = model["creativity"]
+            request.seed = model["seed"]
+            request.color_preservation = str(model["color_preservation"]).lower()
+            if model["grain"]:
+                request.grain = "true"
+                request.grain_model = model["grain_model"]
+                request.grain_strength = model["grain_strength"]
+                request.grain_size = model["grain_size"]
+                request.grain_density = model["grain_density"]
+        else:
+            request.enhancement_strength = model["enhancement_strength"]
+            if model["grain"]:
+                request.grain = "true"
+                request.grain_model = model["grain_model"]
+                request.grain_strength = model["grain_strength"]
+                request.grain_size = model["grain_size"]
+                request.grain_density = model["grain_density"]
+        initial_response = await sync_op(
+            cls,
+            ApiEndpoint(path="/proxy/topaz/image/v1/enhance-gen/async", method="POST"),
+            response_model=ImageAsyncTaskResponse,
+            data=request,
+            content_type="multipart/form-data",
+        )
+        await poll_op(
+            cls,
+            poll_endpoint=ApiEndpoint(path=f"/proxy/topaz/image/v1/status/{initial_response.process_id}"),
+            response_model=ImageStatusResponse,
+            status_extractor=lambda x: x.status,
+            progress_extractor=lambda x: getattr(x, "progress", 0),
+            price_extractor=lambda x: x.credits * (0.08 if model_choice == "Reimagine" else 0.1144),
+            poll_interval=8.0,
+            estimated_duration=60,
+        )
+        results = await sync_op(
+            cls,
+            ApiEndpoint(path=f"/proxy/topaz/image/v1/download/{initial_response.process_id}"),
+            response_model=ImageDownloadResponse,
+            monitor_progress=False,
+        )
+        return IO.NodeOutput(await download_url_to_image_tensor(results.download_url))
+
+
 class TopazVideoEnhance(IO.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -818,6 +1169,7 @@ class TopazExtension(ComfyExtension):
     async def get_node_list(self) -> list[type[IO.ComfyNode]]:
         return [
             TopazImageEnhance,
+            TopazImageEnhanceV2,
             TopazVideoEnhance,
             TopazVideoEnhanceV2,
         ]

comfy_api_nodes/util/__init__.py

@@ -18,6 +18,7 @@ from .conversions import (
     downscale_image_tensor_by_max_side,
     downscale_video_to_max_pixels,
     image_tensor_pair_to_batch,
+    pad_images_to_common_channels,
     pil_to_bytesio,
     resize_mask_to_image,
     tensor_to_base64_string,
@@ -92,6 +93,7 @@ __all__ = [
     "downscale_image_tensor_by_max_side",
     "downscale_video_to_max_pixels",
     "image_tensor_pair_to_batch",
+    "pad_images_to_common_channels",
     "pil_to_bytesio",
     "resize_mask_to_image",
     "tensor_to_base64_string",

comfy_api_nodes/util/conversions.py

@@ -16,12 +16,14 @@ from comfy_api.latest import Input, InputImpl, Types
 from ._helpers import mimetype_to_extension
 
 
-def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch.Tensor:
+def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str | None = None) -> torch.Tensor:
     """Converts image data from BytesIO to a torch.Tensor.
 
     Args:
         image_bytesio: BytesIO object containing the image data.
-        mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA").
+        mode: The PIL mode to convert the image to (e.g., "RGB", "RGBA"). Defaults
+            to RGBA when the decoded image carries transparency and RGB when it
+            does not, so an API that returns no alpha does not get an opaque one.
 
     Returns:
         A torch.Tensor representing the image (1, H, W, C).
@@ -31,6 +33,8 @@ def bytesio_to_image_tensor(image_bytesio: BytesIO, mode: str = "RGBA") -> torch
         ValueError: If the specified mode is invalid.
     """
     image = Image.open(image_bytesio)
+    if mode is None:
+        mode = "RGBA" if "A" in image.getbands() or "transparency" in image.info else "RGB"
     image = image.convert(mode)
     image_array = np.array(image).astype(np.float32) / 255.0
     return torch.from_numpy(image_array).unsqueeze(0)
@@ -53,6 +57,17 @@ def image_tensor_pair_to_batch(image1: torch.Tensor, image2: torch.Tensor) -> to
     return torch.cat((image1, image2), dim=0)
 
 
+def pad_images_to_common_channels(images: list[torch.Tensor]) -> list[torch.Tensor]:
+    """Pads [B, H, W, C] image tensors with opaque alpha so they all share the largest channel count."""
+    channels = max(image.shape[-1] for image in images)
+    return [
+        torch.nn.functional.pad(image, (0, channels - image.shape[-1]), value=1.0)
+        if image.shape[-1] < channels
+        else image
+        for image in images
+    ]
+
+
 def tensor_to_bytesio(
     image: torch.Tensor,
     *,

comfy_execution/jobs.py

@@ -197,6 +197,7 @@ def normalize_queue_item(item: tuple, status: str) -> dict:
         'priority': priority,
         'create_time': create_time,
         'outputs_count': 0,
+        'previewable_outputs_count': 0,
         'workflow_id': workflow_id,
     })
 
@@ -215,6 +216,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
 
     outputs = history_item.get('outputs', {})
     outputs_count, preview_output = get_outputs_summary(outputs)
+    previewable_outputs_count = count_previewable_outputs(outputs)
 
     execution_error = None
     execution_start_time = None
@@ -251,6 +253,7 @@ def normalize_history_item(prompt_id: str, history_item: dict, include_outputs:
         'execution_end_time': execution_end_time,
         'execution_error': execution_error,
         'outputs_count': outputs_count,
+        'previewable_outputs_count': previewable_outputs_count,
         'preview_output': preview_output,
         'workflow_id': workflow_id,
     })
@@ -345,6 +348,33 @@ def get_outputs_summary(outputs: dict) -> tuple[int, Optional[dict]]:
     return count, preview_output or fallback_preview or text_file_fallback or text_fallback
 
 
+def count_previewable_outputs(outputs: dict) -> int:
+    """
+    Count only outputs that would actually render in the expanded asset view,
+    i.e. items is_previewable() accepts (image/video/audio/3D/text). Kept
+    separate from get_outputs_summary()'s outputs_count, which counts every
+    output item regardless of media type, so a job with a non-previewable
+    saved file alongside real media (e.g. SaveLatent's .latent output next to
+    a SaveImage output) doesn't inflate the Media Assets badge beyond what
+    the expanded view shows.
+    """
+    count = 0
+    for node_outputs in outputs.values():
+        if not isinstance(node_outputs, dict):
+            continue
+        for media_type, items in node_outputs.items():
+            if media_type == 'animated' or not isinstance(items, list):
+                continue
+            for item in items:
+                if not isinstance(item, dict):
+                    item = normalize_output_item(item)
+                    if item is None:
+                        continue
+                if is_previewable(media_type, item):
+                    count += 1
+    return count
+
+
 def apply_sorting(jobs: list[dict], sort_by: str, sort_order: str) -> list[dict]:
     """Sort jobs list by specified field and order."""
     reverse = (sort_order == 'desc')

comfy_extras/compositor_blend.py

@@ -0,0 +1,331 @@
+import math
+from typing import NamedTuple, Optional, Union
+
+import numpy as np
+
+EPSILON = 1e-6
+
+LUM_R = 0.2224884
+LUM_G = 0.71690369
+LUM_B = 0.06060791
+
+ArrayLike = Union[np.ndarray, float]
+
+
+def srgb_to_linear(c: ArrayLike) -> np.ndarray:
+    c = np.asarray(c, dtype=np.float32)
+    high = ((np.maximum(c, 0.0) + 0.055) / 1.055) ** 2.4
+    return np.where(c <= 0.04045, c / 12.92, high).astype(np.float32)
+
+
+def linear_to_srgb(c: ArrayLike) -> np.ndarray:
+    c = np.asarray(c, dtype=np.float32)
+    high = 1.055 * np.maximum(c, 0.0) ** (1.0 / 2.4) - 0.055
+    return np.where(c <= 0.0031308, 12.92 * c, high).astype(np.float32)
+
+
+def luminance(rgb: np.ndarray) -> np.ndarray:
+    return rgb[..., 0] * LUM_R + rgb[..., 1] * LUM_G + rgb[..., 2] * LUM_B
+
+
+def safe_div(a: ArrayLike, b: ArrayLike) -> np.ndarray:
+    a, b = np.broadcast_arrays(
+        np.asarray(a, dtype=np.float32), np.asarray(b, dtype=np.float32)
+    )
+    out = np.zeros(b.shape, dtype=np.float32)
+    np.divide(a, b, out=out, where=np.abs(b) >= EPSILON)
+    return out
+
+
+CHANNEL_BLEND = {
+    "normal": lambda i, l: l,
+    "multiply": lambda i, l: i * l,
+    "screen": lambda i, l: 1 - (1 - i) * (1 - l),
+    "overlay": lambda i, l: np.where(i < 0.5, 2 * i * l, 1 - 2 * (1 - l) * (1 - i)),
+    "darken": lambda i, l: np.minimum(i, l),
+    "lighten": lambda i, l: np.maximum(i, l),
+    "color-dodge": lambda i, l: np.where(
+        i <= 0,
+        0.0,
+        np.where(1 - l <= EPSILON, 1.0, np.minimum(safe_div(i, 1 - l), 1.0)),
+    ),
+    "color-burn": lambda i, l: np.where(
+        i >= 1,
+        1.0,
+        np.where(l <= EPSILON, 0.0, 1 - np.minimum(safe_div(1 - i, l), 1.0)),
+    ),
+    "hard-light": lambda i, l: np.where(
+        l > 0.5,
+        np.minimum(1 - (1 - i) * (1 - (l - 0.5) * 2), 1),
+        np.minimum(i * (l * 2), 1),
+    ),
+    "soft-light": lambda i, l: (1 - i) * (i * l) + i * (1 - (1 - i) * (1 - l)),
+    "difference": lambda i, l: np.abs(i - l),
+    "exclusion": lambda i, l: 0.5 - 2 * (i - 0.5) * (l - 0.5),
+    "linear-dodge": lambda i, l: i + l,
+    "linear-burn": lambda i, l: i + l - 1,
+    "vivid-light": lambda i, l: np.where(
+        l <= 0.5,
+        np.where(
+            i >= 1,
+            1.0,
+            np.where(
+                2 * l <= EPSILON,
+                0.0,
+                np.maximum(1 - safe_div(1 - i, 2 * l), 0.0),
+            ),
+        ),
+        np.where(
+            i <= 0,
+            0.0,
+            np.where(
+                2 * (1 - l) <= EPSILON,
+                1.0,
+                np.minimum(safe_div(i, 2 * (1 - l)), 1.0),
+            ),
+        ),
+    ),
+    "pin-light": lambda i, l: np.where(
+        l > 0.5, np.maximum(i, 2 * (l - 0.5)), np.minimum(i, 2 * l)
+    ),
+    "linear-light": lambda i, l: i + 2 * l - 1,
+    "hard-mix": lambda i, l: np.where(i + l < 1, 0.0, 1.0),
+    "subtract": lambda i, l: np.maximum(i - l, 0),
+    "divide": lambda i, l: np.clip(i / np.maximum(l, EPSILON), 0, 1),
+    "grain-extract": lambda i, l: i - l + 0.5,
+    "grain-merge": lambda i, l: i + l - 0.5,
+}
+
+
+def _blend_hue(i: np.ndarray, l: np.ndarray) -> np.ndarray:
+    src_min = l.min(axis=-1)
+    src_max = l.max(axis=-1)
+    src_delta = src_max - src_min
+    achromatic = src_delta <= EPSILON
+    dest_max = i.max(axis=-1)
+    dest_delta = dest_max - i.min(axis=-1)
+    dest_s = np.where(dest_max != 0, dest_delta / np.where(dest_max != 0, dest_max, 1), 0)
+    ratio = np.where(
+        achromatic, 0, dest_s * dest_max / np.where(achromatic, 1, src_delta)
+    )
+    offset = dest_max - src_max * ratio
+    return np.where(achromatic[..., None], i, l * ratio[..., None] + offset[..., None])
+
+
+def _blend_saturation(i: np.ndarray, l: np.ndarray) -> np.ndarray:
+    dest_max = i.max(axis=-1)
+    dest_delta = dest_max - i.min(axis=-1)
+    flat = dest_delta <= EPSILON
+    src_max = l.max(axis=-1)
+    src_delta = src_max - l.min(axis=-1)
+    src_s = np.where(src_max != 0, src_delta / np.where(src_max != 0, src_max, 1), 0)
+    ratio = np.where(flat, 0, src_s * dest_max / np.where(flat, 1, dest_delta))
+    offset = (1 - ratio) * dest_max
+    return np.where(
+        flat[..., None],
+        np.broadcast_to(dest_max[..., None], i.shape),
+        i * ratio[..., None] + offset[..., None],
+    )
+
+
+def _blend_color(i: np.ndarray, l: np.ndarray) -> np.ndarray:
+    dest_l = (i.min(axis=-1) + i.max(axis=-1)) / 2
+    src_l = (l.min(axis=-1) + l.max(axis=-1)) / 2
+    gray = (np.abs(src_l) <= EPSILON) | (np.abs(1 - src_l) <= EPSILON)
+    dest_high = dest_l > 0.5
+    src_high = src_l > 0.5
+    dl = np.minimum(dest_l, 1 - dest_l)
+    sl = np.minimum(src_l, 1 - src_l)
+    ratio = dl / np.where(gray, 1, sl)
+    offset = np.where(dest_high, 1 - 2 * dl, 0) + np.where(src_high, 2 * dl - ratio, 0)
+    return np.where(
+        gray[..., None],
+        np.broadcast_to(dest_l[..., None], i.shape),
+        l * ratio[..., None] + offset[..., None],
+    )
+
+
+def _blend_luminosity(i: np.ndarray, l: np.ndarray) -> np.ndarray:
+    # Scale the backdrop so it carries the layer's luminance. Where the backdrop
+    # has no luminance to scale there is no hue or saturation to preserve either,
+    # so the result is a neutral grey at the layer's luminance - which is also the
+    # analytic limit of i * lum(l)/lum(i) as a grey backdrop approaches black.
+    # Guarding the numerator here instead (returning black) makes a luminosity
+    # layer disappear over dark backdrops; see tests-unit/comfy_extras_test/
+    # compositor_blend_golden.json.
+    lum_i = luminance(i)
+    lum_l = luminance(l)
+    degenerate = lum_i <= EPSILON
+    ratio = np.where(degenerate, 0.0, lum_l / np.where(degenerate, 1.0, lum_i))
+    return np.where(
+        degenerate[..., None],
+        np.broadcast_to(lum_l[..., None], i.shape),
+        i * ratio[..., None],
+    )
+
+
+HSL_BLEND = {
+    "hue": _blend_hue,
+    "saturation": _blend_saturation,
+    "color": _blend_color,
+    "luminosity": _blend_luminosity,
+}
+
+
+def blend_pixel(blend: str, in_rgb: np.ndarray, layer_rgb: np.ndarray) -> np.ndarray:
+    in_rgb = np.asarray(in_rgb, dtype=np.float32)
+    layer_rgb = np.asarray(layer_rgb, dtype=np.float32)
+    hsl = HSL_BLEND.get(blend)
+    if hsl is not None:
+        return np.asarray(hsl(in_rgb, layer_rgb), dtype=np.float32)
+    fn = CHANNEL_BLEND.get(blend, CHANNEL_BLEND["normal"])
+    return np.asarray(fn(in_rgb, layer_rgb), dtype=np.float32)
+
+
+def _composite_union(in_c, layer, comp, cov):
+    in_a = in_c[..., 3]
+    layer_a = layer[..., 3] * cov
+    new_a = layer_a + (1 - layer_a) * in_a
+    ratio = np.where(new_a != 0, layer_a / np.where(new_a != 0, new_a, 1), 0)
+    blended = (
+        ratio[..., None]
+        * (in_a[..., None] * (comp - layer[..., :3]) + layer[..., :3] - in_c[..., :3])
+        + in_c[..., :3]
+    )
+    keep = (layer_a == 0) | (new_a == 0)
+    rgb = np.where(
+        keep[..., None],
+        in_c[..., :3],
+        np.where((in_a == 0)[..., None], layer[..., :3], blended),
+    )
+    return np.concatenate([rgb, new_a[..., None]], axis=-1)
+
+
+def _composite_clip_to_backdrop(in_c, layer, comp, cov):
+    in_a = in_c[..., 3]
+    layer_a = layer[..., 3] * cov
+    mixed = comp * layer_a[..., None] + in_c[..., :3] * (1 - layer_a[..., None])
+    keep = (in_a == 0) | (layer_a == 0)
+    rgb = np.where(keep[..., None], in_c[..., :3], mixed)
+    return np.concatenate([rgb, in_a[..., None]], axis=-1)
+
+
+def _composite_clip_to_layer(in_c, layer, comp, cov):
+    in_a = in_c[..., 3]
+    layer_a = layer[..., 3] * cov
+    mixed = comp * in_a[..., None] + layer[..., :3] * (1 - in_a[..., None])
+    rgb = np.where(
+        (layer_a == 0)[..., None],
+        in_c[..., :3],
+        np.where((in_a == 0)[..., None], layer[..., :3], mixed),
+    )
+    return np.concatenate([rgb, layer_a[..., None]], axis=-1)
+
+
+def _composite_intersection(in_c, layer, comp, cov):
+    new_a = in_c[..., 3] * layer[..., 3] * cov
+    rgb = np.where((new_a == 0)[..., None], in_c[..., :3], comp)
+    return np.concatenate([rgb, new_a[..., None]], axis=-1)
+
+
+_COMPOSITE = {
+    "union": _composite_union,
+    "clip-to-backdrop": _composite_clip_to_backdrop,
+    "clip-to-layer": _composite_clip_to_layer,
+    "intersection": _composite_intersection,
+}
+
+
+def run_composite(mode: str, in_c, layer, comp, cov) -> np.ndarray:
+    fn = _COMPOSITE.get(mode, _composite_union)
+    return fn(in_c, layer, comp, cov)
+
+
+def _to_space(rgb: np.ndarray, space: str) -> np.ndarray:
+    return rgb if space == "linear" else linear_to_srgb(rgb)
+
+
+def _from_space(rgb: np.ndarray, space: str) -> np.ndarray:
+    return rgb if space == "linear" else srgb_to_linear(rgb)
+
+
+class EffectiveMode(NamedTuple):
+    blend: str
+    blend_space: str
+    composite: str
+
+
+_LAYER_MODES = {
+    "normal": ("linear", "union"),
+    "multiply": ("linear", "clip-to-backdrop"),
+    "screen": ("perceptual", "clip-to-backdrop"),
+    "overlay": ("perceptual", "clip-to-backdrop"),
+    "darken": ("linear", "clip-to-backdrop"),
+    "lighten": ("linear", "clip-to-backdrop"),
+    "color-dodge": ("perceptual", "clip-to-backdrop"),
+    "color-burn": ("perceptual", "clip-to-backdrop"),
+    "hard-light": ("perceptual", "clip-to-backdrop"),
+    "soft-light": ("perceptual", "clip-to-backdrop"),
+    "difference": ("perceptual", "clip-to-backdrop"),
+    "exclusion": ("perceptual", "clip-to-backdrop"),
+    "linear-dodge": ("linear", "clip-to-backdrop"),
+    "linear-burn": ("perceptual", "clip-to-backdrop"),
+    "vivid-light": ("perceptual", "clip-to-backdrop"),
+    "pin-light": ("perceptual", "clip-to-backdrop"),
+    "linear-light": ("perceptual", "clip-to-backdrop"),
+    "hard-mix": ("perceptual", "clip-to-backdrop"),
+    "subtract": ("linear", "clip-to-backdrop"),
+    "divide": ("linear", "clip-to-backdrop"),
+    "grain-extract": ("perceptual", "clip-to-backdrop"),
+    "grain-merge": ("perceptual", "clip-to-backdrop"),
+    "hue": ("perceptual", "clip-to-backdrop"),
+    "saturation": ("perceptual", "clip-to-backdrop"),
+    "color": ("perceptual", "clip-to-backdrop"),
+    "luminosity": ("linear", "clip-to-backdrop"),
+}
+
+
+def resolve_mode(blend: str = "normal") -> EffectiveMode:
+    blend_space, composite = _LAYER_MODES.get(blend, _LAYER_MODES["normal"])
+    return EffectiveMode(
+        blend=blend,
+        blend_space=blend_space,
+        composite=composite,
+    )
+
+
+def blend_composite(
+    mode: EffectiveMode,
+    backdrop: np.ndarray,
+    layer: np.ndarray,
+    opacity: float,
+    mask: Optional[ArrayLike] = None,
+) -> np.ndarray:
+    backdrop = np.asarray(backdrop, dtype=np.float32)
+    layer = np.asarray(layer, dtype=np.float32)
+    cov = opacity * (1.0 if mask is None else mask)
+
+    in_b = _to_space(backdrop[..., :3], mode.blend_space)
+    layer_b = _to_space(layer[..., :3], mode.blend_space)
+    comp = _from_space(blend_pixel(mode.blend, in_b, layer_b), mode.blend_space)
+
+    return run_composite(mode.composite, backdrop, layer, comp, cov)
+
+
+def placed_bounds(
+    x: float, y: float, w: float, h: float, rotation: float
+) -> tuple[int, int, int, int]:
+    cx = x + w / 2
+    cy = y + h / 2
+    cos = math.cos(rotation)
+    sin = math.sin(rotation)
+    hw = w / 2
+    hh = h / 2
+    corners = ((-hw, -hh), (hw, -hh), (hw, hh), (-hw, hh))
+    xs = [cx + dx * cos - dy * sin for dx, dy in corners]
+    ys = [cy + dx * sin + dy * cos for dx, dy in corners]
+    bx = math.floor(min(xs))
+    by = math.floor(min(ys))
+    bw = max(1, math.ceil(max(xs)) - bx)
+    bh = max(1, math.ceil(max(ys)) - by)
+    return bx, by, bw, bh

comfy_extras/nodes_compositor.py

@@ -0,0 +1,855 @@
+import hashlib
+import json
+import math
+
+import numpy as np
+import torch
+from PIL import Image
+
+from comfy_api.latest import ComfyExtension, io, UI
+from comfy_extras.compositor_blend import (
+    _LAYER_MODES,
+    blend_composite,
+    linear_to_srgb,
+    placed_bounds,
+    resolve_mode,
+    srgb_to_linear,
+)
+from comfy_extras.color_util import hex_to_rgb
+from comfy_extras.nodes_bounding_boxes import boxes_from_input
+from nodes import MAX_RESOLUTION
+from typing_extensions import override
+
+
+MAX_LAYERS = 50
+
+
+def document_items(doc) -> list[dict]:
+    if not isinstance(doc, dict):
+        return []
+    version = doc.get("version")
+    if version is not None and version != 1:
+        raise ValueError(f"LAYERS document version {version!r} is not supported")
+    items = []
+    for item in doc.get("layers") or []:
+        if not isinstance(item, dict):
+            continue
+        item_type = item.get("type", "raster")
+        if item_type != "raster":
+            raise ValueError(f"LAYERS item type {item_type!r} is not supported yet")
+        if not isinstance(item.get("image"), torch.Tensor):
+            continue
+        blend = item.get("blend_mode")
+        if blend is not None and blend not in _LAYER_MODES:
+            raise ValueError(f"LAYERS item blend_mode {blend!r} is not a known blend mode")
+        items.append(item)
+    return sorted(items, key=lambda item: _int(item.get("z_index"), 0))
+
+
+def document_canvas(doc) -> tuple[int, int] | None:
+    if not isinstance(doc, dict):
+        return None
+    canvas = doc.get("canvas")
+    if not isinstance(canvas, (tuple, list)) or len(canvas) != 2:
+        return None
+    w, h = _int(canvas[0], 0), _int(canvas[1], 0)
+    return (w, h) if w > 0 and h > 0 else None
+
+
+def _int(value, default: int) -> int:
+    return int(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else default
+
+
+def _bbox_list(bboxes, canvas_width: int, canvas_height: int) -> list[dict]:
+    if bboxes is None:
+        return []
+    if isinstance(bboxes, str):
+        text = bboxes.strip()
+        if not text:
+            return []
+        try:
+            bboxes = json.loads(text)
+        except (json.JSONDecodeError, ValueError) as exc:
+            raise ValueError(f"bboxes string input is not valid JSON: {exc}") from exc
+    probe = bboxes if isinstance(bboxes, list) else [bboxes]
+    if probe and isinstance(probe[0], list):
+        probe = probe[0]
+    has_elements = any(
+        isinstance(box, dict) and isinstance(box.get("bbox"), (list, tuple))
+        for box in probe
+    )
+    if has_elements and (canvas_width <= 0 or canvas_height <= 0):
+        raise ValueError(
+            "normalized element boxes need canvas_width and canvas_height to resolve to pixels"
+        )
+    return boxes_from_input(bboxes, canvas_width, canvas_height)
+
+
+def _item_mask_frame(mask, index: int) -> torch.Tensor | None:
+    if not isinstance(mask, torch.Tensor):
+        return None
+    if mask.shape[0] == 1:
+        return mask[:1]
+    if index < mask.shape[0]:
+        return mask[index : index + 1]
+    return None
+
+
+def expand_item_frames(items: list[dict]) -> list[dict]:
+    frames = []
+    for item in items:
+        image = item["image"]
+        for index in range(image.shape[0]):
+            width = _int(item.get("w"), 0)
+            height = _int(item.get("h"), 0)
+            rotation = item.get("rotation")
+            frames.append({
+                "tensor": image[index : index + 1],
+                "mask": _item_mask_frame(item.get("mask"), index),
+                "name": item.get("name") if isinstance(item.get("name"), str) else None,
+                "x": _int(item.get("x"), 0),
+                "y": _int(item.get("y"), 0),
+                "w": width if width > 0 else int(image.shape[2]),
+                "h": height if height > 0 else int(image.shape[1]),
+                "rotation": float(rotation)
+                if isinstance(rotation, (int, float)) and not isinstance(rotation, bool)
+                else 0.0,
+                "opacity": item.get("opacity", 1.0),
+                "blend": item.get("blend_mode", "normal"),
+                "visible": item.get("visible", True),
+                "flip_h": bool(item.get("flip_h", False)),
+                "flip_v": bool(item.get("flip_v", False)),
+            })
+    if len(frames) > MAX_LAYERS:
+        raise ValueError(
+            f"Compositor supports at most {MAX_LAYERS} layers, got {len(frames)}"
+        )
+    return frames
+
+
+def frame_alpha(
+    tensor: torch.Tensor, mask: torch.Tensor | None
+) -> torch.Tensor | None:
+    alpha = tensor[:1, :, :, 3] if tensor.shape[-1] == 4 else None
+    if mask is None:
+        return alpha
+    h, w = tensor.shape[1], tensor.shape[2]
+    m = mask[:1].to(device=tensor.device, dtype=torch.float32)
+    if m.shape[1] != h or m.shape[2] != w:
+        m = torch.nn.functional.interpolate(
+            m.unsqueeze(1), size=(h, w), mode="bilinear"
+        ).squeeze(1)
+    inv = torch.clamp(1.0 - m, 0.0, 1.0)
+    return inv if alpha is None else alpha * inv
+
+
+def layer_preview_tensor(
+    tensor: torch.Tensor, alpha: torch.Tensor | None
+) -> torch.Tensor:
+    rgb = tensor[:1, :, :, :3]
+    if alpha is None:
+        return rgb
+    return torch.cat([rgb, alpha.unsqueeze(-1)], dim=-1)
+
+
+def canvas_extent(frames: list[dict]) -> tuple[int, int]:
+    right = 1
+    bottom = 1
+    for frame in frames:
+        bx, by, bw, bh = placed_bounds(
+            frame["x"], frame["y"], frame["w"], frame["h"], frame["rotation"]
+        )
+        right = max(right, bx + bw)
+        bottom = max(bottom, by + bh)
+    return (right, bottom)
+
+
+def input_fingerprints(
+    frames: list[dict], alphas: list[torch.Tensor | None]
+) -> list[str]:
+    fingerprints = []
+    for frame, alpha in zip(frames, alphas):
+        tensor = frame["tensor"]
+        rgb = tensor[0, :, :, :3].detach().cpu().numpy()
+        rgb8 = np.clip(np.rint(rgb * 255.0), 0, 255).astype(np.uint8)
+        digest = hashlib.sha256()
+        digest.update(repr(tuple(tensor.shape)).encode())
+        digest.update(rgb8.tobytes())
+        if alpha is not None:
+            alpha8 = np.clip(
+                np.rint(alpha[0].detach().cpu().numpy() * 255.0), 0, 255
+            ).astype(np.uint8)
+            digest.update(alpha8.tobytes())
+        digest.update(
+            repr((
+                frame["x"],
+                frame["y"],
+                frame["w"],
+                frame["h"],
+                frame["rotation"],
+                frame["opacity"],
+                frame["blend"],
+                bool(frame["visible"]),
+                frame["flip_h"],
+                frame["flip_v"],
+            )).encode()
+        )
+        fingerprints.append(digest.hexdigest()[:16])
+    return fingerprints
+
+
+def state_from_items(frames: list[dict], canvas: tuple[int, int]) -> dict:
+    layers = []
+    for frame in frames:
+        layers.append({
+            "name": frame["name"],
+            "visible": bool(frame["visible"]),
+            "opacity": frame["opacity"],
+            "blend": frame["blend"],
+            "flipH": frame["flip_h"],
+            "flipV": frame["flip_v"],
+            "transform": {
+                "x": frame["x"],
+                "y": frame["y"],
+                "w": frame["w"],
+                "h": frame["h"],
+                "rotation": frame["rotation"],
+            },
+        })
+    return {
+        "canvas": canvas,
+        "layers": layers,
+        "inputs": None,
+        "background": {"color": "#ffffff", "opacity": 1.0, "visible": False},
+    }
+
+
+def layer_ui_entries(frames: list[dict]) -> list:
+    entries = []
+    for frame in frames:
+        entries.append({
+            "x": frame["x"],
+            "y": frame["y"],
+            "width": int(frame["w"]),
+            "height": int(frame["h"]),
+            "rotation": frame["rotation"],
+            "name": frame["name"],
+            "visible": bool(frame["visible"]),
+            "opacity": frame["opacity"] if isinstance(frame["opacity"], (int, float)) else 1.0,
+            "blend": frame["blend"] if isinstance(frame["blend"], str) else "normal",
+            "flipH": frame["flip_h"],
+            "flipV": frame["flip_v"],
+        })
+    return entries
+
+
+_HEX_DIGITS = set("0123456789abcdef")
+
+
+def _normalize_hex_color(value) -> str:
+    if isinstance(value, str):
+        text = value.strip().lower()
+        if text.startswith("#"):
+            digits = text[1:]
+            if len(digits) == 3 and set(digits) <= _HEX_DIGITS:
+                digits = "".join(ch * 2 for ch in digits)
+            if len(digits) == 6 and set(digits) <= _HEX_DIGITS:
+                return "#" + digits
+    return "#ffffff"
+
+
+def _parse_background(entry) -> dict | None:
+    if not isinstance(entry, dict):
+        return None
+    return {
+        "color": _normalize_hex_color(entry.get("color")),
+        "opacity": min(max(_number(entry, "opacity", 1.0), 0.0), 1.0),
+        "visible": bool(entry.get("visible", True)),
+    }
+
+
+def _parse_order(value, layer_count: int) -> list[int] | None:
+    if not isinstance(value, list) or not value:
+        return None
+    if not all(
+        isinstance(item, int) and not isinstance(item, bool) for item in value
+    ):
+        return None
+    if sorted(value) != list(range(layer_count)):
+        return None
+    return value
+
+
+def layer_state_provided(raw) -> bool:
+    if isinstance(raw, dict):
+        return bool(raw)
+    if isinstance(raw, str):
+        return raw not in ("", "{}")
+    return False
+
+
+def parse_layer_state(raw) -> dict | None:
+    if isinstance(raw, str):
+        if not raw.strip():
+            return None
+        try:
+            raw = json.loads(raw)
+        except (json.JSONDecodeError, ValueError):
+            return None
+    if not isinstance(raw, dict):
+        return None
+    state = raw
+    version = state.get("version")
+    if version is not None and version != 1:
+        return None
+    canvas = state.get("canvas")
+    layers = state.get("layers")
+    if not isinstance(canvas, dict) or not isinstance(layers, list) or not layers:
+        return None
+    try:
+        w = int(round(float(canvas.get("w"))))
+        h = int(round(float(canvas.get("h"))))
+    except (TypeError, ValueError, OverflowError):
+        return None
+    if w <= 0 or h <= 0:
+        return None
+    inputs = state.get("inputs")
+    if (
+        not isinstance(inputs, list)
+        or len(inputs) != len(layers)
+        or not all(isinstance(entry, str) for entry in inputs)
+    ):
+        inputs = None
+    return {
+        "canvas": (w, h),
+        "layers": layers,
+        "inputs": inputs,
+        "background": _parse_background(state.get("background")),
+        "order": _parse_order(state.get("order"), len(layers)),
+    }
+
+
+def _number(source: dict, key: str, default: float) -> float:
+    value = source.get(key, default)
+    if not isinstance(value, (int, float)) or not math.isfinite(value):
+        return float(default)
+    return float(value)
+
+
+def _clamped_size(value: float, natural: int) -> float:
+    return float(natural) if value <= 0 else min(value, float(MAX_RESOLUTION))
+
+
+def _layer_params(entry, natural_w: int, natural_h: int) -> dict:
+    if not isinstance(entry, dict):
+        entry = {}
+    transform = entry.get("transform")
+    if not isinstance(transform, dict):
+        transform = {}
+    blend = entry.get("blend")
+    return {
+        "visible": bool(entry.get("visible", True)),
+        # The layer state is untrusted input: it round-trips through the saved
+        # workflow and can be posted directly to /prompt. An out-of-range opacity
+        # would otherwise reach blend_composite as a raw coverage multiplier and
+        # produce negative or greater-than-white RGB. _parse_background already
+        # clamps the same field.
+        "opacity": min(max(_number(entry, "opacity", 1.0), 0.0), 1.0),
+        "blend": blend if isinstance(blend, str) else "normal",
+        "x": min(max(_number(transform, "x", 0.0), -MAX_RESOLUTION), MAX_RESOLUTION),
+        "y": min(max(_number(transform, "y", 0.0), -MAX_RESOLUTION), MAX_RESOLUTION),
+        "w": _clamped_size(_number(transform, "w", natural_w), natural_w),
+        "h": _clamped_size(_number(transform, "h", natural_h), natural_h),
+        "rotation": _number(transform, "rotation", 0.0),
+        "flip_h": bool(entry.get("flipH", False)),
+        "flip_v": bool(entry.get("flipV", False)),
+    }
+
+
+def _prepare_layer_bitmap(
+    tensor: torch.Tensor, params: dict, alpha: torch.Tensor | None
+) -> Image.Image:
+    frame = tensor[0, :, :, :3].detach().cpu().numpy()
+    rgb8 = np.clip(np.rint(frame * 255.0), 0, 255).astype(np.uint8)
+    if alpha is None:
+        img = Image.fromarray(rgb8, "RGB").convert("RGBA")
+    else:
+        alpha8 = np.clip(
+            np.rint(alpha[0].detach().cpu().numpy() * 255.0), 0, 255
+        ).astype(np.uint8)
+        img = Image.fromarray(np.dstack([rgb8, alpha8]), "RGBA")
+    if params["flip_h"]:
+        img = img.transpose(Image.Transpose.FLIP_LEFT_RIGHT)
+    if params["flip_v"]:
+        img = img.transpose(Image.Transpose.FLIP_TOP_BOTTOM)
+    target = (max(1, round(params["w"])), max(1, round(params["h"])))
+    if img.size != target:
+        img = img.resize(target, Image.Resampling.LANCZOS)
+    if params["rotation"] != 0:
+        img = img.rotate(
+            -math.degrees(params["rotation"]),
+            expand=True,
+            resample=Image.Resampling.BICUBIC,
+            fillcolor=(0, 0, 0, 0),
+        )
+    return img
+
+
+def _place_in_bounds(img: Image.Image, bw: int, bh: int) -> np.ndarray:
+    arr = np.asarray(img, dtype=np.float32) / 255.0
+    rgba = np.concatenate([srgb_to_linear(arr[..., :3]), arr[..., 3:4]], axis=-1)
+    aw, ah = img.size
+    buf = np.zeros((bh, bw, 4), dtype=np.float32)
+    ox = (bw - aw) // 2
+    oy = (bh - ah) // 2
+    dx0, dy0 = max(ox, 0), max(oy, 0)
+    dx1, dy1 = min(ox + aw, bw), min(oy + ah, bh)
+    if dx0 < dx1 and dy0 < dy1:
+        buf[dy0:dy1, dx0:dx1] = rgba[dy0 - oy : dy1 - oy, dx0 - ox : dx1 - ox]
+    return buf
+
+
+def _fill_background(canvas: np.ndarray, background: dict) -> np.ndarray:
+    layer = np.empty(canvas.shape, dtype=np.float32)
+    layer[..., :3] = srgb_to_linear(
+        np.array(hex_to_rgb(background["color"]), dtype=np.float32) / 255.0
+    )
+    layer[..., 3] = 1.0
+    return blend_composite(
+        resolve_mode("normal"), canvas, layer, background["opacity"]
+    )
+
+
+def composite_from_state(
+    tensors: list[torch.Tensor],
+    state: dict,
+    alphas: list[torch.Tensor | None],
+) -> torch.Tensor:
+    cw, ch = state["canvas"]
+    if cw > MAX_RESOLUTION or ch > MAX_RESOLUTION:
+        raise ValueError(
+            f"Compositor canvas {cw}x{ch} exceeds the maximum supported size of "
+            f"{MAX_RESOLUTION}x{MAX_RESOLUTION}"
+        )
+    canvas = np.zeros((ch, cw, 4), dtype=np.float32)
+    background = state.get("background")
+    if background is not None and background["visible"] and background["opacity"] > 0:
+        canvas = _fill_background(canvas, background)
+    layers = state["layers"]
+    order = state.get("order") or range(len(tensors))
+    for index in order:
+        if index < 0 or index >= len(tensors):
+            continue
+        tensor = tensors[index]
+        entry = layers[index] if index < len(layers) else None
+        params = _layer_params(entry, tensor.shape[2], tensor.shape[1])
+        if not params["visible"]:
+            continue
+        img = _prepare_layer_bitmap(
+            tensor, params, alphas[index] if index < len(alphas) else None
+        )
+        bx, by, bw, bh = placed_bounds(
+            params["x"], params["y"], params["w"], params["h"], params["rotation"]
+        )
+        buf = _place_in_bounds(img, bw, bh)
+        x0, y0 = max(bx, 0), max(by, 0)
+        x1, y1 = min(bx + bw, cw), min(by + bh, ch)
+        if x0 >= x1 or y0 >= y1:
+            continue
+        region = buf[y0 - by : y1 - by, x0 - bx : x1 - bx]
+        mode = resolve_mode(params["blend"])
+        canvas[y0:y1, x0:x1] = blend_composite(
+            mode, canvas[y0:y1, x0:x1], region, params["opacity"]
+        )
+    rgb = linear_to_srgb(np.clip(canvas[..., :3], 0.0, 1.0))
+    alpha = np.clip(canvas[..., 3:4], 0.0, 1.0)
+    rgba = np.concatenate([rgb, alpha], axis=-1)
+    return torch.from_numpy(rgba.astype(np.float32)).unsqueeze(0)
+
+
+OPAQUE_EPSILON = 1e-3
+
+
+def composite_outputs(out: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    if out.shape[-1] != 4:
+        return out, torch.zeros(out.shape[:3], dtype=torch.float32)
+    alpha = out[..., 3]
+    if bool((alpha >= 1.0 - OPAQUE_EPSILON).all()):
+        return out[..., :3], torch.zeros_like(alpha)
+    return out, torch.clamp(1.0 - alpha, 0.0, 1.0)
+
+
+class ImageCompositor(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="ImageCompositor",
+            display_name="Create Layered Image",
+            category="image",
+            search_aliases=["compositor", "composite", "layer", "layers", "layer editor", "psd"],
+            is_experimental=True,
+            # both flags on purpose: terminal compositor graphs must execute (the
+            # editor needs a run to open), and cache hits must replay the layer UI
+            is_output_node=True,
+            has_intermediate_output=True,
+            inputs=[
+                io.Layers.Input(
+                    "layers",
+                    tooltip="Layer stack to composite; build it with Add Layer. Items are stacked by z_index, batch frames inside an item expand to consecutive layers, and item placement, opacity, and blend mode define the initial composition. Without an explicit document canvas the size is a best-effort maximum extent of the placed layers. A saved composition that matches the current inputs takes priority.",
+                ),
+                io.Compositor.Input(
+                    "compositor",
+                    tooltip="Layered composition saved by the compositor editor.",
+                ),
+            ],
+            outputs=[
+                io.Image.Output(
+                    tooltip="Composited image. Carries an alpha channel when the composite has transparent areas (e.g. hidden background), otherwise plain RGB."
+                ),
+                io.Mask.Output(
+                    tooltip="Transparency of the composite (1 = fully transparent). All zeros when the composite is opaque."
+                ),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, layers: io.Layers.Type, compositor: io.Compositor.Type = None) -> io.NodeOutput:
+        frames = expand_item_frames(document_items(layers))
+        tensors = [frame["tensor"] for frame in frames]
+        alphas = [frame_alpha(frame["tensor"], frame["mask"]) for frame in frames]
+
+        layer_refs = []
+        for tensor, alpha in zip(tensors, alphas):
+            layer_refs.extend(
+                UI.PreviewImage(layer_preview_tensor(tensor, alpha), cls=cls).values
+            )
+
+        fp = input_fingerprints(frames, alphas)
+        raw_state = compositor
+        state = parse_layer_state(raw_state)
+        replay = bool(state is not None and tensors and state["inputs"] == fp)
+        if replay:
+            out = composite_from_state(tensors, state, alphas)
+        elif tensors:
+            canvas = document_canvas(layers) or canvas_extent(frames)
+            out = composite_from_state(
+                tensors, state_from_items(frames, canvas), alphas
+            )
+        else:
+            out = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
+        state_stale = layer_state_provided(raw_state) and not replay
+        out, mask = composite_outputs(out)
+
+        ui_dict = UI.PreviewImage(out, cls=cls).as_dict()
+        ui_dict["compositor_layers"] = layer_refs
+        ui_dict["compositor_inputs"] = fp
+        ui_dict["compositor_bboxes"] = layer_ui_entries(frames)
+        if state_stale:
+            ui_dict["compositor_state_stale"] = [True]
+        return io.NodeOutput(out, mask, ui=ui_dict)
+
+
+class AddLayer(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="AddLayer",
+            display_name="Add Layer",
+            category="image",
+            is_experimental=True,
+            inputs=[
+                io.Layers.Input(
+                    "layers",
+                    optional=True,
+                    tooltip="Layer stack to append to. Leave unconnected to start a new stack.",
+                ),
+                io.Image.Input(
+                    "image",
+                    tooltip="Layer content at its native size. A batch expands to consecutive layers.",
+                ),
+                io.Mask.Input(
+                    "mask",
+                    optional=True,
+                    tooltip="Transparency mask for this layer. Masked areas (value 1) become transparent, multiplying with any alpha channel the image already carries.",
+                ),
+                io.String.Input(
+                    "name",
+                    optional=True,
+                    default="",
+                    tooltip="Layer name shown in the compositor editor.",
+                ),
+                io.Int.Input(
+                    "x",
+                    optional=True,
+                    default=0,
+                    min=-MAX_RESOLUTION,
+                    max=MAX_RESOLUTION,
+                    tooltip="Initial horizontal placement on the canvas.",
+                ),
+                io.Int.Input(
+                    "y",
+                    optional=True,
+                    default=0,
+                    min=-MAX_RESOLUTION,
+                    max=MAX_RESOLUTION,
+                    tooltip="Initial vertical placement on the canvas.",
+                ),
+                io.Float.Input(
+                    "opacity",
+                    optional=True,
+                    default=1.0,
+                    min=0.0,
+                    max=1.0,
+                    step=0.01,
+                    tooltip="Initial layer opacity.",
+                ),
+                io.Combo.Input(
+                    "blend_mode",
+                    options=list(_LAYER_MODES),
+                    default="normal",
+                    optional=True,
+                    tooltip="Initial blend mode, applied against the layers below. On the bottom layer over the default transparent background, non-normal modes produce transparency.",
+                ),
+                io.Float.Input(
+                    "rotation",
+                    optional=True,
+                    default=0.0,
+                    min=-360.0,
+                    max=360.0,
+                    step=1.0,
+                    tooltip="Initial rotation in degrees, clockwise.",
+                ),
+                io.Int.Input(
+                    "width",
+                    optional=True,
+                    default=0,
+                    min=0,
+                    max=MAX_RESOLUTION,
+                    tooltip="Initial display width. 0 keeps the image's native width.",
+                ),
+                io.Int.Input(
+                    "height",
+                    optional=True,
+                    default=0,
+                    min=0,
+                    max=MAX_RESOLUTION,
+                    tooltip="Initial display height. 0 keeps the image's native height.",
+                ),
+                io.Int.Input(
+                    "z_index",
+                    optional=True,
+                    default=0,
+                    min=-1000,
+                    max=1000,
+                    tooltip="Stacking override. Layers are stable-sorted by z_index; equal values keep their list order.",
+                ),
+                io.Boolean.Input(
+                    "flip_h",
+                    optional=True,
+                    default=False,
+                    tooltip="Flip the layer horizontally.",
+                ),
+                io.Boolean.Input(
+                    "flip_v",
+                    optional=True,
+                    default=False,
+                    tooltip="Flip the layer vertically.",
+                ),
+            ],
+            outputs=[
+                io.Layers.Output(tooltip="The layer stack with this layer appended."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, image: io.Image.Type, layers: io.Layers.Type = None, mask: io.Mask.Type = None, name: str = "", x: int = 0, y: int = 0, opacity: float = 1.0, blend_mode: str = "normal", rotation: float = 0.0, width: int = 0, height: int = 0, z_index: int = 0, flip_h: bool = False, flip_v: bool = False) -> io.NodeOutput:
+        item: dict = {
+            "image": image,
+            "type": "raster",
+            "x": int(x),
+            "y": int(y),
+            "z_index": int(z_index),
+        }
+        if mask is not None:
+            item["mask"] = mask
+        if name:
+            item["name"] = name
+        if opacity != 1.0:
+            item["opacity"] = float(opacity)
+        if blend_mode != "normal":
+            item["blend_mode"] = blend_mode
+        if rotation != 0.0:
+            item["rotation"] = math.radians(rotation)
+        if width > 0:
+            item["w"] = int(width)
+        if height > 0:
+            item["h"] = int(height)
+        if flip_h:
+            item["flip_h"] = True
+        if flip_v:
+            item["flip_v"] = True
+        previous = layers if isinstance(layers, dict) else None
+        document: dict = {
+            "version": 1,
+            "layers": [*(previous.get("layers") or []), item] if previous else [item],
+        }
+        previous_canvas = document_canvas(previous)
+        if previous_canvas:
+            document["canvas"] = previous_canvas
+        return io.NodeOutput(document)
+
+
+class LayersFromBoundingBoxes(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LayersFromBoundingBoxes",
+            display_name="Layers From Bounding Boxes",
+            category="image",
+            is_experimental=True,
+            description=(
+                "Turn an image batch plus its bounding boxes into a layer stack, one layer per frame, "
+                "each placed by its own box. Use this when a node emits layers as a batch - a batch "
+                "carries a single placement for every frame, so the individual positions are otherwise lost."
+            ),
+            inputs=[
+                io.Image.Input(
+                    "image",
+                    tooltip="Image batch; each frame becomes one layer.",
+                ),
+                io.MultiType.Input(
+                    "bboxes",
+                    [io.BoundingBox, io.Array, io.String],
+                    tooltip=(
+                        "Placement boxes, index-aligned with the image batch. Accepts bounding boxes "
+                        "(x, y, width, height), normalized elements (with a 'bbox' - these need "
+                        "canvas_width/canvas_height to resolve to pixels), or a JSON string of either. "
+                        "Frames without a matching box are placed at the origin. A box's width/height "
+                        "scales the layer to fit it. metadata.name (or desc) and metadata.z_index are "
+                        "used when present, and metadata.content_rect (frame-relative) crops the frame "
+                        "to its real content."
+                    ),
+                ),
+                io.Mask.Input(
+                    "mask",
+                    optional=True,
+                    tooltip=(
+                        "Per-frame transparency, index-aligned with the image batch "
+                        "(1 = transparent, LoadImage convention)."
+                    ),
+                ),
+                io.Layers.Input(
+                    "layers",
+                    optional=True,
+                    tooltip="Layer stack to append to. Leave unconnected to start a new stack.",
+                ),
+                io.Boolean.Input(
+                    "crop_to_content",
+                    default=True,
+                    optional=True,
+                    tooltip=(
+                        "Crop each frame to metadata.content_rect where present and place the content "
+                        "at the box position plus the rect offset. Leave on for batches whose frames "
+                        "are padded - it keeps only the real content at its true spot."
+                    ),
+                ),
+                io.Int.Input(
+                    "canvas_width",
+                    default=0,
+                    min=0,
+                    max=MAX_RESOLUTION,
+                    optional=True,
+                    tooltip="Document canvas width. 0 derives it from the placed layers.",
+                ),
+                io.Int.Input(
+                    "canvas_height",
+                    default=0,
+                    min=0,
+                    max=MAX_RESOLUTION,
+                    optional=True,
+                    tooltip="Document canvas height. 0 derives it from the placed layers.",
+                ),
+            ],
+            outputs=[
+                io.Layers.Output(tooltip="The layer stack, ready for Create Layered Image."),
+            ],
+        )
+
+    @classmethod
+    def execute(
+        cls,
+        image: io.Image.Type,
+        bboxes: io.MultiType.Type,
+        mask: io.Mask.Type = None,
+        layers: io.Layers.Type = None,
+        crop_to_content: bool = True,
+        canvas_width: int = 0,
+        canvas_height: int = 0,
+    ) -> io.NodeOutput:
+        boxes = _bbox_list(bboxes, canvas_width, canvas_height)
+        previous = layers if isinstance(layers, dict) else None
+        items: list[dict] = list((previous.get("layers") or []) if previous else [])
+        base_z = max((_int(i.get("z_index"), 0) for i in items), default=-1) + 1
+
+        for index in range(image.shape[0]):
+            box = boxes[index] if index < len(boxes) else {}
+            meta = box.get("metadata") if isinstance(box.get("metadata"), dict) else {}
+            frame = image[index : index + 1]
+            frame_mask = _item_mask_frame(mask, index)
+
+            x, y = _int(box.get("x"), 0), _int(box.get("y"), 0)
+            box_w, box_h = _int(box.get("width"), 0), _int(box.get("height"), 0)
+            cropped = False
+            rect = meta.get("content_rect")
+            if crop_to_content and isinstance(rect, (list, tuple)) and len(rect) == 4:
+                left, top, cw, ch = (_int(v, 0) for v in rect)
+                left = min(max(left, 0), int(frame.shape[2]))
+                top = min(max(top, 0), int(frame.shape[1]))
+                cw = min(max(cw, 0), int(frame.shape[2]) - left)
+                ch = min(max(ch, 0), int(frame.shape[1]) - top)
+                if cw > 0 and ch > 0:
+                    frame = frame[:, top : top + ch, left : left + cw]
+                    if frame_mask is not None:
+                        frame_mask = frame_mask[:, top : top + ch, left : left + cw]
+                    x, y = x + left, y + top
+                    cropped = True
+
+            item: dict = {
+                "image": frame,
+                "type": "raster",
+                "x": x,
+                "y": y,
+                "z_index": _int(meta.get("z_index"), base_z + index),
+            }
+            if not cropped:
+                if box_w > 0:
+                    item["w"] = box_w
+                if box_h > 0:
+                    item["h"] = box_h
+            if frame_mask is not None:
+                item["mask"] = frame_mask
+            name = meta.get("name")
+            if not (isinstance(name, str) and name):
+                name = meta.get("desc")
+            if isinstance(name, str) and name:
+                item["name"] = name
+            items.append(item)
+
+        document: dict = {"version": 1, "layers": items}
+        if canvas_width > 0 and canvas_height > 0:
+            document["canvas"] = (canvas_width, canvas_height)
+        else:
+            inherited = document_canvas(previous)
+            if inherited:
+                document["canvas"] = inherited
+        return io.NodeOutput(document)
+
+
+class CompositorExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [ImageCompositor, AddLayer, LayersFromBoundingBoxes]
+
+
+async def comfy_entrypoint() -> CompositorExtension:
+    return CompositorExtension()

comfy_extras/nodes_custom_sampler.py

@@ -591,7 +591,7 @@ class SamplerER_SDE(io.ComfyNode):
             inputs=[
                 io.Combo.Input("solver_type", options=["ER-SDE", "Reverse-time SDE", "ODE"]),
                 io.Int.Input("max_stage", default=3, min=1, max=3, advanced=True),
-                io.Float.Input("eta", default=1.0, min=0.0, max=100.0, step=0.01, round=False, tooltip="Stochastic strength of reverse-time SDE.\nWhen eta=0, it reduces to deterministic ODE. This setting doesn't apply to ER-SDE solver type.", advanced=True),
+                io.Float.Input("eta", default=1.0, min=0.0, max=10.0, step=0.01, round=False, tooltip="Stochastic strength of SDEs.\nWhen eta=0, they reduce to deterministic ODE.\nLarge eta may cause invalid outputs. If this occurs, try decreasing this value.", advanced=True),
                 io.Float.Input("s_noise", default=1.0, min=0.0, max=100.0, step=0.01, round=False, advanced=True),
             ],
             outputs=[io.Sampler.Output()]
@@ -599,21 +599,35 @@ class SamplerER_SDE(io.ComfyNode):
 
     @classmethod
     def execute(cls, solver_type, max_stage, eta, s_noise) -> io.NodeOutput:
-        if solver_type == "ODE" or (solver_type == "Reverse-time SDE" and eta == 0):
-            eta = 0
-            s_noise = 0
+        # Extend existing noise scalers phi(x) with eta-controlled noise scalers:
+        #   psi(x) = x**(1-eta) * phi(x)**eta
+        # where eta is constant and directly scales the h^2(t) contribution.
 
-        def reverse_time_sde_noise_scaler(x):
+        def er_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
+            return x * ((x ** 0.3).exp() + 10.0) ** eta
+
+        def reverse_time_sde_noise_scaler(x: torch.Tensor) -> torch.Tensor:
             return x ** (eta + 1)
 
-        if solver_type == "ER-SDE":
-            # Use the default one in sample_er_sde()
-            noise_scaler = None
-        else:
-            noise_scaler = reverse_time_sde_noise_scaler
+        def ode_noise_scaler(x: torch.Tensor) -> torch.Tensor:
+            return x
+
+        solver_scalers = {
+            "ER-SDE": er_sde_noise_scaler,
+            "Reverse-time SDE": reverse_time_sde_noise_scaler,
+            "ODE": ode_noise_scaler,
+        }
+
+        if solver_type == "ODE" or eta == 0:
+            s_noise = 0.0
+            solver_type = "ODE"
+        noise_scaler = solver_scalers[solver_type]
 
         sampler_name = "er_sde"
-        sampler = comfy.samplers.ksampler(sampler_name, {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage})
+        sampler = comfy.samplers.ksampler(
+            sampler_name,
+            {"s_noise": s_noise, "noise_scaler": noise_scaler, "max_stage": max_stage},
+        )
         return io.NodeOutput(sampler)
 
     get_sampler = execute
@@ -704,15 +718,7 @@ class Noise_EmptyNoise:
         self.seed = 0
 
     def generate_noise(self, input_latent):
-        latent_image = input_latent["samples"]
-        if latent_image.is_nested:
-            tensors = latent_image.unbind()
-            zeros = []
-            for t in tensors:
-                zeros.append(torch.zeros(t.shape, dtype=t.dtype, layout=t.layout, device="cpu"))
-            return comfy.nested_tensor.NestedTensor(zeros)
-        else:
-            return torch.zeros(latent_image.shape, dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+        return comfy.sample.prepare_empty_noise(input_latent["samples"])
 
 
 class Noise_RandomNoise:

comfy_extras/nodes_dataset.py

@@ -195,7 +195,7 @@ class LoadImageDataSetFromFolderNode(io.ComfyNode):
 
     @classmethod
     def execute(cls, folder):
-        sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
+        sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
         valid_extensions = [".png", ".jpg", ".jpeg", ".webp"]
         image_files = [
             f
@@ -241,7 +241,7 @@ class LoadImageTextDataSetFromFolderNode(io.ComfyNode):
     def execute(cls, folder):
         logging.info(f"Loading images from folder: {folder}")
 
-        sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
+        sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
         valid_extensions = [".png", ".jpg", ".jpeg", ".webp"]
 
         image_files = []
@@ -310,7 +310,7 @@ class LoadVideoDataSetFromFolderNode(io.ComfyNode):
 
     @classmethod
     def execute(cls, folder):
-        sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
+        sub_input_dir = secure_subfolder_path(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)
@@ -357,7 +357,7 @@ class LoadVideoTextDataSetFromFolderNode(io.ComfyNode):
 
     @classmethod
     def execute(cls, folder):
-        sub_input_dir = os.path.join(folder_paths.get_input_directory(), folder)
+        sub_input_dir = secure_subfolder_path(folder_paths.get_input_directory(), folder)
 
         video_files = []
         for item in sorted(os.listdir(sub_input_dir)):

comfy_extras/nodes_gaussian_splat.py

@@ -471,6 +471,12 @@ def _mat_to_quat(m):
 
 
 class SplatToFile3D(IO.ComfyNode):
+    FORMAT_WRITERS = {
+        "ply": _gaussian_ply_bytes,
+        "ksplat": _gaussian_ksplat_bytes,
+        "spz": _gaussian_spz_bytes,
+    }
+
     @classmethod
     def define_schema(cls):
         return IO.Schema(
@@ -482,7 +488,7 @@ class SplatToFile3D(IO.ComfyNode):
                         "Supports one item per batch only.",
             inputs=[
                 IO.Splat.Input("splat"),
-                IO.Combo.Input("format", options=["ply", "ksplat", "spz"],  # TODO: add "splat" when we have a writer for it
+                IO.Combo.Input("format", options=list(cls.FORMAT_WRITERS),  # TODO: add "splat" when we have a writer for it
                                tooltip="ply: standard 3D Gaussian Splat with full spherical harmonics. "
                                        "ksplat: mkkellogg SplatBuffer (level 0, uncompressed), base color only "
                                        "spz: Niantic gzip-compressed (~10x smaller), base color only "
@@ -493,10 +499,13 @@ class SplatToFile3D(IO.ComfyNode):
 
     @classmethod
     def execute(cls, splat, format="ply") -> IO.NodeOutput:
+        writer = cls.FORMAT_WRITERS.get(format)
+        if writer is None:
+            raise ValueError(f"Unsupported splat format: {format!r}")
+
         if splat.positions.shape[0] > 1:
             logging.warning("SplatToFile3D supports one item per batch only. Got %d; using first.", splat.positions.shape[0])
         end = _real_len(splat, 0)
-        writer = {"ksplat": _gaussian_ksplat_bytes, "spz": _gaussian_spz_bytes}.get(format, _gaussian_ply_bytes)
         data = writer(splat.positions[0, :end], splat.scales[0, :end],
                       splat.rotations[0, :end], splat.opacities[0, :end], splat.sh[0, :end])
         return IO.NodeOutput(Types.File3D(BytesIO(data), file_format=format))

comfy_extras/nodes_lt.py

@@ -2,11 +2,14 @@ import nodes
 import node_helpers
 import torch
 import torchaudio
+import comfy.ldm.lightricks.duration_head
 import comfy.model_management
 import comfy.model_sampling
 import comfy.samplers
 import comfy.utils
+import logging
 import math
+import re
 import numpy as np
 import av
 from io import BytesIO
@@ -758,22 +761,60 @@ class LTXVConcatAVLatent(io.ComfyNode):
             ],
         )
 
+    @staticmethod
+    def fit_audio(reference, audio, noise_mask):
+        """Trim or zero-pad the audio stream to the length of the one it replaces.
+
+        The padded tail is left unmasked so the model generates it, which is what a
+        clip shorter than the video should do.
+        """
+        dims = [i for i in range(reference.ndim) if reference.shape[i] != audio.shape[i]]
+        if len(dims) == 0:
+            return audio, noise_mask
+        if len(dims) > 1 or dims[0] < 2:
+            raise ValueError("audio latent {} cannot be fitted to {}".format(tuple(audio.shape), tuple(reference.shape)))
+
+        dim, length = dims[0], reference.shape[dims[0]]
+        if noise_mask is not None:  # masks carry their own shape until sampling resizes them
+            noise_mask = comfy.utils.reshape_mask(noise_mask, audio.shape)
+
+        if audio.shape[dim] > length:
+            audio = audio.narrow(dim, 0, length)
+            if noise_mask is not None:
+                noise_mask = noise_mask.narrow(dim, 0, length)
+        else:
+            pad = torch.zeros_like(audio.narrow(dim, 0, 1)).repeat(
+                [length - audio.shape[dim] if i == dim else 1 for i in range(audio.ndim)])
+            audio = torch.cat([audio, pad], dim=dim)
+            if noise_mask is not None:
+                noise_mask = torch.cat([noise_mask, torch.ones_like(pad)], dim=dim)
+        return audio, noise_mask
+
     @classmethod
     def execute(cls, video_latent, audio_latent) -> io.NodeOutput:
         output = {}
         output.update(video_latent)
         output.update(audio_latent)
+        video_samples = video_latent["samples"]
+        audio_samples = audio_latent["samples"]
         video_noise_mask = video_latent.get("noise_mask", None)
         audio_noise_mask = audio_latent.get("noise_mask", None)
 
+        if video_samples.is_nested:  # already an AV latent: keep its video and swap the audio stream
+            streams = video_samples.unbind()
+            video_samples = streams[0]
+            if video_noise_mask is not None:
+                video_noise_mask = video_noise_mask.unbind()[0]
+            audio_samples, audio_noise_mask = cls.fit_audio(streams[1], audio_samples, audio_noise_mask)
+
         if video_noise_mask is not None or audio_noise_mask is not None:
             if video_noise_mask is None:
-                video_noise_mask = torch.ones_like(video_latent["samples"])
+                video_noise_mask = torch.ones_like(video_samples)
             if audio_noise_mask is None:
-                audio_noise_mask = torch.ones_like(audio_latent["samples"])
+                audio_noise_mask = torch.ones_like(audio_samples)
             output["noise_mask"] = comfy.nested_tensor.NestedTensor((video_noise_mask, audio_noise_mask))
 
-        output["samples"] = comfy.nested_tensor.NestedTensor((video_latent["samples"], audio_latent["samples"]))
+        output["samples"] = comfy.nested_tensor.NestedTensor((video_samples, audio_samples))
 
         return io.NodeOutput(output)
 
@@ -896,6 +937,243 @@ class LTXVReferenceAudio(io.ComfyNode):
         return io.NodeOutput(m, positive, negative)
 
 
+class LTXVSpatioTemporalGuidance(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVSpatioTemporalGuidance",
+            display_name="LTXV Spatio-Temporal Guidance (STG)",
+            category="advanced/guidance",
+            description="Runs one extra pass per step with the self-attention of the selected blocks degraded to a value-passthrough, "
+                        "then guides away from it - improving spatial detail and motion coherence.",
+            inputs=[
+                io.Model.Input("model"),
+                io.Float.Input("scale", default=1.0, min=0.0, max=100.0, step=0.01, round=0.01),
+                io.String.Input("blocks", default="29", tooltip="Comma-separated transformer block indices to perturb."),
+                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001, advanced=True),
+                io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001, advanced=True),
+            ],
+            outputs=[io.Model.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, scale, blocks, start_percent, end_percent) -> io.NodeOutput:
+        block_set = frozenset(int(b) for b in re.findall(r"\d+", blocks))
+
+        m = model.clone()
+        model_sampling = m.get_model_object("model_sampling")
+        sigma_start = model_sampling.percent_to_sigma(start_percent)
+        sigma_end = model_sampling.percent_to_sigma(end_percent)
+
+        def post_cfg_function(args):
+            if scale == 0 or not block_set:
+                return args["denoised"]
+
+            sigma_ = args["sigma"][0].item()
+            if sigma_ > sigma_start or sigma_ < sigma_end:
+                return args["denoised"]
+
+            cond_pred = args["cond_denoised"]
+            cond = args["cond"]
+            cfg_result = args["denoised"]
+            x = args["input"]
+
+            model_options = args["model_options"].copy()
+            transformer_options = model_options.get("transformer_options", {}).copy()
+            transformer_options["stg_self_attn_blocks"] = block_set
+            model_options["transformer_options"] = transformer_options
+
+            (perturbed,) = comfy.samplers.calc_cond_batch(args["model"], [cond], x, args["sigma"], model_options)
+
+            return cfg_result + (cond_pred - perturbed) * scale
+
+        m.set_model_sampler_post_cfg_function(post_cfg_function)
+        return io.NodeOutput(m)
+
+
+class LTXVModalityGuidance(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVModalityGuidance",
+            display_name="LTXV Modality Guidance (A/V coupling)",
+            category="advanced/guidance",
+            description="Cross-modal (audio-video) guidance for LTXV-AV. Runs one extra forward "
+                        "pass per step with the a2v/v2a cross-attention severed, then pushes the "
+                        "result toward the coupled prediction - strengthening audio-visual sync "
+                        "(e.g. lip-sync). Reference default modality_scale is 3.0. Stacks with the "
+                        "dual-CFG guider and STG. Set to 1.0 to disable (no extra pass).",
+            inputs=[
+                io.Model.Input("model"),
+                io.Float.Input("modality_scale", default=3.0, min=1.0, max=100.0, step=0.1, round=0.01),
+                io.Float.Input("start_percent", default=0.0, min=0.0, max=1.0, step=0.001, advanced=True),
+                io.Float.Input("end_percent", default=1.0, min=0.0, max=1.0, step=0.001, advanced=True),
+            ],
+            outputs=[io.Model.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, modality_scale, start_percent, end_percent) -> io.NodeOutput:
+        m = model.clone()
+        model_sampling = m.get_model_object("model_sampling")
+        sigma_start = model_sampling.percent_to_sigma(start_percent)
+        sigma_end = model_sampling.percent_to_sigma(end_percent)
+
+        def post_cfg_function(args):
+            if math.isclose(modality_scale, 1.0):
+                return args["denoised"]
+
+            sigma_ = args["sigma"][0].item()
+            if sigma_ > sigma_start or sigma_ < sigma_end:
+                return args["denoised"]
+
+            cond_pred = args["cond_denoised"]
+            cond = args["cond"]
+            cfg_result = args["denoised"]
+            x = args["input"]
+
+            # Extra pass with audio-video cross-attention severed (both directions)
+            model_options = args["model_options"].copy()
+            transformer_options = model_options.get("transformer_options", {}).copy()
+            transformer_options["a2v_cross_attn"] = False
+            transformer_options["v2a_cross_attn"] = False
+            model_options["transformer_options"] = transformer_options
+
+            (mod_pred,) = comfy.samplers.calc_cond_batch(
+                args["model"], [cond], x, args["sigma"], model_options
+            )
+
+            # (modality_scale - 1) * (cond - uncond_modality), per the reference guider.
+            return cfg_result + (cond_pred - mod_pred) * (modality_scale - 1.0)
+
+        m.set_model_sampler_post_cfg_function(post_cfg_function)
+        return io.NodeOutput(m)
+
+
+class Guider_LTXAVDualCFG(comfy.samplers.CFGGuider):
+    """CFG guider that applies separate guidance scales to the video and audio
+    modalities of a packed LTXV-AV latent.
+    """
+
+    def set_conds(self, positive, negative):
+        self.inner_set_conds({"positive": positive, "negative": negative})
+
+    def set_cfg(self, video_cfg, audio_cfg):
+        self.video_cfg = video_cfg
+        self.audio_cfg = audio_cfg
+        self.cfg = max(video_cfg, audio_cfg)
+
+    def sample(self, noise, latent_image, *args, **kwargs):
+        # Capture the video/audio split from the nested latent before it is packed.
+        self._v_numel = None
+        if getattr(latent_image, "is_nested", False):
+            parts = latent_image.unbind()
+            if len(parts) >= 2:
+                self._v_numel = math.prod(parts[0].shape[1:])
+        return super().sample(noise, latent_image, *args, **kwargs)
+
+    def predict_noise(self, x, timestep, model_options={}, seed=None):
+        v = getattr(self, "_v_numel", None)
+        if v is None or math.isclose(self.video_cfg, self.audio_cfg):
+            # Not an AV latent, or equal scales: fall back to standard single-CFG.
+            self.cfg = self.video_cfg
+            return super().predict_noise(x, timestep, model_options, seed)
+
+        video_cfg, audio_cfg = self.video_cfg, self.audio_cfg
+
+        def dual_cfg(args):
+            # Noise-space: cond = x - cond_pred, uncond = x - uncond_pred; the
+            # returned tensor is subtracted from x by cfg_function.
+            cond, uncond = args["cond"], args["uncond"]
+            out = uncond + (cond - uncond) * video_cfg
+            out[..., v:] = uncond[..., v:] + (cond[..., v:] - uncond[..., v:]) * audio_cfg
+            return out
+
+        # disable_cfg1_optimization so the uncond pass always runs even if one of the two scales is 1.0.
+        model_options = {**model_options, "sampler_cfg_function": dual_cfg, "disable_cfg1_optimization": True}
+        return super().predict_noise(x, timestep, model_options, seed)
+
+
+class LTXVDualCFGGuider(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVDualCFGGuider",
+            display_name="LTXV Dual CFG Guider",
+            category="model/sampling/guiders",
+            description="Separate CFG scales for the video and audio modalities of a packed LTXV-AV latent.",
+            inputs=[
+                io.Model.Input("model"),
+                io.Conditioning.Input("positive"),
+                io.Conditioning.Input("negative"),
+                io.Float.Input("video_cfg", default=3.0, min=0.0, max=100.0, step=0.1, round=0.01),
+                io.Float.Input("audio_cfg", default=7.0, min=0.0, max=100.0, step=0.1, round=0.01),
+            ],
+            outputs=[io.Guider.Output()],
+        )
+
+    @classmethod
+    def execute(cls, model, positive, negative, video_cfg, audio_cfg) -> io.NodeOutput:
+        guider = Guider_LTXAVDualCFG(model)
+        guider.set_conds(positive, negative)
+        guider.set_cfg(video_cfg, audio_cfg)
+        return io.NodeOutput(guider)
+
+
+class LTXVDurationPredictor(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="LTXVDurationPredictor",
+            display_name="LTXV Duration Predictor",
+            category="conditioning/video_models",
+            description="Predicts the natural shot duration for a prompt using the LTX 2.4 duration "
+                        "head (loaded with ModelPatchLoader), and snaps it to the VAE's 8k+1 frame grid.",
+            search_aliases=["auto duration", "duration head", "num_frames"],
+            inputs=[
+                io.Model.Input("model"),
+                io.Conditioning.Input("positive"),
+                io.Custom("MODEL_PATCH").Input("duration_head",
+                                               tooltip="LTX 2.4 duration head loaded with ModelPatchLoader."),
+                io.Float.Input("frame_rate", default=24.0, min=1.0, max=120.0, step=0.01),
+                io.Float.Input("min_seconds", default=1.0, min=0.5, max=120.0, step=0.1),
+                io.Float.Input("max_seconds", default=20.0, min=0.5, max=120.0, step=0.1),
+            ],
+            outputs=[
+                io.Int.Output(display_name="num_frames"),
+                io.Float.Output(display_name="seconds", tooltip="Raw (unclamped) predicted duration."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model, positive, duration_head, frame_rate, min_seconds, max_seconds) -> io.NodeOutput:
+        dm = model.model.diffusion_model
+        head = duration_head.model
+        if not isinstance(head, comfy.ldm.lightricks.duration_head.DurationHead):
+            raise ValueError("The connected model_patch is not an LTX duration head.")
+
+        context = positive[0][0]
+        meta = positive[0][1]
+        if context.shape[0] != 1:
+            context = context[:1]
+
+        # Run the caption connectors exactly the way sampling does.
+        comfy.model_management.load_models_gpu([model, duration_head])
+        device = model.load_device
+        head = head.to(device)
+        with torch.no_grad():
+            context = context.to(device=device, dtype=model.model.get_dtype_inference())
+            processed = dm.preprocess_text_embeds(context, unprocessed=meta.get("unprocessed_ltxav_embeds", False))
+            video_tokens = processed[..., :dm.cross_attention_dim].float()
+            audio_tokens = processed[..., dm.cross_attention_dim:].float()
+            seconds = float(head(video_tokens, audio_tokens)[0])
+
+        num_frames = comfy.ldm.lightricks.duration_head.seconds_to_num_frames(
+            seconds, frame_rate, min_seconds, max_seconds)
+        logging.info("LTXV duration head predicted %.2fs -> %d frames @ %.2f fps", seconds, num_frames, frame_rate)
+        return io.NodeOutput(num_frames, seconds)
+
+
 class LtxvExtension(ComfyExtension):
     @override
     async def get_node_list(self) -> list[type[io.ComfyNode]]:
@@ -913,6 +1191,10 @@ class LtxvExtension(ComfyExtension):
             LTXVConcatAVLatent,
             LTXVSeparateAVLatent,
             LTXVReferenceAudio,
+            LTXVDualCFGGuider,
+            LTXVModalityGuidance,
+            LTXVSpatioTemporalGuidance,
+            LTXVDurationPredictor,
         ]
 
 

comfy_extras/nodes_lt_audio.py

@@ -173,7 +173,7 @@ class LTXAVTextEncoderLoader(io.ComfyNode):
             node_id="LTXAVTextEncoderLoader",
             display_name="Load LTXV Audio Text Encoder",
             category="model/loaders",
-            description="Recipes:\nltxav: gemma 3 12B",
+            description="Recipes:\nltxav: gemma 3 12B or matching gemma 4 model",
             inputs=[
                 io.Combo.Input(
                     "text_encoder",

comfy_extras/nodes_minimax_h3.py

@@ -20,6 +20,7 @@ import comfy.model_sampling
 import comfy.nested_tensor
 import comfy.utils
 import node_helpers
+from comfy.ldm.minimax.model import FRAME_PER_TOKEN, FRAME_RESCALE
 from comfy_api.latest import ComfyExtension, io
 
 CANVAS_MULTIPLE = 32
@@ -67,6 +68,16 @@ def _resize(image, width, height, crop):
     return samples.movedim(1, -1)
 
 
+def _encode_ref_audio(audio_vae, audio):
+    waveform = audio["waveform"]  # [B, C, L]
+    sr = audio["sample_rate"]
+    vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
+    if sr != vae_sr:
+        waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
+    z = audio_vae.encode(waveform[:1].movedim(1, -1))  # [1, 32, 2, T]
+    return z, z.shape[-1]
+
+
 def _empty_av_latent(width, height, length, batch_size=1):
     frame_count, latent_t, audio_t = temporal_shape(length)
     video = torch.zeros([batch_size, 24, latent_t, height // 16, width // 16],
@@ -144,13 +155,87 @@ class MiniMaxH3ImageToVideo(io.ComfyNode):
         if keyframes:
             for kf in keyframes:
                 kf["latent"] = vae.encode(kf.pop("image"))
-            cond = node_helpers.conditioning_set_values(cond, {
-                "minimax_keyframes": keyframes,
-                "minimax_frame_count": frame_count,
-            })
+            cond = node_helpers.conditioning_set_values(cond, {"minimax_keyframes": keyframes})
         return io.NodeOutput(cond, latent)
 
 
+class MiniMaxH3AddGuide(io.ComfyNode):
+    """Anchor image and/or audio guides at an arbitrary pixel frame of the target video."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxH3AddGuide",
+            display_name="Add Guide for MiniMax H3",
+            category="model/conditioning/minimax",
+            description="Anchor an image, a short clip, audio, or a clip with its soundtrack at any frame of a MiniMax H3 video. Chain several nodes to anchor several frames.",
+            inputs=[
+                io.Conditioning.Input("positive"),
+                io.Vae.Input("vae", optional=True, tooltip="Video VAE, needed when an image is connected."),
+                io.Vae.Input("audio_vae", optional=True, tooltip="Audio VAE, needed when an audio is connected."),
+                io.Latent.Input("latent"),
+                io.Image.Input("image", optional=True, tooltip="Image or video frames to anchor. Multi-frame batches are anchored as a clip and cropped down to the model's valid clip lengths: 5, 22, 39... (17k + 5) frames. Batches shorter than 5 frames use only the first image."),
+                io.Audio.Input("audio", optional=True,
+                               tooltip="Soundtrack to anchor starting at the same frame index, cropped to the video's remaining duration."),
+                io.Int.Input("frame_idx", default=0, min=-9999, max=9999,
+                             tooltip="Frame index to anchor the image or the clip's first frame at. Negative values are counted from the end of the video."),
+            ],
+            outputs=[io.Conditioning.Output(display_name="positive")],
+        )
+
+    @classmethod
+    def execute(cls, positive, latent, frame_idx, vae=None, audio_vae=None, image=None, audio=None) -> io.NodeOutput:
+        samples = latent["samples"]
+        if not samples.is_nested or len(samples.tensors) != 2 or samples.tensors[0].ndim != 5 or samples.tensors[0].shape[1] != 24:
+            raise ValueError("MiniMaxH3AddGuide expects a MiniMax H3 AV latent")
+        if image is None and audio is None:
+            raise ValueError("MiniMaxH3AddGuide needs an image or an audio to anchor")
+        video = samples.tensors[0]
+        height = video.shape[3] * 16
+        width = video.shape[4] * 16
+        frame_count = sum(FRAME_PER_TOKEN[k % 5] for k in range(video.shape[2]))
+
+        guide_frames = 1
+        if image is not None:
+            if vae is None:
+                raise ValueError("anchoring guide frames needs the vae input")
+            guide_frames = image.shape[0]
+            if guide_frames < 5:
+                guide_frames = 1
+            else:
+                while guide_frames % 17 != 5:
+                    guide_frames -= 1
+
+        resolved_frame_index = frame_idx if frame_idx >= 0 else frame_count + frame_idx
+        if resolved_frame_index < 0 or resolved_frame_index + guide_frames > frame_count:
+            if guide_frames == 1:
+                raise ValueError("frame_idx {} is outside the video's {} frames".format(frame_idx, frame_count))
+            raise ValueError("a {} frame guide clip at frame_idx {} does not fit in the video's {} frames".format(
+                guide_frames, frame_idx, frame_count))
+
+        keyframe = {"resolved_frame_index": resolved_frame_index}
+        if image is not None:
+            frames = _resize(image[:guide_frames], width, height, "center")
+            keyframe["latent"] = vae.encode(frames)
+
+        if audio is not None:
+            if audio_vae is None:
+                raise ValueError("anchoring guide audio needs the audio_vae input")
+            audio_latent, audio_rt = _encode_ref_audio(audio_vae, audio)
+            # the streams share one time axis: FRAME_RESCALE per pixel frame, 1.0 per audio latent frame
+            max_rt = math.floor(samples.tensors[1].shape[-1] - FRAME_RESCALE * resolved_frame_index)
+            if max_rt < 1:
+                raise ValueError("frame_idx {} is past the end of the video's audio track".format(frame_idx))
+            if audio_rt > max_rt:
+                audio_latent = audio_latent[..., :max_rt].clone()
+            keyframe["audio_latent"] = audio_latent
+
+        keyframes = list(positive[0][1].get("minimax_keyframes", []))
+        keyframes.append(keyframe)
+        positive = node_helpers.conditioning_set_values(positive, {"minimax_keyframes": keyframes})
+        return io.NodeOutput(positive)
+
+
 class MiniMaxH3ReferenceToVideo(io.ComfyNode):
     """ref2va: prompt + reference images / videos / audio -> conditioning + AV latent.
 
@@ -197,16 +282,6 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
             outputs=[io.Conditioning.Output(display_name="positive"), io.Latent.Output()],
         )
 
-    @staticmethod
-    def _encode_ref_audio(audio_vae, audio):
-        waveform = audio["waveform"]  # [B, C, L]
-        sr = audio["sample_rate"]
-        vae_sr = getattr(audio_vae, "audio_sample_rate", 32000)
-        if sr != vae_sr:
-            waveform = torchaudio.functional.resample(waveform, sr, vae_sr)
-        z = audio_vae.encode(waveform[:1].movedim(1, -1))  # [1, 32, 2, T]
-        return z, z.shape[-1]
-
     @classmethod
     def execute(cls, clip, vae, audio_vae, prompt, width, height, length, ref_image_size="match",
                 ref_images=None, ref_videos=None, ref_video_audios=None, ref_audios=None) -> io.NodeOutput:
@@ -254,7 +329,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
             z = vae.encode(frames)
             audio_latent, ref_audio_t = (None, 0)
             if soundtrack is not None:
-                audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, soundtrack)
+                audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, soundtrack)
                 # the soundtrack gets its own <Audio j> label, emitted before <Video k>
                 ref_items.append({"type": "audio"})
             # Qwen sees the video at 2 fps with timestamps
@@ -269,7 +344,7 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
         for audio in (ref_audios or {}).values():
             if audio is None:
                 continue
-            audio_latent, ref_audio_t = cls._encode_ref_audio(audio_vae, audio)
+            audio_latent, ref_audio_t = _encode_ref_audio(audio_vae, audio)
             ref_items.append({"type": "audio"})
             ref_blocks.append({"kind": "audio", "ref_audio_t": ref_audio_t, "audio_latent": audio_latent})
 
@@ -283,9 +358,9 @@ class MiniMaxH3ReferenceToVideo(io.ComfyNode):
 class MiniMaxH3SigmaShift(io.ComfyNode):
     """Set the video/audio flow shifts coherently.
 
-    The video shift drives the sampler's sigma schedule; both values are also
-    handed to the DiT, which inverts the video schedule to the shared base grid
-    and derives the audio schedule from it.
+    The video shift drives the sampler's sigma schedule (ModelSamplingAV); both
+    values are also handed to the DiT, which inverts the video schedule to the
+    shared base grid and derives the audio schedule from it.
     """
 
     @classmethod
@@ -293,7 +368,8 @@ class MiniMaxH3SigmaShift(io.ComfyNode):
         return io.Schema(
             node_id="MiniMaxH3SigmaShift",
             description="Set the video/audio flow shifts.",
-            display_name="MiniMax H3 Sigma Shift",
+            display_name="ModelSamplingMiniMaxH3",
+            search_aliases=["sigma shift", "minimax shift"],
             category="model/patch/minimax",
             inputs=[
                 io.Model.Input("model"),
@@ -307,12 +383,12 @@ class MiniMaxH3SigmaShift(io.ComfyNode):
     def execute(cls, model, shift_video, shift_audio) -> io.NodeOutput:
         m = model.clone()
 
-        class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingDiscreteFlow, comfy.model_sampling.CONST):
+        class ModelSamplingAdvanced(comfy.model_sampling.ModelSamplingAV, comfy.model_sampling.CONST):
             pass
 
         original = m.get_model_object("model_sampling")
         model_sampling = ModelSamplingAdvanced(model.model.model_config)
-        model_sampling.set_parameters(shift=shift_video)
+        model_sampling.set_parameters(shift=shift_video, audio_shift=shift_audio)
         if hasattr(original, "noise_scale"):
             model_sampling.set_noise_scale(original.noise_scale)
         m.add_object_patch("model_sampling", model_sampling)
@@ -328,8 +404,9 @@ class MiniMaxH3Extension(ComfyExtension):
         return [
             EmptyMiniMaxH3LatentAV,
             MiniMaxH3ImageToVideo,
+            MiniMaxH3AddGuide,
             MiniMaxH3ReferenceToVideo,
-            MiniMaxH3SigmaShift
+            MiniMaxH3SigmaShift,
             ]
 
 

comfy_extras/nodes_minimax_music.py

@@ -0,0 +1,77 @@
+import torch
+from typing_extensions import override
+
+import comfy.model_management
+from comfy.ldm.minimax_music.ar import AUDIO_FRAMES_PER_SECOND, CFG_SCALE, CFG_TOP_K, C0_VOCAB_SIZE, MAX_AUDIO_FRAMES
+from comfy.ldm.minimax_music.dit import latent_length
+from comfy_api.latest import ComfyExtension, io
+
+
+class MiniMaxMusic3TextEncode(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="MiniMaxMusic3TextEncode",
+            display_name="MiniMax Music3 Text Encode",
+            category="model/conditioning/minimax music",
+            description="Uses a MiniMax Music3 CLIP model to generate the acoustic conditioning sequence.",
+            inputs=[
+                io.Clip.Input("clip"),
+                io.String.Input("caption", multiline=True, dynamic_prompts=True),
+                io.String.Input("lyrics", multiline=True, dynamic_prompts=True),
+                io.Int.Input("seed", default=0, min=0, max=0xffffffffffffffff, control_after_generate=True),
+                io.Float.Input("max_duration", default=120.0, min=0.04, max=MAX_AUDIO_FRAMES / AUDIO_FRAMES_PER_SECOND, step=0.04, tooltip="Maximum duration in seconds; the model can end the song earlier."),
+                io.Float.Input("cfg_scale", default=CFG_SCALE, min=0.0, max=100.0, step=0.1, round=0.01, advanced=True),
+                io.Int.Input("top_k", default=CFG_TOP_K, min=1, max=C0_VOCAB_SIZE, advanced=True),
+            ],
+            outputs=[
+                io.Conditioning.Output(),
+                io.Float.Output(display_name="seconds"),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, clip, caption, lyrics, seed, max_duration, cfg_scale, top_k):
+        max_audio_frames = min(MAX_AUDIO_FRAMES, max(1, round(max_duration * AUDIO_FRAMES_PER_SECOND)))
+        tokens = clip.tokenize(caption, lyrics=lyrics, seed=seed, max_audio_frames=max_audio_frames, cfg_scale=cfg_scale, top_k=top_k)
+        conditioning = clip.encode_from_tokens_scheduled(tokens)
+        for cond in conditioning:
+            hidden = cond[0]
+            cond[1]["conditioning_scale"] = torch.ones((hidden.shape[0], 1, 1), device=hidden.device, dtype=hidden.dtype)
+        return io.NodeOutput(conditioning, conditioning[0][0].shape[1] / AUDIO_FRAMES_PER_SECOND)
+
+
+class EmptyMiniMaxMusic3LatentAudio(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="EmptyMiniMaxMusic3LatentAudio",
+            display_name="Empty MiniMax Music3 Latent Audio",
+            category="model/latent/minimax music",
+            description="Creates an empty MiniMax Music3 audio latent for the requested duration.",
+            inputs=[
+                io.Float.Input("seconds", default=120.0, min=0.04, max=MAX_AUDIO_FRAMES / AUDIO_FRAMES_PER_SECOND, step=0.04),
+                io.Int.Input("batch_size", default=1, min=1, max=4096),
+            ],
+            outputs=[io.Latent.Output()],
+        )
+
+    @classmethod
+    def execute(cls, seconds, batch_size):
+        audio_frames = min(MAX_AUDIO_FRAMES, max(1, round(seconds * AUDIO_FRAMES_PER_SECOND)))
+        latent = torch.zeros(
+            (batch_size, 128, latent_length(audio_frames)),
+            device=comfy.model_management.intermediate_device(),
+            dtype=comfy.model_management.intermediate_dtype(),
+        )
+        return io.NodeOutput({"samples": latent, "type": "audio", "downscale_ratio_temporal": 512})
+
+
+class MiniMaxMusic3Extension(ComfyExtension):
+    @override
+    async def get_node_list(self):
+        return [MiniMaxMusic3TextEncode, EmptyMiniMaxMusic3LatentAudio]
+
+
+async def comfy_entrypoint():
+    return MiniMaxMusic3Extension()

comfy_extras/nodes_model_advanced.py

@@ -1,6 +1,9 @@
+import logging
+
 import comfy.sd
 import comfy.model_sampling
 import comfy.latent_formats
+import comfy.ldm.modules.attention
 import nodes
 import torch
 import node_helpers
@@ -346,6 +349,39 @@ class ModelComputeDtype:
         return (m, )
 
 
+class ModelAttentionBackend:
+    @classmethod
+    def INPUT_TYPES(s):
+        backends = ["pytorch attention"]
+        if comfy.ldm.modules.attention.COMFY_KITCHEN_INT8_ATTENTION_IS_AVAILABLE:
+            backends.append("comfy kitchen attention")
+        return {"required": {"model": ("MODEL",),
+                             "attention": (backends,),
+                             }}
+
+    @classmethod
+    def VALIDATE_INPUTS(s, attention):
+        return True
+
+    RETURN_TYPES = ("MODEL",)
+    FUNCTION = "patch"
+
+    CATEGORY = "model/patch"
+
+    def patch(self, model, attention):
+        attention_name = {
+            "comfy kitchen attention": "comfy_kitchen_int8",
+            "pytorch attention": "pytorch",
+        }.get(attention)
+        attention_function = comfy.ldm.modules.attention.get_attention_function(attention_name, None)
+        if attention_function is None:
+            logging.warning("Attention backend '%s' is unavailable; using PyTorch attention.", attention)
+            attention_function = comfy.ldm.modules.attention.get_attention_function("pytorch")
+        m = model.clone()
+        m.set_model_optimized_attention(attention_function)
+        return (m, )
+
+
 NODE_CLASS_MAPPINGS = {
     "ModelSamplingDiscrete": ModelSamplingDiscrete,
     "ModelSamplingContinuousEDM": ModelSamplingContinuousEDM,
@@ -357,4 +393,5 @@ NODE_CLASS_MAPPINGS = {
     "ModelNoiseScale": ModelNoiseScale,
     "RescaleCFG": RescaleCFG,
     "ModelComputeDtype": ModelComputeDtype,
+    "ModelAttentionBackend": ModelAttentionBackend,
 }

comfy_extras/nodes_model_patch.py

@@ -10,6 +10,7 @@ import comfy.ldm.lumina.controlnet
 import comfy.ldm.supir.supir_modules
 import comfy.ldm.anima.lllite
 import comfy.ldm.wan.uni3c
+import comfy.ldm.lightricks.duration_head
 from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
 from comfy_api.latest import io
 from comfy.ldm.supir.supir_patch import SUPIRPatch
@@ -296,6 +297,10 @@ class ModelPatchLoader:
                     device=comfy.model_management.unet_offload_device(),
                     dtype=dtype,
                     operations=comfy.ops.manual_cast)
+        elif any(k.endswith("duration_head.attention_pooler.query_tokens") for k in sd) or "attention_pooler.query_tokens" in sd:
+            sd = comfy.ldm.lightricks.duration_head.normalize_state_dict(sd)
+            sd = {k: v.float() for k, v in sd.items()}  # tiny head, keep fp32
+            model = comfy.ldm.lightricks.duration_head.DurationHead()
         elif "audio_proj.proj1.weight" in sd:
             model = MultiTalkModelPatch(
                     audio_window=5, context_tokens=32, vae_scale=4,

comfy_extras/nodes_preview_any.py

@@ -29,7 +29,7 @@ class PreviewAny():
             value = str(source)
         elif source is not None:
             try:
-                value = json.dumps(source, indent=4)
+                value = json.dumps(source, indent=4, ensure_ascii=False)
             except Exception:
                 try:
                     value = str(source)

comfy_extras/nodes_textgen.py

@@ -1,3 +1,4 @@
+import re
 from comfy_api.latest import ComfyExtension, io
 from typing_extensions import override
 
@@ -152,6 +153,64 @@ You are a Creative Assistant writing concise, action-focused image-to-video prom
 Style: realistic - cinematic - The woman glances at her watch and smiles warmly. She speaks in a cheerful, friendly voice, "I think we're right on time!" In the background, a café barista prepares drinks at the counter. The barista calls out in a clear, upbeat tone, "Two cappuccinos ready!" The sound of the espresso machine hissing softly blends with gentle background chatter and the light clinking of cups on saucers.
 """
 
+LTX24_T2V_SYSTEM_PROMPT = """You are given a user's short text-to-video request. Write a single, highly detailed audio-visual caption describing the video that best fulfills that request, in the EXACT style of the training captions used for this video model. The generated video is scored against the user's ORIGINAL request, so preserve every element the user stated; expand faithfully into the full caption style without contradicting or dropping anything they asked for.
+
+Match this captioning style precisely:
+
+1. Begin immediately with the action or visual detail. Do NOT use "The scene opens…", "We see…", "There is…".
+
+2. Objective, observable description only. Do not infer emotions or intentions — describe what is visible and audible (e.g. not "he looks sad" but "his eyebrows angle downward and his lips are pressed together").
+
+3. Full visual detail: environment (materials, textures, lighting, colors), character appearance (clothing, posture, facial details), and the spatial positioning of all elements. When a human appears, identify them specifically (gendered terms when clearly implied; differentiate multiple people consistently) and describe visible physical attributes — apparent gender presentation, skin tone, estimated age group, hair color/length/style, build, clothing and accessories. Do not infer ethnicity, nationality, religion, or culture.
+
+4. Precise motion and cinematic description. For every shot you MUST include, woven naturally into the prose (never as tags or labels):
+   - Shot type (exactly one: extreme wide shot / wide shot / medium shot / medium close-up / close-up / extreme close-up)
+   - Camera motion (always stated; if none, explicitly say the camera remains static). Camera movement is expected and good — match the user if they specified it, otherwise choose the treatment that best presents the requested scene.
+   - Camera viewpoint relative to subject (front-facing / back-facing / side view / over-the-shoulder / top-down / low-angle / high-angle).
+   Express these as flowing prose: "a medium shot frames…, captured from a front-facing angle as the camera slowly pans…". Never as "medium shot, static camera —".
+
+5. Complete soundscape, integrated naturally: any dialogue (quote it exactly, in the original language), tone of voice, background music (type, mood, volume changes), and environmental sounds (footsteps, wind, traffic, animals). If the request implies sound, describe it plausibly.
+
+6. Strict chronological, real-time flow using transitions like "Initially…", "A moment later…", "Simultaneously…". Keep every stated action in motion.
+
+7. One single continuous paragraph. No bullet points, no section headers, no labels like "Audio:" or "Visual:". Exhaustive and lossless — include background elements, subtle movements, lighting, secondary sounds — detailed enough to reconstruct the scene. Aim for a rich, complete paragraph (roughly 150–220 words).
+
+If the user wrote in another language, produce the English caption of the same content. Output ONLY the caption text — no JSON, no preamble.
+
+AESTHETIC QUALITY (in addition to the above, without breaking the objective caption style): render the described scene with strong visual production value — cinematic, film-grade color and contrast, beautiful natural lighting, crisp fine detail and texture, pleasing composition and depth. Weave these quality descriptors naturally into the same observable prose (e.g. "warm cinematic lighting", "richly saturated film-grade color", "crisp high-resolution detail") — describe how the exact requested scene LOOKS at its most visually striking, never adding new objects or actions. Keep everything else (framing triple, soundscape, chronological single paragraph, faithfulness) exactly as specified.
+"""
+
+
+LTX24_I2V_SYSTEM_PROMPT = """You are given a REFERENCE IMAGE (the exact first frame of the video) and a user's short image-to-video request. Write a single, highly detailed audio-visual caption describing the video that BEGINS from this exact reference image and best fulfills that request, in the EXACT style of the training captions used for this video model. The generated video is scored against the user's ORIGINAL request, so preserve every element the user stated; expand faithfully into the full caption style without contradicting or dropping anything they asked for.
+
+FIRST-FRAME / IMAGE GROUNDING (do this first): the opening of your caption must match the reference image exactly — same subject(s), identity, appearance, clothing, setting, lighting, and composition as shown. The video starts on this frame; describe it faithfully, then narrate chronologically as the user's requested action unfolds from it. Never contradict, replace, or invent things not consistent with the image. Single continuous take — no hard cuts.
+
+Match this captioning style precisely:
+
+1. Begin immediately with the action or visual detail. Do NOT use "The scene opens…", "We see…", "There is…".
+
+2. Objective, observable description only. Do not infer emotions or intentions — describe what is visible and audible (e.g. not "he looks sad" but "his eyebrows angle downward and his lips are pressed together").
+
+3. Full visual detail: environment (materials, textures, lighting, colors), character appearance (clothing, posture, facial details), and the spatial positioning of all elements — grounded in and consistent with the reference image. When a human appears, identify them specifically (gendered terms when clearly implied; differentiate multiple people consistently) and describe visible physical attributes — apparent gender presentation, skin tone, estimated age group, hair color/length/style, build, clothing and accessories. Do not infer ethnicity, nationality, religion, or culture.
+
+4. Precise motion and cinematic description. For every shot you MUST include, woven naturally into the prose (never as tags or labels):
+   - Shot type (exactly one: extreme wide shot / wide shot / medium shot / medium close-up / close-up / extreme close-up) — consistent with how the reference image is framed at the start.
+   - Camera motion (always stated; if none, explicitly say the camera remains static). Camera movement is expected and good — match the user if they specified it, otherwise choose the treatment that best presents the requested scene starting from this frame.
+   - Camera viewpoint relative to subject (front-facing / back-facing / side view / over-the-shoulder / top-down / low-angle / high-angle) — matching the reference image's viewpoint at the opening.
+   Express these as flowing prose: "a medium shot frames…, captured from a front-facing angle as the camera slowly pans…". Never as "medium shot, static camera —".
+
+5. Complete soundscape, integrated naturally: any dialogue (quote it exactly, in the original language), tone of voice, background music (type, mood, volume changes), and environmental sounds (footsteps, wind, traffic, animals). If the request implies sound, describe it plausibly.
+
+6. Strict chronological, real-time flow using transitions like "Initially…", "A moment later…", "Simultaneously…". Keep the user's requested motion/action central and in motion throughout.
+
+7. One single continuous paragraph. No bullet points, no section headers, no labels like "Audio:" or "Visual:". Exhaustive and lossless — include background elements, subtle movements, lighting, secondary sounds — detailed enough to reconstruct the scene. Aim for a rich, complete paragraph (roughly 150–220 words).
+
+If the user wrote in another language, produce the English caption of the same content. Output ONLY the caption text — no JSON, no preamble.
+
+AESTHETIC QUALITY (in addition to the above, without breaking the objective caption style or contradicting the reference image): render the described scene with strong visual production value — cinematic, film-grade color and contrast, beautiful natural lighting, crisp fine detail and texture, pleasing composition and depth. Weave these quality descriptors naturally into the same observable prose (e.g. "warm cinematic lighting", "richly saturated film-grade color", "crisp high-resolution detail") — describe how the exact requested scene, starting from this frame, LOOKS at its most visually striking, never adding new objects or actions and never contradicting the first frame. Keep everything else (first-frame grounding, framing triple, soundscape, chronological single paragraph, faithfulness) exactly as specified.
+"""
+
+
 class TextGenerateLTX2Prompt(TextGenerate):
     @classmethod
     def define_schema(cls):
@@ -167,11 +226,40 @@ class TextGenerateLTX2Prompt(TextGenerate):
 
     @classmethod
     def execute(cls, clip, prompt, max_length, sampling_mode, image=None, thinking=False, use_default_template=True, video=None, audio=None) -> io.NodeOutput:
-        if image is None:
-            formatted_prompt = f"<start_of_turn>system\n{LTX2_T2V_SYSTEM_PROMPT.strip()}<end_of_turn>\n<start_of_turn>user\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n<start_of_turn>model\n"
+        # Gemma 3 and Gemma 4 use different chat-turn markers and image tokens.
+        # The Gemma 4 text encoder is the LTX 2.4 path; Gemma 3 is LTX 2.0.
+        is_gemma4 = "gemma4" in getattr(clip.tokenizer, "clip_name", "")
+
+        if is_gemma4:
+            if image is not None:
+                system = LTX24_I2V_SYSTEM_PROMPT.strip()
+                user_text = f"User Raw Input Prompt: {prompt}."
+            else:
+                system = LTX24_T2V_SYSTEM_PROMPT.strip()
+                user_text = f"user prompt: {prompt}"
+            think_prefix = "<|think|>\n" if thinking else ""
+            model_open = "" if thinking else "<|channel>final\n"
+            media = "<|image><|image|><image|>\n\n" if image is not None else ""
+            formatted_prompt = (
+                f"<|turn>system\n{think_prefix}{system}<turn|>\n"
+                f"<|turn>user\n{media}{user_text}<turn|>\n"
+                f"<|turn>model\n{model_open}"
+            )
         else:
-            formatted_prompt = f"<start_of_turn>system\n{LTX2_I2V_SYSTEM_PROMPT.strip()}<end_of_turn>\n<start_of_turn>user\n\n<image_soft_token>\n\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n<start_of_turn>model\n"
-        return super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
+            system = (LTX2_I2V_SYSTEM_PROMPT if image is not None else LTX2_T2V_SYSTEM_PROMPT).strip()
+            media = "\n<image_soft_token>\n" if image is not None else ""
+            formatted_prompt = (
+                f"<start_of_turn>system\n{system}<end_of_turn>\n"
+                f"<start_of_turn>user\n{media}\nUser Raw Input Prompt: {prompt}.<end_of_turn>\n"
+                f"<start_of_turn>model\n"
+            )
+
+        out = super().execute(clip, formatted_prompt, max_length, sampling_mode, image=image, thinking=thinking, use_default_template=use_default_template, video=video, audio=audio)
+
+        # Drop reasoning, including a block left unclosed by max_length. Both system prompts ask
+        # for the original prompt back when there is nothing to give; empty conditions on nothing.
+        text = re.sub(r"<think>.*?(?:</think>|$)", "", out.args[0], flags=re.DOTALL).strip()
+        return io.NodeOutput(text or prompt)
 
 
 class TextgenExtension(ComfyExtension):

comfy_extras/nodes_upscale_model.py

@@ -72,7 +72,7 @@ class ImageUpscaleWithModel(io.ComfyNode):
 
         memory_required = (512 * 512 * 3) * image.element_size() * max(upscale_model.scale, 1.0) * 384.0 #The 384.0 is an estimate of how much some of these models take, TODO: make it more accurate
         memory_required += image.nelement() * image.element_size()
-        model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required)
+        model_management.load_models_gpu([upscale_model.patcher], memory_required=memory_required, force_full_load=True)
 
         in_img = image.movedim(-1,-3).to(device)
 

comfy_extras/nodes_wan.py

@@ -6,6 +6,8 @@ import comfy.model_management
 import comfy.utils
 import comfy.latent_formats
 import comfy.clip_vision
+import comfy.ldm.wan.model_animate2
+import comfy.patcher_extension
 import json
 import numpy as np
 from typing import Tuple, TypedDict
@@ -1248,6 +1250,167 @@ class WanAnimateToVideo(io.ComfyNode):
         out_latent["samples"] = latent
         return io.NodeOutput(positive, negative, out_latent, trim_latent, max(0, ref_motion_latent_length * 4 - 3), video_frame_offset + length)
 
+class WanAnimate2ToVideo(io.ComfyNode):
+
+    CONTINUE_MOTION_FRAMES = 1
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="WanAnimate2ToVideo",
+            category="model/conditioning/wan/animate",
+            description="Animate a character in a reference image using a video, effectively replicating the facial animation, body motion and hands gesture from the video.",
+            inputs=[
+                io.Conditioning.Input("positive"),
+                io.Conditioning.Input("negative"),
+                io.Vae.Input("vae"),
+                io.Int.Input("width", default=832, min=16, max=nodes.MAX_RESOLUTION, step=16, tooltip="Output video width in pixels."),
+                io.Int.Input("height", default=480, min=16, max=nodes.MAX_RESOLUTION, step=16, tooltip="Output video height in pixels."),
+                io.Int.Input("length", default=81, min=1, max=nodes.MAX_RESOLUTION, step=4, tooltip="Number of frames to generate."),
+                io.Int.Input("batch_size", default=1, min=1, max=4096, tooltip="Number of videos to generate simultaneously."),
+                io.Image.Input("reference_image", optional=True, tooltip="The character to animate."),
+                io.Image.Input("pose_video", optional=True, tooltip="The video whose motion is transferred to the reference character."),
+                io.ClipVisionOutput.Input("clip_vision_output", optional=True, tooltip="CLIP vision of the reference image."),
+                io.Conditioning.Input("positive_pose", optional=True, tooltip="Prompt for the pose-video branch, describing the motion rather than the character. Defaults to positive. Used for both the cond and uncond passes."),
+                io.ClipVisionOutput.Input("clip_vision_output_pose", optional=True, tooltip="CLIP vision of the pose video's first frame. Defaults to clip_vision_output."),
+                io.Image.Input("continue_motion", optional=True, tooltip="Previous motion sequence to continue from for temporal consistency."),
+                io.Int.Input("video_frame_offset", default=0, min=0, max=nodes.MAX_RESOLUTION, step=1, tooltip="Frames to seek into the pose video. Connect to the video_frame_offset output of the previous node when extending."),
+                io.Float.Input("pose_strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Scales the pose video's influence on the motion. 1.0 is the trained behavior; below weakens adherence, above amplifies. 0.0 mutes it but does not fully remove it."),
+                io.Float.Input("pose_start_percent", default=0.0, min=0.0, max=1.0, step=0.01, tooltip="Sampling percent at which the pose influence starts. Outside the window the pose branch is skipped entirely, which also speeds those steps up."),
+                io.Float.Input("pose_end_percent", default=1.0, min=0.0, max=1.0, step=0.01, tooltip="Sampling percent at which the pose influence ends. Motion is mostly established early, so e.g. 0.7 can loosen fine detail while keeping the choreography."),
+                io.Float.Input("reference_image_strength", default=1.0, min=0.0, max=10.0, step=0.01, tooltip="Scales how strongly generated frames attend to the reference image's latent frame. Below 1.0 loosens identity/appearance adherence (e.g. to let the prompt restyle), above tightens it against drift."),
+            ],
+            outputs=[
+                io.Conditioning.Output(display_name="positive"),
+                io.Conditioning.Output(display_name="negative"),
+                io.Latent.Output(display_name="latent"),
+                io.Int.Output(display_name="trim_latent", tooltip="Number of latent frames that should be trimmed before decoding."),
+                io.Int.Output(display_name="trim_image", tooltip="Number of overlapping image frames when extending a video."),
+                io.Int.Output(display_name="video_frame_offset", tooltip="Frames to seek into the pose video."),
+            ],
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, positive, negative, vae, width, height, length, batch_size, video_frame_offset, reference_image=None, pose_video=None, clip_vision_output=None, positive_pose=None, clip_vision_output_pose=None, continue_motion=None, pose_strength=1.0, pose_start_percent=0.0, pose_end_percent=1.0, reference_image_strength=1.0) -> io.NodeOutput:
+        if pose_start_percent > pose_end_percent:
+            raise ValueError("pose_start_percent ({}) must not be greater than pose_end_percent ({}).".format(pose_start_percent, pose_end_percent))
+        latent_length = ((length - 1) // 4) + 1
+        latent_width = width // 8
+        latent_height = height // 8
+
+        if reference_image is None:
+            reference_image = torch.zeros((1, height, width, 3))
+
+        ref_image = comfy.utils.common_upscale(reference_image[:1].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
+        ref_latent = vae.encode(ref_image[:, :, :, :3])
+        trim_latent = ref_latent.shape[2]
+
+        ref_motion_latent_length = 0
+        if continue_motion is None:
+            image = torch.ones((length, height, width, 3)) * 0.5
+        else:
+            continue_motion = continue_motion[-cls.CONTINUE_MOTION_FRAMES:]
+            video_frame_offset = max(0, video_frame_offset - continue_motion.shape[0])
+            continue_motion = comfy.utils.common_upscale(continue_motion[-length:].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
+            # 0.5 is mid-grey, matching upstream's zeros in [-1, 1] pixel space
+            image = torch.ones((length, height, width, continue_motion.shape[-1]), device=continue_motion.device, dtype=continue_motion.dtype) * 0.5
+            image[:continue_motion.shape[0]] = continue_motion
+            ref_motion_latent_length += ((continue_motion.shape[0] - 1) // 4) + 1
+
+        concat_latent_image = torch.cat((ref_latent, vae.encode(image[:, :, :, :3])), dim=2)
+
+        # 1-channel, 0 == known; concat_cond inverts and repeats it to the 4 mask channels
+        mask = torch.ones((1, 1, latent_length + trim_latent, latent_height, latent_width), device=concat_latent_image.device, dtype=concat_latent_image.dtype)
+        mask[:, :, :trim_latent + ref_motion_latent_length] = 0.0
+
+        positive = node_helpers.conditioning_set_values(positive, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
+        negative = node_helpers.conditioning_set_values(negative, {"concat_latent_image": concat_latent_image, "concat_mask": mask})
+
+        if clip_vision_output is not None:
+            positive = node_helpers.conditioning_set_values(positive, {"clip_vision_output": clip_vision_output})
+            negative = node_helpers.conditioning_set_values(negative, {"clip_vision_output": clip_vision_output})
+
+        # not windowed with the pose values: the reference frame is part of the latent on every step
+        if reference_image_strength != 1.0:
+            positive = node_helpers.conditioning_set_values(positive, {"reference_strength": reference_image_strength})
+            negative = node_helpers.conditioning_set_values(negative, {"reference_strength": reference_image_strength})
+
+        # set on the negative too: upstream runs the pose branch once, outside the CFG loop, so it never sees the negative prompt
+        pose_values = {}
+        if pose_video is not None:
+            if pose_video.shape[0] <= video_frame_offset:
+                raise ValueError("pose_video has {} frames but video_frame_offset is {} -- nothing left to read.".format(pose_video.shape[0], video_frame_offset))
+            pose_video = pose_video[video_frame_offset:]
+            pose_video = comfy.utils.common_upscale(pose_video[:length].movedim(-1, 1), width, height, "area", "center").movedim(1, -1)
+            if pose_video.shape[0] < length:  # hold the last frame, as upstream pads its clips
+                pose_video = torch.cat((pose_video,) + (pose_video[-1:],) * (length - pose_video.shape[0]), dim=0)
+            pose_values["pose_video_latent"] = vae.encode(pose_video[:, :, :, :3])
+
+        pose_clip = clip_vision_output_pose if clip_vision_output_pose is not None else clip_vision_output
+        if pose_clip is not None:
+            pose_values["clip_vision_output_pose"] = pose_clip
+
+        pose_cond = positive_pose if positive_pose is not None else positive
+        if len(pose_cond) > 0:
+            pose_values["cross_attn_pose"] = pose_cond[0][0]
+
+        if pose_strength != 1.0:
+            pose_values["pose_strength"] = pose_strength
+
+        if pose_start_percent > 0.0 or pose_end_percent < 1.0:
+            # windowed via cond timestep ranges: the pose values ride a cond limited to the window, and complement conds without them cover the rest, where the model runs without the pose branch at all
+            def windowed(cond):
+                parts = node_helpers.conditioning_set_values(cond, {**pose_values, "start_percent": pose_start_percent, "end_percent": pose_end_percent})
+                if pose_start_percent > 0.0:
+                    parts = parts + node_helpers.conditioning_set_values(cond, {"start_percent": 0.0, "end_percent": pose_start_percent})
+                if pose_end_percent < 1.0:
+                    parts = parts + node_helpers.conditioning_set_values(cond, {"start_percent": pose_end_percent, "end_percent": 1.0})
+                return parts
+            positive = windowed(positive)
+            negative = windowed(negative)
+        else:
+            positive = node_helpers.conditioning_set_values(positive, pose_values)
+            negative = node_helpers.conditioning_set_values(negative, pose_values)
+
+        latent = torch.zeros([batch_size, 16, latent_length + trim_latent, latent_height, latent_width], device=comfy.model_management.intermediate_device())
+        out_latent = {}
+        out_latent["samples"] = latent
+        return io.NodeOutput(positive, negative, out_latent, trim_latent, max(0, ref_motion_latent_length * 4 - 3), video_frame_offset + length)
+
+
+class WanAnimate2Cache(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="WanAnimate2Cache",
+            category="model/conditioning/wan/animate",
+            description=(
+                "Caches the pose-video's per-block activations so it runs once instead of on every sampling step. Roughly halves generation time "
+                "Tradeoff is ~12.5 GB of system RAM at 480x832/81 frames in bf16 (scales with resolution and length). "
+                "With context windows each window is cached separately, so RAM scales with the window count; use the static_standard schedule, as uniform schedules shift the windows every step and nothing ever recurs to hit the cache."
+            ),
+            inputs=[
+                io.Model.Input("model"),
+                io.Combo.Input("device", options=["cpu", "gpu"], default="cpu",
+                               tooltip="Where to keep the cache. cpu (RAM) is the safe choice, the cache will not fit in VRAM alongside the model at typical sizes. gpu (VRAM) can be faster if it fits."),
+                io.Combo.Input("dtype", options=["default", "int8", "int4"], default="default",
+                               tooltip="Storage precision. default stores the activations in the model's compute dtype. int8 halves the cache, int4 quarters it, convrot is used to retain accuracy."),
+            ],
+            outputs=[io.Model.Output()],
+            is_experimental=True,
+        )
+
+    @classmethod
+    def execute(cls, model, device, dtype="default") -> io.NodeOutput:
+        store = comfy.model_management.get_torch_device() if device == "gpu" else torch.device("cpu")
+        cache = comfy.ldm.wan.model_animate2.PoseBranchCache(store_device=store, dtype=dtype)
+        m = model.clone()
+        m.model_options["transformer_options"]["animate2_cache"] = cache
+        m.add_callback(comfy.patcher_extension.CallbacksMP.ON_CLEANUP, lambda patcher: cache.free())
+        return io.NodeOutput(m)
+
+
 class Wan22ImageToVideoLatent(io.ComfyNode):
     @classmethod
     def define_schema(cls):
@@ -1475,6 +1638,8 @@ class WanExtension(ComfyExtension):
             WanSoundImageToVideoExtend,
             WanHuMoImageToVideo,
             WanAnimateToVideo,
+            WanAnimate2ToVideo,
+            WanAnimate2Cache,
             Wan22ImageToVideoLatent,
             WanInfiniteTalkToVideo,
         ]

comfyui_version.py

@@ -1,3 +1,3 @@
 # This file is automatically generated by the build process when version is
 # updated in pyproject.toml.
-__version__ = "0.30.0"
+__version__ = "0.33.0"

main.py

@@ -248,7 +248,17 @@ import hook_breaker_ac10a0
 import comfy.memory_management
 import comfy.model_patcher
 
-if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_management.is_nvidia() and not comfy.model_management.is_wsl()):
+
+def dynamic_vram_supported():
+    if comfy.model_management.is_nvidia():
+        return True
+    if comfy.model_management.is_amd():
+        if comfy.model_management.rocm_version >= (7, 14):
+            return True
+    return False
+
+
+if args.enable_dynamic_vram or (enables_dynamic_vram() and dynamic_vram_supported()):
     if (not args.enable_dynamic_vram) and (comfy.model_management.torch_version_numeric < (2, 8)):
         logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows")
     else:

nodes.py

@@ -290,6 +290,9 @@ class ConditioningZeroOut:
             conditioning_lyrics = d.get("conditioning_lyrics", None)
             if conditioning_lyrics is not None:
                 d["conditioning_lyrics"] = torch.zeros_like(conditioning_lyrics)
+            conditioning_scale = d.get("conditioning_scale", None)
+            if conditioning_scale is not None:
+                d["conditioning_scale"] = torch.zeros_like(conditioning_scale)
             n = [torch.zeros_like(t[0]), d]
             c.append(n)
         return (c, )
@@ -364,8 +367,12 @@ class VAEDecodeTiled:
             temporal_size = None
             temporal_overlap = None
 
+        latent = samples["samples"]
+        if latent.is_nested:
+            latent = latent.unbind()[0]
+
         compression = vae.spacial_compression_decode()
-        images = vae.decode_tiled(samples["samples"], tile_x=tile_size // compression, tile_y=tile_size // compression, overlap=overlap // compression, tile_t=temporal_size, overlap_t=temporal_overlap)
+        images = vae.decode_tiled(latent, tile_x=tile_size // compression, tile_y=tile_size // compression, overlap=overlap // compression, tile_t=temporal_size, overlap_t=temporal_overlap)
         if len(images.shape) == 5: #Combine batches
             images = images.reshape(-1, images.shape[-3], images.shape[-2], images.shape[-1])
         return (images, )
@@ -633,15 +640,18 @@ class DiffusersLoader:
     SEARCH_ALIASES = ["load diffusers model"]
 
     @classmethod
-    def INPUT_TYPES(cls):
+    def _model_paths(cls):
         paths = []
         for search_path in folder_paths.get_folder_paths("diffusers"):
             if os.path.exists(search_path):
                 for root, subdir, files in os.walk(search_path, followlinks=True):
                     if "model_index.json" in files:
                         paths.append(os.path.relpath(root, start=search_path))
+        return paths
 
-        return {"required": {"model_path": (paths,), }}
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {"required": {"model_path": (cls._model_paths(),), }}
     RETURN_TYPES = ("MODEL", "CLIP", "VAE")
     FUNCTION = "load_checkpoint"
     DEPRECATED = True
@@ -649,14 +659,20 @@ class DiffusersLoader:
     CATEGORY = "model/loaders"
 
     def load_checkpoint(self, model_path, output_vae=True, output_clip=True):
+        if model_path not in self._model_paths():
+            raise ValueError(f"Invalid diffusers model path: {model_path!r}")
+
+        resolved_model_path = None
         for search_path in folder_paths.get_folder_paths("diffusers"):
             if os.path.exists(search_path):
                 path = os.path.join(search_path, model_path)
-                if os.path.exists(path):
-                    model_path = path
+                if os.path.isfile(os.path.join(path, "model_index.json")):
+                    resolved_model_path = path
                     break
+        if resolved_model_path is None:
+            raise FileNotFoundError(f"Diffusers model {model_path!r} not found.")
 
-        return comfy.diffusers_load.load_diffusers(model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
+        return comfy.diffusers_load.load_diffusers(resolved_model_path, output_vae=output_vae, output_clip=output_clip, embedding_directory=folder_paths.get_folder_paths("embeddings"))
 
 
 class unCLIPCheckpointLoader:
@@ -1002,7 +1018,7 @@ class CLIPLoader:
 
     CATEGORY = "model/loaders"
 
-    DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
+    DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm\nminimax: MiniMax H3 Qwen3-VL or Music3 Qwen/RVQ"
 
     def load_clip(self, clip_name, type="stable_diffusion", device="default"):
         clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
@@ -1557,7 +1573,7 @@ def common_ksampler(model, seed, steps, cfg, sampler_name, scheduler, positive,
     latent_image = comfy.sample.fix_empty_latent_channels(model, latent_image, latent.get("downscale_ratio_spacial", None), latent.get("downscale_ratio_temporal", None))
 
     if disable_noise:
-        noise = torch.zeros(latent_image.size(), dtype=latent_image.dtype, layout=latent_image.layout, device="cpu")
+        noise = comfy.sample.prepare_empty_noise(latent_image)
     else:
         batch_inds = latent["batch_index"] if "batch_index" in latent else None
         noise = comfy.sample.prepare_noise(latent_image, seed, batch_inds)
@@ -2436,6 +2452,7 @@ async def init_builtin_extra_nodes():
         "nodes_mahiro.py",
         "nodes_lt_upsampler.py",
         "nodes_lt_audio.py",
+        "nodes_minimax_music.py",
         "nodes_minimax_h3.py",
         "nodes_lt.py",
         "nodes_hooks.py",
@@ -2492,6 +2509,7 @@ async def init_builtin_extra_nodes():
         "nodes_math.py",
         "nodes_number_convert.py",
         "nodes_painter.py",
+        "nodes_compositor.py",
         "nodes_curve.py",
         "nodes_bg_removal.py",
         "nodes_rtdetr.py",

openapi.yaml

@@ -35,6 +35,10 @@ components:
                     description: Timestamp when the asset was last accessed
                     format: date-time
                     type: string
+                loader_path:
+                    description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
+                    nullable: true
+                    type: string
                 metadata:
                     additionalProperties: true
                     description: System-managed metadata from download sources (HuggingFace, CivitAI, etc.) - read-only, not user-modifiable
@@ -165,6 +169,10 @@ components:
                     format: uuid
                     nullable: true
                     type: string
+                loader_path:
+                    description: The bare value a loader widget consumes for this asset. For models it is the path inside the category folder (e.g. "flux.safetensors" for "models/checkpoints/flux.safetensors"), which is what the model resolver matches. For input/output/temp it is the content hash, because those assets are fetched by hash rather than staged by name — that is the value LoadImage-style widgets must carry. Clients add the "[output]"/"[temp]" annotation from the asset's own type, so it is never included here. Null when no such value can be derived.
+                    nullable: true
+                    type: string
                 mime_type:
                     description: Updated MIME type of the asset
                     type: string
@@ -188,6 +196,31 @@ components:
                 - id
                 - updated_at
             type: object
+        ChurnkeyAuthResponse:
+            description: |
+                Credentials the Churnkey embed requires to launch the cancel flow.
+                `auth_hash` is hex-encoded HMAC-SHA256 of `customer_id` signed with the
+                server's CHURNKEY_HMAC_SECRET; it is bound to that single customer ID
+                and must not be reused for other customers.
+            properties:
+                auth_hash:
+                    description: Hex-encoded HMAC-SHA256(customer_id, CHURNKEY_HMAC_SECRET)
+                    type: string
+                customer_id:
+                    description: Stripe customer ID for the workspace
+                    type: string
+                mode:
+                    description: Churnkey environment matching the configured app
+                    enum:
+                        - live
+                        - test
+                        - sandbox
+                    type: string
+            required:
+                - customer_id
+                - auth_hash
+                - mode
+            type: object
         CreateWorkflowRequest:
             description: Request body for creating a new saved workflow.
             properties:
@@ -511,6 +544,25 @@ components:
             required:
                 - history
             type: object
+        JobAssetsResponse:
+            description: Paginated list of the assets produced by a single job.
+            properties:
+                assets:
+                    description: The job's output assets for the requested page (empty when the job produced none)
+                    items:
+                        $ref: '#/components/schemas/JobOutputAsset'
+                    type: array
+                job_id:
+                    description: ID of the job these assets belong to
+                    format: uuid
+                    type: string
+                pagination:
+                    $ref: '#/components/schemas/PaginationInfo'
+            required:
+                - job_id
+                - assets
+                - pagination
+            type: object
         JobCancelResponse:
             description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops.
             properties:
@@ -565,6 +617,9 @@ components:
                     additionalProperties: true
                     description: Primary preview output (only for terminal states)
                     type: object
+                previewable_outputs_count:
+                    description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
+                    type: integer
                 status:
                     description: User-friendly job status
                     enum:
@@ -597,6 +652,13 @@ components:
                 workflow_id:
                     description: UUID identifying the workflow graph definition
                     type: string
+                workflow_version_id:
+                    description: |
+                        UUID of the cloud workflow version this job is pinned to, if the
+                        submission carried one (see PromptRequest's workflow_version_id).
+                        Absent for jobs submitted without that association, including
+                        every job submitted through the public API v2 today.
+                    type: string
                 workspace_id:
                     description: |
                         ID of the workspace that owns this job. A successful (200)
@@ -645,6 +707,9 @@ components:
                     additionalProperties: true
                     description: Primary preview output (only present for terminal states)
                     type: object
+                previewable_outputs_count:
+                    description: Count of outputs classified as previewable media types (images, video, audio, 3D, text) — a subset of outputs_count (omitted for non-terminal states)
+                    type: integer
                 status:
                     description: User-friendly job status
                     enum:
@@ -662,6 +727,56 @@ components:
                 - status
                 - create_time
             type: object
+        JobOutputAsset:
+            description: |
+                An asset produced by a job, enriched with the per-output node context
+                (`node_id`, `output_key`, `output_index`) correlated from the job's
+                execution outputs by content hash. The node-context fields are null
+                when the asset cannot be matched to an output entry.
+            properties:
+                created_at:
+                    description: Timestamp when the asset was created
+                    format: date-time
+                    type: string
+                hash:
+                    description: Blake3 hash of the asset content.
+                    pattern: ^blake3:[a-f0-9]{64}$
+                    type: string
+                id:
+                    description: Unique identifier for the asset
+                    format: uuid
+                    type: string
+                mime_type:
+                    description: MIME type of the asset
+                    type: string
+                name:
+                    description: Name of the asset file
+                    type: string
+                node_id:
+                    description: ID of the workflow node that produced this asset, if known
+                    nullable: true
+                    type: string
+                output_index:
+                    description: Zero-based index of this asset within the node's output slot, if known
+                    nullable: true
+                    type: integer
+                output_key:
+                    description: Output slot key under the producing node (e.g. "images"), if known
+                    nullable: true
+                    type: string
+                preview_url:
+                    description: Relative URL for asset preview/thumbnail
+                    format: uri-reference
+                    type: string
+                size:
+                    description: Size of the asset in bytes
+                    format: int64
+                    type: integer
+            required:
+                - id
+                - name
+                - created_at
+            type: object
         JobStatusResponse:
             description: Job status information
             properties:
@@ -1521,7 +1636,12 @@ paths:
                 Supports filtering by tags, name, metadata, and sorting options.
             operationId: listAssets
             parameters:
-                - description: Filter assets that have ALL of these tags
+                - deprecated: true
+                  description: |
+                    Deprecated alias for `tags_all`, kept permanently for existing
+                    callers. Filter assets that have ALL of these tags. Combining it
+                    with `tags_all`, or exceeding 100 tags (counted after removing
+                    empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: include_tags
@@ -1530,7 +1650,12 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - deprecated: true
+                  description: |
+                    Deprecated alias for `tags_none`, kept permanently for existing
+                    callers. Exclude assets that have ANY of these tags. Combining it
+                    with `tags_none`, or exceeding 100 tags (counted after removing
+                    empty values and duplicates), returns 400 `INVALID_TAG_FILTER`.
                   explode: false
                   in: query
                   name: exclude_tags
@@ -1539,6 +1664,51 @@ paths:
                         type: string
                     type: array
                   style: form
+                - description: |
+                    Filter assets that have ALL of these tags. Tag values are opaque
+                    byte-strings compared exactly and case-sensitively; unknown tags
+                    are not an error — they simply match nothing. Replaces the
+                    deprecated `include_tags`. Sending both spellings, listing the
+                    same tag here and in `tags_none`, or exceeding 100 tags per list
+                    (counted after removing empty values and duplicates) returns 400
+                    `INVALID_TAG_FILTER`.
+                  explode: false
+                  in: query
+                  name: tags_all
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: |
+                    Filter assets that have AT LEAST ONE of these tags. Combines with
+                    `tags_all`/`tags_none` by intersection (`tags_none` always wins;
+                    overlap with `tags_none` is allowed and leaves a dead term).
+                    Supplying a positive tag filter (`tags_any`, `tags_all`, or
+                    `include_tags`) replaces the default category filter that is
+                    otherwise applied. Lists over 100 tags (counted after removing
+                    empty values and duplicates) return 400 `INVALID_TAG_FILTER`.
+                  explode: false
+                  in: query
+                  name: tags_any
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: |
+                    Exclude assets that have ANY of these tags. Replaces the
+                    deprecated `exclude_tags`. Sending both spellings, or exceeding
+                    100 tags per list (counted after removing empty values and
+                    duplicates), returns 400 `INVALID_TAG_FILTER`.
+                  explode: false
+                  in: query
+                  name: tags_none
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
                 - description: Filter assets where name contains this substring (case-insensitive)
                   in: query
                   name: name_contains
@@ -2312,7 +2482,12 @@ paths:
                 Only returns tags with non-zero counts (tags that exist on matching assets).
             operationId: getAssetTagHistogram
             parameters:
-                - description: Filter assets that have ALL of these tags
+                - deprecated: true
+                  description: |
+                    Deprecated alias for `tags_all`, kept permanently for existing
+                    callers. Filter assets that have ALL of these tags. The same
+                    combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: include_tags
@@ -2321,7 +2496,12 @@ paths:
                         type: string
                     type: array
                   style: form
-                - description: Exclude assets that have ANY of these tags
+                - deprecated: true
+                  description: |
+                    Deprecated alias for `tags_none`, kept permanently for existing
+                    callers. Exclude assets that have ANY of these tags. The same
+                    combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
                   explode: false
                   in: query
                   name: exclude_tags
@@ -2330,6 +2510,43 @@ paths:
                         type: string
                     type: array
                   style: form
+                - description: |
+                    Filter assets that have ALL of these tags. Replaces the deprecated
+                    `include_tags`. The same combination and list-size rules as on
+                    `/api/assets` apply (400 `INVALID_TAG_FILTER`).
+                  explode: false
+                  in: query
+                  name: tags_all
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: |
+                    Filter assets that have AT LEAST ONE of these tags. Combines with
+                    `tags_all`/`tags_none` by intersection (`tags_none` always wins).
+                    The same combination and list-size rules as on `/api/assets` apply
+                    (400 `INVALID_TAG_FILTER`).
+                  explode: false
+                  in: query
+                  name: tags_any
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
+                - description: |
+                    Exclude assets that have ANY of these tags. Replaces the deprecated
+                    `exclude_tags`. The same combination and list-size rules as on
+                    `/api/assets` apply (400 `INVALID_TAG_FILTER`).
+                  explode: false
+                  in: query
+                  name: tags_none
+                  schema:
+                    items:
+                        type: string
+                    type: array
+                  style: form
                 - description: Filter assets where name contains this substring (case-insensitive)
                   in: query
                   name: name_contains
@@ -2382,6 +2599,49 @@ paths:
             summary: Get tag histogram for filtered assets
             tags:
                 - file
+    /api/billing/churnkey/auth:
+        get:
+            description: |
+                Returns the Stripe customer identifier and a server-signed
+                HMAC-SHA256 of the customer ID, used to launch the Churnkey-hosted
+                cancellation flow embed.
+            operationId: getChurnkeyAuth
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ChurnkeyAuthResponse'
+                    description: Success
+                "401":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unauthorized
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Workspace has no Stripe customer (never subscribed)
+                "500":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Internal server error
+                "503":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Churnkey is not configured on the server
+            security:
+                - BearerAuth: []
+            summary: Get Churnkey HMAC auth credentials
+            tags:
+                - billing
     /api/embeddings:
         get:
             description: Returns the list of text-encoder embeddings available on disk.
@@ -2402,9 +2662,10 @@ paths:
                 Returns a list of model folders available in the system.
                 This is an experimental endpoint that replaces the legacy /models endpoint.
                 Each folder's name is the identifier to pass to /api/experiment/models/{folder}.
-                Once the model_type migration is active the names are model_type folder_names
-                (e.g. `ultralytics_bbox`); a folder with no folder_name mapping is returned by
-                its directory path.
+                The folder vocabulary is resolved per request from the caller's identity: where the
+                model_type migration is active for that caller the names are model_type folder_names
+                (e.g. `ultralytics_bbox`), and a folder with no folder_name mapping is returned by its
+                directory path. An authenticated response can therefore differ from an anonymous one.
             operationId: getModelFolders
             responses:
                 "200":
@@ -2421,7 +2682,10 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/ErrorResponse'
                     description: Internal server error
-            security: []
+            security:
+                - ApiKeyAuth: []
+                - BearerAuth: []
+                - {}
             summary: Get available model folders
             tags:
                 - file
@@ -2430,6 +2694,10 @@ paths:
             description: |
                 Returns a list of models available in the specified folder.
                 This is an experimental endpoint that provides enhanced model information.
+                Accepted folder identifiers are those returned by /api/experiment/models for the same
+                caller. That vocabulary is request-scoped, so list folders and fetch a folder's models
+                with the same credentials — a name obtained anonymously may not resolve when
+                authenticated, and vice versa.
             operationId: getModelsInFolder
             parameters:
                 - description: The folder name to list models from
@@ -2460,7 +2728,10 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/ErrorResponse'
                     description: Internal server error
-            security: []
+            security:
+                - ApiKeyAuth: []
+                - BearerAuth: []
+                - {}
             summary: Get models in a specific folder
             tags:
                 - file
@@ -3097,6 +3368,74 @@ paths:
             summary: Get full job details
             tags:
                 - workflow
+    /api/jobs/{job_id}/assets:
+        get:
+            description: |
+                Retrieve a paginated list of the assets produced by a specific job,
+                enriched with the per-output node context (`node_id`, `output_key`,
+                `output_index`) correlated from the job's execution outputs by content
+                hash. Unlike `GET /api/assets?job_ids={id}`, this endpoint is scoped to a
+                single job and carries node-level placement, making it suited to job
+                output views rather than the general asset browser. Returns an empty
+                `assets` array for jobs that produced no assets.
+            operationId: getJobAssets
+            parameters:
+                - description: Job identifier (UUID)
+                  in: path
+                  name: job_id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+                - description: Maximum number of assets to return (1-500)
+                  in: query
+                  name: limit
+                  schema:
+                    default: 20
+                    maximum: 500
+                    minimum: 1
+                    type: integer
+                - description: Number of assets to skip for pagination
+                  in: query
+                  name: offset
+                  schema:
+                    default: 0
+                    minimum: 0
+                    type: integer
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/JobAssetsResponse'
+                    description: Success - Job assets returned
+                "400":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Invalid request parameters
+                "401":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unauthorized - Authentication required
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Job not found or does not belong to the user
+                "500":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Internal server error
+            summary: List a job's output assets
+            tags:
+                - workflow
     /api/jobs/{job_id}/cancel:
         post:
             description: |
@@ -3302,6 +3641,12 @@ paths:
                             schema:
                                 $ref: '#/components/schemas/PromptErrorResponse'
                     description: Payment required - Insufficient credits
+                "403":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/PromptErrorResponse'
+                    description: Workspace governance policy blocks one or more partner providers (error.type PARTNER_NODE_DISABLED; error.class_types lists the offending nodes, error.providers the disabled providers)
                 "413":
                     content:
                         application/json:
@@ -3313,7 +3658,7 @@ paths:
                         application/json:
                             schema:
                                 $ref: '#/components/schemas/PromptErrorResponse'
-                    description: Payment required - User has not paid
+                    description: 'Retryable backpressure. Two distinct causes, disambiguated by the body''s `error.type`, NOT by parsing `error.message`: `PAYMENT_REQUIRED` / `FREE_TIER_UNAVAILABLE` / `FREE_TIER_EXHAUSTED` / `PARTNER_NODE_PAYMENT_REQUIRED` (a billing gate - retrying without paying never succeeds), or `QUEUE_LIMIT` (this workspace''s bounded job queue is full - retrying after some queued jobs complete will succeed).'
                 "500":
                     content:
                         application/json:
@@ -5152,6 +5497,8 @@ tags:
       name: user
     - description: Background task management
       name: task
+    - description: Workspace billing and subscription management
+      name: billing
     - description: Workflow storage and version management
       name: workflows
     - description: Job queue state and control

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "ComfyUI"
-version = "0.30.0"
+version = "0.33.0"
 readme = "README.md"
 license = { file = "LICENSE" }
 requires-python = ">=3.10"