upstream-archive/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-09-21 13:38:08 -05:00
Code Issues Packages Projects Releases Wiki Activity

merge: origin/master into v2-nodes-api

Browse Source
This commit is contained in:
benjcooley
2026-09-02 11:22:46 -07:00
parent eef9033de8 a87667f72f
commit 21f3a1cad1
269 changed files with 54466 additions and 3518 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.coderabbit.yaml
+4
View File
@@ -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
.github/workflows/ci-cursor-review.yml
+2 -2
View File
@@ -23,9 +23,9 @@ jobs:
# SHA-pinned per zizmor `unpinned-uses: hash-pin`. Bump this SHA to pick up
# upstream changes; keep `workflows_ref` matching so prompts/scripts load
# from the same commit as the workflow definition.
uses: Comfy-Org/github-workflows/.github/workflows/cursor-review.yml@047ca48febe3a6647608ed2e0c4331b491cb9d6a # github-workflows#9
uses: Comfy-Org/github-workflows/.github/workflows/cursor-review.yml@964d5aad37cbfb57c5b23961d42c2fd85868bf1d # github-workflows main (964d5aa)
with:
workflows_ref: 047ca48febe3a6647608ed2e0c4331b491cb9d6a
workflows_ref: 964d5aad37cbfb57c5b23961d42c2fd85868bf1d
diff_excludes: >-
:!**/.claude/**
:!**/dist/**
.github/workflows/cla.yml
+6 -3
View File
@@ -35,9 +35,12 @@ jobs:
# For each commit emit the GitHub login when the author/committer email resolves to a GitHub account
# otherwise fall back to the raw git name.
run: |
others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
--jq '.[] | (.author.login // .commit.author.name // empty), (.committer.login // .commit.committer.name // empty)' \
| sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
if ! commit_authors=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \
--jq '.[] | (.author.login // .commit.author.name // empty), (.committer.login // .commit.committer.name // empty)'); then
echo "Failed to fetch pull request commits" >&2
exit 1
fi
others=$(printf '%s\n' "$commit_authors" | sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -)
if [ -n "$others" ]; then
echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT"
else
.github/workflows/notify-on-merge.yml
+30
View File
@@ -0,0 +1,30 @@
name: Notify on Merge
on:
push:
branches:
- master
jobs:
notify:
runs-on: ubuntu-latest
if: github.repository == 'Comfy-Org/ComfyUI'
steps:
- name: Notify downstream
env:
DISPATCH_TOKEN: ${{ secrets.SYNC_DISPATCH_TOKEN }}
TARGET_REPO: ${{ secrets.SYNC_TARGET_REPO }}
COMMIT_SHA: ${{ github.sha }}
run: |
set -euo pipefail
if [ -z "${DISPATCH_TOKEN:-}" ] || [ -z "${TARGET_REPO:-}" ]; then
echo "::notice::SYNC_DISPATCH_TOKEN/SYNC_TARGET_REPO not set; skipping downstream notify."
exit 0
fi
PAYLOAD="$(jq -n --arg sha "$COMMIT_SHA" \
'{ event_type: "upstream-push", client_payload: { sha: $sha } }')"
curl -fsSL --connect-timeout 10 --max-time 60 -X POST \
-H "Accept: application/vnd.github+json" \
-H "Authorization: Bearer ${DISPATCH_TOKEN}" \
"https://api.github.com/repos/${TARGET_REPO}/dispatches" \
-d "$PAYLOAD"
.github/workflows/release-stable-all.yml
+2 -2
View File
@@ -20,7 +20,7 @@ jobs:
git_tag: ${{ inputs.git_tag }}
cache_tag: "cu130"
python_minor: "13"
python_patch: "12"
python_patch: "14"
rel_name: "nvidia"
rel_extra_name: ""
test_release: true
@@ -71,7 +71,7 @@ jobs:
git_tag: ${{ inputs.git_tag }}
cache_tag: "xpu"
python_minor: "13"
python_patch: "12"
python_patch: "14"
rel_name: "intel"
rel_extra_name: ""
test_release: true
.github/workflows/test-ci.yml
+12 -12
View File
@@ -2,18 +2,18 @@
# Results are reported as checkmarks on the commits, as well as onto https://ci.comfy.org/
name: Full Comfy CI Workflow Runs
on:
push:
branches:
- master
- release/**
paths-ignore:
- 'app/**'
- 'input/**'
- 'output/**'
- 'notebooks/**'
- 'script_examples/**'
- '.github/**'
- 'web/**'
# push:
# branches:
# - master
# - release/**
# paths-ignore:
# - 'app/**'
# - 'input/**'
# - 'output/**'
# - 'notebooks/**'
# - 'script_examples/**'
# - '.github/**'
# - 'web/**'
workflow_dispatch:
jobs:
AGENTS.md
+51 -2
View File
@@ -162,8 +162,26 @@
adding parallel code paths. Use `comfy.quant_ops`, `comfy.model_management`,
`comfy.memory_management`, `comfy.pinned_memory`, `comfy_aimdo`, and
`comfy-kitchen` helpers where they already solve the problem.
- Use optimized comfy-kitchen ops in places where they improve performance
without changing the expected dtype, device, memory, or interface behavior.
- Model implementations must use an existing optimized Comfy Kitchen or
ComfyUI operation whenever one supports the required math and tensor layout
without changing expected dtype, device, memory, or interface behavior. This
is the default implementation requirement, not an optional follow-up
optimization.
- Before implementing model math, inspect the operations already exposed by
Comfy Kitchen, `comfy.quant_ops`, and existing ComfyUI model helpers. Check
for optimized single, paired, fused, layout-specific, and quantized variants
before writing a local implementation or composing lower-level torch ops.
- Use the compatible optimized operation first and adapt the model's inputs to
its documented layout while preserving the model's exact math. If several
optimized variants apply, benchmark representative model shapes and select
the fastest valid path.
- Add or retain a local implementation only when no existing optimized
operation supports the required math, layout, dtype, device, autograd, or
patch contract. Keep differentiable or patch-compatible fallbacks when the
optimized inference operation does not provide those contracts.
- Use the existing ComfyUI cast, offload, and cleanup helpers for parameters
passed to optimized operations. Preserve model-specific epsilon, scaling,
layout, dtype, device, and output-shape behavior.
- Prefer ComfyUI's shared optimized kernels and backend dispatchers over
handwritten implementations of the same operation. Remove duplicate local
kernels and adapt inputs to the shared operation's documented layout while
@@ -279,14 +297,45 @@
- Avoid caches that persist across different executions as much as possible.
Persistent caches are acceptable only when they use a very minimal amount of
memory and have a clear ownership and invalidation story.
- When condition-dependent model work would otherwise repeat on every denoising
step and preprocessing it once materially improves performance, expose a
model preprocessing method and call it from `BaseModel.extra_conds`, following
patterns such as LTXAV and Anima. Pass the result through normal conditioning;
do not add model-owned caches, sampler-option caches, or cache-management
wrappers for this work.
- When optimizing, favor small measurable changes: fewer allocations, fewer
device transfers, less peak memory, better batching, or use of a faster
existing backend op.
## User Input Tolerance
- Prefer completing a workflow with the user's supplied values over rejecting
them because they fall outside recommended, UI-advertised, or quality-oriented
limits. If the downstream implementation can consume an input, pass it
through unchanged even when the result may be poor. For example, do not reject
or truncate additional reference images merely because a node advertises a
smaller recommended maximum.
- Do not add validation errors solely to prevent degraded, nonsensical, or
low-quality model output. A bad result is preferable to failing an otherwise
executable workflow.
- Resize, pad, clamp, normalize, or otherwise adapt user input only when passing
it through unchanged would make the existing model or underlying operation
fail. Make the smallest adjustment needed to keep execution running; do not
add a model-level validation failure merely to justify changing the input.
- This permissive policy does not override security boundaries such as path
containment, or integrity checks required to load model formats and
checkpoints safely.
## Nodes and User-Facing Behavior
- 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
CODEOWNERS
+1
View File
@@ -1,5 +1,6 @@
* @comfyanonymous @kosinkadink @guill @alexisrolland @rattus128 @kijai
/CODEOWNERS @comfyanonymous
/AGENTS.md @comfyanonymous
/.ci/ @comfyanonymous
/.github/ @comfyanonymous
README.md
+32 -65
View File
@@ -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,26 @@ 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).
- Support for saving and loading high bit depth images and videos: 16 bit PNG images, 32 bit EXR, 10 bit AVIF are supported and more.
- Support for saving and loading HDR videos and images in various formats.
Workflow examples can be found on the [Examples page](https://comfyanonymous.github.io/ComfyUI_examples/)
## Release Process
@@ -188,11 +147,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 +175,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 +196,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
+174 -45
View File
@@ -2,6 +2,7 @@ import asyncio
import functools
import json
import logging
import mimetypes
import os
import urllib.parse
import uuid
@@ -18,7 +19,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,
@@ -32,6 +33,7 @@ from app.assets.services import (
create_from_hash,
delete_asset_reference,
get_asset_detail,
get_preview_file_paths,
list_assets_page,
list_tags,
remove_tags,
@@ -40,7 +42,7 @@ from app.assets.services import (
upload_from_temp_path,
)
from app.assets.services.cursor import InvalidCursorError
from app.assets.services.path_utils import compute_display_name
from app.assets.services.path_utils import compute_asset_response_paths
from app.assets.services.tagging import list_tag_histogram
ROUTES = web.RouteTableDef()
@@ -117,6 +119,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"
@@ -126,44 +209,62 @@ def _validate_sort_field(requested: str | None) -> str:
return "created_at"
def _build_preview_url_from_view(tags: list[str], user_metadata: dict[str, Any] | None) -> str | None:
"""Build a /api/view preview URL from asset tags and user_metadata filename."""
if not user_metadata:
# What a client can render from the bytes themselves; anything else needs a nominated preview.
PREVIEWABLE_MIME_PREFIXES = ("image/", "video/", "audio/", "text/")
# models is deliberately absent: /api/view has no directory type for it.
VIEWABLE_NAMESPACES = frozenset({"input", "output", "temp"})
def _has_previewable_content(asset: schemas.AssetData | None, file_path: str | None) -> bool:
if asset is None:
return False
# Resolved from the path, not the caller-editable name, so a rename cannot change what previews.
raw = asset.mime_type or mimetypes.guess_type(file_path or "")[0] or ""
return raw.split(";", 1)[0].strip().lower().startswith(PREVIEWABLE_MIME_PREFIXES)
def _build_view_url(file_path: str | None) -> str | None:
# /api/view is a FileResponse: byte-range seeking, no user header, no access write.
if not file_path:
return None
filename = user_metadata.get("filename")
if not filename:
paths = compute_asset_response_paths(file_path)
if not paths:
return None
logical_path, relative_path = paths
namespace = logical_path.split("/", 1)[0]
if namespace not in VIEWABLE_NAMESPACES or not relative_path:
return None
if "input" in tags:
view_type = "input"
elif "output" in tags:
view_type = "output"
else:
return None
subfolder = ""
if "/" in filename:
subfolder, filename = filename.rsplit("/", 1)
encoded_filename = urllib.parse.quote(filename, safe="")
url = f"/api/view?type={view_type}&filename={encoded_filename}"
subfolder, _, filename = relative_path.rpartition("/")
url = f"/api/view?type={namespace}&filename={urllib.parse.quote(filename, safe='')}"
if subfolder:
url += f"&subfolder={urllib.parse.quote(subfolder, safe='')}"
return url
def _build_asset_response(result: schemas.AssetDetailResult | schemas.UploadResult) -> schemas_out.Asset:
"""Build an Asset response from a service result."""
def _resolve_preview_paths(
results: "list[schemas.AssetDetailResult] | list[schemas.AssetSummaryData]",
) -> dict[str, str]:
# A miss means no live preview - that is what keeps a soft-deleted one quiet.
preview_ids = {r.ref.preview_id for r in results if r.ref.preview_id}
return get_preview_file_paths(sorted(preview_ids))
def _build_asset_response(
result: schemas.AssetDetailResult | schemas.UploadResult,
preview_paths: dict[str, str],
) -> schemas_out.Asset:
if result.ref.preview_id:
preview_detail = get_asset_detail(result.ref.preview_id)
if preview_detail:
preview_url = _build_preview_url_from_view(preview_detail.tags, preview_detail.ref.user_metadata)
else:
preview_url = None
# A nominated preview is one whatever it holds, so no media check here.
preview_url = _build_view_url(preview_paths.get(result.ref.preview_id))
elif _has_previewable_content(result.asset, result.ref.file_path):
preview_url = _build_view_url(result.ref.file_path)
else:
preview_url = _build_preview_url_from_view(result.tags, result.ref.user_metadata)
preview_url = None
if result.ref.file_path:
display_name = compute_display_name(result.ref.file_path)
paths = compute_asset_response_paths(result.ref.file_path)
display_name = paths[1] if paths else None
# In-root loader path (model category dropped): what model loaders consume.
loader_path = result.ref.loader_path
else:
@@ -217,6 +318,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 +330,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,
@@ -237,7 +344,8 @@ async def list_assets_route(request: web.Request) -> web.Response:
except InvalidCursorError as e:
return _build_error_response(400, "INVALID_CURSOR", str(e))
summaries = [_build_asset_response(item) for item in result.items]
preview_paths = _resolve_preview_paths(result.items)
summaries = [_build_asset_response(item, preview_paths) for item in result.items]
# has_more semantics differ by mode:
# - cursor mode: a non-empty next_cursor means there are more results.
@@ -276,7 +384,7 @@ async def get_asset_route(request: web.Request) -> web.Response:
{"id": reference_id},
)
payload = _build_asset_response(result)
payload = _build_asset_response(result, _resolve_preview_paths([result]))
except ValueError as e:
return _build_error_response(
404, "ASSET_NOT_FOUND", str(e), {"id": reference_id}
@@ -315,15 +423,29 @@ async def download_asset_content(request: web.Request) -> web.Response:
404, "FILE_NOT_FOUND", "Underlying file not found on disk."
)
# User-controlled asset content must never render inline in the app origin
# User-controlled asset content must not render inline in the app origin
# (stored XSS via SVG/HTML/XML). Force dangerous types to download and
# override any requested inline disposition. Centralised through
# folder_paths.is_dangerous_content_type so this can't drift from /view and
# /userdata (the previous inline set here omitted image/svg+xml and missed
# the charset/casing/+xml-dialect bypasses).
# override any requested inline disposition; SVG loaded into an <img> is
# exempt, see renders_safely_as_image. Centralised through folder_paths so
# this can't drift from /view and /userdata (the previous inline set here
# omitted image/svg+xml and missed the charset/casing/+xml-dialect bypasses).
extra_headers = {}
sec_fetch_dest = request.headers.get("Sec-Fetch-Dest")
if folder_paths.is_dangerous_content_type(content_type):
content_type = "application/octet-stream"
disposition = "attachment"
# This response now depends on a request header, so it must not be
# reused across destinations by a browser or intermediary cache: an
# inline SVG primed by an <img> fetch and replayed to a document
# navigation of the same URL would re-enable the stored XSS.
extra_headers["Vary"] = "Sec-Fetch-Dest"
extra_headers["Cache-Control"] = "no-store"
if not folder_paths.renders_safely_as_image(content_type, sec_fetch_dest):
content_type = "application/octet-stream"
disposition = "attachment"
# mime_type is uploader-supplied and unvalidated, so it can carry
# parameters. aiohttp rejects a charset in the content_type argument with
# ValueError, which would turn a valid inline SVG into a 500.
content_type = content_type.split(";", 1)[0].strip() or "application/octet-stream"
safe_name = (filename or "").replace("\r", "").replace("\n", "")
encoded = urllib.parse.quote(safe_name)
@@ -356,6 +478,7 @@ async def download_asset_content(request: web.Request) -> web.Response:
"Content-Disposition": cd,
"Content-Length": str(file_size),
"X-Content-Type-Options": "nosniff",
**extra_headers,
},
)
@@ -392,7 +515,7 @@ async def create_asset_from_hash_route(request: web.Request) -> web.Response:
404, "ASSET_NOT_FOUND", f"Asset content {body.hash} does not exist"
)
asset = _build_asset_response(result)
asset = _build_asset_response(result, _resolve_preview_paths([result]))
payload_out = schemas_out.AssetCreated(
**asset.model_dump(),
created_new=result.created_new,
@@ -483,7 +606,7 @@ async def upload_asset(request: web.Request) -> web.Response:
logging.exception("upload_asset failed for owner_id=%s", owner_id)
return _build_error_response(500, "INTERNAL", "Unexpected server error.")
asset = _build_asset_response(result)
asset = _build_asset_response(result, _resolve_preview_paths([result]))
payload_out = schemas_out.AssetCreated(
**asset.model_dump(),
created_new=result.created_new,
@@ -513,7 +636,7 @@ async def update_asset_route(request: web.Request) -> web.Response:
owner_id=USER_MANAGER.get_request_user_id(request),
preview_id=body.preview_id,
)
payload = _build_asset_response(result)
payload = _build_asset_response(result, _resolve_preview_paths([result]))
except PermissionError as pe:
return _build_error_response(403, "FORBIDDEN", str(pe), {"id": reference_id})
except ValueError as ve:
@@ -700,10 +823,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
+20 -6
View File
@@ -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/__init__.py
+2
View File
@@ -28,6 +28,7 @@ from app.assets.database.queries.asset_reference import (
get_reference_by_id,
get_reference_with_owner_check,
get_reference_ids_by_ids,
get_reference_paths_by_ids,
get_references_by_paths_and_asset_ids,
get_references_for_prefixes,
get_unenriched_references,
@@ -101,6 +102,7 @@ __all__ = [
"get_reference_by_id",
"get_reference_with_owner_check",
"get_reference_ids_by_ids",
"get_reference_paths_by_ids",
"get_reference_tags",
"get_references_by_paths_and_asset_ids",
"get_references_for_prefixes",
app/assets/database/queries/asset_reference.py
+25 -2
View File
@@ -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)
@@ -1062,6 +1064,27 @@ def get_references_by_paths_and_asset_ids(
return winners
def get_reference_paths_by_ids(
session: Session,
reference_ids: list[str],
) -> dict[str, str]:
"""Map reference id -> file_path for live, file-backed references."""
if not reference_ids:
return {}
paths: dict[str, str] = {}
for chunk in iter_chunks(reference_ids, MAX_BIND_PARAMS):
rows = session.execute(
select(AssetReference.id, AssetReference.file_path).where(
AssetReference.id.in_(chunk),
AssetReference.file_path.is_not(None),
AssetReference.deleted_at.is_(None),
)
)
paths.update({rid: fp for rid, fp in rows})
return paths
def get_reference_ids_by_ids(
session: Session,
reference_ids: list[str],
app/assets/database/queries/common.py
+12 -1
View File
@@ -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
+3 -1
View File
@@ -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/scanner.py
+38 -9
View File
@@ -57,10 +57,11 @@ class _AssetAccumulator(TypedDict):
refs: list[_RefInfo]
# Temp is deliberately absent: it is wiped before every scan, so walking it finds nothing.
RootType = Literal["models", "input", "output"]
def get_prefixes_for_root(root: RootType) -> list[str]:
def get_scan_prefixes_for_root(root: RootType) -> list[str]:
if root == "models":
bases: list[str] = []
for _bucket, paths, _exts in get_comfy_models_folders():
@@ -73,10 +74,15 @@ def get_prefixes_for_root(root: RootType) -> list[str]:
return []
def get_all_known_prefixes() -> list[str]:
"""Get all known asset prefixes across all root types."""
all_roots: tuple[RootType, ...] = ("models", "input", "output")
return [p for root in all_roots for p in get_prefixes_for_root(root)]
def get_owned_prefixes() -> list[str]:
"""Every directory an asset may live in; references outside these are marked missing."""
scan_roots: tuple[RootType, ...] = ("models", "input", "output")
prefixes = [p for root in scan_roots for p in get_scan_prefixes_for_root(root)]
return prefixes + get_temp_prefixes()
def get_temp_prefixes() -> list[str]:
return [os.path.abspath(folder_paths.get_temp_directory())]
def collect_models_files() -> list[str]:
@@ -107,7 +113,21 @@ def sync_references_with_filesystem(
collect_existing_paths: bool = False,
update_missing_tags: bool = False,
) -> set[str] | None:
"""Reconcile asset references with filesystem for a root.
return sync_prefixes_with_filesystem(
session,
get_scan_prefixes_for_root(root),
collect_existing_paths=collect_existing_paths,
update_missing_tags=update_missing_tags,
)
def sync_prefixes_with_filesystem(
session,
prefixes: list[str],
collect_existing_paths: bool = False,
update_missing_tags: bool = False,
) -> set[str] | None:
"""Reconcile asset references with filesystem under the given prefixes.
- Toggle needs_verify per reference using mtime/size stat check
- For hashed assets with at least one stat-unchanged ref: delete stale missing refs
@@ -117,14 +137,13 @@ def sync_references_with_filesystem(
Args:
session: Database session
root: Root type to scan
prefixes: Absolute directory prefixes whose references to reconcile
collect_existing_paths: If True, return set of surviving file paths
update_missing_tags: If True, update 'missing' tags based on file status
Returns:
Set of surviving absolute paths if collect_existing_paths=True, else None
"""
prefixes = get_prefixes_for_root(root)
if not prefixes:
return set() if collect_existing_paths else None
@@ -251,6 +270,16 @@ def sync_root_safely(root: RootType) -> set[str]:
return set()
def sync_temp_references_safely() -> None:
"""Retire temp references whose file is gone; temp is never scanned, so nothing else stats them."""
try:
with create_session() as sess:
sync_prefixes_with_filesystem(sess, get_temp_prefixes())
sess.commit()
except Exception as e:
logging.exception("temp reference sync failed: %s", e)
def mark_missing_outside_prefixes_safely(prefixes: list[str]) -> int:
"""Mark references as missing when outside the given prefixes.
@@ -384,7 +413,7 @@ def get_unenriched_assets_for_roots(
"""
prefixes: list[str] = []
for root in roots:
prefixes.extend(get_prefixes_for_root(root))
prefixes.extend(get_scan_prefixes_for_root(root))
if not prefixes:
return []
app/assets/seeder.py
+7 -5
View File
@@ -15,12 +15,13 @@ from app.assets.scanner import (
build_asset_specs,
collect_paths_for_roots,
enrich_assets_batch,
get_all_known_prefixes,
get_prefixes_for_root,
get_owned_prefixes,
get_scan_prefixes_for_root,
get_unenriched_assets_for_roots,
insert_asset_specs,
mark_missing_outside_prefixes_safely,
sync_root_safely,
sync_temp_references_safely,
)
from app.database.db import dependencies_available
@@ -413,7 +414,7 @@ class _AssetSeeder:
)
return 0
all_prefixes = get_all_known_prefixes()
all_prefixes = get_owned_prefixes()
marked = mark_missing_outside_prefixes_safely(all_prefixes)
if marked > 0:
logging.info("Marked %d references as missing", marked)
@@ -523,7 +524,7 @@ class _AssetSeeder:
os.path.abspath(folder_paths.models_dir),
)
else:
prefixes = get_prefixes_for_root(root)
prefixes = get_scan_prefixes_for_root(root)
if prefixes:
logging.info("Asset scan [%s] directories: %s", root, prefixes)
@@ -548,10 +549,11 @@ class _AssetSeeder:
return
if self._prune_first:
all_prefixes = get_all_known_prefixes()
all_prefixes = get_owned_prefixes()
marked = mark_missing_outside_prefixes_safely(all_prefixes)
if marked > 0:
logging.info("Marked %d refs as missing before scan", marked)
sync_temp_references_safely()
if self._check_pause_and_cancel():
logging.info("Asset scan cancelled after pruning phase")
app/assets/services/__init__.py
+2
View File
@@ -4,6 +4,7 @@ from app.assets.services.asset_management import (
get_asset_by_hash,
get_asset_detail,
list_assets_page,
get_preview_file_paths,
resolve_asset_for_download,
set_asset_preview,
update_asset_metadata,
@@ -83,6 +84,7 @@ __all__ = [
"list_tags",
"cleanup_unreferenced_assets",
"remove_tags",
"get_preview_file_paths",
"resolve_asset_for_download",
"set_asset_preview",
"update_asset_metadata",
app/assets/services/asset_management.py
+12
View File
@@ -21,6 +21,7 @@ from app.assets.database.queries import (
reference_exists_for_asset_id,
delete_reference_by_id,
fetch_reference_and_asset,
get_reference_paths_by_ids,
soft_delete_reference_by_id,
fetch_reference_asset_and_tags,
get_asset_by_hash as queries_get_asset_by_hash,
@@ -279,6 +280,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 +320,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,
@@ -421,6 +425,14 @@ def resolve_hash_to_path(
)
def get_preview_file_paths(preview_ids: list[str]) -> dict[str, str]:
"""Map preview reference id -> file_path, in one query for the whole page."""
if not preview_ids:
return {}
with create_session() as session:
return get_reference_paths_by_ids(session, reference_ids=preview_ids)
def resolve_asset_for_download(
reference_id: str,
owner_id: str = "",
app/assets/services/tagging.py
+3
View File
@@ -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,
app/database/db.py
+53 -5
View File
@@ -4,7 +4,7 @@ import shutil
from app.logger import log_startup_warning
from utils.install_util import get_missing_requirements_message
from filelock import FileLock, Timeout
from comfy.cli_args import args
from comfy.cli_args import args, database_default_path
_DB_AVAILABLE = False
Session = None
@@ -57,19 +57,66 @@ def get_alembic_config():
config = Config(config_path)
config.set_main_option("script_location", scripts_path)
config.set_main_option("sqlalchemy.url", args.database_url)
config.set_main_option("sqlalchemy.url", get_database_url())
return config
def get_database_url():
if args.database_url is not None:
return args.database_url
import folder_paths
db_path = os.path.join(folder_paths.get_user_directory(), "comfyui.db")
return f"sqlite:///{db_path}"
def get_legacy_default_db_path():
return database_default_path
def get_db_path():
url = args.database_url
url = get_database_url()
if url.startswith("sqlite:///"):
return url.split("///")[1]
return url.split("///", 1)[1]
else:
raise ValueError(f"Unsupported database URL '{url}'.")
def copy_legacy_default_db(db_path):
if args.database_url is not None:
return
legacy_db_path = get_legacy_default_db_path()
if legacy_db_path is None:
return
if os.path.abspath(legacy_db_path) == os.path.abspath(db_path):
return
if os.path.exists(db_path) or not os.path.exists(legacy_db_path):
return
backup_path = legacy_db_path + ".bak"
if os.path.exists(backup_path):
return
os.replace(legacy_db_path, backup_path)
shutil.copy(backup_path, db_path)
logging.info(
f"Renamed legacy database '{legacy_db_path}' to '{backup_path}' and copied it to '{db_path}'"
)
def prepare_file_db_path(db_path):
db_dir = os.path.dirname(db_path)
if db_dir:
os.makedirs(db_dir, exist_ok=True)
copy_legacy_default_db(db_path)
_db_lock = None
def _acquire_file_lock(db_path):
@@ -97,7 +144,7 @@ def _is_memory_db(db_url):
def init_db():
db_url = args.database_url
db_url = get_database_url()
logging.debug(f"Database URL: {db_url}")
if _is_memory_db(db_url):
@@ -134,6 +181,7 @@ def _init_memory_db(db_url):
def _init_file_db(db_url):
"""Initialize a file-backed SQLite database using Alembic migrations."""
db_path = get_db_path()
prepare_file_db_path(db_path)
db_exists = os.path.exists(db_path)
config = get_alembic_config()
app/logger.py
+29 -2
View File
@@ -2,9 +2,12 @@ from collections import deque
from datetime import datetime
import io
import logging
import os
import sys
import threading
import comfy.internal_logging
ANSI_NAMED_COLORS = {
'black': '\033[30m',
'red': '\033[31m',
@@ -18,6 +21,7 @@ ANSI_NAMED_COLORS = {
ANSI_LEVEL_COLORS = {
'DEBUG': ANSI_NAMED_COLORS['cyan'],
'DETAIL': ANSI_NAMED_COLORS['blue'],
'INFO': ANSI_NAMED_COLORS['green'],
'WARNING': ANSI_NAMED_COLORS['yellow'],
'ERROR': ANSI_NAMED_COLORS['red'],
@@ -85,7 +89,12 @@ def on_flush(callback):
if stderr_interceptor is not None:
stderr_interceptor.on_flush(callback)
def setup_logger(log_level: str = 'INFO', capacity: int = 300, use_stdout: bool = False):
def get_log_level(level):
return comfy.internal_logging.DETAIL if level == "DETAIL" else logging.getLevelName(level)
def setup_logger(log_level: str = 'INFO', file_outputs=None, capacity: int = 300, use_stdout: bool = False):
global logs
if logs:
return
@@ -99,13 +108,18 @@ def setup_logger(log_level: str = 'INFO', capacity: int = 300, use_stdout: bool
stderr_interceptor = sys.stderr = LogInterceptor(sys.stderr)
# Setup default global logger
if file_outputs is None:
file_outputs = [('DETAIL', 'comfyui_detail.log')]
logger = logging.getLogger()
logger.setLevel(log_level)
console_level = get_log_level(log_level)
file_levels = [get_log_level(level) for level, _ in file_outputs]
logger.setLevel(min([console_level, *file_levels]))
formatter = ColoredFormatter("%(message)s")
stream_handler = logging.StreamHandler()
stream_handler.setFormatter(formatter)
stream_handler.setLevel(console_level)
if use_stdout:
# Only errors and critical to stderr
@@ -114,11 +128,24 @@ def setup_logger(log_level: str = 'INFO', capacity: int = 300, use_stdout: bool
# Lesser to stdout
stdout_handler = logging.StreamHandler(sys.stdout)
stdout_handler.setFormatter(formatter)
stdout_handler.setLevel(console_level)
stdout_handler.addFilter(lambda record: record.levelno < logging.ERROR)
logger.addHandler(stdout_handler)
logger.addHandler(stream_handler)
for output_level, output_path in file_outputs:
output_path = os.path.abspath(output_path)
try:
output_handler = logging.FileHandler(output_path, encoding="utf-8")
except OSError as e:
logging.warning("Could not open %s log %s: %s", output_level, output_path, e)
continue
output_handler.setLevel(get_log_level(output_level))
output_handler.setFormatter(logging.Formatter("[%(asctime)s] [%(levelname)s] %(message)s"))
logger.addHandler(output_handler)
logging.info("%s log: %s", output_level.title(), output_path)
STARTUP_WARNINGS = []
app/user_manager.py
+12 -3
View File
@@ -343,13 +343,22 @@ class UserManager():
# XSS). Content-Disposition: attachment is the load-bearing guard;
# the content-type override and nosniff are defence in depth.
content_type = mimetypes.guess_type(path)[0] or 'application/octet-stream'
if folder_paths.is_dangerous_content_type(content_type):
content_type = 'application/octet-stream'
user_root = self.get_request_user_filepath(request, None, create_dir=False)
is_user_css = path == os.path.abspath(os.path.join(user_root, "user.css"))
if is_user_css:
content_type = "text/css"
disposition = "inline"
else:
if folder_paths.is_dangerous_content_type(content_type):
content_type = 'application/octet-stream'
disposition = "attachment"
return web.FileResponse(path, headers={
"Content-Type": content_type,
"X-Content-Type-Options": "nosniff",
"Content-Disposition": "attachment",
"Content-Disposition": disposition,
})
@routes.post("/userdata/{file}")
blueprints/Image to Layers(Qwen-Image-Layered).json
+95 -5
View File
@@ -1,6 +1,6 @@
{
"revision": 0,
"last_node_id": 176,
"last_node_id": 177,
"last_link_id": 0,
"nodes": [
{
@@ -164,8 +164,8 @@
"version": 1,
"state": {
"lastGroupId": 8,
"lastNodeId": 176,
"lastLinkId": 380,
"lastNodeId": 177,
"lastLinkId": 381,
"lastRerouteId": 0
},
"revision": 0,
@@ -715,6 +715,88 @@
1
]
},
{
"id": 177,
"type": "LatentCut",
"pos": [
830,
-70
],
"size": [
270,
170
],
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"localized_name": "samples",
"name": "samples",
"type": "LATENT",
"link": 142
},
{
"localized_name": "dim",
"name": "dim",
"type": "COMBO",
"widget": {
"name": "dim"
},
"link": null
},
{
"localized_name": "index",
"name": "index",
"type": "INT",
"widget": {
"name": "index"
},
"link": null
},
{
"localized_name": "amount",
"name": "amount",
"type": "INT",
"widget": {
"name": "amount"
},
"link": null
}
],
"outputs": [
{
"localized_name": "LATENT",
"name": "LATENT",
"type": "LATENT",
"links": [
381
]
}
],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.5.1",
"ue_properties": {
"widget_ue_connectable": {},
"input_ue_unconnectable": {},
"version": "7.7"
},
"Node name for S&R": "LatentCut",
"enableTabs": false,
"tabWidth": 65,
"tabXOffset": 10,
"hasSecondTab": false,
"secondTabText": "Send Back",
"secondTabOffset": 80,
"secondTabWidth": 65
},
"widgets_values": [
"t",
1,
16384
]
},
{
"id": 76,
"type": "LatentCutToBatch",
@@ -734,7 +816,7 @@
"localized_name": "samples",
"name": "samples",
"type": "LATENT",
"link": 142
"link": 381
},
{
"localized_name": "dim",
@@ -1434,7 +1516,7 @@
"id": 142,
"origin_id": 3,
"origin_slot": 0,
"target_id": 76,
"target_id": 177,
"target_slot": 0,
"type": "LATENT"
},
@@ -1581,6 +1663,14 @@
"target_id": 39,
"target_slot": 0,
"type": "COMBO"
},
{
"id": 381,
"origin_id": 177,
"origin_slot": 0,
"target_id": 76,
"target_slot": 0,
"type": "LATENT"
}
],
"extra": {
comfy/background_removal/birefnet.py
+10 -13
View File
@@ -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
+32 -4
View File
@@ -33,6 +33,31 @@ class EnumAction(argparse.Action):
setattr(namespace, self.dest, value)
LOG_LEVELS = ('DEBUG', 'DETAIL', 'INFO', 'WARNING', 'ERROR', 'CRITICAL')
class VerboseAction(argparse.Action):
def __call__(self, parser, namespace, values, option_string=None):
if len(values) == 0:
output = ('DEBUG', None)
elif len(values) == 1 and values[0] in LOG_LEVELS:
output = (values[0], None)
elif len(values) == 2 and values[0] in LOG_LEVELS:
output = tuple(values)
else:
parser.error(f"{option_string} expects no values, a console LEVEL, or LEVEL FILE")
setattr(namespace, self.dest, [*getattr(namespace, self.dest, []), output])
def get_console_log_level(outputs):
console_levels = [level for level, path in outputs if path is None]
return min(console_levels, key=LOG_LEVELS.index, default='INFO')
def get_file_log_outputs(outputs):
return [(level, path) for level, path in outputs if path is not None]
parser = argparse.ArgumentParser()
parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0,::", help="Specify the IP address to listen on (default: 127.0.0.1). You can give a list of ip addresses by separating them with a comma like: 127.2.2.2,127.3.3.3 If --listen is provided without an argument, it defaults to 0.0.0.0,:: (listens on all ipv4 and ipv6)")
@@ -49,7 +74,7 @@ parser.add_argument("--temp-directory", type=str, default=None, help="Set the Co
parser.add_argument("--input-directory", type=str, default=None, help="Set the ComfyUI input directory. Overrides --base-directory.")
parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.")
parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.")
parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the ids of cuda devices this instance will use, as a comma-separated list (e.g. '0' or '0,1'). All other devices will not be visible.")
parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the ids of cuda devices this instance will use, as a comma-separated list (e.g. '0' or '0,1'), or 'all' to leave all currently visible devices available. All other devices will not be visible.")
parser.add_argument("--default-device", type=int, default=None, metavar="DEFAULT_DEVICE_ID", help="Set the id of the default device, all other devices will stay visible.")
cm_group = parser.add_mutually_exclusive_group()
cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).")
@@ -112,7 +137,7 @@ parser.add_argument("--preview-method", type=LatentPreviewMethod, default=Latent
parser.add_argument("--preview-size", type=int, default=512, help="Sets the maximum preview size for sampler nodes.")
cache_group = parser.add_mutually_exclusive_group()
cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 10%% of system RAM (min 2GB, max 10GB), inactive 100%% of system RAM (max 96GB).")
cache_group.add_argument("--cache-ram", nargs='*', type=float, default=[], metavar="GB", help="Use RAM pressure caching with the specified headroom thresholds. This is the default caching mode. The first value sets the active-cache threshold; the optional second value sets the inactive-cache/pin threshold. Defaults when no values are provided: active 10%% of system RAM (min 2GB, max 10GB), inactive 100%% of system RAM (max 128GB).")
cache_group.add_argument("--cache-classic", action="store_true", help="Use the old style (aggressive) caching.")
cache_group.add_argument("--cache-lru", type=int, default=0, help="Use LRU caching with a maximum of N node results cached. May use more RAM/VRAM.")
cache_group.add_argument("--cache-none", action="store_true", help="Reduced RAM/VRAM usage at the expense of executing every node for each run.")
@@ -124,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.")
@@ -147,12 +173,14 @@ vram_group.add_argument("--cpu", action="store_true", help="To use the CPU for e
parser.add_argument("--reserve-vram", type=float, default=None, help="Set the amount of vram in GB you want to reserve for use by your OS/other software. By default some amount is reserved depending on your OS.")
parser.add_argument("--vram-headroom", type=float, default=0, help="Set the amount of vram in GB for DynamicVRAM to maintain as extra headroom above default. ComfyUI will try and keep this much VRAM completely free and unused, even counting VRAM from other apps.")
parser.add_argument("--disable-nvml-pressure", action="store_true", help="Use CUDA instead of NVML for DynamicVRAM memory pressure.")
parser.add_argument("--async-offload", nargs='?', const=2, type=int, default=None, metavar="NUM_STREAMS", help="Use async weight offloading. An optional argument controls the amount of offload streams. Default is 2. Enabled by default on Nvidia.")
parser.add_argument("--disable-async-offload", action="store_true", help="Disable async weight offloading.")
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.")
@@ -187,7 +215,7 @@ parser.add_argument("--disable-api-nodes", action="store_true", help="Disable lo
parser.add_argument("--multi-user", action="store_true", help="Enables per-user storage.")
parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level')
parser.add_argument("--verbose", action=VerboseAction, nargs='*', default=[], metavar='LEVEL FILE', help='Set console logging with no values or LEVEL, or add a LEVEL FILE log output. May be repeated.')
parser.add_argument("--log-stdout", action="store_true", help="Send normal process output to stdout instead of stderr (default).")
@@ -240,7 +268,7 @@ parser.add_argument(
database_default_path = os.path.abspath(
os.path.join(os.path.dirname(__file__), "..", "user", "comfyui.db")
)
parser.add_argument("--database-url", type=str, default=f"sqlite:///{database_default_path}", help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite:///:memory:'.")
parser.add_argument("--database-url", type=str, default=None, help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite:///:memory:'. Defaults to 'comfyui.db' in the effective user directory.")
parser.add_argument("--enable-assets", action="store_true", help="Enable the assets system (API routes, database synchronization, and background scanning).")
parser.add_argument("--enable-asset-hashing", action="store_true", help="Compute blake3 content hashes when scanning assets. Hashing enables future asset-portability features (deduplication, cross-machine model resolution) but adds startup cost and per-output cost on large models directories. Off by default; enable to opt in.")
parser.add_argument("--feature-flag", type=str, action='append', default=[], metavar="KEY[=VALUE]", help="Set a server feature flag. Use KEY=VALUE to set an explicit value, or bare KEY to set it to true. Can be specified multiple times. Boolean values (true/false) and numbers are auto-converted. Examples: --feature-flag show_signin_button=true or --feature-flag show_signin_button")
comfy/clip_model.py
+6 -1
View File
@@ -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/clip_vision.py
+12
View File
@@ -10,6 +10,7 @@ import comfy.utils
import comfy.clip_model
import comfy.image_encoders.dino2
import comfy.image_encoders.dino3
from comfy.image_encoders.naf import NAF
class Output:
def __getitem__(self, key):
@@ -53,6 +54,7 @@ class ClipVisionModel():
self.model.eval()
self.patcher = comfy.model_patcher.CoreModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
self.naf = None
def load_sd(self, sd):
return self.model.load_state_dict(sd, strict=False, assign=self.patcher.is_dynamic())
@@ -141,6 +143,8 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json")
elif 'layer.0.mlp.gate_proj.weight' in sd and 'layer.31.norm1.weight' in sd: # Dinov3 ViT-H/16+ (SwiGLU gated MLP, 32 layers)
json_config = comfy.image_encoders.dino3.DINOV3_VITH_CONFIG
elif 'layer.23.attention.o_proj.bias' in sd: # dinov3 large (24 layers)
json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino3_large.json")
else:
return None
@@ -153,6 +157,14 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
for k in keys:
if k not in u:
sd.pop(k)
# NAF feature upsampler bundled into the DINOv3 file under the `naf.` prefix.
naf_keys = [k for k in sd if k.startswith("naf.")]
if naf_keys:
naf_sd = {k[len("naf."):]: sd.pop(k) for k in naf_keys}
naf = NAF(operations=comfy.ops.manual_cast).eval()
naf.load_state_dict(naf_sd)
naf.to(comfy.model_management.text_encoder_dtype(clip.load_device))
clip.naf = comfy.model_patcher.CoreModelPatcher(naf, load_device=clip.load_device, offload_device=comfy.model_management.text_encoder_offload_device())
return clip
def load(ckpt_path):
comfy/controlnet.py
+4 -7
View File
@@ -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/image_encoders/dino3.py
+17 -4
View File
@@ -156,10 +156,11 @@ class DINOv3ViTRopePositionEmbedding(nn.Module):
class DINOv3ViTEmbeddings(nn.Module):
def __init__(self, hidden_size, num_register_tokens, num_channels, patch_size, dtype, device, operations):
def __init__(self, hidden_size, num_register_tokens, num_channels, patch_size, dtype, device, operations, use_mask_token=True):
super().__init__()
self.cls_token = nn.Parameter(torch.empty(1, 1, hidden_size, device=device, dtype=dtype))
self.mask_token = nn.Parameter(torch.empty(1, 1, hidden_size, device=device, dtype=dtype))
self.mask_token = nn.Parameter(torch.empty(1, 1, hidden_size, device=device, dtype=dtype)) if use_mask_token else None
self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, hidden_size, device=device, dtype=dtype))
self.patch_embeddings = operations.Conv2d(
num_channels, hidden_size, kernel_size=patch_size, stride=patch_size, device=device, dtype=dtype
@@ -212,7 +213,7 @@ class DINOv3ViTLayer(nn.Module):
class DINOv3ViTModel(nn.Module):
def __init__(self, config, dtype, device, operations):
def __init__(self, config, dtype, device, operations, use_mask_token=True):
super().__init__()
num_hidden_layers = config["num_hidden_layers"]
hidden_size = config["hidden_size"]
@@ -228,7 +229,7 @@ class DINOv3ViTModel(nn.Module):
self.embeddings = DINOv3ViTEmbeddings(
hidden_size, num_register_tokens, num_channels=num_channels, patch_size=patch_size,
dtype=dtype, device=device, operations=operations
dtype=dtype, device=device, operations=operations, use_mask_token=use_mask_token
)
self.rope_embeddings = DINOv3ViTRopePositionEmbedding(
rope_theta, hidden_size, num_attention_heads, patch_size=patch_size, dtype=dtype, device=device
@@ -240,6 +241,10 @@ class DINOv3ViTModel(nn.Module):
for _ in range(num_hidden_layers)])
self.norm = operations.LayerNorm(hidden_size, eps=layer_norm_eps, dtype=dtype, device=device)
self.patch_size = patch_size
self.embed_dim = self.embed_dims = hidden_size
self.num_prefix_tokens = 1 + num_register_tokens # cls + register
def get_input_embeddings(self):
return self.embeddings.patch_embeddings
@@ -257,3 +262,11 @@ class DINOv3ViTModel(nn.Module):
sequence_output = norm(hidden_states)
pooled_output = sequence_output[:, 0, :]
return sequence_output, None, pooled_output, None
def forward_features(self, pixel_values, **kwargs):
sequence_output = self.forward(pixel_values, **kwargs)[0]
b = pixel_values.shape[0]
h = pixel_values.shape[-2] // self.patch_size
w = pixel_values.shape[-1] // self.patch_size
patches = sequence_output[:, self.num_prefix_tokens:, :]
return patches.reshape(b, h, w, self.embed_dim).permute(0, 3, 1, 2).contiguous()
comfy/image_encoders/dino3_large.json
+23
View File
@@ -0,0 +1,23 @@
{
"model_type": "dinov3",
"hidden_size": 1024,
"image_size": 224,
"initializer_range": 0.02,
"intermediate_size": 4096,
"key_bias": false,
"layer_norm_eps": 1e-05,
"mlp_bias": true,
"num_attention_heads": 16,
"num_channels": 3,
"num_hidden_layers": 24,
"num_register_tokens": 4,
"patch_size": 16,
"pos_embed_rescale": 2.0,
"proj_bias": true,
"query_bias": true,
"rope_theta": 100.0,
"use_gated_mlp": false,
"value_bias": true,
"image_mean": [0.485, 0.456, 0.406],
"image_std": [0.229, 0.224, 0.225]
}
comfy/image_encoders/naf.py
+282
View File
@@ -0,0 +1,282 @@
"""NAF (Neighborhood Attention Filtering) feature upsampler.
Vendored from valeoai/NAF (Apache-2.0):
https://github.com/valeoai/NAF — src/model/naf.py + src/layers/{convolutions,attentions,rope}.py
Used by Pixal3D's shape/texture conditioning to produce
the 2x-upsampled half of the 2048-channel proj feature map.
"""
import math
from typing import Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.ops
# Pure-torch neighborhood attention (replaces natten.na2d / na2d_qk + na2d_av).
def upsample_lr_slice(src_lr: torch.Tensor, lr_dh: int, lr_dw: int,
hr_h_range: Tuple[int, int], hr_w_range: Tuple[int, int]) -> torch.Tensor:
"""Slice a LR-layout tensor [B, h_lr, w_lr, n, C], permute to BCHW, and
nearest-exact upsample only the region covering [hr_h_range, hr_w_range].
Returns BCHW at hr_h_end-hr_h_start x hr_w_end-hr_w_start (no padding for
out-of-bounds regions)."""
B = src_lr.shape[0]
n = src_lr.shape[-2]
C = src_lr.shape[-1]
h_hr_start, h_hr_end = hr_h_range
w_hr_start, w_hr_end = hr_w_range
# LR positions covering [h_hr_start, h_hr_end). Nearest-exact maps HR p → p // D.
lr_h_start = h_hr_start // lr_dh
lr_h_end = (h_hr_end - 1) // lr_dh + 1
lr_w_start = w_hr_start // lr_dw
lr_w_end = (w_hr_end - 1) // lr_dw + 1
lr_slice = src_lr[:, lr_h_start:lr_h_end, lr_w_start:lr_w_end]
lh, lw = lr_slice.shape[1], lr_slice.shape[2]
lr_bcd = lr_slice.permute(0, 3, 4, 1, 2).reshape(B * n, C, lh, lw).contiguous()
up = F.interpolate(lr_bcd, scale_factor=(lr_dh, lr_dw), mode="nearest-exact")
offset_h = h_hr_start - lr_h_start * lr_dh
offset_w = w_hr_start - lr_w_start * lr_dw
return up[:, :, offset_h:offset_h + (h_hr_end - h_hr_start),
offset_w:offset_w + (w_hr_end - w_hr_start)]
def na2d_pure(
q: torch.Tensor, # [B, H, W, n_heads, d_qk] at HR.
k_lr: torch.Tensor, # [B, h_lr, w_lr, n_heads, d_qk] at LR
v_lr: torch.Tensor, # [B, h_lr, w_lr, n_heads, d_v] at LR
kernel_size: Tuple[int, int], # (Kh, Kw) attention window.
dilation: Tuple[int, int], # (Dh, Dw) stride within the unrolled K/V grid; also the LR→HR upsample factor.
scale: float, # 1 / sqrt(d_qk) scaling for the Q·K scores.
tile: int = 128, # Spatial tile size (output positions per tile)
v_chunk: int = 64, # Sub-divide d_v into chunks of this size when computing attn·V. None disables chunking.
output: torch.Tensor = None, # Pre-allocated [B, n_heads, d_v, H, W] buffer (may be on CPU).
) -> torch.Tensor: # [B, n_heads, d_v, H, W] (caller views as BCHW).
"""Neighborhood attention in pure torch via F.unfold + per-tile slicing.
K and V are passed at LR resolution and upsampled (nearest-exact) per-tile only
for the slice the unfold needs. Avoids the [B, n*d, H, W] HR allocations for K
(512 MB) and V (2 GB) at tex_1024 fp16. Spatial tiling bounds the per-tile
F.unfold blob; `v_chunk` further slices d_v so attn·V is computed in C-sized
chunks (attn is reused, computed once from Q/K).
"""
B, H, W, n, d_qk = q.shape
d_v = v_lr.shape[-1]
Kh, Kw = kernel_size
Dh, Dw = dilation
pad_h, pad_w = (Kh // 2) * Dh, (Kw // 2) * Dw
out = output if output is not None else torch.empty((B, n, d_v, H, W), device=q.device, dtype=q.dtype)
th = min(tile, H) if tile else H
tw = min(tile, W) if tile else W
chunk = v_chunk if (v_chunk and v_chunk < d_v) else d_v
for h0 in range(0, H, th):
for w0 in range(0, W, tw):
h1, w1 = min(h0 + th, H), min(w0 + tw, W)
t_h, t_w = h1 - h0, w1 - w0
# Padded HR region the unfold needs (kernel span = (K-1)*D + 1).
h_src_start = max(0, h0 - pad_h)
h_src_end = min(H, h1 + pad_h)
w_src_start = max(0, w0 - pad_w)
w_src_end = min(W, w1 + pad_w)
pad_top = max(0, pad_h - h0)
pad_bot = max(0, (h1 + pad_h) - H)
pad_lft = max(0, pad_w - w0)
pad_rgt = max(0, (w1 + pad_w) - W)
# Upsample only the tile region from k_lr / v_lr.
k_tile = upsample_lr_slice(k_lr, Dh, Dw,
(h_src_start, h_src_end),
(w_src_start, w_src_end))
v_tile = upsample_lr_slice(v_lr, Dh, Dw,
(h_src_start, h_src_end),
(w_src_start, w_src_end))
if pad_top or pad_bot or pad_lft or pad_rgt:
k_tile = F.pad(k_tile, [pad_lft, pad_rgt, pad_top, pad_bot])
v_tile = F.pad(v_tile, [pad_lft, pad_rgt, pad_top, pad_bot])
# Q·K → attention weights (small: KK=81 per output position).
KK = Kh * Kw
k_w = F.unfold(k_tile, kernel_size=(Kh, Kw), dilation=(Dh, Dw), padding=0)
k_w = k_w.view(B, n, d_qk, KK, t_h * t_w).permute(0, 1, 4, 3, 2) # [B, n, t, KK, d_qk]
# q is [B, H, W, n, d_qk]; per-tile slice + permute -> [B, n, t_h*t_w, 1, d_qk].
q_tile = q[:, h0:h1, w0:w1].permute(0, 3, 1, 2, 4).reshape(B, n, t_h * t_w, 1, d_qk)
scores = torch.matmul(q_tile, k_w.transpose(-1, -2)) * scale
attn = scores.softmax(dim=-1)
del k_w, scores, q_tile, k_tile
# attn · V, chunked over d_v.
for c0 in range(0, d_v, chunk):
c1 = min(c0 + chunk, d_v)
v_w = F.unfold(v_tile[:, c0:c1], kernel_size=(Kh, Kw),dilation=(Dh, Dw), padding=0) # [B*n, (c1-c0)*KK, t]
v_w = v_w.view(B, n, c1 - c0, KK, t_h * t_w).permute(0, 1, 4, 3, 2)
out_chunk = torch.matmul(attn, v_w).squeeze(-2) # [B, n, t, c1-c0]
out_chunk = out_chunk.view(B, n, t_h, t_w, c1 - c0).permute(0, 1, 4, 2, 3)
out[:, :, c0:c1, h0:h1, w0:w1] = out_chunk
del v_w, out_chunk
del attn, v_tile
return out # [B, n, d_v, H, W] — sole caller (CrossAttention) views it as BCHW directly.
class CrossAttention(nn.Module):
"""Window-restricted cross-attention. No learnable parameters; the model's
capacity lives entirely in the ImageEncoder convs."""
def __init__(self, dim: int, num_heads: int, kernel_size: Tuple[int, int] = (9, 9)):
super().__init__()
assert dim % num_heads == 0, "dim must be divisible by num_heads"
self.num_heads = num_heads
self.kernel_size = kernel_size
self.scale = (dim // num_heads) ** -0.5
@staticmethod
def _split_heads_lr(x: torch.Tensor, num_heads: int) -> torch.Tensor:
"""[B, n*d, h, w] -> [B, h, w, n, d] at the input resolution (no upsample)."""
B, C, H, W = x.shape
return x.view(B, num_heads, C // num_heads, H, W).permute(0, 3, 4, 1, 2).contiguous()
def forward(self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
output=None) -> torch.Tensor:
hq, wq = q.shape[-2:]
hk, wk = k.shape[-2:]
dilation = (hq // hk, wq // wk)
B, C, _, _ = q.shape
q = q.view(B, self.num_heads, C // self.num_heads, hq, wq).permute(0, 3, 4, 1, 2).contiguous()
k_lr = self._split_heads_lr(k, self.num_heads).to(q.dtype)
v_lr = self._split_heads_lr(v, self.num_heads).to(q.dtype)
out_buf = output.view(B, self.num_heads, v.shape[1] // self.num_heads, hq, wq) if output is not None else None
out = na2d_pure(q, k_lr, v_lr, self.kernel_size, dilation, self.scale, output=out_buf)
return out.view(B, -1, hq, wq)
# RoPE positional embedding
def rope_rotate_half(x: torch.Tensor) -> torch.Tensor:
x1, x2 = x.chunk(2, dim=-1)
return torch.cat([-x2, x1], dim=-1)
class RoPE(nn.Module):
def __init__(self, embed_dim: int, num_heads: int, base: float = 100.0):
super().__init__()
assert embed_dim % (4 * num_heads) == 0
self.num_heads = num_heads
self.D_head = embed_dim // num_heads
self.base = base
self.register_buffer("periods", torch.empty(self.D_head // 4), persistent=True) # loaded from the checkpoint
def _cos_sin(self, H: int, W: int, x: torch.Tensor):
"""cos/sin depend only on (H, W, dtype) and the checkpoint-fixed periods; recomputed per forward."""
periods = comfy.ops.cast_to_input(self.periods, x)
coords_h = torch.arange(0.5, H, device=x.device, dtype=torch.float32) / H
coords_w = torch.arange(0.5, W, device=x.device, dtype=torch.float32) / W
coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing="ij"), dim=-1) # [H, W, 2]
coords = coords.flatten(0, 1) * 2.0 - 1.0 # [HW, 2]
angles = 2 * math.pi * coords[:, :, None] / periods.to(coords.dtype)[None, None, :] # [HW, 2, D//4]
angles = angles.flatten(1, 2).tile(2) # [HW, D]
cos = torch.cos(angles).to(x.dtype)
sin = torch.sin(angles).to(x.dtype)
return cos, sin
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x: [B, n*D_head, H, W]
B, C, H, W = x.shape
n = self.num_heads
D = C // n
x = x.view(B, n, D, H, W).permute(0, 1, 3, 4, 2).reshape(B, n, H * W, D)
cos, sin = self._cos_sin(H, W, x)
x = (x * cos) + (rope_rotate_half(x) * sin)
x = x.view(B, n, H, W, D).permute(0, 1, 4, 2, 3).reshape(B, n * D, H, W)
return x
# Image encoder
class EncBlock(nn.Module):
def __init__(self, channels: int, kernel_size: int, operations, num_groups: int = 8):
super().__init__()
self.norm1 = operations.GroupNorm(num_groups=num_groups, num_channels=channels)
self.conv1 = operations.Conv2d(channels, channels, kernel_size=kernel_size,
padding=kernel_size // 2, padding_mode="reflect", bias=True)
self.norm2 = operations.GroupNorm(num_groups=num_groups, num_channels=channels)
self.conv2 = operations.Conv2d(channels, channels, kernel_size=kernel_size,
padding=kernel_size // 2, padding_mode="reflect", bias=True)
self.activation_fn = nn.SiLU()
def forward(self, x):
x = self.norm1(x)
x = self.activation_fn(x)
x = self.conv1(x)
x = self.norm2(x)
x = self.activation_fn(x)
x = self.conv2(x)
return x # no skip connection
def _encoder(in_dim: int, hidden_dim: int, operations, kernel_size: int = 1, ks_res: int = 1, num_layers: int = 2) -> nn.Sequential:
return nn.Sequential(
operations.Conv2d(in_dim, hidden_dim, kernel_size=kernel_size, padding=kernel_size // 2, padding_mode="reflect", bias=True),
*[EncBlock(hidden_dim, kernel_size=ks_res, operations=operations) for _ in range(num_layers)],
)
class ImageEncoder(nn.Module):
"""Two parallel conv stacks (1x1 + 3x3) producing dim/2 channels each, then concat,
spatial average-pool to target size, RoPE-embed positions."""
def __init__(self, operations, in_channels: int = 3, out_channels: int = 256,
heads_rope: int = 4, rope_base: float = 100.0, img_layers: int = 2):
super().__init__()
half = out_channels // 2
self.encoder = _encoder(in_channels, half, operations=operations, kernel_size=1, ks_res=1, num_layers=img_layers)
self.sem_encoder = _encoder(in_channels, half, operations=operations, kernel_size=3, ks_res=3, num_layers=img_layers)
self.rope = RoPE(embed_dim=out_channels, num_heads=heads_rope, base=rope_base)
def forward(self, x: torch.Tensor, output_size: Tuple[int, int]) -> torch.Tensor:
# Avoid running the conv stacks on >4× the target resolution.
out_h, out_w = output_size
if x.shape[-2] > 4 * out_h or x.shape[-1] > 4 * out_w:
x = F.interpolate(x, size=(min(x.shape[-2], 4 * out_h),
min(x.shape[-1], 4 * out_w)),
mode="bilinear", align_corners=False)
x = torch.cat([self.encoder(x), self.sem_encoder(x)], dim=1)
x = F.adaptive_avg_pool2d(x, output_size=output_size)
x = self.rope(x)
return x
class NAF(nn.Module):
"""NAF feature upsampler."""
def __init__(
self, operations,
dim: int = 256, # internal channel dimension of the ImageEncoder
heads_attn: int = 4, # attention heads in the windowed cross-attn
heads_rope: int = 4, # heads for RoPE position encoding (must divide dim)
kernel_size: int = 9, # square kernel for the neighborhood attention window
rope_base: float = 100.0, # base for RoPE frequency periods
img_layers: int = 2, # number of EncBlocks in each conv stack
):
super().__init__()
self.image_encoder = ImageEncoder(operations=operations, in_channels=3, out_channels=dim, heads_rope=heads_rope, rope_base=rope_base, img_layers=img_layers)
self.upsampler = CrossAttention(dim=dim, num_heads=heads_attn, kernel_size=(kernel_size, kernel_size))
def forward(
self,
image: torch.Tensor, # [B, 3, H_img, W_img] in [0, 1].
features: torch.Tensor, # [B, C, H_feat, W_feat] low-resolution features (any C).
output_size: Tuple[int, int], # (H_out, W_out) target spatial resolution for the upsampled features.
output=None,
) -> torch.Tensor: # [B, C, H_out, W_out] upsampled features.
"""Upsample low-res feature map to output_size, guided by the image."""
q = self.image_encoder(image, output_size=output_size)
k = F.adaptive_avg_pool2d(q, output_size=features.shape[-2:])
return self.upsampler(q, k, features, output=output)
comfy/internal_logging.py
+10
View File
@@ -0,0 +1,10 @@
import logging
DETAIL = 15
logging.addLevelName(DETAIL, "DETAIL")
def detail(message, *args, **kwargs):
kwargs.setdefault("stacklevel", 2)
logging.log(DETAIL, message, *args, **kwargs)
comfy/k_diffusion/sampling.py
+69 -1
View File
@@ -416,6 +416,74 @@ def linear_multistep_coeff(order, t, i, j):
return integrate.quad(fn, t[i], t[i + 1], epsrel=1e-4)[0]
def _sample_cfgpp_history(model, x, sigmas, extra_args=None, callback=None, disable=None, history_weight=0.5, zero_weight=None, zero_order=1, uncond_history_weight=0.0):
"""CFG++ Euler with variable-step AB2 history and optional sigma-zero extrapolation."""
extra_args = {} if extra_args is None else extra_args
model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
s_in = x.new_ones([x.shape[0]])
sigmas_cpu = sigmas.detach().cpu().numpy()
derivatives = []
denoised_history = []
old_uncond_d = None
uncond_denoised = None
def post_cfg_function(args):
nonlocal uncond_denoised
uncond_denoised = args["uncond_denoised"] if args["uncond"] is not None else args["cond_denoised"]
return args["denoised"]
model_options = extra_args.get("model_options", {}).copy()
extra_args["model_options"] = comfy.model_patcher.set_model_options_post_cfg_function(model_options, post_cfg_function)
for i in trange(len(sigmas) - 1, disable=disable):
denoised = model(x, sigmas[i] * s_in, **extra_args)
if callback is not None:
callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
alpha_s = sigmas[i] * lambda_fn(sigmas[i]).exp()
alpha_t = sigmas[i + 1] * lambda_fn(sigmas[i + 1]).exp() if sigmas[i + 1] != 0 else sigmas[i + 1].new_ones([])
current_uncond_d = to_d(x, sigmas[i], alpha_s * uncond_denoised)
uncond_d = current_uncond_d
dt = sigmas[i + 1] - sigmas[i]
if i > 0 and uncond_history_weight:
step_ratio = dt / (sigmas[i] - sigmas[i - 1])
uncond_d = uncond_d + uncond_history_weight * step_ratio * (current_uncond_d - old_uncond_d)
euler_step = alpha_t * denoised + sigmas[i + 1] * uncond_d - x
d = euler_step / dt
derivatives.append(d)
if len(derivatives) > 2:
derivatives.pop(0)
if len(derivatives) == 1:
step = euler_step
else:
coeffs = [linear_multistep_coeff(2, sigmas_cpu, i, j) for j in range(2)]
history_step = sum(coeff * derivative for coeff, derivative in zip(coeffs, reversed(derivatives)))
step = torch.lerp(euler_step, history_step, history_weight)
x = x + step
if sigmas[i + 1] == 0 and zero_weight is not None and denoised_history:
if zero_order == 2 and len(denoised_history) > 1:
sigma_0, sigma_1, sigma_2 = sigmas[i - 2], sigmas[i - 1], sigmas[i]
weight_0 = sigma_1 * sigma_2 / ((sigma_0 - sigma_1) * (sigma_0 - sigma_2))
weight_1 = sigma_0 * sigma_2 / ((sigma_1 - sigma_0) * (sigma_1 - sigma_2))
weight_2 = sigma_0 * sigma_1 / ((sigma_2 - sigma_0) * (sigma_2 - sigma_1))
zero_prediction = weight_0 * denoised_history[-2] + weight_1 * denoised_history[-1] + weight_2 * denoised
else:
denoised_slope = (denoised - denoised_history[-1]) / (sigmas[i] - sigmas[i - 1])
zero_prediction = denoised - sigmas[i] * denoised_slope
x = torch.lerp(x, zero_prediction, zero_weight)
denoised_history.append(denoised)
if len(denoised_history) > 2:
denoised_history.pop(0)
old_uncond_d = current_uncond_d
return x
def sample_cfgpp_ud10_ab(model, x, sigmas, extra_args=None, callback=None, disable=None):
return _sample_cfgpp_history(model, x, sigmas, extra_args=extra_args, callback=callback, disable=disable, history_weight=0.25, zero_weight=1.0, uncond_history_weight=0.1)
@torch.no_grad()
def sample_lms(model, x, sigmas, extra_args=None, callback=None, disable=None, order=4):
extra_args = {} if extra_args is None else extra_args
@@ -1445,7 +1513,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
+241 -2
View File
@@ -1,4 +1,5 @@
import torch
import comfy.nested_tensor
class LatentFormat:
scale_factor = 1.0
@@ -17,6 +18,9 @@ class LatentFormat:
def process_out(self, latent):
return latent / self.scale_factor
def fix_empty_latent(self, latent):
return latent
class SD15(LatentFormat):
def __init__(self, scale_factor=0.18215):
self.scale_factor = scale_factor
@@ -249,6 +253,53 @@ class TripoSplat(LatentFormat):
def process_out(self, latent):
return latent
class Trellis2(LatentFormat):
latent_channels = 32
class Trellis2SLAT(Trellis2):
# Sparse structured latent: per-token feats [N, 32]. process_out denormalizes
# the decoded feats (latent * std + mean); subclasses carry each space's stats.
latents_mean = None
latents_std = None
def process_in(self, latent):
mean = self.latents_mean.to(latent.device, latent.dtype)
std = self.latents_std.to(latent.device, latent.dtype)
return (latent - mean) / std
def process_out(self, latent):
mean = self.latents_mean.to(latent.device, latent.dtype)
std = self.latents_std.to(latent.device, latent.dtype)
return latent * std + mean
class Trellis2ShapeSLAT(Trellis2SLAT):
latents_mean = torch.tensor([
0.781296, 0.018091, -0.495192, -0.558457, 1.060530, 0.093252, 1.518149, -0.933218,
-0.732996, 2.604095, -0.118341, -2.143904, 0.495076, -2.179512, -2.130751, -0.996944,
0.261421, -2.217463, 1.260067, -0.150213, 3.790713, 1.481266, -1.046058, -1.523667,
-0.059621, 2.220780, 1.621212, 0.877230, 0.567247, -3.175944, -3.186688, 1.578665
])[None]
latents_std = torch.tensor([
5.972266, 4.706852, 5.445010, 5.209927, 5.320220, 4.547237, 5.020802, 5.444004,
5.226681, 5.683095, 4.831436, 5.286469, 5.652043, 5.367606, 5.525084, 4.730578,
4.805265, 5.124013, 5.530808, 5.619001, 5.103930, 5.417670, 5.269677, 5.547194,
5.634698, 5.235274, 6.110351, 5.511298, 6.237273, 4.879207, 5.347008, 5.405691
])[None]
class Trellis2TexSLAT(Trellis2SLAT):
latents_mean = torch.tensor([
3.501659, 2.212398, 2.226094, 0.251093, -0.026248, -0.687364, 0.439898, -0.928075,
0.029398, -0.339596, -0.869527, 1.038479, -0.972385, 0.126042, -1.129303, 0.455149,
-1.209521, 2.069067, 0.544735, 2.569128, -0.323407, 2.293000, -1.925608, -1.217717,
1.213905, 0.971588, -0.023631, 0.106750, 2.021786, 0.250524, -0.662387, -0.768862
])[None]
latents_std = torch.tensor([
2.665652, 2.743913, 2.765121, 2.595319, 3.037293, 2.291316, 2.144656, 2.911822,
2.969419, 2.501689, 2.154811, 3.163343, 2.621215, 2.381943, 3.186697, 3.021588,
2.295916, 3.234985, 3.233086, 2.260140, 2.874801, 2.810596, 3.292720, 2.674999,
2.680878, 2.372054, 2.451546, 2.353556, 2.995195, 2.379849, 2.786195, 2.775190
])[None]
class Mochi(LatentFormat):
latent_channels = 12
latent_dimensions = 3
@@ -434,8 +485,191 @@ class LTXV(LatentFormat):
class LTXAV(LTXV):
def __init__(self):
self.latent_rgb_factors = None
self.latent_rgb_factors_bias = None
# video-stream preview factors for the packed AV latent (audio stream is not previewed)
self.latent_rgb_factors = [
[ 0.001135, -0.010555, -0.004925],
[-0.008019, -0.006231, -0.005564],
[ 0.012637, 0.005605, 0.012713],
[ 0.023454, 0.020771, 0.017844],
[-0.011940, -0.000932, 0.009292],
[ 0.018602, 0.011018, 0.013969],
[-0.036369, -0.046631, -0.057898],
[-0.031919, 0.000131, 0.015214],
[ 0.014519, 0.021041, 0.015325],
[ 0.018889, 0.016149, -0.002836],
[-0.003784, -0.006057, -0.008195],
[ 0.013262, 0.030259, 0.029775],
[ 0.050465, 0.050366, 0.025255],
[ 0.018628, 0.007691, 0.002893],
[-0.015698, -0.008451, -0.000676],
[-0.013600, -0.012587, -0.004437],
[ 0.012482, 0.021469, 0.027913],
[-0.018241, -0.013488, -0.010975],
[ 0.013828, 0.012568, 0.021984],
[ 0.017911, 0.006552, 0.005567],
[ 0.026769, 0.006803, -0.009360],
[-0.006794, -0.008447, -0.013921],
[ 0.029708, 0.018671, 0.022811],
[-0.014732, -0.019169, 0.000903],
[ 0.019607, 0.032595, 0.053409],
[-0.003721, 0.003976, 0.010364],
[-0.020193, -0.026076, -0.036068],
[-0.002328, 0.006527, 0.013052],
[ 0.017171, 0.009224, 0.006548],
[ 0.001104, -0.000591, 0.000147],
[-0.000217, 0.011834, 0.017945],
[-0.015329, -0.012463, -0.006178],
[-0.009478, -0.008680, -0.004107],
[-0.005565, -0.006006, -0.001493],
[ 0.009451, 0.008794, 0.013207],
[-0.009989, -0.008027, -0.009568],
[-0.001505, -0.008805, -0.006828],
[ 0.001105, 0.008999, 0.009079],
[ 0.025935, 0.016426, 0.008036],
[ 0.006313, 0.000694, -0.006039],
[-0.001893, -0.006951, -0.009560],
[-0.007082, -0.002566, -0.007152],
[-0.005231, 0.004829, 0.008220],
[-0.004333, 0.001251, -0.004852],
[-0.017024, -0.012730, -0.007457],
[ 0.024988, 0.032963, 0.036556],
[ 0.013697, 0.012278, 0.009979],
[-0.013751, -0.008369, -0.015446],
[-0.009348, -0.001047, 0.007622],
[-0.003135, -0.003350, -0.003766],
[ 0.007436, 0.004957, 0.010480],
[ 0.018315, 0.022066, 0.021104],
[-0.005621, -0.006770, -0.008219],
[-0.007427, 0.001911, -0.001231],
[-0.007413, 0.000486, -0.006039],
[-0.014698, -0.007160, 0.006509],
[ 0.013775, 0.014185, 0.008203],
[ 0.060246, 0.069787, 0.072833],
[ 0.009861, 0.004870, 0.001194],
[-0.003660, 0.003251, 0.008015],
[ 0.003696, -0.003680, -0.008851],
[ 0.014924, 0.006196, 0.005282],
[-0.006740, -0.004319, -0.006729],
[ 0.020635, 0.015163, 0.012385],
[-0.032623, -0.006105, 0.010436],
[-0.058988, -0.030162, -0.037961],
[-0.035614, -0.021929, -0.011062],
[-0.023412, -0.011305, -0.005054],
[-0.002716, -0.005184, -0.004084],
[ 0.014591, 0.015294, 0.014045],
[ 0.008310, 0.002466, -0.003225],
[ 0.005176, 0.001119, 0.000695],
[-0.021569, -0.030886, -0.044732],
[ 0.007517, 0.003891, 0.000551],
[-0.006793, 0.004059, 0.010184],
[-0.086481, -0.082033, -0.083414],
[ 0.004192, 0.000762, -0.008658],
[ 0.010970, 0.009002, 0.007384],
[ 0.004042, -0.006732, -0.011031],
[ 0.012164, 0.006401, 0.007483],
[ 0.029252, 0.013990, 0.011128],
[ 0.048452, 0.034648, 0.016269],
[ 0.024104, 0.012647, 0.011754],
[-0.013216, -0.020192, -0.019752],
[-0.010799, -0.008535, -0.005467],
[ 0.005823, 0.001403, 0.001890],
[ 0.052393, 0.044771, 0.032777],
[ 0.007576, -0.008080, -0.012453],
[ 0.009830, 0.004244, 0.001213],
[-0.025867, -0.013169, -0.010636],
[ 0.008494, 0.003135, 0.000790],
[ 0.003969, -0.002625, -0.010204],
[ 0.006509, 0.008272, 0.020819],
[-0.004943, -0.013424, -0.015351],
[ 0.005541, 0.009136, -0.003666],
[-0.014300, -0.015864, -0.016853],
[ 0.002650, 0.028393, 0.014125],
[-0.027661, -0.045422, -0.064995],
[ 0.009220, 0.015522, 0.010574],
[-0.002236, 0.002915, 0.004557],
[-0.020269, -0.008212, -0.000532],
[ 0.019294, 0.003655, -0.002809],
[ 0.007116, -0.002784, 0.000017],
[ 0.057277, 0.073270, 0.074401],
[-0.002616, -0.001696, -0.000498],
[ 0.007248, 0.009793, 0.022829],
[-0.002590, -0.005601, -0.000436],
[-0.007681, 0.003893, -0.004119],
[-0.057392, -0.045545, -0.025290],
[ 0.045188, 0.047985, 0.054059],
[ 0.000937, -0.008861, -0.038406],
[-0.010192, -0.008036, -0.005385],
[-0.030222, -0.027498, -0.030765],
[-0.008359, 0.013247, 0.010918],
[ 0.004102, 0.002093, 0.006934],
[ 0.039461, 0.027339, 0.008284],
[-0.075747, -0.076340, -0.071625],
[ 0.002692, 0.005096, -0.002247],
[-0.002453, -0.002785, -0.010483],
[ 0.012265, 0.005481, 0.001729],
[ 0.017755, 0.008655, 0.003532],
[ 0.055560, 0.049128, 0.044137],
[-0.025861, -0.023798, -0.018815],
[-0.014876, -0.010770, -0.010713],
[-0.017315, -0.012599, -0.008661],
[-0.008461, -0.006210, -0.007744],
[-0.040175, -0.042255, -0.048119],
[-0.019355, -0.021055, -0.021919],
]
self.latent_rgb_factors_bias = [-0.347892, -0.363814, -0.370287]
class MiniMaxH3Video(LatentFormat):
latent_channels = 24
latent_dimensions = 3
spacial_downscale_ratio = 16
temporal_downscale_ratio = 4
scale_factor = 1.0
taesd_decoder_name = "taeh3"
latent_rgb_factors = [
[-0.018555, 0.024344, -0.017536],
[ 0.150164, 0.137244, 0.129221],
[ 0.027367, -0.050369, -0.208606],
[-0.000793, -0.164622, -0.323161],
[-0.048556, 0.013970, -0.074286],
[ 0.011740, 0.014172, -0.006906],
[ 0.061517, 0.061212, 0.110025],
[ 0.035321, 0.086879, 0.110059],
[-0.017426, 0.002997, 0.035356],
[ 0.531539, 0.548819, 0.624404],
[-0.024968, -0.040234, -0.034302],
[-0.032549, -0.029096, -0.017221],
[ 0.022609, 0.020286, 0.050661],
[-0.084001, -0.038131, -0.020805],
[-0.018830, 0.010412, 0.061120],
[ 0.020777, 0.011196, -0.030994],
[-0.008390, -0.012201, -0.025687],
[-0.013281, -0.002924, 0.006331],
[ 0.000260, 0.001833, -0.011038],
[ 0.105471, 0.100482, 0.132106],
[ 0.016529, 0.015213, 0.009999],
[-0.014015, -0.017438, -0.019134],
[-0.033787, -0.009984, -0.019725],
[ 0.004224, 0.017284, 0.027196],
]
latent_rgb_factors_bias = [ 0.057426, -0.022078, -0.071449]
class MiniMaxH3AV(MiniMaxH3Video):
# max channels across the two streams (video 24, audio 32) so per-stream slices keep both streams whole
latent_channels = 32
def fix_empty_latent(self, latent):
video_latent_channels = MiniMaxH3Video.latent_channels
audio_latent_channels = 32
audio_channels = 2
frames_per_token = (1, 4, 4, 4, 4)
audio_frame_rescale = 5.0 / 3.0
video = latent[:, :video_latent_channels].clone()
frame_count = sum(frames_per_token[i % len(frames_per_token)] for i in range(video.shape[2]))
audio_t = round(frame_count * audio_frame_rescale)
audio = latent.new_zeros((latent.shape[0], audio_latent_channels, audio_channels, audio_t))
return comfy.nested_tensor.NestedTensor((video, audio))
class HunyuanVideo(LatentFormat):
latent_channels = 16
@@ -788,6 +1022,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/ernie/model.py
+12 -5
View File
@@ -5,6 +5,7 @@ import torch.nn.functional as F
from comfy.ldm.modules.attention import optimized_attention
import comfy.model_management
import comfy.ops
import comfy.quant_ops
def rope(pos: torch.Tensor, dim: int, theta: int) -> torch.Tensor:
@@ -111,11 +112,17 @@ class ErnieImageAttention(nn.Module):
query = q_flat.view(B, S, self.heads, self.head_dim)
key = k_flat.view(B, S, self.heads, self.head_dim)
query = self.norm_q(query)
key = self.norm_k(key)
if image_rotary_emb is not None:
query, key = comfy.quant_ops.ck.apply_rope_split_half(query, key, image_rotary_emb)
if image_rotary_emb is not None and not comfy.model_management.in_training:
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.norm_q, query, offloadable=True)
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.norm_k, key, offloadable=True)
query, key = comfy.quant_ops.ck.rms_rope_split_half(query, key, image_rotary_emb, q_scale, k_scale, self.norm_q.eps)
comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_offload_stream)
comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_offload_stream)
else:
query = self.norm_q(query)
key = self.norm_k(key)
if image_rotary_emb is not None:
query, key = comfy.quant_ops.ck.apply_rope_split_half(query, key, image_rotary_emb)
q_flat = query.reshape(B, S, -1)
k_flat = key.reshape(B, S, -1)
comfy/ldm/ideogram4/model.py
+32 -5
View File
@@ -12,10 +12,13 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.model_management
import comfy.ops
import comfy.patcher_extension
import comfy.quant_ops
from comfy.ldm.lumina.model import FeedForward
from comfy.ldm.modules.attention import optimized_attention_masked
from comfy.text_encoders.llama import apply_rope, precompute_freqs_cis
from comfy.text_encoders.llama import precompute_freqs_cis
# Per-token role indicators
SEQUENCE_PADDING_INDICATOR = -1
@@ -25,6 +28,22 @@ LLM_TOKEN_INDICATOR = 3
IMAGE_POSITION_OFFSET = 65536
def _split_half_rope_matrix(freqs_cis):
cos, sin, neg_sin = freqs_cis
half_dim = sin.shape[-1]
matrix = torch.stack(
(cos[..., :half_dim], neg_sin, sin, cos[..., half_dim:]), dim=-1
)
return matrix.reshape(*matrix.shape[:-1], 2, 2).unsqueeze(2)
def _apply_rope_split_half1(x, freqs_cis):
x_dtype = x.dtype
x = x.reshape(*x.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2).to(freqs_cis.dtype)
output = freqs_cis[..., 0] * x[..., 0] + freqs_cis[..., 1] * x[..., 1]
return output.movedim(-1, -2).reshape(*x.shape[:-3], -1).to(x_dtype)
class Ideogram4Attention(nn.Module):
def __init__(self, hidden_size, num_heads, eps=1e-5, dtype=None, device=None, operations=None):
super().__init__()
@@ -42,16 +61,23 @@ class Ideogram4Attention(nn.Module):
qkv = self.qkv(x).view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
q, k, v = qkv.unbind(dim=2)
q = self.norm_q(q)
k = self.norm_k(k)
if comfy.model_management.in_training:
q = _apply_rope_split_half1(self.norm_q(q), freqs_cis)
k = _apply_rope_split_half1(self.norm_k(k), freqs_cis)
else:
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.norm_q, q, offloadable=True)
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.norm_k, k, offloadable=True)
q, k = comfy.quant_ops.ck.rms_rope_split_half(
q, k, freqs_cis, q_scale, k_scale, self.norm_q.eps
)
comfy.ops.uncast_bias_weight(self.norm_q, q_scale, None, q_offload_stream)
comfy.ops.uncast_bias_weight(self.norm_k, k_scale, None, k_offload_stream)
# (B, heads, L, head_dim)
q = q.transpose(1, 2)
k = k.transpose(1, 2)
v = v.transpose(1, 2)
q, k = apply_rope(q, k, freqs_cis)
out = optimized_attention_masked(q, k, v, self.num_heads, attn_mask, skip_reshape=True, transformer_options=transformer_options)
return self.o(out)
@@ -181,6 +207,7 @@ class Ideogram4Transformer(nn.Module):
self.head_dim, position_ids[0].transpose(0, 1), self.rope_theta,
rope_dims=self.mrope_section, interleaved_mrope=True, device=position_ids.device,
)
freqs_cis = _split_half_rope_matrix(freqs_cis)
if attn_mask is not None and attn_mask.dtype == torch.bool:
attn_mask = torch.zeros_like(attn_mask, dtype=h.dtype).masked_fill_(~attn_mask, -torch.finfo(h.dtype).max)
comfy/ldm/joyimage/model.py
+12 -3
View File
@@ -94,12 +94,21 @@ class JoyImageAttention(nn.Module):
txt_k = txt_k.unflatten(-1, (heads, -1))
txt_v = txt_v.unflatten(-1, (heads, -1))
img_q = self.img_attn_q_norm(img_q)
img_k = self.img_attn_k_norm(img_k)
txt_q = self.txt_attn_q_norm(txt_q)
txt_k = self.txt_attn_k_norm(txt_k)
img_q, img_k = comfy_kitchen.apply_rope(img_q, img_k, image_rotary_emb)
img_q_scale, _, img_q_offload_stream = comfy.ops.cast_bias_weight(self.img_attn_q_norm, img_q, offloadable=True)
img_k_scale, _, img_k_offload_stream = comfy.ops.cast_bias_weight(self.img_attn_k_norm, img_k, offloadable=True)
img_q, img_k = comfy_kitchen.rms_rope(
img_q,
img_k,
image_rotary_emb,
img_q_scale,
img_k_scale,
self.img_attn_q_norm.eps,
)
comfy.ops.uncast_bias_weight(self.img_attn_q_norm, img_q_scale, None, img_q_offload_stream)
comfy.ops.uncast_bias_weight(self.img_attn_k_norm, img_k_scale, None, img_k_offload_stream)
joint_q = torch.cat([img_q, txt_q], dim=1)
joint_k = torch.cat([img_k, txt_k], dim=1)
comfy/ldm/krea2/model.py
+124 -23
View File
@@ -15,6 +15,7 @@ from einops import rearrange
import comfy.model_management
import comfy.patcher_extension
import comfy.ldm.common_dit
import comfy.utils
from comfy.ldm.flux.layers import EmbedND, timestep_embedding
from comfy.ldm.flux.math import apply_rope
from comfy.ldm.modules.attention import optimized_attention_masked
@@ -73,11 +74,20 @@ class Attention(nn.Module):
self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype)
def forward(self, x, freqs=None, mask=None, transformer_options={}):
transformer_patches = transformer_options.get("patches", {})
extra_options = transformer_options.copy()
q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x)
q = rearrange(q, "B L (H D) -> B H L D", H=self.heads)
k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads)
v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads)
q, k = self.qknorm(q, k)
if "block_index" in transformer_options and "attn1_patch" in transformer_patches:
for p in transformer_patches["attn1_patch"]:
out = p(q, k, v, pe=freqs, attn_mask=mask, extra_options=extra_options)
q, k, v = out.get("q", q), out.get("k", k), out.get("v", v)
freqs, mask = out.get("pe", freqs), out.get("attn_mask", mask)
if freqs is not None:
q, k = apply_rope(q, k, freqs)
if self.kvheads != self.heads:
@@ -86,6 +96,11 @@ class Attention(nn.Module):
v = v.repeat_interleave(rep, dim=1)
out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True,
transformer_options=transformer_options)
if "block_index" in transformer_options and "attn1_output_patch" in transformer_patches:
for p in transformer_patches["attn1_output_patch"]:
out = p(out, extra_options)
return self.wo(out * F.sigmoid(gate))
@@ -158,8 +173,44 @@ class SingleStreamBlock(nn.Module):
self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations)
self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations)
def forward(self, x, vec, freqs, mask=None, transformer_options={}):
def forward(self, x, vec, freqs, mask=None, timestep_zero_index=None, transformer_options={}):
prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
if timestep_zero_index is not None:
bs = x.shape[0]
ref_prescale = prescale[bs:]
ref_preshift = preshift[bs:]
ref_pregate = pregate[bs:]
ref_postscale = postscale[bs:]
ref_postshift = postshift[bs:]
ref_postgate = postgate[bs:]
prescale = prescale[:bs]
preshift = preshift[:bs]
pregate = pregate[:bs]
postscale = postscale[:bs]
postshift = postshift[:bs]
postgate = postgate[:bs]
pre = self.prenorm(x)
pre[:, :timestep_zero_index].mul_(1 + prescale).add_(preshift)
pre[:, timestep_zero_index:].mul_(1 + ref_prescale).add_(ref_preshift)
attn = self.attn(pre, freqs, mask, transformer_options=transformer_options)
del pre
attn[:, :timestep_zero_index].mul_(pregate)
attn[:, timestep_zero_index:].mul_(ref_pregate)
x = x + attn
del attn
post = self.postnorm(x)
post[:, :timestep_zero_index].mul_(1 + postscale).add_(postshift)
post[:, timestep_zero_index:].mul_(1 + ref_postscale).add_(ref_postshift)
mlp = self.mlp(post)
del post
mlp[:, :timestep_zero_index].mul_(postgate)
mlp[:, timestep_zero_index:].mul_(ref_postgate)
x = x + mlp
del mlp
return x
x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options)
x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
return x
@@ -181,7 +232,7 @@ class LastLayer(nn.Module):
class SingleStreamDiT(nn.Module):
def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4,
layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12,
txtheads=20, txtkvheads=20, image_model=None,
txtheads=20, txtkvheads=20, default_ref_method=None, image_model=None,
device=None, dtype=None, operations=None, **kwargs):
super().__init__()
self.dtype = dtype
@@ -191,6 +242,7 @@ class SingleStreamDiT(nn.Module):
self.heads = heads
self.txtdim = txtdim
self.txtlayers = txtlayers
self.default_ref_method = default_ref_method
headdim = features // heads
axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
@@ -221,61 +273,110 @@ class SingleStreamDiT(nn.Module):
operations.Linear(features, features * 6, device=device, dtype=dtype),
)
def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
def forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
self._forward,
self,
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options),
).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs)
).execute(x, timesteps, context, attention_mask, ref_latents, transformer_options, **kwargs)
def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
def process_img(self, x, index=0):
patch = self.patch
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
h, w = x.shape[-2] // patch, x.shape[-1] // patch
img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
img_ids = torch.zeros(h, w, 3, device=x.device, dtype=torch.float32)
img_ids[..., 0] = index
img_ids[..., 1] = torch.arange(h, device=x.device, dtype=torch.float32)[:, None]
img_ids[..., 2] = torch.arange(w, device=x.device, dtype=torch.float32)[None, :]
return img, img_ids.reshape(1, h * w, 3).repeat(x.shape[0], 1, 1), h, w
def _forward(self, x, timesteps, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
transformer_options = transformer_options.copy()
temporal = x.ndim == 5
if temporal:
b5, c5, t5, h5, w5 = x.shape
x = x.reshape(b5 * t5, c5, h5, w5)
bs, c, H_orig, W_orig = x.shape
bs, _, h_orig, w_orig = x.shape
patch = self.patch
# Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end.
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
H, W = x.shape[-2], x.shape[-1]
h_, w_ = H // patch, W // patch
# context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim).
context = self._unpack_context(context)
img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
img, imgpos, h_, w_ = self.process_img(x)
img_tokens = img.shape[1]
timestep_zero_index = None
ref_method = kwargs.get("ref_latents_method", self.default_ref_method)
if ref_method is not None and ref_latents is not None and len(ref_latents) > 0:
ref_tokens = []
ref_pos = []
ref_num_tokens = []
for index, ref in enumerate(ref_latents, 1):
if ref.ndim == 5:
rb, rc, rt, rh5, rw5 = ref.shape
ref = ref.reshape(rb * rt, rc, rh5, rw5)
ref = comfy.utils.repeat_to_batch_size(ref, bs)
kontext, kontext_ids, _, _ = self.process_img(ref, index=index)
ref_tokens.append(kontext)
ref_pos.append(kontext_ids)
ref_num_tokens.append(kontext.shape[1])
img = torch.cat([img] + ref_tokens, dim=1)
imgpos = torch.cat([imgpos] + ref_pos, dim=1)
del ref_tokens, ref_pos
if ref_method == "index_timestep_zero":
timestep_zero_index = img_tokens
transformer_options["reference_image_num_tokens"] = ref_num_tokens
img = self.first(img)
t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype))
tvec = self.tproj(t)
if timestep_zero_index is not None:
t0 = self.tmlp(timestep_embedding(torch.zeros_like(timesteps), self.tdim).unsqueeze(1).to(img.dtype))
tvec = torch.cat((tvec, self.tproj(t0)), dim=0)
context = self.txtfusion(context, mask=None, transformer_options=transformer_options)
context = self.txtmlp(context)
txtlen, imglen = context.shape[1], img.shape[1]
txtlen = context.shape[1]
device = context.device
txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
patches = transformer_options.get("patches", {})
if "post_input" in patches:
for p in patches["post_input"]:
out = p({"img": img, "txt": context, "img_ids": imgpos, "txt_ids": txtpos, "transformer_options": transformer_options})
img, context = out["img"], out["txt"]
imgpos, txtpos = out["img_ids"], out["txt_ids"]
combined = torch.cat((context, img), dim=1)
del context, img
if timestep_zero_index is not None:
timestep_zero_index += txtlen
# Position ids: text at 0, image at (0, h_idx, w_idx).
device = combined.device
txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
imgids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
imgids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
pos = torch.cat((txtpos, imgpos), dim=1)
del txtpos, imgpos
freqs = self.pe_embedder(pos)
del pos
for block in self.blocks:
combined = block(combined, tvec, freqs, None, transformer_options=transformer_options)
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "single"
transformer_options["img_slice"] = [txtlen, combined.shape[1]]
for i, block in enumerate(self.blocks):
transformer_options["block_index"] = i
combined = block(combined, tvec, freqs, None, timestep_zero_index=timestep_zero_index, transformer_options=transformer_options)
final = self.last(combined, t)
out = final[:, txtlen:txtlen + imglen, :]
del combined
out = final[:, txtlen:txtlen + img_tokens, :]
out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
h=h_, w=w_, ph=patch, pw=patch, c=self.channels)
out = out[:, :, :H_orig, :W_orig] # crop padding back off
out = out[:, :, :h_orig, :w_orig] # crop padding back off
if temporal:
out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2)
out = out.reshape(b5, t5, self.channels, h_orig, w_orig).movedim(1, 2)
return out
def _unpack_context(self, context):
comfy/ldm/lightricks/av_model.py
+39 -11
View File
@@ -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
@@ -275,7 +279,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
@@ -318,7 +325,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
@@ -328,7 +334,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
@@ -347,7 +356,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)
@@ -357,12 +369,14 @@ class BasicAVTransformerBlock(nn.Module):
ax.addcmul_(v2a_out, gate_out_v2a, value=video_to_audio_scale)
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)
@@ -414,12 +428,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
@@ -452,6 +470,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,
@@ -476,7 +496,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,
@@ -607,6 +627,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,
@@ -925,9 +947,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):
@@ -970,7 +998,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,
@@ -994,7 +1022,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
+81
View File
@@ -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
+13 -7
View File
@@ -6,9 +6,8 @@ import torch
from comfy.ldm.lightricks.model import (
CrossAttention,
FeedForward,
freqs_cis_matrix,
generate_freq_grid_np,
interleaved_freqs_cis,
split_freqs_cis,
)
from torch import nn
@@ -51,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,
@@ -75,6 +75,7 @@ class BasicTransformerBlock1D(nn.Module):
dim,
dim_out=dim,
glu=True,
ff_bias=ff_bias,
dtype=dtype,
device=device,
operations=operations,
@@ -124,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,
@@ -149,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,
@@ -244,12 +247,15 @@ class Embeddings1DConnector(nn.Module):
expected_freqs = dim // 2
current_freqs = freqs.shape[-1]
pad_size = expected_freqs - current_freqs
cos_freq, sin_freq = split_freqs_cis(
freqs, pad_size, self.num_attention_heads
)
else:
cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
return cos_freq.to(dtype=out_dtype), sin_freq.to(dtype=out_dtype), self.split_rope
pad_size = dim % n_elem
return freqs_cis_matrix(
freqs,
pad_size,
self.split_rope,
self.num_attention_heads,
out_dtype,
)
def forward(
self,
comfy/ldm/lightricks/latent_upsampler.py
+20 -15
View File
@@ -97,11 +97,11 @@ class SpatialRationalResampler(nn.Module):
For dims==3, work per-frame for spatial scaling (temporal axis untouched).
"""
def __init__(self, mid_channels: int, scale: float):
def __init__(self, mid_channels: int, scale: float, operations):
super().__init__()
self.scale = float(scale)
self.num, self.den = _rational_for_scale(self.scale)
self.conv = nn.Conv2d(
self.conv = operations.Conv2d(
mid_channels, (self.num**2) * mid_channels, kernel_size=3, padding=1
)
self.pixel_shuffle = PixelShuffleND(2, upscale_factors=(self.num, self.num))
@@ -119,18 +119,18 @@ class SpatialRationalResampler(nn.Module):
class ResBlock(nn.Module):
def __init__(
self, channels: int, mid_channels: Optional[int] = None, dims: int = 3
self, channels: int, operations, mid_channels: Optional[int] = None, dims: int = 3
):
super().__init__()
if mid_channels is None:
mid_channels = channels
Conv = nn.Conv2d if dims == 2 else nn.Conv3d
Conv = operations.Conv2d if dims == 2 else operations.Conv3d
self.conv1 = Conv(channels, mid_channels, kernel_size=3, padding=1)
self.norm1 = nn.GroupNorm(32, mid_channels)
self.norm1 = operations.GroupNorm(32, mid_channels)
self.conv2 = Conv(mid_channels, channels, kernel_size=3, padding=1)
self.norm2 = nn.GroupNorm(32, channels)
self.norm2 = operations.GroupNorm(32, channels)
self.activation = nn.SiLU()
def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -159,6 +159,7 @@ class LatentUpsampler(nn.Module):
def __init__(
self,
operations,
in_channels: int = 128,
mid_channels: int = 512,
num_blocks_per_stage: int = 4,
@@ -179,34 +180,34 @@ class LatentUpsampler(nn.Module):
self.spatial_scale = float(spatial_scale)
self.rational_resampler = rational_resampler
Conv = nn.Conv2d if dims == 2 else nn.Conv3d
Conv = operations.Conv2d if dims == 2 else operations.Conv3d
self.initial_conv = Conv(in_channels, mid_channels, kernel_size=3, padding=1)
self.initial_norm = nn.GroupNorm(32, mid_channels)
self.initial_norm = operations.GroupNorm(32, mid_channels)
self.initial_activation = nn.SiLU()
self.res_blocks = nn.ModuleList(
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
[ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)]
)
if spatial_upsample and temporal_upsample:
self.upsampler = nn.Sequential(
nn.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
operations.Conv3d(mid_channels, 8 * mid_channels, kernel_size=3, padding=1),
PixelShuffleND(3),
)
elif spatial_upsample:
if rational_resampler:
self.upsampler = SpatialRationalResampler(
mid_channels=mid_channels, scale=self.spatial_scale
mid_channels=mid_channels, scale=self.spatial_scale, operations=operations
)
else:
self.upsampler = nn.Sequential(
nn.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
operations.Conv2d(mid_channels, 4 * mid_channels, kernel_size=3, padding=1),
PixelShuffleND(2),
)
elif temporal_upsample:
self.upsampler = nn.Sequential(
nn.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
operations.Conv3d(mid_channels, 2 * mid_channels, kernel_size=3, padding=1),
PixelShuffleND(1),
)
else:
@@ -215,11 +216,14 @@ class LatentUpsampler(nn.Module):
)
self.post_upsample_res_blocks = nn.ModuleList(
[ResBlock(mid_channels, dims=dims) for _ in range(num_blocks_per_stage)]
[ResBlock(mid_channels, dims=dims, operations=operations) for _ in range(num_blocks_per_stage)]
)
self.final_conv = Conv(mid_channels, in_channels, kernel_size=3, padding=1)
def get_dtype(self):
return getattr(self.initial_conv, "weight_comfy_model_dtype", self.initial_conv.weight.dtype)
def forward(self, latent: torch.Tensor) -> torch.Tensor:
b, c, f, h, w = latent.shape
@@ -266,7 +270,7 @@ class LatentUpsampler(nn.Module):
return x
@classmethod
def from_config(cls, config):
def from_config(cls, config, operations):
return cls(
in_channels=config.get("in_channels", 4),
mid_channels=config.get("mid_channels", 128),
@@ -276,6 +280,7 @@ class LatentUpsampler(nn.Module):
temporal_upsample=config.get("temporal_upsample", False),
spatial_scale=config.get("spatial_scale", 2.0),
rational_resampler=config.get("rational_resampler", False),
operations=operations,
)
def config(self):
comfy/ldm/lightricks/model.py
+179 -83
View File
@@ -12,6 +12,9 @@ from torch import nn
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
@@ -300,62 +303,68 @@ 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):
cos_freqs, sin_freqs = freqs_cis[0], freqs_cis[1]
split_pe = freqs_cis[2] if len(freqs_cis) > 2 else False
return (
apply_split_rotary_emb(input_tensor, cos_freqs, sin_freqs)
if split_pe else
apply_interleaved_rotary_emb(input_tensor, cos_freqs, sin_freqs)
rotation_matrix, split_pe = freqs_cis
original_shape = input_tensor.shape
input_tensor = input_tensor.reshape(
input_tensor.shape[0], input_tensor.shape[1], rotation_matrix.shape[2], -1
)
def apply_interleaved_rotary_emb(input_tensor, cos_freqs, sin_freqs): # TODO: remove duplicate funcs and pick the best/fastest one
t_dup = rearrange(input_tensor, "... (d r) -> ... d r", r=2)
t1, t2 = t_dup.unbind(dim=-1)
t_dup = torch.stack((-t2, t1), dim=-1)
input_tensor_rot = rearrange(t_dup, "... d r -> ... (d r)")
if comfy.model_management.in_training:
if split_pe:
t = input_tensor.reshape(*input_tensor.shape[:-1], 2, -1).movedim(-2, -1).unsqueeze(-2)
else:
t = input_tensor.reshape(*input_tensor.shape[:-1], -1, 1, 2)
t = t.to(rotation_matrix.dtype)
output = rotation_matrix[..., 0] * t[..., 0] + rotation_matrix[..., 1] * t[..., 1]
if split_pe:
output = output.movedim(-1, -2)
output = output.reshape(input_tensor.shape).type_as(input_tensor)
elif split_pe:
output = comfy.quant_ops.ck.apply_rope_split_half1(input_tensor, rotation_matrix)
else:
output = comfy.quant_ops.ck.apply_rope1(input_tensor, rotation_matrix)
return output.reshape(original_shape)
out = input_tensor * cos_freqs + input_tensor_rot * sin_freqs
def apply_rotary_emb_qk(q, k, freqs_cis):
if comfy.model_management.in_training:
return apply_rotary_emb(q, freqs_cis), apply_rotary_emb(k, freqs_cis)
return out
def apply_split_rotary_emb(input_tensor, cos, sin):
needs_reshape = False
if input_tensor.ndim != 4 and cos.ndim == 4:
B, H, T, _ = cos.shape
input_tensor = input_tensor.reshape(B, T, H, -1).swapaxes(1, 2)
needs_reshape = True
split_input = rearrange(input_tensor, "... (d r) -> ... d r", d=2)
first_half_input = split_input[..., :1, :]
second_half_input = split_input[..., 1:, :]
output = split_input * cos.unsqueeze(-2)
first_half_output = output[..., :1, :]
second_half_output = output[..., 1:, :]
first_half_output.addcmul_(-sin.unsqueeze(-2), second_half_input)
second_half_output.addcmul_(sin.unsqueeze(-2), first_half_input)
output = rearrange(output, "... d r -> ... (d r)")
return output.swapaxes(1, 2).reshape(B, T, -1) if needs_reshape else output
rotation_matrix, split_pe = freqs_cis
q_shape = q.shape
k_shape = k.shape
q = q.reshape(q.shape[0], q.shape[1], rotation_matrix.shape[2], -1)
k = k.reshape(k.shape[0], k.shape[1], rotation_matrix.shape[2], -1)
if split_pe:
q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rotation_matrix)
else:
q, k = comfy.quant_ops.ck.apply_rope(q, k, rotation_matrix)
return q.reshape(q_shape), k.reshape(k_shape)
class GuideAttentionMask:
@@ -453,24 +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)
if pe is not None:
q = apply_rotary_emb(q, pe)
k = apply_rotary_emb(k, pe if k_pe is None else k_pe)
if 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:
@@ -489,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__()
@@ -505,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,
@@ -527,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)
@@ -581,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
@@ -653,36 +680,23 @@ def generate_freqs(indices, indices_grid, max_pos, use_middle_indices_grid):
)
return freqs
def interleaved_freqs_cis(freqs, pad_size):
cos_freq = freqs.cos().repeat_interleave(2, dim=-1)
sin_freq = freqs.sin().repeat_interleave(2, dim=-1)
if pad_size != 0:
cos_padding = torch.ones_like(cos_freq[:, :, : pad_size])
sin_padding = torch.zeros_like(cos_freq[:, :, : pad_size])
cos_freq = torch.cat([cos_padding, cos_freq], dim=-1)
sin_freq = torch.cat([sin_padding, sin_freq], dim=-1)
return cos_freq, sin_freq
def split_freqs_cis(freqs, pad_size, num_attention_heads):
cos_freq = freqs.cos()
sin_freq = freqs.sin()
if pad_size != 0:
cos_padding = torch.ones_like(cos_freq[:, :, :pad_size])
sin_padding = torch.zeros_like(sin_freq[:, :, :pad_size])
cos_freq = torch.concatenate([cos_padding, cos_freq], axis=-1)
sin_freq = torch.concatenate([sin_padding, sin_freq], axis=-1)
# Reshape freqs to be compatible with multi-head attention
B , T, half_HD = cos_freq.shape
def freqs_cis_matrix(freqs, pad_size, split_mode, num_attention_heads, out_dtype):
cos_freq = freqs.cos().to(out_dtype)
sin_freq = freqs.sin().to(out_dtype)
if pad_size:
matrix_pad_size = pad_size if split_mode else pad_size // 2
cos_padding = torch.ones_like(cos_freq[:, :, :matrix_pad_size])
sin_padding = torch.zeros_like(sin_freq[:, :, :matrix_pad_size])
cos_freq = torch.cat((cos_padding, cos_freq), dim=-1)
sin_freq = torch.cat((sin_padding, sin_freq), dim=-1)
B, T, half_HD = cos_freq.shape
cos_freq = cos_freq.reshape(B, T, num_attention_heads, half_HD // num_attention_heads)
sin_freq = sin_freq.reshape(B, T, num_attention_heads, half_HD // num_attention_heads)
cos_freq = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2)
sin_freq = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2)
return cos_freq, sin_freq
rotation_matrix = torch.stack(
(cos_freq, -sin_freq, sin_freq, cos_freq), dim=-1
)
return rotation_matrix.reshape(*rotation_matrix.shape[:-1], 2, 2), split_mode
class LTXBaseModel(torch.nn.Module, ABC):
"""
@@ -709,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,
@@ -738,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
@@ -765,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
)
@@ -885,12 +910,17 @@ class LTXBaseModel(torch.nn.Module, ABC):
expected_freqs = dim // 2
current_freqs = freqs.shape[-1]
pad_size = expected_freqs - current_freqs
cos_freq, sin_freq = split_freqs_cis(freqs, pad_size, num_attention_heads)
else:
# 2 because of cos and sin by 3 for (t, x, y), 1 for temporal only
n_elem = 2 * indices_grid.shape[1]
cos_freq, sin_freq = interleaved_freqs_cis(freqs, dim % n_elem)
return cos_freq.to(out_dtype), sin_freq.to(out_dtype), split_mode
pad_size = dim % n_elem
return freqs_cis_matrix(
freqs,
pad_size,
split_mode,
num_attention_heads,
out_dtype,
)
def _prepare_positional_embeddings(self, pixel_coords, frame_rate, x_dtype):
"""Prepare positional embeddings."""
@@ -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
+2 -3
View File
@@ -1,6 +1,5 @@
import json
from dataclasses import dataclass
import math
import torch
import torchaudio
@@ -185,8 +184,8 @@ class AudioVAE(torch.nn.Module):
self.autoencoder.mel_bins,
)
def num_of_latents_from_frames(self, frames_number: int, frame_rate: int) -> int:
return math.ceil((float(frames_number) / frame_rate) * self.latents_per_second)
def num_of_latents_from_frames(self, frames_number: int, frame_rate: float) -> int:
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/causal_conv3d.py
+8 -2
View File
@@ -49,6 +49,12 @@ class CausalConv3d(nn.Module):
)
self.temporal_cache_state={}
def _empty_output(self, x):
# empty (0 frame) outputs must still have the conv's output channels and spatial dims
h = (x.shape[3] + 2 * self.conv.padding[1] - self.conv.kernel_size[1]) // self.conv.stride[1] + 1
w = (x.shape[4] + 2 * self.conv.padding[2] - self.conv.kernel_size[2]) // self.conv.stride[2] + 1
return x.new_empty((x.shape[0], self.out_channels, 0, h, w))
def forward(self, x, causal: bool = True):
tid = threading.get_ident()
@@ -58,7 +64,7 @@ class CausalConv3d(nn.Module):
if not causal:
padding_length = padding_length // 2
if x.shape[2] == 0:
return x
return self._empty_output(x)
cached = x[:, :, :1, :, :].repeat((1, 1, padding_length, 1, 1))
pieces = [ cached, x ]
if is_end and not causal:
@@ -83,7 +89,7 @@ class CausalConv3d(nn.Module):
elif is_end:
self.temporal_cache_state[tid] = (None, True)
return self.conv(x) if x.shape[2] >= self.time_kernel_size else x[:, :, :0, :, :]
return self.conv(x) if x.shape[2] >= self.time_kernel_size else self._empty_output(x)
@property
def weight(self):
comfy/ldm/lightricks/vae/causal_video_autoencoder.py
+3 -3
View File
@@ -390,10 +390,10 @@ class Decoder(nn.Module):
# Compute output channel to be product of all channel-multiplier blocks
output_channel = base_channels
for block_name, block_params in list(reversed(blocks)):
for block_name, block_params in blocks:
block_params = block_params if isinstance(block_params, dict) else {}
if block_name == "res_x_y":
output_channel = output_channel * block_params.get("multiplier", 2)
output_channel = block_params.get("in_channels", output_channel * block_params.get("multiplier", 2))
if block_name == "compress_all":
output_channel = output_channel * block_params.get("multiplier", 1)
if block_name == "compress_space":
@@ -432,7 +432,7 @@ class Decoder(nn.Module):
spatial_padding_mode=spatial_padding_mode,
)
elif block_name == "res_x_y":
output_channel = output_channel // block_params.get("multiplier", 2)
output_channel = block_params.get("out_channels", output_channel // block_params.get("multiplier", 2))
block = ResnetBlock3D(
dims=dims,
in_channels=input_channel,
comfy/ldm/lightricks/vae/na_diffusion_decoder.py
+520
View File
@@ -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/lumina/model.py
+19 -4
View File
@@ -6,6 +6,9 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.ldm.common_dit
import comfy.model_management
import comfy.ops
import comfy.quant_ops
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
from comfy.ldm.modules.attention import optimized_attention_masked
@@ -97,6 +100,7 @@ class JointAttention(nn.Module):
self.n_local_kv_heads = self.n_kv_heads
self.n_rep = self.n_local_heads // self.n_local_kv_heads
self.head_dim = dim // n_heads
self.qk_norm = qk_norm
self.qkv = operation_settings.get("operations").Linear(
dim,
@@ -151,10 +155,21 @@ class JointAttention(nn.Module):
xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
xq = self.q_norm(xq)
xk = self.k_norm(xk)
xq, xk = apply_rope(xq, xk, freqs_cis)
if self.qk_norm and not comfy.model_management.in_training:
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, xq, offloadable=True)
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, xk, offloadable=True)
epsilon = self.q_norm.eps if self.q_norm.eps is not None else torch.finfo(torch.float32).eps
if self.n_local_heads == self.n_local_kv_heads:
xq, xk = comfy.quant_ops.ck.rms_rope(xq, xk, freqs_cis, q_scale, k_scale, epsilon)
else:
xq = comfy.quant_ops.ck.rms_rope1(xq, freqs_cis, q_scale, epsilon)
xk = comfy.quant_ops.ck.rms_rope1(xk, freqs_cis, k_scale, epsilon)
comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
else:
xq = self.q_norm(xq)
xk = self.k_norm(xk)
xq, xk = apply_rope(xq, xk, freqs_cis)
n_rep = self.n_local_heads // self.n_local_kv_heads
if n_rep >= 1:
comfy/ldm/mage_flow/model.py
+186
View File
@@ -0,0 +1,186 @@
# Mage-Flow (https://github.com/microsoft/Mage) native-resolution MMDiT (MIT)
# Architecture is a 12-layer variant of the Qwen-Image double-stream block with
# patch_size=1 (no 2x2 packing), unrotated text tokens and a bf16-rounded
# timestep frequency table.
import math
import torch
import torch.nn as nn
from typing import Optional, Tuple
from comfy.ldm.lightricks.model import TimestepEmbedding
from comfy.ldm.flux.layers import EmbedND
from comfy.ldm.qwen_image.model import QwenImageTransformerBlock, LastLayer
import comfy.patcher_extension
class MageTimestepProjEmbeddings(nn.Module):
def __init__(self, embedding_dim, dtype=None, device=None, operations=None):
super().__init__()
self.timestep_embedder = TimestepEmbedding(
in_channels=256, time_embed_dim=embedding_dim,
dtype=dtype, device=device, operations=operations
)
def forward(self, timestep, hidden_states):
half_dim = 128
exponent = -math.log(10000) * torch.arange(half_dim, dtype=torch.float32, device=timestep.device) / half_dim
emb = torch.exp(exponent).to(timestep.dtype)
emb = timestep[:, None].float() * emb[None, :]
emb = 1000.0 * emb
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
emb = torch.cat([emb[:, half_dim:], emb[:, :half_dim]], dim=-1) # flip_sin_to_cos
return self.timestep_embedder(emb.to(dtype=hidden_states.dtype))
class MageFlowTransformer2DModel(nn.Module):
def __init__(
self,
in_channels: int = 128,
out_channels: Optional[int] = 128,
num_layers: int = 12,
attention_head_dim: int = 128,
num_attention_heads: int = 24,
joint_attention_dim: int = 2560,
axes_dims_rope: Tuple[int, int, int] = (16, 56, 56),
image_model=None,
dtype=None,
device=None,
operations=None,
):
super().__init__()
self.dtype = dtype
self.patch_size = 1
self.in_channels = in_channels
self.out_channels = out_channels or in_channels
self.inner_dim = num_attention_heads * attention_head_dim
self.pe_embedder = EmbedND(dim=attention_head_dim, theta=10000, axes_dim=list(axes_dims_rope))
self.time_text_embed = MageTimestepProjEmbeddings(embedding_dim=self.inner_dim, dtype=dtype, device=device, operations=operations)
self.txt_norm = operations.RMSNorm(joint_attention_dim, eps=1e-6, dtype=dtype, device=device)
self.img_in = operations.Linear(in_channels, self.inner_dim, dtype=dtype, device=device)
self.txt_in = operations.Linear(joint_attention_dim, self.inner_dim, dtype=dtype, device=device)
self.transformer_blocks = nn.ModuleList([
QwenImageTransformerBlock(
dim=self.inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
dtype=dtype,
device=device,
operations=operations
)
for _ in range(num_layers)
])
self.norm_out = LastLayer(self.inner_dim, self.inner_dim, dtype=dtype, device=device, operations=operations)
self.proj_out = operations.Linear(self.inner_dim, self.out_channels, bias=True, dtype=dtype, device=device)
def process_img(self, x, index=0):
# patch_size=1: tokens are raw latent pixels, no 2x2 packing.
bs, c, h, w = x.shape
hidden_states = x.movedim(1, -1).reshape(bs, h * w, c)
img_ids = torch.zeros((h, w, 3), device=x.device)
# Frame axis: positive image index (0 = target, 1..N = reference images).
img_ids[:, :, 0] = index
# Mage scale_rope centering: positions [-ceil(n/2), floor(n/2)), i.e.
# offset by (n - n//2). Differs from Qwen-Image's -(n//2) for odd sizes.
img_ids[:, :, 1] = img_ids[:, :, 1] + torch.arange(h, device=x.device)[:, None] - (h - h // 2)
img_ids[:, :, 2] = img_ids[:, :, 2] + torch.arange(w, device=x.device)[None, :] - (w - w // 2)
return hidden_states, img_ids.reshape(h * w, 3).unsqueeze(0).expand(bs, -1, -1), (h, w)
def forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, **kwargs):
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
self._forward,
self,
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
).execute(x, timestep, context, attention_mask, ref_latents, transformer_options, **kwargs)
def _forward(self, x, timestep, context, attention_mask=None, ref_latents=None, transformer_options={}, control=None, **kwargs):
if attention_mask is not None and not torch.is_floating_point(attention_mask):
attention_mask = (attention_mask - 1).to(x.dtype) * torch.finfo(x.dtype).max
hidden_states, img_ids, orig_shape = self.process_img(x)
num_embeds = hidden_states.shape[1]
if ref_latents is not None:
ref_num_tokens = []
index = 0
for ref in ref_latents:
index += 1
kontext, kontext_ids, _ = self.process_img(ref, index=index)
hidden_states = torch.cat([hidden_states, kontext], dim=1)
img_ids = torch.cat([img_ids, kontext_ids], dim=1)
ref_num_tokens.append(kontext.shape[1])
transformer_options = transformer_options.copy()
transformer_options["reference_image_num_tokens"] = ref_num_tokens
# Text tokens are not rotated in Mage-Flow: RoPE at position 0 is the
# identity rotation.
txt_ids = torch.zeros((x.shape[0], context.shape[1], 3), device=x.device)
hidden_states = self.img_in(hidden_states)
context = self.txt_norm(context)
context = self.txt_in(context)
temb = self.time_text_embed(timestep, hidden_states)
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
if "post_input" in patches:
for p in patches["post_input"]:
out = p({"img": hidden_states, "txt": context, "img_ids": img_ids, "txt_ids": txt_ids, "transformer_options": transformer_options})
hidden_states = out["img"]
context = out["txt"]
img_ids = out["img_ids"]
txt_ids = out["txt_ids"]
ids = torch.cat((txt_ids, img_ids), dim=1)
image_rotary_emb = self.pe_embedder(ids).contiguous()
del ids, txt_ids, img_ids
transformer_options["total_blocks"] = len(self.transformer_blocks)
transformer_options["block_type"] = "double"
for i, block in enumerate(self.transformer_blocks):
transformer_options["block_index"] = i
if ("double_block", i) in blocks_replace:
def block_wrap(args):
out = {}
out["txt"], out["img"] = block(hidden_states=args["img"], encoder_hidden_states=args["txt"], encoder_hidden_states_mask=attention_mask, temb=args["vec"], image_rotary_emb=args["pe"], transformer_options=args["transformer_options"])
return out
out = blocks_replace[("double_block", i)]({"img": hidden_states, "txt": context, "vec": temb, "pe": image_rotary_emb, "transformer_options": transformer_options}, {"original_block": block_wrap})
hidden_states = out["img"]
context = out["txt"]
else:
context, hidden_states = block(
hidden_states=hidden_states,
encoder_hidden_states=context,
encoder_hidden_states_mask=attention_mask,
temb=temb,
image_rotary_emb=image_rotary_emb,
transformer_options=transformer_options,
)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": hidden_states, "txt": context, "x": x, "block_index": i, "transformer_options": transformer_options})
hidden_states = out["img"]
context = out["txt"]
if control is not None: # Controlnet
control_i = control.get("input")
if i < len(control_i):
add = control_i[i]
if add is not None:
hidden_states[:, :add.shape[1]] += add
hidden_states = self.norm_out(hidden_states, temb)
hidden_states = self.proj_out(hidden_states)
hidden_states = hidden_states[:, :num_embeds]
h, w = orig_shape
return hidden_states.reshape(x.shape[0], h, w, self.out_channels).movedim(-1, 1)
comfy/ldm/mage_flow/vae.py
+477
View File
@@ -0,0 +1,477 @@
# Mage-VAE (https://github.com/microsoft/Mage) (MIT)
# Symmetric one-step diffusion codec: DConvEncoder (image -> 128ch latent) and
# DConvDenoiser + CoD Decoder (latent -> image). 16x downsample, latents in the
# Flux.2-VAE-anchored space (no patch packing, no BN normalization).
# Both encode and decode are single forward passes at t=0.
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.ops
from comfy.ldm.modules.diffusionmodules.model import vae_attention
ops = comfy.ops.disable_weight_init
def nonlinearity(x):
return torch.nn.functional.silu(x)
def Normalize(in_channels):
return ops.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
def modulate(x, shift, scale):
if x.dim() == 4:
b, c = x.shape[:2]
return x * (1 + scale.view(b, c, 1, 1)) + shift.view(b, c, 1, 1)
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
class LayerNorm2d(ops.LayerNorm):
def __init__(self, num_channels, eps=1e-6, affine=True):
super().__init__(num_channels, eps=eps, elementwise_affine=affine)
def forward(self, x):
x = x.permute(0, 2, 3, 1).contiguous()
x = super().forward(x)
return x.permute(0, 3, 1, 2).contiguous()
class TimestepEmbedder(nn.Module):
"""DConv-style timestep MLP (max_period=10000, freq_size=256)."""
def __init__(self, hidden_size, frequency_embedding_size=256):
super().__init__()
self.mlp = nn.Sequential(
ops.Linear(frequency_embedding_size, hidden_size, bias=True),
nn.SiLU(),
ops.Linear(hidden_size, hidden_size, bias=True),
)
self.frequency_embedding_size = frequency_embedding_size
@staticmethod
def timestep_embedding(t, dim, max_period=10000):
half = dim // 2
freqs = torch.exp(
-math.log(max_period) * torch.arange(0, half, dtype=torch.float32) / half
).to(t.device)
args = t[:, None].float() * freqs[None]
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
if dim % 2:
emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
return emb
def forward(self, t, dtype):
emb = self.timestep_embedding(t, self.frequency_embedding_size)
return self.mlp(emb.to(dtype))
class BottleneckPatchEmbed(nn.Module):
"""Image patch embed concatenated with a per-patch conditioning vector."""
def __init__(self, patch_size=16, in_chans=3, pca_dim=128, embed_dim=384, bias=True):
super().__init__()
self.proj1 = ops.Conv2d(in_chans, pca_dim, kernel_size=patch_size, stride=patch_size, bias=False)
self.proj2 = ops.Conv2d(pca_dim + embed_dim, embed_dim, kernel_size=1, bias=bias)
def forward(self, x, cond):
return self.proj2(torch.cat([self.proj1(x), cond], dim=1))
class DiCoBlock(nn.Module):
"""DConv block with adaLN modulation."""
def __init__(self, hidden_size, mlp_ratio=4.0):
super().__init__()
self.conv1 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
self.conv2 = ops.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
self.conv3 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
self.ca = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
ops.Conv2d(hidden_size, hidden_size, 1, bias=True),
nn.Sigmoid(),
)
ffn = int(mlp_ratio * hidden_size)
self.conv4 = ops.Conv2d(hidden_size, ffn, 1, bias=True)
self.conv5 = ops.Conv2d(ffn, hidden_size, 1, bias=True)
self.norm1 = LayerNorm2d(hidden_size, affine=False)
self.norm2 = LayerNorm2d(hidden_size, affine=False)
self.adaLN_modulation = nn.Sequential(
nn.SiLU(),
ops.Linear(hidden_size, 6 * hidden_size, bias=True),
)
def forward(self, inp, c):
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
x = modulate(self.norm1(inp), shift_msa, scale_msa)
x = F.gelu(self.conv2(self.conv1(x)))
x = x * self.ca(x)
x = self.conv3(x)
x = inp + gate_msa[..., None, None] * x
x = x + gate_mlp[..., None, None] * self.conv5(
F.gelu(self.conv4(modulate(self.norm2(x), shift_mlp, scale_mlp)))
)
return x
class EncoderDiCoBlock(nn.Module):
"""DiCoBlock without adaLN, for the encoder head pathway."""
def __init__(self, hidden_size, mlp_ratio=4.0):
super().__init__()
self.conv1 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
self.conv2 = ops.Conv2d(hidden_size, hidden_size, 3, padding=1, groups=hidden_size, bias=True)
self.conv3 = ops.Conv2d(hidden_size, hidden_size, 1, bias=True)
self.ca = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
ops.Conv2d(hidden_size, hidden_size, 1, bias=True),
nn.Sigmoid(),
)
ffn = int(mlp_ratio * hidden_size)
self.conv4 = ops.Conv2d(hidden_size, ffn, 1, bias=True)
self.conv5 = ops.Conv2d(ffn, hidden_size, 1, bias=True)
self.norm1 = LayerNorm2d(hidden_size)
self.norm2 = LayerNorm2d(hidden_size)
def forward(self, inp):
x = self.norm1(inp)
x = F.gelu(self.conv2(self.conv1(x)))
x = x * self.ca(x)
x = self.conv3(x)
x = inp + x
return x + self.conv5(F.gelu(self.conv4(self.norm2(x))))
class NerfEmbedder(nn.Module):
"""Patch-position embedder used by the DConv decoder x-pathway."""
def __init__(self, in_channels, hidden_size_input, max_freqs=8):
super().__init__()
self.max_freqs = max_freqs
self.embedder = nn.Sequential(
ops.Linear(in_channels + max_freqs ** 2, hidden_size_input, bias=True),
)
def fetch_pos(self, patch_size, device, dtype):
pos = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
pos_y, pos_x = torch.meshgrid(pos, pos, indexing="ij")
pos_x = pos_x.reshape(-1, 1, 1)
pos_y = pos_y.reshape(-1, 1, 1)
freqs = torch.linspace(0, self.max_freqs, self.max_freqs, dtype=dtype, device=device)
fx = freqs[None, :, None]
fy = freqs[None, None, :]
coeffs = (1 + fx * fy) ** -1
dct_x = torch.cos(pos_x * fx * torch.pi)
dct_y = torch.cos(pos_y * fy * torch.pi)
return (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
def forward(self, x):
B, P2, _ = x.shape
ps = int(P2 ** 0.5)
dct = self.fetch_pos(ps, x.device, x.dtype).expand(B, -1, -1)
return self.embedder(torch.cat([x, dct], dim=-1))
class NerfFinalLayer(nn.Module):
def __init__(self, hidden_size, out_channels):
super().__init__()
self.norm = ops.RMSNorm(hidden_size, eps=1e-6)
self.linear = ops.Linear(hidden_size, out_channels, bias=True)
def forward(self, x):
return self.linear(self.norm(x))
class MLPResBlock(nn.Module):
def __init__(self, channels):
super().__init__()
self.in_ln = ops.LayerNorm(channels, eps=1e-6)
self.mlp = nn.Sequential(
ops.Linear(channels, channels, bias=True),
nn.SiLU(),
ops.Linear(channels, channels, bias=True),
)
self.adaLN_modulation = nn.Sequential(
nn.SiLU(),
ops.Linear(channels, 3 * channels, bias=True),
)
def forward(self, x, y):
shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
h = self.in_ln(x) * (1 + scale) + shift
return x + gate * self.mlp(h)
class SimpleMLPAdaLN(nn.Module):
"""Final small MLP that maps NerfEmbedder features to per-patch RGB."""
def __init__(self, in_channels, model_channels, out_channels, z_channels, num_res_blocks, patch_size):
super().__init__()
self.in_channels = in_channels
self.model_channels = model_channels
self.out_channels = out_channels
self.num_res_blocks = num_res_blocks
self.patch_size = patch_size
self.cond_embed = ops.Linear(z_channels, patch_size ** 2 * model_channels)
self.input_proj = ops.Linear(in_channels, model_channels)
self.res_blocks = nn.ModuleList(MLPResBlock(model_channels) for _ in range(num_res_blocks))
def forward(self, x, c):
x = self.input_proj(x)
c = self.cond_embed(c).reshape(c.shape[0], self.patch_size ** 2, -1)
for block in self.res_blocks:
x = block(x, c)
return x
class ResnetBlock(nn.Module):
"""GroupNorm + Conv ResBlock used by the CoD Decoder."""
def __init__(self, *, in_channels, out_channels=None):
super().__init__()
out_channels = out_channels or in_channels
self.in_channels = in_channels
self.out_channels = out_channels
self.norm1 = Normalize(in_channels)
self.conv1 = ops.Conv2d(in_channels, out_channels, 3, padding=1)
self.norm2 = Normalize(out_channels)
self.conv2 = ops.Conv2d(out_channels, out_channels, 3, padding=1)
if in_channels != out_channels:
self.nin_shortcut = ops.Conv2d(in_channels, out_channels, 1)
def forward(self, x):
h = self.conv1(nonlinearity(self.norm1(x)))
h = self.conv2(nonlinearity(self.norm2(h)))
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return x + h
class AttnBlock(nn.Module):
"""Patched (windowed) self-attention used by the CoD Decoder."""
def __init__(self, in_channels, patch_size=32):
super().__init__()
self.in_channels = in_channels
self.patch_size = patch_size
self.norm = Normalize(in_channels)
self.q = ops.Conv2d(in_channels, in_channels, 1)
self.k = ops.Conv2d(in_channels, in_channels, 1)
self.v = ops.Conv2d(in_channels, in_channels, 1)
self.proj_out = ops.Conv2d(in_channels, in_channels, 1)
# VAE attention selection: full-precision backends only (no sage/quantized attention)
self.optimized_attention = vae_attention()
def forward(self, x):
h_ = self.norm(x)
Q = self.q(h_)
K = self.k(h_)
V = self.v(h_)
d = self.patch_size
b, c, H, W = Q.shape
pad_h = (d - H % d) % d
pad_w = (d - W % d) % d
if pad_h or pad_w:
Q = F.pad(Q, (0, pad_w, 0, pad_h), mode="replicate")
K = F.pad(K, (0, pad_w, 0, pad_h), mode="replicate")
V = F.pad(V, (0, pad_w, 0, pad_h), mode="replicate")
_, _, H_pad, W_pad = Q.shape
nph, npw = H_pad // d, W_pad // d
np_ = nph * npw
def to_patches(t):
return (t.reshape(b, c, nph, d, npw, d)
.permute(0, 2, 4, 1, 3, 5)
.reshape(b * np_, c, d * d))
# [b*np, c, d*d]: attention over the d*d spatial positions of each window
Q = to_patches(Q)
K = to_patches(K)
V = to_patches(V)
h_ = self.optimized_attention(Q, K, V)
h_ = h_.reshape(b, nph, npw, c, d, d).permute(0, 3, 1, 4, 2, 5).reshape(b, c, H_pad, W_pad)
if pad_h or pad_w:
h_ = h_[:, :, :H, :W]
return x + self.proj_out(h_)
class CoDDecoder(nn.Module):
"""CoD Decoder: latent -> conditioning features for the denoiser (ds=16, light)."""
def __init__(self, out_ch=384, z_ch=128):
super().__init__()
self.conv_in = ops.Conv2d(z_ch, out_ch, kernel_size=3, stride=1, padding=1)
self.block = nn.Sequential(
ResnetBlock(in_channels=out_ch, out_channels=out_ch),
AttnBlock(out_ch, patch_size=32),
ResnetBlock(in_channels=out_ch, out_channels=out_ch),
AttnBlock(out_ch, patch_size=32),
ResnetBlock(in_channels=out_ch, out_channels=out_ch),
)
self.norm_out = Normalize(out_ch)
self.conv_out = ops.Conv2d(out_ch, out_ch, kernel_size=3, stride=1, padding=1)
self.ada = nn.Identity()
def forward(self, z):
h = self.block(self.conv_in(z))
h = self.conv_out(nonlinearity(self.norm_out(h)))
return self.ada(h)
class DConvEncoder(nn.Module):
"""DConvEncoder: image -> packed (mean, logvar) latent."""
def __init__(
self,
z_ch=128,
hidden_size=384,
num_blocks=21,
patch_size=16,
mlp_ratio=4.0,
head_size=768,
num_head_blocks=2,
out_ch_mult=2,
):
super().__init__()
self.z_ch = z_ch
self.patch_size = patch_size
self.patch_cond_embed = ops.Conv2d(3, head_size, kernel_size=patch_size, stride=patch_size, bias=True)
self.head_blocks = nn.ModuleList([
EncoderDiCoBlock(head_size, mlp_ratio=mlp_ratio) for _ in range(num_head_blocks)
])
self.proj_down = ops.Conv2d(head_size, hidden_size, kernel_size=1, bias=True)
self.z_proj = ops.Conv2d(z_ch, hidden_size, kernel_size=1, bias=True)
self.fuse_proj = ops.Conv2d(hidden_size * 2, hidden_size, kernel_size=1, bias=True)
self.t_embedder = TimestepEmbedder(hidden_size)
self.blocks = nn.ModuleList([
DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_blocks)
])
self.norm_out = LayerNorm2d(hidden_size)
self.proj_out = ops.Conv2d(hidden_size, z_ch * out_ch_mult, kernel_size=1, bias=True)
def forward_pred(self, z_t, t, y):
cond = self.patch_cond_embed(y)
for block in self.head_blocks:
cond = block(cond)
cond = self.proj_down(cond)
s = self.fuse_proj(torch.cat([cond, self.z_proj(z_t)], dim=1))
c = self.t_embedder(t.view(-1), y.dtype)
for block in self.blocks:
s = block(s, c)
return self.proj_out(self.norm_out(s))
class YEmbedder(nn.Module):
"""Holds only the CoD decoder (the original Flux2-VAE encoder side is dropped at load)."""
def __init__(self, ch=384, z_ch=128):
super().__init__()
self.decoder = CoDDecoder(out_ch=ch, z_ch=z_ch)
class DConvDenoiser(nn.Module):
"""One-step DConv denoiser: latent (via cond) + zero noise -> reconstructed image."""
def __init__(
self,
patch_size=16,
in_channels=3,
hidden_size=384,
hidden_size_x=32,
mlp_ratio=4.0,
num_blocks=24,
num_cond_blocks=21,
bottleneck_dim=128,
):
super().__init__()
self.in_channels = in_channels
self.patch_size = patch_size
self.hidden_size = hidden_size
self.num_cond_blocks = num_cond_blocks
self.t_embedder = TimestepEmbedder(hidden_size)
self.y_embedder_x = ops.Conv2d(hidden_size, hidden_size_x * patch_size ** 2, 1, 1, 0)
self.x_embedder = NerfEmbedder(in_channels + hidden_size_x, hidden_size_x, max_freqs=8)
self.s_embedder = BottleneckPatchEmbed(patch_size, in_channels, bottleneck_dim, hidden_size, bias=True)
self.blocks = nn.ModuleList([
DiCoBlock(hidden_size, mlp_ratio=mlp_ratio) for _ in range(num_cond_blocks)
])
self.dec_net = SimpleMLPAdaLN(
in_channels=hidden_size_x,
model_channels=hidden_size_x,
out_channels=in_channels,
z_channels=hidden_size,
num_res_blocks=num_blocks - num_cond_blocks,
patch_size=patch_size,
)
self.final_layer = NerfFinalLayer(hidden_size_x, in_channels)
self.y_embedder = YEmbedder(ch=hidden_size, z_ch=bottleneck_dim)
def forward(self, x, t, cond):
b, _, h, w = x.shape
c = self.t_embedder(t.view(-1), x.dtype)
s = self.s_embedder(x, cond)
for block in self.blocks:
s = block(s, c)
length = s.shape[-2] * s.shape[-1]
s = s.permute(0, 2, 3, 1).reshape(-1, self.hidden_size)
x = torch.nn.functional.unfold(x, kernel_size=self.patch_size, stride=self.patch_size)
x = torch.cat([x, self.y_embedder_x(cond).flatten(2)], dim=1)
x = x.reshape(b, -1, self.patch_size ** 2, length).permute(0, 3, 2, 1).flatten(0, 1)
x = self.x_embedder(x)
x = self.dec_net(x, s)
x = self.final_layer(x)
x = x.transpose(1, 2).reshape(b, length, -1)
return torch.nn.functional.fold(
x.transpose(1, 2).contiguous(), (h, w),
kernel_size=self.patch_size, stride=self.patch_size,
)
class MageVAE(nn.Module):
"""
Encode: DConvEncoder (one-step at t=0) -> posterior mean [B, 128, H/16, W/16]
Decode: DConvDenoiser + CoD Decoder -> image [B, 3, H, W] in [-1, 1]
"""
latent_channels = 128
downsample_factor = 16
def __init__(self):
super().__init__()
self.dconv_encoder = DConvEncoder()
self.decoder_model = DConvDenoiser()
def encode(self, x):
B, _, H, W = x.shape
ps = self.dconv_encoder.patch_size
z_t = torch.zeros(B, self.dconv_encoder.z_ch, H // ps, W // ps, device=x.device, dtype=x.dtype)
t = torch.zeros(B, device=x.device, dtype=x.dtype)
out = self.dconv_encoder.forward_pred(z_t, t, x)
return out[:, : self.latent_channels] # posterior mean (sample_posterior=False)
def decode(self, z):
cond = self.decoder_model.y_embedder.decoder(z)
B = z.shape[0]
H = z.shape[2] * self.downsample_factor
W = z.shape[3] * self.downsample_factor
noise = torch.zeros(B, 3, H, W, device=z.device, dtype=z.dtype)
t = torch.zeros(B, device=z.device, dtype=z.dtype)
return self.decoder_model.forward(noise, t, cond)
comfy/ldm/minimax/audio_vae.py
+443
View File
@@ -0,0 +1,443 @@
# MiniMax H3 audio VAE: DAC-lineage waveform encoder + BigVGAN decoder.
# Weight-norm parametrizations are folded into plain conv weights, so this
# module uses ordinary ops.Conv1d / ops.ConvTranspose1d and loads the converted
# checkpoint (plain "*.weight" tensors) with strict=True.
#
# Lineage / licenses of the reference implementation:
# DAC encoder: descript-audio-codec (MIT)
# BigVGAN decoder: NVIDIA BigVGAN (MIT), adapted from hifi-gan (MIT)
# Alias-free ops: junjun3518/alias-free-torch (Apache-2.0), julius (MIT)
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.ops
ops = comfy.ops.disable_weight_init
# Snake activations
def snake(x, alpha, beta):
# x + 1/beta * sin^2(alpha * x)
t = torch.sin(alpha * x)
return t.mul_(t).mul_((beta + 1e-9).reciprocal()).add_(x)
class Snake1d(nn.Module):
"""Snake activation with per-channel alpha (encoder side)."""
def __init__(self, channels):
super().__init__()
self.alpha = nn.Parameter(torch.empty(1, channels, 1))
def forward(self, x):
alpha = comfy.ops.cast_to_input(self.alpha, x)
return snake(x, alpha, alpha)
class SnakeBeta(nn.Module):
"""SnakeBeta := x + 1/beta * sin^2(alpha * x); alpha/beta stored in log scale."""
def __init__(self, in_features):
super().__init__()
self.alpha = nn.Parameter(torch.empty(in_features))
self.beta = nn.Parameter(torch.empty(in_features))
def forward(self, x):
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)
# Alias-free (anti-aliased) activation: kaiser-windowed sinc resampling
def kaiser_sinc_filter1d(cutoff, half_width, kernel_size):
# returns filter [1, 1, kernel_size]
even = kernel_size % 2 == 0
half_size = kernel_size // 2
# kaiser window design
delta_f = 4 * half_width
A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
if A > 50.0:
beta = 0.1102 * (A - 8.7)
elif A >= 21.0:
beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
else:
beta = 0.0
window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
if even:
time = torch.arange(-half_size, half_size) + 0.5
else:
time = torch.arange(kernel_size) - half_size
filter_ = 2 * cutoff * window * torch.sinc(2 * cutoff * time)
# Normalize filter to have sum = 1, otherwise there is a small leakage of
# the constant component in the input signal.
filter_ /= filter_.sum()
return filter_.view(1, 1, kernel_size)
class UpSample1d(nn.Module):
def __init__(self, ratio=2, kernel_size=12):
super().__init__()
self.ratio = ratio
self.stride = ratio
self.pad = kernel_size // ratio - 1
self.pad_left = self.pad * ratio + (kernel_size - ratio) // 2
self.pad_right = self.pad * ratio + (kernel_size - ratio + 1) // 2
self.register_buffer(
"filter",
kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=kernel_size),
)
def forward(self, x):
_, C, _ = x.shape
x = F.pad(x, (self.pad, self.pad), mode="replicate")
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
class LowPassFilter1d(nn.Module):
def __init__(self, cutoff=0.5, half_width=0.6, stride=1, kernel_size=12):
super().__init__()
self.pad_left = kernel_size // 2 - int(kernel_size % 2 == 0)
self.pad_right = kernel_size // 2
self.stride = stride
self.register_buffer("filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size))
def forward(self, x):
_, C, _ = x.shape
x = F.pad(x, (self.pad_left, self.pad_right), mode="replicate")
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):
def __init__(self, ratio=2, kernel_size=12):
super().__init__()
self.ratio = ratio
self.kernel_size = kernel_size
self.lowpass = LowPassFilter1d(
cutoff=0.5 / ratio,
half_width=0.6 / ratio,
stride=ratio,
kernel_size=self.kernel_size,
)
def forward(self, x):
return self.lowpass(x)
class Activation1d(nn.Module):
"""upsample x2 -> pointwise activation -> downsample x2 (anti-aliased)."""
def __init__(self, activation, up_ratio=2, down_ratio=2, up_kernel_size=12, down_kernel_size=12):
super().__init__()
self.act = activation
self.upsample = UpSample1d(up_ratio, up_kernel_size)
self.downsample = DownSample1d(down_ratio, down_kernel_size)
def forward(self, x):
x = self.upsample(x)
x = self.act(x)
x = self.downsample(x)
return x
# DAC encoder
class ResidualUnit(nn.Module):
def __init__(self, dim=16, dilation=1):
super().__init__()
pad = ((7 - 1) * dilation) // 2
self.block = nn.Sequential(
Snake1d(dim),
ops.Conv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
Snake1d(dim),
ops.Conv1d(dim, dim, kernel_size=1),
)
def forward(self, x):
y = self.block(x)
pad = (x.shape[-1] - y.shape[-1]) // 2
if pad > 0:
x = x[..., pad:-pad]
return y.add_(x)
class EncoderBlock(nn.Module):
def __init__(self, dim=16, stride=1):
super().__init__()
self.block = nn.Sequential(
ResidualUnit(dim // 2, dilation=1),
ResidualUnit(dim // 2, dilation=3),
ResidualUnit(dim // 2, dilation=9),
Snake1d(dim // 2),
ops.Conv1d(
dim // 2,
dim,
kernel_size=2 * stride,
stride=stride,
padding=math.ceil(stride / 2),
),
)
def forward(self, x):
return self.block(x)
class Encoder(nn.Module):
def __init__(self, d_model=64, strides=(2, 4, 4, 5, 5), d_latent=2048):
super().__init__()
block = [ops.Conv1d(1, d_model, kernel_size=7, padding=3)]
for stride in strides:
d_model *= 2
block += [EncoderBlock(d_model, stride=stride)]
block += [
Snake1d(d_model),
ops.Conv1d(d_model, d_latent, kernel_size=3, padding=1),
]
self.block = nn.Sequential(*block)
def forward(self, x):
return self.block(x)
# Attention projection (encoder posterior head)
class GeGluMlp(nn.Module):
def __init__(self, in_features, hidden_features):
super().__init__()
self.norm = ops.LayerNorm(in_features)
self.act = nn.GELU(approximate="tanh")
self.w0 = ops.Linear(in_features, hidden_features)
self.w1 = ops.Linear(in_features, hidden_features)
self.w2 = ops.Linear(hidden_features, in_features)
def forward(self, x):
x = self.norm(x)
return self.w2(self.act(self.w0(x)).mul_(self.w1(x)))
class CausalAttention(nn.Module):
def __init__(self, in_dim, out_dim, num_heads):
super().__init__()
self.head_dim = in_dim // num_heads
self.num_heads = num_heads
self.out_dim = out_dim
self.qkv = ops.Linear(in_dim, in_dim * 3, bias=False)
self.q_bias = nn.Parameter(torch.empty(in_dim))
self.v_bias = nn.Parameter(torch.empty(in_dim))
self.register_buffer("zero_k_bias", torch.empty(in_dim))
self.proj = ops.Linear(out_dim, out_dim)
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=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)
# mean over heads then pool down to the latent width (in_dim >> out_dim)
x = comfy.ops.scaled_dot_product_attention(q, k, v, is_causal=True)
x = F.adaptive_avg_pool1d(torch.mean(x, dim=1), self.out_dim)
return self.proj(x)
class AttnProjection(nn.Module):
def __init__(self, in_dim, out_dim, num_heads, mlp_ratio=2):
super().__init__()
self.norm1 = ops.LayerNorm(in_dim)
self.attn = CausalAttention(in_dim, out_dim, num_heads)
self.proj = ops.Linear(in_dim, out_dim)
self.norm3 = ops.LayerNorm(in_dim)
self.norm2 = ops.LayerNorm(out_dim)
hidden_dim = int(out_dim * mlp_ratio)
self.mlp = GeGluMlp(in_features=out_dim, hidden_features=hidden_dim)
def forward(self, x):
# x: [B, T, in_dim]
x = self.proj(self.norm3(x)).add_(self.attn(self.norm1(x)))
return x.add_(self.mlp(self.norm2(x)))
# BigVGAN decoder
def get_padding(kernel_size, dilation=1):
return int((kernel_size * dilation - dilation) / 2)
class AMPBlock1(nn.Module):
def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
super().__init__()
self.convs1 = nn.ModuleList(
[
ops.Conv1d(channels, channels, kernel_size, stride=1, dilation=d, padding=get_padding(kernel_size, d))
for d in dilation
]
)
self.convs2 = nn.ModuleList(
[
ops.Conv1d(channels, channels, kernel_size, stride=1, dilation=1, padding=get_padding(kernel_size, 1))
for _ in range(len(dilation))
]
)
self.num_layers = len(self.convs1) + len(self.convs2)
self.activations = nn.ModuleList(
[Activation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
)
def forward(self, x):
acts1, acts2 = self.activations[::2], self.activations[1::2]
for c1, c2, a1, a2 in zip(self.convs1, self.convs2, acts1, acts2):
xt = a1(x)
xt = c1(xt)
xt = a2(xt)
xt = c2(xt)
x = xt.add_(x)
return x
class BigVGAN(nn.Module):
"""BigVGAN vocoder (MiniMax H3 32 kHz configuration).
use_bias_at_final=False, use_tanh_at_final=False (output clamped to [-1, 1]).
"""
def __init__(
self,
num_mels=2048,
upsample_initial_channel=1024,
upsample_rates=(5, 5, 2, 2, 2, 2, 2),
upsample_kernel_sizes=(9, 9, 4, 4, 4, 4, 4),
resblock_kernel_sizes=(3, 7, 11),
resblock_dilation_sizes=((1, 3, 5), (1, 3, 5), (1, 3, 5)),
):
super().__init__()
self.num_kernels = len(resblock_kernel_sizes)
self.num_upsamples = len(upsample_rates)
self.conv_pre = ops.Conv1d(num_mels, upsample_initial_channel, 7, 1, padding=3)
self.ups = nn.ModuleList()
for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)):
self.ups.append(
nn.ModuleList(
[
ops.ConvTranspose1d(
upsample_initial_channel // (2 ** i),
upsample_initial_channel // (2 ** (i + 1)),
k,
u,
padding=(k - u) // 2,
)
]
)
)
self.resblocks = nn.ModuleList()
for i in range(len(self.ups)):
ch = upsample_initial_channel // (2 ** (i + 1))
for k, d in zip(resblock_kernel_sizes, resblock_dilation_sizes):
self.resblocks.append(AMPBlock1(ch, k, d))
self.activation_post = Activation1d(activation=SnakeBeta(ch))
self.conv_post = ops.Conv1d(ch, 1, 7, 1, padding=3, bias=False)
def forward(self, x):
x = self.conv_pre(x)
for i in range(self.num_upsamples):
for i_up in range(len(self.ups[i])):
x = self.ups[i][i_up](x)
xs = None
for j in range(self.num_kernels):
if xs is None:
xs = self.resblocks[i * self.num_kernels + j](x)
else:
xs += self.resblocks[i * self.num_kernels + j](x)
x = xs.div_(self.num_kernels)
x = self.activation_post(x)
return self.conv_post(x).clamp_(-1.0, 1.0)
# Top-level VAE
class MiniMaxH3AudioVAE(nn.Module):
"""MiniMax H3 stereo audio VAE at 32 kHz.
Latents are [B, 32, 2, T]: 32 channels, 2 stereo channels, T frames at
40 latent frames per second (800 audio samples per latent frame). The
stereo channels are processed independently by the mono encoder/decoder.
Latents are normalized with the stored per-channel latents_mean/std.
"""
def __init__(
self,
encoder_dim=64,
encoder_rates=(2, 4, 4, 5, 5),
latent_dim=2048,
decoder_dim=1024,
vae_latent_channels=32,
):
super().__init__()
self.sample_rate = 32000
self.hop_length = 1
for r in encoder_rates:
self.hop_length *= r
self.samples_per_latent = self.hop_length # 800
self.latents_per_second = self.sample_rate // self.hop_length # 40
self.output_sample_rate = self.sample_rate # read by LTXVAudioVAEDecode
self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
self.pre_block = AttnProjection(latent_dim, vae_latent_channels, num_heads=8)
self.mean_proj = ops.Conv1d(vae_latent_channels, vae_latent_channels, 1)
# logs_proj exists in the checkpoint but is unused at inference
# (encode returns the posterior mean, no sampling).
self.logs_proj = ops.Conv1d(vae_latent_channels, vae_latent_channels, 1)
self.dec_in_proj = ops.Conv1d(vae_latent_channels, latent_dim, 1)
self.decoder = BigVGAN(num_mels=latent_dim, upsample_initial_channel=decoder_dim)
self.register_buffer("latents_mean", torch.empty(vae_latent_channels))
self.register_buffer("latents_std", torch.empty(vae_latent_channels))
def decode(self, z):
"""Decode normalized latents [B, 32, 2, T] to stereo waveforms [B, 2, L] at 32 kHz."""
b, c, s, t = z.shape
z = z.permute(0, 2, 1, 3).reshape(b * s, c, t)
mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
z = z * std + mean
x = self.dec_in_proj(z)
x = self.decoder(x) # [b * s, 1, L], already clamped to [-1, 1]
return x.reshape(b, s, -1)
def encode(self, waveform):
"""Encode stereo waveforms [B, 2, L] at 32 kHz (in [-1, 1]) to normalized latents [B, 32, 2, T].
L is right-padded with zeros to a multiple of 800 samples; the returned
posterior mean is used directly (no sampling).
"""
b, s, length = waveform.shape
right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
waveform = F.pad(waveform, (0, right_pad))
x = waveform.reshape(b * s, 1, -1)
x = self.encoder(x) # [b * s, latent_dim, T]
x = self.pre_block(x.transpose(1, 2)).transpose(1, 2) # [b * s, 32, T]
z = self.mean_proj(x)
mean = self.latents_mean.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
std = self.latents_std.view(1, -1, 1).to(device=z.device, dtype=z.dtype)
z = (z - mean) / std
return z.reshape(b, s, z.shape[1], z.shape[2]).permute(0, 2, 1, 3)
comfy/ldm/minimax/controlnet.py
+85
View File
@@ -0,0 +1,85 @@
"""MiniMax H3 Fun ControlNet-Union model patch."""
import torch
import torch.nn as nn
import comfy.ldm.common_dit
from .model import DiTBlock, patchify_video
class ControlDiTBlock(DiTBlock):
def __init__(self, hidden, heads, head_dim, ffn, t_dim, eps, qk_eps, first_block=False,
apply_silu=True, adaln_dtype=None, dtype=None, device=None, operations=None):
super().__init__(hidden, heads, head_dim, ffn, t_dim, eps, qk_eps, apply_silu=apply_silu,
adaln_dtype=adaln_dtype, dtype=dtype, device=device, operations=operations)
if first_block:
self.before_proj = operations.Linear(hidden, hidden, bias=True, dtype=dtype, device=device)
self.after_proj = operations.Linear(hidden, hidden, bias=True, dtype=dtype, device=device)
class MiniMaxH3FunControl(torch.nn.Module):
def __init__(self, control_in_dim=49, injection_layers=(0, 10, 20, 30, 40), hidden_size=5376,
num_attention_heads=56, attention_head_dim=128, ffn_hidden_size=14336,
time_embed_dim=2688, patch_size=(1, 2, 2), norm_eps=1e-5, qk_norm_eps=1e-5,
use_adaln_curves=False, dtype=None, device=None, operations=None):
super().__init__()
self.dtype = dtype
self.patch_size = tuple(patch_size)
self.injection_layers = tuple(injection_layers)
if not self.injection_layers or self.injection_layers[0] != 0:
raise ValueError("MiniMax H3 Fun control injection layers must start at layer 0")
if self.injection_layers != tuple(sorted(set(self.injection_layers))):
raise ValueError("MiniMax H3 Fun control injection layers must be unique and increasing")
self.control_in_dim = control_in_dim
patch_dim = control_in_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
self.control_proj_in = operations.Linear(patch_dim, hidden_size, bias=True, dtype=torch.float32, device=device)
self.control_blocks = nn.ModuleList([
ControlDiTBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size,
time_embed_dim, norm_eps, qk_norm_eps, first_block=(i == 0),
apply_silu=not use_adaln_curves,
adaln_dtype=torch.float32 if use_adaln_curves else dtype,
dtype=dtype, device=device, operations=operations)
for i in range(len(self.injection_layers))])
def init_stream(self, h, control_latent, layout, t_emb):
adaln_in = self.control_blocks[0].adaln_proj.linear.in_features
if t_emb.shape[-1] != adaln_in:
raise RuntimeError(
"MiniMax H3 controlnet adaln width {} does not match the base model's timestep embedding width {}: "
"the controlnet and base checkpoint use different adaln forms (curve basis vs full), "
"convert the controlnet to match the base model.".format(adaln_in, t_emb.shape[-1]))
patch_dim = self.control_in_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
control_latent = comfy.ldm.common_dit.pad_to_patch_size(control_latent.to(torch.float32), self.patch_size)
target_rows = patchify_video(control_latent, self.patch_size)
if target_rows.shape[1] < patch_dim:
target_rows = torch.nn.functional.pad(target_rows, (0, patch_dim - target_rows.shape[1]))
elif target_rows.shape[1] > patch_dim:
raise ValueError("MiniMax H3 control input has {} columns but the model patch expects {}".format(target_rows.shape[1], patch_dim))
# keyframe/reference conditioning rows get a zero control row
img_update = layout.img_update.to(h.device)
rows = torch.zeros(img_update.shape[0], patch_dim, dtype=torch.float32, device=h.device)
rows[img_update] = target_rows
c = h.clone()
c[layout.img_pos.to(h.device)] = self.control_proj_in(rows).to(h.dtype)
return self.control_blocks[0].before_proj(c).add_(h)
def step(self, index, c, t_emb, mod_segments, rope_freqs, transformer_options):
block = self.control_blocks[index]
c = DiTBlock.forward(block, c, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)
return c, block.after_proj(c)
def is_minimax_h3_fun_state_dict(state_dict):
required = (
"control_proj_in.weight",
"control_blocks.0.adaln_proj.linear.weight",
"control_blocks.0.after_proj.weight",
"control_blocks.0.before_proj.weight",
"control_blocks.0.attn.qkv_proj.weight",
"control_blocks.0.attn.q_norm.weight",
"control_blocks.0.mlp.fc1.weight",
)
return all(key in state_dict for key in required)
comfy/ldm/minimax/model.py
+758
View File
@@ -0,0 +1,758 @@
"""MiniMax H3 audio-video DiT.
Single-stream packed-token transformer denoising video (24ch, patch 1x2x2) and
stereo audio (32ch, 40 Hz) latents jointly, conditioned on Qwen3-VL layer-50 hidden states.
The packed sequence is:
[text | cond rows | audio | video] for t2va/fl2va
[text | reference blocks | audio | video] for ref2va
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 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
import torch
import torch.nn as nn
import comfy.ldm.common_dit
import comfy.model_management
import comfy.model_prefetch
import comfy.ops
import comfy.patcher_extension
import comfy.quant_ops
from comfy.ldm.modules.attention import AttentionTensorContainer, optimized_attention
FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
FRAME_RESCALE = 5.0 / 3.0
VISUAL_COND_TIMESTEP = 0.999
AUDIO_COND_TIMESTEP = 1.0
def time_shift_sigma(sigma, from_shift, to_shift):
# invert sigma = s*b/(1+(s-1)*b) to the base grid, re-apply the other shift
base = sigma / (from_shift + sigma * (1.0 - from_shift))
return to_shift * base / (1.0 + (to_shift - 1.0) * base)
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
pt, ph, pw = patch_size
t, h, w = t_full // pt, h_full // ph, w_full // pw
x = latent.reshape(b, c, t, pt, h, ph, w, pw)
x = torch.einsum("nctrhpwq->nthwcrpq", x)
return x.reshape(b * t * h * w, c * pt * ph * pw)
def unpatchify_video(rows, t, h, w, c=24, patch_size=(1, 2, 2)):
pt, ph, pw = patch_size
x = rows.reshape(-1, t, h, w, c, pt, ph, pw)
x = torch.einsum("nthwcrpq->nctrhpwq", x)
return x.reshape(-1, c, t * pt, h * ph, w * pw)
def pack_audio(latent):
# [B, C=32, ch=2, T] -> [ch*T, 32] channel-major (ch0 t0..T-1, ch1 t0..T-1)
b, c, ch, t = latent.shape
return latent[0].permute(1, 2, 0).reshape(ch * t, c)
def unpack_audio(rows, ch=2):
t = rows.shape[0] // ch
return rows.reshape(ch, t, rows.shape[-1]).permute(2, 0, 1).unsqueeze(0)
def _axis_from_sqrt_area(dim, patch, sqrt_area):
# linspace((1 - ratio) / 2, (1 + ratio) / 2, dim // patch, endpoint=False) * 32
ratio = dim / sqrt_area
n = dim // patch
return (torch.arange(n, dtype=torch.float64) * (ratio / n) + (1.0 - ratio) / 2.0) * 32.0
def mask_row_values(mask, latent_t, lat_h, lat_w):
# [T, H, W] denoise mask (1 = generate) -> per-2x2-patch-row float in [0, 1],
# None when every row fully generates
m = torch.nn.functional.pad(mask, (0, lat_w - mask.shape[-1], 0, lat_h - mask.shape[-2]), mode="replicate")
m = m.reshape(latent_t, lat_h // 2, 2, lat_w // 2, 2).amax(dim=(2, 4))
values = m.reshape(-1)
if bool((values >= 1.0 - 1e-3).all()):
return None
return values
def _frame_grid(h, w):
# area-normalized (h, w) coordinates of one latent frame's 2x2-patch rows
area = math.sqrt(h * w)
hh, ww = torch.meshgrid(_axis_from_sqrt_area(h, 2, area), _axis_from_sqrt_area(w, 2, area), indexing="ij")
return torch.stack([hh.reshape(-1), ww.reshape(-1)], dim=-1), _axis_from_sqrt_area(w, 2, area)
def _video_t_spans(n):
return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(n)]
def _video_t_grid(n, origin):
# origin + exclusive cumsum
spans = torch.tensor(_video_t_spans(n), dtype=torch.float64)
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)
g[:, 0] = (cursor + torch.arange(t, dtype=torch.float64)).repeat(2)
g[:t, 2] = w_low
g[t:, 2] = w_high
return g
def _video_grid(vt, frame, cursor):
g = torch.empty(vt, frame.shape[0], 3, dtype=torch.float64)
g[:, :, 0] = _video_t_grid(vt, cursor)[:, None]
g[:, :, 1:] = frame[None]
return g.reshape(-1, 3)
class TimeEmbedder(nn.Module):
def __init__(self, freq_dim, hidden, out, dtype=None, device=None, operations=None):
super().__init__()
self.freq_dim = freq_dim
self.proj_in = operations.Linear(freq_dim, hidden, bias=True, dtype=dtype, device=device)
self.proj_out = operations.Linear(hidden, out, bias=True, dtype=dtype, device=device)
def forward(self, t):
# t: [M] in [0, 1]; fp32 throughout, cos before sin
half = self.freq_dim // 2
freqs = torch.exp(-math.log(10000.0) * torch.arange(half, dtype=torch.float32, device=t.device) / half)
args = t.to(torch.float32)[:, None] * freqs[None]
emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
return self.proj_out(nn.functional.silu(self.proj_in(emb)))
def rope_rotation_table(angles, dtype):
"""[S, rot_dim] pair angles -> [1, S, 1, rot_dim/2, 2, 2] rotation matrices."""
half = angles.shape[-1] // 2
ang = angles[:, :half] # duplicated halves: [:, :half] == [:, half:]
c, s = torch.cos(ang), torch.sin(ang)
table = torch.stack([c, -s, s, c], dim=-1).reshape(1, angles.shape[0], 1, half, 2, 2)
return table.to(dtype)
class Attention(nn.Module):
def __init__(self, hidden, heads, head_dim, eps, dtype=None, device=None, operations=None):
super().__init__()
self.heads = heads
self.head_dim = head_dim
inner = heads * head_dim
self.qkv_proj = operations.Linear(hidden, inner * 3, bias=False, dtype=dtype, device=device)
self.q_norm = operations.RMSNorm(head_dim, eps=eps, dtype=dtype, device=device)
self.k_norm = operations.RMSNorm(head_dim, eps=eps, dtype=dtype, device=device)
self.out_proj = operations.Linear(inner, hidden, bias=False, dtype=dtype, device=device)
def forward(self, x, rope_freqs=None, transformer_options={}):
s = x.shape[0]
q, k, v = self.qkv_proj(x).split(self.heads * self.head_dim, dim=-1)
v = v.view(s, self.heads, self.head_dim)
if rope_freqs is not None:
# fused per-head RMSNorm + partial split-half rope, in place on the qkv buffer
q = q.view(1, s, self.heads, self.head_dim)
k = k.view(1, s, self.heads, self.head_dim)
qw = comfy.model_management.cast_to(self.q_norm.weight, device=x.device)
kw = comfy.model_management.cast_to(self.k_norm.weight, device=x.device)
rot = rope_freqs.shape[-3] * 2
if comfy.model_management.in_training:
q, k = comfy.quant_ops.ck.rms_rope_split_half(
q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot)
else:
comfy.quant_ops.ck.rms_rope_split_half_(
q, k, rope_freqs, qw, kw, epsilon=self.q_norm.eps, rot_dim=rot)
q = q[0]
k = k[0]
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))
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))
class MLP(nn.Module):
def __init__(self, hidden, ffn, dtype=None, device=None, operations=None):
super().__init__()
self.fc1 = operations.Linear(hidden, ffn * 2, bias=False, dtype=dtype, device=device)
self.fc2 = operations.Linear(ffn, hidden, bias=False, dtype=dtype, device=device)
def forward(self, x):
return comfy.ops.linear_input_act(self.fc2, self.fc1(x), "swiglu")
class AdalnProj(nn.Module):
def __init__(self, t_dim, hidden, expand, modalities, apply_silu=True,
dtype=None, device=None, operations=None):
super().__init__()
self.expand = expand
self.modalities = modalities
self.hidden = hidden
self.apply_silu = apply_silu
self.linear = operations.Linear(t_dim, expand * hidden * modalities, bias=True, dtype=dtype, device=device)
def forward(self, t_emb):
# [M, t_dim] -> expand tensors of [M*modalities, hidden]
x = self.linear(nn.functional.silu(t_emb) if self.apply_silu else t_emb)
x = x.view(x.shape[0] * self.modalities, self.expand * self.hidden)
return x.chunk(self.expand, dim=-1)
def _mod_row(vecs, row, dtype):
# row is a mod-row index, or a per-token LongTensor of mod-row indices
return vecs[row].to(dtype)
def _mod_scale_shift(h, shift, scale, segments):
# segments: [(start, stop, mod_row)] covering h contiguously.
for a, b, row in segments:
h[a:b].mul_(1.0 + _mod_row(scale, row, h.dtype)).add_(_mod_row(shift, row, h.dtype))
return h
def _mod_gate(x, gate, other, segments):
# other is the fresh attn/mlp output: accumulate the gated residual into the stream in place, one fused kernel per segment
for a, b, row in segments:
x[a:b].addcmul_(other[a:b], _mod_row(gate, row, x.dtype))
return x
class RefinerBlock(nn.Module):
def __init__(self, hidden, heads, head_dim, ffn, eps, qk_eps, dtype=None, device=None, operations=None):
super().__init__()
self.norm1 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
self.norm2 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
self.attn = Attention(hidden, heads, head_dim, qk_eps, dtype=dtype, device=device, operations=operations)
self.mlp = MLP(hidden, ffn, dtype=dtype, device=device, operations=operations)
def forward(self, x, transformer_options={}):
# attn/mlp outputs are fresh: accumulate residuals in place
x = self.attn(self.norm1(x), transformer_options=transformer_options).add_(x)
return self.mlp(self.norm2(x)).add_(x)
class TokenRefiner(nn.Module):
def __init__(self, num_layers, hidden, heads, head_dim, ffn, eps, qk_eps, final_eps,
dtype=None, device=None, operations=None):
super().__init__()
self.blocks = nn.ModuleList([
RefinerBlock(hidden, heads, head_dim, ffn, eps, qk_eps, dtype=dtype, device=device, operations=operations)
for _ in range(num_layers)])
self.final_norm = operations.RMSNorm(hidden, eps=final_eps, dtype=dtype, device=device)
def forward(self, x, transformer_options={}):
for block in self.blocks:
x = block(x, transformer_options=transformer_options)
return self.final_norm(x)
class DiTBlock(nn.Module):
def __init__(self, hidden, heads, head_dim, ffn, t_dim, eps, qk_eps,
apply_silu=True, adaln_dtype=None, dtype=None, device=None, operations=None):
super().__init__()
self.norm1 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
self.norm2 = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
self.attn = Attention(hidden, heads, head_dim, qk_eps, dtype=dtype, device=device, operations=operations)
self.mlp = MLP(hidden, ffn, dtype=dtype, device=device, operations=operations)
self.adaln_proj = AdalnProj(t_dim, hidden, 6, 3, apply_silu=apply_silu,
dtype=adaln_dtype if adaln_dtype is not None else dtype,
device=device, operations=operations)
def forward(self, x, t_emb, mod_segments, rope_freqs, transformer_options={}):
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaln_proj(t_emb)
h = _mod_scale_shift(self.norm1(x), shift_msa, scale_msa, mod_segments)
x = _mod_gate(x, gate_msa, self.attn(h, rope_freqs=rope_freqs, transformer_options=transformer_options), mod_segments)
h = _mod_scale_shift(self.norm2(x), shift_mlp, scale_mlp, mod_segments)
return _mod_gate(x, gate_mlp, self.mlp(h), mod_segments)
class FinalLayer(nn.Module):
def __init__(self, hidden, t_dim, video_dim, audio_dim, eps, apply_silu=True, adaln_dtype=None,
dtype=None, device=None, operations=None):
super().__init__()
self.norm = operations.RMSNorm(hidden, eps=eps, dtype=dtype, device=device)
self.adaln_proj = AdalnProj(t_dim, hidden, 2, 1, apply_silu=apply_silu,
dtype=adaln_dtype if adaln_dtype is not None else dtype,
device=device, operations=operations)
# output heads are the checkpoint's fp32 island; norm/adaln are stored at model dtype
self.video_out = operations.Linear(hidden, video_dim, bias=True, dtype=torch.float32, device=device)
self.audio_out = operations.Linear(hidden, audio_dim, bias=True, dtype=torch.float32, device=device)
def forward(self, x, t_emb, video_seg, audio_seg, sigma, sample_sigmas, shifts):
# video_seg / audio_seg: (start, stop, row) of the target streams, where row
# is a mod-row index or a per-token blend (see _mod_row)
shift, scale = self.adaln_proj(t_emb)
def mod(seg):
a, b, row = seg
return (self.norm(x[a:b]) * (1.0 + _mod_row(scale, row, scale.dtype)) + _mod_row(shift, row, shift.dtype)).to(torch.float32)
n = self.video_out.weight.shape[0] // self.video_out.out_features
if n == 1:
return self.video_out(mod(video_seg)), self.audio_out(mod(audio_seg))
# PDD head bank: row block 0 is a full head, later blocks are offsets from it;
# a step consumes the dt-weighted mean of the heads it spans.
if sample_sigmas is None:
raise ValueError("MiniMax H3 PDD heads need the sampler's sigma schedule")
i = int((sample_sigmas - sigma).abs().argmin())
sigma_next = sample_sigmas[min(i + 1, sample_sigmas.shape[0] - 1)]
start, stop = (round(float(1.0 - time_shift_sigma(s, shifts[0], 1.0)) * n) for s in (sigma, sigma_next))
start = min(start, n - 1)
stop = max(stop, start + 1)
return (_pdd_head(self.video_out, mod(video_seg), n, start, stop, shifts[0]),
_pdd_head(self.audio_out, mod(audio_seg), n, start, stop, shifts[1]))
def _pdd_head(head, h, n, start, stop, flow_shift):
grid = torch.linspace(1.0, 0.0, n + 1, dtype=torch.float64)
dt = (1.0 - flow_shift * grid / (1.0 + (flow_shift - 1.0) * grid)).diff()[start:stop]
w = (dt / dt.sum()).to(h)
with comfy.ops.CastBiasWeightContext(head, h, offloadable=True) as (weight, bias):
rows = weight.reshape(n, -1, weight.shape[1])
brows = bias.reshape(n, -1)
first = max(start, 1)
return nn.functional.linear(h, rows[0] + torch.einsum("n,noi->oi", w[first - start:], rows[first:stop]),
brows[0] + torch.einsum("n,no->o", w[first - start:], brows[first:stop]))
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, w_grid = _frame_grid(latent_h, latent_w)
frame_rows = frame.shape[0]
segments = [("text", text_len)] # (kind, n_rows)
g = torch.zeros(text_len, 3, dtype=torch.float64)
g[:, 0] = torch.arange(text_len, dtype=torch.float64)
pos = [g] # per segment: [n, 3] float64 (t, h, w)
img_pos, img_update = [], []
audio_pos, audio_update = [], []
row = text_len
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:
kind = blk["kind"]
if kind == "image":
r_frame, _ = _frame_grid(blk["latent_h"], blk["latent_w"])
n = r_frame.shape[0]
g = torch.empty(n, 3, dtype=torch.float64)
g[:, 0] = cursor
g[:, 1:] = r_frame
segments.append(("ref_img", n))
pos.append(g)
img_pos.append(torch.arange(row, row + n))
img_update.append(torch.zeros(n, dtype=torch.bool))
row += n
cursor += 1.0
elif kind == "audio":
rt = blk["ref_audio_t"]
if rt > 0:
segments.append(("ref_audio", rt * 2))
pos.append(_audio_grid(cursor, 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
cursor += float(rt)
elif kind in ("video", "video_audio"):
# the block's audio rows pack immediately before its video
# rows, both sharing the cursor origin
rt = blk["ref_audio_t"]
vt = blk["latent_t"]
r_frame, r_w_grid = _frame_grid(blk["latent_h"], blk["latent_w"])
if rt > 0:
segments.append(("ref_audio", rt * 2))
pos.append(_audio_grid(cursor, rt, float(r_w_grid[0]), float(r_w_grid[-1])))
audio_pos.append(torch.arange(row, row + rt * 2))
audio_update.append(torch.zeros(rt * 2, dtype=torch.bool))
row += rt * 2
n = vt * r_frame.shape[0]
segments.append(("ref_img", n))
pos.append(_video_grid(vt, r_frame, cursor))
img_pos.append(torch.arange(row, row + n))
img_update.append(torch.zeros(n, dtype=torch.bool))
row += n
cursor += max(float(rt), sum(_video_t_spans(vt)))
# target audio then target video, always the last two segments
segments.append(("audio", audio_t * 2))
pos.append(_audio_grid(cursor, audio_t, *target_audio_w))
audio_pos.append(torch.arange(row, row + audio_t * 2))
audio_update.append(torch.ones(audio_t * 2, dtype=torch.bool))
row += audio_t * 2
n_video = latent_t * frame_rows
segments.append(("video", n_video))
pos.append(_video_grid(latent_t, frame, cursor))
img_pos.append(torch.arange(row, row + n_video))
img_update.append(torch.ones(n_video, dtype=torch.bool))
row += n_video
self.seq_len = row
self.position_ids = torch.cat(pos) # [S, 3] float64
self.img_pos = torch.cat(img_pos)
self.img_update = torch.cat(img_update)
self.audio_pos = torch.cat(audio_pos)
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 / 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 = []
off = 0
for kind, n in segments:
seg_abs.append((off, off + n, kind))
off += n
self.segments = seg_abs
class MiniMaxH3Model(nn.Module):
def __init__(self, hidden_size=5376, num_layers=50, token_refiner_num_layers=2,
num_attention_heads=56, attention_head_dim=128, ffn_hidden_size=14336,
latents_dim=24, audio_latents_dim=32, patch_size=(1, 2, 2), text_dim=5120,
timestep_input_dim=256, time_embed_hidden_size=5376, time_embed_dim=2688,
rope_inv_freq_len=16, norm_eps=1e-5, qk_norm_eps=1e-5, final_norm_eps=1e-5,
sigma_shift_video=12.0, sigma_shift_audio=3.0,
adaln_curve_grid=None,
image_model=None, dtype=None, device=None, operations=None, **kwargs):
super().__init__()
self.dtype = dtype
self.hidden_size = hidden_size
self.patch_size = tuple(patch_size)
self.latents_dim = latents_dim
self.audio_latents_dim = audio_latents_dim
self.sigma_shift_video = sigma_shift_video
self.sigma_shift_audio = sigma_shift_audio
self.use_adaln_curves = adaln_curve_grid is not None
# curve-form checkpoints replace the time embedder and full-width adaln weights with a small shared basis of the time-embedding curve
curve = {"apply_silu": not self.use_adaln_curves,
"adaln_dtype": torch.float32 if self.use_adaln_curves else dtype}
video_patch_dim = latents_dim * self.patch_size[0] * self.patch_size[1] * self.patch_size[2]
self.video_patch_proj = operations.Linear(video_patch_dim, hidden_size, bias=True, dtype=torch.float32, device=device)
self.audio_patch_proj = operations.Linear(audio_latents_dim, hidden_size, bias=True, dtype=torch.float32, device=device)
self.condition_proj = operations.Linear(text_dim, hidden_size, bias=True, dtype=dtype, device=device)
if self.use_adaln_curves:
self.register_buffer("adaln_t_table", torch.empty(adaln_curve_grid, time_embed_dim, dtype=torch.float32))
else:
self.time_embedder = TimeEmbedder(timestep_input_dim, time_embed_hidden_size, time_embed_dim,
dtype=torch.float32, device=device, operations=operations)
self.rope = nn.Module()
self.rope.register_buffer("inv_freq", torch.empty(rope_inv_freq_len, dtype=torch.float32))
self.token_refiner = TokenRefiner(token_refiner_num_layers, hidden_size, num_attention_heads,
attention_head_dim, ffn_hidden_size, norm_eps, qk_norm_eps,
final_norm_eps, dtype=dtype, device=device, operations=operations)
self.blocks = nn.ModuleList([
DiTBlock(hidden_size, num_attention_heads, attention_head_dim, ffn_hidden_size,
time_embed_dim, norm_eps, qk_norm_eps, **curve, dtype=dtype, device=device, operations=operations)
for _ in range(num_layers)])
self.final_layer = FinalLayer(hidden_size, time_embed_dim, video_patch_dim, audio_latents_dim,
final_norm_eps, **curve, dtype=dtype, device=device, operations=operations)
def preprocess_text_embeds(self, text_states):
"""[B, L, text_dim] Qwen states -> [B, L, hidden] refined text embeds."""
if text_states.shape[-1] == self.hidden_size:
return text_states
return self.token_refiner(self.condition_proj(text_states[0])).unsqueeze(0)
def rope_freqs(self, position_ids, device):
# [S, 3] float64 -> [S, 96] fp32
pos = position_ids.to(torch.float32).to(device)
inv = comfy.model_management.cast_to(self.rope.inv_freq, device=device)
per_axis = pos.unsqueeze(-1) * inv.view(1, 1, -1) # [S, 3, 16]
t_f, h_f, w_f = per_axis.unbind(dim=1)
half = torch.cat((t_f, h_f, w_f), dim=-1) # [S, 48]
return torch.cat((half, half), dim=-1) # [S, 96]
def _cond_video_rows(self, payload, device):
"""Concatenated visual condition rows (normalized latents -> patchified), with condition noise augmentation."""
rows = []
aug = payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP)
seed = int(payload.get("seed", 0))
# every condition intentionally restarts the same RNG stream
for z in payload.get("cond_video_latents", []):
r = patchify_video(z.to(torch.float32), self.patch_size)
if aug < 1.0:
gen = torch.Generator("cpu").manual_seed(seed)
noise = torch.randn(r.shape, generator=gen, dtype=torch.float32)
r = aug * r + (1.0 - aug) * noise.to(r.device)
rows.append(r.to(device))
return torch.cat(rows, dim=0) if rows else None
def _cond_audio_rows(self, payload, device):
rows = []
aug = payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP)
seed = int(payload.get("seed", 0)) + 1
for z in payload.get("cond_audio_latents", []):
r = pack_audio(z.to(torch.float32))
if aug < 1.0:
gen = torch.Generator("cpu").manual_seed(seed)
noise = torch.randn(r.shape, generator=gen, dtype=torch.float32)
r = aug * r + (1.0 - aug) * noise.to(r.device)
rows.append(r.to(device))
return torch.cat(rows, dim=0) if rows else None
def forward(self, x, timestep, context, transformer_options={}, minimax_payload=None, denoise_mask=None, audio_denoise_mask=None, **kwargs):
# 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,
denoise_mask=denoise_mask, audio_denoise_mask=audio_denoise_mask, **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, denoise_mask=None, audio_denoise_mask=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]
video_x = comfy.ldm.common_dit.pad_to_patch_size(video_x, self.patch_size)
if video_x.shape[0] != 1:
raise ValueError("MiniMax H3 supports batch size 1")
payload = minimax_payload or {}
device = video_x.device
dtype = context.dtype # compute dtype
latent_t, lat_h, lat_w = video_x.shape[2], video_x.shape[3], video_x.shape[4]
audio_t = audio_x.shape[-1]
text_len = context.shape[1]
# extra_conds prebuilds the layout once per sampling run
layout = payload.get("layout")
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"))
# model_base passes model_sampling.timestep(sigma) = sigma * 1000
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)
t_v = float(1.0 - sigma_v)
t_a = float(1.0 - time_shift_sigma(sigma_v, shift_v, shift_a))
# distinct timesteps are known analytically: text/pad follow video, cond rows pin near 1
vis_aug = float(payload.get("visual_cond_noise_aug", VISUAL_COND_TIMESTEP))
aud_aug = float(payload.get("audio_cond_noise_aug", AUDIO_COND_TIMESTEP))
seg_t = {"text": t_v, "video": t_v, "audio": t_a,
"cond": max(t_v, vis_aug), "ref_img": max(t_v, vis_aug),
"cond_audio": max(t_a, aud_aug), "ref_audio": max(t_a, aud_aug)}
# masked rows run at their own strength: mask value m puts a row at sigma = m * sigma_stream,
# so its label is 1 - m * sigma, clamped at the cond timestep for fully preserved rows
t_pin_v = max(t_v, VISUAL_COND_TIMESTEP)
t_pin_a = max(t_a, AUDIO_COND_TIMESTEP)
video_rows_t = None
audio_rows_t = None
if denoise_mask is not None:
m = mask_row_values(denoise_mask[0, 0].to(torch.float32), latent_t, lat_h, lat_w)
if m is not None:
rows_t = (1.0 - m * sigma_v.to(m.device)).clamp(max=t_pin_v)
if rows_t.unique().numel() == 1:
seg_t["video"] = float(rows_t[0])
else:
video_rows_t = rows_t
if audio_denoise_mask is not None:
m = audio_denoise_mask[0, 0].to(torch.float32).reshape(-1)
if not bool((m >= 1.0 - 1e-3).all()):
sigma_a = 1.0 - t_a
rows_t = (1.0 - m * sigma_a).clamp(max=t_pin_a)
if rows_t.unique().numel() == 1:
seg_t["audio"] = float(rows_t[0])
else:
audio_rows_t = rows_t
unique_t = sorted({t_v, t_a} | {seg_t[k] for _, _, k in layout.segments}
| (set(video_rows_t.unique().tolist()) if video_rows_t is not None else set())
| (set(audio_rows_t.unique().tolist()) if audio_rows_t is not None 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, "cond_audio": 2, "ref_audio": 2}
def rows_to_mod_index(rows_t, tag):
# per-row timestep values -> per-row mod-row indices into the t_emb table
levels = rows_t.unique()
base = torch.tensor([t_row[v] * 3 + tag for v in levels.tolist()],
dtype=torch.long, device=rows_t.device)
return base[torch.searchsorted(levels, rows_t)]
text_tags = payload.get("text_token_tags")
mod_segments = []
for a, b, kind in layout.segments:
row_base = t_row[seg_t[kind]] * 3
if kind == "text" and text_tags is not None:
# the presentation text span mixes tags (vision pads carry the video modality) split into tag runs
tags = text_tags.view(-1).tolist()
run_start = 0
for i in range(1, b - a + 1):
if i == b - a or tags[i] != tags[run_start]:
mod_segments.append((a + run_start, a + i, row_base + int(tags[run_start])))
run_start = i
elif kind == "video" and video_rows_t is not None:
mod_segments.append((a, b, rows_to_mod_index(video_rows_t, seg_tag[kind])))
elif kind == "audio" and audio_rows_t is not None:
mod_segments.append((a, b, rows_to_mod_index(audio_rows_t, seg_tag[kind])))
else:
mod_segments.append((a, b, row_base + seg_tag[kind]))
# embed
img_update = layout.img_update.to(device)
audio_update = layout.audio_update.to(device)
video_rows = patchify_video(video_x.to(torch.float32), self.patch_size)
audio_rows = pack_audio(audio_x.to(torch.float32))
cond_video_rows = self._cond_video_rows(payload, device)
cond_audio_rows = self._cond_audio_rows(payload, device)
all_video_rows = video_rows
if cond_video_rows is not None:
all_video_rows = torch.empty(img_update.shape[0], video_rows.shape[1], dtype=torch.float32, device=device)
all_video_rows[~img_update] = cond_video_rows
all_video_rows[img_update] = video_rows
all_audio_rows = audio_rows
if cond_audio_rows is not None:
all_audio_rows = torch.empty(audio_update.shape[0], audio_rows.shape[1], dtype=torch.float32, device=device)
all_audio_rows[~audio_update] = cond_audio_rows
all_audio_rows[audio_update] = audio_rows
video_embed = self.video_patch_proj(all_video_rows).to(dtype)
audio_embed = self.audio_patch_proj(all_audio_rows).to(dtype)
text_states = context[0]
if text_states.shape[-1] != self.hidden_size:
text_states = self.token_refiner(self.condition_proj(text_states),
transformer_options=transformer_options)
# segments are contiguous: assemble by slices, embed rows follow segment order
h = torch.empty(layout.seq_len, self.hidden_size, dtype=dtype, device=device)
voff = aoff = 0
for a, b, kind in layout.segments:
n = b - a
if kind == "text":
h[a:b] = text_states
elif kind in ("cond", "ref_img", "video"):
h[a:b] = video_embed[voff:voff + n]
voff += n
else: # ref_audio / audio
h[a:b] = audio_embed[aoff:aoff + n]
aoff += n
t_vals = torch.tensor(unique_t, dtype=torch.float32, device=device)
if self.use_adaln_curves:
# adaln projections consume interpolated coordinates of the time-embedding curve
table = comfy.model_management.cast_to(self.adaln_t_table, device=device)
pos = t_vals.clamp(0.0, 1.0) * (table.shape[0] - 1) # t in [0,1] -> fractional grid index, out-of-range t clamps to the curve ends
i0 = pos.floor().long().clamp(max=table.shape[0] - 2) # lower grid row, max-clamp keeps t=1.0 on the last interval instead of reading past the table
t_emb = torch.lerp(table[i0], table[i0 + 1], (pos - i0).unsqueeze(1)) # blend the two rows by the fractional part
else:
t_emb = self.time_embedder(t_vals).to(dtype)
# rotation table computed once per forward, consumed by the kitchen split-half rope
rope_freqs = rope_rotation_table(self.rope_freqs(layout.position_ids, device), dtype)
# blocks
patches_replace = transformer_options.get("patches_replace", {})
blocks_replace = patches_replace.get("dit", {})
prefetch_queue = comfy.model_prefetch.make_prefetch_queue(list(self.blocks), device, transformer_options)
for i, block in enumerate(self.blocks):
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, block)
if ("double_block", i) in blocks_replace:
def block_wrap(args):
return {"img": block(args["img"], args["t_emb"], args["mod_segments"], args["rope_freqs"],
transformer_options=args["transformer_options"])}
h = blocks_replace[("double_block", i)](
{"img": h, "t_emb": t_emb, "mod_segments": mod_segments, "rope_freqs": rope_freqs,
"layout": layout, "transformer_options": transformer_options},
{"original_block": block_wrap})["img"]
else:
h = block(h, t_emb, mod_segments, rope_freqs, transformer_options=transformer_options)
if prefetch_queue is not None:
comfy.model_prefetch.prefetch_queue_pop(prefetch_queue, device, None)
# target streams are single contiguous segments (audio then video, last two)
va, vb, _ = next(s for s in layout.segments if s[2] == "video")
aa, ab, _ = next(s for s in layout.segments if s[2] == "audio")
if video_rows_t is not None:
video_seg = (va, vb, rows_to_mod_index(video_rows_t, 0) // 3)
else:
video_seg = (va, vb, t_row[seg_t["video"]])
if audio_rows_t is not None:
audio_seg = (aa, ab, rows_to_mod_index(audio_rows_t, 0) // 3)
else:
audio_seg = (aa, ab, t_row[seg_t["audio"]])
v, a = self.final_layer(h, t_emb, video_seg, audio_seg, sigma_v, transformer_options.get("sample_sigmas"), (shift_v, shift_a))
video_out = unpatchify_video(v, latent_t, lat_h // 2, lat_w // 2, self.latents_dim, self.patch_size)
video_out = video_out[:, :, :orig_t, :orig_h, :orig_w]
audio_out = unpack_audio(a)
return [-video_out.to(video_x.dtype), -audio_out.to(audio_x.dtype)]
comfy/ldm/minimax/vae.py
+710
View File
@@ -0,0 +1,710 @@
# MiniMax H3 video VAE: 3D causal CNN encoder + ViT3D decoder.
import math
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
from comfy.ldm.modules.attention import optimized_attention
ops = comfy.ops.disable_weight_init
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
LATENTS_MEAN = [
0.858090341091156, -0.9606591463088989, 1.0661640167236328, -0.5090325474739075,
-0.2727581858634949, -1.3675414323806763, -0.2553254961967468, -0.26907554268836975,
-0.5376840829849243, -0.0464097298681736, 0.6657370328903198, 0.19690127670764923,
-0.5460608005523682, -0.4035342037677765, -0.23683024942874908, 0.25928452610969543,
-0.30133944749832153, 0.211341992020607, -1.1206848621368408, 0.3581933379173279,
-0.04225143790245056, 0.2604829967021942, 0.22864092886447906, 0.7056031823158264,
]
LATENTS_STD = [
1.2223774194717407, 1.2767263650894165, 1.68317747116088865, 1.7549455165863037,
1.5636216402053833, 2.194143533706665, 0.96531379222869875, 1.05698859691619875,
0.841948926448822, 0.7729952931404114, 1.8955937623977661, 0.946841835975647,
0.7996809482574463, 0.44988900423049925, 0.7197399735450745, 0.69362932443618775,
2.961095094680786, 2.7694199085235595, 3.0496184825897215, 2.1088054180145265,
3.276226282119751, 3.1627357006073, 2.28168129920959475, 2.6127843856811525,
]
# 3D causal CNN encoder
class CausalConv3d(ops.Conv3d):
# Reflect spatial padding, causal (zeros, front-only) temporal padding.
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0):
super().__init__(in_channels, out_channels, kernel_size=kernel_size, stride=stride)
self.causal_padding = (padding,) * 3 if isinstance(padding, int) else tuple(padding)
def forward(self, x):
if sum(self.causal_padding) == 0:
return super().forward(x)
x = F.pad(x, (self.causal_padding[2], self.causal_padding[2], self.causal_padding[1], self.causal_padding[1], 0, 0), mode="reflect")
if x.shape[2] == 1:
# single frame: the causal front padding is all zeros truncate the temporal taps instead of convolving zero frames
return super().forward(x, autopad="causal_zero")
x = F.pad(x, (0, 0, 0, 0, self.causal_padding[0] * 2, 0), mode="constant")
return super().forward(x)
class TemporalIsolatedGroupNorm(ops.GroupNorm):
# GroupNorm with statistics computed per frame (time merged into batch).
def forward(self, x):
if x.dim() == 5:
b, c, t, h, w = x.shape
x = x.permute(0, 2, 1, 3, 4).contiguous().view(b * t, c, 1, h, w)
x = super().forward(x)
return x.view(b, t, c, h, w).permute(0, 2, 1, 3, 4).contiguous()
return super().forward(x)
def group_norm_3d(num_channels):
return TemporalIsolatedGroupNorm(num_groups=32, num_channels=num_channels, eps=1e-6, affine=True)
class Downsample3D(nn.Module):
def __init__(self, in_channels, out_channels, time_stride=1, space_stride=2):
super().__init__()
self.space_stride = space_stride
self.conv = CausalConv3d(
in_channels,
out_channels,
kernel_size=3,
padding=(1, 0, 0),
stride=(time_stride, space_stride, space_stride),
)
def forward(self, x):
if self.space_stride == 2:
x = F.pad(x, (0, 1, 0, 1, 0, 0), mode="reflect")
return self.conv(x)
class ResnetBlock3D(nn.Module):
def __init__(self, in_channels, out_channels=None):
super().__init__()
self.in_channels = in_channels
out_channels = in_channels if out_channels is None else out_channels
self.out_channels = out_channels
self.norm1 = group_norm_3d(in_channels)
self.norm2 = group_norm_3d(out_channels)
self.conv1 = CausalConv3d(in_channels, out_channels, kernel_size=3, padding=1)
self.conv2 = CausalConv3d(out_channels, out_channels, kernel_size=3, padding=1)
if in_channels != out_channels:
self.nin_shortcut = CausalConv3d(in_channels, out_channels, kernel_size=1)
def forward(self, x):
h = self.conv1(F.silu(self.norm1(x), inplace=True))
h = self.conv2(F.silu(self.norm2(h), inplace=True))
if self.in_channels != self.out_channels:
x = self.nin_shortcut(x)
return h.add_(x)
class EncoderFCN3D(nn.Module):
def __init__(self, ch, ch_mult, space_down, time_down, num_res_blocks, in_channels, z_channels, double_z=True):
super().__init__()
self.num_levels = len(ch_mult)
if isinstance(num_res_blocks, int):
num_res_blocks = [num_res_blocks] * self.num_levels
self.num_res_blocks = num_res_blocks
block_mid = [ch * ch_mult[i] for i in range(self.num_levels)]
block_in = [block_mid[0]] + block_mid[:-1]
block_out = block_mid
self.conv_in = CausalConv3d(in_channels, block_in[0], kernel_size=3, padding=1)
self.down = nn.ModuleList()
for i_level in range(self.num_levels):
down = nn.Module()
down.block = nn.ModuleList()
for i in range(self.num_res_blocks[i_level]):
down.block.append(
ResnetBlock3D(
in_channels=block_in[i_level] if i == 0 else block_mid[i_level],
out_channels=block_mid[i_level],
)
)
if space_down[i_level] * time_down[i_level] > 1:
down.downsample = Downsample3D(
block_mid[i_level],
block_out[i_level],
time_stride=time_down[i_level],
space_stride=space_down[i_level],
)
self.down.append(down)
self.norm_out = group_norm_3d(block_out[-1])
self.conv_out = CausalConv3d(
block_out[-1],
2 * z_channels if double_z else z_channels,
kernel_size=3,
padding=1,
)
def forward(self, x):
h = self.conv_in(x)
for i_level in range(self.num_levels):
for i_block in range(self.num_res_blocks[i_level]):
h = self.down[i_level].block[i_block](h)
if hasattr(self.down[i_level], "downsample"):
h = self.down[i_level].downsample(h)
h = F.silu(self.norm_out(h))
return self.conv_out(h)
# ViT3D decoder
def create_token_ids(patch_dims, device, dtype):
coords_list = []
for dim_size in patch_dims:
coords = torch.arange(0.5, dim_size, dtype=dtype, device=device)
coords = coords / dim_size
coords = 2.0 * coords - 1.0
coords_list.append(coords)
coords = torch.stack(torch.meshgrid(*coords_list, indexing="ij"), dim=-1)
return coords.flatten(0, len(patch_dims) - 1).unsqueeze(0)
class RotaryEmbeddingND(nn.Module):
def __init__(self, dim, rotary_base=100.0, n_dim=3):
super().__init__()
self.n_dim = n_dim
self.angle_scale = 2.0 * math.pi
inv_freq = 1 / rotary_base ** torch.arange(0, 1, 2 * n_dim / dim, dtype=torch.float32)
self.register_buffer("inv_freq", inv_freq, persistent=False)
def forward(self, img_ids):
# [B, S, n_dim] -> [B, S, 1, pairs, 2, 2] rotation table for the kitchen split-half rope
angles = (
self.angle_scale
* img_ids[:, :, :, None].float()
* self.inv_freq.to(img_ids.device)[None, None, None, :]
)
angles = angles.flatten(2, 3)
c, s = torch.cos(angles), torch.sin(angles)
table = torch.stack([c, -s, s, c], dim=-1).reshape(*angles.shape[:2], 1, angles.shape[-1], 2, 2)
return table.to(img_ids.dtype)
class FeedForward(nn.Module):
# Gated SiLU FFN.
def __init__(self, dim, mult=4, bias=True, operations=ops):
super().__init__()
inner_dim = dim * mult
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)
return self.w2(F.silu(gate).mul_(x))
class Attention(nn.Module):
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 = 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
qkv = self.to_qkv(x)
qkv = qkv.view(batch_size, seq_len, -1, 3 * self.dim_head)
query, key, value = torch.chunk(qkv, 3, dim=-1)
query = comfy.rmsnorm.rms_norm(query, self.norm_q.weight, self.norm_q.eps)
key = comfy.rmsnorm.rms_norm(key, self.norm_k.weight, self.norm_k.eps)
if rotary_pos_emb is not None:
rot = rotary_pos_emb.shape[-3] * 2
query[..., :rot], key[..., :rot] = comfy.quant_ops.ck.apply_rope_split_half(
query[..., :rot], key[..., :rot], rotary_pos_emb)
out = optimized_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2),
self.heads, skip_reshape=True).nan_to_num_(0.0)
return self.to_out(out)
class TransformerBlock(nn.Module):
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, 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, 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), 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, operations=ops):
super().__init__()
dim = heads * dim_head
self.patch_size = patch_size
self.patch_size_t = patch_size_t
self.out_channels = out_channels
self.num_register_tokens = num_register_tokens
self.pos_embed = RotaryEmbeddingND(int(dim_head * rope_dim_ratio), rope_theta, n_dim=3)
self.x_embedder = ops.Linear(in_channels, dim)
self.register_tokens = nn.Parameter(torch.empty(1, num_register_tokens, dim))
# unused at inference; kept so the checkpoint loads without leftover keys
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, operations=operations)
for _ in range(num_layers)]
)
self.norm_out = ops.LayerNorm(dim, elementwise_affine=True, eps=eps)
self.proj_out = ops.Linear(dim, out_channels * patch_size_t * patch_size * patch_size)
def forward(self, x):
B, C, latent_T, latent_H, latent_W = x.shape
h = self.x_embedder(x.flatten(2).transpose(1, 2)) # [B, T*H*W, C]
num_patches = h.shape[1]
num_suffix = 1 + self.num_register_tokens
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)
img_ids = torch.cat([img_ids, suffix_ids], dim=1)
rotary_pos_emb = self.pos_embed(img_ids)
for block in self.transformer_blocks:
h = block(h, rotary_pos_emb)
output = self.proj_out(self.norm_out(h))
output = output[:, :num_patches, :]
output = output.view(
B, latent_T, latent_H, latent_W,
self.out_channels, self.patch_size_t, self.patch_size, self.patch_size,
)
output = output.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()
output = output.reshape(
B, self.out_channels,
latent_T * self.patch_size_t,
latent_H * self.patch_size,
latent_W * self.patch_size,
)
return output
# Full VAE
class MiniMaxH3VideoVAE(nn.Module):
comfy_has_chunked_io = True
def __init__(
self,
in_channels=3,
out_ch=3,
ch=128,
embed_dim=24,
z_channels=24,
ch_mult=(1, 2, 2, 4, 4, 8),
num_res_blocks=2,
space_down=(2, 2, 2, 2, 1, 1),
time_down=(1, 2, 2, 1, 1, 1),
clip_length=17,
token_drop=3,
tile_size=256,
tile_overlap_min=64,
tiling=True,
operations=ops,
):
super().__init__()
self.vae_ratio = int(math.prod(space_down))
self.vae_ratio_t = int(math.prod(time_down))
# temporal chunking parameters
self.clip_length = clip_length
self.token_drop = token_drop
self.frame_pre_padding = (-clip_length) % self.vae_ratio_t
self.tokens_chunk_size = math.ceil(clip_length / self.vae_ratio_t)
self.token_overlap = (-token_drop) % self.tokens_chunk_size
self.frame_overlap = max(self.token_overlap * self.vae_ratio_t - self.frame_pre_padding, 0)
# spatial tiling parameters
self.tiling = tiling
self.tile_size = tile_size
self.tile_overlap_min = tile_overlap_min
self.encoder = EncoderFCN3D(
ch=ch,
ch_mult=list(ch_mult),
space_down=list(space_down),
time_down=list(time_down),
num_res_blocks=num_res_blocks,
in_channels=in_channels,
z_channels=z_channels,
double_z=True,
)
self.quant_conv = ops.Conv3d(z_channels * 2, 2 * embed_dim, 1)
self.post_quant_conv = ops.Conv3d(embed_dim, z_channels, 1)
self.decoder = ViT3DDecoder(
patch_size=self.vae_ratio,
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))
self.register_buffer("latents_std", torch.tensor(LATENTS_STD))
self.register_buffer("pixel_mean", torch.tensor(IMAGENET_MEAN).view(1, 3, 1, 1, 1), persistent=False)
self.register_buffer("pixel_std", torch.tensor(IMAGENET_STD).view(1, 3, 1, 1, 1), persistent=False)
# single-shot forward
def _encode_moments(self, x):
return self.quant_conv(self.encoder(x))
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)
return self._encode_moments(x)
def _adaptive_decode(self, z):
if self.tiling:
return self.tiled_decode(z)
return self._decode_pixels(z)
# spatial tiling
def split_tiles(self, input_len):
tile_size = self.tile_size
if tile_size >= input_len:
return [0], [input_len], []
N = math.ceil(input_len / tile_size)
while True:
overlaps = [self.tile_overlap_min] * (N - 1)
remaining = tile_size * N - sum(overlaps) - input_len
if remaining < 0:
N += 1
else:
break
remaining_units = remaining // self.vae_ratio
for i in range(remaining_units):
overlaps[i % (N - 1)] += self.vae_ratio
tile_start_idx = [0]
for i in range(N - 1):
tile_start_idx.append(tile_start_idx[-1] + tile_size - overlaps[i])
return tile_start_idx, [tile_size] * N, overlaps
def blend(self, a, b, blend_extent, dim):
blend_extent = min(a.shape[dim], b.shape[dim], blend_extent)
positions = torch.arange(blend_extent, device=b.device, dtype=b.dtype)
weight_a = 1 - positions / blend_extent
weight_b = positions / blend_extent
shape = [1] * a.ndim
shape[dim] = blend_extent
weight_a = weight_a.view(shape)
weight_b = weight_b.view(shape)
slice_a = [slice(None)] * a.ndim
slice_a[dim] = slice(-blend_extent, None)
slice_b = [slice(None)] * b.ndim
slice_b[dim] = slice(0, blend_extent)
blended = a[tuple(slice_a)] * weight_a + b[tuple(slice_b)] * weight_b
if blend_extent < b.shape[dim]:
slice_b_rest = [slice(None)] * b.ndim
slice_b_rest[dim] = slice(blend_extent, None)
return torch.cat([blended, b[tuple(slice_b_rest)]], dim=dim)
return blended
def tiled_encode(self, x):
height, width = x.shape[-2], x.shape[-1]
y_idx, y_len, y_overlap = self.split_tiles(height)
x_idx, x_len, x_overlap = self.split_tiles(width)
rows = []
for i_pos, i_len in zip(y_idx, y_len):
row = []
for j_pos, j_len in zip(x_idx, x_len):
tile = x[..., i_pos:i_pos + i_len, j_pos:j_pos + j_len]
row.append(self._encode_moments(tile))
rows.append(row)
latent_y_overlap = [o // self.vae_ratio for o in y_overlap]
latent_x_overlap = [o // self.vae_ratio for o in x_overlap]
result_rows = []
for i, row in enumerate(rows):
result_row = []
for j, tile in enumerate(row):
if i > 0:
tile = self.blend(rows[i - 1][j], tile, latent_y_overlap[i - 1], dim=-2)
if j > 0:
tile = self.blend(row[j - 1], tile, latent_x_overlap[j - 1], dim=-1)
if i < len(rows) - 1:
tile = tile[..., :-latent_y_overlap[i], :]
if j < len(row) - 1:
tile = tile[..., :, :-latent_x_overlap[j]]
result_row.append(tile)
result_rows.append(torch.cat(result_row, dim=-1))
return torch.cat(result_rows, dim=-2)
def tiled_decode(self, z):
height, width = z.shape[-2] * self.vae_ratio, z.shape[-1] * self.vae_ratio
y_idx, y_len, y_overlap = self.split_tiles(height)
x_idx, x_len, x_overlap = self.split_tiles(width)
# Blended tiles are written straight into a pre-allocated canvas.
canvas = None
row_tails = []
out_y = 0
for i, (i_pos, i_len) in enumerate(zip(y_idx, y_len)):
zi, zl = i_pos // self.vae_ratio, i_len // self.vae_ratio
new_tails = []
left_tail = None
out_x = 0
for j, (j_pos, j_len) in enumerate(zip(x_idx, x_len)):
zj, zw = j_pos // self.vae_ratio, j_len // self.vae_ratio
tile = self._decode_pixels(z[..., zi:zi + zl, zj:zj + zw])
if i < len(y_idx) - 1:
new_tails.append(tile[..., -y_overlap[i]:, :].clone())
next_left_tail = tile[..., :, -x_overlap[j]:].clone() if j < len(x_idx) - 1 else None
if i > 0:
tile = self.blend(row_tails[j], tile, y_overlap[i - 1], dim=-2)
if j > 0:
tile = self.blend(left_tail, tile, x_overlap[j - 1], dim=-1)
left_tail = next_left_tail
if i < len(y_idx) - 1:
tile = tile[..., :-y_overlap[i], :]
if j < len(x_idx) - 1:
tile = tile[..., :, :-x_overlap[j]]
if canvas is None:
canvas = torch.empty(*tile.shape[:-2], height, width, dtype=tile.dtype, device=tile.device)
canvas[..., out_y:out_y + tile.shape[-2], out_x:out_x + tile.shape[-1]].copy_(tile)
out_x += tile.shape[-1]
row_tails = new_tails
out_y += tile.shape[-2]
return canvas
# temporal chunking
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(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:
z = z[:, :, :-self.token_drop]
return z
def _decode_temporal_pad_frames(self, z_len, pad_tokens):
if pad_tokens <= 0:
return 0
intra_tail = self.clip_length % self.vae_ratio_t
if intra_tail == 0:
return pad_tokens * self.vae_ratio_t
z_len_before_pad = z_len - pad_tokens
return sum(
(intra_tail if (z_len_before_pad + k) % self.tokens_chunk_size == 0
else self.vae_ratio_t)
for k in range(pad_tokens)
)
def _decode_temporal_frame_plan(self, z_len, num_chunks, pad_tokens):
chunk_dec = self.tokens_chunk_size * self.vae_ratio_t
split_count = int(self.token_drop > 0) + 1
total_frames = 0
final_overlap_frames = 0
for i in range(num_chunks):
t_start_idx = i * self.tokens_chunk_size
t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
clip_token_len = max(0, min(t_end_idx, z_len) - min(t_start_idx, z_len))
clip_frame_len = clip_token_len * self.vae_ratio_t
for j in range(split_count):
f_start_idx = j * chunk_dec
f_end_idx = min(f_start_idx + chunk_dec, clip_frame_len)
chunk_frames = max(0, f_end_idx - f_start_idx - self.frame_pre_padding)
if j == 0:
total_frames += chunk_frames
else:
final_overlap_frames = chunk_frames
total_frames += final_overlap_frames
return total_frames - self._decode_temporal_pad_frames(z_len, 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)
dec = output_buffer
dec_overlap = None
write_pos = 0
def write_part(part):
nonlocal write_pos
part_frames = part.shape[2]
if part_frames <= 0:
return
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_(
part[:, :, :copy_frames, :, :]
)
write_pos += copy_frames
for i in range(num_chunks):
t_start_idx = i * self.tokens_chunk_size
t_end_idx = t_start_idx + self.tokens_chunk_size + self.token_overlap
clip_z = z[:, :, t_start_idx:t_end_idx, :, :]
clip_dec = self._adaptive_decode(clip_z)
for j in range(split_count):
f_start_idx = j * chunk_dec
f_end_idx = min(f_start_idx + chunk_dec, clip_dec.shape[2])
clip_dec_chunk = clip_dec[:, :, f_start_idx:f_end_idx, :, :]
clip_dec_chunk = clip_dec_chunk[:, :, self.frame_pre_padding:, :, :]
if j == 0:
if dec_overlap is not None:
clip_dec_chunk = self.blend(
dec_overlap, clip_dec_chunk, self.frame_overlap, dim=-3
)
dec_overlap = None
write_part(clip_dec_chunk)
else:
dec_overlap = clip_dec_chunk.contiguous()
if i == num_chunks - 1 and dec_overlap is not None:
write_part(dec_overlap)
dec_overlap = None
del clip_dec, clip_z
return dec
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)
if device is None:
device = x.device
if x.shape[2] == 1:
moments = self._adaptive_encode(self._normalize_pixels(x.to(device)))
moments = moments[:, :, -1:, :, :]
else:
moments = self.encode_temporal(x, device)
mean = torch.chunk(moments.float(), 2, dim=1)[0]
latents_mean = self.latents_mean.view(1, -1, 1, 1, 1).to(mean)
latents_std = self.latents_std.view(1, -1, 1, 1, 1).to(mean)
return (mean - latents_mean) / latents_std
def encode_tiled(self, x, **kwargs):
# tiling is always on internally with the reference's semantic tile sizes, ignore tiling fallbacks
return self.encode(x)
def decode_tiled(self, z, **kwargs):
return self.decode(z)
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._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/__init__.py
View File
comfy/ldm/minimax_music/ar.py
+343
View File
@@ -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
+137
View File
@@ -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
+213
View File
@@ -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
+70
View File
@@ -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
+120 -27
View File
@@ -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
@@ -90,22 +94,6 @@ def default(val, d):
return val
return d
def _gqa_repeat_factor(query_heads, key_heads, value_heads):
if key_heads != value_heads:
raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
if query_heads == key_heads:
return 1
if query_heads % key_heads != 0:
raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
return query_heads // key_heads
def _repeat_kv_for_gqa(k, v, query_heads, head_dim):
n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
if n_rep > 1:
k = k.repeat_interleave(n_rep, dim=head_dim)
v = v.repeat_interleave(n_rep, dim=head_dim)
return k, v
def _heads_from_dim(tensor, dim_head, name):
inner_dim = tensor.shape[-1]
if inner_dim % dim_head != 0:
@@ -122,10 +110,8 @@ def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_
value_heads = heads
k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head)
v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head)
if enable_gqa:
_gqa_repeat_factor(heads, key_heads, value_heads)
if expand_kv:
k, v = _repeat_kv_for_gqa(k, v, heads, -2)
if enable_gqa and expand_kv:
k, v = comfy.ops.repeat_kv_for_gqa(k, v, heads, -2)
return q, k, v
@@ -163,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:
@@ -174,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
@@ -196,7 +218,7 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
h = heads
if skip_reshape:
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-3], -3)
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head),
(q, k, v),
@@ -262,7 +284,7 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None,
if skip_reshape:
if kwargs.get("enable_gqa", False):
key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3)
key, value = comfy.ops.repeat_kv_for_gqa(key, value, query.shape[-3], -3)
query = query.reshape(b * heads, -1, dim_head)
value = value.reshape(b * heads, -1, dim_head)
key = key.reshape(b * heads, -1, dim_head).movedim(1, 2)
@@ -338,7 +360,7 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
if skip_reshape:
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-3], -3)
q, k, v = map(
lambda t: t.reshape(b * heads, -1, dim_head),
(q, k, v),
@@ -476,7 +498,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh
(q, k, v),
)
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2)
k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-2], -2)
# actually do the reshaping
else:
dim_head //= heads
@@ -563,6 +585,67 @@ 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):
if kwargs.get("low_precision_attention", True) is False and q.dtype == torch.float32:
return attention_pytorch(q, k, v, heads, mask=mask, attn_precision=attn_precision, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **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()
if kwargs.get("low_precision_attention", True) is False and q.dtype == torch.float32:
return attention_pytorch(q, k, v, heads, mask=mask, attn_precision=attn_precision, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs)
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):
@@ -573,7 +656,7 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=
b, _, _, dim_head = q.shape
tensor_layout = "HND"
if kwargs.get("enable_gqa", False):
k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3)
k, v = comfy.ops.repeat_kv_for_gqa(k, v, q.shape[-3], -3)
else:
b, _, dim_head = q.shape
dim_head //= heads
@@ -671,7 +754,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
if skip_reshape:
q_s = q
if kwargs.get("enable_gqa", False):
k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3)
k_s, v_s = comfy.ops.repeat_kv_for_gqa(k, v, H, -3)
else:
k_s, v_s = k, v
else:
@@ -793,10 +876,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/moge/model.py
+2 -3
View File
@@ -69,7 +69,7 @@ class MoGeModelV1(nn.Module):
resize = ((num_tokens * 14 ** 2) / (H * W)) ** 0.5
rh, rw = int(H * resize), int(W * resize)
x = F.interpolate(image, (rh, rw), mode="bicubic", align_corners=False, antialias=True)
x = (x - self.image_mean) / self.image_std
x = (x - comfy.ops.cast_to_input(self.image_mean, x, copy=False)) / comfy.ops.cast_to_input(self.image_std, x, copy=False)
x14 = F.interpolate(x, (rh // 14 * 14, rw // 14 * 14), mode="bilinear", align_corners=False, antialias=True)
n_layers = len(self.backbone.encoder.layer)
@@ -268,7 +268,6 @@ class MoGeModel:
"""Loaded MoGe model + ComfyUI memory management."""
def __init__(self, state_dict: dict):
# text encoder dtype closest match
self.load_device = comfy.model_management.text_encoder_device()
offload_device = comfy.model_management.text_encoder_offload_device()
self.dtype = comfy.model_management.text_encoder_dtype(self.load_device)
@@ -287,7 +286,7 @@ class MoGeModel:
) -> Dict[str, torch.Tensor]:
"""Run a single MoGe forward + post-process pass. image is (B, 3, H, W) in [0, 1]."""
comfy.model_management.load_model_gpu(self.patcher)
image = image.to(device=self.load_device, dtype=self.dtype)
image = image.to(device=self.load_device, dtype=torch.float32)
H, W = image.shape[-2:]
aspect_ratio = W / H
comfy/ldm/moge/modules.py
+1 -1
View File
@@ -136,7 +136,7 @@ class DINOv2Encoder(nn.Module):
def forward(self, image: torch.Tensor, token_rows: int, token_cols: int,
return_class_token: bool = False) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
image_14 = F.interpolate(image, (token_rows * 14, token_cols * 14), mode="bilinear", align_corners=False, antialias=True)
image_14 = (image_14 - self.image_mean) / self.image_std
image_14 = (image_14 - comfy.ops.cast_to_input(self.image_mean, image_14, copy=False)) / comfy.ops.cast_to_input(self.image_std, image_14, copy=False)
feats = self.backbone.get_intermediate_layers(image_14, self.intermediate_layers, apply_norm=True)
x = torch.stack([
proj(feat.permute(0, 2, 1).unflatten(2, (token_rows, token_cols)).contiguous())
comfy/ldm/rt_detr/rtdetr_v4.py
+1 -1
View File
@@ -31,7 +31,7 @@ class ConvBNAct(nn.Module):
super().__init__()
self.conv = operations.Conv2d(ic, oc, k, s, (k - 1) // 2, groups=groups, bias=False, device=device, dtype=dtype)
self.bn = nn.BatchNorm2d(oc, device=device, dtype=dtype)
self.bn = operations.BatchNorm2d(oc, device=device, dtype=dtype)
self.act = nn.ReLU() if use_act else nn.Identity()
def forward(self, x):
comfy/ldm/sam3d_body/mhr/mhr_head.py
+344
View File
@@ -0,0 +1,344 @@
from typing import Optional
import torch
import torch.nn as nn
from comfy.ops import cast_to_input
from ..utils import euler_to_rotmat, rot6d_to_rotmat, rotmat_to_euler, unitquat_to_rotmat
from .mhr_utils import compact_cont_to_model_params_body, compact_cont_to_model_params_hand, mhr_param_hand_idxs
from ..model.transformer import MLP
class MHRHead(nn.Module):
def __init__(self, input_dim: int, mhr_rig, mlp_depth: int = 1, mlp_channel_div_factor: int = 8, enable_hand_model=False,
device=None, dtype=None, operations=None):
super().__init__()
# Store the shared MHRRig as a non-registered Python attribute
object.__setattr__(self, "mhr", mhr_rig)
self.num_shape_comps = 45
self.num_scale_comps = 28
self.num_hand_comps = 54
self.num_face_comps = 72
self.enable_hand_model = enable_hand_model
self.body_cont_dim = 260
self.npose = (
6 # Global Rotation
+ self.body_cont_dim # then body
+ self.num_shape_comps
+ self.num_scale_comps
+ self.num_hand_comps * 2
+ self.num_face_comps
)
self.proj = MLP(
input_dim=input_dim,
hidden_dim=input_dim // mlp_channel_div_factor,
output_dim=self.npose,
num_layers=mlp_depth,
device=device, dtype=dtype, operations=operations,
)
# MHR Parameters
self.num_hand_scale_comps = self.num_scale_comps - 18
self.num_hand_pose_comps = self.num_hand_comps
# Buffers populated by load_state_dict from the safetensors
def _p(*shape, dtype=torch.float32):
return nn.Parameter(torch.empty(*shape, dtype=dtype, device=device), requires_grad=False)
self.joint_rotation = _p(127, 3, 3)
self.scale_mean = _p(68)
self.scale_comps = _p(28, 68)
self.register_buffer("faces", torch.empty(36874, 3, dtype=torch.int64, device=device))
self._faces_np = None
self.hand_pose_mean = _p(54)
self.hand_pose_comps = _p(54, 54)
self.register_buffer("hand_joint_idxs_left", torch.empty(27, dtype=torch.int64, device=device))
self.register_buffer("hand_joint_idxs_right", torch.empty(27, dtype=torch.int64, device=device))
self.keypoint_mapping = _p(308, 18439 + 127)
# Some special buffers for the hand-version
self.right_wrist_coords = _p(3)
self.root_coords = _p(3)
self.local_to_world_wrist = _p(3, 3)
self.register_buffer("nonhand_param_idxs", torch.empty(145, dtype=torch.int64, device=device))
if not enable_hand_model:
self.register_buffer("face_region_rgb", torch.empty(18439, 3, dtype=torch.float32, device=device))
def canonical_vertices(self):
"""Return the T-pose vertices for the mean shape (scaled to meters).
Runs MHR with zero pose / shape / scale / expression so the returned
mesh is the canonical rest pose — fixed per-model
"""
device = self.scale_mean.device
dtype = self.scale_mean.dtype
B = 1
global_trans = torch.zeros(B, 3, device=device, dtype=dtype)
global_rot = torch.zeros(B, 3, device=device, dtype=dtype)
body_pose = torch.zeros(B, 130, device=device, dtype=dtype)
hand_pose = torch.zeros(B, self.num_hand_comps * 2, device=device, dtype=dtype)
scale = torch.zeros(B, self.num_scale_comps, device=device, dtype=dtype)
shape = torch.zeros(B, self.num_shape_comps, device=device, dtype=dtype)
expr = torch.zeros(B, self.num_face_comps, device=device, dtype=dtype)
verts = self.mhr_forward(
global_trans=global_trans,
global_rot=global_rot,
body_pose_params=body_pose,
hand_pose_params=hand_pose,
scale_params=scale,
shape_params=shape,
expr_params=expr,
) # single-tensor shape (1, N_v, 3) in meters
return verts[0]
def faces_np(self):
"""Static topology — cached so the per-layer pose_output doesn't force a D2H sync."""
if self._faces_np is None:
self._faces_np = self.faces.cpu().numpy()
return self._faces_np
def replace_hands_in_pose(self, full_pose_params, hand_pose_params):
assert full_pose_params.shape[1] == 136
# This drops in the hand poses from hand_pose_params (PCA 6D) into full_pose_params.
# Split into left and right hands
left_hand_params, right_hand_params = torch.split(
hand_pose_params,
[self.num_hand_pose_comps, self.num_hand_pose_comps],
dim=1,
)
# Change from cont to model params
left_hand_params_model_params = compact_cont_to_model_params_hand(
cast_to_input(self.hand_pose_mean, left_hand_params, copy=False)
+ torch.einsum("da,ab->db", left_hand_params, cast_to_input(self.hand_pose_comps, left_hand_params, copy=False))
)
right_hand_params_model_params = compact_cont_to_model_params_hand(
cast_to_input(self.hand_pose_mean, right_hand_params, copy=False)
+ torch.einsum("da,ab->db", right_hand_params, cast_to_input(self.hand_pose_comps, right_hand_params, copy=False))
)
# Drop it in
full_pose_params[:, self.hand_joint_idxs_left.to(full_pose_params.device)] = left_hand_params_model_params
full_pose_params[:, self.hand_joint_idxs_right.to(full_pose_params.device)] = right_hand_params_model_params
return full_pose_params # B x 207
def mhr_forward(
self,
global_trans,
global_rot,
body_pose_params,
hand_pose_params,
scale_params,
shape_params,
expr_params=None,
return_keypoints=False,
return_joint_coords=False,
return_model_params=False,
return_joint_rotations=False,
):
# Align everything to the static buffers
dt = self.scale_mean.dtype
global_trans = global_trans.to(dt)
global_rot = global_rot.to(dt)
body_pose_params = body_pose_params.to(dt)
if hand_pose_params is not None:
hand_pose_params = hand_pose_params.to(dt)
scale_params = scale_params.to(dt)
shape_params = shape_params.to(dt)
if expr_params is not None:
expr_params = expr_params.to(dt)
if self.enable_hand_model:
# Transfer wrist-centric predictions to the body.
global_rot_ori = global_rot.clone()
global_trans_ori = global_trans.clone()
global_rot = rotmat_to_euler(
"xyz",
euler_to_rotmat("xyz", global_rot_ori) @ cast_to_input(self.local_to_world_wrist, global_rot_ori, copy=False),
)
right_wrist_coords = cast_to_input(self.right_wrist_coords, global_rot, copy=False)
root_coords = cast_to_input(self.root_coords, global_rot, copy=False)
global_trans = (
-(
euler_to_rotmat("xyz", global_rot)
@ (right_wrist_coords - root_coords)
+ root_coords
)
+ global_trans_ori
)
body_pose_params = body_pose_params[..., :130]
# Convert from scale and shape params to actual scales and vertices
# Add singleton batches in case...
if len(scale_params.shape) == 1:
scale_params = scale_params[None]
if len(shape_params.shape) == 1:
shape_params = shape_params[None]
# Convert scale...
scale_mean = cast_to_input(self.scale_mean, scale_params, copy=False)
scale_comps = cast_to_input(self.scale_comps, scale_params, copy=False)
scales = scale_mean[None, :] + scale_params @ scale_comps
# Now, figure out the pose.
## 10 here is because it's more stable to optimize global translation in meters.
full_pose_params = torch.cat([global_trans * 10, global_rot, body_pose_params], dim=1) # B x 127
## Put in hands
if hand_pose_params is not None:
full_pose_params = self.replace_hands_in_pose(
full_pose_params, hand_pose_params
)
model_params = torch.cat([full_pose_params, scales], dim=1)
if self.enable_hand_model:
# Zero out non-hand parameters
model_params[:, self.nonhand_param_idxs.to(model_params.device)] = 0
curr_skinned_verts, curr_skel_state = self.mhr(
shape_params, model_params, expr_params
)
curr_joint_coords, curr_joint_quats, _ = torch.split(
curr_skel_state, [3, 4, 1], dim=2
)
curr_skinned_verts = curr_skinned_verts / 100
curr_joint_coords = curr_joint_coords / 100
curr_joint_rots = unitquat_to_rotmat(curr_joint_quats)
# Prepare returns
to_return = [curr_skinned_verts]
if return_keypoints:
# Get sapiens 308 keypoints
model_vert_joints = torch.cat(
[curr_skinned_verts, curr_joint_coords], dim=1
) # B x (num_verts + 127) x 3
kp_map = cast_to_input(self.keypoint_mapping, model_vert_joints, copy=False)
model_keypoints_pred = (
(kp_map @ model_vert_joints.permute(1, 0, 2).flatten(1, 2))
.reshape(-1, model_vert_joints.shape[0], 3)
.permute(1, 0, 2)
)
if self.enable_hand_model:
# Zero out everything except for the right hand
model_keypoints_pred[:, :21] = 0
model_keypoints_pred[:, 42:] = 0
to_return = to_return + [model_keypoints_pred]
if return_joint_coords:
to_return = to_return + [curr_joint_coords]
if return_model_params:
to_return = to_return + [model_params]
if return_joint_rotations:
to_return = to_return + [curr_joint_rots]
if isinstance(to_return, list) and len(to_return) == 1:
return to_return[0]
else:
return tuple(to_return)
def forward(self, x: torch.Tensor, init_estimate: Optional[torch.Tensor] = None, intermediate: bool = False):
"""
Args:
x: pose token with shape [B, C], usually C=DECODER.DIM
init_estimate: [B, self.npose]
intermediate: when True, the caller only needs the keypoints/pose
outputs needed by the per-layer keypoint-token update path —
vertex output is suppressed so `camera_project` skips the
18439-vertex perspective projection on intermediate decoder
layers. The final layer must call with intermediate=False.
"""
batch_size = x.shape[0]
pred = self.proj(x)
if init_estimate is not None:
pred = pred + init_estimate
# From pred, we want to pull out individual predictions.
## First, get globals
### Global rotation is first 6.
count = 6
global_rot_6d = pred[:, :count]
global_rot_rotmat = rot6d_to_rotmat(global_rot_6d) # B x 3 x 3
global_rot_euler = rotmat_to_euler("ZYX", global_rot_rotmat) # B x 3
global_trans = torch.zeros_like(global_rot_euler)
## Next, get body pose.
### Hold onto raw, continuous version for iterative correction.
pred_pose_cont = pred[:, count : count + self.body_cont_dim]
count += self.body_cont_dim
### Convert to eulers (and trans)
pred_pose_euler = compact_cont_to_model_params_body(pred_pose_cont)
### Zero-out hands
pred_pose_euler[:, mhr_param_hand_idxs] = 0
### Zero-out jaw
pred_pose_euler[:, -3:] = 0
## Get remaining parameters
pred_shape = pred[:, count : count + self.num_shape_comps]
count += self.num_shape_comps
pred_scale = pred[:, count : count + self.num_scale_comps]
count += self.num_scale_comps
pred_hand = pred[:, count : count + self.num_hand_comps * 2]
count += self.num_hand_comps * 2
pred_face = pred[:, count : count + self.num_face_comps] * 0
count += self.num_face_comps
# Run everything through mhr
output = self.mhr_forward(
global_trans=global_trans,
global_rot=global_rot_euler,
body_pose_params=pred_pose_euler,
hand_pose_params=pred_hand,
scale_params=pred_scale,
shape_params=pred_shape,
expr_params=pred_face,
return_keypoints=True,
return_joint_coords=True,
return_model_params=True,
return_joint_rotations=True,
)
# Some existing code to get joints and fix camera system
verts, j3d, jcoords, mhr_model_params, joint_global_rots = output
j3d = j3d[:, :70] # 308 --> 70 keypoints
# Intermediate decoder layers only consume pred_keypoints_3d via the
# keypoint-token update path; suppress verts so camera_project skips
# the 18439-vertex perspective projection.
if intermediate:
verts = None
if verts is not None:
verts[..., [1, 2]] *= -1 # Camera system difference
j3d[..., [1, 2]] *= -1 # Camera system difference
if jcoords is not None:
jcoords[..., [1, 2]] *= -1
# Head-MLP outputs are promoted to fp32 here so the external
# pose_output["mhr"] contract has a stable dtype regardless of what
# the head ran at (fp16/bf16 for speed). MHR-derived outputs are
# already fp32 from MHR's math layers.
output = {
"pred_pose_raw": torch.cat([global_rot_6d, pred_pose_cont], dim=1).float(),
"pred_pose_rotmat": None,
"global_rot": global_rot_euler.float(),
"body_pose": pred_pose_euler.float(),
"shape": pred_shape.float(),
"scale": pred_scale.float(),
"hand": pred_hand.float(),
"face": pred_face.float(),
"pred_keypoints_3d": j3d.reshape(batch_size, -1, 3),
"pred_vertices": verts.reshape(batch_size, -1, 3) if verts is not None else None,
"pred_joint_coords": jcoords.reshape(batch_size, -1, 3) if jcoords is not None else None,
"joint_global_rots": joint_global_rots,
"mhr_model_params": mhr_model_params,
}
return output
comfy/ldm/sam3d_body/mhr/mhr_rig.py
+246
View File
@@ -0,0 +1,246 @@
# Adapted from facebookresearch/MHR (Apache 2.0):
# https://github.com/facebookresearch/MHR/blob/main/mhr/mhr.py
# Skinning ops follow facebookincubator/momentum (Apache 2.0) — formulas
# verbatim from the upstream mhr_model.pt
# (pymomentum.{skel_state,quaternion,backend.skel_state_backend}).
# Original Copyright (c) Meta Platforms, Inc. and affiliates.
import torch
import torch.nn as nn
import torch.nn.functional as F
from comfy.ops import cast_to_input
from .mhr_utils import batch6DFromXYZ
_LN2 = 0.6931471824645996
# Half-angle cos/sin are computed on the
# whole (..., 3) at once and concatenated to [cr, cp, cy, sr, sp, sy]; _EQ_I then
# picks the three factors of each term, reproducing:
# x = sr*cp*cy - cr*sp*sy z = cr*cp*sy - sr*sp*cy
# y = cr*sp*cy + sr*cp*sy w = cr*cp*cy + sr*sp*sy
def _euler_xyz_to_quat(angles):
"""(roll, pitch, yaw) -> quaternion (x, y, z, w). Matches pymomentum.quaternion.euler_xyz_to_quaternion."""
half = angles * 0.5
c = torch.cos(half)
s = torch.sin(half)
cr, cp, cy = c.unbind(-1)
sr, sp, sy = s.unbind(-1)
return torch.stack([
sr * cp * cy - cr * sp * sy,
cr * sp * cy + sr * cp * sy,
cr * cp * sy - sr * sp * cy,
cr * cp * cy + sr * sp * sy,
], dim=-1)
# Hamilton product as gather + 3 adds. Each output component is a 4-term sum;
# _QM_P1/_QM_P2 pick the operands and _QM_S the signs, reproducing:
# x = w1*x2 + x1*w2 + y1*z2 - z1*y2
# y = w1*y2 - x1*z2 + y1*w2 + z1*x2
# z = w1*z2 + x1*y2 - y1*x2 + z1*w2
# w = w1*w2 - x1*x2 - y1*y2 - z1*z2
def _quat_multiply(q1, q2):
x1, y1, z1, w1 = q1.unbind(-1)
x2, y2, z2, w2 = q2.unbind(-1)
return torch.stack([
w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2,
w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2,
w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2,
w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2,
], dim=-1)
def _quat_rotate(q, v):
"""Rotate v by unit quaternion q (xyzw). v + 2 * (axis x v * w + axis x (axis x v))."""
axis = q[..., :3]
r = q[..., 3:4]
av = torch.cross(axis, v, dim=-1)
aav = torch.cross(axis, av, dim=-1)
return v + 2.0 * (av * r + aav)
def _skel_multiply(s1, s2):
"""Compose two skel states (..., 8). Returns parent ∘ child.
Mirrors pymomentum.skel_state.multiply: both quaternions are renormalized
before composition. With many FK levels the previously-normalized quats
drift in ULPs; upstream renormalizes defensively, so we do too to stay
bit-close to its outputs.
"""
t1, sc1 = s1[..., :3], s1[..., 7:8]
t2, sc2 = s2[..., :3], s2[..., 7:8]
q1 = F.normalize(s1[..., 3:7], p=2, dim=-1, eps=1e-12)
q2 = F.normalize(s2[..., 3:7], p=2, dim=-1, eps=1e-12)
t_res = t1 + sc1 * _quat_rotate(q1, t2)
q_res = _quat_multiply(q1, q2)
s_res = sc1 * sc2
return torch.cat([t_res, q_res, s_res], dim=-1)
def _skel_transform_points(skel_state, points):
"""Apply skel_state (..., 8) to points (..., 3): t + q * (s * points).
Assumes the quaternion in skel_state is already unit-norm. Callers that
can't guarantee that should normalize first.
"""
t = skel_state[..., :3]
q = skel_state[..., 3:7]
s = skel_state[..., 7:8]
return t + _quat_rotate(q, s * points)
def _global_skel_state_from_local(local, pmi_levels):
"""FK walk in fp64 (matches upstream's use_double_precision=True path).
`pmi_levels` is a precomputed list of (source_idx, target_idx) tensor pairs,
one per BFS level. Avoids per-call torch.split + tolist() sync.
"""
orig_dtype = local.dtype
g = local.to(torch.float64).clone()
for source, target in pmi_levels:
parent = g.index_select(-2, target)
child = g.index_select(-2, source)
g.index_copy_(-2, source, _skel_multiply(parent, child))
return g.to(orig_dtype)
class MHRRig(nn.Module):
"""Plain-PyTorch reimplementation of Meta's MHR rig.
All math runs in fp32 (FK upcast to fp64 internally, matching upstream's
use_double_precision=True backend) regardless of the host model's dtype.
"""
NUM_VERTS = 18439
NUM_JOINTS = 127
NUM_LBS_TRIPLETS = 51337
NUM_IDENTITY = 45
NUM_EXPR = 72
PARAM_TRANSFORM_IN = 249 # = model_parameters(204) + identity_coeffs(45)
PARAM_TRANSFORM_OUT = 889 # = NUM_JOINTS * 7
POSE_CORR_IN = 750 # = (NUM_JOINTS - 2) * 6
POSE_CORR_HIDDEN = 3000
POSE_CORR_SPARSE_NNZ = 53136
def __init__(self, device=None):
super().__init__()
# All buffers are populated by load_state_dict from the `mhr.*` keys
def _p(*shape, dtype=torch.float32):
return nn.Parameter(torch.empty(*shape, dtype=dtype, device=device), requires_grad=False)
def _b(name, *shape, dtype):
self.register_buffer(name, torch.empty(*shape, dtype=dtype, device=device))
self.base_shape = _p(self.NUM_VERTS, 3)
self.identity_basis = _p(self.NUM_IDENTITY, self.NUM_VERTS, 3)
self.expr_basis = _p(self.NUM_EXPR, self.NUM_VERTS, 3)
self.param_transform = _p(self.PARAM_TRANSFORM_OUT, self.PARAM_TRANSFORM_IN)
self.skel_joint_translation_offsets = _p(self.NUM_JOINTS, 3)
self.skel_joint_prerotations = _p(self.NUM_JOINTS, 4)
_b("skel_joint_parents", self.NUM_JOINTS, dtype=torch.int32)
_b("skel_pmi", 2, 266, dtype=torch.int64)
_b("skel_pmi_buffer_sizes", 4, dtype=torch.int64)
self.lbs_inverse_bind_pose = _p(self.NUM_JOINTS, 8)
self.lbs_skin_weights = _p(self.NUM_LBS_TRIPLETS)
_b("lbs_skin_indices", self.NUM_LBS_TRIPLETS, dtype=torch.int32)
_b("lbs_vert_indices", self.NUM_LBS_TRIPLETS, dtype=torch.int64)
_b("pose_corr_sparse_indices", 2, self.POSE_CORR_SPARSE_NNZ, dtype=torch.int64)
self.pose_corr_sparse_weight = _p(self.POSE_CORR_SPARSE_NNZ)
_b("pose_corr_sparse_shape", 2, dtype=torch.int64)
self.pose_corr_weight = _p(self.NUM_VERTS * 3, self.POSE_CORR_HIDDEN)
self.pose_corr_bias = None
self._pmi_sizes = None
self._pose_corr_shape = None
self.register_load_state_dict_post_hook(self._set_pmi_sizes)
def _set_pmi_sizes(self, module, incompatible_keys):
self._pmi_sizes = tuple(self.skel_pmi_buffer_sizes.tolist())
self._pose_corr_shape = tuple(self.pose_corr_sparse_shape.tolist())
def forward(self, identity_coeffs, model_parameters, expr_coeffs, apply_correctives: bool = True):
dtype = self.base_shape.dtype
identity_coeffs = identity_coeffs.to(dtype)
model_parameters = model_parameters.to(dtype)
expr_coeffs = expr_coeffs.to(dtype)
B = identity_coeffs.shape[0]
base_shape = cast_to_input(self.base_shape, identity_coeffs, copy=False)
identity_basis = cast_to_input(self.identity_basis, identity_coeffs, copy=False)
identity_rest = base_shape + torch.einsum("nvd,bn->bvd", identity_basis, identity_coeffs)
cat_in = torch.cat([model_parameters, torch.zeros_like(identity_coeffs)], dim=1)
joint_parameters = torch.einsum("dn,bn->bd", cast_to_input(self.param_transform, cat_in, copy=False), cat_in)
jp = joint_parameters.view(B, self.NUM_JOINTS, 7)
local_t = jp[..., :3] + cast_to_input(self.skel_joint_translation_offsets, jp, copy=False).unsqueeze(0)
local_q = _euler_xyz_to_quat(jp[..., 3:6])
local_q = _quat_multiply(cast_to_input(self.skel_joint_prerotations, local_q, copy=False).unsqueeze(0), local_q)
local_s = torch.exp(jp[..., 6:7] * _LN2)
local_state = torch.cat([local_t, local_q, local_s], dim=-1)
skel_state = _global_skel_state_from_local(local_state, self._pmi_levels(local_state.device))
face_expr = torch.einsum("nvd,bn->bvd", cast_to_input(self.expr_basis, expr_coeffs, copy=False), expr_coeffs)
unposed = identity_rest + face_expr
if apply_correctives:
unposed = unposed + self._pose_correctives(joint_parameters)
verts = self._skin(skel_state, unposed)
return verts, skel_state
def _pose_correctives(self, joint_parameters):
B = joint_parameters.shape[0]
jp = joint_parameters.view(B, self.NUM_JOINTS, 7)
# Joints [2:] only — root and one more skipped. Take Euler XYZ (cols 3:6).
feat = batch6DFromXYZ(jp[:, 2:, 3:6], return_9D=False) # (B, 125, 6)
feat[..., 0] -= 1.0
feat[..., 4] -= 1.0
feat = feat.flatten(1, 2) # (B, 750)
h = (self._sparse_w(feat) @ feat.T).T # (B, 3000)
h = F.relu(h)
out = F.linear(h, cast_to_input(self.pose_corr_weight, h, copy=False), self.pose_corr_bias) # (B, 55317)
return out.view(B, self.NUM_VERTS, 3)
def _pmi_levels(self, device):
if self._pmi_sizes is None:
raise RuntimeError("MHR rig weights have not been loaded")
pmi = self.skel_pmi.to(device=device)
return [(part[0], part[1]) for part in torch.split(pmi, self._pmi_sizes, dim=1)]
def _sparse_w(self, ref):
if self._pose_corr_shape is None:
raise RuntimeError("MHR rig weights have not been loaded")
w = cast_to_input(self.pose_corr_sparse_weight, ref, copy=False)
# PyTorch 2.12 warns unless invariant checking is explicitly scoped.
with torch.sparse.check_sparse_tensor_invariants():
return torch.sparse_coo_tensor(
self.pose_corr_sparse_indices.to(device=ref.device),
w,
self._pose_corr_shape,
).coalesce()
def _skin(self, skel_state, rest_verts):
B = skel_state.shape[0]
ibp = cast_to_input(self.lbs_inverse_bind_pose, skel_state, copy=False).unsqueeze(0).expand(B, self.NUM_JOINTS, 8)
joint_xform = _skel_multiply(skel_state, ibp)
norm_q = F.normalize(joint_xform[..., 3:7], p=2, dim=-1, eps=1e-12)
joint_xform = torch.cat([joint_xform[..., :3], norm_q, joint_xform[..., 7:8]], dim=-1)
sk_idx = self.lbs_skin_indices.to(device=rest_verts.device, dtype=torch.long)
v_idx = self.lbs_vert_indices.to(device=rest_verts.device)
w = cast_to_input(self.lbs_skin_weights, rest_verts, copy=False)
per_triplet_xform = joint_xform.index_select(-2, sk_idx) # (B, 51337, 8)
per_triplet_rest = rest_verts.index_select(-2, v_idx) # (B, 51337, 3)
contrib = _skel_transform_points(per_triplet_xform, per_triplet_rest) * w.unsqueeze(0).unsqueeze(-1)
out = torch.zeros(B, self.NUM_VERTS, 3, dtype=rest_verts.dtype, device=rest_verts.device)
out.index_add_(-2, v_idx, contrib)
return out
comfy/ldm/sam3d_body/mhr/mhr_utils.py
+215
View File
@@ -0,0 +1,215 @@
# MHR (Meta Human Rig) parameter packing/unpacking. The 6D-rotation helpers
# (batch6DFromXYZ, batchXYZfrom6D) are the continuity
# representation from Zhou et al., "On the Continuity of Rotation
# Representations in Neural Networks" (CVPR 2019, https://arxiv.org/abs/1812.07035),
# implementations from papagina/RotationContinuity:
# https://github.com/papagina/RotationContinuity/blob/758b0ce5/shapenet/code/tools.py
# The compact_cont_to_model_params_* functions are MHR-rig-specific glue.
import torch
import torch.nn.functional as F
def rotation_angle_difference(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor:
"""
Compute the angle difference (magnitude) between two batches of SO(3) rotation matrices.
Args:
A: Tensor of shape (*, 3, 3), batch of rotation matrices.
B: Tensor of shape (*, 3, 3), batch of rotation matrices.
Returns:
Tensor of shape (*,), angle differences in radians.
"""
# Compute relative rotation matrix
R_rel = torch.matmul(A, B.transpose(-2, -1)) # (B, 3, 3)
# Compute trace of relative rotation
trace = R_rel[..., 0, 0] + R_rel[..., 1, 1] + R_rel[..., 2, 2] # (B,)
# Compute angle using the trace formula
cos_theta = (trace - 1) / 2
# Clamp for numerical stability
cos_theta_clamped = torch.clamp(cos_theta, -1.0, 1.0)
# Compute angle difference
angle = torch.acos(cos_theta_clamped)
return angle
def fix_wrist_euler(
wrist_xzy, limits_x=(-2.2, 1.0), limits_z=(-2.2, 1.5), limits_y=(-1.2, 1.5)
):
"""
wrist_xzy: B x 2 x 3 (X, Z, Y angles)
Returns: Fixed angles within joint limits
"""
x, z, y = wrist_xzy[..., 0], wrist_xzy[..., 1], wrist_xzy[..., 2]
x_alt = torch.atan2(torch.sin(x + torch.pi), torch.cos(x + torch.pi))
z_alt = torch.atan2(torch.sin(-(z + torch.pi)), torch.cos(-(z + torch.pi)))
y_alt = torch.atan2(torch.sin(y + torch.pi), torch.cos(y + torch.pi))
# Calculate L2 violation distance
def calc_violation(val, limits):
below = torch.clamp(limits[0] - val, min=0.0)
above = torch.clamp(val - limits[1], min=0.0)
return below**2 + above**2
violation_orig = (
calc_violation(x, limits_x)
+ calc_violation(z, limits_z)
+ calc_violation(y, limits_y)
)
violation_alt = (
calc_violation(x_alt, limits_x)
+ calc_violation(z_alt, limits_z)
+ calc_violation(y_alt, limits_y)
)
# Use alternative where it has lower L2 violation
use_alt = violation_alt < violation_orig
# Stack alternative and apply mask
wrist_xzy_alt = torch.stack([x_alt, z_alt, y_alt], dim=-1)
result = torch.where(use_alt.unsqueeze(-1), wrist_xzy_alt, wrist_xzy)
return result
# https://github.com/papagina/RotationContinuity/blob/758b0ce551c06372cab7022d4c0bdf331c89c696/shapenet/code/tools.py
def batch6DFromXYZ(r, return_9D=False):
"""
Generate a matrix representing a rotation defined by a XYZ-Euler
rotation.
Args:
r: ... x 3 rotation vectors
Returns:
... x 6
"""
rc = torch.cos(r)
rs = torch.sin(r)
cx = rc[..., 0]
cy = rc[..., 1]
cz = rc[..., 2]
sx = rs[..., 0]
sy = rs[..., 1]
sz = rs[..., 2]
result = torch.empty(list(r.shape[:-1]) + [3, 3], dtype=r.dtype, device=r.device)
result[..., 0, 0] = cy * cz
result[..., 0, 1] = -cx * sz + sx * sy * cz
result[..., 0, 2] = sx * sz + cx * sy * cz
result[..., 1, 0] = cy * sz
result[..., 1, 1] = cx * cz + sx * sy * sz
result[..., 1, 2] = -sx * cz + cx * sy * sz
result[..., 2, 0] = -sy
result[..., 2, 1] = sx * cy
result[..., 2, 2] = cx * cy
if not return_9D:
return torch.cat([result[..., :, 0], result[..., :, 1]], dim=-1)
else:
return result
# https://github.com/papagina/RotationContinuity/blob/758b0ce551c06372cab7022d4c0bdf331c89c696/shapenet/code/tools.py#L82
def batchXYZfrom6D(poses):
# Args: poses: ... x 6, where "6" is the combined first and second columns
# First, get the rotaiton matrix
x_raw = poses[..., :3]
y_raw = poses[..., 3:]
x = F.normalize(x_raw, dim=-1)
z = torch.cross(x, y_raw, dim=-1)
z = F.normalize(z, dim=-1)
y = torch.cross(z, x, dim=-1)
matrix = torch.stack([x, y, z], dim=-1) # ... x 3 x 3
# Now get it into euler
# https://github.com/papagina/RotationContinuity/blob/758b0ce551c06372cab7022d4c0bdf331c89c696/shapenet/code/tools.py#L412
sy = torch.sqrt(
matrix[..., 0, 0] * matrix[..., 0, 0] + matrix[..., 1, 0] * matrix[..., 1, 0]
)
singular = sy < 1e-6
singular = singular.float()
x = torch.atan2(matrix[..., 2, 1], matrix[..., 2, 2])
y = torch.atan2(-matrix[..., 2, 0], sy)
z = torch.atan2(matrix[..., 1, 0], matrix[..., 0, 0])
xs = torch.atan2(-matrix[..., 1, 2], matrix[..., 1, 1])
ys = torch.atan2(-matrix[..., 2, 0], sy)
zs = matrix[..., 1, 0] * 0
out_euler = torch.zeros_like(matrix[..., 0])
out_euler[..., 0] = x * (1 - singular) + xs * singular
out_euler[..., 1] = y * (1 - singular) + ys * singular
out_euler[..., 2] = z * (1 - singular) + zs * singular
return out_euler
_HAND_DOFS = (3, 1, 1, 3, 1, 1, 3, 1, 1, 3, 1, 1, 2, 3, 1, 1)
_HAND_CONT_3DOF_MASK = [k == 3 for k in _HAND_DOFS for _ in range(2 * k)]
_HAND_CONT_1DOF_MASK = [k in (1, 2) for k in _HAND_DOFS for _ in range(2 * k)]
_HAND_MODEL_3DOF_MASK = [k == 3 for k in _HAND_DOFS for _ in range(k)]
_HAND_MODEL_1DOF_MASK = [k in (1, 2) for k in _HAND_DOFS for _ in range(k)]
def compact_cont_to_model_params_hand(hand_cont):
# These are ordered by joint, not model params ^^
# Convert hand_cont to eulers
## First for 3DoFs
hand_cont_threedofs = hand_cont[..., _HAND_CONT_3DOF_MASK].unflatten(-1, (-1, 6))
hand_model_params_threedofs = batchXYZfrom6D(hand_cont_threedofs).flatten(-2, -1)
## Next for 1DoFs
hand_cont_onedofs = hand_cont[..., _HAND_CONT_1DOF_MASK].unflatten(
-1, (-1, 2)
) # (sincos)
hand_model_params_onedofs = torch.atan2(
hand_cont_onedofs[..., -2], hand_cont_onedofs[..., -1]
)
# Finally, assemble into a 27-dim vector, ordered by joint, then XYZ.
hand_model_params = torch.zeros(*hand_cont.shape[:-1], 27, dtype=hand_cont.dtype, device=hand_cont.device)
hand_model_params[..., _HAND_MODEL_3DOF_MASK] = hand_model_params_threedofs
hand_model_params[..., _HAND_MODEL_1DOF_MASK] = hand_model_params_onedofs
return hand_model_params
# fmt: off
_BODY_3DOF_IDXS = ((0, 2, 4), (6, 8, 10), (12, 13, 14), (15, 16, 17), (18, 19, 20), (21, 22, 23), (24, 25, 26), (27, 28, 29), (34, 35, 36), (37, 38, 39), (44, 45, 46), (53, 54, 55), (64, 65, 66), (85, 69, 73), (86, 70, 79), (87, 71, 82), (88, 72, 76), (91, 92, 93), (112, 96, 100), (113, 97, 106), (114, 98, 109), (115, 99, 103), (130, 131, 132))
_BODY_1DOF_ROT_IDXS = (1, 3, 5, 7, 9, 11, 30, 31, 32, 33, 40, 41, 42, 43, 47, 48, 49, 50, 51, 52, 56, 57, 58, 59, 60, 61, 62, 63, 67, 68, 74, 75, 77, 78, 80, 81, 83, 84, 89, 90, 94, 95, 101, 102, 104, 105, 107, 108, 110, 111, 116, 117, 118, 119, 120, 121, 122, 123)
_BODY_1DOF_TRANS_IDXS = (124, 125, 126, 127, 128, 129)
_BODY_3DOF_FLAT_IDXS = tuple(i for group in _BODY_3DOF_IDXS for i in group)
# fmt: on
def compact_cont_to_model_params_body(body_pose_cont):
num_3dof_angles = len(_BODY_3DOF_IDXS) * 3
num_1dof_angles = len(_BODY_1DOF_ROT_IDXS)
# Get subsets
body_cont_3dofs = body_pose_cont[..., : 2 * num_3dof_angles]
body_cont_1dofs = body_pose_cont[..., 2 * num_3dof_angles : 2 * num_3dof_angles + 2 * num_1dof_angles]
body_cont_trans = body_pose_cont[..., 2 * num_3dof_angles + 2 * num_1dof_angles :]
# Convert conts to model params
## First for 3dofs
body_cont_3dofs = body_cont_3dofs.unflatten(-1, (-1, 6))
body_params_3dofs = batchXYZfrom6D(body_cont_3dofs).flatten(-2, -1)
## Next for 1dofs
body_cont_1dofs = body_cont_1dofs.unflatten(-1, (-1, 2)) # (sincos)
body_params_1dofs = torch.atan2(body_cont_1dofs[..., -2], body_cont_1dofs[..., -1])
## Nothing to do for trans
body_params_trans = body_cont_trans
# Put them together
body_pose_params = torch.zeros(*body_pose_cont.shape[:-1], 133, dtype=body_pose_cont.dtype, device=body_pose_cont.device)
body_pose_params[..., list(_BODY_3DOF_FLAT_IDXS)] = body_params_3dofs
body_pose_params[..., list(_BODY_1DOF_ROT_IDXS)] = body_params_1dofs
body_pose_params[..., list(_BODY_1DOF_TRANS_IDXS)] = body_params_trans
return body_pose_params
# Hand indices into the 133-dim body-pose vector.
mhr_param_hand_idxs = list(range(62, 116))
comfy/ldm/sam3d_body/model/camera_modules.py
+139
View File
@@ -0,0 +1,139 @@
import math
import torch
import torch.nn.functional as F
from comfy.ldm.cascade.common import LayerNorm2d_op
from torch import nn
from typing import List, Optional, Tuple, Union
from ..utils import perspective_projection
from .transformer import MLP
class CameraEncoder(nn.Module):
def __init__(self, embed_dim: int, patch_size: int = 14, device=None, dtype=None, operations=None):
super().__init__()
self.patch_size = patch_size
self.embed_dim = embed_dim
self.camera = FourierPositionEncoding(n=3, num_bands=16, max_resolution=64)
self.conv = operations.Conv2d(embed_dim + 99, embed_dim, kernel_size=1, bias=False, device=device, dtype=dtype)
self.norm = LayerNorm2d_op(operations)(embed_dim, device=device, dtype=dtype)
def forward(self, img_embeddings: torch.Tensor, rays: torch.Tensor):
B, D, _h, _w = img_embeddings.shape
scale = 1 / self.patch_size
rays = F.interpolate(rays, scale_factor=(scale, scale), mode="bilinear", align_corners=False, antialias=True)
rays = rays.permute(0, 2, 3, 1).contiguous() # [b, h, w, 2]
rays = torch.cat([rays, torch.ones_like(rays[..., :1])], dim=-1)
rays_embeddings = self.camera(pos=rays.reshape(B, -1, 3)) # (bs, N, 99): rays fourier embedding
rays_embeddings = rays_embeddings.reshape(B, _h, _w, -1).permute(0, 3, 1, 2).contiguous()
z = torch.cat([img_embeddings, rays_embeddings], dim=1)
return self.norm(self.conv(z))
class FourierPositionEncoding(nn.Module):
"""Sin/cos Fourier features for ray positions"""
def __init__(self, n: int, num_bands: int, max_resolution: int):
super().__init__()
self.num_bands = num_bands
self.max_resolution = [max_resolution] * n
def forward(self, pos: torch.Tensor):
fourier_pos_enc = _generate_fourier_features(pos, num_bands=self.num_bands, max_resolution=self.max_resolution)
return fourier_pos_enc
def _generate_fourier_features(pos: torch.Tensor, num_bands: int, max_resolution: List[int], min_freq: float = 1.0):
b, n = pos.shape[:2]
freq_bands = torch.stack([torch.linspace(start=min_freq, end=res / 2, steps=num_bands, device=pos.device, dtype=pos.dtype) for res in max_resolution], dim=0)
per_pos_features = pos.unsqueeze(-1) * freq_bands.unsqueeze(0).unsqueeze(0)
per_pos_features = per_pos_features.reshape(b, n, -1)
# Sin-Cos
per_pos_features = torch.cat([torch.sin(math.pi * per_pos_features), torch.cos(math.pi * per_pos_features)], dim=-1)
# Concat with initial pos
per_pos_features = torch.cat([pos, per_pos_features], dim=-1)
return per_pos_features
class PerspectiveHead(nn.Module):
"""
Predict camera translation (s, tx, ty) and perform full-perspective 2D reprojection (CLIFF/CameraHMR setup).
"""
def __init__(self, input_dim: int, img_size: Union[int, Tuple[int, int]], # model input size (W, H)
mlp_depth: int = 1, mlp_channel_div_factor: int = 8, default_scale_factor: float = 1.0,
device=None, dtype=None, operations=None
):
super().__init__()
# Metadata to compute 3D skeleton and 2D reprojection
self.img_size = img_size if isinstance(img_size, tuple) else (img_size, img_size)
self.ncam = 3 # (s, tx, ty)
self.default_scale_factor = default_scale_factor
self.proj = MLP(
input_dim=input_dim,
hidden_dim=input_dim // mlp_channel_div_factor,
output_dim=self.ncam,
num_layers=mlp_depth,
device=device,
dtype=dtype,
operations=operations,
)
def forward(self, x: torch.Tensor, init_estimate: Optional[torch.Tensor] = None):
"""
Args:
x: pose token with shape [B, C], usually C=DECODER.DIM
init_estimate: [B, self.ncam]
"""
pred_cam = self.proj(x)
if init_estimate is not None:
pred_cam = pred_cam + init_estimate
return pred_cam
def perspective_projection(
self,
points_3d: torch.Tensor,
pred_cam: torch.Tensor,
bbox_center: torch.Tensor, # [N, 2], in original image space (w, h)
bbox_size: torch.Tensor, # [N,], in original image space
cam_int: torch.Tensor, # [B, 3, 3]
):
batch_size = points_3d.shape[0]
pred_cam = pred_cam.clone()
pred_cam[..., [0, 2]] *= -1 # Camera system difference
# Compute camera translation: (scale, x, y) --> (x, y, depth)
# depth ~= f / s, Note that f is in the NDC space
s, tx, ty = pred_cam[:, 0], pred_cam[:, 1], pred_cam[:, 2]
bs = bbox_size * s * self.default_scale_factor + 1e-8
focal_length = cam_int[:, 0, 0]
tz = 2 * focal_length / bs
cx = 2 * (bbox_center[:, 0] - cam_int[:, 0, 2]) / bs
cy = 2 * (bbox_center[:, 1] - cam_int[:, 1, 2]) / bs
pred_cam_t = torch.stack([tx + cx, ty + cy, tz], dim=-1)
# Compute camera translation
j3d_cam = points_3d + pred_cam_t.unsqueeze(1)
# Projection to the image plane, note that the projection output is in original image space now.
j2d = perspective_projection(j3d_cam, cam_int)
return {
"pred_keypoints_2d": j2d.reshape(batch_size, -1, 2),
"pred_keypoints_2d_depth": j3d_cam.reshape(batch_size, -1, 3)[:, :, 2],
"pred_cam_t": pred_cam_t, "focal_length": focal_length,
}
comfy/ldm/sam3d_body/model/model.py
+1160
View File
File diff suppressed because it is too large Load Diff
comfy/ldm/sam3d_body/model/prompt.py
+261
View File
@@ -0,0 +1,261 @@
"""SAM 3D Body prompt pipeline: encode (keypoint, mask) prompts and run them
through a cross-attention transformer decoder over (token, image) pairs.
Both adapted from the SAM-style prompt path (Meta, Apache 2.0):
https://github.com/facebookresearch/segment-anything
"""
from typing import Optional, Tuple
import torch
import torch.nn as nn
from comfy.ldm.cascade.common import LayerNorm2d_op
from comfy.ops import cast_to_input
from comfy.ldm.sam3.sam import PositionEmbeddingRandom
from .transformer import TransformerDecoderLayer
class PromptEncoder(nn.Module):
def __init__(
self,
embed_dim: int,
num_body_joints: int,
device=None,
dtype=None,
operations=None,
) -> None:
"""
Encodes prompts for input to SAM's mask decoder.
"""
super().__init__()
self.embed_dim = embed_dim
self.num_body_joints = num_body_joints
# Keypoint prompts
self.pe_layer = PositionEmbeddingRandom(embed_dim // 2)
self.point_embeddings = nn.ModuleList(
[operations.Embedding(1, embed_dim, device=device, dtype=dtype) for _ in range(self.num_body_joints)]
)
self.not_a_point_embed = operations.Embedding(1, embed_dim, device=device, dtype=dtype)
self.invalid_point_embed = operations.Embedding(1, embed_dim, device=device, dtype=dtype)
# Mask prompt: 5-stage 2x2 strided conv downscaling to embed_dim.
LN2d = LayerNorm2d_op(operations)
mask_in_chans = 256
self.mask_downscaling = nn.Sequential(
operations.Conv2d(1, mask_in_chans // 64, kernel_size=2, stride=2, device=device, dtype=dtype),
LN2d(mask_in_chans // 64, device=device, dtype=dtype),
nn.GELU(),
operations.Conv2d(mask_in_chans // 64, mask_in_chans // 16, kernel_size=2, stride=2, device=device, dtype=dtype),
LN2d(mask_in_chans // 16, device=device, dtype=dtype),
nn.GELU(),
operations.Conv2d(mask_in_chans // 16, mask_in_chans // 4, kernel_size=2, stride=2, device=device, dtype=dtype),
LN2d(mask_in_chans // 4, device=device, dtype=dtype),
nn.GELU(),
operations.Conv2d(mask_in_chans // 4, mask_in_chans, kernel_size=2, stride=2, device=device, dtype=dtype),
LN2d(mask_in_chans, device=device, dtype=dtype),
nn.GELU(),
operations.Conv2d(mask_in_chans, embed_dim, kernel_size=1, device=device, dtype=dtype),
)
# Trained values for the gating conv and no_mask_embed are loaded from the state dict
self.no_mask_embed = operations.Embedding(1, embed_dim, device=device, dtype=dtype)
def get_dense_pe(self, size: Tuple[int, int]) -> torch.Tensor:
"""Positional encoding over the image-embedding grid; (1, C, H, W)."""
return self.pe_layer(size)
def _embed_keypoints(self, points: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
"""
Embeds point prompts.
Assuming points have been normalized to [0, 1].
Output shape [B, N, C], mask shape [B, N]
"""
assert points.min() >= 0 and points.max() <= 1
# PE compute in fp32 for precision (sin/cos of large coords), then cast back to the embedding weight dtype
weight_dtype = self.invalid_point_embed.weight.dtype
point_embedding = self.pe_layer._encode(points.to(torch.float)).to(weight_dtype)
# One gather over the stacked joint table.
joint_w = cast_to_input(torch.cat([e.weight for e in self.point_embeddings], dim=0), point_embedding, copy=False)
idx = labels.long().clamp(0, self.num_body_joints - 1)
is_joint = ((labels >= 0) & (labels < self.num_body_joints)).unsqueeze(-1)
point_embedding = point_embedding + joint_w[idx] * is_joint.to(point_embedding.dtype)
# -2/-1 zero the PE first, so the embedding replaces it outright.
invalid_w = cast_to_input(self.invalid_point_embed.weight, point_embedding, copy=False)
not_a_point_w = cast_to_input(self.not_a_point_embed.weight, point_embedding, copy=False)
point_embedding = torch.where((labels == -2).unsqueeze(-1), invalid_w, point_embedding)
point_embedding = torch.where((labels == -1).unsqueeze(-1), not_a_point_w, point_embedding)
point_mask = labels > -2
return point_embedding, point_mask
def _get_batch_size(self, keypoints: Optional[torch.Tensor], boxes: Optional[torch.Tensor], masks: Optional[torch.Tensor]) -> int:
if keypoints is not None:
return keypoints.shape[0]
elif boxes is not None:
return boxes.shape[0]
elif masks is not None:
return masks.shape[0]
else:
return 1
def forward(
self,
keypoints: Optional[torch.Tensor],
boxes: Optional[torch.Tensor] = None,
masks: Optional[torch.Tensor] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""
Embeds different types of prompts, returning both sparse and dense
embeddings.
Arguments:
keypoints (torchTensor or none): point coordinates and labels to embed.
boxes (torch.Tensor or none): boxes to embed
masks (torch.Tensor or none): masks to embed
Returns:
torch.Tensor: sparse embeddings for the points and boxes, with shape
BxNx(embed_dim), where N is determined by the number of input points
and boxes.
torch.Tensor: dense embeddings for the masks, in the shape
Bx(embed_dim)x(embed_H)x(embed_W)
"""
bs = self._get_batch_size(keypoints, boxes, masks)
ref = keypoints if keypoints is not None else boxes if boxes is not None else masks
device = ref.device if ref is not None else self.point_embeddings[0].weight.device
weight_dtype = self.invalid_point_embed.weight.dtype
sparse_embeddings = torch.empty((bs, 0, self.embed_dim), device=device, dtype=weight_dtype)
sparse_masks = torch.empty((bs, 0), device=device)
if keypoints is not None:
coords = keypoints[:, :, :2]
labels = keypoints[:, :, -1]
point_embeddings, point_mask = self._embed_keypoints(coords, labels)
sparse_embeddings = torch.cat([sparse_embeddings, point_embeddings], dim=1)
sparse_masks = torch.cat([sparse_masks, point_mask], dim=1)
return sparse_embeddings, sparse_masks
def get_mask_embeddings(self, masks: torch.Tensor, bs: int = 1, size: Tuple[int, int] = (16, 16)) -> torch.Tensor:
"""Embeds mask inputs. Caller casts both outputs to its working dtype."""
no_mask_embeddings = cast_to_input(self.no_mask_embed.weight, masks).reshape(1, -1, 1, 1).expand(bs, -1, size[0], size[1])
mask_embeddings = self.mask_downscaling(masks)
return mask_embeddings, no_mask_embeddings
class PromptableDecoder(nn.Module):
"""Cross-attention transformer decoder over (token, image) pairs."""
def __init__(
self,
dims: int,
context_dims: int,
depth: int,
num_heads: int = 8,
head_dims: int = 64,
mlp_dims: int = 1024,
repeat_pe: bool = False,
do_interm_preds: bool = False,
keypoint_token_update: bool = False,
device=None, dtype=None, operations=None,
):
super().__init__()
self.layers = nn.ModuleList(
TransformerDecoderLayer(
token_dims=dims,
context_dims=context_dims,
num_heads=num_heads,
head_dims=head_dims,
mlp_dims=mlp_dims,
repeat_pe=repeat_pe,
skip_first_pe=(i == 0),
device=device,
dtype=dtype,
operations=operations,
)
for i in range(depth)
)
self.norm_final = operations.LayerNorm(dims, eps=1e-6, device=device, dtype=dtype)
self.do_interm_preds = do_interm_preds
self.keypoint_token_update = keypoint_token_update
def forward(
self,
token_embedding: torch.Tensor,
image_embedding: torch.Tensor,
token_augment: Optional[torch.Tensor] = None,
image_augment: Optional[torch.Tensor] = None,
token_mask: Optional[torch.Tensor] = None,
token_to_pose_output_fn=None,
keypoint_token_update_fn=None,
hand_embeddings=None,
hand_augment=None,
):
"""
Args:
token_embedding: [B, N, C]
image_embedding: [B, C, H, W] -- flattened to [B, HW, C] inline
"""
# Channels-last for the transformer.
image_embedding = image_embedding.flatten(2).permute(0, 2, 1)
if image_augment is not None:
image_augment = image_augment.flatten(2).permute(0, 2, 1)
if hand_embeddings is not None:
hand_embeddings = hand_embeddings.flatten(2).permute(0, 2, 1)
hand_augment = hand_augment.flatten(2).permute(0, 2, 1)
if len(hand_augment) == 1:
# inflate batch dimension
assert len(hand_augment.shape) == 3
hand_augment = hand_augment.repeat(len(hand_embeddings), 1, 1)
all_pose_outputs = [] if self.do_interm_preds else None
if self.do_interm_preds:
assert token_to_pose_output_fn is not None
layer_idx = 0
for layer_idx, layer in enumerate(self.layers):
if hand_embeddings is None:
token_embedding, image_embedding = layer(
token_embedding, image_embedding,
token_augment, image_augment, token_mask,
)
else:
token_embedding, image_embedding = layer(
token_embedding,
torch.cat([image_embedding, hand_embeddings], dim=1),
token_augment,
torch.cat([image_augment, hand_augment], dim=1),
token_mask,
)
image_embedding = image_embedding[:, : image_augment.shape[1]]
if self.do_interm_preds and layer_idx < len(self.layers) - 1:
curr = token_to_pose_output_fn(
self.norm_final(token_embedding),
prev_pose_output=all_pose_outputs[-1] if all_pose_outputs else None,
layer_idx=layer_idx,
)
all_pose_outputs.append(curr)
if self.keypoint_token_update:
assert keypoint_token_update_fn is not None
token_embedding, token_augment, _, _ = keypoint_token_update_fn(
token_embedding, token_augment, curr, layer_idx,
)
out = self.norm_final(token_embedding)
if self.do_interm_preds:
curr = token_to_pose_output_fn(
out,
prev_pose_output=all_pose_outputs[-1] if all_pose_outputs else None,
layer_idx=layer_idx,
)
all_pose_outputs.append(curr)
return out, all_pose_outputs
return out
comfy/ldm/sam3d_body/model/transformer.py
+104
View File
@@ -0,0 +1,104 @@
from typing import Optional
import torch
import torch.nn as nn
from comfy.ldm.modules.attention import optimized_attention
class MLP(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim, num_layers, act_layer=nn.ReLU, device=None, dtype=None, operations=None):
super().__init__()
dims = [input_dim] + [hidden_dim] * (num_layers - 1) + [output_dim]
self.layers = nn.ModuleList(
operations.Linear(dims[i], dims[i + 1], device=device, dtype=dtype)
for i in range(num_layers)
)
self.act = act_layer()
def forward(self, x):
for i, layer in enumerate(self.layers):
x = self.act(layer(x)) if i < len(self.layers) - 1 else layer(x)
return x
class Attention(nn.Module):
def __init__(self, embed_dims, num_heads, query_dims=None, key_dims=None, value_dims=None, qkv_bias=True, proj_bias=True,
device=None, dtype=None, operations=None):
super().__init__()
self.query_dims = query_dims or embed_dims
self.key_dims = key_dims or embed_dims
self.value_dims = value_dims or embed_dims
self.embed_dims = embed_dims
self.num_heads = num_heads
self.head_dims = embed_dims // num_heads
lin = lambda i, o, b: operations.Linear(i, o, bias=b, device=device, dtype=dtype)
self.q_proj = lin(self.query_dims, embed_dims, qkv_bias)
self.k_proj = lin(self.key_dims, embed_dims, qkv_bias)
self.v_proj = lin(self.value_dims, embed_dims, qkv_bias)
self.proj = lin(embed_dims, self.query_dims, proj_bias)
def _split(self, x: torch.Tensor) -> torch.Tensor:
b, n, _ = x.shape
return x.reshape(b, n, self.num_heads, self.head_dims).transpose(1, 2)
def forward(self, q, k, v, attn_mask: Optional[torch.Tensor] = None):
q, k, v = self._split(self.q_proj(q)), self._split(self.k_proj(k)), self._split(self.v_proj(v))
x = optimized_attention(q, k, v, self.num_heads, mask=attn_mask, skip_reshape=True, low_precision_attention=False)
return self.proj(x)
class TransformerDecoderLayer(nn.Module):
def __init__(self, token_dims, context_dims, num_heads=8, head_dims=64, mlp_dims=1024,
repeat_pe=False, skip_first_pe=False, device=None, dtype=None, operations=None):
super().__init__()
self.repeat_pe = repeat_pe
self.skip_first_pe = skip_first_pe
ln = lambda d: operations.LayerNorm(d, eps=1e-6, device=device, dtype=dtype)
attn_dim = num_heads * head_dims
attn_kwargs = dict(embed_dims=attn_dim, num_heads=num_heads, device=device, dtype=dtype, operations=operations)
if repeat_pe:
self.ln_pe_1, self.ln_pe_2 = ln(token_dims), ln(context_dims)
self.ln1 = ln(token_dims)
self.self_attn = Attention(query_dims=token_dims, key_dims=token_dims, value_dims=token_dims, **attn_kwargs)
self.ln2_1, self.ln2_2 = ln(token_dims), ln(context_dims)
self.cross_attn = Attention(query_dims=token_dims, key_dims=context_dims, value_dims=context_dims, **attn_kwargs)
self.ln3 = ln(token_dims)
self.ffn = MLP(token_dims, mlp_dims, token_dims, num_layers=2, act_layer=nn.GELU, device=device, dtype=dtype, operations=operations)
def forward(self, x, context, x_pe=None, context_pe=None, x_mask=None):
"""x: [B, N_tokens, C], context: [B, N_ctx, C], x_mask: [B, N_tokens] or None."""
# LaPE-style PE re-norm per layer.
if self.repeat_pe and context_pe is not None:
x_pe = self.ln_pe_1(x_pe)
context_pe = self.ln_pe_2(context_pe)
# Self-attn over tokens.
if self.repeat_pe and not self.skip_first_pe and x_pe is not None:
q = k = self.ln1(x) + x_pe
v = self.ln1(x)
else:
q = k = v = self.ln1(x)
attn_mask = None
if x_mask is not None:
attn_mask = x_mask[:, :, None] @ x_mask[:, None, :]
attn_mask.diagonal(dim1=1, dim2=2).fill_(1) # avoid all-invalid rows -> nan
attn_mask = attn_mask > 0
x = x + self.self_attn(q, k, v, attn_mask=attn_mask)
# Cross-attn: tokens attend to image context.
if self.repeat_pe and context_pe is not None:
q = self.ln2_1(x) + x_pe
k = self.ln2_2(context) + context_pe
v = self.ln2_2(context)
else:
q = self.ln2_1(x)
k = v = self.ln2_2(context)
x = x + self.cross_attn(q, k, v)
x = x + self.ffn(self.ln3(x))
return x, context
comfy/ldm/sam3d_body/utils.py
+339
View File
@@ -0,0 +1,339 @@
# The bbox/affine math (xyxy<->cs, get_warp_matrices) is the standard
# top-down pose-estimation crop pipeline from MMPose (Apache 2.0):
# https://github.com/open-mmlab/mmpose — same algorithm as UDP (CVPR 2020).
from typing import Dict, Tuple
import torch
import torch.nn.functional as F
# Bbox + affine math
# All `output_size` / image-shape tuples in this block are (H, W) to match
# the torch.Size convention used everywhere else in the codebase.
def bbox_xyxy2cs(bbox, padding: float) -> Tuple[torch.Tensor, torch.Tensor]:
"""xyxy bbox -> (center, scale) with optional padding multiplier."""
bbox = torch.as_tensor(bbox, dtype=torch.float32)
dim = bbox.dim()
if dim == 1:
bbox = bbox.unsqueeze(0)
x1, y1, x2, y2 = bbox[:, 0:1], bbox[:, 1:2], bbox[:, 2:3], bbox[:, 3:4]
center = torch.cat([x1 + x2, y1 + y2], dim=1) * 0.5
scale = torch.cat([x2 - x1, y2 - y1], dim=1) * padding
if dim == 1:
return center[0], scale[0]
return center, scale
def fix_aspect_ratio(bbox_scale, aspect_ratio: float) -> torch.Tensor:
"""Pad whichever side is too narrow to hit `aspect_ratio` (w/h)."""
bbox_scale = torch.as_tensor(bbox_scale, dtype=torch.float32)
dim = bbox_scale.dim()
if dim == 1:
bbox_scale = bbox_scale.unsqueeze(0)
w, h = bbox_scale[:, 0:1], bbox_scale[:, 1:2]
out = torch.where(
w > h * aspect_ratio,
torch.cat([w, w / aspect_ratio], dim=1),
torch.cat([h * aspect_ratio, h], dim=1),
)
return out[0] if dim == 1 else out
def get_warp_matrices(centers, scales, output_size: Tuple[int, int]) -> torch.Tensor:
"""Batched 2x3 affine matrices mapping each (center, scale) bbox region to
the output box. `output_size` is (H_out, W_out). With rot=0 the MMPose
3-point fit reduces to a closed-form isotropic scale + translate.
"""
centers = torch.as_tensor(centers, dtype=torch.float32)
scales = torch.as_tensor(scales, dtype=torch.float32)
if centers.dim() == 1:
centers = centers.unsqueeze(0)
scales = scales.unsqueeze(0)
n = centers.shape[0]
src_w = scales[:, 0]
dst_h = float(output_size[0])
dst_w = float(output_size[1])
# With rot=0 the warp is just scale + translate (uniform x/y scale based
# on src_w/dst_w). The closed form drops out of MMPose's 3-point solve.
s = dst_w / src_w # (N,)
mats = torch.zeros((n, 2, 3), dtype=centers.dtype, device=centers.device)
mats[:, 0, 0] = s
mats[:, 1, 1] = s
mats[:, 0, 2] = dst_w * 0.5 - s * centers[:, 0]
mats[:, 1, 2] = dst_h * 0.5 - s * centers[:, 1]
return mats # (N, 2, 3)
def warp_affine_batched(
src_t: torch.Tensor, # (N, C, H_src, W_src) float
mats: torch.Tensor, # (N, 2, 3) float
output_size: Tuple[int, int] # (H_out, W_out)
) -> torch.Tensor:
"""Apply N forward (src->dst) 2x3 affine warps to N source images in one
grid_sample call. Kept generic over arbitrary affines (not specialized to
the scale+translate produced by `get_warp_matrices`) so callers can pass
rotated/sheared affines; the per-crop 3x3 invert is O(N) of trivial work."""
H_out, W_out = int(output_size[0]), int(output_size[1])
N, _, H_src, W_src = src_t.shape
device = src_t.device
# Invert each forward affine; grid_sample needs dst->src.
mats_t = mats.to(device=device, dtype=torch.float32)
bottom = mats_t.new_tensor([0.0, 0.0, 1.0]).expand(N, 1, 3)
mats_3 = torch.cat([mats_t, bottom], dim=1) # (N, 3, 3)
mats_inv = torch.linalg.inv(mats_3)[:, :2, :] # (N, 2, 3)
# Output pixel-center grid (i+0.5, j+0.5).
ys, xs = torch.meshgrid(
torch.arange(H_out, dtype=torch.float32, device=device) + 0.5,
torch.arange(W_out, dtype=torch.float32, device=device) + 0.5,
indexing="ij",
)
homo = torch.stack([xs, ys, torch.ones_like(xs)], dim=-1) # (H_out, W_out, 3)
src_pos = torch.einsum("nkl,ijl->nijk", mats_inv, homo) # (N, H_out, W_out, 2)
# Normalize to [-1, 1] grid_sample coords (align_corners=False).
src_pos[..., 0] = src_pos[..., 0] / W_src * 2 - 1
src_pos[..., 1] = src_pos[..., 1] / H_src * 2 - 1
return F.grid_sample(src_t, src_pos, mode="bilinear", padding_mode="zeros", align_corners=False)
# Batch construction (one prediction over N person crops from a single image)
def prepare_batch(
img, # (H, W, 3) uint8 torch tensor or list of such tensors
boxes, # (N, 4) xyxy (numpy or torch)
input_size: Tuple[int, int], # (W, H) of the model crop
bbox_padding: float = 1.25, # xyxy->cs padding multiplier (1.25 body, 0.9 hand)
aspect_ratio: float = 0.75, # w/h of the crop (0.75 matches HMR2/Sapiens)
masks=None, # optional per-person masks
masks_score=None, # optional per-person mask scores
cam_int=None, # optional camera intrinsics
) -> Dict:
"""Build the batch dict the SAM3DBody forward expects, doing the N crops in one batched `grid_sample` call."""
is_multi_image = isinstance(img, list)
if is_multi_image:
assert len(img) == boxes.shape[0]
height, width = img[0].shape[:2]
else:
height, width = img.shape[:2]
n = int(boxes.shape[0])
assert n > 0, "prepare_batch needs at least one box"
W_out, H_out = int(input_size[0]), int(input_size[1])
# Per-box bbox math (cheap, vectorized, CPU).
centers, scales = bbox_xyxy2cs(boxes, padding=bbox_padding)
# Two passes: first hits the upstream bbox aspect (e.g. 0.75 HMR2/Sapiens
# convention), second pads further if the model crop's W_out/H_out differs
# from that. When they match (common case) the second call is a no-op.
scales = fix_aspect_ratio(scales, aspect_ratio)
scales = fix_aspect_ratio(scales, W_out / H_out)
mats = get_warp_matrices(centers, scales, (H_out, W_out)) # (N, 2, 3)
# Stack source images into a contiguous (N, 3, H, W) tensor on CPU.
if is_multi_image:
src_t = torch.stack(list(img), dim=0)
else:
src_t = img.unsqueeze(0).expand(n, -1, -1, -1)
src_t = src_t.permute(0, 3, 1, 2).contiguous().float() # (N, 3, H, W) in [0, 255]
warped_t = warp_affine_batched(src_t, mats, (H_out, W_out)) # (N, 3, H_out, W_out)
# Float warp -> floor (matches the legacy uint8 round-trip) -> /255.
img_t = torch.floor(warped_t).clamp_(0.0, 255.0) / 255.0 # (N, 3, H_out, W_out) in [0, 1]
# Masks: zero-init when missing, otherwise stack and warp through the same matrices.
boxes_t = torch.as_tensor(boxes, dtype=torch.float32)
if masks is None:
mask_t = torch.zeros((n, H_out, W_out), dtype=torch.float32)
mask_score_t = torch.zeros((n,), dtype=torch.float32)
else:
# masks is an array of N items, each (H, W) or (H, W, 1).
masks_t = torch.stack([torch.as_tensor(masks[i]) for i in range(n)], dim=0)
if masks_t.dim() == 4 and masks_t.shape[-1] == 1:
masks_t = masks_t[..., 0]
masks_src_t = masks_t.float().unsqueeze(1) # (N, 1, H, W) in [0, 255]
warped_masks = warp_affine_batched(masks_src_t, mats, (H_out, W_out))
mask_t = torch.floor(warped_masks.squeeze(1)).clamp_(0.0, 255.0)
if masks_score is not None:
mask_score_t = torch.as_tensor([masks_score[i] for i in range(n)], dtype=torch.float32)
else:
mask_score_t = torch.ones((n,), dtype=torch.float32)
img_size_t = torch.tensor([W_out, H_out], dtype=torch.float32).expand(n, 2).contiguous()
batch = {
"img": img_t.unsqueeze(0), # (1, N, 3, H_out, W_out)
"img_size": img_size_t.unsqueeze(0), # (1, N, 2)
"bbox_center": centers.unsqueeze(0), # (1, N, 2)
"bbox_scale": scales.unsqueeze(0), # (1, N, 2)
"bbox": boxes_t.unsqueeze(0), # (1, N, 4)
"affine_trans": mats.unsqueeze(0), # (1, N, 2, 3)
"mask": mask_t.unsqueeze(0).unsqueeze(2), # (1, N, 1, H_out, W_out)
"mask_score": mask_score_t.unsqueeze(0), # (1, N)
"person_valid": torch.ones((1, n), dtype=torch.float32),
}
if cam_int is not None:
batch["cam_int"] = cam_int.to(batch["img"])
else:
# Default intrinsics: focal = sqrt(W^2 + H^2), principal point = image center.
f = (height ** 2 + width ** 2) ** 0.5
batch["cam_int"] = torch.tensor(
[[[f, 0, width / 2.0], [0, f, height / 2.0], [0, 0, 1]]],
).to(batch["img"])
return batch
# Geometry utils
def rot6d_to_rotmat(
x: torch.Tensor # (B, 6) batch of 6-D rotation representations.
) -> torch.Tensor: # (B, 3, 3) rotation matrices.
"""6D continuous rotation rep (Zhou et al., CVPR 2019) -> 3x3 rotation matrix."""
x = x.reshape(-1, 2, 3).permute(0, 2, 1).contiguous()
a1, a2 = x[:, :, 0], x[:, :, 1]
b1 = F.normalize(a1)
b2 = F.normalize(a2 - torch.einsum("bi,bi->b", b1, a2).unsqueeze(-1) * b1)
b3 = torch.linalg.cross(b1, b2)
return torch.stack((b1, b2, b3), dim=-1)
def perspective_projection(
x: torch.Tensor, # (B, N, 3) 3D points in camera coords.
K: torch.Tensor # (B, 3, 3) camera intrinsics.
) -> torch.Tensor: # (B, N, 2) 2D image-plane projections.
"""Project 3D points (already in camera frame) through intrinsics K."""
y = x / x[:, :, -1].unsqueeze(-1) # perspective divide
y = torch.einsum("bij,bkj->bki", K, y) # apply intrinsics
return y[:, :, :2]
# Rotation conversions, behavior mirrors the roma library (https://github.com/naver/roma)
def _axis_rotmat(axis: str, angle: torch.Tensor) -> torch.Tensor:
"""Rotation matrices around a single coordinate axis. Shape (..., 3, 3)."""
cos = torch.cos(angle)
sin = torch.sin(angle)
one = torch.ones_like(angle)
zero = torch.zeros_like(angle)
if axis == "X":
flat = (one, zero, zero,
zero, cos, -sin,
zero, sin, cos)
elif axis == "Y":
flat = (cos, zero, sin,
zero, one, zero,
-sin, zero, cos)
elif axis == "Z":
flat = (cos, -sin, zero,
sin, cos, zero,
zero, zero, one)
else:
raise ValueError(f"Invalid axis {axis!r}; expected X/Y/Z.")
return torch.stack(flat, dim=-1).reshape(angle.shape + (3, 3))
def euler_to_rotmat(convention: str, angles: torch.Tensor) -> torch.Tensor:
"""Euler angles -> rotation matrix, matching roma's case-keyed convention."""
axes = convention.upper()
R0 = _axis_rotmat(axes[0], angles[..., 0])
R1 = _axis_rotmat(axes[1], angles[..., 1])
R2 = _axis_rotmat(axes[2], angles[..., 2])
if convention.islower():
return R2 @ R1 @ R0
return R0 @ R1 @ R2
def _index_from_letter(letter: str) -> int:
return {"X": 0, "Y": 1, "Z": 2}[letter]
def _angle_from_tan(
axis: str,
other_axis: str,
data: torch.Tensor,
horizontal: bool,
tait_bryan: bool,
) -> torch.Tensor:
"""Extract an outer Euler angle from a row/column of a rotation matrix.
Adapted from PyTorch3D's matrix_to_euler_angles helper.
"""
i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
if horizontal:
i2, i1 = i1, i2
even = (axis + other_axis) in ("XY", "YZ", "ZX")
if horizontal == even:
return torch.atan2(data[..., i1], data[..., i2])
if tait_bryan:
return torch.atan2(-data[..., i2], data[..., i1])
return torch.atan2(data[..., i2], -data[..., i1])
def _matrix_to_euler_intrinsic(matrix: torch.Tensor, convention: str) -> torch.Tensor:
"""Decompose a rotation matrix into intrinsic Euler angles (uppercase abc).
Adapted from PyTorch3D's matrix_to_euler_angles.
"""
i0 = _index_from_letter(convention[0])
i2 = _index_from_letter(convention[2])
tait_bryan = i0 != i2
if tait_bryan:
sign = -1.0 if (i0 - i2) in (-1, 2) else 1.0
central = torch.asin(matrix[..., i0, i2] * sign)
else:
central = torch.acos(matrix[..., i0, i0])
out = (
_angle_from_tan(convention[0], convention[1], matrix[..., i2], False, tait_bryan),
central,
_angle_from_tan(convention[2], convention[1], matrix[..., i0, :], True, tait_bryan),
)
return torch.stack(out, dim=-1)
def rotmat_to_euler(convention: str, matrix: torch.Tensor) -> torch.Tensor:
"""Rotation matrix -> Euler angles, inverse of :func:`euler_to_rotmat`.
PyTorch3D's matrix_to_euler_angles uses the convention R = R_a R_b R_c for
convention "abc"; that matches roma's UPPERCASE ordering directly. For
roma's lowercase, the matrix is reversed (R_c R_b R_a), so we decompose
with the reversed convention and flip the angles back to axis order.
"""
if matrix.shape[-2:] != (3, 3):
raise ValueError(f"Expected (..., 3, 3) rotation matrix, got {tuple(matrix.shape)}.")
if convention.isupper():
return _matrix_to_euler_intrinsic(matrix, convention)
decomposed = _matrix_to_euler_intrinsic(matrix, convention.upper()[::-1])
return decomposed.flip(-1)
def unitquat_to_rotmat(quat: torch.Tensor) -> torch.Tensor:
"""Unit quaternion (x, y, z, w) -> rotation matrix.
Matches roma.unitquat_to_rotmat (scalar-last). The quaternion is assumed to be normalized.
Args:
quat: (..., 4) unit quaternion.
Returns:
(..., 3, 3) rotation matrix.
"""
x, y, z, w = quat.unbind(dim=-1)
tx, ty, tz = 2 * x, 2 * y, 2 * z
twx, twy, twz = tx * w, ty * w, tz * w
txx, txy, txz = tx * x, ty * x, tz * x
tyy, tyz, tzz = ty * y, tz * y, tz * z
one = torch.ones_like(w)
flat = (
one - (tyy + tzz), txy - twz, txz + twy,
txy + twz, one - (txx + tzz), tyz - twx,
txz - twy, tyz + twx, one - (txx + tyy),
)
return torch.stack(flat, dim=-1).reshape(quat.shape[:-1] + (3, 3))
comfy/ldm/sensenova/__init__.py
+1
View File
@@ -0,0 +1 @@
"""SenseNova U1.5 model implementation."""
comfy/ldm/sensenova/conditioning.py
+135
View File
@@ -0,0 +1,135 @@
import torch
IMAGE_CONTEXT_ID = 151669
IMAGE_START_ID = 151670
IMAGE_END_ID = 151671
IM_START_ID = 151644
IM_END_ID = 151645
USER_ID = 872
ASSISTANT_ID = 77091
NEWLINE_ID = 198
IMAGE_LABEL_ID = 1906
HYPHEN_ID = 12
DIGIT_ZERO_ID = 15
COLON_ID = 25
def preprocess_reference(image):
if image.ndim == 3:
image = image.unsqueeze(0)
image = image[:, :, :, :3].movedim(-1, 1).float()
if image.shape[1] == 0:
image = image.new_zeros((image.shape[0], 3, *image.shape[-2:]))
elif image.shape[1] < 3:
repeats = (3 + image.shape[1] - 1) // image.shape[1]
image = image.repeat(1, repeats, 1, 1)[:, :3]
mean = image.new_tensor((0.485, 0.456, 0.406)).view(1, 3, 1, 1)
std = image.new_tensor((0.229, 0.224, 0.225)).view(1, 3, 1, 1)
return (image - mean) / std
def split_reference_batches(images):
references = []
for image in images:
if image.ndim == 3:
image = image.unsqueeze(0)
references.extend(image[index : index + 1] for index in range(image.shape[0]))
return references
def preprocess_references(images):
return [preprocess_reference(image) for image in split_reference_batches(images)]
def _image_tokens(token_height, token_width):
return (
[IMAGE_START_ID]
+ [IMAGE_CONTEXT_ID] * (token_height * token_width)
+ [IMAGE_END_ID]
)
def _image_label_tokens(index):
digits = (DIGIT_ZERO_ID + int(digit) for digit in str(index + 1))
return (IMAGE_LABEL_ID, HYPHEN_ID, *digits, COLON_ID)
def conditioned_input_length(input_length, reference_grids, image_only=False):
image_token_count = sum(height * width for height, width in reference_grids)
if image_only:
return image_token_count + 9 + 2 * len(reference_grids)
label_count = (
sum(len(_image_label_tokens(index)) for index in range(len(reference_grids)))
if len(reference_grids) > 1
else 0
)
return input_length + image_token_count + 3 * len(reference_grids) + label_count
def condition_input_ids(input_ids, reference_grids, image_only=False):
image_blocks = [_image_tokens(height, width) for height, width in reference_grids]
if image_only:
values = (
[IM_START_ID, USER_ID, NEWLINE_ID]
+ [token for block in image_blocks for token in block]
+ [
IM_END_ID,
NEWLINE_ID,
IM_START_ID,
ASSISTANT_ID,
NEWLINE_ID,
IMAGE_START_ID,
]
)
return torch.tensor([values], dtype=torch.long, device=input_ids.device)
values = input_ids[0].tolist()
starts = [index for index, value in enumerate(values) if value == IM_START_ID]
insert_at = starts[1] + 3 if len(starts) > 1 else len(values)
inserted = []
for index, block in enumerate(image_blocks):
if len(image_blocks) > 1:
inserted.extend(_image_label_tokens(index))
inserted.extend(block)
inserted.append(NEWLINE_ID)
values[insert_at:insert_at] = inserted
return torch.tensor([values], dtype=torch.long, device=input_ids.device)
def thw_indexes(input_ids, reference_grids):
values = input_ids[0]
image_start_shift = torch.cat(
(
torch.zeros(1, dtype=torch.long, device=values.device),
(values == IMAGE_START_ID).long(),
)
)[:-1]
not_image = (values != IMAGE_CONTEXT_ID).long()
time_indexes = (image_start_shift + not_image).cumsum(0) - 1
height_indexes = torch.zeros_like(time_indexes)
width_indexes = torch.zeros_like(time_indexes)
selected = values == IMAGE_CONTEXT_ID
height_positions = []
width_positions = []
for token_height, token_width in reference_grids:
positions = torch.arange(
token_height * token_width, dtype=torch.long, device=values.device
)
height_positions.append(positions // token_width)
width_positions.append(positions % token_width)
if height_positions:
height_indexes[selected] = torch.cat(height_positions)
width_indexes[selected] = torch.cat(width_positions)
return torch.stack((time_indexes, height_indexes, width_indexes)).unsqueeze(0)
def block_causal_mask(time_indexes, dtype=torch.float32):
values = time_indexes[0, 0]
length = values.shape[0]
same_block = values[:, None] == values[None, :]
positions = torch.arange(length, device=values.device)
causal = positions[None, :] <= positions[:, None]
allowed = same_block | causal
mask = torch.zeros((1, 1, length, length), dtype=dtype, device=values.device)
return mask.masked_fill_(~allowed[None, None], float("-inf"))
comfy/ldm/sensenova/model.py
+622
View File
@@ -0,0 +1,622 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
import comfy.patcher_extension
import comfy.utils
from comfy.ldm.common_dit import pad_to_patch_size
from comfy.ldm.flux.math import apply_rope1
from comfy.ldm.modules.attention import optimized_attention
from comfy.ldm.modules.diffusionmodules.mmdit import TimestepEmbedder
from .sampling import resolution_noise_scale
HIDDEN_SIZE = 4096
INTERMEDIATE_SIZE = 12288
NUM_LAYERS = 42
NUM_HEADS = 32
NUM_KV_HEADS = 8
HEAD_DIM = 128
MERGED_PATCH_SIZE = 32
VOCAB_SIZE = 151936
def _pad_to_merged_patch_size(value):
height, width = value.shape[-2:]
height_pad = max(16 - height, 0)
width_pad = max(16 - width, 0)
if height_pad or width_pad:
value = F.pad(
value,
(0, width_pad, 0, height_pad),
mode="replicate" if height > 0 and width > 0 else "constant",
)
return pad_to_patch_size(value, (MERGED_PATCH_SIZE, MERGED_PATCH_SIZE))
def _generation_batch_size(total_batch, prefix_batch):
if prefix_batch < 1 or total_batch < 1 or total_batch % prefix_batch != 0:
raise ValueError(
"SenseNova generation batch must be a positive multiple of the prefix batch "
f"(generation={total_batch}, prefix={prefix_batch})"
)
return total_batch // prefix_batch
def _match_prefix_batch(total_batch, text_input_ids, prefix_indexes, prefix_mask):
prefix_batch = text_input_ids.shape[0]
if prefix_batch > 0 and total_batch % prefix_batch:
text_input_ids = comfy.utils.resize_to_batch_size(text_input_ids, total_batch)
if prefix_indexes is not None:
prefix_indexes = comfy.utils.resize_to_batch_size(
prefix_indexes, total_batch
)
if prefix_mask is not None:
prefix_mask = comfy.utils.resize_to_batch_size(prefix_mask, total_batch)
return text_input_ids, prefix_indexes, prefix_mask
def _expand_prefix_batch(value, generation_batch):
"""Repeat each guidance branch's prefix KV for its generated variants."""
if generation_batch == 1:
return value
prefix_batch = value.shape[0]
return (
value.unsqueeze(1)
.expand(prefix_batch, generation_batch, *value.shape[1:])
.reshape(prefix_batch * generation_batch, *value.shape[1:])
)
def _prepare_llm_rope(positions, dim, theta, device, dtype):
frequencies = theta ** (
-torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim
)
positions = positions.to(device=device, dtype=torch.float32)
if positions.ndim == 1:
positions = positions.unsqueeze(0)
angles = positions.unsqueeze(-1) * frequencies
embedding = torch.cat((angles, angles), dim=-1).unsqueeze(1)
return embedding.cos().to(dtype), embedding.sin().to(dtype)
def _prepare_mrope(indexes, device, dtype):
return (
_prepare_llm_rope(indexes[0], HEAD_DIM // 2, 5000000.0, device, dtype),
_prepare_llm_rope(indexes[1], HEAD_DIM // 4, 10000.0, device, dtype),
_prepare_llm_rope(indexes[2], HEAD_DIM // 4, 10000.0, device, dtype),
)
def _apply_llm_rope(query, key, rope):
cosine, sine = rope
def rotate_half(value):
first, second = value.chunk(2, dim=-1)
return torch.cat((-second, first), dim=-1)
# Keep this split-half RoPE on the reference PyTorch formula. The
# comfy-kitchen CUDA kernel is selected automatically on CUDA 13 builds;
# on Blackwell it can return finite but numerically incorrect values, which
# corrupts the generated image without raising an execution error.
return (
query * cosine + rotate_half(query) * sine,
key * cosine + rotate_half(key) * sine,
)
def _apply_interleaved_rope(value, positions, theta):
dim = value.shape[-1]
frequencies = theta ** (
-torch.arange(0, dim, 2, dtype=torch.float32, device=value.device) / dim
)
angles = (
positions.to(device=value.device, dtype=torch.float32).unsqueeze(-1)
* frequencies
)
cosine = angles.cos()
sine = angles.sin()
# comfy-kitchen acceleration backends use the canonical four-dimensional
# input and six-dimensional rotation layout. SenseNova's vision patches
# have no head axis, so add a singleton one instead of relying on the eager
# backend's more permissive rank handling.
rotation = torch.stack((cosine, -sine, sine, cosine), dim=-1).reshape(
1, 1, *angles.shape, 2, 2
)
return apply_rope1(value.float().unsqueeze(1), rotation).squeeze(1)
class VisionEmbeddings(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.patch_embedding = operations.Conv2d(
3, 1024, kernel_size=16, stride=16, device=device, dtype=dtype
)
self.dense_embedding = operations.Conv2d(
1024, HIDDEN_SIZE, kernel_size=2, stride=2, device=device, dtype=dtype
)
self.gelu = nn.GELU()
def forward(self, image):
patches = self.gelu(self.patch_embedding(image))
batch, channels, height, width = patches.shape
patches = patches.flatten(2).transpose(1, 2)
indexes = torch.arange(height * width, device=patches.device)
x_positions = indexes % width
y_positions = indexes // width
first = _apply_interleaved_rope(
patches[..., : channels // 2], x_positions, 10000.0
)
second = _apply_interleaved_rope(
patches[..., channels // 2 :], y_positions, 10000.0
)
patches = torch.cat((first, second), dim=-1).to(image.dtype)
patches = patches.transpose(1, 2).reshape(batch, channels, height, width)
patches = self.dense_embedding(patches)
return patches.flatten(2).transpose(1, 2)
class VisionModel(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.embeddings = VisionEmbeddings(
device=device, dtype=dtype, operations=operations
)
def forward(self, image):
return self.embeddings(image)
class MLP(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.gate_proj = operations.Linear(
HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False, device=device, dtype=dtype
)
self.up_proj = operations.Linear(
HIDDEN_SIZE, INTERMEDIATE_SIZE, bias=False, device=device, dtype=dtype
)
self.down_proj = operations.Linear(
INTERMEDIATE_SIZE, HIDDEN_SIZE, bias=False, device=device, dtype=dtype
)
def forward(self, hidden_states):
return self.down_proj(
F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)
)
class Attention(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.q_proj = operations.Linear(
HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
)
self.q_proj_mot_gen = operations.Linear(
HIDDEN_SIZE, NUM_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
)
self.k_proj = operations.Linear(
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
)
self.k_proj_mot_gen = operations.Linear(
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
)
self.v_proj = operations.Linear(
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
)
self.v_proj_mot_gen = operations.Linear(
HIDDEN_SIZE, NUM_KV_HEADS * HEAD_DIM, bias=False, device=device, dtype=dtype
)
self.o_proj = operations.Linear(
NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, bias=False, device=device, dtype=dtype
)
self.o_proj_mot_gen = operations.Linear(
NUM_HEADS * HEAD_DIM, HIDDEN_SIZE, bias=False, device=device, dtype=dtype
)
self.q_norm = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.q_norm_mot_gen = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.q_norm_hw = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.q_norm_hw_mot_gen = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.k_norm = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.k_norm_mot_gen = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.k_norm_hw = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
self.k_norm_hw_mot_gen = operations.RMSNorm(
HEAD_DIM // 2, eps=1e-6, device=device, dtype=dtype
)
def _project(self, hidden_states, rope, generation):
batch, length, _ = hidden_states.shape
if generation:
query = self.q_proj_mot_gen(hidden_states).view(
batch, length, NUM_HEADS, HEAD_DIM
)
key = self.k_proj_mot_gen(hidden_states).view(
batch, length, NUM_KV_HEADS, HEAD_DIM
)
value = (
self.v_proj_mot_gen(hidden_states)
.view(batch, length, NUM_KV_HEADS, HEAD_DIM)
.transpose(1, 2)
)
query_t, query_hw = query.chunk(2, dim=-1)
key_t, key_hw = key.chunk(2, dim=-1)
query_t = self.q_norm_mot_gen(query_t).transpose(1, 2)
query_hw = self.q_norm_hw_mot_gen(query_hw).transpose(1, 2)
key_t = self.k_norm_mot_gen(key_t).transpose(1, 2)
key_hw = self.k_norm_hw_mot_gen(key_hw).transpose(1, 2)
else:
query = self.q_proj(hidden_states).view(batch, length, NUM_HEADS, HEAD_DIM)
key = self.k_proj(hidden_states).view(batch, length, NUM_KV_HEADS, HEAD_DIM)
value = (
self.v_proj(hidden_states)
.view(batch, length, NUM_KV_HEADS, HEAD_DIM)
.transpose(1, 2)
)
query_t, query_hw = query.chunk(2, dim=-1)
key_t, key_hw = key.chunk(2, dim=-1)
query_t = self.q_norm(query_t).transpose(1, 2)
query_hw = self.q_norm_hw(query_hw).transpose(1, 2)
key_t = self.k_norm(key_t).transpose(1, 2)
key_hw = self.k_norm_hw(key_hw).transpose(1, 2)
query_h, query_w = query_hw.chunk(2, dim=-1)
key_h, key_w = key_hw.chunk(2, dim=-1)
query_t, key_t = _apply_llm_rope(query_t, key_t, rope[0])
query_h, key_h = _apply_llm_rope(query_h, key_h, rope[1])
query_w, key_w = _apply_llm_rope(query_w, key_w, rope[2])
query = torch.cat((query_t, query_h, query_w), dim=-1)
key = torch.cat((key_t, key_h, key_w), dim=-1)
return query, key, value
def forward_prefix(
self, hidden_states, rope, attention_mask, transformer_options
):
query, key, value = self._project(hidden_states, rope, False)
output = optimized_attention(
query,
key,
value,
NUM_HEADS,
mask=attention_mask,
skip_reshape=True,
transformer_options=transformer_options,
enable_gqa=True,
)
return self.o_proj(output), key, value
def forward_generation(
self, hidden_states, rope, prefix_key, prefix_value, transformer_options
):
query, key, value = self._project(hidden_states, rope, True)
key = torch.cat((prefix_key, key), dim=2)
value = torch.cat((prefix_value, value), dim=2)
output = optimized_attention(
query,
key,
value,
NUM_HEADS,
mask=None,
skip_reshape=True,
transformer_options=transformer_options,
enable_gqa=True,
)
return self.o_proj_mot_gen(output)
class DecoderLayer(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.self_attn = Attention(device=device, dtype=dtype, operations=operations)
self.mlp = MLP(device=device, dtype=dtype, operations=operations)
self.mlp_mot_gen = MLP(device=device, dtype=dtype, operations=operations)
self.input_layernorm = operations.RMSNorm(
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
)
self.input_layernorm_mot_gen = operations.RMSNorm(
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
)
self.post_attention_layernorm = operations.RMSNorm(
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
)
self.post_attention_layernorm_mot_gen = operations.RMSNorm(
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
)
def forward_prefix(self, prefix, prefix_rope, prefix_mask, transformer_options):
prefix_attention, prefix_key, prefix_value = self.self_attn.forward_prefix(
self.input_layernorm(prefix),
prefix_rope,
prefix_mask,
transformer_options,
)
prefix = prefix + prefix_attention
prefix = prefix + self.mlp(self.post_attention_layernorm(prefix))
return prefix, prefix_key, prefix_value
def forward_generation(
self, image, image_rope, prefix_key, prefix_value, transformer_options
):
image_attention = self.self_attn.forward_generation(
self.input_layernorm_mot_gen(image),
image_rope,
prefix_key,
prefix_value,
transformer_options,
)
image = image + image_attention
image = image + self.mlp_mot_gen(self.post_attention_layernorm_mot_gen(image))
return image
class LanguageBackbone(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.embed_tokens = operations.Embedding(
VOCAB_SIZE, HIDDEN_SIZE, padding_idx=151643, device=device, dtype=dtype
)
self.layers = nn.ModuleList(
DecoderLayer(device=device, dtype=dtype, operations=operations)
for _ in range(NUM_LAYERS)
)
self.norm = operations.RMSNorm(
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
)
self.norm_mot_gen = operations.RMSNorm(
HIDDEN_SIZE, eps=1e-6, device=device, dtype=dtype
)
class LanguageModel(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.model = LanguageBackbone(device=device, dtype=dtype, operations=operations)
class ConvDecoder(nn.Module):
def __init__(self, device=None, dtype=None, operations=None):
super().__init__()
self.ps1 = nn.PixelShuffle(2)
self.conv1 = operations.Conv2d(
1024, 1024, kernel_size=3, padding=1, device=device, dtype=dtype
)
self.act1 = nn.GELU()
self.ps2 = nn.PixelShuffle(2)
self.conv2 = operations.Conv2d(
256, 192, kernel_size=3, padding=1, device=device, dtype=dtype
)
self.ps3 = nn.PixelShuffle(8)
def forward(self, hidden_states):
hidden_states = self.act1(self.conv1(self.ps1(hidden_states)))
return self.ps3(self.conv2(self.ps2(hidden_states)))
class SenseNovaU15(nn.Module):
def __init__(
self, image_model=None, dtype=None, device=None, operations=None, **kwargs
):
super().__init__()
self.dtype = dtype
self.vision_model = VisionModel(
device=device, dtype=dtype, operations=operations
)
self.language_model = LanguageModel(
device=device, dtype=dtype, operations=operations
)
self.fm_modules = nn.ModuleDict(
{
"vision_model_mot_gen": VisionModel(
device=device, dtype=dtype, operations=operations
),
"timestep_embedder": TimestepEmbedder(
HIDDEN_SIZE, device=device, dtype=dtype, operations=operations
),
"fm_head": ConvDecoder(
device=device, dtype=dtype, operations=operations
),
"noise_scale_embedder": TimestepEmbedder(
HIDDEN_SIZE, device=device, dtype=dtype, operations=operations
),
}
)
def forward(self, x, timesteps, context=None, transformer_options={}, **kwargs):
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
self._forward,
self,
comfy.patcher_extension.get_all_wrappers(
comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options
),
).execute(x, timesteps, context, transformer_options, **kwargs)
def _prepare_prefix(
self, text_input_ids, reference_images, prefix_indexes, prefix_mask
):
prefix = self.language_model.model.embed_tokens(text_input_ids)
if reference_images:
reference_embeds = [
self.vision_model(_pad_to_merged_patch_size(reference))
for reference in reference_images
]
selected = text_input_ids == 151669
prefix = prefix.clone()
prefix[selected] = torch.cat(reference_embeds, dim=1).reshape(
-1, HIDDEN_SIZE
)
prefix_length = text_input_ids.shape[1]
if prefix_indexes is None:
prefix_positions = torch.arange(
prefix_length, dtype=torch.long, device=prefix.device
)
zeros = torch.zeros_like(prefix_positions)
prefix_indexes = torch.stack((prefix_positions, zeros, zeros))
prefix_mask = torch.full(
(prefix_length, prefix_length),
float("-inf"),
dtype=prefix.dtype,
device=prefix.device,
).triu(1)
prefix_time = torch.full(
(prefix.shape[0],),
prefix_length,
dtype=torch.long,
device=prefix.device,
)
else:
prefix_indexes = prefix_indexes.transpose(0, 1)
prefix_time = prefix_indexes[0].amax(dim=-1) + 1
return prefix, prefix_indexes, prefix_mask, prefix_time
def preprocess_prefix(
self,
text_input_ids,
reference_images=None,
prefix_indexes=None,
prefix_mask=None,
):
prefix, prefix_indexes, prefix_mask, prefix_time = self._prepare_prefix(
text_input_ids, reference_images, prefix_indexes, prefix_mask
)
prefix_keys = []
prefix_values = []
prefix_rope = _prepare_mrope(prefix_indexes, prefix.device, prefix.dtype)
transformer_options = {}
for layer_index, layer in enumerate(self.language_model.model.layers):
transformer_options["block_index"] = layer_index
prefix, prefix_key, prefix_value = layer.forward_prefix(
prefix,
prefix_rope,
prefix_mask,
transformer_options,
)
prefix_keys.append(prefix_key)
prefix_values.append(prefix_value)
return prefix_keys, prefix_values, prefix_time
def _forward(
self,
x,
timesteps,
context=None,
transformer_options={},
text_input_ids=None,
reference_images=None,
prefix_indexes=None,
prefix_mask=None,
prefix_keys=None,
prefix_values=None,
prefix_time=None,
**kwargs,
):
if text_input_ids is None and prefix_keys is None:
raise ValueError("SenseNova-U1.5 requires text conditioning")
original_height, original_width = x.shape[-2:]
x = _pad_to_merged_patch_size(x)
batch, _, height, width = x.shape
if prefix_keys is None:
text_input_ids, prefix_indexes, prefix_mask = _match_prefix_batch(
batch, text_input_ids, prefix_indexes, prefix_mask
)
prefix_batch = text_input_ids.shape[0]
if reference_images:
reference_images = [
comfy.utils.resize_to_batch_size(reference, prefix_batch)
for reference in reference_images
]
else:
reference_images = None
else:
prefix_batch = prefix_keys[0].shape[0]
if prefix_batch > 0 and batch % prefix_batch:
prefix_keys = [
comfy.utils.resize_to_batch_size(value, batch)
for value in prefix_keys
]
prefix_values = [
comfy.utils.resize_to_batch_size(value, batch)
for value in prefix_values
]
prefix_time = comfy.utils.resize_to_batch_size(prefix_time, batch)
prefix_batch = batch
generation_batch = _generation_batch_size(batch, prefix_batch)
token_height = height // MERGED_PATCH_SIZE
token_width = width // MERGED_PATCH_SIZE
image_length = token_height * token_width
image = self.fm_modules["vision_model_mot_gen"](x)
time_embedding = self.fm_modules["timestep_embedder"](timesteps, image.dtype)
noise_scale = resolution_noise_scale(height, width) / 16.0
scale_timesteps = torch.full_like(timesteps, noise_scale)
time_embedding = time_embedding + self.fm_modules["noise_scale_embedder"](
scale_timesteps, image.dtype
)
image = image + time_embedding[:, None, :]
if prefix_keys is None:
prefix, prefix_indexes, prefix_mask, prefix_time = self._prepare_prefix(
text_input_ids, reference_images, prefix_indexes, prefix_mask
)
prefix_rope = _prepare_mrope(prefix_indexes, prefix.device, prefix.dtype)
image_time = prefix_time.repeat_interleave(generation_batch)
image_positions = torch.arange(image_length, dtype=torch.long, device=x.device)
image_indexes = torch.stack(
(
image_time[:, None].expand(batch, image_length),
(image_positions // token_width)[None].expand(batch, image_length),
(image_positions % token_width)[None].expand(batch, image_length),
)
)
image_rope = _prepare_mrope(image_indexes, image.device, image.dtype)
for layer_index, layer in enumerate(self.language_model.model.layers):
transformer_options["block_index"] = layer_index
if prefix_keys is None:
prefix, prefix_key, prefix_value = layer.forward_prefix(
prefix,
prefix_rope,
prefix_mask,
transformer_options,
)
else:
prefix_key = prefix_keys[layer_index]
prefix_value = prefix_values[layer_index]
generation_prefix_key = _expand_prefix_batch(prefix_key, generation_batch)
generation_prefix_value = _expand_prefix_batch(
prefix_value, generation_batch
)
image = layer.forward_generation(
image,
image_rope,
generation_prefix_key,
generation_prefix_value,
transformer_options,
)
image = self.language_model.model.norm_mot_gen(image)
image = image.view(batch, token_height, token_width, HIDDEN_SIZE).permute(
0, 3, 1, 2
)
predicted = self.fm_modules["fm_head"](image)
denominator = (1.0 - timesteps).clamp_min(0.02).view(batch, 1, 1, 1)
velocity = (x - predicted) / denominator
return velocity[..., :original_height, :original_width]
comfy/ldm/sensenova/sampling.py
+69
View File
@@ -0,0 +1,69 @@
import math
import torch
import comfy.model_sampling
def time_snr_shift(shift, value):
if shift == 1.0:
return value
return shift * value / (1.0 + (shift - 1.0) * value)
def inverse_time_snr_shift(shift, value):
if shift == 1.0:
return value
return value / (shift - (shift - 1.0) * value)
def upstream_timesteps(steps, shift, device=None):
base = torch.linspace(0.0, 1.0, steps + 1, device=device)
return 1.0 - time_snr_shift(shift, 1.0 - base)
def upstream_sigmas(steps, shift, device=None):
return 1.0 - upstream_timesteps(steps, shift, device=device)
def resolution_noise_scale(
height, width, base_seq_len=64, noise_scale=1.0, maximum=16.0
):
token_height = math.ceil(height / 32)
token_width = math.ceil(width / 32)
scale = math.sqrt(token_height * token_width / base_seq_len) * noise_scale
return min(scale, maximum)
class SenseNovaModelSampling(
comfy.model_sampling.ModelSamplingDiscreteFlow, comfy.model_sampling.CONST
):
def set_parameters(self, shift=1.0, timesteps=1000, multiplier=1000):
self.shift = shift
self.multiplier = multiplier
base_timesteps = torch.linspace(multiplier, 0.0, timesteps + 1)
self.register_buffer("sigmas", self.sigma(base_timesteps))
def timestep(self, sigma):
base_sigma = inverse_time_snr_shift(self.shift, sigma)
return (1.0 - base_sigma) * self.multiplier
def sigma(self, timestep):
base_sigma = 1.0 - timestep / self.multiplier
return time_snr_shift(self.shift, base_sigma)
def percent_to_sigma(self, percent):
if percent <= 0.0:
return 1.0
if percent >= 1.0:
return 0.0
return float(time_snr_shift(self.shift, 1.0 - percent))
def noise_scaling(self, sigma, noise, latent_image, max_denoise=False):
sigma = comfy.model_sampling.reshape_sigma(sigma, noise.ndim)
scale = resolution_noise_scale(
latent_image.shape[-2],
latent_image.shape[-1],
noise_scale=self.noise_scale,
)
return sigma * (scale * noise) + (1.0 - sigma) * latent_image
comfy/ldm/trellis2/flexgemm.py
+154
View File
@@ -0,0 +1,154 @@
from typing import Optional, Tuple
import torch
import comfy.model_management
def compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device):
"""Kernel spatial offsets in the same order as the CUDA/Triton kernels."""
offsets = []
for vx in range(Kw):
for vy in range(Kh):
for vz in range(Kd):
offsets.append((vx * Dw, vy * Dh, vz * Dd))
return torch.tensor(offsets, device=device, dtype=torch.int32)
class TorchHashMap:
"""Sorted-array hashmap backed by torch.searchsorted."""
def __init__(self, keys: torch.Tensor, values: torch.Tensor):
self.sorted_keys, order = torch.sort(keys.to(torch.long))
self.sorted_vals = values.to(torch.long)[order]
self._n = self.sorted_keys.numel()
# Chunk size for lookup_flat, caps each transient to ~CHUNK rows.
_LOOKUP_CHUNK = 1 << 23 # 8M rows ≈ 64 MB per int64 temp
def lookup_flat(self, flat_keys: torch.Tensor) -> torch.Tensor:
N = flat_keys.shape[0]
out = torch.full((N,), -1, device=flat_keys.device, dtype=torch.int32)
if self._n == 0 or N == 0:
return out
for s in range(0, N, self._LOOKUP_CHUNK):
e = min(s + self._LOOKUP_CHUNK, N)
flat_chunk = flat_keys[s:e].to(torch.long)
idx = torch.searchsorted(self.sorted_keys, flat_chunk)
in_range = idx < self._n
idx.clamp_(max=self._n - 1) # reuse idx as the "safe" index
found = in_range & (self.sorted_keys[idx] == flat_chunk)
if found.any():
found_idx = found.nonzero(as_tuple=True)[0]
out[s + found_idx] = self.sorted_vals[idx[found_idx]].to(torch.int32)
return out
def build_submanifold_neighbor_map(
hashmap,
coords: torch.Tensor,
W, H, D,
Kw, Kh, Kd,
Dw, Dh, Dd,
):
# neighbor[i, v] = index of the voxel at voxel i's coord + kernel-offset v, or -1.
# Chunked over voxels so the [chunk, V, 3] candidate transient stays bounded.
device = coords.device
M = coords.shape[0]
offsets = compute_kernel_offsets(Kw, Kh, Kd, Dw, Dh, Dd, device).long() # [V, 3]
V = offsets.shape[0]
center = torch.tensor([(Kw // 2) * Dw, (Kh // 2) * Dh, (Kd // 2) * Dd], device=device)
WHD, HD = W * H * D, H * D
neighbor = torch.empty((M, V), dtype=torch.int32, device=device)
# ~V*40 bytes/voxel of transient (int64 cand + flat + masks); cap at ~0.5 GB.
chunk = max(1, min(M, int(0.5 * (1024 ** 3) / (V * 40))))
for s in range(0, M, chunk):
e = min(s + chunk, M)
b = coords[s:e, 0].long()
cand = coords[s:e, 1:4].long()[:, None, :] + offsets[None, :, :] - center # [c, V, 3]
x, y, z = cand[..., 0], cand[..., 1], cand[..., 2]
in_bounds = (x >= 0) & (x < W) & (y >= 0) & (y < H) & (z >= 0) & (z < D) # [c, V]
flat = b[:, None] * WHD + x * HD + y * D + z # [c, V]
flat = torch.where(in_bounds, flat, torch.full_like(flat, -1)) # OOB -> guaranteed miss
neighbor[s:e] = hashmap.lookup_flat(flat.reshape(-1)).view(e - s, V)
return neighbor
def get_recommended_chunk_mem(
device=None,
safety_fraction: float = 0.2,
min_gb: float = 0.25,
max_gb: float = 2.0,
):
"""Pick a chunk-memory budget (in GB) for sparse conv batching."""
free_gb = comfy.model_management.get_free_memory(device) / (1024 ** 3)
return max(min_gb, min(free_gb * safety_fraction, max_gb))
def sparse_submanifold_conv3d(
feats: torch.Tensor,
coords: torch.Tensor,
shape: tuple,
weight: torch.Tensor,
bias: Optional[torch.Tensor],
neighbor_cache: Optional[torch.Tensor],
dilation: tuple,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
if feats.shape[0] == 0:
Co = weight.shape[0]
return torch.empty((0, Co), device=feats.device, dtype=feats.dtype), None
W, H, D = shape
Co, Kw, Kh, Kd, Ci = weight.shape
V = Kw * Kh * Kd
device = feats.device
if neighbor_cache is None:
b_stride = W * H * D
x_stride = H * D
y_stride = D
z_stride = 1
flat_keys = (coords[:, 0].long() * b_stride +
coords[:, 1].long() * x_stride +
coords[:, 2].long() * y_stride +
coords[:, 3].long() * z_stride)
vals = torch.arange(coords.shape[0], dtype=torch.int32, device=device)
hashmap = TorchHashMap(flat_keys, vals)
neighbor = build_submanifold_neighbor_map(
hashmap, coords, W, H, D, Kw, Kh, Kd,
dilation[0], dilation[1], dilation[2]
)
else:
neighbor = neighbor_cache
N_pts = feats.shape[0]
weight_T = weight.view(Co, V * Ci).T
output = torch.empty(N_pts, Co, device=device, dtype=feats.dtype)
# Zero row at index N_pts; missing neighbors (-1) gather it -> no separate masking.
feats_padded = torch.cat([feats, feats.new_zeros(1, Ci)], dim=0)
# Chunk over voxels to bound the (chunk, V, Ci) gather.
max_chunk_mem_gb = get_recommended_chunk_mem(device)
mem_per_row = V * Ci * feats.element_size()
max_chunk_mem = max_chunk_mem_gb * (1024 ** 3)
chunk_size = max(1, int(max_chunk_mem / mem_per_row))
chunk_size = min(chunk_size, N_pts)
for start in range(0, N_pts, chunk_size):
end = min(start + chunk_size, N_pts)
actual_chunk = end - start
chunk_idx = torch.where(neighbor[start:end] < 0, N_pts, neighbor[start:end]) # -1 -> zero row
gathered = feats_padded[chunk_idx] # (chunk, V, Ci)
gathered_flat = gathered.view(actual_chunk, V * Ci)
output[start:end] = torch.matmul(gathered_flat, weight_T) # (chunk, V*Ci) @ (V*Ci, Co)
if bias is not None:
output += bias.unsqueeze(0).to(output.dtype)
return output, neighbor
comfy/ldm/trellis2/model.py
+1151
View File
File diff suppressed because it is too large Load Diff
comfy/ldm/trellis2/vae.py
+1092
View File
File diff suppressed because it is too large Load Diff
comfy/ldm/wan/model.py
+61 -2
View File
@@ -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")))
@@ -552,6 +561,7 @@ class WanModel(torch.nn.Module):
List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
"""
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
transformer_options["grid_sizes"] = grid_sizes
@@ -564,11 +574,13 @@ class WanModel(torch.nn.Module):
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
full_ref = None
img_offset = 0
if self.ref_conv is not None:
full_ref = kwargs.get("reference_latent", None)
if full_ref is not None:
full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
x = torch.concat((full_ref, x), dim=1)
img_offset = full_ref.shape[1]
# In-context reference (Bernini)
context_latents = kwargs.get("context_latents", None)
@@ -589,6 +601,7 @@ class WanModel(torch.nn.Module):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -604,6 +617,11 @@ class WanModel(torch.nn.Module):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
@@ -777,6 +795,7 @@ class VaceWanModel(WanModel):
**kwargs,
):
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
transformer_options["grid_sizes"] = grid_sizes
@@ -807,6 +826,7 @@ class VaceWanModel(WanModel):
x_orig = x
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -822,6 +842,11 @@ class VaceWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
ii = self.vace_layers_mapping.get(i, None)
if ii is not None:
for iii in range(len(c)):
@@ -887,6 +912,7 @@ class CameraWanModel(WanModel):
**kwargs,
):
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
if self.control_adapter is not None and camera_conditions is not None:
x = x + self.control_adapter(camera_conditions).to(x.dtype)
@@ -909,6 +935,7 @@ class CameraWanModel(WanModel):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -924,6 +951,11 @@ class CameraWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
@@ -1335,6 +1367,7 @@ class WanModel_S2V(WanModel):
# embeddings
bs, _, time, height, width = x.shape
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
if control_video is not None:
x = x + self.cond_encoder(control_video)
@@ -1379,6 +1412,7 @@ class WanModel_S2V(WanModel):
context = self.text_embedding(context)
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -1393,6 +1427,12 @@ class WanModel_S2V(WanModel):
x = out["img"]
else:
x = block(x, e=e0, freqs=freqs, context=context, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
if audio_emb is not None:
x = self.audio_injector(x, i, audio_emb, audio_emb_global, seq_len)
# head
@@ -1599,6 +1639,7 @@ class HumoWanModel(WanModel):
bs, _, time, height, width = x.shape
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
grid_sizes = x.shape[2:]
x = x.flatten(2).transpose(1, 2)
@@ -1630,6 +1671,7 @@ class HumoWanModel(WanModel):
audio = None
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -1645,6 +1687,11 @@ class HumoWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, audio=audio, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": 0, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
@@ -1660,8 +1707,14 @@ class SCAILWanModel(WanModel):
def forward_orig(self, x, t, context, clip_fea=None, freqs=None, transformer_options={}, pose_latents=None, reference_latent=None, ref_mask_latents=None, sam_latents=None, **kwargs):
x_input = x
img_offset = 0
if reference_latent is not None:
x = torch.cat((reference_latent, x), dim=2)
img_offset = (reference_latent.shape[2] // self.patch_size[0]) * \
(reference_latent.shape[3] // self.patch_size[1]) * \
(reference_latent.shape[4] // self.patch_size[2])
# embeddings
x = self.patch_embedding(x.float()).to(x.dtype)
@@ -1697,6 +1750,7 @@ class SCAILWanModel(WanModel):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -1712,6 +1766,11 @@ class SCAILWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
# head
x = self.head(x, e)
comfy/ldm/wan/model_animate.py
+9
View File
@@ -493,6 +493,7 @@ class AnimateWanModel(WanModel):
**kwargs,
):
# embeddings
x_input = x
x = self.patch_embedding(x.float()).to(x.dtype)
x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
grid_sizes = x.shape[2:]
@@ -505,11 +506,13 @@ class AnimateWanModel(WanModel):
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
full_ref = None
img_offset = 0
if self.ref_conv is not None:
full_ref = kwargs.get("reference_latent", None)
if full_ref is not None:
full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
x = torch.concat((full_ref, x), dim=1)
img_offset = full_ref.shape[1]
# context
context = self.text_embedding(context)
@@ -522,6 +525,7 @@ class AnimateWanModel(WanModel):
context_img_len = clip_fea.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -537,6 +541,11 @@ class AnimateWanModel(WanModel):
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
if i % 5 == 0 and motion_vec is not None:
x = x + self.face_adapter.fuser_blocks[i // 5](x, motion_vec)
comfy/ldm/wan/model_animate2.py
+387
View File
@@ -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/model_wandancer.py
+10
View File
@@ -111,6 +111,7 @@ class WanDancerModel(WanModel):
def forward_orig(self, x, t, context, clip_fea=None, clip_fea_ref=None, freqs=None, audio_embed=None, fps=30, audio_inject_scale=1.0, transformer_options={}, **kwargs):
# embeddings
x_input = x
if int(fps + 0.5) != 30:
x = self.patch_embedding_global(x.float()).to(x.dtype)
else:
@@ -128,11 +129,13 @@ class WanDancerModel(WanModel):
e0 = self.time_projection(e).unflatten(2, (6, self.dim))
full_ref = None
img_offset = 0
if self.ref_conv is not None: # model has the weight, but this wasn't used in the original pipeline
full_ref = kwargs.get("reference_latent", None)
if full_ref is not None:
full_ref = self.ref_conv(full_ref).flatten(2).transpose(1, 2)
x = torch.concat((full_ref, x), dim=1)
img_offset = full_ref.shape[1]
# context
context = self.text_embedding(context)
@@ -163,6 +166,7 @@ class WanDancerModel(WanModel):
context_img_len += clip_fea_ref.shape[-2]
patches_replace = transformer_options.get("patches_replace", {})
patches = transformer_options.get("patches", {})
blocks_replace = patches_replace.get("dit", {})
transformer_options["total_blocks"] = len(self.blocks)
transformer_options["block_type"] = "double"
@@ -177,6 +181,12 @@ class WanDancerModel(WanModel):
x = out["img"]
else:
x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
if "double_block" in patches:
for p in patches["double_block"]:
out = p({"img": x, "x": x_input, "vec": e, "block_index": i, "img_offset": img_offset, "transformer_options": transformer_options})
x = out["img"]
if audio_emb is not None:
x = self.music_injector(x, i, audio_emb, audio_emb_global=None, seq_len=seq_len, scale=audio_inject_scale)
comfy/ldm/wan/uni3c.py
+149
View File
@@ -0,0 +1,149 @@
# Uni3C controlnet for Wan 2.1: https://github.com/ewrfcas/Uni3C
# Converted from the original diffusers based implementation.
import torch
import torch.nn as nn
from comfy.ldm.flux.layers import EmbedND
from .model import WanFeedForward, WanSelfAttention
class Uni3CLayerNormZero(nn.Module):
def __init__(
self,
conditioning_dim,
embedding_dim,
eps=1e-5,
device=None, dtype=None, operations=None
):
super().__init__()
self.silu = nn.SiLU()
self.linear = operations.Linear(conditioning_dim, 3 * embedding_dim, device=device, dtype=dtype)
self.norm = operations.LayerNorm(embedding_dim, eps=eps, elementwise_affine=True, device=device, dtype=dtype)
def forward(self, x, temb):
shift, scale, gate = self.linear(self.silu(temb)).chunk(3, dim=1)
x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
return x, gate[:, None, :]
class Uni3CAttentionBlock(nn.Module):
def __init__(
self,
dim,
ffn_dim,
num_heads,
time_embed_dim=5120,
eps=1e-6,
device=None, dtype=None, operations=None
):
super().__init__()
operation_settings = {"operations": operations, "device": device, "dtype": dtype}
self.norm1 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
self.self_attn = WanSelfAttention(dim, num_heads, qk_norm=True, eps=eps, operation_settings=operation_settings)
self.norm2 = Uni3CLayerNormZero(time_embed_dim, dim, device=device, dtype=dtype, operations=operations)
self.ffn = WanFeedForward(
operations.Linear(dim, ffn_dim, device=device, dtype=dtype), nn.GELU(approximate='tanh'),
operations.Linear(ffn_dim, dim, device=device, dtype=dtype))
def forward(self, x, temb, freqs):
norm_x, gate_msa = self.norm1(x, temb)
x = x + gate_msa * self.self_attn(norm_x, freqs)
norm_x, gate_ff = self.norm2(x, temb)
x = x + gate_ff * self.ffn(norm_x)
return x
class MaskCamEmbed(nn.Module):
def __init__(
self,
add_channels=7,
mid_channels=256,
conv_out_dim=5120,
device=None, dtype=None, operations=None
):
super().__init__()
self.mask_padding = [0, 0, 0, 0, 3, 0] # first frame conditioning
self.mask_proj = nn.Sequential(
operations.Conv3d(add_channels, mid_channels, kernel_size=(4, 8, 8), stride=(4, 8, 8), device=device, dtype=dtype),
operations.GroupNorm(mid_channels // 8, mid_channels, device=device, dtype=dtype),
nn.SiLU())
self.mask_zero_proj = operations.Conv3d(mid_channels, conv_out_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2), device=device, dtype=dtype)
def forward(self, add_inputs):
add_padded = torch.nn.functional.pad(add_inputs, self.mask_padding, mode="constant", value=0)
add_embeds = self.mask_proj(add_padded)
add_embeds = self.mask_zero_proj(add_embeds)
add_embeds = add_embeds.flatten(2).transpose(1, 2)
return add_embeds
class WanUni3CControlnet(nn.Module):
def __init__(
self,
in_channels=36,
conv_out_dim=5120,
dim=1024,
ffn_dim=8192,
num_heads=16,
num_layers=20,
time_embed_dim=5120,
out_proj_dim=5120,
add_channels=7,
mid_channels=256,
device=None, dtype=None, operations=None
):
super().__init__()
patch_size = (1, 2, 2)
self.num_layers = num_layers
self.controlnet_patch_embedding = operations.Conv3d(
in_channels, conv_out_dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=torch.float32)
self.controlnet_mask_embedding = MaskCamEmbed(add_channels, mid_channels, conv_out_dim, device=device, dtype=dtype, operations=operations)
if conv_out_dim != dim:
self.proj_in = operations.Linear(conv_out_dim, dim, device=device, dtype=dtype)
else:
self.proj_in = nn.Identity()
self.controlnet_blocks = nn.ModuleList([
Uni3CAttentionBlock(dim, ffn_dim, num_heads, time_embed_dim, device=device, dtype=dtype, operations=operations)
for _ in range(num_layers)])
self.proj_out = nn.ModuleList([
operations.Linear(dim, out_proj_dim, device=device, dtype=dtype)
for _ in range(num_layers)])
head_dim = dim // num_heads
self.rope_embedder = EmbedND(dim=head_dim, theta=10000.0, axes_dim=[head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)])
def rope_encode(self, t_len, h_len, w_len, device=None, dtype=None):
img_ids = torch.zeros((t_len, h_len, w_len, 3), device=device, dtype=dtype)
img_ids[:, :, :, 0] = img_ids[:, :, :, 0] + torch.arange(t_len, device=device, dtype=dtype).reshape(-1, 1, 1)
img_ids[:, :, :, 1] = img_ids[:, :, :, 1] + torch.arange(h_len, device=device, dtype=dtype).reshape(1, -1, 1)
img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.arange(w_len, device=device, dtype=dtype).reshape(1, 1, -1)
img_ids = img_ids.reshape(1, -1, img_ids.shape[-1])
freqs = self.rope_embedder(img_ids).movedim(1, 2)
return freqs
def process_input(self, control_input, render_mask=None, camera_embedding=None):
# render_mask/camera_embedding are the checkpoint's extra conditioning path, not wired up yet
hidden = self.controlnet_patch_embedding(control_input.float()).to(control_input.dtype)
t_len, h_len, w_len = hidden.shape[2:]
freqs = self.rope_encode(t_len, h_len, w_len, device=hidden.device, dtype=hidden.dtype)
hidden = hidden.flatten(2).transpose(1, 2)
add_inputs = None
if camera_embedding is not None and render_mask is not None:
add_inputs = torch.cat([render_mask, camera_embedding], dim=1)
elif render_mask is not None:
add_inputs = render_mask
if add_inputs is not None:
hidden = hidden + self.controlnet_mask_embedding(add_inputs.to(hidden.dtype))
hidden = self.proj_in(hidden)
return hidden, freqs
def forward_block(self, block_index, hidden, temb, freqs):
hidden = self.controlnet_blocks[block_index](hidden, temb, freqs)
residual = self.proj_out[block_index](hidden)
return hidden, residual
comfy/model_base.py
+391 -8
View File
@@ -21,6 +21,9 @@ import comfy.ldm.hunyuan3dv2_1.hunyuandit
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
@@ -44,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
@@ -58,11 +62,13 @@ import comfy.ldm.omnigen.omnigen2
import comfy.ldm.seedvr.model
import comfy.ldm.boogu.model
import comfy.ldm.qwen_image.model
import comfy.ldm.mage_flow.model
import comfy.ldm.joyimage.model
import comfy.ldm.ideogram4.model
import comfy.ldm.krea2.model
import comfy.ldm.kandinsky5.model
import comfy.ldm.anima.model
import comfy.ldm.trellis2.model
import comfy.ldm.ace.ace_step15
import comfy.ldm.cogvideo.model
import comfy.ldm.rt_detr.rtdetr_v4
@@ -70,6 +76,9 @@ import comfy.ldm.ernie.model
import comfy.ldm.sam3.detector
import comfy.ldm.hidream_o1.model
from comfy.ldm.hidream_o1.conditioning import build_extra_conds
import comfy.ldm.sensenova.conditioning
import comfy.ldm.sensenova.model
from comfy.ldm.sensenova.sampling import SenseNovaModelSampling, time_snr_shift
import comfy.ldm.depth_anything_3.model
import comfy.model_management
@@ -98,6 +107,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):
@@ -134,6 +144,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
@@ -168,6 +181,7 @@ class BaseModel(torch.nn.Module):
else:
operations = model_config.custom_operations
self.diffusion_model = unet_model(**unet_config, device=device, operations=operations)
self.diffusion_model.requires_grad_(False)
self.diffusion_model.eval()
if comfy.model_management.force_channels_last():
self.diffusion_model.to(memory_format=torch.channels_last)
@@ -177,6 +191,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:
@@ -1144,6 +1159,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):
@@ -1204,6 +1223,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):
@@ -1805,6 +1828,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)
@@ -1867,6 +1925,23 @@ class WAN22(WAN21):
def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
return latent_image
class Trellis2(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None, unet_model=comfy.ldm.trellis2.model.Trellis2):
super().__init__(model_config, model_type, device, unet_model)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
embeds = kwargs.get("embeds")
out["embeds"] = comfy.conds.CONDRegular(embeds)
# CONDConstant: shared across pos/neg
for k in ("trellis2_coords", "trellis2_coord_counts",
"trellis2_generation_mode", "trellis2_shape_slat",
"trellis2_proj_feats", "trellis2_model_frame"):
v = kwargs.get(k)
if v is not None:
out[k] = comfy.conds.CONDConstant(v)
return out
class WAN21_FlowRVS(WAN21):
def __init__(self, model_config, model_type=ModelType.IMG_TO_IMG_FLOW, image_to_video=False, device=None):
model_config.unet_config["model_type"] = "t2v"
@@ -2025,11 +2100,11 @@ class WAN22_WanDancer(WAN21):
fps = kwargs.get("fps", None)
if fps is not None:
out['fps'] = comfy.conds.CONDRegular(torch.FloatTensor([fps]))
out['fps'] = comfy.conds.CONDConstant(fps)
audio_inject_scale = kwargs.get("audio_inject_scale", None)
if audio_inject_scale is not None:
out['audio_inject_scale'] = comfy.conds.CONDRegular(torch.FloatTensor([audio_inject_scale]))
out['audio_inject_scale'] = comfy.conds.CONDConstant(audio_inject_scale)
return out
class Hunyuan3Dv2(BaseModel):
@@ -2062,6 +2137,144 @@ class Hunyuan3Dv2_1(BaseModel):
out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
return out
class MiniMaxH3(BaseModel):
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)
if cross_attn is not None:
# run condition_proj + token refiner once per sampling instead of per step
cross_attn = self.diffusion_model.preprocess_text_embeds(
cross_attn.to(device=kwargs["device"], dtype=self.get_dtype_inference()))
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
latent_shapes = kwargs.get("latent_shapes", None)
if latent_shapes is not None:
out['latent_shapes'] = comfy.conds.CONDConstant(latent_shapes)
# Everything H3-specific rides in one dict so _apply_model's dtype cast
# (which would flatten fp32 cond latents and long tags to bf16) skips it.
payload = {}
tags = kwargs.get("minimax_token_tags", None)
if tags is not None:
payload["text_token_tags"] = tags
keyframes = kwargs.get("minimax_keyframes", None)
if keyframes is not None:
payload["keyframes"] = 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"] = 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()
denoise_mask = kwargs.get("denoise_mask", None)
if denoise_mask is not None:
out.update(self._denoise_mask_conds(denoise_mask, latent_shapes))
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"))
out['minimax_payload'] = comfy.conds.CONDConstant(payload)
return out
def _pool_masks_to_token_grid(self, masks):
# pool the per-pixel masks to the label grid with amax: video per 2x2 DiT patch, audio per latent frame
video_mask = masks[0]
h, w = video_mask.shape[-2:]
ph, pw = self.diffusion_model.patch_size[1:]
lead = video_mask.shape[:-2]
video_mask = torch.nn.functional.pad(video_mask.reshape((-1,) + video_mask.shape[-3:]), (0, -w % pw, 0, -h % ph), mode="replicate")
video_mask = video_mask.reshape(lead + video_mask.shape[-2:])
video_mask = video_mask.reshape(video_mask.shape[:-2] + (video_mask.shape[-2] // ph, ph, video_mask.shape[-1] // pw, pw)).amax(dim=(-3, -1))
pooled = [video_mask.repeat_interleave(ph, dim=-2).repeat_interleave(pw, dim=-1)[..., :h, :w]]
if len(masks) > 1:
audio_mask = masks[1].amax(dim=1, keepdim=True)
pooled.append(audio_mask.expand_as(masks[1]).contiguous())
return pooled
def _token_grid_masks(self, denoise_mask, latent_shapes):
masks = utils.unpack_latents(denoise_mask, latent_shapes)
return [torch.ceil(mask * 256.0) / 256.0 for mask in self._pool_masks_to_token_grid(masks)]
def _denoise_mask_values(self, denoise_mask, latent_shapes):
if latent_shapes is None or len(latent_shapes) < 2:
return {}
masks = self._token_grid_masks(denoise_mask, latent_shapes)
out = {}
if torch.amin(masks[0]).item() < 1.0 - 1e-3:
out['denoise_mask'] = masks[0][:1, :1].clone()
if torch.amin(masks[1]).item() < 1.0 - 1e-3:
out['audio_denoise_mask'] = masks[1][:1].amax(dim=1, keepdim=True)
return out
def _denoise_mask_conds(self, denoise_mask, latent_shapes):
return {name: comfy.conds.CONDRegular(value) for name, value in self._denoise_mask_values(denoise_mask, latent_shapes).items()}
def scale_latent_inpaint(self, sigma, noise, latent_image, x=None, denoise_mask=None, **kwargs):
# preserved regions run at the cond timestep, inject them at cond strength
shapes = self.latent_shapes
if shapes is None or len(shapes) < 2:
return super().scale_latent_inpaint(sigma=sigma, noise=noise, latent_image=latent_image, **kwargs)
cleans = utils.unpack_latents(latent_image, shapes)
noises = utils.unpack_latents(noise, shapes)
aug = comfy.ldm.minimax.model.VISUAL_COND_TIMESTEP # H3's video timestep is 0.999 by default
cleans[0] = aug * cleans[0] + (1.0 - aug) * noises[0]
scale = self.audio_scale()
if scale != 1.0:
# the sampler carries audio as (sigma_v / sigma_a) * x_audio and latent_image
# holds audio_scale * x_audio, so rescale for the model to see it clean
model_sampling = self.model_sampling
sigma_v = sigma.clamp(min=1e-6)
sigma_a = comfy.ldm.minimax.model.time_shift_sigma(sigma_v, model_sampling.shift, model_sampling.audio_shift)
factor = (sigma_v / sigma_a) / scale
cleans[1] = cleans[1] * factor.view(factor.shape[:1] + (1,) * (cleans[1].ndim - 1)).to(cleans[1].dtype)
injected = utils.pack_latents(cleans)[0]
if x is None or denoise_mask is None:
return injected
token_grid_mask = utils.pack_latents(self._token_grid_masks(denoise_mask, shapes))[0]
x_blend_weight = (token_grid_mask - denoise_mask) / (1.0 - denoise_mask).clamp(min=1e-6)
x_blend_weight = torch.where(denoise_mask < 1.0, x_blend_weight.clamp(0.0, 1.0), torch.zeros_like(x_blend_weight))
return injected + x_blend_weight.to(injected.dtype) * (x - injected)
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)
@@ -2133,6 +2346,134 @@ class HiDreamO1(BaseModel):
out[k] = cls(v)
return out
class SenseNovaSharedRegular(comfy.conds.CONDRegular):
"""Keep the shared text/reference prefix at one copy per guidance branch."""
def process_cond(self, batch_size, **kwargs):
return self._copy_with(self.cond)
class SenseNovaSharedList(comfy.conds.CONDList):
def process_cond(self, batch_size, **kwargs):
return self._copy_with(self.cond)
class SenseNovaU15(BaseModel):
PATCH_SIZE = 32
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.sensenova.model.SenseNovaU15)
self.model_sampling = SenseNovaModelSampling(model_config)
self.memory_usage_factor_conds = ("reference_images",)
def process_timestep(self, timestep, **kwargs):
base_timestep = timestep / self.model_sampling.multiplier
return 1.0 - time_snr_shift(self.model_sampling.shift, 1.0 - base_timestep)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
text_input_ids = kwargs.get("text_input_ids")
if text_input_ids is not None:
device = kwargs["device"]
reference_images = kwargs.get("reference_latents")
if reference_images is not None:
reference_images = comfy.ldm.sensenova.conditioning.preprocess_references(reference_images)
image_only = kwargs.get("prompt_type") == "negative"
indexes = None
prefix_mask = None
if reference_images:
reference_grids = [
(
max(1, math.ceil(image.shape[-2] / self.PATCH_SIZE)),
max(1, math.ceil(image.shape[-1] / self.PATCH_SIZE)),
)
for image in reference_images
]
text_input_ids = comfy.ldm.sensenova.conditioning.condition_input_ids(
text_input_ids,
reference_grids,
image_only=image_only,
)
indexes = comfy.ldm.sensenova.conditioning.thw_indexes(text_input_ids, reference_grids)
prefix_mask = comfy.ldm.sensenova.conditioning.block_causal_mask(
indexes, dtype=self.get_dtype_inference()
)
if kwargs.get("hooks") is None:
dtype = self.get_dtype_inference()
prefix_keys, prefix_values, prefix_time = (
self.diffusion_model.preprocess_prefix(
text_input_ids.to(device=device),
[
image.to(device=device, dtype=dtype)
for image in reference_images
]
if reference_images
else None,
indexes.to(device=device) if indexes is not None else None,
prefix_mask.to(device=device)
if prefix_mask is not None
else None,
)
)
out["prefix_keys"] = SenseNovaSharedList(prefix_keys)
out["prefix_values"] = SenseNovaSharedList(prefix_values)
out["prefix_time"] = SenseNovaSharedRegular(prefix_time)
else:
if reference_images:
out["prefix_indexes"] = SenseNovaSharedRegular(indexes)
out["prefix_mask"] = SenseNovaSharedRegular(prefix_mask)
out["reference_images"] = SenseNovaSharedList(reference_images)
out["text_input_ids"] = SenseNovaSharedRegular(text_input_ids)
return out
def extra_conds_shapes(self, **kwargs):
images = kwargs.get("reference_latents")
images = comfy.ldm.sensenova.conditioning.split_reference_batches(images) if images is not None else []
reference_grids = [
(
max(1, math.ceil(image.shape[-3] / self.PATCH_SIZE)),
max(1, math.ceil(image.shape[-2] / self.PATCH_SIZE)),
)
for image in images
]
reference_pixels = sum(
height * width * self.PATCH_SIZE**2
for height, width in reference_grids
)
out = {}
if reference_pixels:
out["reference_images"] = [1, 3, reference_pixels]
text_input_ids = kwargs.get("text_input_ids")
if text_input_ids is not None:
if reference_grids:
length = comfy.ldm.sensenova.conditioning.conditioned_input_length(
text_input_ids.shape[1],
reference_grids,
image_only=kwargs.get("prompt_type") == "negative",
)
else:
length = text_input_ids.shape[1]
out["prefix_mask"] = [1, 1, length, length]
if kwargs.get("hooks") is None:
prefix_shape = [
1,
comfy.ldm.sensenova.model.NUM_KV_HEADS,
comfy.ldm.sensenova.model.NUM_LAYERS
* length
* comfy.ldm.sensenova.model.HEAD_DIM,
]
out["prefix_keys"] = prefix_shape
out["prefix_values"] = prefix_shape
return out
def memory_required(self, input_shape, cond_shapes={}):
memory = super().memory_required(input_shape, cond_shapes)
dtype_size = comfy.model_management.dtype_size(self.get_dtype_inference())
return memory + sum(
math.prod(shape) * dtype_size
for key in ("prefix_mask", "prefix_keys", "prefix_values")
for shape in cond_shapes.get(key, ())
)
class Chroma(Flux):
def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.chroma.model.Chroma):
super().__init__(model_config, model_type, device=device, unet_model=unet_model)
@@ -2211,6 +2552,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)
@@ -2228,10 +2581,7 @@ class Omnigen2(BaseModel):
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
latents = []
for lat in ref_latents:
latents.append(self.process_latent_in(lat))
out['ref_latents'] = comfy.conds.CONDList(latents)
out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
return out
def extra_conds_shapes(self, **kwargs):
@@ -2247,8 +2597,8 @@ class Boogu(Omnigen2):
self.memory_usage_factor_conds = ("ref_latents",)
class QwenImage(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel)
def __init__(self, model_config, model_type=ModelType.FLUX, device=None, unet_model=comfy.ldm.qwen_image.model.QwenImageTransformer2DModel):
super().__init__(model_config, model_type, device=device, unet_model=unet_model)
self.memory_usage_factor_conds = ("ref_latents",)
def extra_conds(self, **kwargs):
@@ -2278,6 +2628,21 @@ class QwenImage(BaseModel):
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
return out
class MageFlow(QwenImage):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.mage_flow.model.MageFlowTransformer2DModel)
def process_timestep(self, timestep, **kwargs):
# Mage runs in bf16 and rounds its timestep frequency table to the timestep dtype, keep that on fp32 devices.
return timestep.to(torch.bfloat16)
def extra_conds_shapes(self, **kwargs):
out = {}
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
out['ref_latents'] = list([1, 128, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 128])
return out
class JoyImage(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel)
@@ -2318,12 +2683,30 @@ class Ideogram4(BaseModel):
class Krea2(BaseModel):
def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.krea2.model.SingleStreamDiT)
self.memory_usage_factor_conds = ("ref_latents",)
def extra_conds(self, **kwargs):
out = super().extra_conds(**kwargs)
cross_attn = kwargs.get("cross_attn", None)
if cross_attn is not None:
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
latents = []
for lat in ref_latents:
latents.append(self.process_latent_in(lat))
out['ref_latents'] = comfy.conds.CONDList(latents)
ref_latents_method = kwargs.get("reference_latents_method", None)
if ref_latents_method is not None:
out['ref_latents_method'] = comfy.conds.CONDConstant(ref_latents_method)
return out
def extra_conds_shapes(self, **kwargs):
out = {}
ref_latents = kwargs.get("reference_latents", None)
if ref_latents is not None:
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
return out
class HunyuanImage21(BaseModel):
comfy/model_detection.py
+83 -2
View File
@@ -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]
@@ -113,6 +120,27 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
unet_config['block_repeat'] = [[1, 1, 1, 1], [2, 2, 2, 2]]
return unet_config
shape_key = '{}img2shape.t_embedder.mlp.0.weight'.format(key_prefix)
tex_key = '{}shape2txt.t_embedder.mlp.0.weight'.format(key_prefix)
if shape_key in state_dict_keys or tex_key in state_dict_keys: # trellis2 / pixal3d
has_shape = shape_key in state_dict_keys
has_tex = tex_key in state_dict_keys
unet_config = {
"image_model": "trellis2",
"resolution": 32 if (metadata or {}).get("is_512") else 64,
"init_txt_model": has_tex,
"txt_only": has_tex and not has_shape,
}
# Per-submodel projection head (Pixal3D adds `proj_linear`; Trellis2 doesn't).
for sub, name, proj_in_channels in (("img2shape", "shape", 2048),
("shape2txt", "texture", 2048),
("structure_model", "structure", 1024)):
key = '{}{}.blocks.0.cross_attn.proj_linear.weight'.format(key_prefix, sub)
if key in state_dict_keys:
unet_config["image_attn_mode_{}".format(name)] = "proj"
unet_config["proj_in_channels_{}".format(name)] = proj_in_channels
return unet_config
if '{}transformer.rotary_pos_emb.inv_freq'.format(key_prefix) in state_dict_keys: #stable audio dit
unet_config = {}
unet_config["audio_model"] = "dit1.0"
@@ -359,6 +387,35 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
# PixArt diffusers
return None
if '{}video_patch_proj.weight'.format(key_prefix) in state_dict_keys and '{}audio_patch_proj.weight'.format(key_prefix) in state_dict_keys: # MiniMax H3
dit_config = {}
dit_config["image_model"] = "minimax_h3"
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
dit_config["token_refiner_num_layers"] = count_blocks(state_dict_keys, '{}token_refiner.blocks.'.format(key_prefix) + '{}.')
dit_config["hidden_size"] = state_dict['{}video_patch_proj.weight'.format(key_prefix)].shape[0]
dit_config["latents_dim"] = state_dict['{}final_layer.video_out.weight'.format(key_prefix)].shape[0] // 4 # patch 1x2x2
dit_config["audio_latents_dim"] = state_dict['{}final_layer.audio_out.weight'.format(key_prefix)].shape[0]
dit_config["attention_head_dim"] = state_dict['{}blocks.0.attn.q_norm.weight'.format(key_prefix)].shape[0]
qkv = state_dict['{}blocks.0.attn.qkv_proj.weight'.format(key_prefix)]
dit_config["num_attention_heads"] = qkv.shape[0] // (3 * dit_config["attention_head_dim"])
dit_config["ffn_hidden_size"] = state_dict['{}blocks.0.mlp.fc1.weight'.format(key_prefix)].shape[0] // 2
dit_config["text_dim"] = state_dict['{}condition_proj.weight'.format(key_prefix)].shape[1]
table_key = '{}adaln_t_table'.format(key_prefix)
if table_key in state_dict_keys:
# adaln shipped over a precomputed curve basis: the adaln linears span a small shared basis of the time-embedding curve (no time embedder)
table = state_dict[table_key].shape # [grid, k]
dit_config["adaln_curve_grid"] = table[0]
dit_config["time_embed_dim"] = table[1]
else:
te = state_dict['{}time_embedder.proj_in.weight'.format(key_prefix)]
dit_config["timestep_input_dim"] = te.shape[1]
dit_config["time_embed_hidden_size"] = te.shape[0]
dit_config["time_embed_dim"] = state_dict['{}time_embedder.proj_out.weight'.format(key_prefix)].shape[0]
dit_config["rope_inv_freq_len"] = state_dict['{}rope.inv_freq'.format(key_prefix)].shape[0]
if metadata is not None and "config" in metadata:
dit_config.update(json.loads(metadata["config"]).get("transformer", {}))
return dit_config
if '{}adaln_single.emb.timestep_embedder.linear_1.bias'.format(key_prefix) in state_dict_keys: #Lightricks ltxv
dit_config = {}
dit_config["image_model"] = "ltxav" if f'{key_prefix}audio_adaln_single.linear.weight' in state_dict_keys else "ltxv"
@@ -368,6 +425,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
@@ -757,6 +815,16 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
if '{}t_embedder1.mlp.0.weight'.format(key_prefix) in state_dict_keys and '{}x_embedder.proj1.weight'.format(key_prefix) in state_dict_keys: # HiDream-O1
return {"image_model": "hidream_o1"}
vision_key = f"{key_prefix}fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight"
query_key = f"{key_prefix}language_model.model.layers.0.self_attn.q_proj_mot_gen.weight"
if (
vision_key in state_dict
and query_key in state_dict
and state_dict[vision_key].shape[0] == 1024
and state_dict[query_key].shape[0] == 4096
): # SenseNova U1.5
return {"image_model": "sensenova_u15"}
if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys: # HiDream
dit_config = {}
dit_config["image_model"] = "hidream"
@@ -800,11 +868,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:
@@ -884,6 +951,13 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
"selected_layer_index": selected_layer_index,
}
if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys and '{}proj_out.weight'.format(key_prefix) in state_dict_keys and state_dict['{}txt_norm.weight'.format(key_prefix)].shape[0] == 2560 and state_dict['{}proj_out.weight'.format(key_prefix)].shape[0] == 128: # Mage-Flow (Qwen Image txt_norm/proj_out are 3584/64)
dit_config = {}
dit_config["image_model"] = "mage_flow"
dit_config["in_channels"] = 128
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}transformer_blocks.'.format(key_prefix) + '{}.')
return dit_config
if '{}txt_norm.weight'.format(key_prefix) in state_dict_keys: # Qwen Image
dit_config = {}
dit_config["image_model"] = "qwen_image"
@@ -1237,6 +1311,13 @@ def unet_prefix_from_state_dict(state_dict):
if any(k.startswith("detector.") for k in state_dict) and any(k.startswith("tracker.") for k in state_dict):
return ""
# SenseNova checkpoints store the diffusion and language backbones at top level.
if (
"fm_modules.vision_model_mot_gen.embeddings.patch_embedding.weight" in state_dict
and "language_model.model.layers.0.self_attn.q_proj_mot_gen.weight" in state_dict
):
return ""
candidates = ["model.diffusion_model.", #ldm/sgm models
"model.model.", #audio models
"net.", #cosmos
comfy/model_management.py
+142 -53
View File
@@ -30,10 +30,12 @@ import gc
import os
from contextlib import contextmanager, nullcontext
import comfy.memory_management
import comfy.system_memory
import comfy.utils
import comfy.quant_ops
import comfy_aimdo.host_buffer
import comfy_aimdo.vram_buffer
from comfy.internal_logging import detail
from typing import TYPE_CHECKING
if TYPE_CHECKING:
@@ -403,7 +405,7 @@ def get_total_memory(dev=None, torch_total_too=False):
dev = get_torch_device()
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
mem_total = psutil.virtual_memory().total
mem_total = comfy.system_memory.virtual_memory_total()
mem_total_torch = mem_total
else:
if directml_enabled:
@@ -446,8 +448,11 @@ def mac_version():
return None
total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
total_ram = psutil.virtual_memory().total / (1024 * 1024)
total_ram = comfy.system_memory.virtual_memory_total() / (1024 * 1024)
logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
cgroup_ram_limit = comfy.system_memory.cgroup_memory_limit()
if cgroup_ram_limit is not None:
logging.info("RAM limited by cgroup to {:0.0f} MB (host has {:0.0f} MB)".format(cgroup_ram_limit / (1024 * 1024), psutil.virtual_memory().total / (1024 * 1024)))
try:
logging.info("pytorch version: {}".format(torch_version))
@@ -559,7 +564,7 @@ except:
SUPPORT_FP8_OPS = args.supports_fp8_compute
AMD_RDNA2_AND_OLDER_ARCH = ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]
AMD_RDNA2_AND_OLDER_ARCH = ["gfx1030", "gfx1031", "gfx1035", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]
AMD_ENABLE_MIOPEN_ENV = 'COMFYUI_ENABLE_MIOPEN'
try:
@@ -575,30 +580,38 @@ 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
if any((a in arch) for a in ["gfx90a", "gfx942", "gfx950", "gfx1100", "gfx1101", "gfx1150", "gfx1151", "gfx1170"]): # 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
if any((a in arch) for a in ["gfx1200", "gfx1201", "gfx950", "gfx1170"]): # TODO: more arches, "gfx942" gives error on pytorch nightly 2.10 1013 rocm7.0
SUPPORT_FP8_OPS = True
except:
@@ -718,39 +731,88 @@ def mark_mmap_dirty(storage):
if mmap_refs is not None:
DIRTY_MMAPS.add(mmap_refs[0])
def free_pins(size, evict_active=False):
PIN_SUBSETS = [ "weights", "patches" ]
LOADED_PIN_SUBSETS = [ "weights-loaded", "patches-loaded" ]
def models_for_pin_eviction(active, current_prompt=None):
for loaded_model in current_loaded_models:
model = loaded_model.model
if model is None or not model.is_dynamic():
continue
pin_state = model.model.dynamic_pins[model.load_device]
if ((active is None or pin_state["active"] == active) and
(current_prompt is None or pin_state["current_prompt"] == current_prompt)):
yield model
def free_model_pins(size, subsets, current_prompt, active, registrations=False):
freed_total = 0
for loaded_model in reversed(current_loaded_models):
for model in models_for_pin_eviction(active, current_prompt=current_prompt):
if size <= 0:
return freed_total
model = loaded_model.model
if model is not None and model.is_dynamic() and (evict_active or not model.model.dynamic_pins[model.load_device]["active"]):
freed = model.partially_unload_ram(size)
freed_total += freed
size -= freed
if registrations:
freed = model.unregister_inactive_pins(size, subsets=subsets)
else:
freed = model.partially_unload_ram(size, subsets=subsets)
freed_total += freed
size -= freed
return freed_total
def pin_eviction_tiers(loaded, evict_active):
tiers = [
(PIN_SUBSETS, False, None),
(LOADED_PIN_SUBSETS, False, None),
(LOADED_PIN_SUBSETS, True, None),
]
if not loaded:
tiers.append((PIN_SUBSETS, True, False))
if evict_active:
tiers.append((PIN_SUBSETS, True, True))
return tiers
def registration_eviction_tiers(evict_active):
subsets = PIN_SUBSETS + LOADED_PIN_SUBSETS
tiers = [
(subsets, False, False),
(subsets, True, False),
]
if evict_active:
tiers.extend([
(subsets, False, True),
(subsets, True, True),
])
return tiers
def free_pins(size, evict_active=False, loaded=False):
freed = 0
for subsets, current_prompt, active in pin_eviction_tiers(loaded, evict_active):
freed += free_model_pins(size - freed, subsets, current_prompt, active)
return freed
def should_free_pins_for_ram_pressure(shortfall):
if shortfall <= 0:
return False
if not WINDOWS:
return True
if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
if comfy.system_memory.virtual_memory_available() < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
return True
try:
return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
except RuntimeError as err:
logging.warning("Could not read Windows swap usage; falling back to RAM-pressure pin eviction: %s", err)
return True
return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
def ensure_pin_budget(size, evict_active=False):
def ensure_pin_budget(size, evict_active=False, loaded=False):
if args.high_ram:
return True
if args.fast_disk:
shortfall = TOTAL_PINNED_MEMORY + size - MAX_PINNED_MEMORY
else:
shortfall = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - psutil.virtual_memory().available
shortfall = size + max(comfy.memory_management.RAM_CACHE_HEADROOM / 2, 2048 * 1024 ** 2) - comfy.system_memory.virtual_memory_available()
if shortfall <= 0:
return True
to_free = shortfall + PIN_PRESSURE_HYSTERESIS
return free_pins(to_free, evict_active=evict_active) >= shortfall
return free_pins(to_free, evict_active=evict_active, loaded=loaded) >= shortfall
def free_registrations(shortfall, evict_active=True):
if MAX_PINNED_MEMORY <= 0:
@@ -759,19 +821,8 @@ def free_registrations(shortfall, evict_active=True):
return True
shortfall += REGISTERABLE_PIN_HYSTERESIS
for loaded_model in reversed(current_loaded_models):
model = loaded_model.model
if model is not None and model.is_dynamic() and not model.model.dynamic_pins[model.load_device]["active"]:
shortfall -= model.unregister_inactive_pins(shortfall)
if shortfall <= 0:
return True
if evict_active:
for loaded_model in current_loaded_models:
model = loaded_model.model
if model is not None and model.is_dynamic() and model.model.dynamic_pins[model.load_device]["active"]:
shortfall -= model.unregister_inactive_pins(shortfall)
if shortfall <= 0:
return True
for subsets, current_prompt, active in registration_eviction_tiers(evict_active):
shortfall -= free_model_pins(shortfall, subsets, current_prompt, active, registrations=True)
return shortfall <= REGISTERABLE_PIN_HYSTERESIS
def ensure_pin_registerable(size, evict_active=True):
@@ -901,6 +952,8 @@ def minimum_inference_memory():
def free_memory(memory_required, device, keep_loaded=[], for_dynamic=False, pins_required=0, ram_required=0):
cleanup_models_gc()
if not for_dynamic:
detail("Non dynamic memory free called! memory_required=%s pins_required=%s ram_required=%s", memory_required, pins_required, ram_required)
unloaded_model = []
can_unload = []
unloaded_models = []
@@ -1039,6 +1092,9 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
lowvram_model_memory = 0.1
loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
vram_used = 0 if is_device_cpu(torch_dev) else loaded_model.model_loaded_memory()
ram_used = model.loaded_ram_size() if model.is_dynamic() else loaded_model.model_memory() - vram_used
detail("Model loaded: patcher=%s model=%s ram_mb=%.1f vram_mb=%.1f", model.__class__.__name__, model.model.__class__.__name__, ram_used / (1024 ** 2), vram_used / (1024 ** 2))
current_loaded_models.insert(0, loaded_model)
return
@@ -1402,9 +1458,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
@@ -1459,21 +1520,27 @@ 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():
pin_state = model.model.dynamic_pins[model.load_device]
if pin_state["active"]:
*_, buckets = pin_state["weights"]
for size, bucket in list(buckets.items()):
bucket[:] = [ entry for entry in bucket if entry[-1] is not None ]
if not bucket:
del buckets[size]
for subset in ("weights", "weights-loaded"):
*_, buckets = pin_state[subset]
for size, bucket in list(buckets.items()):
bucket[:] = [ entry for entry in bucket if entry[-1] is not None ]
if not bucket:
del buckets[size]
pin_state["active"] = False
model.partially_unload_ram(1e30, subsets=[ "patches" ])
model.model.dynamic_pins[model.load_device]["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {})
model.partially_unload_ram(1e30, subsets=[ "patches", "patches-loaded" ])
for subset in ("patches", "patches-loaded"):
pin_state[subset] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, pinned_hostbuf_size(model.model_size())), [], [-1], [0], [0], {})
STREAM_CAST_BUFFERS.clear()
STREAM_AIMDO_CAST_BUFFERS.clear()
@@ -1583,13 +1650,32 @@ 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
swap = 0 if comfy.system_memory.cgroup_memory_limit() is not None else get_disk_swap_total()
MAX_PINNED_MEMORY = max(ram * 0.40, min(ram * 0.90, ram - 4 * 1024 ** 3, ram + swap - 16 * 1024 ** 3))
logging.info("Enabled pinned memory {}".format(MAX_PINNED_MEMORY // (1024 * 1024)))
PINNING_ALLOWED_TYPES = set(["Tensor", "Parameter", "QuantizedTensor"])
@@ -1680,6 +1766,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
@@ -1753,7 +1842,7 @@ def get_free_memory(dev=None, torch_free_too=False):
dev = get_torch_device()
if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
mem_free_total = psutil.virtual_memory().available
mem_free_total = comfy.system_memory.virtual_memory_available()
mem_free_torch = mem_free_total
else:
if directml_enabled:
comfy/model_patcher.py
+151 -25
View File
@@ -22,6 +22,7 @@ import collections
import inspect
import logging
import math
import time
import uuid
from typing import Callable, Optional
@@ -37,11 +38,58 @@ import comfy.patcher_extension
import comfy.utils
import comfy_aimdo.host_buffer
from comfy.comfy_types import UnetWrapperFunction
from comfy.internal_logging import detail
from comfy.quant_ops import QuantizedTensor
from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP
import comfy_aimdo.model_vbar
def is_model_patcher_output(output):
return isinstance(output, ModelPatcher) or isinstance(getattr(output, "patcher", None), ModelPatcher)
class PromptModelTracker:
def __init__(self):
self.models = {}
def start(self):
self.end()
def add(self, outputs):
if isinstance(outputs, collections.abc.Mapping):
outputs = outputs.values()
elif not isinstance(outputs, (list, tuple)):
outputs = (outputs,)
for output in outputs:
if isinstance(output, (collections.abc.Mapping, list, tuple)):
self.add(output)
continue
models = []
if isinstance(output, ModelPatcher):
models.append(output)
models.extend(output.model_patches_models())
models.extend(output.get_nested_additional_models())
else:
patcher = getattr(output, "patcher", None)
if isinstance(patcher, ModelPatcher):
models.append(patcher)
get_models = getattr(output, "get_models", None)
if callable(get_models):
models.extend(get_models())
for model in models:
if not isinstance(model, ModelPatcher) or not model.is_dynamic():
continue
key = (id(model.model), model.load_device)
self.models[key] = model
model.set_in_use_by_current_prompt(True)
def end(self):
for model in self.models.values():
model.set_in_use_by_current_prompt(False)
self.models.clear()
def set_model_options_patch_replace(model_options, patch, name, block_name, number, transformer_index=None):
to = model_options["transformer_options"].copy()
@@ -510,12 +558,9 @@ class ModelPatcher:
new_multigpu_models = []
for mm in multigpu_models:
# clone main model, but bring over relevant props from existing multigpu clone
n = self.clone()
n = self.clone(model_override=mm.get_clone_model_override())
n.load_device = mm.load_device
n.backup = mm.backup
n.object_patches_backup = mm.object_patches_backup
n.hook_backup = mm.hook_backup
n.model = mm.model
n.is_multigpu_base_clone = mm.is_multigpu_base_clone
if "torch_compile_kwargs" in n.model_options:
from comfy_api.torch_helpers.torch_compile import (
@@ -645,6 +690,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")
@@ -972,14 +1025,14 @@ class ModelPatcher:
m.bias_function = []
if weight_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[weight_key], m.weight, weight_key) != m.weight.shape:
self.patch_weight_to_device(weight_key)
else:
_, set_func, convert_func = get_key_weight(self.model, weight_key)
m.weight_function = [LowVramPatch(weight_key, self.patches, convert_func, set_func)]
patch_counter += 1
if bias_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[bias_key], m.bias, bias_key) != m.bias.shape:
self.patch_weight_to_device(bias_key)
else:
_, set_func, convert_func = get_key_weight(self.model, bias_key)
@@ -1169,14 +1222,14 @@ class ModelPatcher:
module_mem += move_weight_functions(m, device_to)
if lowvram_possible:
if weight_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[weight_key], m.weight, weight_key) != m.weight.shape:
self.patch_weight_to_device(weight_key)
else:
_, set_func, convert_func = get_key_weight(self.model, weight_key)
m.weight_function.append(LowVramPatch(weight_key, self.patches, convert_func, set_func))
patch_counter += 1
if bias_key in self.patches:
if force_patch_weights:
if force_patch_weights or comfy.lora.calculate_shape(self.patches[bias_key], m.bias, bias_key) != m.bias.shape:
self.patch_weight_to_device(bias_key)
else:
_, set_func, convert_func = get_key_weight(self.model, bias_key)
@@ -1736,6 +1789,9 @@ class ModelPatcherDynamic(ModelPatcher):
self.register_load_device(self.load_device)
self.non_dynamic_delegate_model = None
assert load_device is not None
if not hasattr(self.model, "dynamic_patchers"):
self.model.dynamic_patchers = set()
self.model.dynamic_patchers.add(id(self))
def register_load_device(self, device):
"""Ensure dynamic_pins has an entry for *device*.
@@ -1750,14 +1806,20 @@ class ModelPatcherDynamic(ModelPatcher):
self.model.dynamic_pins[device] = {
"weights": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"patches": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"weights-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"patches-loaded": (comfy_aimdo.host_buffer.HostBuffer(0, 0, 0), [], [-1], [0], [0], {}),
"hostbufs_initialized": False,
"failed": False,
"active": False,
"current_prompt": False,
}
def is_dynamic(self):
return True
def set_in_use_by_current_prompt(self, in_use):
self.model.dynamic_pins[self.load_device]["current_prompt"] = in_use
def _vbar_get(self, create=False):
if self.load_device == torch.device("cpu"):
return None
@@ -1785,6 +1847,18 @@ class ModelPatcherDynamic(ModelPatcher):
def unpin_all_weights(self):
self.partially_unload_ram(1e32)
def __del__(self):
model = getattr(self, "model", None)
dynamic_patchers = getattr(model, "dynamic_patchers", None)
if dynamic_patchers is None or id(self) not in dynamic_patchers:
return
dynamic_patchers.discard(id(self))
try:
if not dynamic_patchers:
self.unpin_all_weights()
finally:
self.detach(unpatch_all=False)
def memory_required(self, input_shape):
#Pad this significantly. We are trying to get away from precise estimates. This
#estimate is only used when using the ModelPatcherDynamic after ModelPatcher. If you
@@ -1828,6 +1902,8 @@ class ModelPatcherDynamic(ModelPatcher):
hostbuf_size = comfy.model_management.pinned_hostbuf_size(self.model_size())
pin_state["weights"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["patches"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["weights-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 64 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["patches-loaded"] = (comfy_aimdo.host_buffer.HostBuffer(0, 8 * 1024 * 1024, hostbuf_size), [], [-1], [0], [0], {})
pin_state["hostbufs_initialized"] = True
pin_state["failed"] = False
pin_state["active"] = True
@@ -1837,8 +1913,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"):
@@ -1931,6 +2028,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
@@ -1961,20 +2065,37 @@ class ModelPatcherDynamic(ModelPatcher):
assert self.load_device != torch.device("cpu")
vbar = self._vbar_get()
freed = 0 if vbar is None else vbar.free_memory(memory_to_free)
vbar_freed = 0 if vbar is None else vbar.free_memory(memory_to_free)
freed = vbar_freed
backup_freed = 0
if freed < memory_to_free:
freed += self.restore_loaded_backups()
backup_freed = self.restore_loaded_backups()
freed += backup_freed
method = "vbar+backups" if vbar_freed and backup_freed else "vbar" if vbar_freed else "backups" if backup_freed else "none"
free_methods = getattr(self, "_free_methods", {})
free_methods[method] = free_methods.get(method, 0) + 1
self._free_methods = free_methods
now = time.monotonic()
if now - getattr(self, "_last_free_log_time", 0) >= 5:
requested = "all" if memory_to_free >= 1e30 else f"{memory_to_free / (1024 ** 2):.1f}MB"
prevailing_method = max(free_methods, key=free_methods.get)
detail("AIMDO free: model=%s device=%s prevailing_method=%s methods=%s requested=%s vbar_mb=%.1f backups_mb=%.1f", self.model.__class__.__name__, self.load_device, prevailing_method, free_methods, requested, vbar_freed / (1024 ** 2), backup_freed / (1024 ** 2))
self._free_methods = {}
self._last_free_log_time = now
return freed
def loaded_ram_size(self):
return (self.model.dynamic_pins[self.load_device]["weights"][0].size)
pin_state = self.model.dynamic_pins[self.load_device]
return pin_state["weights"][0].size + pin_state["weights-loaded"][0].size
def pinned_memory_size(self):
return (self.model.dynamic_pins[self.load_device]["weights"][3][0])
pin_state = self.model.dynamic_pins[self.load_device]
return pin_state["weights"][3][0] + pin_state["weights-loaded"][3][0]
def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights", "patches" ]):
def unregister_inactive_pins(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]):
freed = 0
pin_state = self.model.dynamic_pins[self.load_device]
for subset in subsets:
@@ -1982,15 +2103,17 @@ class ModelPatcherDynamic(ModelPatcher):
split = stack_split[0]
while split >= 0:
module, offset = stack[split]
module_pin = module._pins[subset]
split -= 1
stack_split[0] = split
if not module._pin_registered:
if not module_pin["registered"]:
continue
size = module._pin.numel() * module._pin.element_size()
if torch.cuda.cudart().cudaHostUnregister(module._pin.data_ptr()) != 0:
pin = module_pin["pin"]
size = pin.numel() * pin.element_size()
if torch.cuda.cudart().cudaHostUnregister(pin.data_ptr()) != 0:
comfy.model_management.discard_cuda_async_error()
continue
module._pin_registered = False
module_pin["registered"] = False
comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
pinned_size[0] = max(0, pinned_size[0] - size)
freed += size
@@ -1999,20 +2122,23 @@ class ModelPatcherDynamic(ModelPatcher):
return freed
return freed
def partially_unload_ram(self, ram_to_unload, subsets=[ "weights", "patches" ]):
def partially_unload_ram(self, ram_to_unload, subsets=[ "weights-loaded", "patches-loaded", "weights", "patches" ]):
freed = 0
pin_state = self.model.dynamic_pins[self.load_device]
for subset in subsets:
hostbuf, stack, stack_split, pinned_size, *_ = pin_state[subset]
while len(stack) > 0:
module, offset = stack.pop()
size = module._pin.numel() * module._pin.element_size()
module._pin_balancer_entry[-1] = None
del module._pin_balancer_entry
del module._pin
hostbuf.truncate(offset, do_unregister=module._pin_registered)
module_pin = module._pins[subset]
pin = module_pin["pin"]
size = pin.numel() * pin.element_size()
module_pin["balancer_entry"][-1] = None
del module_pin["balancer_entry"]
del module_pin["pin"]
registered = module_pin["registered"]
hostbuf.truncate(offset, do_unregister=registered)
stack_split[0] = min(stack_split[0], len(stack) - 1)
if module._pin_registered:
if registered:
comfy.model_management.TOTAL_PINNED_MEMORY = max(0, comfy.model_management.TOTAL_PINNED_MEMORY - size)
pinned_size[0] = max(0, pinned_size[0] - size)
freed += size
comfy/model_prefetch.py
+98 -12
View File
@@ -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
+21
View File
@@ -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
+6
View File
@@ -51,6 +51,9 @@ class NestedTensor:
def float(self):
return self.to(dtype=torch.float)
def cpu(self):
return self.to(device="cpu")
def chunk(self, *args, **kwargs):
return self.apply_operation(None, lambda x, y: x.chunk(*args, **kwargs))
@@ -80,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
+266 -155
View File
@@ -19,6 +19,7 @@
import torch
import logging
import contextlib
import inspect
import comfy.model_management
from comfy.cli_args import args, PerformanceFeature
import comfy.float
@@ -36,27 +37,61 @@ def run_every_op():
comfy.model_management.throw_exception_if_processing_interrupted()
def gqa_repeat_factor(query_heads, key_heads, value_heads):
if key_heads != value_heads:
raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}")
if query_heads == key_heads:
return 1
if query_heads % key_heads != 0:
raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}")
return query_heads // key_heads
def repeat_kv_for_gqa(k, v, query_heads, head_dim):
n_rep = gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim])
if n_rep > 1:
k = k.repeat_interleave(n_rep, dim=head_dim)
v = v.repeat_interleave(n_rep, dim=head_dim)
return k, v
def 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:
k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
kwargs["enable_gqa"] = False
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
try:
if torch.cuda.is_available() and comfy.model_management.WINDOWS:
if torch.cuda.is_available():
from torch.nn.attention import SDPBackend, sdpa_kernel
import inspect
if "set_priority" in inspect.signature(sdpa_kernel).parameters:
SDPA_BACKEND_PRIORITY = [
SDPBackend.FLASH_ATTENTION,
SDPBackend.CUDNN_ATTENTION,
SDPBackend.EFFICIENT_ATTENTION,
SDPBackend.MATH,
]
SDPA_BACKEND_PRIORITY.insert(0, SDPBackend.CUDNN_ATTENTION)
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)
kwargs["enable_gqa"] = False
with sdpa_kernel(SDPA_BACKEND_PRIORITY, set_priority=True):
if kwargs.get("enable_gqa", False) and attn_mask is not None and q.shape[-3] != k.shape[-3]:
dropout_p = args[1] if len(args) > 1 else kwargs.get("dropout_p", 0.0)
is_causal = args[2] if len(args) > 2 else kwargs.get("is_causal", False)
params = torch.backends.cuda.SDPAParams(q, k, v, attn_mask, dropout_p, is_causal, True)
supports_native_gqa = (
torch.backends.cuda.can_use_flash_attention(params)
or torch.backends.cuda.can_use_cudnn_attention(params)
or torch.backends.cuda.can_use_efficient_attention(params)
)
if not supports_native_gqa:
k, v = repeat_kv_for_gqa(k, v, q.shape[-3], -3)
kwargs["enable_gqa"] = False
return torch.nn.functional.scaled_dot_product_attention(q, k, v, *args, **kwargs)
else:
logging.warning("Torch version too old to set sdpa backend priority.")
@@ -88,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
@@ -128,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,
@@ -144,8 +183,13 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
needs_cast = False
xfer_source = [ s.weight, s.bias ]
pin = comfy.pinned_memory.get_pin(s)
subset = "weights"
pin = comfy.pinned_memory.get_pin(s, subset=subset)
if pin is None and not args.fast_disk:
loaded_pin = comfy.pinned_memory.get_pin(s, subset="weights-loaded")
if loaded_pin is not None or signature is not None:
subset = "weights-loaded"
pin = loaded_pin
if pin is not None:
xfer_source = [ pin ]
@@ -182,12 +226,12 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
if pin is not None:
cast_maybe_lowvram_patch([pin], dest, offload_stream)
return
if signature is None or args.high_ram:
if signature is None or not args.fast_disk or args.high_ram:
comfy.pinned_memory.pin_memory(m, subset=subset, size=size)
pin = comfy.pinned_memory.get_pin(m, subset=subset)
cast_maybe_lowvram_patch(source, pin, offload_stream, xfer_dest2=dest)
handle_pin(s, pin, xfer_source, xfer_dest, size=dest_size)
handle_pin(s, pin, xfer_source, xfer_dest, subset=subset, size=dest_size)
for param_key in ("weight", "bias"):
lowvram_source = getattr(s, param_key + "_lowvram_function", None)
@@ -197,8 +241,16 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin
lowvram_dest = get_cast_buffer(lowvram_size)
lowvram_source.prepare(lowvram_dest, None, copy=False, commit=True)
pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches")
handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset="patches", size=lowvram_size)
subset = "patches"
pin = comfy.pinned_memory.get_pin(lowvram_source, subset=subset)
if pin is None:
loaded_pin = comfy.pinned_memory.get_pin(lowvram_source, subset="patches-loaded")
if loaded_pin is not None:
subset = "patches-loaded"
pin = loaded_pin
elif signature is not None and not args.fast_disk:
subset = "patches-loaded"
handle_pin(lowvram_source, pin, lowvram_source, lowvram_dest, subset=subset, size=lowvram_size)
prefetch["xfer_dest"] = xfer_dest
@@ -207,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)
@@ -250,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)
@@ -277,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):
@@ -404,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
@@ -448,8 +522,7 @@ class disable_weight_init:
def __init__(self, in_features, out_features, bias=True, device=None, dtype=None):
# don't trust subclasses that BYO state dict loader to call us.
if (not comfy.model_management.WINDOWS
or not comfy.memory_management.aimdo_enabled
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):
super().__init__(in_features, out_features, bias, device, dtype)
return
@@ -471,8 +544,7 @@ class disable_weight_init:
def _load_from_state_dict(self, state_dict, prefix, local_metadata,
strict, missing_keys, unexpected_keys, error_msgs):
if (not comfy.model_management.WINDOWS
or not comfy.memory_management.aimdo_enabled
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Linear._load_from_state_dict):
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs)
@@ -492,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()
@@ -509,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()
@@ -526,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()
@@ -554,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()
@@ -571,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()
@@ -588,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()
@@ -607,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()
@@ -630,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()
@@ -657,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()
@@ -681,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()
@@ -700,8 +742,7 @@ class disable_weight_init:
norm_type=2.0, scale_grad_by_freq=False, sparse=False, _weight=None,
_freeze=False, device=None, dtype=None):
# don't trust subclasses that BYO state dict loader to call us.
if (not comfy.model_management.WINDOWS
or not comfy.memory_management.aimdo_enabled
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):
super().__init__(num_embeddings, embedding_dim, padding_idx, max_norm,
norm_type, scale_grad_by_freq, sparse, _weight,
@@ -728,8 +769,7 @@ class disable_weight_init:
def _load_from_state_dict(self, state_dict, prefix, local_metadata,
strict, missing_keys, unexpected_keys, error_msgs):
if (not comfy.model_management.WINDOWS
or not comfy.memory_management.aimdo_enabled
if (not comfy.memory_management.aimdo_enabled
or type(self)._load_from_state_dict is not disable_weight_init.Embedding._load_from_state_dict):
return super()._load_from_state_dict(state_dict, prefix, local_metadata, strict,
missing_keys, unexpected_keys, error_msgs)
@@ -751,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):
@@ -830,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)
@@ -839,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
@@ -867,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:
@@ -886,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()
@@ -901,13 +935,61 @@ if CUBLAS_IS_AVAILABLE:
# ==============================================================================
# Mixed Precision Operations
# ==============================================================================
from . import quant_ops
from .quant_ops import (
QuantizedTensor,
QUANT_ALGOS,
TensorCoreFP8Layout,
TensorWiseINT8Layout,
get_layout_class,
)
def _swiglu_eager(x):
gate, up = x.chunk(2, dim=-1)
return torch.nn.functional.silu(gate).mul_(up)
INPUT_ACT_EAGER = {
"gelu_tanh": lambda x: torch.nn.functional.gelu(x, approximate="tanh"),
"swiglu": _swiglu_eager,
}
def linear_input_act(linear, x, input_act):
"""``linear(act(x))``, with ``act`` folded into an INT8 activation quantizer.
An INT8 linear quantizes its input anyway, so an elementwise activation can
ride along inside that kernel instead of writing a full-size intermediate to
HBM and reading it straight back. Worth it for an MLP's down-projection,
where the intermediate is several times the hidden size.
"""
weight = linear.weight
if (comfy.model_management.in_training
or not isinstance(weight, QuantizedTensor)
or weight._layout_cls != "TensorWiseINT8Layout"
or getattr(weight._params, "transposed", False)):
return linear(INPUT_ACT_EAGER[input_act](x))
# want_requant keeps a vbar-streamed layer on the INT8 path when a LoRA is
# patched in on the fly; without it the cast hands back a dequantized weight.
weight, bias, offload_stream = cast_bias_weight(
linear, x, offloadable=True, compute_dtype=x.dtype, want_requant=True)
try:
if not isinstance(weight, QuantizedTensor):
# A LoRA weight_function, or activations whose dtype differs from the
# weight's, make the cast hand back a dequantized tensor.
return torch.nn.functional.linear(INPUT_ACT_EAGER[input_act](x), weight, bias)
qdata, scale = TensorWiseINT8Layout.get_plain_tensors(weight)
return quant_ops.ck.int8_linear(
x, qdata, scale, bias, x.dtype,
convrot=getattr(weight._params, "convrot", False),
convrot_groupsize=getattr(weight._params, "convrot_groupsize", 256),
input_act=input_act,
)
finally:
uncast_bias_weight(linear, weight, bias, offload_stream)
class QuantLinearFunc(torch.autograd.Function):
"""Custom autograd function for quantized linear: quantized forward, optionally FP8 backward.
@@ -1119,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}")
@@ -1170,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)
@@ -1217,7 +1322,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
def state_dict(self, *args, destination=None, prefix="", **kwargs):
sd = destination if destination is not None else {}
return _quantized_weight_state_dict(self, sd, prefix, extra_quant_params=("input_scale",))
return _quantized_weight_state_dict(self, sd, prefix, extra_quant_params=("input_scale", "pre_quant_scale"))
def _forward(self, input, weight, bias):
return torch.nn.functional.linear(input, weight, bias)
@@ -1229,33 +1334,37 @@ 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()
# ModelOpt AWQ-style smoothing
pre_quant_scale = getattr(self, 'pre_quant_scale', None)
if pre_quant_scale is not None:
input = input * comfy.model_management.cast_to_device(pre_quant_scale, input.device, input.dtype)
input_shape = input.shape
reshaped_nd = False
#If cast needs to apply lora, it should be done in the compute dtype
@@ -1271,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:
@@ -1400,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)."""
@@ -1414,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):
@@ -1469,12 +1568,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
if layer_conf is not None:
layer_conf = json.loads(layer_conf.numpy().tobytes())
# Only fp8 makes sense for embeddings (per-row dequant via index select).
# Only fp8 and int8_tensorwise support per-row dequant via index select.
# Block-scaled formats (NVFP4, MXFP8) can't do per-row lookup efficiently.
quant_format = layer_conf.get("format") if layer_conf is not None else None
manually_loaded_keys = []
if quant_format in ("float8_e4m3fn", "float8_e5m2") and weight_key in state_dict:
if quant_format in ("float8_e4m3fn", "float8_e5m2", "int8_tensorwise") and weight_key in state_dict:
self.quant_format = quant_format
qconfig = QUANT_ALGOS[quant_format]
self.layout_type = qconfig["comfy_tensor_layout"]
@@ -1488,10 +1587,16 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
scale = scale.float()
manually_loaded_keys.append(scale_key)
extra = {}
if quant_format == "int8_tensorwise" and layer_conf.get("convrot", False):
# rotated embedding table: record it so the forward un-rotates after lookup
extra["convrot"] = True
extra["convrot_groupsize"] = int(layer_conf.get("convrot_groupsize", 256))
params = layout_cls.Params(
scale=scale if scale is not None else torch.ones((), dtype=torch.float32),
orig_dtype=MixedPrecisionOps._compute_dtype,
orig_shape=(self.num_embeddings, self.embedding_dim),
**extra,
)
self.weight = torch.nn.Parameter(
QuantizedTensor(weight.to(dtype=qconfig["storage_t"]), qconfig["comfy_tensor_layout"], params),
@@ -1513,22 +1618,28 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
def forward_comfy_cast_weights(self, input, out_dtype=None):
weight = self.weight
# Optimized path: lookup in fp8, dequantize only the selected rows.
# Optimized path: lookup in fp8/int8, dequantize only the selected rows.
if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
qdata, _, offload_stream = cast_bias_weight(self, device=input.device, dtype=weight.dtype, offloadable=True)
if isinstance(qdata, QuantizedTensor):
scale = qdata._params.scale
qdata = qdata._qdata
else:
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
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)
# 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)
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/pinned_memory.py
+36 -28
View File
@@ -9,14 +9,14 @@ import torch
from comfy.cli_args import args
def _add_to_bucket(module, buckets, size, priority):
def _add_to_bucket(module, module_pin, buckets, size, priority):
bucket = buckets.setdefault(size, [])
entry = [-priority, 0, module]
entry[1] = id(entry)
bisect.insort(bucket, entry)
module._pin_balancer_entry = entry
module_pin["balancer_entry"] = entry
def _steal_pin(module, stack, buckets, size, priority):
def _steal_pin(module, stack, buckets, size, priority, subset):
bucket = buckets.get(size)
if bucket is None:
return False
@@ -31,22 +31,27 @@ def _steal_pin(module, stack, buckets, size, priority):
return False
*_, victim = bucket.pop()
module._pin = victim._pin
module._pin_registered = victim._pin_registered
module._pin_stack_index = victim._pin_stack_index
stack[module._pin_stack_index] = (module, stack[module._pin_stack_index][1])
module_pin = module._pins[subset]
victim_pin = victim._pins[subset]
module_pin["pin"] = victim_pin["pin"]
module_pin["registered"] = victim_pin["registered"]
module_pin["stack_index"] = victim_pin["stack_index"]
stack_index = module_pin["stack_index"]
stack[stack_index] = (module, stack[stack_index][1])
victim._pin_registered = False
del victim._pin
del victim._pin_stack_index
del victim._pin_balancer_entry
victim_pin["registered"] = False
del victim_pin["pin"]
del victim_pin["stack_index"]
del victim_pin["balancer_entry"]
_add_to_bucket(module, buckets, size, priority)
_add_to_bucket(module, module_pin, buckets, size, priority)
return True
def get_pin(module, subset="weights"):
pin = getattr(module, "_pin", None)
if pin is None or module._pin_registered or args.disable_pinned_memory:
pins = module.__dict__.get("_pins")
module_pin = None if pins is None else pins.get(subset)
pin = None if module_pin is None else module_pin.get("pin")
if pin is None or module_pin["registered"] or args.disable_pinned_memory:
return pin
_, _, stack_split, pinned_size, *_ = module._pin_state[subset]
@@ -57,8 +62,8 @@ def get_pin(module, subset="weights"):
comfy.model_management.discard_cuda_async_error()
return pin
module._pin_registered = True
stack_split[0] = max(stack_split[0], module._pin_stack_index)
module_pin["registered"] = True
stack_split[0] = max(stack_split[0], module_pin["stack_index"])
comfy.model_management.TOTAL_PINNED_MEMORY += size
pinned_size[0] += size
return pin
@@ -72,23 +77,26 @@ def pin_memory(module, subset="weights", size=None):
if pin is not None:
return
pins = module.__dict__.setdefault("_pins", {})
module_pin = pins.setdefault(subset, {})
hostbuf, stack, stack_split, pinned_size, counter, buckets = pin_state[subset]
if size is None:
size = comfy.memory_management.vram_aligned_size([ module.weight, module.bias ])
offset = hostbuf.size
registerable_size = size
priority = getattr(module, "_pin_balancer_priority", None)
loaded = subset.endswith("-loaded")
priority = module_pin.get("balancer_priority")
if priority is None:
priority = comfy.utils.bit_reverse_range(counter[0], 16)
counter[0] += 1
module._pin_balancer_priority = priority
module_pin["balancer_priority"] = priority
comfy.memory_management.extra_ram_release(comfy.memory_management.RAM_CACHE_HEADROOM)
if (not comfy.model_management.ensure_pin_budget(size) or
if (not comfy.model_management.ensure_pin_budget(size, loaded=loaded) or
not comfy.model_management.ensure_pin_registerable(registerable_size)):
return _steal_pin(module, stack, buckets, size, priority)
return _steal_pin(module, stack, buckets, size, priority, subset)
offset = hostbuf.size
extended = False
try:
hostbuf.extend(size=size, register=False)
@@ -102,18 +110,18 @@ def pin_memory(module, subset="weights", size=None):
comfy.model_management.discard_cuda_async_error()
del pin
hostbuf.truncate(offset, do_unregister=False)
return _steal_pin(module, stack, buckets, size, priority)
return _steal_pin(module, stack, buckets, size, priority, subset)
except RuntimeError:
if extended:
hostbuf.truncate(offset, do_unregister=False)
return _steal_pin(module, stack, buckets, size, priority)
return _steal_pin(module, stack, buckets, size, priority, subset)
module._pin = pin
module_pin["pin"] = pin
stack.append((module, offset))
module._pin_registered = True
module._pin_stack_index = len(stack) - 1
stack_split[0] = max(stack_split[0], module._pin_stack_index)
module_pin["registered"] = True
module_pin["stack_index"] = len(stack) - 1
stack_split[0] = max(stack_split[0], module_pin["stack_index"])
comfy.model_management.TOTAL_PINNED_MEMORY += size
pinned_size[0] += size
_add_to_bucket(module, buckets, size, priority)
_add_to_bucket(module, module_pin, buckets, size, priority)
return True
comfy/quant_ops.py
+16 -3
View File
@@ -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": {
@@ -240,7 +245,7 @@ QUANT_ALGOS = {
},
"nvfp4": {
"storage_t": torch.uint8,
"parameters": {"weight_scale", "weight_scale_2", "input_scale"},
"parameters": {"weight_scale", "weight_scale_2", "input_scale", "pre_quant_scale"},
"comfy_tensor_layout": "TensorCoreNVFP4Layout",
"group_size": 16,
},
@@ -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
+9 -1
View File
@@ -37,10 +37,15 @@ 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
latent_format = model.get_model_object("latent_format") #Resize the empty latent image so it has the right number of channels
latent_format = model.get_model_object("latent_format")
is_empty = torch.count_nonzero(latent_image) == 0
if is_empty:
if latent_format.latent_channels != latent_image.shape[1]:
@@ -59,6 +64,9 @@ def fix_empty_latent_channels(model, latent_image, downscale_ratio_spacial=None,
new_t = max(1, round(latent_image.shape[2] * ratio))
latent_image = comfy.utils.repeat_to_batch_size(latent_image, new_t, dim=2)
if is_empty:
latent_image = latent_format.fix_empty_latent(latent_image)
return latent_image
def prepare_sampling(model, noise_shape, positive, negative, noise_mask):
comfy/samplers.py
+24 -5
View File
@@ -20,6 +20,7 @@ import comfy.hooks
import comfy.context_windows
import comfy.multigpu
import comfy.utils
from comfy.internal_logging import detail
import scipy.stats
import numpy
@@ -635,7 +636,7 @@ class KSamplerX0Inpaint:
if "denoise_mask_function" in model_options:
denoise_mask = model_options["denoise_mask_function"](sigma, denoise_mask, extra_options={"model": self.inner_model, "sigmas": self.sigmas})
latent_mask = 1. - denoise_mask
x = x * denoise_mask + self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image) * latent_mask
x = x * denoise_mask + self.inner_model.inner_model.scale_latent_inpaint(x=x, sigma=sigma, noise=self.noise, latent_image=self.latent_image, denoise_mask=denoise_mask) * latent_mask
out = self.inner_model(x, sigma, model_options=model_options, seed=seed)
if denoise_mask is not None:
out = out * denoise_mask + self.latent_image * latent_mask
@@ -970,7 +971,7 @@ class Sampler:
KSAMPLER_NAMES = ["euler", "euler_cfg_pp", "euler_ancestral", "euler_ancestral_cfg_pp", "heun", "heunpp2", "exp_heun_2_x0", "exp_heun_2_x0_sde", "dpm_2", "dpm_2_ancestral",
"lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_2s_ancestral_cfg_pp", "dpmpp_sde", "dpmpp_sde_gpu",
"dpmpp_2m", "dpmpp_2m_cfg_pp", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_2m_sde_heun", "dpmpp_2m_sde_heun_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm",
"ipndm", "ipndm_v", "deis", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp",
"ipndm", "ipndm_v", "deis", "cfgpp_ud10_ab", "res_multistep", "res_multistep_cfg_pp", "res_multistep_ancestral", "res_multistep_ancestral_cfg_pp",
"gradient_estimation", "gradient_estimation_cfg_pp", "er_sde", "seeds_2", "seeds_3", "sa_solver", "sa_solver_pece"]
class KSAMPLER(Sampler):
@@ -991,10 +992,15 @@ class KSAMPLER(Sampler):
noise = model_wrap.inner_model.model_sampling.noise_scaling(sigmas[0], noise, latent_image, self.max_denoise(model_wrap, sigmas))
k_callback = None
total_steps = len(sigmas) - 1
if callback is not None:
k_callback = lambda x: callback(x["i"], x["denoised"], x["x"], total_steps)
first_step = True
def k_callback(x):
nonlocal first_step
if first_step:
detail("First sampler step: model=%s sampler=%s step=%s total_steps=%s cfg=%s seed=%s sigma=%s sigma_hat=%s latent_shape=%s denoised_shape=%s", model_wrap.model_patcher.model.__class__.__name__, getattr(self.sampler_function, "__name__", "unknown"), x["i"], total_steps, model_wrap.cfg, extra_args.get("seed"), x.get("sigma"), x.get("sigma_hat"), tuple(x["x"].shape), tuple(x["denoised"].shape))
first_step = False
if callback is not None:
callback(x["i"], x["denoised"], x["x"], total_steps)
samples = self.sampler_function(model_k, noise, sigmas, extra_args=extra_args, callback=k_callback, disable=disable_pbar, **self.extra_options)
samples = model_wrap.inner_model.model_sampling.inverse_noise_scaling(sigmas[-1], samples)
@@ -1212,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)
@@ -1299,10 +1307,21 @@ class CFGGuider:
return latent_image
if latent_image.is_nested:
sampler_shapes = [tuple(x.shape) for x in latent_image.unbind()]
latent_image, latent_shapes = comfy.utils.pack_latents(latent_image.unbind())
noise, _ = comfy.utils.pack_latents(noise.unbind())
else:
latent_shapes = [latent_image.shape]
sampler_shapes = [tuple(latent_image.shape)]
detail("Sampler: model=%s latent_shapes=%s", self.model_patcher.model.__class__.__name__, sampler_shapes)
if len(latent_shapes) > 1 and callback is not None:
# samplers run on the flat pack, hand callbacks (previews, x0 output) the nested view
packed_callback = callback
def callback(step, x0, x, total_steps):
x0 = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x0, latent_shapes))
x = comfy.nested_tensor.NestedTensor(comfy.utils.unpack_latents(x, latent_shapes))
return packed_callback(step, x0, x, total_steps)
if denoise_mask is not None:
if denoise_mask.is_nested:
comfy/sd.py
+238 -14
View File
@@ -11,18 +11,22 @@ 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
import comfy.ldm.trellis2.vae
import comfy.ldm.wan.vae2_2
import comfy.ldm.hunyuan3d.vae
import comfy.ldm.seedvr.vae
import comfy.ldm.mage_flow.vae
import comfy.ldm.triposplat.vae
import comfy.ldm.ace.vae.music_dcae_pipeline
import comfy.ldm.cogvideo.vae
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
@@ -30,6 +34,7 @@ import math
import os
import comfy.utils
import comfy.ops
from . import clip_vision
from . import gligen
@@ -60,6 +65,7 @@ import comfy.text_encoders.qwen_image
import comfy.text_encoders.hunyuan_image
import comfy.text_encoders.z_image
import comfy.text_encoders.krea2
import comfy.text_encoders.mage_flow
import comfy.text_encoders.ideogram4
import comfy.text_encoders.ovis
import comfy.text_encoders.kandinsky5
@@ -70,6 +76,10 @@ import comfy.text_encoders.ace15
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
import comfy.text_encoders.ernie
import comfy.text_encoders.gemma4
@@ -512,7 +522,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]
@@ -554,6 +579,16 @@ class VAE:
self.first_stage_model = StageC_coder()
self.downscale_ratio = 32
self.latent_channels = 16
elif "shape_dec.blocks.1.16.to_subdiv.weight" in sd: # trellis2 shape vae (struct_dec + shape_dec)
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
self.memory_used_decode = lambda shape, dtype: (2500 * math.prod(shape[2:])) * model_management.dtype_size(dtype)
self.memory_used_encode = lambda shape, dtype: (2500 * math.prod(shape[2:])) * model_management.dtype_size(dtype)
self.first_stage_model = comfy.ldm.trellis2.vae.ShapeVae()
elif "txt_dec.blocks.3.4.conv2.weight" in sd: # trellis2 texture vae
self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
self.memory_used_decode = lambda shape, dtype: (2500 * math.prod(shape[2:])) * model_management.dtype_size(dtype)
self.memory_used_encode = lambda shape, dtype: (2500 * math.prod(shape[2:])) * model_management.dtype_size(dtype)
self.first_stage_model = comfy.ldm.trellis2.vae.TextureVae()
elif "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd: # seedvr2
self.first_stage_model = comfy.ldm.seedvr.vae.VideoAutoencoderKLWrapper()
self.latent_channels = comfy.ldm.seedvr.vae.SEEDVR2_LATENT_CHANNELS
@@ -570,6 +605,33 @@ class VAE:
self.upscale_index_formula = (4, 8, 8)
self.process_input = lambda image: image * 2.0 - 1.0
self.crop_input = False
elif "student.dconv_encoder.proj_out.weight" in sd: # Mage-VAE (one-step diffusion codec, Flux2-anchored 128ch/16x latents)
sd = comfy.utils.state_dict_prefix_replace(sd, {"student.dconv_encoder.": "dconv_encoder.", "pipeline.": "decoder_model."})
# Drop the unused Flux2-VAE anchor encoder carried in the checkpoint.
sd = {k: v for k, v in sd.items() if not k.startswith("decoder_model.y_embedder.encoder.") and not k.startswith("decoder_model.y_embedder.bottleneck.")}
self.first_stage_model = comfy.ldm.mage_flow.vae.MageVAE()
self.latent_channels = 128
self.downscale_ratio = 16
self.upscale_ratio = 16
self.working_dtypes = [torch.bfloat16, torch.float32]
self.memory_used_encode = lambda shape, dtype: (400 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
self.memory_used_decode = lambda shape, dtype: (1000 * shape[2] * shape[3] * 16 * 16) * model_management.dtype_size(dtype)
elif "decoder.conv_in_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}
@@ -859,7 +921,14 @@ class VAE:
self.upscale_index_formula = (4, 16, 16)
self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 16, 16)
self.downscale_index_formula = (4, 16, 16)
if self.latent_channels in [48, 128]: # Wan 2.2 and LTX2
if self.latent_channels == 24 and sd["decoder.22.bias"].shape[0] == 12: # MiniMax H3
self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=None)
self.process_input = self.process_output = lambda image: image
self.upscale_ratio = (lambda a: max(1, (a - 2) // 5 * 17 + 5), 16, 16)
self.downscale_ratio = (lambda a: max(1, (a - 1) // 17 * 5 + 2) if a > 1 else 1, 16, 16)
self.memory_used_encode = lambda shape, dtype: (400 * ((shape[-3] + 16) // 17) * shape[-2] * shape[-1] * model_management.dtype_size(dtype))
self.memory_used_decode = lambda shape, dtype: ((260 * 16 * 16 + shape[1] * shape[-3]) * shape[-2] * shape[-1] * model_management.dtype_size(dtype))
elif self.latent_channels in [48, 128]: # Wan 2.2 and LTX2
self.first_stage_model = comfy.taesd.taehv.TAEHV(latent_channels=self.latent_channels, latent_format=None) # taehv doesn't need scaling
self.process_input = self.process_output = lambda image: image
self.process_output = lambda image: image
@@ -926,6 +995,62 @@ class VAE:
#Force cast it for --disable-dynamic-vram users until there is a true core fix.
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
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
self.upscale_ratio = (lambda a: max(1, (a - 2) // 5 * 17 + 5), 16, 16)
self.upscale_index_formula = (4, 16, 16)
self.downscale_ratio = (lambda a: max(1, (a - 5) // 17 * 5 + 2) if a > 1 else 1, 16, 16)
self.downscale_index_formula = (4, 16, 16)
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)
self.memory_used_decode = lambda shape, dtype: estimate_decode_memory(self.upscale_ratio[0](shape[2]), shape[3] * self.upscale_ratio[1], shape[4] * self.upscale_ratio[2], dtype)
elif "pre_block.attn.zero_k_bias" in sd: # MiniMax H3 audio VAE (DAC encoder + BigVGAN decoder)
self.first_stage_model = comfy.ldm.minimax.audio_vae.MiniMaxH3AudioVAE()
self.latent_channels = 32
self.output_channels = 2
self.pad_channel_value = "replicate"
self.audio_sample_rate = 32000
self.upscale_ratio = 800
self.downscale_ratio = 800
self.latent_dim = 2 # [B, 32, stereo 2, T]
self.process_output = lambda audio: audio
self.process_input = lambda audio: audio
self.working_dtypes = [torch.float32]
# encode gets the waveform shape [B, 2, samples], decode the latent shape [B, 32, 2, T]
def estimate_encode_memory(samples, dtype):
return (900 * samples + 105_000_000) * model_management.dtype_size(dtype) * 1.03
def estimate_decode_memory(samples, dtype):
return max(42_000_000, 220 * samples + 20_000_000) * model_management.dtype_size(dtype) * 1.03
self.memory_used_encode = lambda shape, dtype: estimate_encode_memory(shape[2], dtype)
self.memory_used_decode = lambda shape, dtype: estimate_decode_memory(shape[-1] * self.upscale_ratio, dtype)
elif "gs.base_offset_scale" in sd and "octree.out_proj.weight" in sd: # TripoSplat octree gaussian decoder
self.first_stage_model = comfy.ldm.triposplat.vae.OctreeGaussianDecoder()
self.latent_channels = 16
@@ -1137,6 +1262,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:
@@ -1152,16 +1278,57 @@ 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 prepare_decode(self, sample_shape, memory_required=None):
"""For VAEs whose real decode entry point bypasses decode()"""
if memory_required is None:
memory_required = self.memory_used_decode(sample_shape, self.vae_dtype)
memory_required = max(1, int(memory_required))
model_management.load_models_gpu([self.patcher], memory_required=memory_required, force_full_load=self.disable_offload)
free_memory = self.patcher.get_free_memory(self.device)
return max(1, int(free_memory / memory_required))
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 = {}
@@ -1382,6 +1549,8 @@ class CLIPType(Enum):
BOOGU = 31
KREA2 = 32
JOYIMAGE = 33
MAGE = 34
MINIMAX = 35
@@ -1437,6 +1606,8 @@ class TEModel(Enum):
GPT_OSS_20B = 33
QWEN3VL_4B = 34
QWEN3VL_8B = 35
GEMMA_4_12B = 36
QWEN3VL_32B = 37
def detect_te_model(sd):
@@ -1466,6 +1637,9 @@ def detect_te_model(sd):
if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
if 'model.layers.59.self_attn.q_norm.weight' in sd:
return TEModel.GEMMA_4_31B
# Gemma4 12B Unified: 48 layers, encoder-free; global layers drop v_proj (attention_k_eq_v).
if 'model.layers.47.self_attn.q_norm.weight' in sd and 'model.layers.5.self_attn.v_proj.weight' not in sd:
return TEModel.GEMMA_4_12B
if 'model.layers.41.self_attn.q_norm.weight' in sd and 'model.layers.47.self_attn.q_norm.weight' not in sd:
return TEModel.GEMMA_4_E4B
if 'model.layers.34.self_attn.q_norm.weight' in sd and 'model.layers.41.self_attn.q_norm.weight' not in sd:
@@ -1500,6 +1674,9 @@ def detect_te_model(sd):
return TEModel.QWEN35_2B
if "model.visual.deepstack_merger_list.0.norm.weight" in sd: # DeepStack is unique to Qwen3-VL
return TEModel.QWEN3VL_4B if sd["model.visual.merger.linear_fc2.weight"].shape[0] == 2560 else TEModel.QWEN3VL_8B
if "visual.deepstack_merger_list.0.norm.weight" in sd and "model.layers.49.self_attn.q_proj.weight" in sd:
# MiniMax H3 conditioning encoder: Qwen3-VL-32B, truncated to 50 layers
return TEModel.QWEN3VL_32B
if "model.layers.0.post_attention_layernorm.weight" in sd:
weight = sd['model.layers.0.post_attention_layernorm.weight']
if 'model.layers.0.self_attn.q_norm.weight' in sd:
@@ -1563,7 +1740,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
@@ -1621,12 +1807,22 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
clip_target.clip = comfy.text_encoders.sa3.SAT5GemmaModel
clip_target.tokenizer = comfy.text_encoders.sa3.SAT5GemmaTokenizer
tokenizer_data["spiece_model"] = clip_data[0].get("spiece_model", None)
elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B):
variant = {TEModel.GEMMA_4_E4B: comfy.text_encoders.gemma4.Gemma4_E4B,
TEModel.GEMMA_4_E2B: comfy.text_encoders.gemma4.Gemma4_E2B,
TEModel.GEMMA_4_31B: comfy.text_encoders.gemma4.Gemma4_31B}[te_model]
clip_target.clip = comfy.text_encoders.gemma4.gemma4_te(**llama_detect(clip_data), model_class=variant)
clip_target.tokenizer = variant.tokenizer
elif te_model in (TEModel.GEMMA_4_E4B, TEModel.GEMMA_4_E2B, TEModel.GEMMA_4_31B, TEModel.GEMMA_4_12B):
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:
@@ -1711,6 +1907,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
elif clip_type == CLIPType.MAGE and te_model == TEModel.QWEN3VL_4B: # Mage-Flow: full Qwen3-VL-4B, last hidden state, Qwen-Image-style templates.
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
clip_target.clip = comfy.text_encoders.mage_flow.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.mage_flow.MageFlowTokenizer
elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B: # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx.
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data))
@@ -1724,6 +1924,9 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
qwen3vl_type = {TEModel.QWEN3VL_4B: "qwen3vl_4b", TEModel.QWEN3VL_8B: "qwen3vl_8b"}[te_model]
clip_target.clip = comfy.text_encoders.qwen3vl.te(**llama_detect(clip_data), model_type=qwen3vl_type)
clip_target.tokenizer = comfy.text_encoders.qwen3vl.tokenizer(model_type=qwen3vl_type)
elif te_model == TEModel.QWEN3VL_32B:
clip_target.clip = comfy.text_encoders.minimax.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.minimax.MiniMaxH3Tokenizer
elif te_model == TEModel.QWEN3_06B:
clip_target.clip = comfy.text_encoders.anima.te(**llama_detect(clip_data))
clip_target.tokenizer = comfy.text_encoders.anima.AnimaTokenizer
@@ -1787,9 +1990,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
+166
View File
@@ -15,6 +15,8 @@ import comfy.text_encoders.flux
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
@@ -27,6 +29,7 @@ import comfy.text_encoders.z_image
import comfy.text_encoders.ideogram4
import comfy.text_encoders.boogu
import comfy.text_encoders.krea2
import comfy.text_encoders.mage_flow
import comfy.text_encoders.joyimage
import comfy.text_encoders.anima
import comfy.text_encoders.ace15
@@ -34,6 +37,7 @@ import comfy.text_encoders.longcat_image
import comfy.text_encoders.ernie
import comfy.text_encoders.cogvideo
import comfy.text_encoders.hidream_o1
import comfy.text_encoders.sensenova
import comfy.text_encoders.pixeldit
from . import supported_models_base
@@ -954,6 +958,34 @@ class LTXAV(LTXV):
out = model_base.LTXAV(self, device=device)
return out
class MiniMaxH3(supported_models_base.BASE):
unet_config = {
"image_model": "minimax_h3",
}
sampling_settings = {
"shift": 12.0,
"audio_shift": 3.0,
}
unet_extra_config = {}
latent_format = latent_formats.MiniMaxH3AV
memory_usage_factor = 0.114
supported_inference_dtypes = [torch.bfloat16, torch.float32]
vae_key_prefix = ["vae."]
text_encoder_key_prefix = ["text_encoders."]
def get_model(self, state_dict, prefix="", device=None):
return model_base.MiniMaxH3(self, device=device)
def clip_target(self, state_dict={}, prefix=""):
pref = self.text_encoder_key_prefix[0]
detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_32b.transformer.".format(pref))
return supported_models_base.ClipTarget(comfy.text_encoders.minimax.MiniMaxH3Tokenizer, comfy.text_encoders.minimax.te(**detect))
class HunyuanVideo(supported_models_base.BASE):
unet_config = {
"image_model": "hunyuan_video",
@@ -1420,6 +1452,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",
@@ -1433,6 +1479,31 @@ class WAN22_T2V(WAN21_T2V):
out = model_base.WAN22(self, image_to_video=True, device=device)
return out
class Trellis2(supported_models_base.BASE):
unet_config = {
"image_model": "trellis2"
}
unet_extra_config = {"num_heads": 12}
sampling_settings = {
"shift": 3.0,
}
memory_usage_factor = 6
latent_format = latent_formats.Trellis2
vae_key_prefix = ["vae."]
clip_vision_prefix = "conditioner.main_image_encoder.model."
# this is only needed for the texture model
supported_inference_dtypes = [torch.bfloat16, torch.float32]
def get_model(self, state_dict, prefix="", device=None):
return model_base.Trellis2(self, device=device)
def clip_target(self, state_dict={}):
return None
class WAN21_FlowRVS(WAN21_T2V):
unet_config = {
"image_model": "wan2.1",
@@ -1650,6 +1721,44 @@ class HiDreamO1(supported_models_base.BASE):
comfy.text_encoders.hidream_o1.HiDreamO1TE,
)
class SenseNovaU15(supported_models_base.BASE):
unet_config = {
"image_model": "sensenova_u15",
}
sampling_settings = {
"shift": 3.0,
"noise_scale": 1.0,
}
latent_format = latent_formats.HiDreamO1Pixel
memory_usage_factor = 0.033
supported_inference_dtypes = [torch.bfloat16, torch.float32]
vae_key_prefix = ["vae."]
text_encoder_key_prefix = ["text_encoders."]
optimizations = {"fp8": False}
def get_model(self, state_dict, prefix="", device=None):
return model_base.SenseNovaU15(self, device=device)
def process_unet_state_dict(self, state_dict):
state_dict.pop("language_model.lm_head.weight", None)
return state_dict
def process_vae_state_dict(self, state_dict):
return {"pixel_space_vae": torch.tensor(1.0)}
def process_clip_state_dict(self, state_dict):
return {"_sensenova_te_sentinel": torch.zeros(1)}
def clip_target(self, state_dict={}):
return supported_models_base.ClipTarget(
comfy.text_encoders.sensenova.SenseNovaTokenizer,
comfy.text_encoders.sensenova.SenseNovaTextEncoder,
)
class Chroma(supported_models_base.BASE):
unet_config = {
"image_model": "chroma",
@@ -1883,6 +1992,35 @@ class Krea2(supported_models_base.BASE):
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref))
return supported_models_base.ClipTarget(comfy.text_encoders.krea2.Krea2Tokenizer, comfy.text_encoders.krea2.te(**hunyuan_detect))
class MageFlow(supported_models_base.BASE):
unet_config = {
"image_model": "mage_flow",
}
sampling_settings = {
"multiplier": 1.0,
"shift": 6.0,
}
memory_usage_factor = 6.5
unet_extra_config = {}
latent_format = latent_formats.Flux2
supported_inference_dtypes = [torch.bfloat16, torch.float32]
vae_key_prefix = ["vae."]
text_encoder_key_prefix = ["text_encoders."]
def get_model(self, state_dict, prefix="", device=None):
out = model_base.MageFlow(self, device=device)
return out
def clip_target(self, state_dict={}):
pref = self.text_encoder_key_prefix[0]
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref))
return supported_models_base.ClipTarget(comfy.text_encoders.mage_flow.MageFlowTokenizer, comfy.text_encoders.mage_flow.te(**hunyuan_detect))
class QwenImage(supported_models_base.BASE):
unet_config = {
"image_model": "qwen_image",
@@ -2127,6 +2265,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 = {
@@ -2377,6 +2537,7 @@ models = [
GenmoMochi,
LTXV,
LTXAV,
MiniMaxH3,
HunyuanVideo15_SR_Distilled,
HunyuanVideo15,
HunyuanImage21Refiner,
@@ -2404,6 +2565,7 @@ models = [
WAN22_S2V,
WAN21_HuMo,
WAN22_Animate,
WAN_Animate2,
WAN21_FlowRVS,
WAN21_SCAIL,
WAN21_SCAIL2,
@@ -2414,13 +2576,16 @@ models = [
TripoSplat,
HiDream,
HiDreamO1,
SenseNovaU15,
Chroma,
SeedVR2,
ChromaRadiance,
ACEStep,
ACEStep15,
MiniMaxMusic3,
Omnigen2,
Boogu,
MageFlow,
QwenImage,
JoyImage,
Ideogram4,
@@ -2438,5 +2603,6 @@ models = [
CogVideoX_I2V,
CogVideoX_T2V,
SVD_img2vid,
Trellis2,
DepthAnything3,
]
comfy/system_memory.py
+128
View File
@@ -0,0 +1,128 @@
import os
import sys
import psutil
CGROUP_V2_ROOT = "/sys/fs/cgroup"
CGROUP_V1_MEMORY_ROOT = "/sys/fs/cgroup/memory"
PROC_SELF_CGROUP = "/proc/self/cgroup"
_cgroup_dirs = None
def _read_text(path):
try:
with open(path, encoding="utf-8") as f:
return f.read().strip()
except (OSError, ValueError):
return None
def _read_int(path):
raw = _read_text(path)
if raw is None:
return None
try:
return int(raw)
except ValueError:
return None
def _read_stat(path, key):
raw = _read_text(path)
if raw is None:
return None
for line in raw.splitlines():
parts = line.split()
if len(parts) == 2 and parts[0] == key:
try:
return int(parts[1])
except ValueError:
return None
return None
def _lineage(root, path):
dirs = [root]
for part in path.split("/"):
if part:
dirs.append(os.path.join(dirs[-1], part))
return dirs[::-1]
def _own_cgroup_dirs():
raw = _read_text(PROC_SELF_CGROUP)
if raw is None:
return []
dirs = []
for line in raw.splitlines():
parts = line.split(":", 2)
if len(parts) != 3:
continue
controllers, path = parts[1], parts[2]
if controllers == "":
root = CGROUP_V2_ROOT
elif "memory" in controllers.split(","):
root = CGROUP_V1_MEMORY_ROOT
else:
continue
for directory in _lineage(root, path):
if directory not in dirs:
dirs.append(directory)
return dirs
def _cgroup_directories():
global _cgroup_dirs
if _cgroup_dirs is None:
_cgroup_dirs = _own_cgroup_dirs() or [CGROUP_V2_ROOT, CGROUP_V1_MEMORY_ROOT]
return _cgroup_dirs
def _limit_in(directory):
for name in ("memory.max", "memory.limit_in_bytes"):
value = _read_int(os.path.join(directory, name))
if value is not None and value > 0:
return value
return None
def _working_set_in(directory):
usage = _read_int(os.path.join(directory, "memory.current"))
key = "inactive_file"
if usage is None:
usage = _read_int(os.path.join(directory, "memory.usage_in_bytes"))
key = "total_inactive_file"
if usage is None:
return None
inactive_file = _read_stat(os.path.join(directory, "memory.stat"), key) or 0
return max(0, usage - inactive_file)
def _limited(host_total):
if not sys.platform.startswith("linux"):
return []
limited = []
for directory in _cgroup_directories():
limit = _limit_in(directory)
if limit is not None and limit < host_total:
limited.append((limit, directory))
return limited
def cgroup_memory_limit():
return min((limit for limit, _ in _limited(psutil.virtual_memory().total)), default=None)
def virtual_memory_total():
host = psutil.virtual_memory()
return min((limit for limit, _ in _limited(host.total)), default=host.total)
def virtual_memory_available():
host = psutil.virtual_memory()
available = host.available
for limit, directory in _limited(host.total):
used = _working_set_in(directory)
available = min(available, limit if used is None else limit - used)
return max(0, available)
comfy/taesd/taehv.py
+26 -1
View File
@@ -131,10 +131,11 @@ class TAEHV(nn.Module):
self.latent_channels = latent_channels
self.parallel = parallel
self.latent_format = latent_format
self.is_h3 = self.latent_channels == 24
self.show_progress_bar = show_progress_bar
self.process_in = latent_format().process_in if latent_format is not None else (lambda x: x)
self.process_out = latent_format().process_out if latent_format is not None else (lambda x: x)
if self.latent_channels in [48, 32]: # Wan 2.2 and HunyuanVideo1.5
if self.latent_channels in [48, 32, 24]: # Wan 2.2, HunyuanVideo1.5 and MiniMax H3
self.patch_size = 2
elif self.latent_channels == 128: # LTX2
self.patch_size, self.latent_channels, encoder_time_downscale, decoder_time_upscale = 4, 128, (True, True, True), (True, True, True)
@@ -176,6 +177,21 @@ class TAEHV(nn.Module):
def encode(self, x, **kwargs):
x = x.movedim(2, 1) # [B, C, T, H, W] -> [B, T, C, H, W]
if self.is_h3:
single_frame = x.shape[1] == 1
batch = x.shape[0]
x = torch.cat([x, x[:, -1:].expand(-1, -x.shape[1] % 17, -1, -1, -1)], dim=1)
x = F.pad(x.reshape(batch, -1, 17, *x.shape[2:]), (0, 0, 0, 0, 0, 0, 3, 0))
if self.parallel:
x = apply_model_with_memblocks(self.encoder, x.flatten(0, 1), True, self.show_progress_bar,
patch_size=self.patch_size)
x = x.reshape(batch, -1, *x.shape[2:])
else:
x = torch.cat([apply_model_with_memblocks(self.encoder, chunk, False, False,
patch_size=self.patch_size)
for chunk in tqdm(x.unbind(1), disable=not self.show_progress_bar)], dim=1)
x = x[:, :1] if single_frame else x[:, :-3]
return self.process_out(x.movedim(2, 1))
if x.shape[1] % self.t_downscale != 0:
# pad at end to multiple of t_downscale
n_pad = self.t_downscale - x.shape[1] % self.t_downscale
@@ -189,7 +205,16 @@ class TAEHV(nn.Module):
x = x.unsqueeze(0) if x.ndim == 4 else x # [T, C, H, W] -> [1, T, C, H, W]
x = x.movedim(1, 2) if x.shape[1] != self.latent_channels else x # [B, T, C, H, W] or [B, C, T, H, W]
x = self.process_in(x).movedim(2, 1) # [B, C, T, H, W] -> [B, T, C, H, W]
if self.is_h3:
single_frame = x.shape[1] == 1
x = apply_model_with_memblocks(self.decoder, x, self.parallel, self.show_progress_bar,
output_device=comfy.model_management.intermediate_device(),
patch_size=self.patch_size, decode=True)
if self.is_h3:
x.clamp_(0, 1)
if not single_frame:
chunk_frames = 5 * self.t_upscale
x = F.pad(x, (0, 0, 0, 0, 0, 0, 0, -x.shape[1] % chunk_frames))
x = x.unflatten(1, (-1, chunk_frames))[:, :, self.frames_to_trim:].flatten(1, 2)
return x[:, :-3 * self.t_upscale].movedim(2, 1)
return x[:, self.frames_to_trim:].movedim(2, 1)

Some files were not shown because too many files have changed in this diff Show More