Doing some experimentation

.github/workflows/test-assets.yml

@@ -1,173 +0,0 @@
-name: Asset System Tests
-
-on:
-  push:
-    paths:
-      - 'app/**'
-      - 'tests-assets/**'
-      - '.github/workflows/test-assets.yml'
-      - 'requirements.txt'
-  pull_request:
-    branches: [master]
-  workflow_dispatch:
-
-permissions:
-  contents: read
-
-env:
-  PIP_DISABLE_PIP_VERSION_CHECK: '1'
-  PYTHONUNBUFFERED: '1'
-
-jobs:
-  sqlite:
-    name: SQLite (${{ matrix.sqlite_mode }}) • Python ${{ matrix.python }}
-    runs-on: ubuntu-latest
-    timeout-minutes: 40
-    strategy:
-      fail-fast: false
-      matrix:
-        python: ['3.9', '3.12']
-        sqlite_mode: ['memory', 'file']
-
-    steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python }}
-
-      - name: Install dependencies
-        run: |
-          python -m pip install -U pip wheel
-          pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-          pip install -r requirements.txt
-          pip install pytest pytest-aiohttp pytest-asyncio
-
-      - name: Set deterministic test base dir
-        id: basedir
-        shell: bash
-        run: |
-          BASE="$RUNNER_TEMP/comfyui-assets-tests-${{ matrix.python }}-${{ matrix.sqlite_mode }}-${{ github.run_id }}-${{ github.run_attempt }}"
-          echo "ASSETS_TEST_BASE_DIR=$BASE" >> "$GITHUB_ENV"
-          echo "ASSETS_TEST_LOGS=$BASE/logs" >> "$GITHUB_ENV"
-          mkdir -p "$BASE/logs"
-          echo "ASSETS_TEST_BASE_DIR=$BASE"
-
-      - name: Set DB URL for SQLite
-        id: setdb
-        shell: bash
-        run: |
-          if [ "${{ matrix.sqlite_mode }}" = "memory" ]; then
-            echo "ASSETS_TEST_DB_URL=sqlite+aiosqlite:///:memory:" >> "$GITHUB_ENV"
-          else
-            DBFILE="$RUNNER_TEMP/assets-tests.sqlite"
-            mkdir -p "$(dirname "$DBFILE")"
-            echo "ASSETS_TEST_DB_URL=sqlite+aiosqlite:///$DBFILE" >> "$GITHUB_ENV"
-          fi
-
-      - name: Run tests
-        run: python -m pytest tests-assets
-
-      - name: Show ComfyUI logs
-        if: always()
-        shell: bash
-        run: |
-          echo "==== ASSETS_TEST_BASE_DIR: $ASSETS_TEST_BASE_DIR ===="
-          echo "==== ASSETS_TEST_LOGS: $ASSETS_TEST_LOGS ===="
-          ls -la "$ASSETS_TEST_LOGS" || true
-          for f in "$ASSETS_TEST_LOGS"/stdout.log "$ASSETS_TEST_LOGS"/stderr.log; do
-            if [ -f "$f" ]; then
-              echo "----- BEGIN $f -----"
-              sed -n '1,400p' "$f"
-              echo "----- END $f -----"
-            fi
-          done
-
-      - name: Upload ComfyUI logs
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: asset-logs-sqlite-${{ matrix.sqlite_mode }}-py${{ matrix.python }}
-          path: ${{ env.ASSETS_TEST_LOGS }}/*.log
-          if-no-files-found: warn
-
-  postgres:
-    name: PostgreSQL ${{ matrix.pgsql }} • Python ${{ matrix.python }}
-    runs-on: ubuntu-latest
-    timeout-minutes: 40
-    strategy:
-      fail-fast: false
-      matrix:
-        python: ['3.9', '3.12']
-        pgsql: ['16', '18']
-
-    services:
-      postgres:
-        image: postgres:${{ matrix.pgsql }}
-        env:
-          POSTGRES_DB: assets
-          POSTGRES_USER: postgres
-          POSTGRES_PASSWORD: postgres
-        ports:
-          - 5432:5432
-        options: >-
-          --health-cmd "pg_isready -U postgres -d assets"
-          --health-interval 10s
-          --health-timeout 5s
-          --health-retries 12
-
-    steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python
-        uses: actions/setup-python@v5
-        with:
-          python-version: ${{ matrix.python }}
-
-      - name: Install dependencies
-        run: |
-          python -m pip install -U pip wheel
-          pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-          pip install -r requirements.txt
-          pip install pytest pytest-aiohttp pytest-asyncio
-          pip install greenlet psycopg
-
-      - name: Set deterministic test base dir
-        id: basedir
-        shell: bash
-        run: |
-          BASE="$RUNNER_TEMP/comfyui-assets-tests-${{ matrix.python }}-${{ matrix.sqlite_mode }}-${{ github.run_id }}-${{ github.run_attempt }}"
-          echo "ASSETS_TEST_BASE_DIR=$BASE" >> "$GITHUB_ENV"
-          echo "ASSETS_TEST_LOGS=$BASE/logs" >> "$GITHUB_ENV"
-          mkdir -p "$BASE/logs"
-          echo "ASSETS_TEST_BASE_DIR=$BASE"
-
-      - name: Set DB URL for PostgreSQL
-        shell: bash
-        run: |
-          echo "ASSETS_TEST_DB_URL=postgresql+psycopg://postgres:postgres@localhost:5432/assets" >> "$GITHUB_ENV"
-
-      - name: Run tests
-        run: python -m pytest tests-assets
-
-      - name: Show ComfyUI logs
-        if: always()
-        shell: bash
-        run: |
-          echo "==== ASSETS_TEST_BASE_DIR: $ASSETS_TEST_BASE_DIR ===="
-          echo "==== ASSETS_TEST_LOGS: $ASSETS_TEST_LOGS ===="
-          ls -la "$ASSETS_TEST_LOGS" || true
-          for f in "$ASSETS_TEST_LOGS"/stdout.log "$ASSETS_TEST_LOGS"/stderr.log; do
-            if [ -f "$f" ]; then
-              echo "----- BEGIN $f -----"
-              sed -n '1,400p' "$f"
-              echo "----- END $f -----"
-            fi
-          done
-
-      - name: Upload ComfyUI logs
-        if: always()
-        uses: actions/upload-artifact@v4
-        with:
-          name: asset-logs-pgsql-${{ matrix.pgsql }}-py${{ matrix.python }}
-          path: ${{ env.ASSETS_TEST_LOGS }}/*.log
-          if-no-files-found: warn

.github/workflows/test-execution.yml

@@ -1,30 +0,0 @@
-name: Execution Tests
-
-on:
-  push:
-    branches: [ main, master ]
-  pull_request:
-    branches: [ main, master ]
-
-jobs:
-  test:
-    strategy:
-      matrix:
-        os: [ubuntu-latest, windows-latest, macos-latest]
-    runs-on: ${{ matrix.os }}
-    continue-on-error: true
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python      
-      uses: actions/setup-python@v4
-      with:
-        python-version: '3.12'
-    - name: Install requirements
-      run: |
-        python -m pip install --upgrade pip
-        pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
-        pip install -r requirements.txt
-        pip install -r tests-unit/requirements.txt
-    - name: Run Execution Tests
-      run: |
-        python -m pytest tests/execution -v --skip-timing-checks

.github/workflows/test-unit.yml

@@ -10,7 +10,7 @@ jobs:
   test:
     strategy:
       matrix:
-        os: [ubuntu-latest, windows-2022, macos-latest]
+        os: [ubuntu-latest, windows-latest, macos-latest]
     runs-on: ${{ matrix.os }}
     continue-on-error: true
     steps:

CODEOWNERS

@@ -1,3 +1,25 @@
 # Admins
 * @comfyanonymous
-* @kosinkadink
+
+# Note: Github teams syntax cannot be used here as the repo is not owned by Comfy-Org.
+# Inlined the team members for now.
+
+# Maintainers
+*.md @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/tests/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/tests-unit/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/notebooks/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/script_examples/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/.github/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/requirements.txt @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+/pyproject.toml @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @Kosinkadink @christian-byrne @guill
+
+# Python web server
+/api_server/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne @guill
+/app/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne @guill
+/utils/ @yoland68 @robinjhuang @webfiltered @pythongosssss @ltdrdata @christian-byrne @guill
+
+# Node developers
+/comfy_extras/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne @guill
+/comfy/comfy_types/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne @guill
+/comfy_api_nodes/ @yoland68 @robinjhuang @pythongosssss @ltdrdata @Kosinkadink @webfiltered @christian-byrne @guill

README.md

@@ -66,7 +66,6 @@ See what ComfyUI can do with the [example workflows](https://comfyanonymous.gith
    - [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/)
 - 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)

alembic.ini

@@ -3,7 +3,7 @@
 [alembic]
 # path to migration scripts
 # Use forward slashes (/) also on windows to provide an os agnostic path
-script_location = app/alembic_db
+script_location = alembic_db
 
 # template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
 # Uncomment the line below if you want the files to be prepended with date and time

alembic_db/env.py

@@ -2,12 +2,13 @@ from sqlalchemy import engine_from_config
 from sqlalchemy import pool
 
 from alembic import context
-from app.assets.database.models import Base
 
 # this is the Alembic Config object, which provides
 # access to the values within the .ini file in use.
 config = context.config
 
+
+from app.database.models import Base
 target_metadata = Base.metadata
 
 # other values from the config, defined by the needs of env.py,

app/alembic_db/versions/0001_assets.py

@@ -1,175 +0,0 @@
-"""initial assets schema
-
-Revision ID: 0001_assets
-Revises:
-Create Date: 2025-08-20 00:00:00
-"""
-
-from alembic import op
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql
-
-revision = "0001_assets"
-down_revision = None
-branch_labels = None
-depends_on = None
-
-
-def upgrade() -> None:
-    # ASSETS: content identity
-    op.create_table(
-        "assets",
-        sa.Column("id", sa.String(length=36), primary_key=True),
-        sa.Column("hash", sa.String(length=256), nullable=True),
-        sa.Column("size_bytes", sa.BigInteger(), nullable=False, server_default="0"),
-        sa.Column("mime_type", sa.String(length=255), nullable=True),
-        sa.Column("created_at", sa.DateTime(timezone=False), nullable=False),
-        sa.CheckConstraint("size_bytes >= 0", name="ck_assets_size_nonneg"),
-    )
-    op.create_index("uq_assets_hash", "assets", ["hash"], unique=True)
-    op.create_index("ix_assets_mime_type", "assets", ["mime_type"])
-
-    # ASSETS_INFO: user-visible references
-    op.create_table(
-        "assets_info",
-        sa.Column("id", sa.String(length=36), primary_key=True),
-        sa.Column("owner_id", sa.String(length=128), nullable=False, server_default=""),
-        sa.Column("name", sa.String(length=512), nullable=False),
-        sa.Column("asset_id", sa.String(length=36), sa.ForeignKey("assets.id", ondelete="RESTRICT"), nullable=False),
-        sa.Column("preview_id", sa.String(length=36), sa.ForeignKey("assets.id", ondelete="SET NULL"), nullable=True),
-        sa.Column("user_metadata", sa.JSON(), nullable=True),
-        sa.Column("created_at", sa.DateTime(timezone=False), nullable=False),
-        sa.Column("updated_at", sa.DateTime(timezone=False), nullable=False),
-        sa.Column("last_access_time", sa.DateTime(timezone=False), nullable=False),
-        sa.UniqueConstraint("asset_id", "owner_id", "name", name="uq_assets_info_asset_owner_name"),
-    )
-    op.create_index("ix_assets_info_owner_id", "assets_info", ["owner_id"])
-    op.create_index("ix_assets_info_asset_id", "assets_info", ["asset_id"])
-    op.create_index("ix_assets_info_name", "assets_info", ["name"])
-    op.create_index("ix_assets_info_created_at", "assets_info", ["created_at"])
-    op.create_index("ix_assets_info_last_access_time", "assets_info", ["last_access_time"])
-    op.create_index("ix_assets_info_owner_name", "assets_info", ["owner_id", "name"])
-
-    # TAGS: normalized tag vocabulary
-    op.create_table(
-        "tags",
-        sa.Column("name", sa.String(length=512), primary_key=True),
-        sa.Column("tag_type", sa.String(length=32), nullable=False, server_default="user"),
-        sa.CheckConstraint("name = lower(name)", name="ck_tags_lowercase"),
-    )
-    op.create_index("ix_tags_tag_type", "tags", ["tag_type"])
-
-    # ASSET_INFO_TAGS: many-to-many for tags on AssetInfo
-    op.create_table(
-        "asset_info_tags",
-        sa.Column("asset_info_id", sa.String(length=36), sa.ForeignKey("assets_info.id", ondelete="CASCADE"), nullable=False),
-        sa.Column("tag_name", sa.String(length=512), sa.ForeignKey("tags.name", ondelete="RESTRICT"), nullable=False),
-        sa.Column("origin", sa.String(length=32), nullable=False, server_default="manual"),
-        sa.Column("added_at", sa.DateTime(timezone=False), nullable=False),
-        sa.PrimaryKeyConstraint("asset_info_id", "tag_name", name="pk_asset_info_tags"),
-    )
-    op.create_index("ix_asset_info_tags_tag_name", "asset_info_tags", ["tag_name"])
-    op.create_index("ix_asset_info_tags_asset_info_id", "asset_info_tags", ["asset_info_id"])
-
-    # ASSET_CACHE_STATE: N:1 local cache rows per Asset
-    op.create_table(
-        "asset_cache_state",
-        sa.Column("id", sa.Integer(), primary_key=True, autoincrement=True),
-        sa.Column("asset_id", sa.String(length=36), sa.ForeignKey("assets.id", ondelete="CASCADE"), nullable=False),
-        sa.Column("file_path", sa.Text(), nullable=False),  # absolute local path to cached file
-        sa.Column("mtime_ns", sa.BigInteger(), nullable=True),
-        sa.Column("needs_verify", sa.Boolean(), nullable=False, server_default=sa.text("false")),
-        sa.CheckConstraint("(mtime_ns IS NULL) OR (mtime_ns >= 0)", name="ck_acs_mtime_nonneg"),
-        sa.UniqueConstraint("file_path", name="uq_asset_cache_state_file_path"),
-    )
-    op.create_index("ix_asset_cache_state_file_path", "asset_cache_state", ["file_path"])
-    op.create_index("ix_asset_cache_state_asset_id", "asset_cache_state", ["asset_id"])
-
-    # ASSET_INFO_META: typed KV projection of user_metadata for filtering/sorting
-    op.create_table(
-        "asset_info_meta",
-        sa.Column("asset_info_id", sa.String(length=36), sa.ForeignKey("assets_info.id", ondelete="CASCADE"), nullable=False),
-        sa.Column("key", sa.String(length=256), nullable=False),
-        sa.Column("ordinal", sa.Integer(), nullable=False, server_default="0"),
-        sa.Column("val_str", sa.String(length=2048), nullable=True),
-        sa.Column("val_num", sa.Numeric(38, 10), nullable=True),
-        sa.Column("val_bool", sa.Boolean(), nullable=True),
-        sa.Column("val_json", sa.JSON().with_variant(postgresql.JSONB(), 'postgresql'), nullable=True),
-        sa.PrimaryKeyConstraint("asset_info_id", "key", "ordinal", name="pk_asset_info_meta"),
-    )
-    op.create_index("ix_asset_info_meta_key", "asset_info_meta", ["key"])
-    op.create_index("ix_asset_info_meta_key_val_str", "asset_info_meta", ["key", "val_str"])
-    op.create_index("ix_asset_info_meta_key_val_num", "asset_info_meta", ["key", "val_num"])
-    op.create_index("ix_asset_info_meta_key_val_bool", "asset_info_meta", ["key", "val_bool"])
-
-    # Tags vocabulary
-    tags_table = sa.table(
-        "tags",
-        sa.column("name", sa.String(length=512)),
-        sa.column("tag_type", sa.String()),
-    )
-    op.bulk_insert(
-        tags_table,
-        [
-            {"name": "models", "tag_type": "system"},
-            {"name": "input", "tag_type": "system"},
-            {"name": "output", "tag_type": "system"},
-
-            {"name": "configs", "tag_type": "system"},
-            {"name": "checkpoints", "tag_type": "system"},
-            {"name": "loras", "tag_type": "system"},
-            {"name": "vae", "tag_type": "system"},
-            {"name": "text_encoders", "tag_type": "system"},
-            {"name": "diffusion_models", "tag_type": "system"},
-            {"name": "clip_vision", "tag_type": "system"},
-            {"name": "style_models", "tag_type": "system"},
-            {"name": "embeddings", "tag_type": "system"},
-            {"name": "diffusers", "tag_type": "system"},
-            {"name": "vae_approx", "tag_type": "system"},
-            {"name": "controlnet", "tag_type": "system"},
-            {"name": "gligen", "tag_type": "system"},
-            {"name": "upscale_models", "tag_type": "system"},
-            {"name": "hypernetworks", "tag_type": "system"},
-            {"name": "photomaker", "tag_type": "system"},
-            {"name": "classifiers", "tag_type": "system"},
-
-            {"name": "encoder", "tag_type": "system"},
-            {"name": "decoder", "tag_type": "system"},
-
-            {"name": "missing", "tag_type": "system"},
-            {"name": "rescan", "tag_type": "system"},
-        ],
-    )
-
-
-def downgrade() -> None:
-    op.drop_index("ix_asset_info_meta_key_val_bool", table_name="asset_info_meta")
-    op.drop_index("ix_asset_info_meta_key_val_num", table_name="asset_info_meta")
-    op.drop_index("ix_asset_info_meta_key_val_str", table_name="asset_info_meta")
-    op.drop_index("ix_asset_info_meta_key", table_name="asset_info_meta")
-    op.drop_table("asset_info_meta")
-
-    op.drop_index("ix_asset_cache_state_asset_id", table_name="asset_cache_state")
-    op.drop_index("ix_asset_cache_state_file_path", table_name="asset_cache_state")
-    op.drop_constraint("uq_asset_cache_state_file_path", table_name="asset_cache_state")
-    op.drop_table("asset_cache_state")
-
-    op.drop_index("ix_asset_info_tags_asset_info_id", table_name="asset_info_tags")
-    op.drop_index("ix_asset_info_tags_tag_name", table_name="asset_info_tags")
-    op.drop_table("asset_info_tags")
-
-    op.drop_index("ix_tags_tag_type", table_name="tags")
-    op.drop_table("tags")
-
-    op.drop_constraint("uq_assets_info_asset_owner_name", table_name="assets_info")
-    op.drop_index("ix_assets_info_owner_name", table_name="assets_info")
-    op.drop_index("ix_assets_info_last_access_time", table_name="assets_info")
-    op.drop_index("ix_assets_info_created_at", table_name="assets_info")
-    op.drop_index("ix_assets_info_name", table_name="assets_info")
-    op.drop_index("ix_assets_info_asset_id", table_name="assets_info")
-    op.drop_index("ix_assets_info_owner_id", table_name="assets_info")
-    op.drop_table("assets_info")
-
-    op.drop_index("uq_assets_hash", table_name="assets")
-    op.drop_index("ix_assets_mime_type", table_name="assets")
-    op.drop_table("assets")

app/assets/__init__.py

@@ -1,4 +0,0 @@
-from .api.routes import register_assets_system
-from .scanner import sync_seed_assets
-
-__all__ = ["sync_seed_assets", "register_assets_system"]

app/assets/_helpers.py

@@ -1,225 +0,0 @@
-import contextlib
-import os
-import uuid
-from datetime import datetime, timezone
-from pathlib import Path
-from typing import Literal, Optional, Sequence
-
-import folder_paths
-
-from .api import schemas_in
-
-
-def get_comfy_models_folders() -> list[tuple[str, list[str]]]:
-    """Build a list of (folder_name, base_paths[]) categories that are configured for model locations.
-
-    We trust `folder_paths.folder_names_and_paths` and include a category if
-    *any* of its base paths lies under the Comfy `models_dir`.
-    """
-    targets: list[tuple[str, list[str]]] = []
-    models_root = os.path.abspath(folder_paths.models_dir)
-    for name, (paths, _exts) in folder_paths.folder_names_and_paths.items():
-        if any(os.path.abspath(p).startswith(models_root + os.sep) for p in paths):
-            targets.append((name, paths))
-    return targets
-
-
-def get_relative_to_root_category_path_of_asset(file_path: str) -> tuple[Literal["input", "output", "models"], str]:
-    """Given an absolute or relative file path, determine which root category the path belongs to:
-      - 'input' if the file resides under `folder_paths.get_input_directory()`
-      - 'output' if the file resides under `folder_paths.get_output_directory()`
-      - 'models' if the file resides under any base path of categories returned by `get_comfy_models_folders()`
-
-    Returns:
-        (root_category, relative_path_inside_that_root)
-        For 'models', the relative path is prefixed with the category name:
-            e.g. ('models', 'vae/test/sub/ae.safetensors')
-
-    Raises:
-        ValueError: if the path does not belong to input, output, or configured model bases.
-    """
-    fp_abs = os.path.abspath(file_path)
-
-    def _is_within(child: str, parent: str) -> bool:
-        try:
-            return os.path.commonpath([child, parent]) == parent
-        except Exception:
-            return False
-
-    def _rel(child: str, parent: str) -> str:
-        return os.path.relpath(os.path.join(os.sep, os.path.relpath(child, parent)), os.sep)
-
-    # 1) input
-    input_base = os.path.abspath(folder_paths.get_input_directory())
-    if _is_within(fp_abs, input_base):
-        return "input", _rel(fp_abs, input_base)
-
-    # 2) output
-    output_base = os.path.abspath(folder_paths.get_output_directory())
-    if _is_within(fp_abs, output_base):
-        return "output", _rel(fp_abs, output_base)
-
-    # 3) models (check deepest matching base to avoid ambiguity)
-    best: Optional[tuple[int, str, str]] = None  # (base_len, bucket, rel_inside_bucket)
-    for bucket, bases in get_comfy_models_folders():
-        for b in bases:
-            base_abs = os.path.abspath(b)
-            if not _is_within(fp_abs, base_abs):
-                continue
-            cand = (len(base_abs), bucket, _rel(fp_abs, base_abs))
-            if best is None or cand[0] > best[0]:
-                best = cand
-
-    if best is not None:
-        _, bucket, rel_inside = best
-        combined = os.path.join(bucket, rel_inside)
-        return "models", os.path.relpath(os.path.join(os.sep, combined), os.sep)
-
-    raise ValueError(f"Path is not within input, output, or configured model bases: {file_path}")
-
-
-def get_name_and_tags_from_asset_path(file_path: str) -> tuple[str, list[str]]:
-    """Return a tuple (name, tags) derived from a filesystem path.
-
-    Semantics:
-      - Root category is determined by `get_relative_to_root_category_path_of_asset`.
-      - The returned `name` is the base filename with extension from the relative path.
-      - The returned `tags` are:
-            [root_category] + parent folders of the relative path (in order)
-        For 'models', this means:
-            file '/.../ModelsDir/vae/test_tag/ae.safetensors'
-            -> root_category='models', some_path='vae/test_tag/ae.safetensors'
-            -> name='ae.safetensors', tags=['models', 'vae', 'test_tag']
-
-    Raises:
-        ValueError: if the path does not belong to input, output, or configured model bases.
-    """
-    root_category, some_path = get_relative_to_root_category_path_of_asset(file_path)
-    p = Path(some_path)
-    parent_parts = [part for part in p.parent.parts if part not in (".", "..", p.anchor)]
-    return p.name, list(dict.fromkeys(normalize_tags([root_category, *parent_parts])))
-
-
-def normalize_tags(tags: Optional[Sequence[str]]) -> list[str]:
-    return [t.strip().lower() for t in (tags or []) if (t or "").strip()]
-
-
-def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
-    """Validates and maps tags -> (base_dir, subdirs_for_fs)"""
-    root = tags[0]
-    if root == "models":
-        if len(tags) < 2:
-            raise ValueError("at least two tags required for model asset")
-        try:
-            bases = folder_paths.folder_names_and_paths[tags[1]][0]
-        except KeyError:
-            raise ValueError(f"unknown model category '{tags[1]}'")
-        if not bases:
-            raise ValueError(f"no base path configured for category '{tags[1]}'")
-        base_dir = os.path.abspath(bases[0])
-        raw_subdirs = tags[2:]
-    else:
-        base_dir = os.path.abspath(
-            folder_paths.get_input_directory() if root == "input" else folder_paths.get_output_directory()
-        )
-        raw_subdirs = tags[1:]
-    for i in raw_subdirs:
-        if i in (".", ".."):
-            raise ValueError("invalid path component in tags")
-
-    return base_dir, raw_subdirs if raw_subdirs else []
-
-
-def ensure_within_base(candidate: str, base: str) -> None:
-    cand_abs = os.path.abspath(candidate)
-    base_abs = os.path.abspath(base)
-    try:
-        if os.path.commonpath([cand_abs, base_abs]) != base_abs:
-            raise ValueError("destination escapes base directory")
-    except Exception:
-        raise ValueError("invalid destination path")
-
-
-def compute_relative_filename(file_path: str) -> Optional[str]:
-    """
-    Return the model's path relative to the last well-known folder (the model category),
-    using forward slashes, eg:
-      /.../models/checkpoints/flux/123/flux.safetensors -> "flux/123/flux.safetensors"
-      /.../models/text_encoders/clip_g.safetensors -> "clip_g.safetensors"
-
-    For non-model paths, returns None.
-    NOTE: this is a temporary helper, used only for initializing metadata["filename"] field.
-    """
-    try:
-        root_category, rel_path = get_relative_to_root_category_path_of_asset(file_path)
-    except ValueError:
-        return None
-
-    p = Path(rel_path)
-    parts = [seg for seg in p.parts if seg not in (".", "..", p.anchor)]
-    if not parts:
-        return None
-
-    if root_category == "models":
-        # parts[0] is the category ("checkpoints", "vae", etc) – drop it
-        inside = parts[1:] if len(parts) > 1 else [parts[0]]
-        return "/".join(inside)
-    return "/".join(parts)  # input/output: keep all parts
-
-
-def list_tree(base_dir: str) -> list[str]:
-    out: list[str] = []
-    base_abs = os.path.abspath(base_dir)
-    if not os.path.isdir(base_abs):
-        return out
-    for dirpath, _subdirs, filenames in os.walk(base_abs, topdown=True, followlinks=False):
-        for name in filenames:
-            out.append(os.path.abspath(os.path.join(dirpath, name)))
-    return out
-
-
-def prefixes_for_root(root: schemas_in.RootType) -> list[str]:
-    if root == "models":
-        bases: list[str] = []
-        for _bucket, paths in get_comfy_models_folders():
-            bases.extend(paths)
-        return [os.path.abspath(p) for p in bases]
-    if root == "input":
-        return [os.path.abspath(folder_paths.get_input_directory())]
-    if root == "output":
-        return [os.path.abspath(folder_paths.get_output_directory())]
-    return []
-
-
-def ts_to_iso(ts: Optional[float]) -> Optional[str]:
-    if ts is None:
-        return None
-    try:
-        return datetime.fromtimestamp(float(ts), tz=timezone.utc).replace(tzinfo=None).isoformat()
-    except Exception:
-        return None
-
-
-def new_scan_id(root: schemas_in.RootType) -> str:
-    return f"scan-{root}-{uuid.uuid4().hex[:8]}"
-
-
-def collect_models_files() -> list[str]:
-    out: list[str] = []
-    for folder_name, bases in get_comfy_models_folders():
-        rel_files = folder_paths.get_filename_list(folder_name) or []
-        for rel_path in rel_files:
-            abs_path = folder_paths.get_full_path(folder_name, rel_path)
-            if not abs_path:
-                continue
-            abs_path = os.path.abspath(abs_path)
-            allowed = False
-            for b in bases:
-                base_abs = os.path.abspath(b)
-                with contextlib.suppress(Exception):
-                    if os.path.commonpath([abs_path, base_abs]) == base_abs:
-                        allowed = True
-                        break
-            if allowed:
-                out.append(abs_path)
-    return out

app/assets/api/routes.py

@@ -1,544 +0,0 @@
-import contextlib
-import logging
-import os
-import urllib.parse
-import uuid
-from typing import Optional
-
-from aiohttp import web
-from pydantic import ValidationError
-
-import folder_paths
-
-from ... import user_manager
-from .. import manager, scanner
-from . import schemas_in, schemas_out
-
-ROUTES = web.RouteTableDef()
-USER_MANAGER: Optional[user_manager.UserManager] = None
-LOGGER = logging.getLogger(__name__)
-
-# UUID regex (canonical hyphenated form, case-insensitive)
-UUID_RE = r"[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}"
-
-
-@ROUTES.head("/api/assets/hash/{hash}")
-async def head_asset_by_hash(request: web.Request) -> web.Response:
-    hash_str = request.match_info.get("hash", "").strip().lower()
-    if not hash_str or ":" not in hash_str:
-        return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:<hex>'")
-    algo, digest = hash_str.split(":", 1)
-    if algo != "blake3" or not digest or any(c for c in digest if c not in "0123456789abcdef"):
-        return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:<hex>'")
-    exists = await manager.asset_exists(asset_hash=hash_str)
-    return web.Response(status=200 if exists else 404)
-
-
-@ROUTES.get("/api/assets")
-async def list_assets(request: web.Request) -> web.Response:
-    qp = request.rel_url.query
-    query_dict = {}
-    if "include_tags" in qp:
-        query_dict["include_tags"] = qp.getall("include_tags")
-    if "exclude_tags" in qp:
-        query_dict["exclude_tags"] = qp.getall("exclude_tags")
-    for k in ("name_contains", "metadata_filter", "limit", "offset", "sort", "order"):
-        v = qp.get(k)
-        if v is not None:
-            query_dict[k] = v
-
-    try:
-        q = schemas_in.ListAssetsQuery.model_validate(query_dict)
-    except ValidationError as ve:
-        return _validation_error_response("INVALID_QUERY", ve)
-
-    payload = await manager.list_assets(
-        include_tags=q.include_tags,
-        exclude_tags=q.exclude_tags,
-        name_contains=q.name_contains,
-        metadata_filter=q.metadata_filter,
-        limit=q.limit,
-        offset=q.offset,
-        sort=q.sort,
-        order=q.order,
-        owner_id=USER_MANAGER.get_request_user_id(request),
-    )
-    return web.json_response(payload.model_dump(mode="json"))
-
-
-@ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}/content")
-async def download_asset_content(request: web.Request) -> web.Response:
-    disposition = request.query.get("disposition", "attachment").lower().strip()
-    if disposition not in {"inline", "attachment"}:
-        disposition = "attachment"
-
-    try:
-        abs_path, content_type, filename = await manager.resolve_asset_content_for_download(
-            asset_info_id=str(uuid.UUID(request.match_info["id"])),
-            owner_id=USER_MANAGER.get_request_user_id(request),
-        )
-    except ValueError as ve:
-        return _error_response(404, "ASSET_NOT_FOUND", str(ve))
-    except NotImplementedError as nie:
-        return _error_response(501, "BACKEND_UNSUPPORTED", str(nie))
-    except FileNotFoundError:
-        return _error_response(404, "FILE_NOT_FOUND", "Underlying file not found on disk.")
-
-    quoted = (filename or "").replace("\r", "").replace("\n", "").replace('"', "'")
-    cd = f'{disposition}; filename="{quoted}"; filename*=UTF-8\'\'{urllib.parse.quote(filename)}'
-
-    resp = web.FileResponse(abs_path)
-    resp.content_type = content_type
-    resp.headers["Content-Disposition"] = cd
-    return resp
-
-
-@ROUTES.post("/api/assets/from-hash")
-async def create_asset_from_hash(request: web.Request) -> web.Response:
-    try:
-        payload = await request.json()
-        body = schemas_in.CreateFromHashBody.model_validate(payload)
-    except ValidationError as ve:
-        return _validation_error_response("INVALID_BODY", ve)
-    except Exception:
-        return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.")
-
-    result = await manager.create_asset_from_hash(
-        hash_str=body.hash,
-        name=body.name,
-        tags=body.tags,
-        user_metadata=body.user_metadata,
-        owner_id=USER_MANAGER.get_request_user_id(request),
-    )
-    if result is None:
-        return _error_response(404, "ASSET_NOT_FOUND", f"Asset content {body.hash} does not exist")
-    return web.json_response(result.model_dump(mode="json"), status=201)
-
-
-@ROUTES.post("/api/assets")
-async def upload_asset(request: web.Request) -> web.Response:
-    """Multipart/form-data endpoint for Asset uploads."""
-
-    if not (request.content_type or "").lower().startswith("multipart/"):
-        return _error_response(415, "UNSUPPORTED_MEDIA_TYPE", "Use multipart/form-data for uploads.")
-
-    reader = await request.multipart()
-
-    file_present = False
-    file_client_name: Optional[str] = None
-    tags_raw: list[str] = []
-    provided_name: Optional[str] = None
-    user_metadata_raw: Optional[str] = None
-    provided_hash: Optional[str] = None
-    provided_hash_exists: Optional[bool] = None
-
-    file_written = 0
-    tmp_path: Optional[str] = None
-    while True:
-        field = await reader.next()
-        if field is None:
-            break
-
-        fname = getattr(field, "name", "") or ""
-
-        if fname == "hash":
-            try:
-                s = ((await field.text()) or "").strip().lower()
-            except Exception:
-                return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:<hex>'")
-
-            if s:
-                if ":" not in s:
-                    return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:<hex>'")
-                algo, digest = s.split(":", 1)
-                if algo != "blake3" or not digest or any(c for c in digest if c not in "0123456789abcdef"):
-                    return _error_response(400, "INVALID_HASH", "hash must be like 'blake3:<hex>'")
-                provided_hash = f"{algo}:{digest}"
-                try:
-                    provided_hash_exists = await manager.asset_exists(asset_hash=provided_hash)
-                except Exception:
-                    provided_hash_exists = None  # do not fail the whole request here
-
-        elif fname == "file":
-            file_present = True
-            file_client_name = (field.filename or "").strip()
-
-            if provided_hash and provided_hash_exists is True:
-                # If client supplied a hash that we know exists, drain but do not write to disk
-                try:
-                    while True:
-                        chunk = await field.read_chunk(8 * 1024 * 1024)
-                        if not chunk:
-                            break
-                        file_written += len(chunk)
-                except Exception:
-                    return _error_response(500, "UPLOAD_IO_ERROR", "Failed to receive uploaded file.")
-                continue  # Do not create temp file; we will create AssetInfo from the existing content
-
-            # Otherwise, store to temp for hashing/ingest
-            uploads_root = os.path.join(folder_paths.get_temp_directory(), "uploads")
-            unique_dir = os.path.join(uploads_root, uuid.uuid4().hex)
-            os.makedirs(unique_dir, exist_ok=True)
-            tmp_path = os.path.join(unique_dir, ".upload.part")
-
-            try:
-                with open(tmp_path, "wb") as f:
-                    while True:
-                        chunk = await field.read_chunk(8 * 1024 * 1024)
-                        if not chunk:
-                            break
-                        f.write(chunk)
-                        file_written += len(chunk)
-            except Exception:
-                try:
-                    if os.path.exists(tmp_path or ""):
-                        os.remove(tmp_path)
-                finally:
-                    return _error_response(500, "UPLOAD_IO_ERROR", "Failed to receive and store uploaded file.")
-        elif fname == "tags":
-            tags_raw.append((await field.text()) or "")
-        elif fname == "name":
-            provided_name = (await field.text()) or None
-        elif fname == "user_metadata":
-            user_metadata_raw = (await field.text()) or None
-
-    # If client did not send file, and we are not doing a from-hash fast path -> error
-    if not file_present and not (provided_hash and provided_hash_exists):
-        return _error_response(400, "MISSING_FILE", "Form must include a 'file' part or a known 'hash'.")
-
-    if file_present and file_written == 0 and not (provided_hash and provided_hash_exists):
-        # Empty upload is only acceptable if we are fast-pathing from existing hash
-        try:
-            if tmp_path and os.path.exists(tmp_path):
-                os.remove(tmp_path)
-        finally:
-            return _error_response(400, "EMPTY_UPLOAD", "Uploaded file is empty.")
-
-    try:
-        spec = schemas_in.UploadAssetSpec.model_validate({
-            "tags": tags_raw,
-            "name": provided_name,
-            "user_metadata": user_metadata_raw,
-            "hash": provided_hash,
-        })
-    except ValidationError as ve:
-        try:
-            if tmp_path and os.path.exists(tmp_path):
-                os.remove(tmp_path)
-        finally:
-            return _validation_error_response("INVALID_BODY", ve)
-
-    # Validate models category against configured folders (consistent with previous behavior)
-    if spec.tags and spec.tags[0] == "models":
-        if len(spec.tags) < 2 or spec.tags[1] not in folder_paths.folder_names_and_paths:
-            if tmp_path and os.path.exists(tmp_path):
-                os.remove(tmp_path)
-            return _error_response(
-                400, "INVALID_BODY", f"unknown models category '{spec.tags[1] if len(spec.tags) >= 2 else ''}'"
-            )
-
-    owner_id = USER_MANAGER.get_request_user_id(request)
-
-    # Fast path: if a valid provided hash exists, create AssetInfo without writing anything
-    if spec.hash and provided_hash_exists is True:
-        try:
-            result = await manager.create_asset_from_hash(
-                hash_str=spec.hash,
-                name=spec.name or (spec.hash.split(":", 1)[1]),
-                tags=spec.tags,
-                user_metadata=spec.user_metadata or {},
-                owner_id=owner_id,
-            )
-        except Exception:
-            LOGGER.exception("create_asset_from_hash failed for hash=%s, owner_id=%s", spec.hash, owner_id)
-            return _error_response(500, "INTERNAL", "Unexpected server error.")
-
-        if result is None:
-            return _error_response(404, "ASSET_NOT_FOUND", f"Asset content {spec.hash} does not exist")
-
-        # Drain temp if we accidentally saved (e.g., hash field came after file)
-        if tmp_path and os.path.exists(tmp_path):
-            with contextlib.suppress(Exception):
-                os.remove(tmp_path)
-
-        status = 200 if (not result.created_new) else 201
-        return web.json_response(result.model_dump(mode="json"), status=status)
-
-    # Otherwise, we must have a temp file path to ingest
-    if not tmp_path or not os.path.exists(tmp_path):
-        # The only case we reach here without a temp file is: client sent a hash that does not exist and no file
-        return _error_response(404, "ASSET_NOT_FOUND", "Provided hash not found and no file uploaded.")
-
-    try:
-        created = await manager.upload_asset_from_temp_path(
-            spec,
-            temp_path=tmp_path,
-            client_filename=file_client_name,
-            owner_id=owner_id,
-            expected_asset_hash=spec.hash,
-        )
-        status = 201 if created.created_new else 200
-        return web.json_response(created.model_dump(mode="json"), status=status)
-    except ValueError as e:
-        if tmp_path and os.path.exists(tmp_path):
-            os.remove(tmp_path)
-        msg = str(e)
-        if "HASH_MISMATCH" in msg or msg.strip().upper() == "HASH_MISMATCH":
-            return _error_response(
-                400,
-                "HASH_MISMATCH",
-                "Uploaded file hash does not match provided hash.",
-            )
-        return _error_response(400, "BAD_REQUEST", "Invalid inputs.")
-    except Exception:
-        if tmp_path and os.path.exists(tmp_path):
-            os.remove(tmp_path)
-        LOGGER.exception("upload_asset_from_temp_path failed for tmp_path=%s, owner_id=%s", tmp_path, owner_id)
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-
-
-@ROUTES.get(f"/api/assets/{{id:{UUID_RE}}}")
-async def get_asset(request: web.Request) -> web.Response:
-    asset_info_id = str(uuid.UUID(request.match_info["id"]))
-    try:
-        result = await manager.get_asset(
-            asset_info_id=asset_info_id,
-            owner_id=USER_MANAGER.get_request_user_id(request),
-        )
-    except ValueError as ve:
-        return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
-    except Exception:
-        LOGGER.exception(
-            "get_asset failed for asset_info_id=%s, owner_id=%s",
-            asset_info_id,
-            USER_MANAGER.get_request_user_id(request),
-        )
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-    return web.json_response(result.model_dump(mode="json"), status=200)
-
-
-@ROUTES.put(f"/api/assets/{{id:{UUID_RE}}}")
-async def update_asset(request: web.Request) -> web.Response:
-    asset_info_id = str(uuid.UUID(request.match_info["id"]))
-    try:
-        body = schemas_in.UpdateAssetBody.model_validate(await request.json())
-    except ValidationError as ve:
-        return _validation_error_response("INVALID_BODY", ve)
-    except Exception:
-        return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.")
-
-    try:
-        result = await manager.update_asset(
-            asset_info_id=asset_info_id,
-            name=body.name,
-            tags=body.tags,
-            user_metadata=body.user_metadata,
-            owner_id=USER_MANAGER.get_request_user_id(request),
-        )
-    except (ValueError, PermissionError) as ve:
-        return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
-    except Exception:
-        LOGGER.exception(
-            "update_asset failed for asset_info_id=%s, owner_id=%s",
-            asset_info_id,
-            USER_MANAGER.get_request_user_id(request),
-        )
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-    return web.json_response(result.model_dump(mode="json"), status=200)
-
-
-@ROUTES.put(f"/api/assets/{{id:{UUID_RE}}}/preview")
-async def set_asset_preview(request: web.Request) -> web.Response:
-    asset_info_id = str(uuid.UUID(request.match_info["id"]))
-    try:
-        body = schemas_in.SetPreviewBody.model_validate(await request.json())
-    except ValidationError as ve:
-        return _validation_error_response("INVALID_BODY", ve)
-    except Exception:
-        return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.")
-
-    try:
-        result = await manager.set_asset_preview(
-            asset_info_id=asset_info_id,
-            preview_asset_id=body.preview_id,
-            owner_id=USER_MANAGER.get_request_user_id(request),
-        )
-    except (PermissionError, ValueError) as ve:
-        return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
-    except Exception:
-        LOGGER.exception(
-            "set_asset_preview failed for asset_info_id=%s, owner_id=%s",
-            asset_info_id,
-            USER_MANAGER.get_request_user_id(request),
-        )
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-    return web.json_response(result.model_dump(mode="json"), status=200)
-
-
-@ROUTES.delete(f"/api/assets/{{id:{UUID_RE}}}")
-async def delete_asset(request: web.Request) -> web.Response:
-    asset_info_id = str(uuid.UUID(request.match_info["id"]))
-    delete_content = request.query.get("delete_content")
-    delete_content = True if delete_content is None else delete_content.lower() not in {"0", "false", "no"}
-
-    try:
-        deleted = await manager.delete_asset_reference(
-            asset_info_id=asset_info_id,
-            owner_id=USER_MANAGER.get_request_user_id(request),
-            delete_content_if_orphan=delete_content,
-        )
-    except Exception:
-        LOGGER.exception(
-            "delete_asset_reference failed for asset_info_id=%s, owner_id=%s",
-            asset_info_id,
-            USER_MANAGER.get_request_user_id(request),
-        )
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-
-    if not deleted:
-        return _error_response(404, "ASSET_NOT_FOUND", f"AssetInfo {asset_info_id} not found.")
-    return web.Response(status=204)
-
-
-@ROUTES.get("/api/tags")
-async def get_tags(request: web.Request) -> web.Response:
-    query_map = dict(request.rel_url.query)
-
-    try:
-        query = schemas_in.TagsListQuery.model_validate(query_map)
-    except ValidationError as ve:
-        return web.json_response(
-            {"error": {"code": "INVALID_QUERY", "message": "Invalid query parameters", "details": ve.errors()}},
-            status=400,
-        )
-
-    result = await manager.list_tags(
-        prefix=query.prefix,
-        limit=query.limit,
-        offset=query.offset,
-        order=query.order,
-        include_zero=query.include_zero,
-        owner_id=USER_MANAGER.get_request_user_id(request),
-    )
-    return web.json_response(result.model_dump(mode="json"))
-
-
-@ROUTES.post(f"/api/assets/{{id:{UUID_RE}}}/tags")
-async def add_asset_tags(request: web.Request) -> web.Response:
-    asset_info_id = str(uuid.UUID(request.match_info["id"]))
-    try:
-        payload = await request.json()
-        data = schemas_in.TagsAdd.model_validate(payload)
-    except ValidationError as ve:
-        return _error_response(400, "INVALID_BODY", "Invalid JSON body for tags add.", {"errors": ve.errors()})
-    except Exception:
-        return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.")
-
-    try:
-        result = await manager.add_tags_to_asset(
-            asset_info_id=asset_info_id,
-            tags=data.tags,
-            origin="manual",
-            owner_id=USER_MANAGER.get_request_user_id(request),
-        )
-    except (ValueError, PermissionError) as ve:
-        return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
-    except Exception:
-        LOGGER.exception(
-            "add_tags_to_asset failed for asset_info_id=%s, owner_id=%s",
-            asset_info_id,
-            USER_MANAGER.get_request_user_id(request),
-        )
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-
-    return web.json_response(result.model_dump(mode="json"), status=200)
-
-
-@ROUTES.delete(f"/api/assets/{{id:{UUID_RE}}}/tags")
-async def delete_asset_tags(request: web.Request) -> web.Response:
-    asset_info_id = str(uuid.UUID(request.match_info["id"]))
-    try:
-        payload = await request.json()
-        data = schemas_in.TagsRemove.model_validate(payload)
-    except ValidationError as ve:
-        return _error_response(400, "INVALID_BODY", "Invalid JSON body for tags remove.", {"errors": ve.errors()})
-    except Exception:
-        return _error_response(400, "INVALID_JSON", "Request body must be valid JSON.")
-
-    try:
-        result = await manager.remove_tags_from_asset(
-            asset_info_id=asset_info_id,
-            tags=data.tags,
-            owner_id=USER_MANAGER.get_request_user_id(request),
-        )
-    except ValueError as ve:
-        return _error_response(404, "ASSET_NOT_FOUND", str(ve), {"id": asset_info_id})
-    except Exception:
-        LOGGER.exception(
-            "remove_tags_from_asset failed for asset_info_id=%s, owner_id=%s",
-            asset_info_id,
-            USER_MANAGER.get_request_user_id(request),
-        )
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-
-    return web.json_response(result.model_dump(mode="json"), status=200)
-
-
-@ROUTES.post("/api/assets/scan/seed")
-async def seed_assets(request: web.Request) -> web.Response:
-    try:
-        payload = await request.json()
-    except Exception:
-        payload = {}
-
-    try:
-        body = schemas_in.ScheduleAssetScanBody.model_validate(payload)
-    except ValidationError as ve:
-        return _validation_error_response("INVALID_BODY", ve)
-
-    try:
-        await scanner.sync_seed_assets(body.roots)
-    except Exception:
-        LOGGER.exception("sync_seed_assets failed for roots=%s", body.roots)
-        return _error_response(500, "INTERNAL", "Unexpected server error.")
-    return web.json_response({"synced": True, "roots": body.roots}, status=200)
-
-
-@ROUTES.post("/api/assets/scan/schedule")
-async def schedule_asset_scan(request: web.Request) -> web.Response:
-    try:
-        payload = await request.json()
-    except Exception:
-        payload = {}
-
-    try:
-        body = schemas_in.ScheduleAssetScanBody.model_validate(payload)
-    except ValidationError as ve:
-        return _validation_error_response("INVALID_BODY", ve)
-
-    states = await scanner.schedule_scans(body.roots)
-    return web.json_response(states.model_dump(mode="json"), status=202)
-
-
-@ROUTES.get("/api/assets/scan")
-async def get_asset_scan_status(request: web.Request) -> web.Response:
-    root = request.query.get("root", "").strip().lower()
-    states = scanner.current_statuses()
-    if root in {"models", "input", "output"}:
-        states = [s for s in states.scans if s.root == root]  # type: ignore
-        states = schemas_out.AssetScanStatusResponse(scans=states)
-    return web.json_response(states.model_dump(mode="json"), status=200)
-
-
-def register_assets_system(app: web.Application, user_manager_instance: user_manager.UserManager) -> None:
-    global USER_MANAGER
-    USER_MANAGER = user_manager_instance
-    app.add_routes(ROUTES)
-
-
-def _error_response(status: int, code: str, message: str, details: Optional[dict] = None) -> web.Response:
-    return web.json_response({"error": {"code": code, "message": message, "details": details or {}}}, status=status)
-
-
-def _validation_error_response(code: str, ve: ValidationError) -> web.Response:
-    return _error_response(400, code, "Validation failed.", {"errors": ve.json()})

app/assets/api/schemas_in.py

@@ -1,297 +0,0 @@
-import json
-import uuid
-from typing import Any, Literal, Optional
-
-from pydantic import (
-    BaseModel,
-    ConfigDict,
-    Field,
-    conint,
-    field_validator,
-    model_validator,
-)
-
-
-class ListAssetsQuery(BaseModel):
-    include_tags: list[str] = Field(default_factory=list)
-    exclude_tags: list[str] = Field(default_factory=list)
-    name_contains: Optional[str] = None
-
-    # Accept either a JSON string (query param) or a dict
-    metadata_filter: Optional[dict[str, Any]] = None
-
-    limit: conint(ge=1, le=500) = 20
-    offset: conint(ge=0) = 0
-
-    sort: Literal["name", "created_at", "updated_at", "size", "last_access_time"] = "created_at"
-    order: Literal["asc", "desc"] = "desc"
-
-    @field_validator("include_tags", "exclude_tags", mode="before")
-    @classmethod
-    def _split_csv_tags(cls, v):
-        # Accept "a,b,c" or ["a","b"] (we are liberal in what we accept)
-        if v is None:
-            return []
-        if isinstance(v, str):
-            return [t.strip() for t in v.split(",") if t.strip()]
-        if isinstance(v, list):
-            out: list[str] = []
-            for item in v:
-                if isinstance(item, str):
-                    out.extend([t.strip() for t in item.split(",") if t.strip()])
-            return out
-        return v
-
-    @field_validator("metadata_filter", mode="before")
-    @classmethod
-    def _parse_metadata_json(cls, v):
-        if v is None or isinstance(v, dict):
-            return v
-        if isinstance(v, str) and v.strip():
-            try:
-                parsed = json.loads(v)
-            except Exception as e:
-                raise ValueError(f"metadata_filter must be JSON: {e}") from e
-            if not isinstance(parsed, dict):
-                raise ValueError("metadata_filter must be a JSON object")
-            return parsed
-        return None
-
-
-class UpdateAssetBody(BaseModel):
-    name: Optional[str] = None
-    tags: Optional[list[str]] = None
-    user_metadata: Optional[dict[str, Any]] = None
-
-    @model_validator(mode="after")
-    def _at_least_one(self):
-        if self.name is None and self.tags is None and self.user_metadata is None:
-            raise ValueError("Provide at least one of: name, tags, user_metadata.")
-        if self.tags is not None:
-            if not isinstance(self.tags, list) or not all(isinstance(t, str) for t in self.tags):
-                raise ValueError("Field 'tags' must be an array of strings.")
-        return self
-
-
-class CreateFromHashBody(BaseModel):
-    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)
-
-    hash: str
-    name: str
-    tags: list[str] = Field(default_factory=list)
-    user_metadata: dict[str, Any] = Field(default_factory=dict)
-
-    @field_validator("hash")
-    @classmethod
-    def _require_blake3(cls, v):
-        s = (v or "").strip().lower()
-        if ":" not in s:
-            raise ValueError("hash must be 'blake3:<hex>'")
-        algo, digest = s.split(":", 1)
-        if algo != "blake3":
-            raise ValueError("only canonical 'blake3:<hex>' is accepted here")
-        if not digest or any(c for c in digest if c not in "0123456789abcdef"):
-            raise ValueError("hash digest must be lowercase hex")
-        return s
-
-    @field_validator("tags", mode="before")
-    @classmethod
-    def _tags_norm(cls, v):
-        if v is None:
-            return []
-        if isinstance(v, list):
-            out = [str(t).strip().lower() for t in v if str(t).strip()]
-            seen = set()
-            dedup = []
-            for t in out:
-                if t not in seen:
-                    seen.add(t)
-                    dedup.append(t)
-            return dedup
-        if isinstance(v, str):
-            return [t.strip().lower() for t in v.split(",") if t.strip()]
-        return []
-
-
-class TagsListQuery(BaseModel):
-    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)
-
-    prefix: Optional[str] = Field(None, min_length=1, max_length=256)
-    limit: int = Field(100, ge=1, le=1000)
-    offset: int = Field(0, ge=0, le=10_000_000)
-    order: Literal["count_desc", "name_asc"] = "count_desc"
-    include_zero: bool = True
-
-    @field_validator("prefix")
-    @classmethod
-    def normalize_prefix(cls, v: Optional[str]) -> Optional[str]:
-        if v is None:
-            return v
-        v = v.strip()
-        return v.lower() or None
-
-
-class TagsAdd(BaseModel):
-    model_config = ConfigDict(extra="ignore")
-    tags: list[str] = Field(..., min_length=1)
-
-    @field_validator("tags")
-    @classmethod
-    def normalize_tags(cls, v: list[str]) -> list[str]:
-        out = []
-        for t in v:
-            if not isinstance(t, str):
-                raise TypeError("tags must be strings")
-            tnorm = t.strip().lower()
-            if tnorm:
-                out.append(tnorm)
-        seen = set()
-        deduplicated = []
-        for x in out:
-            if x not in seen:
-                seen.add(x)
-                deduplicated.append(x)
-        return deduplicated
-
-
-class TagsRemove(TagsAdd):
-    pass
-
-
-RootType = Literal["models", "input", "output"]
-ALLOWED_ROOTS: tuple[RootType, ...] = ("models", "input", "output")
-
-
-class ScheduleAssetScanBody(BaseModel):
-    roots: list[RootType] = Field(..., min_length=1)
-
-
-class UploadAssetSpec(BaseModel):
-    """Upload Asset operation.
-    - tags: ordered; first is root ('models'|'input'|'output');
-            if root == 'models', second must be a valid category from folder_paths.folder_names_and_paths
-    - name: display name
-    - user_metadata: arbitrary JSON object (optional)
-    - hash: optional canonical 'blake3:<hex>' provided by the client for validation / fast-path
-
-    Files created via this endpoint are stored on disk using the **content hash** as the filename stem
-    and the original extension is preserved when available.
-    """
-    model_config = ConfigDict(extra="ignore", str_strip_whitespace=True)
-
-    tags: list[str] = Field(..., min_length=1)
-    name: Optional[str] = Field(default=None, max_length=512, description="Display Name")
-    user_metadata: dict[str, Any] = Field(default_factory=dict)
-    hash: Optional[str] = Field(default=None)
-
-    @field_validator("hash", mode="before")
-    @classmethod
-    def _parse_hash(cls, v):
-        if v is None:
-            return None
-        s = str(v).strip().lower()
-        if not s:
-            return None
-        if ":" not in s:
-            raise ValueError("hash must be 'blake3:<hex>'")
-        algo, digest = s.split(":", 1)
-        if algo != "blake3":
-            raise ValueError("only canonical 'blake3:<hex>' is accepted here")
-        if not digest or any(c for c in digest if c not in "0123456789abcdef"):
-            raise ValueError("hash digest must be lowercase hex")
-        return f"{algo}:{digest}"
-
-    @field_validator("tags", mode="before")
-    @classmethod
-    def _parse_tags(cls, v):
-        """
-        Accepts a list of strings (possibly multiple form fields),
-        where each string can be:
-          - JSON array (e.g., '["models","loras","foo"]')
-          - comma-separated ('models, loras, foo')
-          - single token ('models')
-        Returns a normalized, deduplicated, ordered list.
-        """
-        items: list[str] = []
-        if v is None:
-            return []
-        if isinstance(v, str):
-            v = [v]
-
-        if isinstance(v, list):
-            for item in v:
-                if item is None:
-                    continue
-                s = str(item).strip()
-                if not s:
-                    continue
-                if s.startswith("["):
-                    try:
-                        arr = json.loads(s)
-                        if isinstance(arr, list):
-                            items.extend(str(x) for x in arr)
-                            continue
-                    except Exception:
-                        pass  # fallback to CSV parse below
-                items.extend([p for p in s.split(",") if p.strip()])
-        else:
-            return []
-
-        # normalize + dedupe
-        norm = []
-        seen = set()
-        for t in items:
-            tnorm = str(t).strip().lower()
-            if tnorm and tnorm not in seen:
-                seen.add(tnorm)
-                norm.append(tnorm)
-        return norm
-
-    @field_validator("user_metadata", mode="before")
-    @classmethod
-    def _parse_metadata_json(cls, v):
-        if v is None or isinstance(v, dict):
-            return v or {}
-        if isinstance(v, str):
-            s = v.strip()
-            if not s:
-                return {}
-            try:
-                parsed = json.loads(s)
-            except Exception as e:
-                raise ValueError(f"user_metadata must be JSON: {e}") from e
-            if not isinstance(parsed, dict):
-                raise ValueError("user_metadata must be a JSON object")
-            return parsed
-        return {}
-
-    @model_validator(mode="after")
-    def _validate_order(self):
-        if not self.tags:
-            raise ValueError("tags must be provided and non-empty")
-        root = self.tags[0]
-        if root not in {"models", "input", "output"}:
-            raise ValueError("first tag must be one of: models, input, output")
-        if root == "models":
-            if len(self.tags) < 2:
-                raise ValueError("models uploads require a category tag as the second tag")
-        return self
-
-
-class SetPreviewBody(BaseModel):
-    """Set or clear the preview for an AssetInfo. Provide an Asset.id or null."""
-    preview_id: Optional[str] = None
-
-    @field_validator("preview_id", mode="before")
-    @classmethod
-    def _norm_uuid(cls, v):
-        if v is None:
-            return None
-        s = str(v).strip()
-        if not s:
-            return None
-        try:
-            uuid.UUID(s)
-        except Exception:
-            raise ValueError("preview_id must be a UUID")
-        return s

app/assets/api/schemas_out.py

@@ -1,115 +0,0 @@
-from datetime import datetime
-from typing import Any, Literal, Optional
-
-from pydantic import BaseModel, ConfigDict, Field, field_serializer
-
-
-class AssetSummary(BaseModel):
-    id: str
-    name: str
-    asset_hash: Optional[str]
-    size: Optional[int] = None
-    mime_type: Optional[str] = None
-    tags: list[str] = Field(default_factory=list)
-    preview_url: Optional[str] = None
-    created_at: Optional[datetime] = None
-    updated_at: Optional[datetime] = None
-    last_access_time: Optional[datetime] = None
-
-    model_config = ConfigDict(from_attributes=True)
-
-    @field_serializer("created_at", "updated_at", "last_access_time")
-    def _ser_dt(self, v: Optional[datetime], _info):
-        return v.isoformat() if v else None
-
-
-class AssetsList(BaseModel):
-    assets: list[AssetSummary]
-    total: int
-    has_more: bool
-
-
-class AssetUpdated(BaseModel):
-    id: str
-    name: str
-    asset_hash: Optional[str]
-    tags: list[str] = Field(default_factory=list)
-    user_metadata: dict[str, Any] = Field(default_factory=dict)
-    updated_at: Optional[datetime] = None
-
-    model_config = ConfigDict(from_attributes=True)
-
-    @field_serializer("updated_at")
-    def _ser_updated(self, v: Optional[datetime], _info):
-        return v.isoformat() if v else None
-
-
-class AssetDetail(BaseModel):
-    id: str
-    name: str
-    asset_hash: Optional[str]
-    size: Optional[int] = None
-    mime_type: Optional[str] = None
-    tags: list[str] = Field(default_factory=list)
-    user_metadata: dict[str, Any] = Field(default_factory=dict)
-    preview_id: Optional[str] = None
-    created_at: Optional[datetime] = None
-    last_access_time: Optional[datetime] = None
-
-    model_config = ConfigDict(from_attributes=True)
-
-    @field_serializer("created_at", "last_access_time")
-    def _ser_dt(self, v: Optional[datetime], _info):
-        return v.isoformat() if v else None
-
-
-class AssetCreated(AssetDetail):
-    created_new: bool
-
-
-class TagUsage(BaseModel):
-    name: str
-    count: int
-    type: str
-
-
-class TagsList(BaseModel):
-    tags: list[TagUsage] = Field(default_factory=list)
-    total: int
-    has_more: bool
-
-
-class TagsAdd(BaseModel):
-    model_config = ConfigDict(str_strip_whitespace=True)
-    added: list[str] = Field(default_factory=list)
-    already_present: list[str] = Field(default_factory=list)
-    total_tags: list[str] = Field(default_factory=list)
-
-
-class TagsRemove(BaseModel):
-    model_config = ConfigDict(str_strip_whitespace=True)
-    removed: list[str] = Field(default_factory=list)
-    not_present: list[str] = Field(default_factory=list)
-    total_tags: list[str] = Field(default_factory=list)
-
-
-class AssetScanError(BaseModel):
-    path: str
-    message: str
-    at: Optional[str] = Field(None, description="ISO timestamp")
-
-
-class AssetScanStatus(BaseModel):
-    scan_id: str
-    root: Literal["models", "input", "output"]
-    status: Literal["scheduled", "running", "completed", "failed", "cancelled"]
-    scheduled_at: Optional[str] = None
-    started_at: Optional[str] = None
-    finished_at: Optional[str] = None
-    discovered: int = 0
-    processed: int = 0
-    file_errors: list[AssetScanError] = Field(default_factory=list)
-
-
-class AssetScanStatusResponse(BaseModel):
-    scans: list[AssetScanStatus] = Field(default_factory=list)

app/assets/database/helpers/__init__.py

@@ -1,25 +0,0 @@
-from .bulk_ops import seed_from_paths_batch
-from .escape_like import escape_like_prefix
-from .fast_check import fast_asset_file_check
-from .filters import apply_metadata_filter, apply_tag_filters
-from .ownership import visible_owner_clause
-from .projection import is_scalar, project_kv
-from .tags import (
-    add_missing_tag_for_asset_id,
-    ensure_tags_exist,
-    remove_missing_tag_for_asset_id,
-)
-
-__all__ = [
-    "apply_tag_filters",
-    "apply_metadata_filter",
-    "escape_like_prefix",
-    "fast_asset_file_check",
-    "is_scalar",
-    "project_kv",
-    "ensure_tags_exist",
-    "add_missing_tag_for_asset_id",
-    "remove_missing_tag_for_asset_id",
-    "seed_from_paths_batch",
-    "visible_owner_clause",
-]

app/assets/database/helpers/bulk_ops.py

@@ -1,230 +0,0 @@
-import os
-import uuid
-from typing import Iterable, Sequence
-
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql as d_pg
-from sqlalchemy.dialects import sqlite as d_sqlite
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from ..models import Asset, AssetCacheState, AssetInfo, AssetInfoMeta, AssetInfoTag
-from ..timeutil import utcnow
-
-MAX_BIND_PARAMS = 800
-
-
-async def seed_from_paths_batch(
-    session: AsyncSession,
-    *,
-    specs: Sequence[dict],
-    owner_id: str = "",
-) -> dict:
-    """Each spec is a dict with keys:
-      - abs_path: str
-      - size_bytes: int
-      - mtime_ns: int
-      - info_name: str
-      - tags: list[str]
-      - fname: Optional[str]
-    """
-    if not specs:
-        return {"inserted_infos": 0, "won_states": 0, "lost_states": 0}
-
-    now = utcnow()
-    dialect = session.bind.dialect.name
-    if dialect not in ("sqlite", "postgresql"):
-        raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-
-    asset_rows: list[dict] = []
-    state_rows: list[dict] = []
-    path_to_asset: dict[str, str] = {}
-    asset_to_info: dict[str, dict] = {}  # asset_id -> prepared info row
-    path_list: list[str] = []
-
-    for sp in specs:
-        ap = os.path.abspath(sp["abs_path"])
-        aid = str(uuid.uuid4())
-        iid = str(uuid.uuid4())
-        path_list.append(ap)
-        path_to_asset[ap] = aid
-
-        asset_rows.append(
-            {
-                "id": aid,
-                "hash": None,
-                "size_bytes": sp["size_bytes"],
-                "mime_type": None,
-                "created_at": now,
-            }
-        )
-        state_rows.append(
-            {
-                "asset_id": aid,
-                "file_path": ap,
-                "mtime_ns": sp["mtime_ns"],
-            }
-        )
-        asset_to_info[aid] = {
-            "id": iid,
-            "owner_id": owner_id,
-            "name": sp["info_name"],
-            "asset_id": aid,
-            "preview_id": None,
-            "user_metadata": {"filename": sp["fname"]} if sp["fname"] else None,
-            "created_at": now,
-            "updated_at": now,
-            "last_access_time": now,
-            "_tags": sp["tags"],
-            "_filename": sp["fname"],
-        }
-
-    # insert all seed Assets (hash=NULL)
-    ins_asset = d_sqlite.insert(Asset) if dialect == "sqlite" else d_pg.insert(Asset)
-    for chunk in _iter_chunks(asset_rows, _rows_per_stmt(5)):
-        await session.execute(ins_asset, chunk)
-
-    # try to claim AssetCacheState (file_path)
-    winners_by_path: set[str] = set()
-    if dialect == "sqlite":
-        ins_state = (
-            d_sqlite.insert(AssetCacheState)
-            .on_conflict_do_nothing(index_elements=[AssetCacheState.file_path])
-            .returning(AssetCacheState.file_path)
-        )
-    else:
-        ins_state = (
-            d_pg.insert(AssetCacheState)
-            .on_conflict_do_nothing(index_elements=[AssetCacheState.file_path])
-            .returning(AssetCacheState.file_path)
-        )
-    for chunk in _iter_chunks(state_rows, _rows_per_stmt(3)):
-        winners_by_path.update((await session.execute(ins_state, chunk)).scalars().all())
-
-    all_paths_set = set(path_list)
-    losers_by_path = all_paths_set - winners_by_path
-    lost_assets = [path_to_asset[p] for p in losers_by_path]
-    if lost_assets:  # losers get their Asset removed
-        for id_chunk in _iter_chunks(lost_assets, MAX_BIND_PARAMS):
-            await session.execute(sa.delete(Asset).where(Asset.id.in_(id_chunk)))
-
-    if not winners_by_path:
-        return {"inserted_infos": 0, "won_states": 0, "lost_states": len(losers_by_path)}
-
-    # insert AssetInfo only for winners
-    winner_info_rows = [asset_to_info[path_to_asset[p]] for p in winners_by_path]
-    if dialect == "sqlite":
-        ins_info = (
-            d_sqlite.insert(AssetInfo)
-            .on_conflict_do_nothing(index_elements=[AssetInfo.asset_id, AssetInfo.owner_id, AssetInfo.name])
-            .returning(AssetInfo.id)
-        )
-    else:
-        ins_info = (
-            d_pg.insert(AssetInfo)
-            .on_conflict_do_nothing(index_elements=[AssetInfo.asset_id, AssetInfo.owner_id, AssetInfo.name])
-            .returning(AssetInfo.id)
-        )
-
-    inserted_info_ids: set[str] = set()
-    for chunk in _iter_chunks(winner_info_rows, _rows_per_stmt(9)):
-        inserted_info_ids.update((await session.execute(ins_info, chunk)).scalars().all())
-
-    # build and insert tag + meta rows for the AssetInfo
-    tag_rows: list[dict] = []
-    meta_rows: list[dict] = []
-    if inserted_info_ids:
-        for row in winner_info_rows:
-            iid = row["id"]
-            if iid not in inserted_info_ids:
-                continue
-            for t in row["_tags"]:
-                tag_rows.append({
-                    "asset_info_id": iid,
-                    "tag_name": t,
-                    "origin": "automatic",
-                    "added_at": now,
-                })
-            if row["_filename"]:
-                meta_rows.append(
-                    {
-                        "asset_info_id": iid,
-                        "key": "filename",
-                        "ordinal": 0,
-                        "val_str": row["_filename"],
-                        "val_num": None,
-                        "val_bool": None,
-                        "val_json": None,
-                    }
-                )
-
-    await bulk_insert_tags_and_meta(session, tag_rows=tag_rows, meta_rows=meta_rows, max_bind_params=MAX_BIND_PARAMS)
-    return {
-        "inserted_infos": len(inserted_info_ids),
-        "won_states": len(winners_by_path),
-        "lost_states": len(losers_by_path),
-    }
-
-
-async def bulk_insert_tags_and_meta(
-    session: AsyncSession,
-    *,
-    tag_rows: list[dict],
-    meta_rows: list[dict],
-    max_bind_params: int,
-) -> None:
-    """Batch insert into asset_info_tags and asset_info_meta with ON CONFLICT DO NOTHING.
-    - tag_rows keys: asset_info_id, tag_name, origin, added_at
-    - meta_rows keys: asset_info_id, key, ordinal, val_str, val_num, val_bool, val_json
-    """
-    dialect = session.bind.dialect.name
-    if tag_rows:
-        if dialect == "sqlite":
-            ins_links = (
-                d_sqlite.insert(AssetInfoTag)
-                .on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
-            )
-        elif dialect == "postgresql":
-            ins_links = (
-                d_pg.insert(AssetInfoTag)
-                .on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
-            )
-        else:
-            raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-        for chunk in _chunk_rows(tag_rows, cols_per_row=4, max_bind_params=max_bind_params):
-            await session.execute(ins_links, chunk)
-    if meta_rows:
-        if dialect == "sqlite":
-            ins_meta = (
-                d_sqlite.insert(AssetInfoMeta)
-                .on_conflict_do_nothing(
-                    index_elements=[AssetInfoMeta.asset_info_id, AssetInfoMeta.key, AssetInfoMeta.ordinal]
-                )
-            )
-        elif dialect == "postgresql":
-            ins_meta = (
-                d_pg.insert(AssetInfoMeta)
-                .on_conflict_do_nothing(
-                    index_elements=[AssetInfoMeta.asset_info_id, AssetInfoMeta.key, AssetInfoMeta.ordinal]
-                )
-            )
-        else:
-            raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-        for chunk in _chunk_rows(meta_rows, cols_per_row=7, max_bind_params=max_bind_params):
-            await session.execute(ins_meta, chunk)
-
-
-def _chunk_rows(rows: list[dict], cols_per_row: int, max_bind_params: int) -> Iterable[list[dict]]:
-    if not rows:
-        return []
-    rows_per_stmt = max(1, max_bind_params // max(1, cols_per_row))
-    for i in range(0, len(rows), rows_per_stmt):
-        yield rows[i:i + rows_per_stmt]
-
-
-def _iter_chunks(seq, n: int):
-    for i in range(0, len(seq), n):
-        yield seq[i:i + n]
-
-
-def _rows_per_stmt(cols: int) -> int:
-    return max(1, MAX_BIND_PARAMS // max(1, cols))

app/assets/database/helpers/escape_like.py

@@ -1,7 +0,0 @@
-def escape_like_prefix(s: str, escape: str = "!") -> tuple[str, str]:
-    """Escapes %, _ and the escape char itself in a LIKE prefix.
-    Returns (escaped_prefix, escape_char). Caller should append '%' and pass escape=escape_char to .like().
-    """
-    s = s.replace(escape, escape + escape)  # escape the escape char first
-    s = s.replace("%", escape + "%").replace("_", escape + "_")  # escape LIKE wildcards
-    return s, escape

app/assets/database/helpers/fast_check.py

@@ -1,19 +0,0 @@
-import os
-from typing import Optional
-
-
-def fast_asset_file_check(
-    *,
-    mtime_db: Optional[int],
-    size_db: Optional[int],
-    stat_result: os.stat_result,
-) -> bool:
-    if mtime_db is None:
-        return False
-    actual_mtime_ns = getattr(stat_result, "st_mtime_ns", int(stat_result.st_mtime * 1_000_000_000))
-    if int(mtime_db) != int(actual_mtime_ns):
-        return False
-    sz = int(size_db or 0)
-    if sz > 0:
-        return int(stat_result.st_size) == sz
-    return True

app/assets/database/helpers/filters.py

@@ -1,87 +0,0 @@
-from typing import Optional, Sequence
-
-import sqlalchemy as sa
-from sqlalchemy import exists
-
-from ..._helpers import normalize_tags
-from ..models import AssetInfo, AssetInfoMeta, AssetInfoTag
-
-
-def apply_tag_filters(
-    stmt: sa.sql.Select,
-    include_tags: Optional[Sequence[str]],
-    exclude_tags: Optional[Sequence[str]],
-) -> sa.sql.Select:
-    """include_tags: every tag must be present; exclude_tags: none may be present."""
-    include_tags = normalize_tags(include_tags)
-    exclude_tags = normalize_tags(exclude_tags)
-
-    if include_tags:
-        for tag_name in include_tags:
-            stmt = stmt.where(
-                exists().where(
-                    (AssetInfoTag.asset_info_id == AssetInfo.id)
-                    & (AssetInfoTag.tag_name == tag_name)
-                )
-            )
-
-    if exclude_tags:
-        stmt = stmt.where(
-            ~exists().where(
-                (AssetInfoTag.asset_info_id == AssetInfo.id)
-                & (AssetInfoTag.tag_name.in_(exclude_tags))
-            )
-        )
-    return stmt
-
-
-def apply_metadata_filter(
-    stmt: sa.sql.Select,
-    metadata_filter: Optional[dict],
-) -> sa.sql.Select:
-    """Apply filters using asset_info_meta projection table."""
-    if not metadata_filter:
-        return stmt
-
-    def _exists_for_pred(key: str, *preds) -> sa.sql.ClauseElement:
-        return sa.exists().where(
-            AssetInfoMeta.asset_info_id == AssetInfo.id,
-            AssetInfoMeta.key == key,
-            *preds,
-        )
-
-    def _exists_clause_for_value(key: str, value) -> sa.sql.ClauseElement:
-        if value is None:
-            no_row_for_key = sa.not_(
-                sa.exists().where(
-                    AssetInfoMeta.asset_info_id == AssetInfo.id,
-                    AssetInfoMeta.key == key,
-                )
-            )
-            null_row = _exists_for_pred(
-                key,
-                AssetInfoMeta.val_json.is_(None),
-                AssetInfoMeta.val_str.is_(None),
-                AssetInfoMeta.val_num.is_(None),
-                AssetInfoMeta.val_bool.is_(None),
-            )
-            return sa.or_(no_row_for_key, null_row)
-
-        if isinstance(value, bool):
-            return _exists_for_pred(key, AssetInfoMeta.val_bool == bool(value))
-        if isinstance(value, (int, float)):
-            from decimal import Decimal
-            num = value if isinstance(value, Decimal) else Decimal(str(value))
-            return _exists_for_pred(key, AssetInfoMeta.val_num == num)
-        if isinstance(value, str):
-            return _exists_for_pred(key, AssetInfoMeta.val_str == value)
-        return _exists_for_pred(key, AssetInfoMeta.val_json == value)
-
-    for k, v in metadata_filter.items():
-        if isinstance(v, list):
-            ors = [_exists_clause_for_value(k, elem) for elem in v]
-            if ors:
-                stmt = stmt.where(sa.or_(*ors))
-        else:
-            stmt = stmt.where(_exists_clause_for_value(k, v))
-    return stmt

app/assets/database/helpers/ownership.py

@@ -1,12 +0,0 @@
-import sqlalchemy as sa
-
-from ..models import AssetInfo
-
-
-def visible_owner_clause(owner_id: str) -> sa.sql.ClauseElement:
-    """Build owner visibility predicate for reads. Owner-less rows are visible to everyone."""
-
-    owner_id = (owner_id or "").strip()
-    if owner_id == "":
-        return AssetInfo.owner_id == ""
-    return AssetInfo.owner_id.in_(["", owner_id])

app/assets/database/helpers/projection.py

@@ -1,64 +0,0 @@
-from decimal import Decimal
-
-
-def is_scalar(v):
-    if v is None:
-        return True
-    if isinstance(v, bool):
-        return True
-    if isinstance(v, (int, float, Decimal, str)):
-        return True
-    return False
-
-
-def project_kv(key: str, value):
-    """
-    Turn a metadata key/value into typed projection rows.
-    Returns list[dict] with keys:
-      key, ordinal, and one of val_str / val_num / val_bool / val_json (others None)
-    """
-    rows: list[dict] = []
-
-    def _null_row(ordinal: int) -> dict:
-        return {
-            "key": key, "ordinal": ordinal,
-            "val_str": None, "val_num": None, "val_bool": None, "val_json": None
-        }
-
-    if value is None:
-        rows.append(_null_row(0))
-        return rows
-
-    if is_scalar(value):
-        if isinstance(value, bool):
-            rows.append({"key": key, "ordinal": 0, "val_bool": bool(value)})
-        elif isinstance(value, (int, float, Decimal)):
-            num = value if isinstance(value, Decimal) else Decimal(str(value))
-            rows.append({"key": key, "ordinal": 0, "val_num": num})
-        elif isinstance(value, str):
-            rows.append({"key": key, "ordinal": 0, "val_str": value})
-        else:
-            rows.append({"key": key, "ordinal": 0, "val_json": value})
-        return rows
-
-    if isinstance(value, list):
-        if all(is_scalar(x) for x in value):
-            for i, x in enumerate(value):
-                if x is None:
-                    rows.append(_null_row(i))
-                elif isinstance(x, bool):
-                    rows.append({"key": key, "ordinal": i, "val_bool": bool(x)})
-                elif isinstance(x, (int, float, Decimal)):
-                    num = x if isinstance(x, Decimal) else Decimal(str(x))
-                    rows.append({"key": key, "ordinal": i, "val_num": num})
-                elif isinstance(x, str):
-                    rows.append({"key": key, "ordinal": i, "val_str": x})
-                else:
-                    rows.append({"key": key, "ordinal": i, "val_json": x})
-            return rows
-        for i, x in enumerate(value):
-            rows.append({"key": key, "ordinal": i, "val_json": x})
-        return rows
-
-    rows.append({"key": key, "ordinal": 0, "val_json": value})
-    return rows

app/assets/database/helpers/tags.py

@@ -1,90 +0,0 @@
-from typing import Iterable
-
-import sqlalchemy as sa
-from sqlalchemy.dialects import postgresql as d_pg
-from sqlalchemy.dialects import sqlite as d_sqlite
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from ..._helpers import normalize_tags
-from ..models import AssetInfo, AssetInfoTag, Tag
-from ..timeutil import utcnow
-
-
-async def ensure_tags_exist(session: AsyncSession, names: Iterable[str], tag_type: str = "user") -> None:
-    wanted = normalize_tags(list(names))
-    if not wanted:
-        return
-    rows = [{"name": n, "tag_type": tag_type} for n in list(dict.fromkeys(wanted))]
-    dialect = session.bind.dialect.name
-    if dialect == "sqlite":
-        ins = (
-            d_sqlite.insert(Tag)
-            .values(rows)
-            .on_conflict_do_nothing(index_elements=[Tag.name])
-        )
-    elif dialect == "postgresql":
-        ins = (
-            d_pg.insert(Tag)
-            .values(rows)
-            .on_conflict_do_nothing(index_elements=[Tag.name])
-        )
-    else:
-        raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-    await session.execute(ins)
-
-
-async def add_missing_tag_for_asset_id(
-    session: AsyncSession,
-    *,
-    asset_id: str,
-    origin: str = "automatic",
-) -> None:
-    select_rows = (
-        sa.select(
-            AssetInfo.id.label("asset_info_id"),
-            sa.literal("missing").label("tag_name"),
-            sa.literal(origin).label("origin"),
-            sa.literal(utcnow()).label("added_at"),
-        )
-        .where(AssetInfo.asset_id == asset_id)
-        .where(
-            sa.not_(
-                sa.exists().where((AssetInfoTag.asset_info_id == AssetInfo.id) & (AssetInfoTag.tag_name == "missing"))
-            )
-        )
-    )
-    dialect = session.bind.dialect.name
-    if dialect == "sqlite":
-        ins = (
-            d_sqlite.insert(AssetInfoTag)
-            .from_select(
-                ["asset_info_id", "tag_name", "origin", "added_at"],
-                select_rows,
-            )
-            .on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
-        )
-    elif dialect == "postgresql":
-        ins = (
-            d_pg.insert(AssetInfoTag)
-            .from_select(
-                ["asset_info_id", "tag_name", "origin", "added_at"],
-                select_rows,
-            )
-            .on_conflict_do_nothing(index_elements=[AssetInfoTag.asset_info_id, AssetInfoTag.tag_name])
-        )
-    else:
-        raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-    await session.execute(ins)
-
-
-async def remove_missing_tag_for_asset_id(
-    session: AsyncSession,
-    *,
-    asset_id: str,
-) -> None:
-    await session.execute(
-        sa.delete(AssetInfoTag).where(
-            AssetInfoTag.asset_info_id.in_(sa.select(AssetInfo.id).where(AssetInfo.asset_id == asset_id)),
-            AssetInfoTag.tag_name == "missing",
-        )
-    )

app/assets/database/models.py

@@ -1,251 +0,0 @@
-import uuid
-from datetime import datetime
-from typing import Any, Optional
-
-from sqlalchemy import (
-    JSON,
-    BigInteger,
-    Boolean,
-    CheckConstraint,
-    DateTime,
-    ForeignKey,
-    Index,
-    Integer,
-    Numeric,
-    String,
-    Text,
-    UniqueConstraint,
-)
-from sqlalchemy.dialects.postgresql import JSONB
-from sqlalchemy.orm import DeclarativeBase, Mapped, foreign, mapped_column, relationship
-
-from .timeutil import utcnow
-
-JSONB_V = JSON(none_as_null=True).with_variant(JSONB(none_as_null=True), 'postgresql')
-
-
-class Base(DeclarativeBase):
-    pass
-
-
-def to_dict(obj: Any, include_none: bool = False) -> dict[str, Any]:
-    fields = obj.__table__.columns.keys()
-    out: dict[str, Any] = {}
-    for field in fields:
-        val = getattr(obj, field)
-        if val is None and not include_none:
-            continue
-        if isinstance(val, datetime):
-            out[field] = val.isoformat()
-        else:
-            out[field] = val
-    return out
-
-
-class Asset(Base):
-    __tablename__ = "assets"
-
-    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
-    hash: Mapped[Optional[str]] = mapped_column(String(256), nullable=True)
-    size_bytes: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
-    mime_type: Mapped[Optional[str]] = mapped_column(String(255))
-    created_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=False), nullable=False, default=utcnow
-    )
-
-    infos: Mapped[list["AssetInfo"]] = relationship(
-        "AssetInfo",
-        back_populates="asset",
-        primaryjoin=lambda: Asset.id == foreign(AssetInfo.asset_id),
-        foreign_keys=lambda: [AssetInfo.asset_id],
-        cascade="all,delete-orphan",
-        passive_deletes=True,
-    )
-
-    preview_of: Mapped[list["AssetInfo"]] = relationship(
-        "AssetInfo",
-        back_populates="preview_asset",
-        primaryjoin=lambda: Asset.id == foreign(AssetInfo.preview_id),
-        foreign_keys=lambda: [AssetInfo.preview_id],
-        viewonly=True,
-    )
-
-    cache_states: Mapped[list["AssetCacheState"]] = relationship(
-        back_populates="asset",
-        cascade="all, delete-orphan",
-        passive_deletes=True,
-    )
-
-    __table_args__ = (
-        Index("uq_assets_hash", "hash", unique=True),
-        Index("ix_assets_mime_type", "mime_type"),
-        CheckConstraint("size_bytes >= 0", name="ck_assets_size_nonneg"),
-    )
-
-    def to_dict(self, include_none: bool = False) -> dict[str, Any]:
-        return to_dict(self, include_none=include_none)
-
-    def __repr__(self) -> str:
-        return f"<Asset id={self.id} hash={(self.hash or '')[:12]}>"
-
-
-class AssetCacheState(Base):
-    __tablename__ = "asset_cache_state"
-
-    id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True)
-    asset_id: Mapped[str] = mapped_column(String(36), ForeignKey("assets.id", ondelete="CASCADE"), nullable=False)
-    file_path: Mapped[str] = mapped_column(Text, nullable=False)
-    mtime_ns: Mapped[Optional[int]] = mapped_column(BigInteger, nullable=True)
-    needs_verify: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
-
-    asset: Mapped["Asset"] = relationship(back_populates="cache_states")
-
-    __table_args__ = (
-        Index("ix_asset_cache_state_file_path", "file_path"),
-        Index("ix_asset_cache_state_asset_id", "asset_id"),
-        CheckConstraint("(mtime_ns IS NULL) OR (mtime_ns >= 0)", name="ck_acs_mtime_nonneg"),
-        UniqueConstraint("file_path", name="uq_asset_cache_state_file_path"),
-    )
-
-    def to_dict(self, include_none: bool = False) -> dict[str, Any]:
-        return to_dict(self, include_none=include_none)
-
-    def __repr__(self) -> str:
-        return f"<AssetCacheState id={self.id} asset_id={self.asset_id} path={self.file_path!r}>"
-
-
-class AssetInfo(Base):
-    __tablename__ = "assets_info"
-
-    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
-    owner_id: Mapped[str] = mapped_column(String(128), nullable=False, default="")
-    name: Mapped[str] = mapped_column(String(512), nullable=False)
-    asset_id: Mapped[str] = mapped_column(String(36), ForeignKey("assets.id", ondelete="RESTRICT"), nullable=False)
-    preview_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("assets.id", ondelete="SET NULL"))
-    user_metadata: Mapped[Optional[dict[str, Any]]] = mapped_column(JSON(none_as_null=True))
-    created_at: Mapped[datetime] = mapped_column(DateTime(timezone=False), nullable=False, default=utcnow)
-    updated_at: Mapped[datetime] = mapped_column(DateTime(timezone=False), nullable=False, default=utcnow)
-    last_access_time: Mapped[datetime] = mapped_column(DateTime(timezone=False), nullable=False, default=utcnow)
-
-    asset: Mapped[Asset] = relationship(
-        "Asset",
-        back_populates="infos",
-        foreign_keys=[asset_id],
-        lazy="selectin",
-    )
-    preview_asset: Mapped[Optional[Asset]] = relationship(
-        "Asset",
-        back_populates="preview_of",
-        foreign_keys=[preview_id],
-    )
-
-    metadata_entries: Mapped[list["AssetInfoMeta"]] = relationship(
-        back_populates="asset_info",
-        cascade="all,delete-orphan",
-        passive_deletes=True,
-    )
-
-    tag_links: Mapped[list["AssetInfoTag"]] = relationship(
-        back_populates="asset_info",
-        cascade="all,delete-orphan",
-        passive_deletes=True,
-        overlaps="tags,asset_infos",
-    )
-
-    tags: Mapped[list["Tag"]] = relationship(
-        secondary="asset_info_tags",
-        back_populates="asset_infos",
-        lazy="selectin",
-        viewonly=True,
-        overlaps="tag_links,asset_info_links,asset_infos,tag",
-    )
-
-    __table_args__ = (
-        UniqueConstraint("asset_id", "owner_id", "name", name="uq_assets_info_asset_owner_name"),
-        Index("ix_assets_info_owner_name", "owner_id", "name"),
-        Index("ix_assets_info_owner_id", "owner_id"),
-        Index("ix_assets_info_asset_id", "asset_id"),
-        Index("ix_assets_info_name", "name"),
-        Index("ix_assets_info_created_at", "created_at"),
-        Index("ix_assets_info_last_access_time", "last_access_time"),
-    )
-
-    def to_dict(self, include_none: bool = False) -> dict[str, Any]:
-        data = to_dict(self, include_none=include_none)
-        data["tags"] = [t.name for t in self.tags]
-        return data
-
-    def __repr__(self) -> str:
-        return f"<AssetInfo id={self.id} name={self.name!r} asset_id={self.asset_id}>"
-
-
-class AssetInfoMeta(Base):
-    __tablename__ = "asset_info_meta"
-
-    asset_info_id: Mapped[str] = mapped_column(
-        String(36), ForeignKey("assets_info.id", ondelete="CASCADE"), primary_key=True
-    )
-    key: Mapped[str] = mapped_column(String(256), primary_key=True)
-    ordinal: Mapped[int] = mapped_column(Integer, primary_key=True, default=0)
-
-    val_str: Mapped[Optional[str]] = mapped_column(String(2048), nullable=True)
-    val_num: Mapped[Optional[float]] = mapped_column(Numeric(38, 10), nullable=True)
-    val_bool: Mapped[Optional[bool]] = mapped_column(Boolean, nullable=True)
-    val_json: Mapped[Optional[Any]] = mapped_column(JSONB_V, nullable=True)
-
-    asset_info: Mapped["AssetInfo"] = relationship(back_populates="metadata_entries")
-
-    __table_args__ = (
-        Index("ix_asset_info_meta_key", "key"),
-        Index("ix_asset_info_meta_key_val_str", "key", "val_str"),
-        Index("ix_asset_info_meta_key_val_num", "key", "val_num"),
-        Index("ix_asset_info_meta_key_val_bool", "key", "val_bool"),
-    )
-
-
-class AssetInfoTag(Base):
-    __tablename__ = "asset_info_tags"
-
-    asset_info_id: Mapped[str] = mapped_column(
-        String(36), ForeignKey("assets_info.id", ondelete="CASCADE"), primary_key=True
-    )
-    tag_name: Mapped[str] = mapped_column(
-        String(512), ForeignKey("tags.name", ondelete="RESTRICT"), primary_key=True
-    )
-    origin: Mapped[str] = mapped_column(String(32), nullable=False, default="manual")
-    added_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=False), nullable=False, default=utcnow
-    )
-
-    asset_info: Mapped["AssetInfo"] = relationship(back_populates="tag_links")
-    tag: Mapped["Tag"] = relationship(back_populates="asset_info_links")
-
-    __table_args__ = (
-        Index("ix_asset_info_tags_tag_name", "tag_name"),
-        Index("ix_asset_info_tags_asset_info_id", "asset_info_id"),
-    )
-
-
-class Tag(Base):
-    __tablename__ = "tags"
-
-    name: Mapped[str] = mapped_column(String(512), primary_key=True)
-    tag_type: Mapped[str] = mapped_column(String(32), nullable=False, default="user")
-
-    asset_info_links: Mapped[list["AssetInfoTag"]] = relationship(
-        back_populates="tag",
-        overlaps="asset_infos,tags",
-    )
-    asset_infos: Mapped[list["AssetInfo"]] = relationship(
-        secondary="asset_info_tags",
-        back_populates="tags",
-        viewonly=True,
-        overlaps="asset_info_links,tag_links,tags,asset_info",
-    )
-
-    __table_args__ = (
-        Index("ix_tags_tag_type", "tag_type"),
-    )
-
-    def __repr__(self) -> str:
-        return f"<Tag {self.name}>"

app/assets/database/services/__init__.py

@@ -1,57 +0,0 @@
-from .content import (
-    check_fs_asset_exists_quick,
-    compute_hash_and_dedup_for_cache_state,
-    ingest_fs_asset,
-    list_cache_states_with_asset_under_prefixes,
-    list_unhashed_candidates_under_prefixes,
-    list_verify_candidates_under_prefixes,
-    redirect_all_references_then_delete_asset,
-    touch_asset_infos_by_fs_path,
-)
-from .info import (
-    add_tags_to_asset_info,
-    create_asset_info_for_existing_asset,
-    delete_asset_info_by_id,
-    fetch_asset_info_and_asset,
-    fetch_asset_info_asset_and_tags,
-    get_asset_tags,
-    list_asset_infos_page,
-    list_tags_with_usage,
-    remove_tags_from_asset_info,
-    replace_asset_info_metadata_projection,
-    set_asset_info_preview,
-    set_asset_info_tags,
-    touch_asset_info_by_id,
-    update_asset_info_full,
-)
-from .queries import (
-    asset_exists_by_hash,
-    asset_info_exists_for_asset_id,
-    get_asset_by_hash,
-    get_asset_info_by_id,
-    get_cache_state_by_asset_id,
-    list_cache_states_by_asset_id,
-    pick_best_live_path,
-)
-
-__all__ = [
-    # queries
-    "asset_exists_by_hash", "get_asset_by_hash", "get_asset_info_by_id", "asset_info_exists_for_asset_id",
-    "get_cache_state_by_asset_id",
-    "list_cache_states_by_asset_id",
-    "pick_best_live_path",
-    # info
-    "list_asset_infos_page", "create_asset_info_for_existing_asset", "set_asset_info_tags",
-    "update_asset_info_full", "replace_asset_info_metadata_projection",
-    "touch_asset_info_by_id", "delete_asset_info_by_id",
-    "add_tags_to_asset_info", "remove_tags_from_asset_info",
-    "get_asset_tags", "list_tags_with_usage", "set_asset_info_preview",
-    "fetch_asset_info_and_asset", "fetch_asset_info_asset_and_tags",
-    # content
-    "check_fs_asset_exists_quick",
-    "redirect_all_references_then_delete_asset",
-    "compute_hash_and_dedup_for_cache_state",
-    "list_unhashed_candidates_under_prefixes", "list_verify_candidates_under_prefixes",
-    "ingest_fs_asset", "touch_asset_infos_by_fs_path",
-    "list_cache_states_with_asset_under_prefixes",
-]

app/assets/database/services/content.py

@@ -1,721 +0,0 @@
-import contextlib
-import logging
-import os
-from datetime import datetime
-from typing import Any, Optional, Sequence, Union
-
-import sqlalchemy as sa
-from sqlalchemy import select
-from sqlalchemy.dialects import postgresql as d_pg
-from sqlalchemy.dialects import sqlite as d_sqlite
-from sqlalchemy.exc import IntegrityError
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import noload
-
-from ..._helpers import compute_relative_filename
-from ...storage import hashing as hashing_mod
-from ..helpers import (
-    ensure_tags_exist,
-    escape_like_prefix,
-    remove_missing_tag_for_asset_id,
-)
-from ..models import Asset, AssetCacheState, AssetInfo, AssetInfoTag, Tag
-from ..timeutil import utcnow
-from .info import replace_asset_info_metadata_projection
-from .queries import list_cache_states_by_asset_id, pick_best_live_path
-
-
-async def check_fs_asset_exists_quick(
-    session: AsyncSession,
-    *,
-    file_path: str,
-    size_bytes: Optional[int] = None,
-    mtime_ns: Optional[int] = None,
-) -> bool:
-    """Returns True if we already track this absolute path with a HASHED asset and the cached mtime/size match."""
-    locator = os.path.abspath(file_path)
-
-    stmt = (
-        sa.select(sa.literal(True))
-        .select_from(AssetCacheState)
-        .join(Asset, Asset.id == AssetCacheState.asset_id)
-        .where(
-            AssetCacheState.file_path == locator,
-            Asset.hash.isnot(None),
-            AssetCacheState.needs_verify.is_(False),
-        )
-        .limit(1)
-    )
-
-    conds = []
-    if mtime_ns is not None:
-        conds.append(AssetCacheState.mtime_ns == int(mtime_ns))
-    if size_bytes is not None:
-        conds.append(sa.or_(Asset.size_bytes == 0, Asset.size_bytes == int(size_bytes)))
-    if conds:
-        stmt = stmt.where(*conds)
-    return (await session.execute(stmt)).first() is not None
-
-
-async def redirect_all_references_then_delete_asset(
-    session: AsyncSession,
-    *,
-    duplicate_asset_id: str,
-    canonical_asset_id: str,
-) -> None:
-    """
-    Safely migrate all references from duplicate_asset_id to canonical_asset_id.
-
-    - If an AssetInfo for (owner_id, name) already exists on the canonical asset,
-      merge tags, metadata, times, and preview, then delete the duplicate AssetInfo.
-    - Otherwise, simply repoint the AssetInfo.asset_id.
-    - Always retarget AssetCacheState rows.
-    - Finally delete the duplicate Asset row.
-    """
-    if duplicate_asset_id == canonical_asset_id:
-        return
-
-    # 1) Migrate AssetInfo rows one-by-one to avoid UNIQUE conflicts.
-    dup_infos = (
-        await session.execute(
-            select(AssetInfo).options(noload(AssetInfo.tags)).where(AssetInfo.asset_id == duplicate_asset_id)
-        )
-    ).unique().scalars().all()
-
-    for info in dup_infos:
-        # Try to find an existing collision on canonical
-        existing = (
-            await session.execute(
-                select(AssetInfo)
-                .options(noload(AssetInfo.tags))
-                .where(
-                    AssetInfo.asset_id == canonical_asset_id,
-                    AssetInfo.owner_id == info.owner_id,
-                    AssetInfo.name == info.name,
-                )
-                .limit(1)
-            )
-        ).unique().scalars().first()
-
-        if existing:
-            merged_meta = dict(existing.user_metadata or {})
-            other_meta = info.user_metadata or {}
-            for k, v in other_meta.items():
-                if k not in merged_meta:
-                    merged_meta[k] = v
-            if merged_meta != (existing.user_metadata or {}):
-                await replace_asset_info_metadata_projection(
-                    session,
-                    asset_info_id=existing.id,
-                    user_metadata=merged_meta,
-                )
-
-            existing_tags = {
-                t for (t,) in (
-                    await session.execute(
-                        select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == existing.id)
-                    )
-                ).all()
-            }
-            from_tags = {
-                t for (t,) in (
-                    await session.execute(
-                        select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == info.id)
-                    )
-                ).all()
-            }
-            to_add = sorted(from_tags - existing_tags)
-            if to_add:
-                await ensure_tags_exist(session, to_add, tag_type="user")
-                now = utcnow()
-                session.add_all([
-                    AssetInfoTag(asset_info_id=existing.id, tag_name=t, origin="automatic", added_at=now)
-                    for t in to_add
-                ])
-                await session.flush()
-
-            if existing.preview_id is None and info.preview_id is not None:
-                existing.preview_id = info.preview_id
-            if info.last_access_time and (
-                existing.last_access_time is None or info.last_access_time > existing.last_access_time
-            ):
-                existing.last_access_time = info.last_access_time
-            existing.updated_at = utcnow()
-            await session.flush()
-
-            # Delete the duplicate AssetInfo (cascades will clean its tags/meta)
-            await session.delete(info)
-            await session.flush()
-        else:
-            # Simple retarget
-            info.asset_id = canonical_asset_id
-            info.updated_at = utcnow()
-            await session.flush()
-
-    # 2) Repoint cache states and previews
-    await session.execute(
-        sa.update(AssetCacheState)
-        .where(AssetCacheState.asset_id == duplicate_asset_id)
-        .values(asset_id=canonical_asset_id)
-    )
-    await session.execute(
-        sa.update(AssetInfo)
-        .where(AssetInfo.preview_id == duplicate_asset_id)
-        .values(preview_id=canonical_asset_id)
-    )
-
-    # 3) Remove duplicate Asset
-    dup = await session.get(Asset, duplicate_asset_id)
-    if dup:
-        await session.delete(dup)
-    await session.flush()
-
-
-async def compute_hash_and_dedup_for_cache_state(
-    session: AsyncSession,
-    *,
-    state_id: int,
-) -> Optional[str]:
-    """
-    Compute hash for the given cache state, deduplicate, and settle verify cases.
-
-    Returns the asset_id that this state ends up pointing to, or None if file disappeared.
-    """
-    state = await session.get(AssetCacheState, state_id)
-    if not state:
-        return None
-
-    path = state.file_path
-    try:
-        if not os.path.isfile(path):
-            # File vanished: drop the state. If the Asset has hash=NULL and has no other states, drop the Asset too.
-            asset = await session.get(Asset, state.asset_id)
-            await session.delete(state)
-            await session.flush()
-
-            if asset and asset.hash is None:
-                remaining = (
-                    await session.execute(
-                        sa.select(sa.func.count())
-                        .select_from(AssetCacheState)
-                        .where(AssetCacheState.asset_id == asset.id)
-                    )
-                ).scalar_one()
-                if int(remaining or 0) == 0:
-                    await session.delete(asset)
-                    await session.flush()
-                else:
-                    await _recompute_and_apply_filename_for_asset(session, asset_id=asset.id)
-            return None
-
-        digest = await hashing_mod.blake3_hash(path)
-        new_hash = f"blake3:{digest}"
-
-        st = os.stat(path, follow_symlinks=True)
-        new_size = int(st.st_size)
-        mtime_ns = getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000))
-
-        # Current asset of this state
-        this_asset = await session.get(Asset, state.asset_id)
-
-        # If the state got orphaned somehow (race), just reattach appropriately.
-        if not this_asset:
-            canonical = (
-                await session.execute(sa.select(Asset).where(Asset.hash == new_hash).limit(1))
-            ).scalars().first()
-            if canonical:
-                state.asset_id = canonical.id
-            else:
-                now = utcnow()
-                new_asset = Asset(hash=new_hash, size_bytes=new_size, mime_type=None, created_at=now)
-                session.add(new_asset)
-                await session.flush()
-                state.asset_id = new_asset.id
-            state.mtime_ns = mtime_ns
-            state.needs_verify = False
-            with contextlib.suppress(Exception):
-                await remove_missing_tag_for_asset_id(session, asset_id=state.asset_id)
-            await session.flush()
-            return state.asset_id
-
-        # 1) Seed asset case (hash is NULL): claim or merge into canonical
-        if this_asset.hash is None:
-            canonical = (
-                await session.execute(sa.select(Asset).where(Asset.hash == new_hash).limit(1))
-            ).scalars().first()
-
-            if canonical and canonical.id != this_asset.id:
-                # Merge seed asset into canonical (safe, collision-aware)
-                await redirect_all_references_then_delete_asset(
-                    session,
-                    duplicate_asset_id=this_asset.id,
-                    canonical_asset_id=canonical.id,
-                )
-                state = await session.get(AssetCacheState, state_id)
-                if state:
-                    state.mtime_ns = mtime_ns
-                    state.needs_verify = False
-                    with contextlib.suppress(Exception):
-                        await remove_missing_tag_for_asset_id(session, asset_id=canonical.id)
-                    await _recompute_and_apply_filename_for_asset(session, asset_id=canonical.id)
-                    await session.flush()
-                return canonical.id
-
-            # No canonical: try to claim the hash; handle races with a SAVEPOINT
-            try:
-                async with session.begin_nested():
-                    this_asset.hash = new_hash
-                    if int(this_asset.size_bytes or 0) == 0 and new_size > 0:
-                        this_asset.size_bytes = new_size
-                    await session.flush()
-            except IntegrityError:
-                # Someone else claimed it concurrently; fetch canonical and merge
-                canonical = (
-                    await session.execute(sa.select(Asset).where(Asset.hash == new_hash).limit(1))
-                ).scalars().first()
-                if canonical and canonical.id != this_asset.id:
-                    await redirect_all_references_then_delete_asset(
-                        session,
-                        duplicate_asset_id=this_asset.id,
-                        canonical_asset_id=canonical.id,
-                    )
-                    state = await session.get(AssetCacheState, state_id)
-                    if state:
-                        state.mtime_ns = mtime_ns
-                        state.needs_verify = False
-                        with contextlib.suppress(Exception):
-                            await remove_missing_tag_for_asset_id(session, asset_id=canonical.id)
-                        await _recompute_and_apply_filename_for_asset(session, asset_id=canonical.id)
-                        await session.flush()
-                    return canonical.id
-                # If we got here, the integrity error was not about hash uniqueness
-                raise
-
-            # Claimed successfully
-            state.mtime_ns = mtime_ns
-            state.needs_verify = False
-            with contextlib.suppress(Exception):
-                await remove_missing_tag_for_asset_id(session, asset_id=this_asset.id)
-            await _recompute_and_apply_filename_for_asset(session, asset_id=this_asset.id)
-            await session.flush()
-            return this_asset.id
-
-        # 2) Verify case for hashed assets
-        if this_asset.hash == new_hash:
-            if int(this_asset.size_bytes or 0) == 0 and new_size > 0:
-                this_asset.size_bytes = new_size
-            state.mtime_ns = mtime_ns
-            state.needs_verify = False
-            with contextlib.suppress(Exception):
-                await remove_missing_tag_for_asset_id(session, asset_id=this_asset.id)
-            await _recompute_and_apply_filename_for_asset(session, asset_id=this_asset.id)
-            await session.flush()
-            return this_asset.id
-
-        # Content changed on this path only: retarget THIS state, do not move AssetInfo rows
-        canonical = (
-            await session.execute(sa.select(Asset).where(Asset.hash == new_hash).limit(1))
-        ).scalars().first()
-        if canonical:
-            target_id = canonical.id
-        else:
-            now = utcnow()
-            new_asset = Asset(hash=new_hash, size_bytes=new_size, mime_type=None, created_at=now)
-            session.add(new_asset)
-            await session.flush()
-            target_id = new_asset.id
-
-        state.asset_id = target_id
-        state.mtime_ns = mtime_ns
-        state.needs_verify = False
-        with contextlib.suppress(Exception):
-            await remove_missing_tag_for_asset_id(session, asset_id=target_id)
-        await _recompute_and_apply_filename_for_asset(session, asset_id=target_id)
-        await session.flush()
-        return target_id
-    except Exception:
-        raise
-
-
-async def list_unhashed_candidates_under_prefixes(session: AsyncSession, *, prefixes: list[str]) -> list[int]:
-    if not prefixes:
-        return []
-
-    conds = []
-    for p in prefixes:
-        base = os.path.abspath(p)
-        if not base.endswith(os.sep):
-            base += os.sep
-        escaped, esc = escape_like_prefix(base)
-        conds.append(AssetCacheState.file_path.like(escaped + "%", escape=esc))
-
-    path_filter = sa.or_(*conds) if len(conds) > 1 else conds[0]
-    if session.bind.dialect.name == "postgresql":
-        stmt = (
-            sa.select(AssetCacheState.id)
-            .join(Asset, Asset.id == AssetCacheState.asset_id)
-            .where(Asset.hash.is_(None), path_filter)
-            .order_by(AssetCacheState.asset_id.asc(), AssetCacheState.id.asc())
-            .distinct(AssetCacheState.asset_id)
-        )
-    else:
-        first_id = sa.func.min(AssetCacheState.id).label("first_id")
-        stmt = (
-            sa.select(first_id)
-            .join(Asset, Asset.id == AssetCacheState.asset_id)
-            .where(Asset.hash.is_(None), path_filter)
-            .group_by(AssetCacheState.asset_id)
-            .order_by(first_id.asc())
-        )
-    return [int(x) for x in (await session.execute(stmt)).scalars().all()]
-
-
-async def list_verify_candidates_under_prefixes(
-    session: AsyncSession, *, prefixes: Sequence[str]
-) -> Union[list[int], Sequence[int]]:
-    if not prefixes:
-        return []
-    conds = []
-    for p in prefixes:
-        base = os.path.abspath(p)
-        if not base.endswith(os.sep):
-            base += os.sep
-        escaped, esc = escape_like_prefix(base)
-        conds.append(AssetCacheState.file_path.like(escaped + "%", escape=esc))
-
-    return (
-        await session.execute(
-            sa.select(AssetCacheState.id)
-            .where(AssetCacheState.needs_verify.is_(True))
-            .where(sa.or_(*conds))
-            .order_by(AssetCacheState.id.asc())
-        )
-    ).scalars().all()
-
-
-async def ingest_fs_asset(
-    session: AsyncSession,
-    *,
-    asset_hash: str,
-    abs_path: str,
-    size_bytes: int,
-    mtime_ns: int,
-    mime_type: Optional[str] = None,
-    info_name: Optional[str] = None,
-    owner_id: str = "",
-    preview_id: Optional[str] = None,
-    user_metadata: Optional[dict] = None,
-    tags: Sequence[str] = (),
-    tag_origin: str = "manual",
-    require_existing_tags: bool = False,
-) -> dict:
-    """
-    Idempotently upsert:
-      - Asset by content hash (create if missing)
-      - AssetCacheState(file_path) pointing to asset_id
-      - Optionally AssetInfo + tag links and metadata projection
-    Returns flags and ids.
-    """
-    locator = os.path.abspath(abs_path)
-    now = utcnow()
-    dialect = session.bind.dialect.name
-
-    if preview_id:
-        if not await session.get(Asset, preview_id):
-            preview_id = None
-
-    out: dict[str, Any] = {
-        "asset_created": False,
-        "asset_updated": False,
-        "state_created": False,
-        "state_updated": False,
-        "asset_info_id": None,
-    }
-
-    # 1) Asset by hash
-    asset = (
-        await session.execute(select(Asset).where(Asset.hash == asset_hash).limit(1))
-    ).scalars().first()
-    if not asset:
-        vals = {
-            "hash": asset_hash,
-            "size_bytes": int(size_bytes),
-            "mime_type": mime_type,
-            "created_at": now,
-        }
-        if dialect == "sqlite":
-            res = await session.execute(
-                d_sqlite.insert(Asset)
-                .values(**vals)
-                .on_conflict_do_nothing(index_elements=[Asset.hash])
-            )
-            if int(res.rowcount or 0) > 0:
-                out["asset_created"] = True
-            asset = (
-                await session.execute(
-                    select(Asset).where(Asset.hash == asset_hash).limit(1)
-                )
-            ).scalars().first()
-        elif dialect == "postgresql":
-            res = await session.execute(
-                d_pg.insert(Asset)
-                .values(**vals)
-                .on_conflict_do_nothing(
-                    index_elements=[Asset.hash],
-                    index_where=Asset.__table__.c.hash.isnot(None),
-                )
-                .returning(Asset.id)
-            )
-            inserted_id = res.scalar_one_or_none()
-            if inserted_id:
-                out["asset_created"] = True
-                asset = await session.get(Asset, inserted_id)
-            else:
-                asset = (
-                    await session.execute(
-                        select(Asset).where(Asset.hash == asset_hash).limit(1)
-                    )
-                ).scalars().first()
-        else:
-            raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-        if not asset:
-            raise RuntimeError("Asset row not found after upsert.")
-    else:
-        changed = False
-        if asset.size_bytes != int(size_bytes) and int(size_bytes) > 0:
-            asset.size_bytes = int(size_bytes)
-            changed = True
-        if mime_type and asset.mime_type != mime_type:
-            asset.mime_type = mime_type
-            changed = True
-        if changed:
-            out["asset_updated"] = True
-
-    # 2) AssetCacheState upsert by file_path (unique)
-    vals = {
-        "asset_id": asset.id,
-        "file_path": locator,
-        "mtime_ns": int(mtime_ns),
-    }
-    if dialect == "sqlite":
-        ins = (
-            d_sqlite.insert(AssetCacheState)
-            .values(**vals)
-            .on_conflict_do_nothing(index_elements=[AssetCacheState.file_path])
-        )
-    elif dialect == "postgresql":
-        ins = (
-            d_pg.insert(AssetCacheState)
-            .values(**vals)
-            .on_conflict_do_nothing(index_elements=[AssetCacheState.file_path])
-        )
-    else:
-        raise NotImplementedError(f"Unsupported database dialect: {dialect}")
-
-    res = await session.execute(ins)
-    if int(res.rowcount or 0) > 0:
-        out["state_created"] = True
-    else:
-        upd = (
-            sa.update(AssetCacheState)
-            .where(AssetCacheState.file_path == locator)
-            .where(
-                sa.or_(
-                    AssetCacheState.asset_id != asset.id,
-                    AssetCacheState.mtime_ns.is_(None),
-                    AssetCacheState.mtime_ns != int(mtime_ns),
-                )
-            )
-            .values(asset_id=asset.id, mtime_ns=int(mtime_ns))
-        )
-        res2 = await session.execute(upd)
-        if int(res2.rowcount or 0) > 0:
-            out["state_updated"] = True
-
-    # 3) Optional AssetInfo + tags + metadata
-    if info_name:
-        try:
-            async with session.begin_nested():
-                info = AssetInfo(
-                    owner_id=owner_id,
-                    name=info_name,
-                    asset_id=asset.id,
-                    preview_id=preview_id,
-                    created_at=now,
-                    updated_at=now,
-                    last_access_time=now,
-                )
-                session.add(info)
-                await session.flush()
-                out["asset_info_id"] = info.id
-        except IntegrityError:
-            pass
-
-        existing_info = (
-            await session.execute(
-                select(AssetInfo)
-                .where(
-                    AssetInfo.asset_id == asset.id,
-                    AssetInfo.name == info_name,
-                    (AssetInfo.owner_id == owner_id),
-                )
-                .limit(1)
-            )
-        ).unique().scalar_one_or_none()
-        if not existing_info:
-            raise RuntimeError("Failed to update or insert AssetInfo.")
-
-        if preview_id and existing_info.preview_id != preview_id:
-            existing_info.preview_id = preview_id
-
-        existing_info.updated_at = now
-        if existing_info.last_access_time < now:
-            existing_info.last_access_time = now
-        await session.flush()
-        out["asset_info_id"] = existing_info.id
-
-        norm = [t.strip().lower() for t in (tags or []) if (t or "").strip()]
-        if norm and out["asset_info_id"] is not None:
-            if not require_existing_tags:
-                await ensure_tags_exist(session, norm, tag_type="user")
-
-            existing_tag_names = set(
-                name for (name,) in (await session.execute(select(Tag.name).where(Tag.name.in_(norm)))).all()
-            )
-            missing = [t for t in norm if t not in existing_tag_names]
-            if missing and require_existing_tags:
-                raise ValueError(f"Unknown tags: {missing}")
-
-            existing_links = set(
-                tag_name
-                for (tag_name,) in (
-                    await session.execute(
-                        select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == out["asset_info_id"])
-                    )
-                ).all()
-            )
-            to_add = [t for t in norm if t in existing_tag_names and t not in existing_links]
-            if to_add:
-                session.add_all(
-                    [
-                        AssetInfoTag(
-                            asset_info_id=out["asset_info_id"],
-                            tag_name=t,
-                            origin=tag_origin,
-                            added_at=now,
-                        )
-                        for t in to_add
-                    ]
-                )
-                await session.flush()
-
-        # metadata["filename"] hack
-        if out["asset_info_id"] is not None:
-            primary_path = pick_best_live_path(await list_cache_states_by_asset_id(session, asset_id=asset.id))
-            computed_filename = compute_relative_filename(primary_path) if primary_path else None
-
-            current_meta = existing_info.user_metadata or {}
-            new_meta = dict(current_meta)
-            if user_metadata is not None:
-                for k, v in user_metadata.items():
-                    new_meta[k] = v
-            if computed_filename:
-                new_meta["filename"] = computed_filename
-
-            if new_meta != current_meta:
-                await replace_asset_info_metadata_projection(
-                    session,
-                    asset_info_id=out["asset_info_id"],
-                    user_metadata=new_meta,
-                )
-
-    try:
-        await remove_missing_tag_for_asset_id(session, asset_id=asset.id)
-    except Exception:
-        logging.exception("Failed to clear 'missing' tag for asset %s", asset.id)
-    return out
-
-
-async def touch_asset_infos_by_fs_path(
-    session: AsyncSession,
-    *,
-    file_path: str,
-    ts: Optional[datetime] = None,
-    only_if_newer: bool = True,
-) -> None:
-    locator = os.path.abspath(file_path)
-    ts = ts or utcnow()
-    stmt = sa.update(AssetInfo).where(
-        sa.exists(
-            sa.select(sa.literal(1))
-            .select_from(AssetCacheState)
-            .where(
-                AssetCacheState.asset_id == AssetInfo.asset_id,
-                AssetCacheState.file_path == locator,
-            )
-        )
-    )
-    if only_if_newer:
-        stmt = stmt.where(
-            sa.or_(
-                AssetInfo.last_access_time.is_(None),
-                AssetInfo.last_access_time < ts,
-            )
-        )
-    await session.execute(stmt.values(last_access_time=ts))
-
-
-async def list_cache_states_with_asset_under_prefixes(
-    session: AsyncSession,
-    *,
-    prefixes: Sequence[str],
-) -> list[tuple[AssetCacheState, Optional[str], int]]:
-    """Return (AssetCacheState, asset_hash, size_bytes) for rows under any prefix."""
-    if not prefixes:
-        return []
-
-    conds = []
-    for p in prefixes:
-        if not p:
-            continue
-        base = os.path.abspath(p)
-        if not base.endswith(os.sep):
-            base = base + os.sep
-        escaped, esc = escape_like_prefix(base)
-        conds.append(AssetCacheState.file_path.like(escaped + "%", escape=esc))
-
-    if not conds:
-        return []
-
-    rows = (
-        await session.execute(
-            select(AssetCacheState, Asset.hash, Asset.size_bytes)
-            .join(Asset, Asset.id == AssetCacheState.asset_id)
-            .where(sa.or_(*conds))
-            .order_by(AssetCacheState.id.asc())
-        )
-    ).all()
-    return [(r[0], r[1], int(r[2] or 0)) for r in rows]
-
-
-async def _recompute_and_apply_filename_for_asset(session: AsyncSession, *, asset_id: str) -> None:
-    """Compute filename from the first *existing* cache state path and apply it to all AssetInfo (if changed)."""
-    try:
-        primary_path = pick_best_live_path(await list_cache_states_by_asset_id(session, asset_id=asset_id))
-        if not primary_path:
-            return
-        new_filename = compute_relative_filename(primary_path)
-        if not new_filename:
-            return
-        infos = (
-            await session.execute(select(AssetInfo).where(AssetInfo.asset_id == asset_id))
-        ).scalars().all()
-        for info in infos:
-            current_meta = info.user_metadata or {}
-            if current_meta.get("filename") == new_filename:
-                continue
-            updated = dict(current_meta)
-            updated["filename"] = new_filename
-            await replace_asset_info_metadata_projection(session, asset_info_id=info.id, user_metadata=updated)
-    except Exception:
-        logging.exception("Failed to recompute filename metadata for asset %s", asset_id)

app/assets/database/services/info.py

@@ -1,586 +0,0 @@
-from collections import defaultdict
-from datetime import datetime
-from typing import Any, Optional, Sequence
-
-import sqlalchemy as sa
-from sqlalchemy import delete, func, select
-from sqlalchemy.exc import IntegrityError
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import contains_eager, noload
-
-from ..._helpers import compute_relative_filename, normalize_tags
-from ..helpers import (
-    apply_metadata_filter,
-    apply_tag_filters,
-    ensure_tags_exist,
-    escape_like_prefix,
-    project_kv,
-    visible_owner_clause,
-)
-from ..models import Asset, AssetInfo, AssetInfoMeta, AssetInfoTag, Tag
-from ..timeutil import utcnow
-from .queries import (
-    get_asset_by_hash,
-    list_cache_states_by_asset_id,
-    pick_best_live_path,
-)
-
-
-async def list_asset_infos_page(
-    session: AsyncSession,
-    *,
-    owner_id: str = "",
-    include_tags: Optional[Sequence[str]] = None,
-    exclude_tags: Optional[Sequence[str]] = None,
-    name_contains: Optional[str] = None,
-    metadata_filter: Optional[dict] = None,
-    limit: int = 20,
-    offset: int = 0,
-    sort: str = "created_at",
-    order: str = "desc",
-) -> tuple[list[AssetInfo], dict[str, list[str]], int]:
-    base = (
-        select(AssetInfo)
-        .join(Asset, Asset.id == AssetInfo.asset_id)
-        .options(contains_eager(AssetInfo.asset), noload(AssetInfo.tags))
-        .where(visible_owner_clause(owner_id))
-    )
-
-    if name_contains:
-        escaped, esc = escape_like_prefix(name_contains)
-        base = base.where(AssetInfo.name.ilike(f"%{escaped}%", escape=esc))
-
-    base = apply_tag_filters(base, include_tags, exclude_tags)
-    base = apply_metadata_filter(base, metadata_filter)
-
-    sort = (sort or "created_at").lower()
-    order = (order or "desc").lower()
-    sort_map = {
-        "name": AssetInfo.name,
-        "created_at": AssetInfo.created_at,
-        "updated_at": AssetInfo.updated_at,
-        "last_access_time": AssetInfo.last_access_time,
-        "size": Asset.size_bytes,
-    }
-    sort_col = sort_map.get(sort, AssetInfo.created_at)
-    sort_exp = sort_col.desc() if order == "desc" else sort_col.asc()
-
-    base = base.order_by(sort_exp).limit(limit).offset(offset)
-
-    count_stmt = (
-        select(func.count())
-        .select_from(AssetInfo)
-        .join(Asset, Asset.id == AssetInfo.asset_id)
-        .where(visible_owner_clause(owner_id))
-    )
-    if name_contains:
-        escaped, esc = escape_like_prefix(name_contains)
-        count_stmt = count_stmt.where(AssetInfo.name.ilike(f"%{escaped}%", escape=esc))
-    count_stmt = apply_tag_filters(count_stmt, include_tags, exclude_tags)
-    count_stmt = apply_metadata_filter(count_stmt, metadata_filter)
-
-    total = int((await session.execute(count_stmt)).scalar_one() or 0)
-
-    infos = (await session.execute(base)).unique().scalars().all()
-
-    id_list: list[str] = [i.id for i in infos]
-    tag_map: dict[str, list[str]] = defaultdict(list)
-    if id_list:
-        rows = await session.execute(
-            select(AssetInfoTag.asset_info_id, Tag.name)
-            .join(Tag, Tag.name == AssetInfoTag.tag_name)
-            .where(AssetInfoTag.asset_info_id.in_(id_list))
-        )
-        for aid, tag_name in rows.all():
-            tag_map[aid].append(tag_name)
-
-    return infos, tag_map, total
-
-
-async def fetch_asset_info_and_asset(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    owner_id: str = "",
-) -> Optional[tuple[AssetInfo, Asset]]:
-    stmt = (
-        select(AssetInfo, Asset)
-        .join(Asset, Asset.id == AssetInfo.asset_id)
-        .where(
-            AssetInfo.id == asset_info_id,
-            visible_owner_clause(owner_id),
-        )
-        .limit(1)
-        .options(noload(AssetInfo.tags))
-    )
-    row = await session.execute(stmt)
-    pair = row.first()
-    if not pair:
-        return None
-    return pair[0], pair[1]
-
-
-async def fetch_asset_info_asset_and_tags(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    owner_id: str = "",
-) -> Optional[tuple[AssetInfo, Asset, list[str]]]:
-    stmt = (
-        select(AssetInfo, Asset, Tag.name)
-        .join(Asset, Asset.id == AssetInfo.asset_id)
-        .join(AssetInfoTag, AssetInfoTag.asset_info_id == AssetInfo.id, isouter=True)
-        .join(Tag, Tag.name == AssetInfoTag.tag_name, isouter=True)
-        .where(
-            AssetInfo.id == asset_info_id,
-            visible_owner_clause(owner_id),
-        )
-        .options(noload(AssetInfo.tags))
-        .order_by(Tag.name.asc())
-    )
-
-    rows = (await session.execute(stmt)).all()
-    if not rows:
-        return None
-
-    first_info, first_asset, _ = rows[0]
-    tags: list[str] = []
-    seen: set[str] = set()
-    for _info, _asset, tag_name in rows:
-        if tag_name and tag_name not in seen:
-            seen.add(tag_name)
-            tags.append(tag_name)
-    return first_info, first_asset, tags
-
-
-async def create_asset_info_for_existing_asset(
-    session: AsyncSession,
-    *,
-    asset_hash: str,
-    name: str,
-    user_metadata: Optional[dict] = None,
-    tags: Optional[Sequence[str]] = None,
-    tag_origin: str = "manual",
-    owner_id: str = "",
-) -> AssetInfo:
-    """Create or return an existing AssetInfo for an Asset identified by asset_hash."""
-    now = utcnow()
-    asset = await get_asset_by_hash(session, asset_hash=asset_hash)
-    if not asset:
-        raise ValueError(f"Unknown asset hash {asset_hash}")
-
-    info = AssetInfo(
-        owner_id=owner_id,
-        name=name,
-        asset_id=asset.id,
-        preview_id=None,
-        created_at=now,
-        updated_at=now,
-        last_access_time=now,
-    )
-    try:
-        async with session.begin_nested():
-            session.add(info)
-            await session.flush()
-    except IntegrityError:
-        existing = (
-            await session.execute(
-                select(AssetInfo)
-                .options(noload(AssetInfo.tags))
-                .where(
-                    AssetInfo.asset_id == asset.id,
-                    AssetInfo.name == name,
-                    AssetInfo.owner_id == owner_id,
-                )
-                .limit(1)
-            )
-        ).unique().scalars().first()
-        if not existing:
-            raise RuntimeError("AssetInfo upsert failed to find existing row after conflict.")
-        return existing
-
-    # metadata["filename"] hack
-    new_meta = dict(user_metadata or {})
-    computed_filename = None
-    try:
-        p = pick_best_live_path(await list_cache_states_by_asset_id(session, asset_id=asset.id))
-        if p:
-            computed_filename = compute_relative_filename(p)
-    except Exception:
-        computed_filename = None
-    if computed_filename:
-        new_meta["filename"] = computed_filename
-    if new_meta:
-        await replace_asset_info_metadata_projection(
-            session,
-            asset_info_id=info.id,
-            user_metadata=new_meta,
-        )
-
-    if tags is not None:
-        await set_asset_info_tags(
-            session,
-            asset_info_id=info.id,
-            tags=tags,
-            origin=tag_origin,
-        )
-    return info
-
-
-async def set_asset_info_tags(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    tags: Sequence[str],
-    origin: str = "manual",
-) -> dict:
-    desired = normalize_tags(tags)
-
-    current = set(
-        tag_name for (tag_name,) in (
-            await session.execute(select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == asset_info_id))
-        ).all()
-    )
-
-    to_add = [t for t in desired if t not in current]
-    to_remove = [t for t in current if t not in desired]
-
-    if to_add:
-        await ensure_tags_exist(session, to_add, tag_type="user")
-        session.add_all([
-            AssetInfoTag(asset_info_id=asset_info_id, tag_name=t, origin=origin, added_at=utcnow())
-            for t in to_add
-        ])
-        await session.flush()
-
-    if to_remove:
-        await session.execute(
-            delete(AssetInfoTag)
-            .where(AssetInfoTag.asset_info_id == asset_info_id, AssetInfoTag.tag_name.in_(to_remove))
-        )
-        await session.flush()
-
-    return {"added": to_add, "removed": to_remove, "total": desired}
-
-
-async def update_asset_info_full(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    name: Optional[str] = None,
-    tags: Optional[Sequence[str]] = None,
-    user_metadata: Optional[dict] = None,
-    tag_origin: str = "manual",
-    asset_info_row: Any = None,
-) -> AssetInfo:
-    if not asset_info_row:
-        info = await session.get(AssetInfo, asset_info_id)
-        if not info:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-    else:
-        info = asset_info_row
-
-    touched = False
-    if name is not None and name != info.name:
-        info.name = name
-        touched = True
-
-    computed_filename = None
-    try:
-        p = pick_best_live_path(await list_cache_states_by_asset_id(session, asset_id=info.asset_id))
-        if p:
-            computed_filename = compute_relative_filename(p)
-    except Exception:
-        computed_filename = None
-
-    if user_metadata is not None:
-        new_meta = dict(user_metadata)
-        if computed_filename:
-            new_meta["filename"] = computed_filename
-        await replace_asset_info_metadata_projection(
-            session, asset_info_id=asset_info_id, user_metadata=new_meta
-        )
-        touched = True
-    else:
-        if computed_filename:
-            current_meta = info.user_metadata or {}
-            if current_meta.get("filename") != computed_filename:
-                new_meta = dict(current_meta)
-                new_meta["filename"] = computed_filename
-                await replace_asset_info_metadata_projection(
-                    session, asset_info_id=asset_info_id, user_metadata=new_meta
-                )
-                touched = True
-
-    if tags is not None:
-        await set_asset_info_tags(
-            session,
-            asset_info_id=asset_info_id,
-            tags=tags,
-            origin=tag_origin,
-        )
-        touched = True
-
-    if touched and user_metadata is None:
-        info.updated_at = utcnow()
-        await session.flush()
-
-    return info
-
-
-async def replace_asset_info_metadata_projection(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    user_metadata: Optional[dict],
-) -> None:
-    info = await session.get(AssetInfo, asset_info_id)
-    if not info:
-        raise ValueError(f"AssetInfo {asset_info_id} not found")
-
-    info.user_metadata = user_metadata or {}
-    info.updated_at = utcnow()
-    await session.flush()
-
-    await session.execute(delete(AssetInfoMeta).where(AssetInfoMeta.asset_info_id == asset_info_id))
-    await session.flush()
-
-    if not user_metadata:
-        return
-
-    rows: list[AssetInfoMeta] = []
-    for k, v in user_metadata.items():
-        for r in project_kv(k, v):
-            rows.append(
-                AssetInfoMeta(
-                    asset_info_id=asset_info_id,
-                    key=r["key"],
-                    ordinal=int(r["ordinal"]),
-                    val_str=r.get("val_str"),
-                    val_num=r.get("val_num"),
-                    val_bool=r.get("val_bool"),
-                    val_json=r.get("val_json"),
-                )
-            )
-    if rows:
-        session.add_all(rows)
-        await session.flush()
-
-
-async def touch_asset_info_by_id(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    ts: Optional[datetime] = None,
-    only_if_newer: bool = True,
-) -> None:
-    ts = ts or utcnow()
-    stmt = sa.update(AssetInfo).where(AssetInfo.id == asset_info_id)
-    if only_if_newer:
-        stmt = stmt.where(
-            sa.or_(AssetInfo.last_access_time.is_(None), AssetInfo.last_access_time < ts)
-        )
-    await session.execute(stmt.values(last_access_time=ts))
-
-
-async def delete_asset_info_by_id(session: AsyncSession, *, asset_info_id: str, owner_id: str) -> bool:
-    stmt = sa.delete(AssetInfo).where(
-        AssetInfo.id == asset_info_id,
-        visible_owner_clause(owner_id),
-    )
-    return int((await session.execute(stmt)).rowcount or 0) > 0
-
-
-async def add_tags_to_asset_info(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    tags: Sequence[str],
-    origin: str = "manual",
-    create_if_missing: bool = True,
-    asset_info_row: Any = None,
-) -> dict:
-    if not asset_info_row:
-        info = await session.get(AssetInfo, asset_info_id)
-        if not info:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-
-    norm = normalize_tags(tags)
-    if not norm:
-        total = await get_asset_tags(session, asset_info_id=asset_info_id)
-        return {"added": [], "already_present": [], "total_tags": total}
-
-    if create_if_missing:
-        await ensure_tags_exist(session, norm, tag_type="user")
-
-    current = {
-        tag_name
-        for (tag_name,) in (
-            await session.execute(
-                sa.select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == asset_info_id)
-            )
-        ).all()
-    }
-
-    want = set(norm)
-    to_add = sorted(want - current)
-
-    if to_add:
-        async with session.begin_nested() as nested:
-            try:
-                session.add_all(
-                    [
-                        AssetInfoTag(
-                            asset_info_id=asset_info_id,
-                            tag_name=t,
-                            origin=origin,
-                            added_at=utcnow(),
-                        )
-                        for t in to_add
-                    ]
-                )
-                await session.flush()
-            except IntegrityError:
-                await nested.rollback()
-
-    after = set(await get_asset_tags(session, asset_info_id=asset_info_id))
-    return {
-        "added": sorted(((after - current) & want)),
-        "already_present": sorted(want & current),
-        "total_tags": sorted(after),
-    }
-
-
-async def remove_tags_from_asset_info(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    tags: Sequence[str],
-) -> dict:
-    info = await session.get(AssetInfo, asset_info_id)
-    if not info:
-        raise ValueError(f"AssetInfo {asset_info_id} not found")
-
-    norm = normalize_tags(tags)
-    if not norm:
-        total = await get_asset_tags(session, asset_info_id=asset_info_id)
-        return {"removed": [], "not_present": [], "total_tags": total}
-
-    existing = {
-        tag_name
-        for (tag_name,) in (
-            await session.execute(
-                sa.select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == asset_info_id)
-            )
-        ).all()
-    }
-
-    to_remove = sorted(set(t for t in norm if t in existing))
-    not_present = sorted(set(t for t in norm if t not in existing))
-
-    if to_remove:
-        await session.execute(
-            delete(AssetInfoTag)
-            .where(
-                AssetInfoTag.asset_info_id == asset_info_id,
-                AssetInfoTag.tag_name.in_(to_remove),
-            )
-        )
-        await session.flush()
-
-    total = await get_asset_tags(session, asset_info_id=asset_info_id)
-    return {"removed": to_remove, "not_present": not_present, "total_tags": total}
-
-
-async def list_tags_with_usage(
-    session: AsyncSession,
-    *,
-    prefix: Optional[str] = None,
-    limit: int = 100,
-    offset: int = 0,
-    include_zero: bool = True,
-    order: str = "count_desc",
-    owner_id: str = "",
-) -> tuple[list[tuple[str, str, int]], int]:
-    counts_sq = (
-        select(
-            AssetInfoTag.tag_name.label("tag_name"),
-            func.count(AssetInfoTag.asset_info_id).label("cnt"),
-        )
-        .select_from(AssetInfoTag)
-        .join(AssetInfo, AssetInfo.id == AssetInfoTag.asset_info_id)
-        .where(visible_owner_clause(owner_id))
-        .group_by(AssetInfoTag.tag_name)
-        .subquery()
-    )
-
-    q = (
-        select(
-            Tag.name,
-            Tag.tag_type,
-            func.coalesce(counts_sq.c.cnt, 0).label("count"),
-        )
-        .select_from(Tag)
-        .join(counts_sq, counts_sq.c.tag_name == Tag.name, isouter=True)
-    )
-
-    if prefix:
-        escaped, esc = escape_like_prefix(prefix.strip().lower())
-        q = q.where(Tag.name.like(escaped + "%", escape=esc))
-
-    if not include_zero:
-        q = q.where(func.coalesce(counts_sq.c.cnt, 0) > 0)
-
-    if order == "name_asc":
-        q = q.order_by(Tag.name.asc())
-    else:
-        q = q.order_by(func.coalesce(counts_sq.c.cnt, 0).desc(), Tag.name.asc())
-
-    total_q = select(func.count()).select_from(Tag)
-    if prefix:
-        escaped, esc = escape_like_prefix(prefix.strip().lower())
-        total_q = total_q.where(Tag.name.like(escaped + "%", escape=esc))
-    if not include_zero:
-        total_q = total_q.where(
-            Tag.name.in_(select(AssetInfoTag.tag_name).group_by(AssetInfoTag.tag_name))
-        )
-
-    rows = (await session.execute(q.limit(limit).offset(offset))).all()
-    total = (await session.execute(total_q)).scalar_one()
-
-    rows_norm = [(name, ttype, int(count or 0)) for (name, ttype, count) in rows]
-    return rows_norm, int(total or 0)
-
-
-async def get_asset_tags(session: AsyncSession, *, asset_info_id: str) -> list[str]:
-    return [
-        tag_name
-        for (tag_name,) in (
-            await session.execute(
-                sa.select(AssetInfoTag.tag_name).where(AssetInfoTag.asset_info_id == asset_info_id)
-            )
-        ).all()
-    ]
-
-
-async def set_asset_info_preview(
-    session: AsyncSession,
-    *,
-    asset_info_id: str,
-    preview_asset_id: Optional[str],
-) -> None:
-    """Set or clear preview_id and bump updated_at. Raises on unknown IDs."""
-    info = await session.get(AssetInfo, asset_info_id)
-    if not info:
-        raise ValueError(f"AssetInfo {asset_info_id} not found")
-
-    if preview_asset_id is None:
-        info.preview_id = None
-    else:
-        # validate preview asset exists
-        if not await session.get(Asset, preview_asset_id):
-            raise ValueError(f"Preview Asset {preview_asset_id} not found")
-        info.preview_id = preview_asset_id
-
-    info.updated_at = utcnow()
-    await session.flush()

app/assets/database/services/queries.py

@@ -1,76 +0,0 @@
-import os
-from typing import Optional, Sequence, Union
-
-import sqlalchemy as sa
-from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from ..models import Asset, AssetCacheState, AssetInfo
-
-
-async def asset_exists_by_hash(session: AsyncSession, *, asset_hash: str) -> bool:
-    row = (
-        await session.execute(
-            select(sa.literal(True)).select_from(Asset).where(Asset.hash == asset_hash).limit(1)
-        )
-    ).first()
-    return row is not None
-
-
-async def get_asset_by_hash(session: AsyncSession, *, asset_hash: str) -> Optional[Asset]:
-    return (
-        await session.execute(select(Asset).where(Asset.hash == asset_hash).limit(1))
-    ).scalars().first()
-
-
-async def get_asset_info_by_id(session: AsyncSession, *, asset_info_id: str) -> Optional[AssetInfo]:
-    return await session.get(AssetInfo, asset_info_id)
-
-
-async def asset_info_exists_for_asset_id(session: AsyncSession, *, asset_id: str) -> bool:
-    q = (
-        select(sa.literal(True))
-        .select_from(AssetInfo)
-        .where(AssetInfo.asset_id == asset_id)
-        .limit(1)
-    )
-    return (await session.execute(q)).first() is not None
-
-
-async def get_cache_state_by_asset_id(session: AsyncSession, *, asset_id: str) -> Optional[AssetCacheState]:
-    return (
-        await session.execute(
-            select(AssetCacheState)
-            .where(AssetCacheState.asset_id == asset_id)
-            .order_by(AssetCacheState.id.asc())
-            .limit(1)
-        )
-    ).scalars().first()
-
-
-async def list_cache_states_by_asset_id(
-    session: AsyncSession, *, asset_id: str
-) -> Union[list[AssetCacheState], Sequence[AssetCacheState]]:
-    return (
-        await session.execute(
-            select(AssetCacheState)
-            .where(AssetCacheState.asset_id == asset_id)
-            .order_by(AssetCacheState.id.asc())
-        )
-    ).scalars().all()
-
-
-def pick_best_live_path(states: Union[list[AssetCacheState], Sequence[AssetCacheState]]) -> str:
-    """
-    Return the best on-disk path among cache states:
-      1) Prefer a path that exists with needs_verify == False (already verified).
-      2) Otherwise, pick the first path that exists.
-      3) Otherwise return empty string.
-    """
-    alive = [s for s in states if getattr(s, "file_path", None) and os.path.isfile(s.file_path)]
-    if not alive:
-        return ""
-    for s in alive:
-        if not getattr(s, "needs_verify", False):
-            return s.file_path
-    return alive[0].file_path

app/assets/database/timeutil.py

@@ -1,6 +0,0 @@
-from datetime import datetime, timezone
-
-
-def utcnow() -> datetime:
-    """Naive UTC timestamp (no tzinfo). We always treat DB datetimes as UTC."""
-    return datetime.now(timezone.utc).replace(tzinfo=None)

app/assets/manager.py

@@ -1,556 +0,0 @@
-import contextlib
-import logging
-import mimetypes
-import os
-from typing import Optional, Sequence
-
-from comfy_api.internal import async_to_sync
-
-from ..db import create_session
-from ._helpers import (
-    ensure_within_base,
-    get_name_and_tags_from_asset_path,
-    resolve_destination_from_tags,
-)
-from .api import schemas_in, schemas_out
-from .database.models import Asset
-from .database.services import (
-    add_tags_to_asset_info,
-    asset_exists_by_hash,
-    asset_info_exists_for_asset_id,
-    check_fs_asset_exists_quick,
-    create_asset_info_for_existing_asset,
-    delete_asset_info_by_id,
-    fetch_asset_info_and_asset,
-    fetch_asset_info_asset_and_tags,
-    get_asset_by_hash,
-    get_asset_info_by_id,
-    get_asset_tags,
-    ingest_fs_asset,
-    list_asset_infos_page,
-    list_cache_states_by_asset_id,
-    list_tags_with_usage,
-    pick_best_live_path,
-    remove_tags_from_asset_info,
-    set_asset_info_preview,
-    touch_asset_info_by_id,
-    touch_asset_infos_by_fs_path,
-    update_asset_info_full,
-)
-from .storage import hashing
-
-
-async def asset_exists(*, asset_hash: str) -> bool:
-    async with await create_session() as session:
-        return await asset_exists_by_hash(session, asset_hash=asset_hash)
-
-
-def populate_db_with_asset(file_path: str, tags: Optional[list[str]] = None) -> None:
-    if tags is None:
-        tags = []
-    try:
-        asset_name, path_tags = get_name_and_tags_from_asset_path(file_path)
-        async_to_sync.AsyncToSyncConverter.run_async_in_thread(
-            add_local_asset,
-            tags=list(dict.fromkeys([*path_tags, *tags])),
-            file_name=asset_name,
-            file_path=file_path,
-        )
-    except ValueError as e:
-        logging.warning("Skipping non-asset path %s: %s", file_path, e)
-
-
-async def add_local_asset(tags: list[str], file_name: str, file_path: str) -> None:
-    abs_path = os.path.abspath(file_path)
-    size_bytes, mtime_ns = _get_size_mtime_ns(abs_path)
-    if not size_bytes:
-        return
-
-    async with await create_session() as session:
-        if await check_fs_asset_exists_quick(session, file_path=abs_path, size_bytes=size_bytes, mtime_ns=mtime_ns):
-            await touch_asset_infos_by_fs_path(session, file_path=abs_path)
-            await session.commit()
-            return
-
-    asset_hash = hashing.blake3_hash_sync(abs_path)
-
-    async with await create_session() as session:
-        await ingest_fs_asset(
-            session,
-            asset_hash="blake3:" + asset_hash,
-            abs_path=abs_path,
-            size_bytes=size_bytes,
-            mtime_ns=mtime_ns,
-            mime_type=None,
-            info_name=file_name,
-            tag_origin="automatic",
-            tags=tags,
-        )
-        await session.commit()
-
-
-async def list_assets(
-    *,
-    include_tags: Optional[Sequence[str]] = None,
-    exclude_tags: Optional[Sequence[str]] = None,
-    name_contains: Optional[str] = None,
-    metadata_filter: Optional[dict] = None,
-    limit: int = 20,
-    offset: int = 0,
-    sort: str = "created_at",
-    order: str = "desc",
-    owner_id: str = "",
-) -> schemas_out.AssetsList:
-    sort = _safe_sort_field(sort)
-    order = "desc" if (order or "desc").lower() not in {"asc", "desc"} else order.lower()
-
-    async with await create_session() as session:
-        infos, tag_map, total = await list_asset_infos_page(
-            session,
-            owner_id=owner_id,
-            include_tags=include_tags,
-            exclude_tags=exclude_tags,
-            name_contains=name_contains,
-            metadata_filter=metadata_filter,
-            limit=limit,
-            offset=offset,
-            sort=sort,
-            order=order,
-        )
-
-    summaries: list[schemas_out.AssetSummary] = []
-    for info in infos:
-        asset = info.asset
-        tags = tag_map.get(info.id, [])
-        summaries.append(
-            schemas_out.AssetSummary(
-                id=info.id,
-                name=info.name,
-                asset_hash=asset.hash if asset else None,
-                size=int(asset.size_bytes) if asset else None,
-                mime_type=asset.mime_type if asset else None,
-                tags=tags,
-                preview_url=f"/api/assets/{info.id}/content",
-                created_at=info.created_at,
-                updated_at=info.updated_at,
-                last_access_time=info.last_access_time,
-            )
-        )
-
-    return schemas_out.AssetsList(
-        assets=summaries,
-        total=total,
-        has_more=(offset + len(summaries)) < total,
-    )
-
-
-async def get_asset(*, asset_info_id: str, owner_id: str = "") -> schemas_out.AssetDetail:
-    async with await create_session() as session:
-        res = await fetch_asset_info_asset_and_tags(session, asset_info_id=asset_info_id, owner_id=owner_id)
-        if not res:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-        info, asset, tag_names = res
-        preview_id = info.preview_id
-
-    return schemas_out.AssetDetail(
-        id=info.id,
-        name=info.name,
-        asset_hash=asset.hash if asset else None,
-        size=int(asset.size_bytes) if asset and asset.size_bytes is not None else None,
-        mime_type=asset.mime_type if asset else None,
-        tags=tag_names,
-        user_metadata=info.user_metadata or {},
-        preview_id=preview_id,
-        created_at=info.created_at,
-        last_access_time=info.last_access_time,
-    )
-
-
-async def resolve_asset_content_for_download(
-    *,
-    asset_info_id: str,
-    owner_id: str = "",
-) -> tuple[str, str, str]:
-    async with await create_session() as session:
-        pair = await fetch_asset_info_and_asset(session, asset_info_id=asset_info_id, owner_id=owner_id)
-        if not pair:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-
-        info, asset = pair
-        states = await list_cache_states_by_asset_id(session, asset_id=asset.id)
-        abs_path = pick_best_live_path(states)
-        if not abs_path:
-            raise FileNotFoundError
-
-        await touch_asset_info_by_id(session, asset_info_id=asset_info_id)
-        await session.commit()
-
-    ctype = asset.mime_type or mimetypes.guess_type(info.name or abs_path)[0] or "application/octet-stream"
-    download_name = info.name or os.path.basename(abs_path)
-    return abs_path, ctype, download_name
-
-
-async def upload_asset_from_temp_path(
-    spec: schemas_in.UploadAssetSpec,
-    *,
-    temp_path: str,
-    client_filename: Optional[str] = None,
-    owner_id: str = "",
-    expected_asset_hash: Optional[str] = None,
-) -> schemas_out.AssetCreated:
-    try:
-        digest = await hashing.blake3_hash(temp_path)
-    except Exception as e:
-        raise RuntimeError(f"failed to hash uploaded file: {e}")
-    asset_hash = "blake3:" + digest
-
-    if expected_asset_hash and asset_hash != expected_asset_hash.strip().lower():
-        raise ValueError("HASH_MISMATCH")
-
-    async with await create_session() as session:
-        existing = await get_asset_by_hash(session, asset_hash=asset_hash)
-        if existing is not None:
-            with contextlib.suppress(Exception):
-                if temp_path and os.path.exists(temp_path):
-                    os.remove(temp_path)
-
-            display_name = _safe_filename(spec.name or (client_filename or ""), fallback=digest)
-            info = await create_asset_info_for_existing_asset(
-                session,
-                asset_hash=asset_hash,
-                name=display_name,
-                user_metadata=spec.user_metadata or {},
-                tags=spec.tags or [],
-                tag_origin="manual",
-                owner_id=owner_id,
-            )
-            tag_names = await get_asset_tags(session, asset_info_id=info.id)
-            await session.commit()
-
-            return schemas_out.AssetCreated(
-                id=info.id,
-                name=info.name,
-                asset_hash=existing.hash,
-                size=int(existing.size_bytes) if existing.size_bytes is not None else None,
-                mime_type=existing.mime_type,
-                tags=tag_names,
-                user_metadata=info.user_metadata or {},
-                preview_id=info.preview_id,
-                created_at=info.created_at,
-                last_access_time=info.last_access_time,
-                created_new=False,
-            )
-
-    base_dir, subdirs = resolve_destination_from_tags(spec.tags)
-    dest_dir = os.path.join(base_dir, *subdirs) if subdirs else base_dir
-    os.makedirs(dest_dir, exist_ok=True)
-
-    src_for_ext = (client_filename or spec.name or "").strip()
-    _ext = os.path.splitext(os.path.basename(src_for_ext))[1] if src_for_ext else ""
-    ext = _ext if 0 < len(_ext) <= 16 else ""
-    hashed_basename = f"{digest}{ext}"
-    dest_abs = os.path.abspath(os.path.join(dest_dir, hashed_basename))
-    ensure_within_base(dest_abs, base_dir)
-
-    content_type = (
-        mimetypes.guess_type(os.path.basename(src_for_ext), strict=False)[0]
-        or mimetypes.guess_type(hashed_basename, strict=False)[0]
-        or "application/octet-stream"
-    )
-
-    try:
-        os.replace(temp_path, dest_abs)
-    except Exception as e:
-        raise RuntimeError(f"failed to move uploaded file into place: {e}")
-
-    try:
-        size_bytes, mtime_ns = _get_size_mtime_ns(dest_abs)
-    except OSError as e:
-        raise RuntimeError(f"failed to stat destination file: {e}")
-
-    async with await create_session() as session:
-        result = await ingest_fs_asset(
-            session,
-            asset_hash=asset_hash,
-            abs_path=dest_abs,
-            size_bytes=size_bytes,
-            mtime_ns=mtime_ns,
-            mime_type=content_type,
-            info_name=_safe_filename(spec.name or (client_filename or ""), fallback=digest),
-            owner_id=owner_id,
-            preview_id=None,
-            user_metadata=spec.user_metadata or {},
-            tags=spec.tags,
-            tag_origin="manual",
-            require_existing_tags=False,
-        )
-        info_id = result["asset_info_id"]
-        if not info_id:
-            raise RuntimeError("failed to create asset metadata")
-
-        pair = await fetch_asset_info_and_asset(session, asset_info_id=info_id, owner_id=owner_id)
-        if not pair:
-            raise RuntimeError("inconsistent DB state after ingest")
-        info, asset = pair
-        tag_names = await get_asset_tags(session, asset_info_id=info.id)
-        await session.commit()
-
-    return schemas_out.AssetCreated(
-        id=info.id,
-        name=info.name,
-        asset_hash=asset.hash,
-        size=int(asset.size_bytes),
-        mime_type=asset.mime_type,
-        tags=tag_names,
-        user_metadata=info.user_metadata or {},
-        preview_id=info.preview_id,
-        created_at=info.created_at,
-        last_access_time=info.last_access_time,
-        created_new=result["asset_created"],
-    )
-
-
-async def update_asset(
-    *,
-    asset_info_id: str,
-    name: Optional[str] = None,
-    tags: Optional[list[str]] = None,
-    user_metadata: Optional[dict] = None,
-    owner_id: str = "",
-) -> schemas_out.AssetUpdated:
-    async with await create_session() as session:
-        info_row = await get_asset_info_by_id(session, asset_info_id=asset_info_id)
-        if not info_row:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-        if info_row.owner_id and info_row.owner_id != owner_id:
-            raise PermissionError("not owner")
-
-        info = await update_asset_info_full(
-            session,
-            asset_info_id=asset_info_id,
-            name=name,
-            tags=tags,
-            user_metadata=user_metadata,
-            tag_origin="manual",
-            asset_info_row=info_row,
-        )
-
-        tag_names = await get_asset_tags(session, asset_info_id=asset_info_id)
-        await session.commit()
-
-    return schemas_out.AssetUpdated(
-        id=info.id,
-        name=info.name,
-        asset_hash=info.asset.hash if info.asset else None,
-        tags=tag_names,
-        user_metadata=info.user_metadata or {},
-        updated_at=info.updated_at,
-    )
-
-
-async def set_asset_preview(
-    *,
-    asset_info_id: str,
-    preview_asset_id: Optional[str],
-    owner_id: str = "",
-) -> schemas_out.AssetDetail:
-    async with await create_session() as session:
-        info_row = await get_asset_info_by_id(session, asset_info_id=asset_info_id)
-        if not info_row:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-        if info_row.owner_id and info_row.owner_id != owner_id:
-            raise PermissionError("not owner")
-
-        await set_asset_info_preview(
-            session,
-            asset_info_id=asset_info_id,
-            preview_asset_id=preview_asset_id,
-        )
-
-        res = await fetch_asset_info_asset_and_tags(session, asset_info_id=asset_info_id, owner_id=owner_id)
-        if not res:
-            raise RuntimeError("State changed during preview update")
-        info, asset, tags = res
-        await session.commit()
-
-    return schemas_out.AssetDetail(
-        id=info.id,
-        name=info.name,
-        asset_hash=asset.hash if asset else None,
-        size=int(asset.size_bytes) if asset and asset.size_bytes is not None else None,
-        mime_type=asset.mime_type if asset else None,
-        tags=tags,
-        user_metadata=info.user_metadata or {},
-        preview_id=info.preview_id,
-        created_at=info.created_at,
-        last_access_time=info.last_access_time,
-    )
-
-
-async def delete_asset_reference(*, asset_info_id: str, owner_id: str, delete_content_if_orphan: bool = True) -> bool:
-    async with await create_session() as session:
-        info_row = await get_asset_info_by_id(session, asset_info_id=asset_info_id)
-        asset_id = info_row.asset_id if info_row else None
-        deleted = await delete_asset_info_by_id(session, asset_info_id=asset_info_id, owner_id=owner_id)
-        if not deleted:
-            await session.commit()
-            return False
-
-        if not delete_content_if_orphan or not asset_id:
-            await session.commit()
-            return True
-
-        still_exists = await asset_info_exists_for_asset_id(session, asset_id=asset_id)
-        if still_exists:
-            await session.commit()
-            return True
-
-        states = await list_cache_states_by_asset_id(session, asset_id=asset_id)
-        file_paths = [s.file_path for s in (states or []) if getattr(s, "file_path", None)]
-
-        asset_row = await session.get(Asset, asset_id)
-        if asset_row is not None:
-            await session.delete(asset_row)
-
-        await session.commit()
-        for p in file_paths:
-            with contextlib.suppress(Exception):
-                if p and os.path.isfile(p):
-                    os.remove(p)
-    return True
-
-
-async def create_asset_from_hash(
-    *,
-    hash_str: str,
-    name: str,
-    tags: Optional[list[str]] = None,
-    user_metadata: Optional[dict] = None,
-    owner_id: str = "",
-) -> Optional[schemas_out.AssetCreated]:
-    canonical = hash_str.strip().lower()
-    async with await create_session() as session:
-        asset = await get_asset_by_hash(session, asset_hash=canonical)
-        if not asset:
-            return None
-
-        info = await create_asset_info_for_existing_asset(
-            session,
-            asset_hash=canonical,
-            name=_safe_filename(name, fallback=canonical.split(":", 1)[1]),
-            user_metadata=user_metadata or {},
-            tags=tags or [],
-            tag_origin="manual",
-            owner_id=owner_id,
-        )
-        tag_names = await get_asset_tags(session, asset_info_id=info.id)
-        await session.commit()
-
-    return schemas_out.AssetCreated(
-        id=info.id,
-        name=info.name,
-        asset_hash=asset.hash,
-        size=int(asset.size_bytes),
-        mime_type=asset.mime_type,
-        tags=tag_names,
-        user_metadata=info.user_metadata or {},
-        preview_id=info.preview_id,
-        created_at=info.created_at,
-        last_access_time=info.last_access_time,
-        created_new=False,
-    )
-
-
-async def list_tags(
-    *,
-    prefix: Optional[str] = None,
-    limit: int = 100,
-    offset: int = 0,
-    order: str = "count_desc",
-    include_zero: bool = True,
-    owner_id: str = "",
-) -> schemas_out.TagsList:
-    limit = max(1, min(1000, limit))
-    offset = max(0, offset)
-
-    async with await create_session() as session:
-        rows, total = await list_tags_with_usage(
-            session,
-            prefix=prefix,
-            limit=limit,
-            offset=offset,
-            include_zero=include_zero,
-            order=order,
-            owner_id=owner_id,
-        )
-
-    tags = [schemas_out.TagUsage(name=name, count=count, type=tag_type) for (name, tag_type, count) in rows]
-    return schemas_out.TagsList(tags=tags, total=total, has_more=(offset + len(tags)) < total)
-
-
-async def add_tags_to_asset(
-    *,
-    asset_info_id: str,
-    tags: list[str],
-    origin: str = "manual",
-    owner_id: str = "",
-) -> schemas_out.TagsAdd:
-    async with await create_session() as session:
-        info_row = await get_asset_info_by_id(session, asset_info_id=asset_info_id)
-        if not info_row:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-        if info_row.owner_id and info_row.owner_id != owner_id:
-            raise PermissionError("not owner")
-        data = await add_tags_to_asset_info(
-            session,
-            asset_info_id=asset_info_id,
-            tags=tags,
-            origin=origin,
-            create_if_missing=True,
-            asset_info_row=info_row,
-        )
-        await session.commit()
-    return schemas_out.TagsAdd(**data)
-
-
-async def remove_tags_from_asset(
-    *,
-    asset_info_id: str,
-    tags: list[str],
-    owner_id: str = "",
-) -> schemas_out.TagsRemove:
-    async with await create_session() as session:
-        info_row = await get_asset_info_by_id(session, asset_info_id=asset_info_id)
-        if not info_row:
-            raise ValueError(f"AssetInfo {asset_info_id} not found")
-        if info_row.owner_id and info_row.owner_id != owner_id:
-            raise PermissionError("not owner")
-
-        data = await remove_tags_from_asset_info(
-            session,
-            asset_info_id=asset_info_id,
-            tags=tags,
-        )
-        await session.commit()
-    return schemas_out.TagsRemove(**data)
-
-
-def _safe_sort_field(requested: Optional[str]) -> str:
-    if not requested:
-        return "created_at"
-    v = requested.lower()
-    if v in {"name", "created_at", "updated_at", "size", "last_access_time"}:
-        return v
-    return "created_at"
-
-
-def _get_size_mtime_ns(path: str) -> tuple[int, int]:
-    st = os.stat(path, follow_symlinks=True)
-    return st.st_size, getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000))
-
-
-def _safe_filename(name: Optional[str], fallback: str) -> str:
-    n = os.path.basename((name or "").strip() or fallback)
-    if n:
-        return n
-    return fallback

app/assets/scanner.py

@@ -1,501 +0,0 @@
-import asyncio
-import contextlib
-import logging
-import os
-import time
-from dataclasses import dataclass, field
-from typing import Literal, Optional
-
-import sqlalchemy as sa
-
-import folder_paths
-
-from ..db import create_session
-from ._helpers import (
-    collect_models_files,
-    compute_relative_filename,
-    get_comfy_models_folders,
-    get_name_and_tags_from_asset_path,
-    list_tree,
-    new_scan_id,
-    prefixes_for_root,
-    ts_to_iso,
-)
-from .api import schemas_in, schemas_out
-from .database.helpers import (
-    add_missing_tag_for_asset_id,
-    ensure_tags_exist,
-    escape_like_prefix,
-    fast_asset_file_check,
-    remove_missing_tag_for_asset_id,
-    seed_from_paths_batch,
-)
-from .database.models import Asset, AssetCacheState, AssetInfo
-from .database.services import (
-    compute_hash_and_dedup_for_cache_state,
-    list_cache_states_by_asset_id,
-    list_cache_states_with_asset_under_prefixes,
-    list_unhashed_candidates_under_prefixes,
-    list_verify_candidates_under_prefixes,
-)
-
-LOGGER = logging.getLogger(__name__)
-
-SLOW_HASH_CONCURRENCY = 1
-
-
-@dataclass
-class ScanProgress:
-    scan_id: str
-    root: schemas_in.RootType
-    status: Literal["scheduled", "running", "completed", "failed", "cancelled"] = "scheduled"
-    scheduled_at: float = field(default_factory=lambda: time.time())
-    started_at: Optional[float] = None
-    finished_at: Optional[float] = None
-    discovered: int = 0
-    processed: int = 0
-    file_errors: list[dict] = field(default_factory=list)
-
-
-@dataclass
-class SlowQueueState:
-    queue: asyncio.Queue
-    workers: list[asyncio.Task] = field(default_factory=list)
-    closed: bool = False
-
-
-RUNNING_TASKS: dict[schemas_in.RootType, asyncio.Task] = {}
-PROGRESS_BY_ROOT: dict[schemas_in.RootType, ScanProgress] = {}
-SLOW_STATE_BY_ROOT: dict[schemas_in.RootType, SlowQueueState] = {}
-
-
-def current_statuses() -> schemas_out.AssetScanStatusResponse:
-    scans = []
-    for root in schemas_in.ALLOWED_ROOTS:
-        prog = PROGRESS_BY_ROOT.get(root)
-        if not prog:
-            continue
-        scans.append(_scan_progress_to_scan_status_model(prog))
-    return schemas_out.AssetScanStatusResponse(scans=scans)
-
-
-async def schedule_scans(roots: list[schemas_in.RootType]) -> schemas_out.AssetScanStatusResponse:
-    results: list[ScanProgress] = []
-    for root in roots:
-        if root in RUNNING_TASKS and not RUNNING_TASKS[root].done():
-            results.append(PROGRESS_BY_ROOT[root])
-            continue
-
-        prog = ScanProgress(scan_id=new_scan_id(root), root=root, status="scheduled")
-        PROGRESS_BY_ROOT[root] = prog
-        state = SlowQueueState(queue=asyncio.Queue())
-        SLOW_STATE_BY_ROOT[root] = state
-        RUNNING_TASKS[root] = asyncio.create_task(
-            _run_hash_verify_pipeline(root, prog, state),
-            name=f"asset-scan:{root}",
-        )
-        results.append(prog)
-    return _status_response_for(results)
-
-
-async def sync_seed_assets(roots: list[schemas_in.RootType]) -> None:
-    t_total = time.perf_counter()
-    created = 0
-    skipped_existing = 0
-    paths: list[str] = []
-    try:
-        existing_paths: set[str] = set()
-        for r in roots:
-            try:
-                survivors = await _fast_db_consistency_pass(r, collect_existing_paths=True, update_missing_tags=True)
-                if survivors:
-                    existing_paths.update(survivors)
-            except Exception as ex:
-                LOGGER.exception("fast DB reconciliation failed for %s: %s", r, ex)
-
-        if "models" in roots:
-            paths.extend(collect_models_files())
-        if "input" in roots:
-            paths.extend(list_tree(folder_paths.get_input_directory()))
-        if "output" in roots:
-            paths.extend(list_tree(folder_paths.get_output_directory()))
-
-        specs: list[dict] = []
-        tag_pool: set[str] = set()
-        for p in paths:
-            ap = os.path.abspath(p)
-            if ap in existing_paths:
-                skipped_existing += 1
-                continue
-            try:
-                st = os.stat(ap, follow_symlinks=True)
-            except OSError:
-                continue
-            if not st.st_size:
-                continue
-            name, tags = get_name_and_tags_from_asset_path(ap)
-            specs.append(
-                {
-                    "abs_path": ap,
-                    "size_bytes": st.st_size,
-                    "mtime_ns": getattr(st, "st_mtime_ns", int(st.st_mtime * 1_000_000_000)),
-                    "info_name": name,
-                    "tags": tags,
-                    "fname": compute_relative_filename(ap),
-                }
-            )
-            for t in tags:
-                tag_pool.add(t)
-
-        if not specs:
-            return
-        async with await create_session() as sess:
-            if tag_pool:
-                await ensure_tags_exist(sess, tag_pool, tag_type="user")
-
-            result = await seed_from_paths_batch(sess, specs=specs, owner_id="")
-            created += result["inserted_infos"]
-            await sess.commit()
-    finally:
-        LOGGER.info(
-            "Assets scan(roots=%s) completed in %.3fs (created=%d, skipped_existing=%d, total_seen=%d)",
-            roots,
-            time.perf_counter() - t_total,
-            created,
-            skipped_existing,
-            len(paths),
-        )
-
-
-def _status_response_for(progresses: list[ScanProgress]) -> schemas_out.AssetScanStatusResponse:
-    return schemas_out.AssetScanStatusResponse(scans=[_scan_progress_to_scan_status_model(p) for p in progresses])
-
-
-def _scan_progress_to_scan_status_model(progress: ScanProgress) -> schemas_out.AssetScanStatus:
-    return schemas_out.AssetScanStatus(
-        scan_id=progress.scan_id,
-        root=progress.root,
-        status=progress.status,
-        scheduled_at=ts_to_iso(progress.scheduled_at),
-        started_at=ts_to_iso(progress.started_at),
-        finished_at=ts_to_iso(progress.finished_at),
-        discovered=progress.discovered,
-        processed=progress.processed,
-        file_errors=[
-            schemas_out.AssetScanError(
-                path=e.get("path", ""),
-                message=e.get("message", ""),
-                at=e.get("at"),
-            )
-            for e in (progress.file_errors or [])
-        ],
-    )
-
-
-async def _run_hash_verify_pipeline(root: schemas_in.RootType, prog: ScanProgress, state: SlowQueueState) -> None:
-    prog.status = "running"
-    prog.started_at = time.time()
-    try:
-        prefixes = prefixes_for_root(root)
-
-        await _fast_db_consistency_pass(root)
-
-        # collect candidates from DB
-        async with await create_session() as sess:
-            verify_ids = await list_verify_candidates_under_prefixes(sess, prefixes=prefixes)
-            unhashed_ids = await list_unhashed_candidates_under_prefixes(sess, prefixes=prefixes)
-        # dedupe: prioritize verification first
-        seen = set()
-        ordered: list[int] = []
-        for lst in (verify_ids, unhashed_ids):
-            for sid in lst:
-                if sid not in seen:
-                    seen.add(sid)
-                    ordered.append(sid)
-
-        prog.discovered = len(ordered)
-
-        # queue up work
-        for sid in ordered:
-            await state.queue.put(sid)
-        state.closed = True
-        _start_state_workers(root, prog, state)
-        await _await_state_workers_then_finish(root, prog, state)
-    except asyncio.CancelledError:
-        prog.status = "cancelled"
-        raise
-    except Exception as exc:
-        _append_error(prog, path="", message=str(exc))
-        prog.status = "failed"
-        prog.finished_at = time.time()
-        LOGGER.exception("Asset scan failed for %s", root)
-    finally:
-        RUNNING_TASKS.pop(root, None)
-
-
-async def _reconcile_missing_tags_for_root(root: schemas_in.RootType, prog: ScanProgress) -> None:
-    """
-    Detect missing files quickly and toggle 'missing' tag per asset_id.
-
-    Rules:
-      - Only hashed assets (assets.hash != NULL) participate in missing tagging.
-      - We consider ALL cache states of the asset (across roots) before tagging.
-    """
-    if root == "models":
-        bases: list[str] = []
-        for _bucket, paths in get_comfy_models_folders():
-            bases.extend(paths)
-    elif root == "input":
-        bases = [folder_paths.get_input_directory()]
-    else:
-        bases = [folder_paths.get_output_directory()]
-
-    try:
-        async with await create_session() as sess:
-            # state + hash + size for the current root
-            rows = await list_cache_states_with_asset_under_prefixes(sess, prefixes=bases)
-
-            # Track fast_ok within the scanned root and whether the asset is hashed
-            by_asset: dict[str, dict[str, bool]] = {}
-            for state, a_hash, size_db in rows:
-                aid = state.asset_id
-                acc = by_asset.get(aid)
-                if acc is None:
-                    acc = {"any_fast_ok_here": False, "hashed": (a_hash is not None), "size_db": int(size_db or 0)}
-                    by_asset[aid] = acc
-                try:
-                    if acc["hashed"]:
-                        st = os.stat(state.file_path, follow_symlinks=True)
-                        if fast_asset_file_check(mtime_db=state.mtime_ns, size_db=acc["size_db"], stat_result=st):
-                            acc["any_fast_ok_here"] = True
-                except FileNotFoundError:
-                    pass
-                except OSError as e:
-                    _append_error(prog, path=state.file_path, message=str(e))
-
-            # Decide per asset, considering ALL its states (not just this root)
-            for aid, acc in by_asset.items():
-                try:
-                    if not acc["hashed"]:
-                        # Never tag seed assets as missing
-                        continue
-
-                    any_fast_ok_global = acc["any_fast_ok_here"]
-                    if not any_fast_ok_global:
-                        # Check other states outside this root
-                        others = await list_cache_states_by_asset_id(sess, asset_id=aid)
-                        for st in others:
-                            try:
-                                any_fast_ok_global = fast_asset_file_check(
-                                    mtime_db=st.mtime_ns,
-                                    size_db=acc["size_db"],
-                                    stat_result=os.stat(st.file_path, follow_symlinks=True),
-                                )
-                            except OSError:
-                                continue
-
-                    if any_fast_ok_global:
-                        await remove_missing_tag_for_asset_id(sess, asset_id=aid)
-                    else:
-                        await add_missing_tag_for_asset_id(sess, asset_id=aid, origin="automatic")
-                except Exception as ex:
-                    _append_error(prog, path="", message=f"reconcile {aid[:8]}: {ex}")
-
-            await sess.commit()
-    except Exception as e:
-        _append_error(prog, path="", message=f"reconcile failed: {e}")
-
-
-def _start_state_workers(root: schemas_in.RootType, prog: ScanProgress, state: SlowQueueState) -> None:
-    if state.workers:
-        return
-
-    async def _worker(_wid: int):
-        while True:
-            sid = await state.queue.get()
-            try:
-                if sid is None:
-                    return
-                try:
-                    async with await create_session() as sess:
-                        # Optional: fetch path for better error messages
-                        st = await sess.get(AssetCacheState, sid)
-                        try:
-                            await compute_hash_and_dedup_for_cache_state(sess, state_id=sid)
-                            await sess.commit()
-                        except Exception as e:
-                            path = st.file_path if st else f"state:{sid}"
-                            _append_error(prog, path=path, message=str(e))
-                            raise
-                except Exception:
-                    pass
-                finally:
-                    prog.processed += 1
-            finally:
-                state.queue.task_done()
-
-    state.workers = [
-        asyncio.create_task(_worker(i), name=f"asset-hash:{root}:{i}")
-        for i in range(SLOW_HASH_CONCURRENCY)
-    ]
-
-    async def _close_when_ready():
-        while not state.closed:
-            await asyncio.sleep(0.05)
-        for _ in range(SLOW_HASH_CONCURRENCY):
-            await state.queue.put(None)
-
-    asyncio.create_task(_close_when_ready())
-
-
-async def _await_state_workers_then_finish(
-    root: schemas_in.RootType, prog: ScanProgress, state: SlowQueueState
-) -> None:
-    if state.workers:
-        await asyncio.gather(*state.workers, return_exceptions=True)
-    await _reconcile_missing_tags_for_root(root, prog)
-    prog.finished_at = time.time()
-    prog.status = "completed"
-
-
-def _append_error(prog: ScanProgress, *, path: str, message: str) -> None:
-    prog.file_errors.append({
-        "path": path,
-        "message": message,
-        "at": ts_to_iso(time.time()),
-    })
-
-
-async def _fast_db_consistency_pass(
-    root: schemas_in.RootType,
-    *,
-    collect_existing_paths: bool = False,
-    update_missing_tags: bool = False,
-) -> Optional[set[str]]:
-    """Fast DB+FS pass for a root:
-      - Toggle needs_verify per state using fast check
-      - For hashed assets with at least one fast-ok state in this root: delete stale missing states
-      - For seed assets with all states missing: delete Asset and its AssetInfos
-      - Optionally add/remove 'missing' tags based on fast-ok in this root
-      - Optionally return surviving absolute paths
-    """
-    prefixes = prefixes_for_root(root)
-    if not prefixes:
-        return set() if collect_existing_paths else None
-
-    conds = []
-    for p in prefixes:
-        base = os.path.abspath(p)
-        if not base.endswith(os.sep):
-            base += os.sep
-        escaped, esc = escape_like_prefix(base)
-        conds.append(AssetCacheState.file_path.like(escaped + "%", escape=esc))
-
-    async with await create_session() as sess:
-        rows = (
-            await sess.execute(
-                sa.select(
-                    AssetCacheState.id,
-                    AssetCacheState.file_path,
-                    AssetCacheState.mtime_ns,
-                    AssetCacheState.needs_verify,
-                    AssetCacheState.asset_id,
-                    Asset.hash,
-                    Asset.size_bytes,
-                )
-                .join(Asset, Asset.id == AssetCacheState.asset_id)
-                .where(sa.or_(*conds))
-                .order_by(AssetCacheState.asset_id.asc(), AssetCacheState.id.asc())
-            )
-        ).all()
-
-        by_asset: dict[str, dict] = {}
-        for sid, fp, mtime_db, needs_verify, aid, a_hash, a_size in rows:
-            acc = by_asset.get(aid)
-            if acc is None:
-                acc = {"hash": a_hash, "size_db": int(a_size or 0), "states": []}
-                by_asset[aid] = acc
-
-            fast_ok = False
-            try:
-                exists = True
-                fast_ok = fast_asset_file_check(
-                    mtime_db=mtime_db,
-                    size_db=acc["size_db"],
-                    stat_result=os.stat(fp, follow_symlinks=True),
-                )
-            except FileNotFoundError:
-                exists = False
-            except OSError:
-                exists = False
-
-            acc["states"].append({
-                "sid": sid,
-                "fp": fp,
-                "exists": exists,
-                "fast_ok": fast_ok,
-                "needs_verify": bool(needs_verify),
-            })
-
-        to_set_verify: list[int] = []
-        to_clear_verify: list[int] = []
-        stale_state_ids: list[int] = []
-        survivors: set[str] = set()
-
-        for aid, acc in by_asset.items():
-            a_hash = acc["hash"]
-            states = acc["states"]
-            any_fast_ok = any(s["fast_ok"] for s in states)
-            all_missing = all(not s["exists"] for s in states)
-
-            for s in states:
-                if not s["exists"]:
-                    continue
-                if s["fast_ok"] and s["needs_verify"]:
-                    to_clear_verify.append(s["sid"])
-                if not s["fast_ok"] and not s["needs_verify"]:
-                    to_set_verify.append(s["sid"])
-
-            if a_hash is None:
-                if states and all_missing:  # remove seed Asset completely, if no valid AssetCache exists
-                    await sess.execute(sa.delete(AssetInfo).where(AssetInfo.asset_id == aid))
-                    asset = await sess.get(Asset, aid)
-                    if asset:
-                        await sess.delete(asset)
-                else:
-                    for s in states:
-                        if s["exists"]:
-                            survivors.add(os.path.abspath(s["fp"]))
-                continue
-
-            if any_fast_ok:  # if Asset has at least one valid AssetCache record, remove any invalid AssetCache records
-                for s in states:
-                    if not s["exists"]:
-                        stale_state_ids.append(s["sid"])
-                if update_missing_tags:
-                    with contextlib.suppress(Exception):
-                        await remove_missing_tag_for_asset_id(sess, asset_id=aid)
-            elif update_missing_tags:
-                with contextlib.suppress(Exception):
-                    await add_missing_tag_for_asset_id(sess, asset_id=aid, origin="automatic")
-
-            for s in states:
-                if s["exists"]:
-                    survivors.add(os.path.abspath(s["fp"]))
-
-        if stale_state_ids:
-            await sess.execute(sa.delete(AssetCacheState).where(AssetCacheState.id.in_(stale_state_ids)))
-        if to_set_verify:
-            await sess.execute(
-                sa.update(AssetCacheState)
-                .where(AssetCacheState.id.in_(to_set_verify))
-                .values(needs_verify=True)
-            )
-        if to_clear_verify:
-            await sess.execute(
-                sa.update(AssetCacheState)
-                .where(AssetCacheState.id.in_(to_clear_verify))
-                .values(needs_verify=False)
-            )
-        await sess.commit()
-        return survivors if collect_existing_paths else None

app/assets/storage/hashing.py

@@ -1,72 +0,0 @@
-import asyncio
-import os
-from typing import IO, Union
-
-from blake3 import blake3
-
-DEFAULT_CHUNK = 8 * 1024 * 1024  # 8 MiB
-
-
-def _hash_file_obj_sync(file_obj: IO[bytes], chunk_size: int) -> str:
-    """Hash an already-open binary file object by streaming in chunks.
-    - Seeks to the beginning before reading (if supported).
-    - Restores the original position afterward (if tell/seek are supported).
-    """
-    if chunk_size <= 0:
-        chunk_size = DEFAULT_CHUNK
-
-    orig_pos = None
-    if hasattr(file_obj, "tell"):
-        orig_pos = file_obj.tell()
-
-    try:
-        if hasattr(file_obj, "seek"):
-            file_obj.seek(0)
-
-        h = blake3()
-        while True:
-            chunk = file_obj.read(chunk_size)
-            if not chunk:
-                break
-            h.update(chunk)
-        return h.hexdigest()
-    finally:
-        if hasattr(file_obj, "seek") and orig_pos is not None:
-            file_obj.seek(orig_pos)
-
-
-def blake3_hash_sync(
-    fp: Union[str, bytes, os.PathLike[str], os.PathLike[bytes], IO[bytes]],
-    chunk_size: int = DEFAULT_CHUNK,
-) -> str:
-    """Returns a BLAKE3 hex digest for ``fp``, which may be:
-      - a filename (str/bytes) or PathLike
-      - an open binary file object
-
-    If ``fp`` is a file object, it must be opened in **binary** mode and support
-    ``read``, ``seek``, and ``tell``. The function will seek to the start before
-    reading and will attempt to restore the original position afterward.
-    """
-    if hasattr(fp, "read"):
-        return _hash_file_obj_sync(fp, chunk_size)
-
-    with open(os.fspath(fp), "rb") as f:
-        return _hash_file_obj_sync(f, chunk_size)
-
-
-async def blake3_hash(
-    fp: Union[str, bytes, os.PathLike[str], os.PathLike[bytes], IO[bytes]],
-    chunk_size: int = DEFAULT_CHUNK,
-) -> str:
-    """Async wrapper for ``blake3_hash_sync``.
-    Uses a worker thread so the event loop remains responsive.
-    """
-    # If it is a path, open inside the worker thread to keep I/O off the loop.
-    if hasattr(fp, "read"):
-        return await asyncio.to_thread(blake3_hash_sync, fp, chunk_size)
-
-    def _worker() -> str:
-        with open(os.fspath(fp), "rb") as f:
-            return _hash_file_obj_sync(f, chunk_size)
-
-    return await asyncio.to_thread(_worker)

app/database/db.py

@@ -0,0 +1,112 @@
+import logging
+import os
+import shutil
+from app.logger import log_startup_warning
+from utils.install_util import get_missing_requirements_message
+from comfy.cli_args import args
+
+_DB_AVAILABLE = False
+Session = None
+
+
+try:
+    from alembic import command
+    from alembic.config import Config
+    from alembic.runtime.migration import MigrationContext
+    from alembic.script import ScriptDirectory
+    from sqlalchemy import create_engine
+    from sqlalchemy.orm import sessionmaker
+
+    _DB_AVAILABLE = True
+except ImportError as e:
+    log_startup_warning(
+        f"""
+------------------------------------------------------------------------
+Error importing dependencies: {e}
+{get_missing_requirements_message()}
+This error is happening because ComfyUI now uses a local sqlite database.
+------------------------------------------------------------------------
+""".strip()
+    )
+
+
+def dependencies_available():
+    """
+    Temporary function to check if the dependencies are available
+    """
+    return _DB_AVAILABLE
+
+
+def can_create_session():
+    """
+    Temporary function to check if the database is available to create a session
+    During initial release there may be environmental issues (or missing dependencies) that prevent the database from being created
+    """
+    return dependencies_available() and Session is not None
+
+
+def get_alembic_config():
+    root_path = os.path.join(os.path.dirname(__file__), "../..")
+    config_path = os.path.abspath(os.path.join(root_path, "alembic.ini"))
+    scripts_path = os.path.abspath(os.path.join(root_path, "alembic_db"))
+
+    config = Config(config_path)
+    config.set_main_option("script_location", scripts_path)
+    config.set_main_option("sqlalchemy.url", args.database_url)
+
+    return config
+
+
+def get_db_path():
+    url = args.database_url
+    if url.startswith("sqlite:///"):
+        return url.split("///")[1]
+    else:
+        raise ValueError(f"Unsupported database URL '{url}'.")
+
+
+def init_db():
+    db_url = args.database_url
+    logging.debug(f"Database URL: {db_url}")
+    db_path = get_db_path()
+    db_exists = os.path.exists(db_path)
+
+    config = get_alembic_config()
+
+    # Check if we need to upgrade
+    engine = create_engine(db_url)
+    conn = engine.connect()
+
+    context = MigrationContext.configure(conn)
+    current_rev = context.get_current_revision()
+
+    script = ScriptDirectory.from_config(config)
+    target_rev = script.get_current_head()
+
+    if target_rev is None:
+        logging.warning("No target revision found.")
+    elif current_rev != target_rev:
+        # Backup the database pre upgrade
+        backup_path = db_path + ".bkp"
+        if db_exists:
+            shutil.copy(db_path, backup_path)
+        else:
+            backup_path = None
+
+        try:
+            command.upgrade(config, target_rev)
+            logging.info(f"Database upgraded from {current_rev} to {target_rev}")
+        except Exception as e:
+            if backup_path:
+                # Restore the database from backup if upgrade fails
+                shutil.copy(backup_path, db_path)
+                os.remove(backup_path)
+            logging.exception("Error upgrading database: ")
+            raise e
+
+    global Session
+    Session = sessionmaker(bind=engine)
+
+
+def create_session():
+    return Session()

app/database/models.py

@@ -0,0 +1,14 @@
+from sqlalchemy.orm import declarative_base
+
+Base = declarative_base()
+
+
+def to_dict(obj):
+    fields = obj.__table__.columns.keys()
+    return {
+        field: (val.to_dict() if hasattr(val, "to_dict") else val)
+        for field in fields
+        if (val := getattr(obj, field))
+    }
+
+# TODO: Define models here

app/db.py

@@ -1,255 +0,0 @@
-import logging
-import os
-import shutil
-from contextlib import asynccontextmanager
-from typing import Optional
-
-from alembic import command
-from alembic.config import Config
-from alembic.runtime.migration import MigrationContext
-from alembic.script import ScriptDirectory
-from sqlalchemy import create_engine, text
-from sqlalchemy.engine import make_url
-from sqlalchemy.ext.asyncio import (
-    AsyncEngine,
-    AsyncSession,
-    async_sessionmaker,
-    create_async_engine,
-)
-
-from comfy.cli_args import args
-
-LOGGER = logging.getLogger(__name__)
-ENGINE: Optional[AsyncEngine] = None
-SESSION: Optional[async_sessionmaker] = None
-
-
-def _root_paths():
-    """Resolve alembic.ini and migrations script folder."""
-    root_path = os.path.abspath(os.path.dirname(__file__))
-    config_path = os.path.abspath(os.path.join(root_path, "../alembic.ini"))
-    scripts_path = os.path.abspath(os.path.join(root_path, "alembic_db"))
-    return config_path, scripts_path
-
-
-def _absolutize_sqlite_url(db_url: str) -> str:
-    """Make SQLite database path absolute. No-op for non-SQLite URLs."""
-    try:
-        u = make_url(db_url)
-    except Exception:
-        return db_url
-
-    if not u.drivername.startswith("sqlite"):
-        return db_url
-
-    db_path: str = u.database or ""
-    if isinstance(db_path, str) and db_path.startswith("file:"):
-        return str(u)  # Do not touch SQLite URI databases like: "file:xxx?mode=memory&cache=shared"
-    if not os.path.isabs(db_path):
-        db_path = os.path.abspath(os.path.join(os.getcwd(), db_path))
-        u = u.set(database=db_path)
-    return str(u)
-
-
-def _normalize_sqlite_memory_url(db_url: str) -> tuple[str, bool]:
-    """
-    If db_url points at an in-memory SQLite DB (":memory:" or file:... mode=memory),
-    rewrite it to a *named* shared in-memory URI and ensure 'uri=true' is present.
-    Returns: (normalized_url, is_memory)
-    """
-    try:
-        u = make_url(db_url)
-    except Exception:
-        return db_url, False
-    if not u.drivername.startswith("sqlite"):
-        return db_url, False
-
-    db = u.database or ""
-    if db == ":memory:":
-        u = u.set(database=f"file:comfyui_db_{os.getpid()}?mode=memory&cache=shared&uri=true")
-        return str(u), True
-    if isinstance(db, str) and db.startswith("file:") and "mode=memory" in db:
-        if "uri=true" not in db:
-            u = u.set(database=(db + ("&" if "?" in db else "?") + "uri=true"))
-        return str(u), True
-    return str(u), False
-
-
-def _get_sqlite_file_path(sync_url: str) -> Optional[str]:
-    """Return the on-disk path for a SQLite URL, else None."""
-    try:
-        u = make_url(sync_url)
-    except Exception:
-        return None
-
-    if not u.drivername.startswith("sqlite"):
-        return None
-    db_path = u.database
-    if isinstance(db_path, str) and db_path.startswith("file:"):
-        return None  # Not a real file if it is a URI like "file:...?"
-    return db_path
-
-
-def _get_alembic_config(sync_url: str) -> Config:
-    """Prepare Alembic Config with script location and DB URL."""
-    config_path, scripts_path = _root_paths()
-    cfg = Config(config_path)
-    cfg.set_main_option("script_location", scripts_path)
-    cfg.set_main_option("sqlalchemy.url", sync_url)
-    return cfg
-
-
-async def init_db_engine() -> None:
-    """Initialize async engine + sessionmaker and run migrations to head.
-
-    This must be called once on application startup before any DB usage.
-    """
-    global ENGINE, SESSION
-
-    if ENGINE is not None:
-        return
-
-    raw_url = args.database_url
-    if not raw_url:
-        raise RuntimeError("Database URL is not configured.")
-
-    db_url, is_mem = _normalize_sqlite_memory_url(raw_url)
-    db_url = _absolutize_sqlite_url(db_url)
-
-    # Prepare async engine
-    connect_args = {}
-    if db_url.startswith("sqlite"):
-        connect_args = {
-            "check_same_thread": False,
-            "timeout": 12,
-        }
-        if is_mem:
-            connect_args["uri"] = True
-
-    ENGINE = create_async_engine(
-        db_url,
-        connect_args=connect_args,
-        pool_pre_ping=True,
-        future=True,
-    )
-
-    # Enforce SQLite pragmas on the async engine
-    if db_url.startswith("sqlite"):
-        async with ENGINE.begin() as conn:
-            if not is_mem:
-                # WAL for concurrency and durability, Foreign Keys for referential integrity
-                current_mode = (await conn.execute(text("PRAGMA journal_mode;"))).scalar()
-                if str(current_mode).lower() != "wal":
-                    new_mode = (await conn.execute(text("PRAGMA journal_mode=WAL;"))).scalar()
-                    if str(new_mode).lower() != "wal":
-                        raise RuntimeError("Failed to set SQLite journal mode to WAL.")
-                    LOGGER.info("SQLite journal mode set to WAL.")
-
-            await conn.execute(text("PRAGMA foreign_keys = ON;"))
-            await conn.execute(text("PRAGMA synchronous = NORMAL;"))
-
-    await _run_migrations(database_url=db_url, connect_args=connect_args)
-
-    SESSION = async_sessionmaker(
-        bind=ENGINE,
-        class_=AsyncSession,
-        expire_on_commit=False,
-        autoflush=False,
-        autocommit=False,
-    )
-
-
-async def _run_migrations(database_url: str, connect_args: dict) -> None:
-    if database_url.find("postgresql+psycopg") == -1:
-        """SQLite: Convert an async SQLAlchemy URL to a sync URL for Alembic."""
-        u = make_url(database_url)
-        driver = u.drivername
-        if not driver.startswith("sqlite+aiosqlite"):
-            raise ValueError(f"Unsupported DB driver: {driver}")
-        database_url, is_mem = _normalize_sqlite_memory_url(str(u.set(drivername="sqlite")))
-        database_url = _absolutize_sqlite_url(database_url)
-
-    cfg = _get_alembic_config(database_url)
-    engine = create_engine(database_url, future=True, connect_args=connect_args)
-    with engine.connect() as conn:
-        context = MigrationContext.configure(conn)
-        current_rev = context.get_current_revision()
-
-    script = ScriptDirectory.from_config(cfg)
-    target_rev = script.get_current_head()
-
-    if target_rev is None:
-        LOGGER.warning("Alembic: no target revision found.")
-        return
-
-    if current_rev == target_rev:
-        LOGGER.debug("Alembic: database already at head %s", target_rev)
-        return
-
-    LOGGER.info("Alembic: upgrading database from %s to %s", current_rev, target_rev)
-
-    # Optional backup for SQLite file DBs
-    backup_path = None
-    sqlite_path = _get_sqlite_file_path(database_url)
-    if sqlite_path and os.path.exists(sqlite_path):
-        backup_path = sqlite_path + ".bkp"
-        try:
-            shutil.copy(sqlite_path, backup_path)
-        except Exception as exc:
-            LOGGER.warning("Failed to create SQLite backup before migration: %s", exc)
-
-    try:
-        command.upgrade(cfg, target_rev)
-    except Exception:
-        if backup_path and os.path.exists(backup_path):
-            LOGGER.exception("Error upgrading database, attempting restore from backup.")
-            try:
-                shutil.copy(backup_path, sqlite_path)  # restore
-                os.remove(backup_path)
-            except Exception as re:
-                LOGGER.error("Failed to restore SQLite backup: %s", re)
-        else:
-            LOGGER.exception("Error upgrading database, backup is not available.")
-        raise
-
-
-def get_engine():
-    """Return the global async engine (initialized after init_db_engine())."""
-    if ENGINE is None:
-        raise RuntimeError("Engine is not initialized. Call init_db_engine() first.")
-    return ENGINE
-
-
-def get_session_maker():
-    """Return the global async_sessionmaker (initialized after init_db_engine())."""
-    if SESSION is None:
-        raise RuntimeError("Session maker is not initialized. Call init_db_engine() first.")
-    return SESSION
-
-
-@asynccontextmanager
-async def session_scope():
-    """Async context manager for a unit of work:
-
-    async with session_scope() as sess:
-        ... use sess ...
-    """
-    maker = get_session_maker()
-    async with maker() as sess:
-        try:
-            yield sess
-            await sess.commit()
-        except Exception:
-            await sess.rollback()
-            raise
-
-
-async def create_session():
-    """Convenience helper to acquire a single AsyncSession instance.
-
-    Typical usage:
-        async with (await create_session()) as sess:
-            ...
-    """
-    maker = get_session_maker()
-    return maker()

app/frontend_management.py

@@ -196,17 +196,6 @@ def download_release_asset_zip(release: Release, destination_path: str) -> None:
 
 
 class FrontendManager:
-    """
-    A class to manage ComfyUI frontend versions and installations.
-
-    This class handles the initialization and management of different frontend versions,
-    including the default frontend from the pip package and custom frontend versions
-    from GitHub repositories.
-
-    Attributes:
-        CUSTOM_FRONTENDS_ROOT (str): The root directory where custom frontend versions are stored.
-    """
-
     CUSTOM_FRONTENDS_ROOT = str(Path(__file__).parents[1] / "web_custom_versions")
 
     @classmethod
@@ -216,15 +205,6 @@ class FrontendManager:
 
     @classmethod
     def default_frontend_path(cls) -> str:
-        """
-        Get the path to the default frontend installation from the pip package.
-
-        Returns:
-            str: The path to the default frontend static files.
-
-        Raises:
-            SystemExit: If the comfyui-frontend-package is not installed.
-        """
         try:
             import comfyui_frontend_package
 
@@ -245,15 +225,6 @@ comfyui-frontend-package is not installed.
 
     @classmethod
     def templates_path(cls) -> str:
-        """
-        Get the path to the workflow templates.
-
-        Returns:
-            str: The path to the workflow templates directory.
-
-        Raises:
-            SystemExit: If the comfyui-workflow-templates package is not installed.
-        """
         try:
             import comfyui_workflow_templates
 
@@ -289,16 +260,11 @@ comfyui-workflow-templates is not installed.
     @classmethod
     def parse_version_string(cls, value: str) -> tuple[str, str, str]:
         """
-        Parse a version string into its components.
-
-        The version string should be in the format: 'owner/repo@version'
-        where version can be either a semantic version (v1.2.3) or 'latest'.
-
         Args:
             value (str): The version string to parse.
 
         Returns:
-            tuple[str, str, str]: A tuple containing (owner, repo, version).
+            tuple[str, str]: A tuple containing provider name and version.
 
         Raises:
             argparse.ArgumentTypeError: If the version string is invalid.
@@ -315,22 +281,18 @@ comfyui-workflow-templates is not installed.
         cls, version_string: str, provider: Optional[FrontEndProvider] = None
     ) -> str:
         """
-        Initialize a frontend version without error handling.
-
-        This method attempts to initialize a specific frontend version, either from
-        the default pip package or from a custom GitHub repository. It will download
-        and extract the frontend files if necessary.
+        Initializes the frontend for the specified version.
 
         Args:
-            version_string (str): The version string specifying which frontend to use.
-            provider (FrontEndProvider, optional): The provider to use for custom frontends.
+            version_string (str): The version string.
+            provider (FrontEndProvider, optional): The provider to use. Defaults to None.
 
         Returns:
             str: The path to the initialized frontend.
 
         Raises:
-            Exception: If there is an error during initialization (e.g., network timeout,
-                      invalid URL, or missing assets).
+            Exception: If there is an error during the initialization process.
+            main error source might be request timeout or invalid URL.
         """
         if version_string == DEFAULT_VERSION_STRING:
             check_frontend_version()
@@ -382,17 +344,13 @@ comfyui-workflow-templates is not installed.
     @classmethod
     def init_frontend(cls, version_string: str) -> str:
         """
-        Initialize a frontend version with error handling.
-
-        This is the main method to initialize a frontend version. It wraps init_frontend_unsafe
-        with error handling, falling back to the default frontend if initialization fails.
+        Initializes the frontend with the specified version string.
 
         Args:
-            version_string (str): The version string specifying which frontend to use.
+            version_string (str): The version string to initialize the frontend with.
 
         Returns:
-            str: The path to the initialized frontend. If initialization fails,
-                 returns the path to the default frontend.
+            str: The path of the initialized frontend.
         """
         try:
             return cls.init_frontend_unsafe(version_string)

comfy/audio_encoders/audio_encoders.py

@@ -1,5 +1,4 @@
 from .wav2vec2 import Wav2Vec2Model
-from .whisper import WhisperLargeV3
 import comfy.model_management
 import comfy.ops
 import comfy.utils
@@ -12,18 +11,7 @@ class AudioEncoderModel():
         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)
-        model_type = config.pop("model_type")
-        model_config = dict(config)
-        model_config.update({
-            "dtype": self.dtype,
-            "device": offload_device,
-            "operations": comfy.ops.manual_cast
-        })
-
-        if model_type == "wav2vec2":
-            self.model = Wav2Vec2Model(**model_config)
-        elif model_type == "whisper3":
-            self.model = WhisperLargeV3(**model_config)
+        self.model = Wav2Vec2Model(dtype=self.dtype, device=offload_device, operations=comfy.ops.manual_cast)
         self.model.eval()
         self.patcher = comfy.model_patcher.ModelPatcher(self.model, load_device=self.load_device, offload_device=offload_device)
         self.model_sample_rate = 16000
@@ -41,51 +29,14 @@ class AudioEncoderModel():
         outputs = {}
         outputs["encoded_audio"] = out
         outputs["encoded_audio_all_layers"] = all_layers
-        outputs["audio_samples"] = audio.shape[2]
         return outputs
 
 
 def load_audio_encoder_from_sd(sd, prefix=""):
+    audio_encoder = AudioEncoderModel(None)
     sd = comfy.utils.state_dict_prefix_replace(sd, {"wav2vec2.": ""})
-    if "encoder.layer_norm.bias" in sd: #wav2vec2
-        embed_dim = sd["encoder.layer_norm.bias"].shape[0]
-        if embed_dim == 1024:# large
-            config = {
-                "model_type": "wav2vec2",
-                "embed_dim": 1024,
-                "num_heads": 16,
-                "num_layers": 24,
-                "conv_norm": True,
-                "conv_bias": True,
-                "do_normalize": True,
-                "do_stable_layer_norm": True
-                }
-        elif embed_dim == 768: # base
-            config = {
-                "model_type": "wav2vec2",
-                "embed_dim": 768,
-                "num_heads": 12,
-                "num_layers": 12,
-                "conv_norm": False,
-                "conv_bias": False,
-                "do_normalize": False, # chinese-wav2vec2-base has this False
-                "do_stable_layer_norm": False
-            }
-        else:
-            raise RuntimeError("ERROR: audio encoder file is invalid or unsupported embed_dim: {}".format(embed_dim))
-    elif "model.encoder.embed_positions.weight" in sd:
-        sd = comfy.utils.state_dict_prefix_replace(sd, {"model.": ""})
-        config = {
-            "model_type": "whisper3",
-        }
-    else:
-        raise RuntimeError("ERROR: audio encoder not supported.")
-
-    audio_encoder = AudioEncoderModel(config)
     m, u = audio_encoder.load_sd(sd)
     if len(m) > 0:
         logging.warning("missing audio encoder: {}".format(m))
-    if len(u) > 0:
-        logging.warning("unexpected audio encoder: {}".format(u))
 
     return audio_encoder

comfy/audio_encoders/wav2vec2.py

@@ -13,49 +13,19 @@ class LayerNormConv(nn.Module):
         x = self.conv(x)
         return torch.nn.functional.gelu(self.layer_norm(x.transpose(-2, -1)).transpose(-2, -1))
 
-class LayerGroupNormConv(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)
-        self.layer_norm = operations.GroupNorm(num_groups=out_channels, num_channels=out_channels, affine=True, device=device, dtype=dtype)
-
-    def forward(self, x):
-        x = self.conv(x)
-        return torch.nn.functional.gelu(self.layer_norm(x))
-
-class ConvNoNorm(nn.Module):
-    def __init__(self, in_channels, out_channels, kernel_size, stride, bias=False, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.conv = operations.Conv1d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, device=device, dtype=dtype)
-
-    def forward(self, x):
-        x = self.conv(x)
-        return torch.nn.functional.gelu(x)
-
 
 class ConvFeatureEncoder(nn.Module):
-    def __init__(self, conv_dim, conv_bias=False, conv_norm=True, dtype=None, device=None, operations=None):
+    def __init__(self, conv_dim, dtype=None, device=None, operations=None):
         super().__init__()
-        if conv_norm:
-            self.conv_layers = nn.ModuleList([
-                LayerNormConv(1, conv_dim, kernel_size=10, stride=5, bias=True, device=device, dtype=dtype, operations=operations),
-                LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-            ])
-        else:
-            self.conv_layers = nn.ModuleList([
-                LayerGroupNormConv(1, conv_dim, kernel_size=10, stride=5, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                ConvNoNorm(conv_dim, conv_dim, kernel_size=3, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                ConvNoNorm(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-                ConvNoNorm(conv_dim, conv_dim, kernel_size=2, stride=2, bias=conv_bias, device=device, dtype=dtype, operations=operations),
-            ])
+        self.conv_layers = nn.ModuleList([
+            LayerNormConv(1, conv_dim, kernel_size=10, stride=5, bias=True, device=device, dtype=dtype, operations=operations),
+            LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
+            LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
+            LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
+            LayerNormConv(conv_dim, conv_dim, kernel_size=3, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
+            LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
+            LayerNormConv(conv_dim, conv_dim, kernel_size=2, stride=2, bias=True, device=device, dtype=dtype, operations=operations),
+        ])
 
     def forward(self, x):
         x = x.unsqueeze(1)
@@ -106,7 +76,6 @@ class TransformerEncoder(nn.Module):
         num_heads=12,
         num_layers=12,
         mlp_ratio=4.0,
-        do_stable_layer_norm=True,
         dtype=None, device=None, operations=None
     ):
         super().__init__()
@@ -117,25 +86,20 @@ class TransformerEncoder(nn.Module):
                 embed_dim=embed_dim,
                 num_heads=num_heads,
                 mlp_ratio=mlp_ratio,
-                do_stable_layer_norm=do_stable_layer_norm,
                 device=device, dtype=dtype, operations=operations
             )
             for _ in range(num_layers)
         ])
 
         self.layer_norm = operations.LayerNorm(embed_dim, eps=1e-05, device=device, dtype=dtype)
-        self.do_stable_layer_norm = do_stable_layer_norm
 
     def forward(self, x, mask=None):
         x = x + self.pos_conv_embed(x)
         all_x = ()
-        if not self.do_stable_layer_norm:
-            x = self.layer_norm(x)
         for layer in self.layers:
             all_x += (x,)
             x = layer(x, mask)
-        if self.do_stable_layer_norm:
-            x = self.layer_norm(x)
+        x = self.layer_norm(x)
         all_x += (x,)
         return x, all_x
 
@@ -181,7 +145,6 @@ class TransformerEncoderLayer(nn.Module):
         embed_dim=768,
         num_heads=12,
         mlp_ratio=4.0,
-        do_stable_layer_norm=True,
         dtype=None, device=None, operations=None
     ):
         super().__init__()
@@ -191,19 +154,15 @@ class TransformerEncoderLayer(nn.Module):
         self.layer_norm = operations.LayerNorm(embed_dim, device=device, dtype=dtype)
         self.feed_forward = FeedForward(embed_dim, mlp_ratio, device=device, dtype=dtype, operations=operations)
         self.final_layer_norm = operations.LayerNorm(embed_dim, device=device, dtype=dtype)
-        self.do_stable_layer_norm = do_stable_layer_norm
 
     def forward(self, x, mask=None):
         residual = x
-        if self.do_stable_layer_norm:
-            x = self.layer_norm(x)
+        x = self.layer_norm(x)
         x = self.attention(x, mask=mask)
         x = residual + x
-        if not self.do_stable_layer_norm:
-            x = self.layer_norm(x)
-            return self.final_layer_norm(x + self.feed_forward(x))
-        else:
-            return x + self.feed_forward(self.final_layer_norm(x))
+
+        x = x + self.feed_forward(self.final_layer_norm(x))
+        return x
 
 
 class Wav2Vec2Model(nn.Module):
@@ -215,38 +174,34 @@ class Wav2Vec2Model(nn.Module):
         final_dim=256,
         num_heads=16,
         num_layers=24,
-        conv_norm=True,
-        conv_bias=True,
-        do_normalize=True,
-        do_stable_layer_norm=True,
         dtype=None, device=None, operations=None
     ):
         super().__init__()
 
         conv_dim = 512
-        self.feature_extractor = ConvFeatureEncoder(conv_dim, conv_norm=conv_norm, conv_bias=conv_bias, device=device, dtype=dtype, operations=operations)
+        self.feature_extractor = ConvFeatureEncoder(conv_dim, device=device, dtype=dtype, operations=operations)
         self.feature_projection = FeatureProjection(conv_dim, embed_dim, device=device, dtype=dtype, operations=operations)
 
         self.masked_spec_embed = nn.Parameter(torch.empty(embed_dim, device=device, dtype=dtype))
-        self.do_normalize = do_normalize
 
         self.encoder = TransformerEncoder(
             embed_dim=embed_dim,
             num_heads=num_heads,
             num_layers=num_layers,
-            do_stable_layer_norm=do_stable_layer_norm,
             device=device, dtype=dtype, operations=operations
         )
 
     def forward(self, x, mask_time_indices=None, return_dict=False):
+
         x = torch.mean(x, dim=1)
 
-        if self.do_normalize:
-            x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7)
+        x = (x - x.mean()) / torch.sqrt(x.var() + 1e-7)
 
         features = self.feature_extractor(x)
         features = self.feature_projection(features)
+
         batch_size, seq_len, _ = features.shape
 
         x, all_x = self.encoder(features)
+
         return x, all_x

comfy/audio_encoders/whisper.py

@@ -1,186 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import torchaudio
-from typing import Optional
-from comfy.ldm.modules.attention import optimized_attention_masked
-import comfy.ops
-
-class WhisperFeatureExtractor(nn.Module):
-    def __init__(self, n_mels=128, device=None):
-        super().__init__()
-        self.sample_rate = 16000
-        self.n_fft = 400
-        self.hop_length = 160
-        self.n_mels = n_mels
-        self.chunk_length = 30
-        self.n_samples = 480000
-
-        self.mel_spectrogram = torchaudio.transforms.MelSpectrogram(
-            sample_rate=self.sample_rate,
-            n_fft=self.n_fft,
-            hop_length=self.hop_length,
-            n_mels=self.n_mels,
-            f_min=0,
-            f_max=8000,
-            norm="slaney",
-            mel_scale="slaney",
-        ).to(device)
-
-    def __call__(self, audio):
-        audio = torch.mean(audio, dim=1)
-        batch_size = audio.shape[0]
-        processed_audio = []
-
-        for i in range(batch_size):
-            aud = audio[i]
-            if aud.shape[0] > self.n_samples:
-                aud = aud[:self.n_samples]
-            elif aud.shape[0] < self.n_samples:
-                aud = F.pad(aud, (0, self.n_samples - aud.shape[0]))
-            processed_audio.append(aud)
-
-        audio = torch.stack(processed_audio)
-
-        mel_spec = self.mel_spectrogram(audio.to(self.mel_spectrogram.spectrogram.window.device))[:, :, :-1].to(audio.device)
-
-        log_mel_spec = torch.clamp(mel_spec, min=1e-10).log10()
-        log_mel_spec = torch.maximum(log_mel_spec, log_mel_spec.max() - 8.0)
-        log_mel_spec = (log_mel_spec + 4.0) / 4.0
-
-        return log_mel_spec
-
-
-class MultiHeadAttention(nn.Module):
-    def __init__(self, d_model: int, n_heads: int, dtype=None, device=None, operations=None):
-        super().__init__()
-        assert d_model % n_heads == 0
-
-        self.d_model = d_model
-        self.n_heads = n_heads
-        self.d_k = d_model // n_heads
-
-        self.q_proj = operations.Linear(d_model, d_model, dtype=dtype, device=device)
-        self.k_proj = operations.Linear(d_model, d_model, bias=False, dtype=dtype, device=device)
-        self.v_proj = operations.Linear(d_model, d_model, dtype=dtype, device=device)
-        self.out_proj = operations.Linear(d_model, d_model, dtype=dtype, device=device)
-
-    def forward(
-        self,
-        query: torch.Tensor,
-        key: torch.Tensor,
-        value: torch.Tensor,
-        mask: Optional[torch.Tensor] = None,
-    ) -> torch.Tensor:
-        batch_size, seq_len, _ = query.shape
-
-        q = self.q_proj(query)
-        k = self.k_proj(key)
-        v = self.v_proj(value)
-
-        attn_output = optimized_attention_masked(q, k, v, self.n_heads, mask)
-        attn_output = self.out_proj(attn_output)
-
-        return attn_output
-
-
-class EncoderLayer(nn.Module):
-    def __init__(self, d_model: int, n_heads: int, d_ff: int, dtype=None, device=None, operations=None):
-        super().__init__()
-
-        self.self_attn = MultiHeadAttention(d_model, n_heads, dtype=dtype, device=device, operations=operations)
-        self.self_attn_layer_norm = operations.LayerNorm(d_model, dtype=dtype, device=device)
-
-        self.fc1 = operations.Linear(d_model, d_ff, dtype=dtype, device=device)
-        self.fc2 = operations.Linear(d_ff, d_model, dtype=dtype, device=device)
-        self.final_layer_norm = operations.LayerNorm(d_model, dtype=dtype, device=device)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        attention_mask: Optional[torch.Tensor] = None
-    ) -> torch.Tensor:
-        residual = x
-        x = self.self_attn_layer_norm(x)
-        x = self.self_attn(x, x, x, attention_mask)
-        x = residual + x
-
-        residual = x
-        x = self.final_layer_norm(x)
-        x = self.fc1(x)
-        x = F.gelu(x)
-        x = self.fc2(x)
-        x = residual + x
-
-        return x
-
-
-class AudioEncoder(nn.Module):
-    def __init__(
-        self,
-        n_mels: int = 128,
-        n_ctx: int = 1500,
-        n_state: int = 1280,
-        n_head: int = 20,
-        n_layer: int = 32,
-        dtype=None,
-        device=None,
-        operations=None
-    ):
-        super().__init__()
-
-        self.conv1 = operations.Conv1d(n_mels, n_state, kernel_size=3, padding=1, dtype=dtype, device=device)
-        self.conv2 = operations.Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1, dtype=dtype, device=device)
-
-        self.embed_positions = operations.Embedding(n_ctx, n_state, dtype=dtype, device=device)
-
-        self.layers = nn.ModuleList([
-            EncoderLayer(n_state, n_head, n_state * 4, dtype=dtype, device=device, operations=operations)
-            for _ in range(n_layer)
-        ])
-
-        self.layer_norm = operations.LayerNorm(n_state, dtype=dtype, device=device)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = F.gelu(self.conv1(x))
-        x = F.gelu(self.conv2(x))
-
-        x = x.transpose(1, 2)
-
-        x = x + comfy.ops.cast_to_input(self.embed_positions.weight[:, :x.shape[1]], x)
-
-        all_x = ()
-        for layer in self.layers:
-            all_x += (x,)
-            x = layer(x)
-
-        x = self.layer_norm(x)
-        all_x += (x,)
-        return x, all_x
-
-
-class WhisperLargeV3(nn.Module):
-    def __init__(
-        self,
-        n_mels: int = 128,
-        n_audio_ctx: int = 1500,
-        n_audio_state: int = 1280,
-        n_audio_head: int = 20,
-        n_audio_layer: int = 32,
-        dtype=None,
-        device=None,
-        operations=None
-    ):
-        super().__init__()
-
-        self.feature_extractor = WhisperFeatureExtractor(n_mels=n_mels, device=device)
-
-        self.encoder = AudioEncoder(
-            n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer,
-            dtype=dtype, device=device, operations=operations
-        )
-
-    def forward(self, audio):
-        mel = self.feature_extractor(audio)
-        x, all_x = self.encoder(mel)
-        return x, all_x

comfy/cli_args.py

@@ -143,9 +143,8 @@ class PerformanceFeature(enum.Enum):
     Fp16Accumulation = "fp16_accumulation"
     Fp8MatrixMultiplication = "fp8_matrix_mult"
     CublasOps = "cublas_ops"
-    AutoTune = "autotune"
 
-parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: {}".format(" ".join(map(lambda c: c.value, PerformanceFeature))))
+parser.add_argument("--fast", nargs="*", type=PerformanceFeature, help="Enable some untested and potentially quality deteriorating optimizations. --fast with no arguments enables everything. You can pass a list specific optimizations if you only want to enable specific ones. Current valid optimizations: fp16_accumulation fp8_matrix_mult cublas_ops")
 
 parser.add_argument("--mmap-torch-files", action="store_true", help="Use mmap when loading ckpt/pt files.")
 parser.add_argument("--disable-mmap", action="store_true", help="Don't use mmap when loading safetensors.")
@@ -212,8 +211,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+aiosqlite:///{database_default_path}", help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite+aiosqlite:///:memory:'.")
-parser.add_argument("--disable-assets-autoscan", action="store_true", help="Disable asset scanning on startup for database synchronization.")
+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:'.")
 
 if comfy.options.args_parsing:
     args = parser.parse_args()

comfy/clip_model.py

@@ -61,12 +61,8 @@ class CLIPEncoder(torch.nn.Module):
     def forward(self, x, mask=None, intermediate_output=None):
         optimized_attention = optimized_attention_for_device(x.device, mask=mask is not None, small_input=True)
 
-        all_intermediate = None
         if intermediate_output is not None:
-            if intermediate_output == "all":
-                all_intermediate = []
-                intermediate_output = None
-            elif intermediate_output < 0:
+            if intermediate_output < 0:
                 intermediate_output = len(self.layers) + intermediate_output
 
         intermediate = None
@@ -74,12 +70,6 @@ class CLIPEncoder(torch.nn.Module):
             x = l(x, mask, optimized_attention)
             if i == intermediate_output:
                 intermediate = x.clone()
-            if all_intermediate is not None:
-                all_intermediate.append(x.unsqueeze(1).clone())
-
-        if all_intermediate is not None:
-            intermediate = torch.cat(all_intermediate, dim=1)
-
         return x, intermediate
 
 class CLIPEmbeddings(torch.nn.Module):

comfy/clip_vision.py

@@ -50,13 +50,7 @@ class ClipVisionModel():
         self.image_size = config.get("image_size", 224)
         self.image_mean = config.get("image_mean", [0.48145466, 0.4578275, 0.40821073])
         self.image_std = config.get("image_std", [0.26862954, 0.26130258, 0.27577711])
-        model_type = config.get("model_type", "clip_vision_model")
-        model_class = IMAGE_ENCODERS.get(model_type)
-        if model_type == "siglip_vision_model":
-            self.return_all_hidden_states = True
-        else:
-            self.return_all_hidden_states = False
-
+        model_class = IMAGE_ENCODERS.get(config.get("model_type", "clip_vision_model"))
         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)
@@ -74,18 +68,12 @@ class ClipVisionModel():
     def encode_image(self, image, crop=True):
         comfy.model_management.load_model_gpu(self.patcher)
         pixel_values = clip_preprocess(image.to(self.load_device), size=self.image_size, mean=self.image_mean, std=self.image_std, crop=crop).float()
-        out = self.model(pixel_values=pixel_values, intermediate_output='all' if self.return_all_hidden_states else -2)
+        out = self.model(pixel_values=pixel_values, intermediate_output=-2)
 
         outputs = Output()
         outputs["last_hidden_state"] = out[0].to(comfy.model_management.intermediate_device())
         outputs["image_embeds"] = out[2].to(comfy.model_management.intermediate_device())
-        if self.return_all_hidden_states:
-            all_hs = out[1].to(comfy.model_management.intermediate_device())
-            outputs["penultimate_hidden_states"] = all_hs[:, -2]
-            outputs["all_hidden_states"] = all_hs
-        else:
-            outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
-
+        outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device())
         outputs["mm_projected"] = out[3]
         return outputs
 
@@ -136,12 +124,8 @@ def load_clipvision_from_sd(sd, prefix="", convert_keys=False):
                 json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl_336.json")
         else:
             json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "clip_vision_config_vitl.json")
-
-    # Dinov2
-    elif 'encoder.layer.39.layer_scale2.lambda1' in sd:
+    elif "embeddings.patch_embeddings.projection.weight" in sd:
         json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_giant.json")
-    elif 'encoder.layer.23.layer_scale2.lambda1' in sd:
-        json_config = os.path.join(os.path.join(os.path.dirname(os.path.realpath(__file__)), "image_encoders"), "dino2_large.json")
     else:
         return None
 

comfy/controlnet.py

@@ -253,10 +253,7 @@ class ControlNet(ControlBase):
                 to_concat = []
                 for c in self.extra_concat_orig:
                     c = c.to(self.cond_hint.device)
-                    c = comfy.utils.common_upscale(c, self.cond_hint.shape[-1], self.cond_hint.shape[-2], self.upscale_algorithm, "center")
-                    if c.ndim < self.cond_hint.ndim:
-                        c = c.unsqueeze(2)
-                        c = comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[2], dim=2)
+                    c = comfy.utils.common_upscale(c, self.cond_hint.shape[3], self.cond_hint.shape[2], self.upscale_algorithm, "center")
                     to_concat.append(comfy.utils.repeat_to_batch_size(c, self.cond_hint.shape[0]))
                 self.cond_hint = torch.cat([self.cond_hint] + to_concat, dim=1)
 
@@ -588,18 +585,11 @@ def load_controlnet_flux_instantx(sd, model_options={}):
 
 def load_controlnet_qwen_instantx(sd, model_options={}):
     model_config, operations, load_device, unet_dtype, manual_cast_dtype, offload_device = controlnet_config(sd, model_options=model_options)
-    control_latent_channels = sd.get("controlnet_x_embedder.weight").shape[1]
-
-    extra_condition_channels = 0
-    concat_mask = False
-    if control_latent_channels == 68: #inpaint controlnet
-        extra_condition_channels = control_latent_channels - 64
-        concat_mask = True
-    control_model = comfy.ldm.qwen_image.controlnet.QwenImageControlNetModel(extra_condition_channels=extra_condition_channels, operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
+    control_model = comfy.ldm.qwen_image.controlnet.QwenImageControlNetModel(operations=operations, device=offload_device, dtype=unet_dtype, **model_config.unet_config)
     control_model = controlnet_load_state_dict(control_model, sd)
     latent_format = comfy.latent_formats.Wan21()
     extra_conds = []
-    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, concat_mask=concat_mask, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
+    control = ControlNet(control_model, compression_ratio=1, latent_format=latent_format, load_device=load_device, manual_cast_dtype=manual_cast_dtype, extra_conds=extra_conds)
     return control
 
 def convert_mistoline(sd):

comfy/image_encoders/dino2.py

@@ -31,20 +31,6 @@ class LayerScale(torch.nn.Module):
     def forward(self, x):
         return x * comfy.model_management.cast_to_device(self.lambda1, x.device, x.dtype)
 
-class Dinov2MLP(torch.nn.Module):
-    def __init__(self, hidden_size: int, dtype, device, operations):
-        super().__init__()
-
-        mlp_ratio = 4
-        hidden_features = int(hidden_size * mlp_ratio)
-        self.fc1 = operations.Linear(hidden_size, hidden_features, bias = True, device=device, dtype=dtype)
-        self.fc2 = operations.Linear(hidden_features, hidden_size, bias = True, device=device, dtype=dtype)
-
-    def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
-        hidden_state = self.fc1(hidden_state)
-        hidden_state = torch.nn.functional.gelu(hidden_state)
-        hidden_state = self.fc2(hidden_state)
-        return hidden_state
 
 class SwiGLUFFN(torch.nn.Module):
     def __init__(self, dim, dtype, device, operations):
@@ -64,15 +50,12 @@ class SwiGLUFFN(torch.nn.Module):
 
 
 class Dino2Block(torch.nn.Module):
-    def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations, use_swiglu_ffn):
+    def __init__(self, dim, num_heads, layer_norm_eps, dtype, device, operations):
         super().__init__()
         self.attention = Dino2AttentionBlock(dim, num_heads, layer_norm_eps, dtype, device, operations)
         self.layer_scale1 = LayerScale(dim, dtype, device, operations)
         self.layer_scale2 = LayerScale(dim, dtype, device, operations)
-        if use_swiglu_ffn:
-            self.mlp = SwiGLUFFN(dim, dtype, device, operations)
-        else:
-            self.mlp = Dinov2MLP(dim, dtype, device, operations)
+        self.mlp = SwiGLUFFN(dim, dtype, device, operations)
         self.norm1 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
         self.norm2 = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
 
@@ -83,10 +66,9 @@ class Dino2Block(torch.nn.Module):
 
 
 class Dino2Encoder(torch.nn.Module):
-    def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations, use_swiglu_ffn):
+    def __init__(self, dim, num_heads, layer_norm_eps, num_layers, dtype, device, operations):
         super().__init__()
-        self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations, use_swiglu_ffn = use_swiglu_ffn)
-                                          for _ in range(num_layers)])
+        self.layer = torch.nn.ModuleList([Dino2Block(dim, num_heads, layer_norm_eps, dtype, device, operations) for _ in range(num_layers)])
 
     def forward(self, x, intermediate_output=None):
         optimized_attention = optimized_attention_for_device(x.device, False, small_input=True)
@@ -96,8 +78,8 @@ class Dino2Encoder(torch.nn.Module):
                 intermediate_output = len(self.layer) + intermediate_output
 
         intermediate = None
-        for i, layer in enumerate(self.layer):
-            x = layer(x, optimized_attention)
+        for i, l in enumerate(self.layer):
+            x = l(x, optimized_attention)
             if i == intermediate_output:
                 intermediate = x.clone()
         return x, intermediate
@@ -146,10 +128,9 @@ class Dinov2Model(torch.nn.Module):
         dim = config_dict["hidden_size"]
         heads = config_dict["num_attention_heads"]
         layer_norm_eps = config_dict["layer_norm_eps"]
-        use_swiglu_ffn = config_dict["use_swiglu_ffn"]
 
         self.embeddings = Dino2Embeddings(dim, dtype, device, operations)
-        self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations, use_swiglu_ffn = use_swiglu_ffn)
+        self.encoder = Dino2Encoder(dim, heads, layer_norm_eps, num_layers, dtype, device, operations)
         self.layernorm = operations.LayerNorm(dim, eps=layer_norm_eps, dtype=dtype, device=device)
 
     def forward(self, pixel_values, attention_mask=None, intermediate_output=None):

comfy/image_encoders/dino2_large.json

@@ -1,22 +0,0 @@
-{
-  "hidden_size": 1024,
-  "use_mask_token": true,
-  "patch_size": 14,
-  "image_size": 518,
-  "num_channels": 3,
-  "num_attention_heads": 16,
-  "initializer_range": 0.02,
-  "attention_probs_dropout_prob": 0.0,
-  "hidden_dropout_prob": 0.0,
-  "hidden_act": "gelu",
-  "mlp_ratio": 4,
-  "model_type": "dinov2",
-  "num_hidden_layers": 24,
-  "layer_norm_eps": 1e-6,
-  "qkv_bias": true,
-  "use_swiglu_ffn": false,
-  "layerscale_value": 1.0,
-  "drop_path_rate": 0.0,
-  "image_mean": [0.485, 0.456, 0.406],
-  "image_std": [0.229, 0.224, 0.225]
-}

comfy/k_diffusion/sampling.py

@@ -86,24 +86,24 @@ class BatchedBrownianTree:
     """A wrapper around torchsde.BrownianTree that enables batches of entropy."""
 
     def __init__(self, x, t0, t1, seed=None, **kwargs):
-        self.cpu_tree = kwargs.pop("cpu", True)
+        self.cpu_tree = True
+        if "cpu" in kwargs:
+            self.cpu_tree = kwargs.pop("cpu")
         t0, t1, self.sign = self.sort(t0, t1)
-        w0 = kwargs.pop('w0', None)
-        if w0 is None:
-            w0 = torch.zeros_like(x)
-        self.batched = False
+        w0 = kwargs.get('w0', torch.zeros_like(x))
         if seed is None:
-            seed = (torch.randint(0, 2 ** 63 - 1, ()).item(),)
-        elif isinstance(seed, (tuple, list)):
-            if len(seed) != x.shape[0]:
-                raise ValueError("Passing a list or tuple of seeds to BatchedBrownianTree requires a length matching the batch size.")
-            self.batched = True
+            seed = torch.randint(0, 2 ** 63 - 1, []).item()
+        self.batched = True
+        try:
+            assert len(seed) == x.shape[0]
             w0 = w0[0]
-        else:
-            seed = (seed,)
+        except TypeError:
+            seed = [seed]
+            self.batched = False
         if self.cpu_tree:
-            t0, w0, t1 = t0.detach().cpu(), w0.detach().cpu(), t1.detach().cpu()
-        self.trees = tuple(torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed)
+            self.trees = [torchsde.BrownianTree(t0.cpu(), w0.cpu(), t1.cpu(), entropy=s, **kwargs) for s in seed]
+        else:
+            self.trees = [torchsde.BrownianTree(t0, w0, t1, entropy=s, **kwargs) for s in seed]
 
     @staticmethod
     def sort(a, b):
@@ -111,10 +111,11 @@ class BatchedBrownianTree:
 
     def __call__(self, t0, t1):
         t0, t1, sign = self.sort(t0, t1)
-        device, dtype = t0.device, t0.dtype
         if self.cpu_tree:
-            t0, t1 = t0.detach().cpu().float(), t1.detach().cpu().float()
-        w = torch.stack([tree(t0, t1) for tree in self.trees]).to(device=device, dtype=dtype) * (self.sign * sign)
+            w = torch.stack([tree(t0.cpu().float(), t1.cpu().float()).to(t0.dtype).to(t0.device) for tree in self.trees]) * (self.sign * sign)
+        else:
+            w = torch.stack([tree(t0, t1) for tree in self.trees]) * (self.sign * sign)
+
         return w if self.batched else w[0]
 
 
@@ -170,16 +171,6 @@ def offset_first_sigma_for_snr(sigmas, model_sampling, percent_offset=1e-4):
     return sigmas
 
 
-def ei_h_phi_1(h: torch.Tensor) -> torch.Tensor:
-    """Compute the result of h*phi_1(h) in exponential integrator methods."""
-    return torch.expm1(h)
-
-
-def ei_h_phi_2(h: torch.Tensor) -> torch.Tensor:
-    """Compute the result of h*phi_2(h) in exponential integrator methods."""
-    return (torch.expm1(h) - h) / h
-
-
 @torch.no_grad()
 def sample_euler(model, x, sigmas, extra_args=None, callback=None, disable=None, s_churn=0., s_tmin=0., s_tmax=float('inf'), s_noise=1.):
     """Implements Algorithm 2 (Euler steps) from Karras et al. (2022)."""
@@ -1559,12 +1550,13 @@ def sample_er_sde(model, x, sigmas, extra_args=None, callback=None, disable=None
 @torch.no_grad()
 def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r=0.5):
     """SEEDS-2 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 2.
-    arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
+    arXiv: https://arxiv.org/abs/2305.14267
     """
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
     noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
+
     inject_noise = eta > 0 and s_noise > 0
 
     model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
@@ -1572,53 +1564,55 @@ def sample_seeds_2(model, x, sigmas, extra_args=None, callback=None, disable=Non
     lambda_fn = partial(sigma_to_half_log_snr, model_sampling=model_sampling)
     sigmas = offset_first_sigma_for_snr(sigmas, model_sampling)
 
-    fac = 1 / (2 * r)
-
     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})
-
         if sigmas[i + 1] == 0:
             x = denoised
-            continue
+        else:
+            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+            h = lambda_t - lambda_s
+            h_eta = h * (eta + 1)
+            lambda_s_1 = lambda_s + r * h
+            fac = 1 / (2 * r)
+            sigma_s_1 = sigma_fn(lambda_s_1)
 
-        lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-        h = lambda_t - lambda_s
-        h_eta = h * (eta + 1)
-        lambda_s_1 = torch.lerp(lambda_s, lambda_t, r)
-        sigma_s_1 = sigma_fn(lambda_s_1)
+            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
+            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+            alpha_t = sigmas[i + 1] * lambda_t.exp()
 
-        alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-        alpha_t = sigmas[i + 1] * lambda_t.exp()
+            coeff_1, coeff_2 = (-r * h_eta).expm1(), (-h_eta).expm1()
+            if inject_noise:
+                # 0 < r < 1
+                noise_coeff_1 = (-2 * r * h * eta).expm1().neg().sqrt()
+                noise_coeff_2 = (-r * h * eta).exp() * (-2 * (1 - r) * h * eta).expm1().neg().sqrt()
+                noise_1, noise_2 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigmas[i + 1])
 
-        # Step 1
-        x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * ei_h_phi_1(-r * h_eta) * denoised
-        if inject_noise:
-            sde_noise = (-2 * r * h * eta).expm1().neg().sqrt() * noise_sampler(sigmas[i], sigma_s_1)
-            x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
-        denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+            # Step 1
+            x_2 = sigma_s_1 / sigmas[i] * (-r * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
+            if inject_noise:
+                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
+            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
-        # Step 2
-        denoised_d = torch.lerp(denoised, denoised_2, fac)
-        x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * ei_h_phi_1(-h_eta) * denoised_d
-        if inject_noise:
-            segment_factor = (r - 1) * h * eta
-            sde_noise = sde_noise * segment_factor.exp()
-            sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_1, sigmas[i + 1])
-            x = x + sde_noise * sigmas[i + 1] * s_noise
+            # Step 2
+            denoised_d = (1 - fac) * denoised + fac * denoised_2
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_2 * denoised_d
+            if inject_noise:
+                x = x + sigmas[i + 1] * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
     return x
 
 
 @torch.no_grad()
 def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=None, eta=1., s_noise=1., noise_sampler=None, r_1=1./3, r_2=2./3):
     """SEEDS-3 - Stochastic Explicit Exponential Derivative-free Solvers (VP Data Prediction) stage 3.
-    arXiv: https://arxiv.org/abs/2305.14267 (NeurIPS 2023)
+    arXiv: https://arxiv.org/abs/2305.14267
     """
     extra_args = {} if extra_args is None else extra_args
     seed = extra_args.get("seed", None)
     noise_sampler = default_noise_sampler(x, seed=seed) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
+
     inject_noise = eta > 0 and s_noise > 0
 
     model_sampling = model.inner_model.model_patcher.get_model_object('model_sampling')
@@ -1630,49 +1624,45 @@ def sample_seeds_3(model, x, sigmas, extra_args=None, callback=None, disable=Non
         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})
-
         if sigmas[i + 1] == 0:
             x = denoised
-            continue
+        else:
+            lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
+            h = lambda_t - lambda_s
+            h_eta = h * (eta + 1)
+            lambda_s_1 = lambda_s + r_1 * h
+            lambda_s_2 = lambda_s + r_2 * h
+            sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
 
-        lambda_s, lambda_t = lambda_fn(sigmas[i]), lambda_fn(sigmas[i + 1])
-        h = lambda_t - lambda_s
-        h_eta = h * (eta + 1)
-        lambda_s_1 = torch.lerp(lambda_s, lambda_t, r_1)
-        lambda_s_2 = torch.lerp(lambda_s, lambda_t, r_2)
-        sigma_s_1, sigma_s_2 = sigma_fn(lambda_s_1), sigma_fn(lambda_s_2)
+            # alpha_t = sigma_t * exp(log(alpha_t / sigma_t)) = sigma_t * exp(lambda_t)
+            alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
+            alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
+            alpha_t = sigmas[i + 1] * lambda_t.exp()
 
-        alpha_s_1 = sigma_s_1 * lambda_s_1.exp()
-        alpha_s_2 = sigma_s_2 * lambda_s_2.exp()
-        alpha_t = sigmas[i + 1] * lambda_t.exp()
+            coeff_1, coeff_2, coeff_3 = (-r_1 * h_eta).expm1(), (-r_2 * h_eta).expm1(), (-h_eta).expm1()
+            if inject_noise:
+                # 0 < r_1 < r_2 < 1
+                noise_coeff_1 = (-2 * r_1 * h * eta).expm1().neg().sqrt()
+                noise_coeff_2 = (-r_1 * h * eta).exp() * (-2 * (r_2 - r_1) * h * eta).expm1().neg().sqrt()
+                noise_coeff_3 = (-r_2 * h * eta).exp() * (-2 * (1 - r_2) * h * eta).expm1().neg().sqrt()
+                noise_1, noise_2, noise_3 = noise_sampler(sigmas[i], sigma_s_1), noise_sampler(sigma_s_1, sigma_s_2), noise_sampler(sigma_s_2, sigmas[i + 1])
 
-        # Step 1
-        x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * ei_h_phi_1(-r_1 * h_eta) * denoised
-        if inject_noise:
-            sde_noise = (-2 * r_1 * h * eta).expm1().neg().sqrt() * noise_sampler(sigmas[i], sigma_s_1)
-            x_2 = x_2 + sde_noise * sigma_s_1 * s_noise
-        denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
+            # Step 1
+            x_2 = sigma_s_1 / sigmas[i] * (-r_1 * h * eta).exp() * x - alpha_s_1 * coeff_1 * denoised
+            if inject_noise:
+                x_2 = x_2 + sigma_s_1 * (noise_coeff_1 * noise_1) * s_noise
+            denoised_2 = model(x_2, sigma_s_1 * s_in, **extra_args)
 
-        # Step 2
-        a3_2 = r_2 / r_1 * ei_h_phi_2(-r_2 * h_eta)
-        a3_1 = ei_h_phi_1(-r_2 * h_eta) - a3_2
-        x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * (a3_1 * denoised + a3_2 * denoised_2)
-        if inject_noise:
-            segment_factor = (r_1 - r_2) * h * eta
-            sde_noise = sde_noise * segment_factor.exp()
-            sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_1, sigma_s_2)
-            x_3 = x_3 + sde_noise * sigma_s_2 * s_noise
-        denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
+            # Step 2
+            x_3 = sigma_s_2 / sigmas[i] * (-r_2 * h * eta).exp() * x - alpha_s_2 * coeff_2 * denoised + (r_2 / r_1) * alpha_s_2 * (coeff_2 / (r_2 * h_eta) + 1) * (denoised_2 - denoised)
+            if inject_noise:
+                x_3 = x_3 + sigma_s_2 * (noise_coeff_2 * noise_1 + noise_coeff_1 * noise_2) * s_noise
+            denoised_3 = model(x_3, sigma_s_2 * s_in, **extra_args)
 
-        # Step 3
-        b3 = ei_h_phi_2(-h_eta) / r_2
-        b1 = ei_h_phi_1(-h_eta) - b3
-        x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * (b1 * denoised + b3 * denoised_3)
-        if inject_noise:
-            segment_factor = (r_2 - 1) * h * eta
-            sde_noise = sde_noise * segment_factor.exp()
-            sde_noise = sde_noise + segment_factor.mul(2).expm1().neg().sqrt() * noise_sampler(sigma_s_2, sigmas[i + 1])
-            x = x + sde_noise * sigmas[i + 1] * s_noise
+            # Step 3
+            x = sigmas[i + 1] / sigmas[i] * (-h * eta).exp() * x - alpha_t * coeff_3 * denoised + (1. / r_2) * alpha_t * (coeff_3 / h_eta + 1) * (denoised_3 - denoised)
+            if inject_noise:
+                x = x + sigmas[i + 1] * (noise_coeff_3 * noise_1 + noise_coeff_2 * noise_2 + noise_coeff_1 * noise_3) * s_noise
     return x
 
 

comfy/latent_formats.py

@@ -533,94 +533,11 @@ class Wan22(Wan21):
                 0.3971, 1.0600, 0.3943, 0.5537, 0.5444, 0.4089, 0.7468, 0.7744
             ]).view(1, self.latent_channels, 1, 1, 1)
 
-class HunyuanImage21(LatentFormat):
-    latent_channels = 64
-    latent_dimensions = 2
-    scale_factor = 0.75289
-
-    latent_rgb_factors = [
-        [-0.0154, -0.0397, -0.0521],
-        [ 0.0005,  0.0093,  0.0006],
-        [-0.0805, -0.0773, -0.0586],
-        [-0.0494, -0.0487, -0.0498],
-        [-0.0212, -0.0076, -0.0261],
-        [-0.0179, -0.0417, -0.0505],
-        [ 0.0158,  0.0310,  0.0239],
-        [ 0.0409,  0.0516,  0.0201],
-        [ 0.0350,  0.0553,  0.0036],
-        [-0.0447, -0.0327, -0.0479],
-        [-0.0038, -0.0221, -0.0365],
-        [-0.0423, -0.0718, -0.0654],
-        [ 0.0039,  0.0368,  0.0104],
-        [ 0.0655,  0.0217,  0.0122],
-        [ 0.0490,  0.1638,  0.2053],
-        [ 0.0932,  0.0829,  0.0650],
-        [-0.0186, -0.0209, -0.0135],
-        [-0.0080, -0.0076, -0.0148],
-        [-0.0284, -0.0201,  0.0011],
-        [-0.0642, -0.0294, -0.0777],
-        [-0.0035,  0.0076, -0.0140],
-        [ 0.0519,  0.0731,  0.0887],
-        [-0.0102,  0.0095,  0.0704],
-        [ 0.0068,  0.0218, -0.0023],
-        [-0.0726, -0.0486, -0.0519],
-        [ 0.0260,  0.0295,  0.0263],
-        [ 0.0250,  0.0333,  0.0341],
-        [ 0.0168, -0.0120, -0.0174],
-        [ 0.0226,  0.1037,  0.0114],
-        [ 0.2577,  0.1906,  0.1604],
-        [-0.0646, -0.0137, -0.0018],
-        [-0.0112,  0.0309,  0.0358],
-        [-0.0347,  0.0146, -0.0481],
-        [ 0.0234,  0.0179,  0.0201],
-        [ 0.0157,  0.0313,  0.0225],
-        [ 0.0423,  0.0675,  0.0524],
-        [-0.0031,  0.0027, -0.0255],
-        [ 0.0447,  0.0555,  0.0330],
-        [-0.0152,  0.0103,  0.0299],
-        [-0.0755, -0.0489, -0.0635],
-        [ 0.0853,  0.0788,  0.1017],
-        [-0.0272, -0.0294, -0.0471],
-        [ 0.0440,  0.0400, -0.0137],
-        [ 0.0335,  0.0317, -0.0036],
-        [-0.0344, -0.0621, -0.0984],
-        [-0.0127, -0.0630, -0.0620],
-        [-0.0648,  0.0360,  0.0924],
-        [-0.0781, -0.0801, -0.0409],
-        [ 0.0363,  0.0613,  0.0499],
-        [ 0.0238,  0.0034,  0.0041],
-        [-0.0135,  0.0258,  0.0310],
-        [ 0.0614,  0.1086,  0.0589],
-        [ 0.0428,  0.0350,  0.0205],
-        [ 0.0153,  0.0173, -0.0018],
-        [-0.0288, -0.0455, -0.0091],
-        [ 0.0344,  0.0109, -0.0157],
-        [-0.0205, -0.0247, -0.0187],
-        [ 0.0487,  0.0126,  0.0064],
-        [-0.0220, -0.0013,  0.0074],
-        [-0.0203, -0.0094, -0.0048],
-        [-0.0719,  0.0429, -0.0442],
-        [ 0.1042,  0.0497,  0.0356],
-        [-0.0659, -0.0578, -0.0280],
-        [-0.0060, -0.0322, -0.0234]]
-
-    latent_rgb_factors_bias = [0.0007, -0.0256, -0.0206]
-
-class HunyuanImage21Refiner(LatentFormat):
-    latent_channels = 64
-    latent_dimensions = 3
-    scale_factor = 1.03682
-
 class Hunyuan3Dv2(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1
     scale_factor = 0.9990943042622529
 
-class Hunyuan3Dv2_1(LatentFormat):
-    scale_factor = 1.0039506158752403
-    latent_channels = 64
-    latent_dimensions = 1
-
 class Hunyuan3Dv2mini(LatentFormat):
     latent_channels = 64
     latent_dimensions = 1
@@ -629,20 +546,3 @@ class Hunyuan3Dv2mini(LatentFormat):
 class ACEAudio(LatentFormat):
     latent_channels = 8
     latent_dimensions = 2
-
-class ChromaRadiance(LatentFormat):
-    latent_channels = 3
-
-    def __init__(self):
-        self.latent_rgb_factors = [
-            # R    G    B
-            [ 1.0, 0.0, 0.0 ],
-            [ 0.0, 1.0, 0.0 ],
-            [ 0.0, 0.0, 1.0 ]
-        ]
-
-    def process_in(self, latent):
-        return latent
-
-    def process_out(self, latent):
-        return latent

comfy/ldm/audio/dit.py

@@ -635,7 +635,7 @@ class ContinuousTransformer(nn.Module):
         # Attention layers
 
         if self.rotary_pos_emb is not None:
-            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=torch.float, device=x.device)
+            rotary_pos_emb = self.rotary_pos_emb.forward_from_seq_len(x.shape[1], dtype=x.dtype, device=x.device)
         else:
             rotary_pos_emb = None
 

comfy/ldm/chroma/model.py

@@ -151,6 +151,8 @@ class Chroma(nn.Module):
         attn_mask: Tensor = None,
     ) -> Tensor:
         patches_replace = transformer_options.get("patches_replace", {})
+        if img.ndim != 3 or txt.ndim != 3:
+            raise ValueError("Input img and txt tensors must have 3 dimensions.")
 
         # running on sequences img
         img = self.img_in(img)
@@ -252,9 +254,8 @@ class Chroma(nn.Module):
                             img[:, txt.shape[1] :, ...] += add
 
         img = img[:, txt.shape[1] :, ...]
-        if hasattr(self, "final_layer"):
-            final_mod = self.get_modulations(mod_vectors, "final")
-            img = self.final_layer(img, vec=final_mod)  # (N, T, patch_size ** 2 * out_channels)
+        final_mod = self.get_modulations(mod_vectors, "final")
+        img = self.final_layer(img, vec=final_mod)  # (N, T, patch_size ** 2 * out_channels)
         return img
 
     def forward(self, x, timestep, context, guidance, control=None, transformer_options={}, **kwargs):
@@ -270,9 +271,6 @@ class Chroma(nn.Module):
 
         img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=self.patch_size, pw=self.patch_size)
 
-        if img.ndim != 3 or context.ndim != 3:
-            raise ValueError("Input img and txt tensors must have 3 dimensions.")
-
         h_len = ((h + (self.patch_size // 2)) // self.patch_size)
         w_len = ((w + (self.patch_size // 2)) // self.patch_size)
         img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)

comfy/ldm/chroma_radiance/layers.py

@@ -1,206 +0,0 @@
-# Adapted from https://github.com/lodestone-rock/flow
-from functools import lru_cache
-
-import torch
-from torch import nn
-
-from comfy.ldm.flux.layers import RMSNorm
-
-
-class NerfEmbedder(nn.Module):
-    """
-    An embedder module that combines input features with a 2D positional
-    encoding that mimics the Discrete Cosine Transform (DCT).
-
-    This module takes an input tensor of shape (B, P^2, C), where P is the
-    patch size, and enriches it with positional information before projecting
-    it to a new hidden size.
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        hidden_size_input: int,
-        max_freqs: int,
-        dtype=None,
-        device=None,
-        operations=None,
-    ):
-        """
-        Initializes the NerfEmbedder.
-
-        Args:
-            in_channels (int): The number of channels in the input tensor.
-            hidden_size_input (int): The desired dimension of the output embedding.
-            max_freqs (int): The number of frequency components to use for both
-                             the x and y dimensions of the positional encoding.
-                             The total number of positional features will be max_freqs^2.
-        """
-        super().__init__()
-        self.dtype = dtype
-        self.max_freqs = max_freqs
-        self.hidden_size_input = hidden_size_input
-
-        # A linear layer to project the concatenated input features and
-        # positional encodings to the final output dimension.
-        self.embedder = nn.Sequential(
-            operations.Linear(in_channels + max_freqs**2, hidden_size_input, dtype=dtype, device=device)
-        )
-
-    @lru_cache(maxsize=4)
-    def fetch_pos(self, patch_size: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
-        """
-        Generates and caches 2D DCT-like positional embeddings for a given patch size.
-
-        The LRU cache is a performance optimization that avoids recomputing the
-        same positional grid on every forward pass.
-
-        Args:
-            patch_size (int): The side length of the square input patch.
-            device: The torch device to create the tensors on.
-            dtype: The torch dtype for the tensors.
-
-        Returns:
-            A tensor of shape (1, patch_size^2, max_freqs^2) containing the
-            positional embeddings.
-        """
-        # Create normalized 1D coordinate grids from 0 to 1.
-        pos_x = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
-        pos_y = torch.linspace(0, 1, patch_size, device=device, dtype=dtype)
-
-        # Create a 2D meshgrid of coordinates.
-        pos_y, pos_x = torch.meshgrid(pos_y, pos_x, indexing="ij")
-
-        # Reshape positions to be broadcastable with frequencies.
-        # Shape becomes (patch_size^2, 1, 1).
-        pos_x = pos_x.reshape(-1, 1, 1)
-        pos_y = pos_y.reshape(-1, 1, 1)
-
-        # Create a 1D tensor of frequency values from 0 to max_freqs-1.
-        freqs = torch.linspace(0, self.max_freqs - 1, self.max_freqs, dtype=dtype, device=device)
-
-        # Reshape frequencies to be broadcastable for creating 2D basis functions.
-        # freqs_x shape: (1, max_freqs, 1)
-        # freqs_y shape: (1, 1, max_freqs)
-        freqs_x = freqs[None, :, None]
-        freqs_y = freqs[None, None, :]
-
-        # A custom weighting coefficient, not part of standard DCT.
-        # This seems to down-weight the contribution of higher-frequency interactions.
-        coeffs = (1 + freqs_x * freqs_y) ** -1
-
-        # Calculate the 1D cosine basis functions for x and y coordinates.
-        # This is the core of the DCT formulation.
-        dct_x = torch.cos(pos_x * freqs_x * torch.pi)
-        dct_y = torch.cos(pos_y * freqs_y * torch.pi)
-
-        # Combine the 1D basis functions to create 2D basis functions by element-wise
-        # multiplication, and apply the custom coefficients. Broadcasting handles the
-        # combination of all (pos_x, freqs_x) with all (pos_y, freqs_y).
-        # The result is flattened into a feature vector for each position.
-        dct = (dct_x * dct_y * coeffs).view(1, -1, self.max_freqs ** 2)
-
-        return dct
-
-    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
-        """
-        Forward pass for the embedder.
-
-        Args:
-            inputs (Tensor): The input tensor of shape (B, P^2, C).
-
-        Returns:
-            Tensor: The output tensor of shape (B, P^2, hidden_size_input).
-        """
-        # Get the batch size, number of pixels, and number of channels.
-        B, P2, C = inputs.shape
-
-        # Infer the patch side length from the number of pixels (P^2).
-        patch_size = int(P2 ** 0.5)
-
-        input_dtype = inputs.dtype
-        inputs = inputs.to(dtype=self.dtype)
-
-        # Fetch the pre-computed or cached positional embeddings.
-        dct = self.fetch_pos(patch_size, inputs.device, self.dtype)
-
-        # Repeat the positional embeddings for each item in the batch.
-        dct = dct.repeat(B, 1, 1)
-
-        # Concatenate the original input features with the positional embeddings
-        # along the feature dimension.
-        inputs = torch.cat((inputs, dct), dim=-1)
-
-        # Project the combined tensor to the target hidden size.
-        return self.embedder(inputs).to(dtype=input_dtype)
-
-
-class NerfGLUBlock(nn.Module):
-    """
-    A NerfBlock using a Gated Linear Unit (GLU) like MLP.
-    """
-    def __init__(self, hidden_size_s: int, hidden_size_x: int, mlp_ratio, dtype=None, device=None, operations=None):
-        super().__init__()
-        # The total number of parameters for the MLP is increased to accommodate
-        # the gate, value, and output projection matrices.
-        # We now need to generate parameters for 3 matrices.
-        total_params = 3 * hidden_size_x**2 * mlp_ratio
-        self.param_generator = operations.Linear(hidden_size_s, total_params, dtype=dtype, device=device)
-        self.norm = RMSNorm(hidden_size_x, dtype=dtype, device=device, operations=operations)
-        self.mlp_ratio = mlp_ratio
-
-
-    def forward(self, x: torch.Tensor, s: torch.Tensor) -> torch.Tensor:
-        batch_size, num_x, hidden_size_x = x.shape
-        mlp_params = self.param_generator(s)
-
-        # Split the generated parameters into three parts for the gate, value, and output projection.
-        fc1_gate_params, fc1_value_params, fc2_params = mlp_params.chunk(3, dim=-1)
-
-        # Reshape the parameters into matrices for batch matrix multiplication.
-        fc1_gate = fc1_gate_params.view(batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio)
-        fc1_value = fc1_value_params.view(batch_size, hidden_size_x, hidden_size_x * self.mlp_ratio)
-        fc2 = fc2_params.view(batch_size, hidden_size_x * self.mlp_ratio, hidden_size_x)
-
-        # Normalize the generated weight matrices as in the original implementation.
-        fc1_gate = torch.nn.functional.normalize(fc1_gate, dim=-2)
-        fc1_value = torch.nn.functional.normalize(fc1_value, dim=-2)
-        fc2 = torch.nn.functional.normalize(fc2, dim=-2)
-
-        res_x = x
-        x = self.norm(x)
-
-        # Apply the final output projection.
-        x = torch.bmm(torch.nn.functional.silu(torch.bmm(x, fc1_gate)) * torch.bmm(x, fc1_value), fc2)
-
-        return x + res_x
-
-
-class NerfFinalLayer(nn.Module):
-    def __init__(self, hidden_size, out_channels, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.norm = RMSNorm(hidden_size, dtype=dtype, device=device, operations=operations)
-        self.linear = operations.Linear(hidden_size, out_channels, dtype=dtype, device=device)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # RMSNorm normalizes over the last dimension, but our channel dim (C) is at dim=1.
-        # So we temporarily move the channel dimension to the end for the norm operation.
-        return self.linear(self.norm(x.movedim(1, -1))).movedim(-1, 1)
-
-
-class NerfFinalLayerConv(nn.Module):
-    def __init__(self, hidden_size: int, out_channels: int, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.norm = RMSNorm(hidden_size, dtype=dtype, device=device, operations=operations)
-        self.conv = operations.Conv2d(
-            in_channels=hidden_size,
-            out_channels=out_channels,
-            kernel_size=3,
-            padding=1,
-            dtype=dtype,
-            device=device,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # RMSNorm normalizes over the last dimension, but our channel dim (C) is at dim=1.
-        # So we temporarily move the channel dimension to the end for the norm operation.
-        return self.conv(self.norm(x.movedim(1, -1)).movedim(-1, 1))

comfy/ldm/chroma_radiance/model.py

@@ -1,329 +0,0 @@
-# Credits:
-# Original Flux code can be found on: https://github.com/black-forest-labs/flux
-# Chroma Radiance adaption referenced from https://github.com/lodestone-rock/flow
-
-from dataclasses import dataclass
-from typing import Optional
-
-import torch
-from torch import Tensor, nn
-from einops import repeat
-import comfy.ldm.common_dit
-
-from comfy.ldm.flux.layers import EmbedND
-
-from comfy.ldm.chroma.model import Chroma, ChromaParams
-from comfy.ldm.chroma.layers import (
-    DoubleStreamBlock,
-    SingleStreamBlock,
-    Approximator,
-)
-from .layers import (
-    NerfEmbedder,
-    NerfGLUBlock,
-    NerfFinalLayer,
-    NerfFinalLayerConv,
-)
-
-
-@dataclass
-class ChromaRadianceParams(ChromaParams):
-    patch_size: int
-    nerf_hidden_size: int
-    nerf_mlp_ratio: int
-    nerf_depth: int
-    nerf_max_freqs: int
-    # Setting nerf_tile_size to 0 disables tiling.
-    nerf_tile_size: int
-    # Currently one of linear (legacy) or conv.
-    nerf_final_head_type: str
-    # None means use the same dtype as the model.
-    nerf_embedder_dtype: Optional[torch.dtype]
-
-
-class ChromaRadiance(Chroma):
-    """
-    Transformer model for flow matching on sequences.
-    """
-
-    def __init__(self, image_model=None, final_layer=True, dtype=None, device=None, operations=None, **kwargs):
-        if operations is None:
-            raise RuntimeError("Attempt to create ChromaRadiance object without setting operations")
-        nn.Module.__init__(self)
-        self.dtype = dtype
-        params = ChromaRadianceParams(**kwargs)
-        self.params = params
-        self.patch_size = params.patch_size
-        self.in_channels = params.in_channels
-        self.out_channels = params.out_channels
-        if params.hidden_size % params.num_heads != 0:
-            raise ValueError(
-                f"Hidden size {params.hidden_size} must be divisible by num_heads {params.num_heads}"
-            )
-        pe_dim = params.hidden_size // params.num_heads
-        if sum(params.axes_dim) != pe_dim:
-            raise ValueError(f"Got {params.axes_dim} but expected positional dim {pe_dim}")
-        self.hidden_size = params.hidden_size
-        self.num_heads = params.num_heads
-        self.in_dim = params.in_dim
-        self.out_dim = params.out_dim
-        self.hidden_dim = params.hidden_dim
-        self.n_layers = params.n_layers
-        self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
-        self.img_in_patch = operations.Conv2d(
-            params.in_channels,
-            params.hidden_size,
-            kernel_size=params.patch_size,
-            stride=params.patch_size,
-            bias=True,
-            dtype=dtype,
-            device=device,
-        )
-        self.txt_in = operations.Linear(params.context_in_dim, self.hidden_size, dtype=dtype, device=device)
-        # set as nn identity for now, will overwrite it later.
-        self.distilled_guidance_layer = Approximator(
-                    in_dim=self.in_dim,
-                    hidden_dim=self.hidden_dim,
-                    out_dim=self.out_dim,
-                    n_layers=self.n_layers,
-                    dtype=dtype, device=device, operations=operations
-                )
-
-
-        self.double_blocks = nn.ModuleList(
-            [
-                DoubleStreamBlock(
-                    self.hidden_size,
-                    self.num_heads,
-                    mlp_ratio=params.mlp_ratio,
-                    qkv_bias=params.qkv_bias,
-                    dtype=dtype, device=device, operations=operations
-                )
-                for _ in range(params.depth)
-            ]
-        )
-
-        self.single_blocks = nn.ModuleList(
-            [
-                SingleStreamBlock(
-                    self.hidden_size,
-                    self.num_heads,
-                    mlp_ratio=params.mlp_ratio,
-                    dtype=dtype, device=device, operations=operations,
-                )
-                for _ in range(params.depth_single_blocks)
-            ]
-        )
-
-        # pixel channel concat with DCT
-        self.nerf_image_embedder = NerfEmbedder(
-            in_channels=params.in_channels,
-            hidden_size_input=params.nerf_hidden_size,
-            max_freqs=params.nerf_max_freqs,
-            dtype=params.nerf_embedder_dtype or dtype,
-            device=device,
-            operations=operations,
-        )
-
-        self.nerf_blocks = nn.ModuleList([
-            NerfGLUBlock(
-                hidden_size_s=params.hidden_size,
-                hidden_size_x=params.nerf_hidden_size,
-                mlp_ratio=params.nerf_mlp_ratio,
-                dtype=dtype,
-                device=device,
-                operations=operations,
-            ) for _ in range(params.nerf_depth)
-        ])
-
-        if params.nerf_final_head_type == "linear":
-            self.nerf_final_layer = NerfFinalLayer(
-                params.nerf_hidden_size,
-                out_channels=params.in_channels,
-                dtype=dtype,
-                device=device,
-                operations=operations,
-            )
-        elif params.nerf_final_head_type == "conv":
-            self.nerf_final_layer_conv = NerfFinalLayerConv(
-                params.nerf_hidden_size,
-                out_channels=params.in_channels,
-                dtype=dtype,
-                device=device,
-                operations=operations,
-            )
-        else:
-            errstr = f"Unsupported nerf_final_head_type {params.nerf_final_head_type}"
-            raise ValueError(errstr)
-
-        self.skip_mmdit = []
-        self.skip_dit = []
-        self.lite = False
-
-    @property
-    def _nerf_final_layer(self) -> nn.Module:
-        if self.params.nerf_final_head_type == "linear":
-            return self.nerf_final_layer
-        if self.params.nerf_final_head_type == "conv":
-            return self.nerf_final_layer_conv
-        # Impossible to get here as we raise an error on unexpected types on initialization.
-        raise NotImplementedError
-
-    def img_in(self, img: Tensor) -> Tensor:
-        img = self.img_in_patch(img) # -> [B, Hidden, H/P, W/P]
-        # flatten into a sequence for the transformer.
-        return img.flatten(2).transpose(1, 2) # -> [B, NumPatches, Hidden]
-
-    def forward_nerf(
-        self,
-        img_orig: Tensor,
-        img_out: Tensor,
-        params: ChromaRadianceParams,
-    ) -> Tensor:
-        B, C, H, W = img_orig.shape
-        num_patches = img_out.shape[1]
-        patch_size = params.patch_size
-
-        # Store the raw pixel values of each patch for the NeRF head later.
-        # unfold creates patches: [B, C * P * P, NumPatches]
-        nerf_pixels = nn.functional.unfold(img_orig, kernel_size=patch_size, stride=patch_size)
-        nerf_pixels = nerf_pixels.transpose(1, 2) # -> [B, NumPatches, C * P * P]
-
-        if params.nerf_tile_size > 0 and num_patches > params.nerf_tile_size:
-            # Enable tiling if nerf_tile_size isn't 0 and we actually have more patches than
-            # the tile size.
-            img_dct = self.forward_tiled_nerf(img_out, nerf_pixels, B, C, num_patches, patch_size, params)
-        else:
-            # Reshape for per-patch processing
-            nerf_hidden = img_out.reshape(B * num_patches, params.hidden_size)
-            nerf_pixels = nerf_pixels.reshape(B * num_patches, C, patch_size**2).transpose(1, 2)
-
-            # Get DCT-encoded pixel embeddings [pixel-dct]
-            img_dct = self.nerf_image_embedder(nerf_pixels)
-
-            # Pass through the dynamic MLP blocks (the NeRF)
-            for block in self.nerf_blocks:
-                img_dct = block(img_dct, nerf_hidden)
-
-        # Reassemble the patches into the final image.
-        img_dct = img_dct.transpose(1, 2) # -> [B*NumPatches, C, P*P]
-        # Reshape to combine with batch dimension for fold
-        img_dct = img_dct.reshape(B, num_patches, -1) # -> [B, NumPatches, C*P*P]
-        img_dct = img_dct.transpose(1, 2) # -> [B, C*P*P, NumPatches]
-        img_dct = nn.functional.fold(
-            img_dct,
-            output_size=(H, W),
-            kernel_size=patch_size,
-            stride=patch_size,
-        )
-        return self._nerf_final_layer(img_dct)
-
-    def forward_tiled_nerf(
-        self,
-        nerf_hidden: Tensor,
-        nerf_pixels: Tensor,
-        batch: int,
-        channels: int,
-        num_patches: int,
-        patch_size: int,
-        params: ChromaRadianceParams,
-    ) -> Tensor:
-        """
-        Processes the NeRF head in tiles to save memory.
-        nerf_hidden has shape [B, L, D]
-        nerf_pixels has shape [B, L, C * P * P]
-        """
-        tile_size = params.nerf_tile_size
-        output_tiles = []
-        # Iterate over the patches in tiles. The dimension L (num_patches) is at index 1.
-        for i in range(0, num_patches, tile_size):
-            end = min(i + tile_size, num_patches)
-
-            # Slice the current tile from the input tensors
-            nerf_hidden_tile = nerf_hidden[:, i:end, :]
-            nerf_pixels_tile = nerf_pixels[:, i:end, :]
-
-            # Get the actual number of patches in this tile (can be smaller for the last tile)
-            num_patches_tile = nerf_hidden_tile.shape[1]
-
-            # Reshape the tile for per-patch processing
-            # [B, NumPatches_tile, D] -> [B * NumPatches_tile, D]
-            nerf_hidden_tile = nerf_hidden_tile.reshape(batch * num_patches_tile, params.hidden_size)
-            # [B, NumPatches_tile, C*P*P] -> [B*NumPatches_tile, C, P*P] -> [B*NumPatches_tile, P*P, C]
-            nerf_pixels_tile = nerf_pixels_tile.reshape(batch * num_patches_tile, channels, patch_size**2).transpose(1, 2)
-
-            # get DCT-encoded pixel embeddings [pixel-dct]
-            img_dct_tile = self.nerf_image_embedder(nerf_pixels_tile)
-
-            # pass through the dynamic MLP blocks (the NeRF)
-            for block in self.nerf_blocks:
-                img_dct_tile = block(img_dct_tile, nerf_hidden_tile)
-
-            output_tiles.append(img_dct_tile)
-
-        # Concatenate the processed tiles along the patch dimension
-        return torch.cat(output_tiles, dim=0)
-
-    def radiance_get_override_params(self, overrides: dict) -> ChromaRadianceParams:
-        params = self.params
-        if not overrides:
-            return params
-        params_dict = {k: getattr(params, k) for k in params.__dataclass_fields__}
-        nullable_keys = frozenset(("nerf_embedder_dtype",))
-        bad_keys = tuple(k for k in overrides if k not in params_dict)
-        if bad_keys:
-            e = f"Unknown key(s) in transformer_options chroma_radiance_options: {', '.join(bad_keys)}"
-            raise ValueError(e)
-        bad_keys = tuple(
-            k
-            for k, v in overrides.items()
-            if type(v) != type(getattr(params, k)) and (v is not None or k not in nullable_keys)
-        )
-        if bad_keys:
-            e = f"Invalid value(s) in transformer_options chroma_radiance_options: {', '.join(bad_keys)}"
-            raise ValueError(e)
-        # At this point it's all valid keys and values so we can merge with the existing params.
-        params_dict |= overrides
-        return params.__class__(**params_dict)
-
-    def _forward(
-        self,
-        x: Tensor,
-        timestep: Tensor,
-        context: Tensor,
-        guidance: Optional[Tensor],
-        control: Optional[dict]=None,
-        transformer_options: dict={},
-        **kwargs: dict,
-    ) -> Tensor:
-        bs, c, h, w = x.shape
-        img = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
-
-        if img.ndim != 4:
-            raise ValueError("Input img tensor must be in [B, C, H, W] format.")
-        if context.ndim != 3:
-            raise ValueError("Input txt tensors must have 3 dimensions.")
-
-        params = self.radiance_get_override_params(transformer_options.get("chroma_radiance_options", {}))
-
-        h_len = (img.shape[-2] // self.patch_size)
-        w_len = (img.shape[-1] // self.patch_size)
-
-        img_ids = torch.zeros((h_len, w_len, 3), device=x.device, dtype=x.dtype)
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
-        img_ids[:, :, 2] = img_ids[:, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
-        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=bs)
-        txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
-
-        img_out = self.forward_orig(
-            img,
-            img_ids,
-            context,
-            txt_ids,
-            timestep,
-            guidance,
-            control,
-            transformer_options,
-            attn_mask=kwargs.get("attention_mask", None),
-        )
-        return self.forward_nerf(img, img_out, params)[:, :, :h, :w]

comfy/ldm/flux/math.py

@@ -35,10 +35,11 @@ def rope(pos: Tensor, dim: int, theta: int) -> Tensor:
     out = rearrange(out, "b n d (i j) -> b n d i j", i=2, j=2)
     return out.to(dtype=torch.float32, device=pos.device)
 
-def apply_rope1(x: Tensor, freqs_cis: Tensor):
-    x_ = x.to(dtype=freqs_cis.dtype).reshape(*x.shape[:-1], -1, 1, 2)
-    x_out = freqs_cis[..., 0] * x_[..., 0] + freqs_cis[..., 1] * x_[..., 1]
-    return x_out.reshape(*x.shape).type_as(x)
 
 def apply_rope(xq: Tensor, xk: Tensor, freqs_cis: Tensor):
-    return apply_rope1(xq, freqs_cis), apply_rope1(xk, freqs_cis)
+    xq_ = xq.to(dtype=freqs_cis.dtype).reshape(*xq.shape[:-1], -1, 1, 2)
+    xk_ = xk.to(dtype=freqs_cis.dtype).reshape(*xk.shape[:-1], -1, 1, 2)
+    xq_out = freqs_cis[..., 0] * xq_[..., 0] + freqs_cis[..., 1] * xq_[..., 1]
+    xk_out = freqs_cis[..., 0] * xk_[..., 0] + freqs_cis[..., 1] * xk_[..., 1]
+    return xq_out.reshape(*xq.shape).type_as(xq), xk_out.reshape(*xk.shape).type_as(xk)
+

comfy/ldm/flux/model.py

@@ -106,7 +106,6 @@ class Flux(nn.Module):
         if y is None:
             y = torch.zeros((img.shape[0], self.params.vec_in_dim), device=img.device, dtype=img.dtype)
 
-        patches = transformer_options.get("patches", {})
         patches_replace = transformer_options.get("patches_replace", {})
         if img.ndim != 3 or txt.ndim != 3:
             raise ValueError("Input img and txt tensors must have 3 dimensions.")
@@ -118,17 +117,9 @@ class Flux(nn.Module):
             if guidance is not None:
                 vec = vec + self.guidance_in(timestep_embedding(guidance, 256).to(img.dtype))
 
-        vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
+        vec = vec + self.vector_in(y[:,:self.params.vec_in_dim])
         txt = self.txt_in(txt)
 
-        if "post_input" in patches:
-            for p in patches["post_input"]:
-                out = p({"img": img, "txt": txt, "img_ids": img_ids, "txt_ids": txt_ids})
-                img = out["img"]
-                txt = out["txt"]
-                img_ids = out["img_ids"]
-                txt_ids = out["txt_ids"]
-
         if img_ids is not None:
             ids = torch.cat((txt_ids, img_ids), dim=1)
             pe = self.pe_embedder(ids)
@@ -137,6 +128,7 @@ class Flux(nn.Module):
 
         blocks_replace = patches_replace.get("dit", {})
         for i, block in enumerate(self.double_blocks):
+            transformer_options["block"] = ("double_block", i, 2)
             if ("double_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -178,6 +170,7 @@ class Flux(nn.Module):
         img = torch.cat((txt, img), 1)
 
         for i, block in enumerate(self.single_blocks):
+            transformer_options["block"] = ("single_block", i, 1)
             if ("single_block", i) in blocks_replace:
                 def block_wrap(args):
                     out = {}
@@ -247,18 +240,12 @@ class Flux(nn.Module):
             h = 0
             w = 0
             index = 0
-            ref_latents_method = kwargs.get("ref_latents_method", "offset")
+            index_ref_method = kwargs.get("ref_latents_method", "offset") == "index"
             for ref in ref_latents:
-                if ref_latents_method == "index":
+                if index_ref_method:
                     index += 1
                     h_offset = 0
                     w_offset = 0
-                elif ref_latents_method == "uxo":
-                    index = 0
-                    h_offset = h_len * patch_size + h
-                    w_offset = w_len * patch_size + w
-                    h += ref.shape[-2]
-                    w += ref.shape[-1]
                 else:
                     index = 1
                     h_offset = 0

comfy/ldm/hunyuan3d/vae.py

@@ -4,458 +4,81 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
+
+from typing import Union, Tuple, List, Callable, Optional
+
 import numpy as np
-import math
+from einops import repeat, rearrange
 from tqdm import tqdm
-
-from typing import Optional
-
 import logging
 
 import comfy.ops
 ops = comfy.ops.disable_weight_init
 
-def fps(src: torch.Tensor, batch: torch.Tensor, sampling_ratio: float, start_random: bool = True):
-
-    # manually create the pointer vector
-    assert src.size(0) == batch.numel()
-
-    batch_size = int(batch.max()) + 1
-    deg = src.new_zeros(batch_size, dtype = torch.long)
-
-    deg.scatter_add_(0, batch, torch.ones_like(batch))
-
-    ptr_vec = deg.new_zeros(batch_size + 1)
-    torch.cumsum(deg, 0, out=ptr_vec[1:])
-
-    #return fps_sampling(src, ptr_vec, ratio)
-    sampled_indicies = []
-
-    for b in range(batch_size):
-        # start and the end of each batch
-        start, end = ptr_vec[b].item(), ptr_vec[b + 1].item()
-        # points from the point cloud
-        points = src[start:end]
-
-        num_points = points.size(0)
-        num_samples = max(1, math.ceil(num_points * sampling_ratio))
-
-        selected = torch.zeros(num_samples, device = src.device, dtype = torch.long)
-        distances = torch.full((num_points,), float("inf"), device = src.device)
-
-        # select a random start point
-        if start_random:
-            farthest = torch.randint(0, num_points, (1,), device = src.device)
-        else:
-            farthest = torch.tensor([0], device = src.device, dtype = torch.long)
-
-        for i in range(num_samples):
-            selected[i] = farthest
-            centroid = points[farthest].squeeze(0)
-            dist = torch.norm(points - centroid, dim = 1) # compute euclidean distance
-            distances = torch.minimum(distances, dist)
-            farthest = torch.argmax(distances)
-
-        sampled_indicies.append(torch.arange(start, end)[selected])
-
-    return torch.cat(sampled_indicies, dim = 0)
-class PointCrossAttention(nn.Module):
-    def __init__(self,
-        num_latents: int,
-        downsample_ratio: float,
-        pc_size: int,
-        pc_sharpedge_size: int,
-        point_feats: int,
-        width: int,
-        heads: int,
-        layers: int,
-        fourier_embedder,
-        normal_pe: bool = False,
-        qkv_bias: bool = False,
-        use_ln_post: bool = True,
-        qk_norm: bool = True):
-
-        super().__init__()
-
-        self.fourier_embedder = fourier_embedder
-
-        self.pc_size = pc_size
-        self.normal_pe = normal_pe
-        self.downsample_ratio = downsample_ratio
-        self.pc_sharpedge_size = pc_sharpedge_size
-        self.num_latents = num_latents
-        self.point_feats = point_feats
-
-        self.input_proj = nn.Linear(self.fourier_embedder.out_dim + point_feats, width)
-
-        self.cross_attn = ResidualCrossAttentionBlock(
-            width = width,
-            heads = heads,
-            qkv_bias = qkv_bias,
-            qk_norm = qk_norm
-        )
-
-        self.self_attn = None
-        if layers > 0:
-            self.self_attn = Transformer(
-                width = width,
-                heads = heads,
-                qkv_bias = qkv_bias,
-                qk_norm = qk_norm,
-                layers = layers
-            )
-
-        if use_ln_post:
-            self.ln_post = nn.LayerNorm(width)
-        else:
-            self.ln_post = None
-
-    def sample_points_and_latents(self, point_cloud: torch.Tensor, features: torch.Tensor):
-
-        """
-        Subsample points randomly from the point cloud (input_pc)
-        Further sample the subsampled points to get query_pc
-        take the fourier embeddings for both input and query pc
-
-        Mental Note: FPS-sampled points (query_pc) act as latent tokens that attend to and learn from the broader context in input_pc.
-        Goal: get a smaller represenation (query_pc) to represent the entire scence structure by learning from a broader subset (input_pc).
-        More computationally efficient.
-
-        Features are additional information for each point in the cloud
-        """
-
-        B, _, D = point_cloud.shape
-
-        num_latents = int(self.num_latents)
-
-        num_random_query = self.pc_size / (self.pc_size + self.pc_sharpedge_size) * num_latents
-        num_sharpedge_query = num_latents - num_random_query
-
-        # Split random and sharpedge surface points
-        random_pc, sharpedge_pc = torch.split(point_cloud, [self.pc_size, self.pc_sharpedge_size], dim=1)
-
-        # assert statements
-        assert random_pc.shape[1] <= self.pc_size, "Random surface points size must be less than or equal to pc_size"
-        assert sharpedge_pc.shape[1] <= self.pc_sharpedge_size, "Sharpedge surface points size must be less than or equal to pc_sharpedge_size"
-
-        input_random_pc_size = int(num_random_query * self.downsample_ratio)
-        random_query_pc, random_input_pc, random_idx_pc, random_idx_query = \
-            self.subsample(pc = random_pc, num_query = num_random_query, input_pc_size = input_random_pc_size)
-
-        input_sharpedge_pc_size = int(num_sharpedge_query * self.downsample_ratio)
-
-        if input_sharpedge_pc_size == 0:
-            sharpedge_input_pc = torch.zeros(B, 0, D, dtype = random_input_pc.dtype).to(point_cloud.device)
-            sharpedge_query_pc = torch.zeros(B, 0, D, dtype= random_query_pc.dtype).to(point_cloud.device)
-
-        else:
-            sharpedge_query_pc, sharpedge_input_pc, sharpedge_idx_pc, sharpedge_idx_query = \
-            self.subsample(pc = sharpedge_pc, num_query = num_sharpedge_query, input_pc_size = input_sharpedge_pc_size)
-
-        # concat the random and sharpedges
-        query_pc = torch.cat([random_query_pc, sharpedge_query_pc], dim = 1)
-        input_pc = torch.cat([random_input_pc, sharpedge_input_pc], dim = 1)
-
-        query = self.fourier_embedder(query_pc)
-        data = self.fourier_embedder(input_pc)
-
-        if self.point_feats > 0:
-            random_surface_features, sharpedge_surface_features = torch.split(features, [self.pc_size, self.pc_sharpedge_size], dim = 1)
-
-            input_random_surface_features, query_random_features = \
-                self.handle_features(features = random_surface_features, idx_pc = random_idx_pc, batch_size = B,
-                                     input_pc_size = input_random_pc_size, idx_query = random_idx_query)
-
-            if input_sharpedge_pc_size == 0:
-                input_sharpedge_surface_features = torch.zeros(B, 0, self.point_feats,
-                                                               dtype = input_random_surface_features.dtype, device = point_cloud.device)
-
-                query_sharpedge_features = torch.zeros(B, 0, self.point_feats,
-                                                       dtype = query_random_features.dtype, device = point_cloud.device)
-            else:
-
-                input_sharpedge_surface_features, query_sharpedge_features = \
-                    self.handle_features(idx_pc = sharpedge_idx_pc, features = sharpedge_surface_features,
-                                         batch_size = B, idx_query = sharpedge_idx_query, input_pc_size = input_sharpedge_pc_size)
-
-            query_features = torch.cat([query_random_features, query_sharpedge_features], dim = 1)
-            input_features = torch.cat([input_random_surface_features, input_sharpedge_surface_features], dim = 1)
-
-            if self.normal_pe:
-                # apply the fourier embeddings on the first 3 dims (xyz)
-                input_features_pe = self.fourier_embedder(input_features[..., :3])
-                query_features_pe = self.fourier_embedder(query_features[..., :3])
-                # replace the first 3 dims with the new PE ones
-                input_features = torch.cat([input_features_pe, input_features[..., :3]], dim = -1)
-                query_features = torch.cat([query_features_pe, query_features[..., :3]], dim = -1)
-
-            # concat at the channels dim
-            query = torch.cat([query, query_features], dim = -1)
-            data = torch.cat([data, input_features], dim = -1)
-
-        # don't return pc_info to avoid unnecessary memory usuage
-        return query.view(B, -1, query.shape[-1]), data.view(B, -1, data.shape[-1])
-
-    def forward(self, point_cloud: torch.Tensor, features: torch.Tensor):
-
-        query, data = self.sample_points_and_latents(point_cloud = point_cloud, features = features)
-
-        # apply projections
-        query = self.input_proj(query)
-        data = self.input_proj(data)
-
-        # apply cross attention between query and data
-        latents = self.cross_attn(query, data)
-
-        if self.self_attn is not None:
-            latents = self.self_attn(latents)
-
-        if self.ln_post is not None:
-            latents = self.ln_post(latents)
-
-        return latents
-
-
-    def subsample(self, pc, num_query, input_pc_size: int):
-
-        """
-        num_query: number of points to keep after FPS
-        input_pc_size: number of points to select before FPS
-        """
-
-        B, _, D = pc.shape
-        query_ratio = num_query / input_pc_size
-
-        # random subsampling of points inside the point cloud
-        idx_pc = torch.randperm(pc.shape[1], device = pc.device)[:input_pc_size]
-        input_pc = pc[:, idx_pc, :]
-
-        # flatten to allow applying fps across the whole batch
-        flattent_input_pc = input_pc.view(B * input_pc_size, D)
-
-        # construct a batch_down tensor to tell fps
-        # which points belong to which batch
-        N_down = int(flattent_input_pc.shape[0] / B)
-        batch_down = torch.arange(B).to(pc.device)
-        batch_down = torch.repeat_interleave(batch_down, N_down)
-
-        idx_query = fps(flattent_input_pc, batch_down, sampling_ratio = query_ratio)
-        query_pc = flattent_input_pc[idx_query].view(B, -1, D)
-
-        return query_pc, input_pc, idx_pc, idx_query
-
-    def handle_features(self, features, idx_pc, input_pc_size, batch_size: int, idx_query):
-
-        B = batch_size
-
-        input_surface_features = features[:, idx_pc, :]
-        flattent_input_features = input_surface_features.view(B * input_pc_size, -1)
-        query_features = flattent_input_features[idx_query].view(B, -1,
-                                                                 flattent_input_features.shape[-1])
-
-        return input_surface_features, query_features
-
-def normalize_mesh(mesh, scale = 0.9999):
-    """Normalize mesh to fit in [-scale, scale]. Translate mesh so its center is [0,0,0]"""
-
-    bbox = mesh.bounds
-    center = (bbox[1] + bbox[0]) / 2
-
-    max_extent = (bbox[1] - bbox[0]).max()
-    mesh.apply_translation(-center)
-    mesh.apply_scale((2 * scale) / max_extent)
-
-    return mesh
-
-def sample_pointcloud(mesh, num = 200000):
-    """ Uniformly sample points from the surface of the mesh """
-
-    points, face_idx = mesh.sample(num, return_index = True)
-    normals = mesh.face_normals[face_idx]
-    return torch.from_numpy(points.astype(np.float32)), torch.from_numpy(normals.astype(np.float32))
-
-def detect_sharp_edges(mesh, threshold=0.985):
-    """Return edge indices (a, b) that lie on sharp boundaries of the mesh."""
-
-    V, F = mesh.vertices, mesh.faces
-    VN, FN = mesh.vertex_normals, mesh.face_normals
-
-    sharp_mask = np.ones(V.shape[0])
-    for i in range(3):
-        indices = F[:, i]
-        alignment = np.einsum('ij,ij->i', VN[indices], FN)
-        dot_stack = np.stack((sharp_mask[indices], alignment), axis=-1)
-        sharp_mask[indices] = np.min(dot_stack, axis=-1)
-
-    edge_a = np.concatenate([F[:, 0], F[:, 1], F[:, 2]])
-    edge_b = np.concatenate([F[:, 1], F[:, 2], F[:, 0]])
-    sharp_edges = (sharp_mask[edge_a] < threshold) & (sharp_mask[edge_b] < threshold)
-
-    return edge_a[sharp_edges], edge_b[sharp_edges]
-
-
-def sharp_sample_pointcloud(mesh, num = 16384):
-    """ Sample points preferentially from sharp edges in the mesh. """
-
-    edge_a, edge_b = detect_sharp_edges(mesh)
-    V, VN = mesh.vertices, mesh.vertex_normals
-
-    va, vb = V[edge_a], V[edge_b]
-    na, nb = VN[edge_a], VN[edge_b]
-
-    edge_lengths = np.linalg.norm(vb - va, axis=-1)
-    weights = edge_lengths / edge_lengths.sum()
-
-    indices = np.searchsorted(np.cumsum(weights), np.random.rand(num))
-    t = np.random.rand(num, 1)
-
-    samples = t * va[indices] + (1 - t) * vb[indices]
-    normals = t * na[indices] + (1 - t) * nb[indices]
-
-    return samples.astype(np.float32), normals.astype(np.float32)
-
-def load_surface_sharpedge(mesh, num_points=4096, num_sharp_points=4096, sharpedge_flag = True, device = "cuda"):
-    """Load a surface with optional sharp-edge annotations from a trimesh mesh."""
-
-    import trimesh
-
-    try:
-        mesh_full = trimesh.util.concatenate(mesh.dump())
-    except Exception:
-        mesh_full = trimesh.util.concatenate(mesh)
-
-    mesh_full = normalize_mesh(mesh_full)
-
-    faces = mesh_full.faces
-    vertices = mesh_full.vertices
-    origin_face_count = faces.shape[0]
-
-    mesh_surface = trimesh.Trimesh(vertices=vertices, faces=faces[:origin_face_count])
-    mesh_fill = trimesh.Trimesh(vertices=vertices, faces=faces[origin_face_count:])
-
-    area_surface = mesh_surface.area
-    area_fill = mesh_fill.area
-    total_area = area_surface + area_fill
-
-    sample_num = 499712 // 2
-    fill_ratio = area_fill / total_area if total_area > 0 else 0
-
-    num_fill = int(sample_num * fill_ratio)
-    num_surface = sample_num - num_fill
-
-    surf_pts, surf_normals = sample_pointcloud(mesh_surface, num_surface)
-    fill_pts, fill_normals = (torch.zeros(0, 3), torch.zeros(0, 3)) if num_fill == 0 else sample_pointcloud(mesh_fill, num_fill)
-
-    sharp_pts, sharp_normals = sharp_sample_pointcloud(mesh_surface, sample_num)
-
-    def assemble_tensor(points, normals, label=None):
-
-        data = torch.cat([points, normals], dim=1).half().to(device)
-
-        if label is not None:
-            label_tensor = torch.full((data.shape[0], 1), float(label), dtype=torch.float16).to(device)
-            data = torch.cat([data, label_tensor], dim=1)
-
-        return data
-
-    surface = assemble_tensor(torch.cat([surf_pts.to(device), fill_pts.to(device)], dim=0),
-                              torch.cat([surf_normals.to(device), fill_normals.to(device)], dim=0),
-                              label = 0 if sharpedge_flag else None)
-
-    sharp_surface = assemble_tensor(torch.from_numpy(sharp_pts), torch.from_numpy(sharp_normals),
-                                    label = 1 if sharpedge_flag else None)
-
-    rng = np.random.default_rng()
-
-    surface = surface[rng.choice(surface.shape[0], num_points, replace = False)]
-    sharp_surface = sharp_surface[rng.choice(sharp_surface.shape[0], num_sharp_points, replace = False)]
-
-    full = torch.cat([surface, sharp_surface], dim = 0).unsqueeze(0)
-
-    return full
-
-class SharpEdgeSurfaceLoader:
-    """ Load mesh surface and sharp edge samples. """
-
-    def __init__(self, num_uniform_points = 8192, num_sharp_points = 8192):
-
-        self.num_uniform_points = num_uniform_points
-        self.num_sharp_points = num_sharp_points
-        self.total_points = num_uniform_points + num_sharp_points
-
-    def __call__(self, mesh_input, device = "cuda"):
-        mesh = self._load_mesh(mesh_input)
-        return load_surface_sharpedge(mesh, self.num_uniform_points, self.num_sharp_points, device = device)
-
-    @staticmethod
-    def _load_mesh(mesh_input):
-        import trimesh
-
-        if isinstance(mesh_input, str):
-            mesh = trimesh.load(mesh_input, force="mesh", merge_primitives = True)
-        else:
-            mesh = mesh_input
-
-        if isinstance(mesh, trimesh.Scene):
-            combined = None
-            for obj in mesh.geometry.values():
-                combined = obj if combined is None else combined + obj
-            return combined
-
-        return mesh
-
-class DiagonalGaussianDistribution:
-    def __init__(self, params: torch.Tensor, feature_dim: int = -1):
-
-        # divide quant channels (8) into mean and log variance
-        self.mean, self.logvar = torch.chunk(params, 2, dim = feature_dim)
-
-        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
-        self.std = torch.exp(0.5 * self.logvar)
-
-    def sample(self):
-
-        eps = torch.randn_like(self.std)
-        z = self.mean + eps * self.std
-
-        return z
-
-################################################
-# Volume Decoder
-################################################
-
-class VanillaVolumeDecoder():
+def generate_dense_grid_points(
+    bbox_min: np.ndarray,
+    bbox_max: np.ndarray,
+    octree_resolution: int,
+    indexing: str = "ij",
+):
+    length = bbox_max - bbox_min
+    num_cells = octree_resolution
+
+    x = np.linspace(bbox_min[0], bbox_max[0], int(num_cells) + 1, dtype=np.float32)
+    y = np.linspace(bbox_min[1], bbox_max[1], int(num_cells) + 1, dtype=np.float32)
+    z = np.linspace(bbox_min[2], bbox_max[2], int(num_cells) + 1, dtype=np.float32)
+    [xs, ys, zs] = np.meshgrid(x, y, z, indexing=indexing)
+    xyz = np.stack((xs, ys, zs), axis=-1)
+    grid_size = [int(num_cells) + 1, int(num_cells) + 1, int(num_cells) + 1]
+
+    return xyz, grid_size, length
+
+
+class VanillaVolumeDecoder:
     @torch.no_grad()
-    def __call__(self, latents: torch.Tensor, geo_decoder: callable, octree_resolution: int, bounds = 1.01,
-                 num_chunks: int = 10_000, enable_pbar: bool = True, **kwargs):
+    def __call__(
+        self,
+        latents: torch.FloatTensor,
+        geo_decoder: Callable,
+        bounds: Union[Tuple[float], List[float], float] = 1.01,
+        num_chunks: int = 10000,
+        octree_resolution: int = None,
+        enable_pbar: bool = True,
+        **kwargs,
+    ):
+        device = latents.device
+        dtype = latents.dtype
+        batch_size = latents.shape[0]
 
+        # 1. generate query points
         if isinstance(bounds, float):
             bounds = [-bounds, -bounds, -bounds, bounds, bounds, bounds]
 
-        bbox_min, bbox_max = torch.tensor(bounds[:3]), torch.tensor(bounds[3:])
-
-        x = torch.linspace(bbox_min[0], bbox_max[0], int(octree_resolution) + 1, dtype = torch.float32)
-        y = torch.linspace(bbox_min[1], bbox_max[1], int(octree_resolution) + 1, dtype = torch.float32)
-        z = torch.linspace(bbox_min[2], bbox_max[2], int(octree_resolution) + 1, dtype = torch.float32)
-
-        [xs, ys, zs] = torch.meshgrid(x, y, z, indexing = "ij")
-        xyz = torch.stack((xs, ys, zs), axis=-1).to(latents.device, dtype = latents.dtype).contiguous().reshape(-1, 3)
-        grid_size = [int(octree_resolution) + 1, int(octree_resolution) + 1, int(octree_resolution) + 1]
+        bbox_min, bbox_max = np.array(bounds[0:3]), np.array(bounds[3:6])
+        xyz_samples, grid_size, length = generate_dense_grid_points(
+            bbox_min=bbox_min,
+            bbox_max=bbox_max,
+            octree_resolution=octree_resolution,
+            indexing="ij"
+        )
+        xyz_samples = torch.from_numpy(xyz_samples).to(device, dtype=dtype).contiguous().reshape(-1, 3)
 
+        # 2. latents to 3d volume
         batch_logits = []
-        for start in tqdm(range(0, xyz.shape[0], num_chunks), desc="Volume Decoding",
+        for start in tqdm(range(0, xyz_samples.shape[0], num_chunks), desc="Volume Decoding",
                           disable=not enable_pbar):
-
-            chunk_queries = xyz[start: start + num_chunks, :]
-            chunk_queries = chunk_queries.unsqueeze(0).repeat(latents.shape[0], 1, 1)
-            logits = geo_decoder(queries = chunk_queries, latents = latents)
+            chunk_queries = xyz_samples[start: start + num_chunks, :]
+            chunk_queries = repeat(chunk_queries, "p c -> b p c", b=batch_size)
+            logits = geo_decoder(queries=chunk_queries, latents=latents)
             batch_logits.append(logits)
 
-        grid_logits = torch.cat(batch_logits, dim = 1)
-        grid_logits = grid_logits.view((latents.shape[0], *grid_size)).float()
+        grid_logits = torch.cat(batch_logits, dim=1)
+        grid_logits = grid_logits.view((batch_size, *grid_size)).float()
 
         return grid_logits
 
+
 class FourierEmbedder(nn.Module):
     """The sin/cosine positional embedding. Given an input tensor `x` of shape [n_batch, ..., c_dim], it converts
     each feature dimension of `x[..., i]` into:
@@ -552,11 +175,13 @@ class FourierEmbedder(nn.Module):
         else:
             return x
 
+
 class CrossAttentionProcessor:
     def __call__(self, attn, q, k, v):
         out = comfy.ops.scaled_dot_product_attention(q, k, v)
         return out
 
+
 class DropPath(nn.Module):
     """Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
     """
@@ -607,42 +232,39 @@ class MLP(nn.Module):
     def forward(self, x):
         return self.drop_path(self.c_proj(self.gelu(self.c_fc(x))))
 
+
 class QKVMultiheadCrossAttention(nn.Module):
     def __init__(
         self,
+        *,
         heads: int,
-        n_data = None,
         width=None,
         qk_norm=False,
         norm_layer=ops.LayerNorm
     ):
         super().__init__()
         self.heads = heads
-        self.n_data = n_data
         self.q_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
         self.k_norm = norm_layer(width // heads, elementwise_affine=True, eps=1e-6) if qk_norm else nn.Identity()
 
-    def forward(self, q, kv):
+        self.attn_processor = CrossAttentionProcessor()
 
+    def forward(self, q, kv):
         _, n_ctx, _ = q.shape
         bs, n_data, width = kv.shape
-
         attn_ch = width // self.heads // 2
         q = q.view(bs, n_ctx, self.heads, -1)
-
         kv = kv.view(bs, n_data, self.heads, -1)
         k, v = torch.split(kv, attn_ch, dim=-1)
 
         q = self.q_norm(q)
         k = self.k_norm(k)
-
-        q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)]
-        out = F.scaled_dot_product_attention(q, k, v)
-
+        q, k, v = map(lambda t: rearrange(t, 'b n h d -> b h n d', h=self.heads), (q, k, v))
+        out = self.attn_processor(self, q, k, v)
         out = out.transpose(1, 2).reshape(bs, n_ctx, -1)
-
         return out
 
+
 class MultiheadCrossAttention(nn.Module):
     def __init__(
         self,
@@ -684,6 +306,7 @@ class MultiheadCrossAttention(nn.Module):
         x = self.c_proj(x)
         return x
 
+
 class ResidualCrossAttentionBlock(nn.Module):
     def __init__(
         self,
@@ -743,7 +366,7 @@ class QKVMultiheadAttention(nn.Module):
         q = self.q_norm(q)
         k = self.k_norm(k)
 
-        q, k, v = [t.permute(0, 2, 1, 3) for t in (q, k, v)]
+        q, k, v = map(lambda t: rearrange(t, 'b n h d -> b h n d', h=self.heads), (q, k, v))
         out = F.scaled_dot_product_attention(q, k, v).transpose(1, 2).reshape(bs, n_ctx, -1)
         return out
 
@@ -760,7 +383,8 @@ class MultiheadAttention(nn.Module):
         drop_path_rate: float = 0.0
     ):
         super().__init__()
-
+        self.width = width
+        self.heads = heads
         self.c_qkv = ops.Linear(width, width * 3, bias=qkv_bias)
         self.c_proj = ops.Linear(width, width)
         self.attention = QKVMultiheadAttention(
@@ -867,7 +491,7 @@ class CrossAttentionDecoder(nn.Module):
         self.query_proj = ops.Linear(self.fourier_embedder.out_dim, width)
         if self.downsample_ratio != 1:
             self.latents_proj = ops.Linear(width * downsample_ratio, width)
-        if not self.enable_ln_post:
+        if self.enable_ln_post == False:
             qk_norm = False
         self.cross_attn_decoder = ResidualCrossAttentionBlock(
             width=width,
@@ -898,44 +522,28 @@ class CrossAttentionDecoder(nn.Module):
 
 class ShapeVAE(nn.Module):
     def __init__(
-            self,
-            *,
-            num_latents: int = 4096,
-            embed_dim: int = 64,
-            width: int = 1024,
-            heads: int = 16,
-            num_decoder_layers: int = 16,
-            num_encoder_layers: int = 8,
-            pc_size: int = 81920,
-            pc_sharpedge_size: int = 0,
-            point_feats: int = 4,
-            downsample_ratio: int = 20,
-            geo_decoder_downsample_ratio: int = 1,
-            geo_decoder_mlp_expand_ratio: int = 4,
-            geo_decoder_ln_post: bool = True,
-            num_freqs: int = 8,
-            qkv_bias: bool = False,
-            qk_norm: bool = True,
-            drop_path_rate: float = 0.0,
-            include_pi: bool = False,
-            scale_factor: float = 1.0039506158752403,
-            label_type: str = "binary",
+        self,
+        *,
+        embed_dim: int,
+        width: int,
+        heads: int,
+        num_decoder_layers: int,
+        geo_decoder_downsample_ratio: int = 1,
+        geo_decoder_mlp_expand_ratio: int = 4,
+        geo_decoder_ln_post: bool = True,
+        num_freqs: int = 8,
+        include_pi: bool = True,
+        qkv_bias: bool = True,
+        qk_norm: bool = False,
+        label_type: str = "binary",
+        drop_path_rate: float = 0.0,
+        scale_factor: float = 1.0,
     ):
         super().__init__()
         self.geo_decoder_ln_post = geo_decoder_ln_post
 
         self.fourier_embedder = FourierEmbedder(num_freqs=num_freqs, include_pi=include_pi)
 
-        self.encoder = PointCrossAttention(layers = num_encoder_layers,
-                                    num_latents = num_latents,
-                                    downsample_ratio = downsample_ratio,
-                                    heads = heads,
-                                    pc_size = pc_size,
-                                    width = width,
-                                    point_feats = point_feats,
-                                    fourier_embedder = self.fourier_embedder,
-                                    pc_sharpedge_size = pc_sharpedge_size)
-
         self.post_kl = ops.Linear(embed_dim, width)
 
         self.transformer = Transformer(
@@ -975,14 +583,5 @@ class ShapeVAE(nn.Module):
         grid_logits = self.volume_decoder(latents, self.geo_decoder, bounds=bounds, num_chunks=num_chunks, octree_resolution=octree_resolution, enable_pbar=enable_pbar)
         return grid_logits.movedim(-2, -1)
 
-    def encode(self, surface):
-
-        pc, feats = surface[:, :, :3], surface[:, :, 3:]
-        latents = self.encoder(pc, feats)
-
-        moments = self.pre_kl(latents)
-        posterior = DiagonalGaussianDistribution(moments, feature_dim = -1)
-
-        latents = posterior.sample()
-
-        return latents
+    def encode(self, x):
+        return None

comfy/ldm/hunyuan3dv2_1/hunyuandit.py

@@ -1,659 +0,0 @@
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from comfy.ldm.modules.attention import optimized_attention
-import comfy.model_management
-
-class GELU(nn.Module):
-
-    def __init__(self, dim_in: int, dim_out: int, operations, device, dtype):
-        super().__init__()
-        self.proj = operations.Linear(dim_in, dim_out, device = device, dtype = dtype)
-
-    def gelu(self, gate: torch.Tensor) -> torch.Tensor:
-
-        if gate.device.type == "mps":
-            return F.gelu(gate.to(dtype = torch.float32)).to(dtype = gate.dtype)
-
-        return F.gelu(gate)
-
-    def forward(self, hidden_states):
-
-        hidden_states = self.proj(hidden_states)
-        hidden_states = self.gelu(hidden_states)
-
-        return hidden_states
-
-class FeedForward(nn.Module):
-
-    def __init__(self, dim: int, dim_out = None, mult: int = 4,
-                dropout: float = 0.0, inner_dim = None, operations = None, device = None, dtype = None):
-
-        super().__init__()
-        if inner_dim is None:
-            inner_dim = int(dim * mult)
-
-        dim_out = dim_out if dim_out is not None else dim
-
-        act_fn = GELU(dim, inner_dim, operations = operations, device = device, dtype = dtype)
-
-        self.net = nn.ModuleList([])
-        self.net.append(act_fn)
-
-        self.net.append(nn.Dropout(dropout))
-        self.net.append(operations.Linear(inner_dim, dim_out, device = device, dtype = dtype))
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        for module in self.net:
-            hidden_states = module(hidden_states)
-        return hidden_states
-
-class AddAuxLoss(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, x, loss):
-        # do nothing in forward (no computation)
-        ctx.requires_aux_loss = loss.requires_grad
-        ctx.dtype = loss.dtype
-
-        return x
-
-    @staticmethod
-    def backward(ctx, grad_output):
-        # add the aux loss gradients
-        grad_loss = None
-        # put the aux grad the same as the main grad loss
-        # aux grad contributes equally
-        if ctx.requires_aux_loss:
-            grad_loss = torch.ones(1, dtype = ctx.dtype, device = grad_output.device)
-
-        return grad_output, grad_loss
-
-class MoEGate(nn.Module):
-
-    def __init__(self, embed_dim, num_experts=16, num_experts_per_tok=2, aux_loss_alpha=0.01, device = None, dtype = None):
-
-        super().__init__()
-        self.top_k = num_experts_per_tok
-        self.n_routed_experts = num_experts
-
-        self.alpha = aux_loss_alpha
-
-        self.gating_dim = embed_dim
-        self.weight = nn.Parameter(torch.empty((self.n_routed_experts, self.gating_dim), device = device, dtype = dtype))
-
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-
-        # flatten hidden states
-        hidden_states = hidden_states.view(-1, hidden_states.size(-1))
-
-        # get logits and pass it to softmax
-        logits = F.linear(hidden_states, comfy.model_management.cast_to(self.weight, dtype=hidden_states.dtype, device=hidden_states.device), bias = None)
-        scores = logits.softmax(dim = -1)
-
-        topk_weight, topk_idx = torch.topk(scores, k = self.top_k, dim = -1, sorted = False)
-
-        if self.training and self.alpha > 0.0:
-            scores_for_aux = scores
-
-            # used bincount instead of one hot encoding
-            counts = torch.bincount(topk_idx.view(-1), minlength = self.n_routed_experts).float()
-            ce = counts / topk_idx.numel()  # normalized expert usage
-
-            # mean expert score
-            Pi = scores_for_aux.mean(0)
-
-            # expert balance loss
-            aux_loss = (Pi * ce * self.n_routed_experts).sum() * self.alpha
-        else:
-            aux_loss = None
-
-        return topk_idx, topk_weight, aux_loss
-
-class MoEBlock(nn.Module):
-    def __init__(self, dim, num_experts: int = 6, moe_top_k: int = 2, dropout: float = 0.0,
-                 ff_inner_dim: int = None, operations = None, device = None, dtype = None):
-        super().__init__()
-
-        self.moe_top_k = moe_top_k
-        self.num_experts = num_experts
-
-        self.experts = nn.ModuleList([
-            FeedForward(dim, dropout = dropout, inner_dim = ff_inner_dim, operations = operations, device = device, dtype = dtype)
-            for _ in range(num_experts)
-        ])
-
-        self.gate = MoEGate(dim, num_experts = num_experts, num_experts_per_tok = moe_top_k, device = device, dtype = dtype)
-        self.shared_experts = FeedForward(dim, dropout = dropout, inner_dim = ff_inner_dim, operations = operations, device = device, dtype = dtype)
-
-    def forward(self, hidden_states) -> torch.Tensor:
-
-        identity = hidden_states
-        orig_shape = hidden_states.shape
-        topk_idx, topk_weight, aux_loss = self.gate(hidden_states)
-
-        hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-        flat_topk_idx = topk_idx.view(-1)
-
-        if self.training:
-
-            hidden_states = hidden_states.repeat_interleave(self.moe_top_k, dim = 0)
-            y = torch.empty_like(hidden_states, dtype = hidden_states.dtype)
-
-            for i, expert in enumerate(self.experts):
-                tmp = expert(hidden_states[flat_topk_idx == i])
-                y[flat_topk_idx == i] = tmp.to(hidden_states.dtype)
-
-            y = (y.view(*topk_weight.shape, -1) * topk_weight.unsqueeze(-1)).sum(dim = 1)
-            y =  y.view(*orig_shape)
-
-            y = AddAuxLoss.apply(y, aux_loss)
-        else:
-            y = self.moe_infer(hidden_states, flat_expert_indices = flat_topk_idx,flat_expert_weights = topk_weight.view(-1, 1)).view(*orig_shape)
-
-        y = y + self.shared_experts(identity)
-
-        return y
-
-    @torch.no_grad()
-    def moe_infer(self, x, flat_expert_indices, flat_expert_weights):
-
-        expert_cache = torch.zeros_like(x)
-        idxs = flat_expert_indices.argsort()
-
-        # no need for .numpy().cpu() here
-        tokens_per_expert = flat_expert_indices.bincount().cumsum(0)
-        token_idxs = idxs // self.moe_top_k
-
-        for i, end_idx in enumerate(tokens_per_expert):
-
-            start_idx = 0 if i == 0 else tokens_per_expert[i-1]
-
-            if start_idx == end_idx:
-                continue
-
-            expert = self.experts[i]
-            exp_token_idx = token_idxs[start_idx:end_idx]
-
-            expert_tokens = x[exp_token_idx]
-            expert_out = expert(expert_tokens)
-
-            expert_out.mul_(flat_expert_weights[idxs[start_idx:end_idx]])
-
-            # use index_add_ with a 1-D index tensor directly avoids building a large [N, D] index map and extra memcopy required by scatter_reduce_
-            # + avoid dtype conversion
-            expert_cache.index_add_(0, exp_token_idx, expert_out)
-
-        return expert_cache
-
-class Timesteps(nn.Module):
-    def __init__(self, num_channels: int, downscale_freq_shift: float = 0.0,
-                 scale: float = 1.0, max_period: int = 10000):
-        super().__init__()
-
-        self.num_channels = num_channels
-        half_dim = num_channels // 2
-
-        # precompute the “inv_freq” vector once
-        exponent = -math.log(max_period) * torch.arange(
-            half_dim, dtype=torch.float32
-        ) / (half_dim - downscale_freq_shift)
-
-        inv_freq = torch.exp(exponent)
-
-        # pad
-        if num_channels % 2 == 1:
-            # we’ll pad a zero at the end of the cos-half
-            inv_freq = torch.cat([inv_freq, inv_freq.new_zeros(1)])
-
-        # register to buffer so it moves with the device
-        self.register_buffer("inv_freq", inv_freq, persistent = False)
-        self.scale = scale
-
-    def forward(self, timesteps: torch.Tensor):
-
-        x = timesteps.float().unsqueeze(1) * self.inv_freq.to(timesteps.device).unsqueeze(0)
-
-
-        # fused CUDA kernels for sin and cos
-        sin_emb = x.sin()
-        cos_emb = x.cos()
-
-        emb = torch.cat([sin_emb, cos_emb], dim = 1)
-
-        # scale factor
-        if self.scale != 1.0:
-            emb = emb * self.scale
-
-        # If we padded inv_freq for odd, emb is already wide enough; otherwise:
-        if emb.shape[1] > self.num_channels:
-            emb = emb[:, :self.num_channels]
-
-        return emb
-
-class TimestepEmbedder(nn.Module):
-    def __init__(self, hidden_size, frequency_embedding_size = 256, cond_proj_dim = None, operations = None, device = None, dtype = None):
-        super().__init__()
-
-        self.mlp = nn.Sequential(
-            operations.Linear(hidden_size, frequency_embedding_size, bias=True, device = device, dtype = dtype),
-            nn.GELU(),
-            operations.Linear(frequency_embedding_size, hidden_size, bias=True, device = device, dtype = dtype),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-        if cond_proj_dim is not None:
-            self.cond_proj = operations.Linear(cond_proj_dim, frequency_embedding_size, bias=False, device = device, dtype = dtype)
-
-        self.time_embed = Timesteps(hidden_size)
-
-    def forward(self, timesteps, condition):
-
-        timestep_embed = self.time_embed(timesteps).type(self.mlp[0].weight.dtype)
-
-        if condition is not None:
-            cond_embed = self.cond_proj(condition)
-            timestep_embed = timestep_embed + cond_embed
-
-        time_conditioned = self.mlp(timestep_embed)
-
-        # for broadcasting with image tokens
-        return time_conditioned.unsqueeze(1)
-
-class MLP(nn.Module):
-    def __init__(self, *, width: int, operations = None, device = None, dtype = None):
-        super().__init__()
-        self.width = width
-        self.fc1 = operations.Linear(width, width * 4, device = device, dtype = dtype)
-        self.fc2 = operations.Linear(width * 4, width, device = device, dtype = dtype)
-        self.gelu = nn.GELU()
-
-    def forward(self, x):
-        return self.fc2(self.gelu(self.fc1(x)))
-
-class CrossAttention(nn.Module):
-    def __init__(
-        self,
-        qdim,
-        kdim,
-        num_heads,
-        qkv_bias=True,
-        qk_norm=False,
-        norm_layer=nn.LayerNorm,
-        use_fp16: bool = False,
-        operations = None,
-        dtype = None,
-        device = None,
-        **kwargs,
-    ):
-        super().__init__()
-        self.qdim = qdim
-        self.kdim = kdim
-
-        self.num_heads = num_heads
-        self.head_dim = self.qdim // num_heads
-
-        self.scale = self.head_dim ** -0.5
-
-        self.to_q = operations.Linear(qdim, qdim, bias=qkv_bias, device = device, dtype = dtype)
-        self.to_k = operations.Linear(kdim, qdim, bias=qkv_bias, device = device, dtype = dtype)
-        self.to_v = operations.Linear(kdim, qdim, bias=qkv_bias, device = device, dtype = dtype)
-
-        if use_fp16:
-            eps = 1.0 / 65504
-        else:
-            eps = 1e-6
-
-        if norm_layer == nn.LayerNorm:
-            norm_layer = operations.LayerNorm
-        else:
-            norm_layer = operations.RMSNorm
-
-        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
-        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
-        self.out_proj = operations.Linear(qdim, qdim, bias=True, device = device, dtype = dtype)
-
-    def forward(self, x, y):
-
-        b, s1, _ = x.shape
-        _, s2, _ = y.shape
-
-        y = y.to(next(self.to_k.parameters()).dtype)
-
-        q = self.to_q(x)
-        k = self.to_k(y)
-        v = self.to_v(y)
-
-        kv = torch.cat((k, v), dim=-1)
-        split_size = kv.shape[-1] // self.num_heads // 2
-
-        kv = kv.view(1, -1, self.num_heads, split_size * 2)
-        k, v = torch.split(kv, split_size, dim=-1)
-
-        q = q.view(b, s1, self.num_heads, self.head_dim)
-        k = k.view(b, s2, self.num_heads, self.head_dim)
-        v = v.reshape(b, s2, self.num_heads * self.head_dim)
-
-        q = self.q_norm(q)
-        k = self.k_norm(k)
-
-        x = optimized_attention(
-            q.reshape(b, s1, self.num_heads * self.head_dim),
-            k.reshape(b, s2, self.num_heads * self.head_dim),
-            v,
-            heads=self.num_heads,
-        )
-
-        out = self.out_proj(x)
-
-        return out
-
-class Attention(nn.Module):
-
-    def __init__(
-        self,
-        dim,
-        num_heads,
-        qkv_bias = True,
-        qk_norm = False,
-        norm_layer = nn.LayerNorm,
-        use_fp16: bool = False,
-        operations = None,
-        device = None,
-        dtype = None
-    ):
-        super().__init__()
-        self.dim = dim
-        self.num_heads = num_heads
-        self.head_dim = self.dim // num_heads
-        self.scale = self.head_dim ** -0.5
-
-        self.to_q = operations.Linear(dim, dim, bias = qkv_bias, device = device, dtype = dtype)
-        self.to_k = operations.Linear(dim, dim, bias = qkv_bias, device = device, dtype = dtype)
-        self.to_v = operations.Linear(dim, dim, bias = qkv_bias, device = device, dtype = dtype)
-
-        if use_fp16:
-            eps = 1.0 / 65504
-        else:
-            eps = 1e-6
-
-        if norm_layer == nn.LayerNorm:
-            norm_layer = operations.LayerNorm
-        else:
-            norm_layer = operations.RMSNorm
-
-        self.q_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
-        self.k_norm = norm_layer(self.head_dim, elementwise_affine=True, eps = eps, device = device, dtype = dtype) if qk_norm else nn.Identity()
-        self.out_proj = operations.Linear(dim, dim, device = device, dtype = dtype)
-
-    def forward(self, x):
-        B, N, _ = x.shape
-
-        query = self.to_q(x)
-        key = self.to_k(x)
-        value = self.to_v(x)
-
-        qkv_combined = torch.cat((query, key, value), dim=-1)
-        split_size = qkv_combined.shape[-1] // self.num_heads // 3
-
-        qkv = qkv_combined.view(1, -1, self.num_heads, split_size * 3)
-        query, key, value = torch.split(qkv, split_size, dim=-1)
-
-        query = query.reshape(B, N, self.num_heads, self.head_dim)
-        key = key.reshape(B, N, self.num_heads, self.head_dim)
-        value = value.reshape(B, N, self.num_heads * self.head_dim)
-
-        query = self.q_norm(query)
-        key = self.k_norm(key)
-
-        x = optimized_attention(
-            query.reshape(B, N, self.num_heads * self.head_dim),
-            key.reshape(B, N, self.num_heads * self.head_dim),
-            value,
-            heads=self.num_heads,
-        )
-
-        x = self.out_proj(x)
-        return x
-
-class HunYuanDiTBlock(nn.Module):
-    def __init__(
-        self,
-        hidden_size,
-        c_emb_size,
-        num_heads,
-        text_states_dim=1024,
-        qk_norm=False,
-        norm_layer=nn.LayerNorm,
-        qk_norm_layer=True,
-        qkv_bias=True,
-        skip_connection=True,
-        timested_modulate=False,
-        use_moe: bool = False,
-        num_experts: int = 8,
-        moe_top_k: int = 2,
-        use_fp16: bool = False,
-        operations = None,
-        device = None, dtype = None
-    ):
-        super().__init__()
-
-        # eps can't be 1e-6 in fp16 mode because of numerical stability issues
-        if use_fp16:
-            eps = 1.0 / 65504
-        else:
-            eps = 1e-6
-
-        self.norm1 = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
-
-        self.attn1 = Attention(hidden_size, num_heads=num_heads, qkv_bias=qkv_bias, qk_norm=qk_norm,
-                               norm_layer=qk_norm_layer, use_fp16 = use_fp16, device = device, dtype = dtype, operations = operations)
-
-        self.norm2 = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
-
-        self.timested_modulate = timested_modulate
-        if self.timested_modulate:
-            self.default_modulation = nn.Sequential(
-                nn.SiLU(),
-                operations.Linear(c_emb_size, hidden_size, bias=True, device = device, dtype = dtype)
-            )
-
-        self.attn2 = CrossAttention(hidden_size, text_states_dim, num_heads=num_heads, qkv_bias=qkv_bias,
-                                    qk_norm=qk_norm, norm_layer=qk_norm_layer, use_fp16 = use_fp16,
-                                    device = device, dtype = dtype, operations = operations)
-
-        self.norm3 = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
-
-        if skip_connection:
-            self.skip_norm = norm_layer(hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
-            self.skip_linear = operations.Linear(2 * hidden_size, hidden_size, device = device, dtype = dtype)
-        else:
-            self.skip_linear = None
-
-        self.use_moe = use_moe
-
-        if self.use_moe:
-            self.moe = MoEBlock(
-                hidden_size,
-                num_experts = num_experts,
-                moe_top_k = moe_top_k,
-                dropout = 0.0,
-                ff_inner_dim = int(hidden_size * 4.0),
-                device = device, dtype = dtype,
-                operations = operations
-            )
-        else:
-            self.mlp = MLP(width=hidden_size, operations=operations, device = device, dtype = dtype)
-
-    def forward(self, hidden_states, conditioning=None, text_states=None, skip_tensor=None):
-
-        if self.skip_linear is not None:
-            combined = torch.cat([skip_tensor, hidden_states], dim=-1)
-            hidden_states = self.skip_linear(combined)
-            hidden_states = self.skip_norm(hidden_states)
-
-        # self attention
-        if self.timested_modulate:
-            modulation_shift = self.default_modulation(conditioning).unsqueeze(dim=1)
-            hidden_states = hidden_states + modulation_shift
-
-        self_attn_out = self.attn1(self.norm1(hidden_states))
-        hidden_states = hidden_states + self_attn_out
-
-        # cross attention
-        hidden_states = hidden_states + self.attn2(self.norm2(hidden_states), text_states)
-
-        # MLP Layer
-        mlp_input = self.norm3(hidden_states)
-
-        if self.use_moe:
-            hidden_states = hidden_states + self.moe(mlp_input)
-        else:
-            hidden_states = hidden_states + self.mlp(mlp_input)
-
-        return hidden_states
-
-class FinalLayer(nn.Module):
-
-    def __init__(self, final_hidden_size, out_channels, operations, use_fp16: bool = False, device = None, dtype = None):
-        super().__init__()
-
-        if use_fp16:
-            eps = 1.0 / 65504
-        else:
-            eps = 1e-6
-
-        self.norm_final = operations.LayerNorm(final_hidden_size, elementwise_affine = True, eps = eps, device = device, dtype = dtype)
-        self.linear = operations.Linear(final_hidden_size, out_channels, bias = True, device = device, dtype = dtype)
-
-    def forward(self, x):
-        x = self.norm_final(x)
-        x = x[:, 1:]
-        x = self.linear(x)
-        return x
-
-class HunYuanDiTPlain(nn.Module):
-
-    # init with the defaults values from https://huggingface.co/tencent/Hunyuan3D-2.1/blob/main/hunyuan3d-dit-v2-1/config.yaml
-    def __init__(
-        self,
-        in_channels: int = 64,
-        hidden_size: int = 2048,
-        context_dim: int = 1024,
-        depth: int = 21,
-        num_heads: int = 16,
-        qk_norm: bool = True,
-        qkv_bias: bool = False,
-        num_moe_layers: int = 6,
-        guidance_cond_proj_dim = 2048,
-        norm_type = 'layer',
-        num_experts: int = 8,
-        moe_top_k: int = 2,
-        use_fp16: bool = False,
-        dtype = None,
-        device = None,
-        operations = None,
-        **kwargs
-        ):
-
-        self.dtype = dtype
-
-        super().__init__()
-
-        self.depth = depth
-
-        self.in_channels = in_channels
-        self.out_channels = in_channels
-
-        self.num_heads = num_heads
-        self.hidden_size = hidden_size
-
-        norm = operations.LayerNorm if norm_type == 'layer' else operations.RMSNorm
-        qk_norm = operations.RMSNorm
-
-        self.context_dim = context_dim
-        self.guidance_cond_proj_dim = guidance_cond_proj_dim
-
-        self.x_embedder = operations.Linear(in_channels, hidden_size, bias = True, device = device, dtype = dtype)
-        self.t_embedder = TimestepEmbedder(hidden_size, hidden_size * 4, cond_proj_dim = guidance_cond_proj_dim, device = device, dtype = dtype, operations = operations)
-
-
-        # HUnYuanDiT Blocks
-        self.blocks = nn.ModuleList([
-            HunYuanDiTBlock(hidden_size=hidden_size,
-                            c_emb_size=hidden_size,
-                            num_heads=num_heads,
-                            text_states_dim=context_dim,
-                            qk_norm=qk_norm,
-                            norm_layer = norm,
-                            qk_norm_layer = qk_norm,
-                            skip_connection=layer > depth // 2,
-                            qkv_bias=qkv_bias,
-                            use_moe=True if depth - layer <= num_moe_layers else False,
-                            num_experts=num_experts,
-                            moe_top_k=moe_top_k,
-                            use_fp16 = use_fp16,
-                            device = device, dtype = dtype, operations = operations)
-            for layer in range(depth)
-        ])
-
-        self.depth = depth
-
-        self.final_layer = FinalLayer(hidden_size, self.out_channels, use_fp16 = use_fp16, operations = operations, device = device, dtype = dtype)
-
-    def forward(self, x, t, context, transformer_options = {}, **kwargs):
-
-        x = x.movedim(-1, -2)
-        uncond_emb, cond_emb = context.chunk(2, dim = 0)
-
-        context = torch.cat([cond_emb, uncond_emb], dim = 0)
-        main_condition = context
-
-        t = 1.0 - t
-
-        time_embedded = self.t_embedder(t, condition = kwargs.get('guidance_cond'))
-
-        x = x.to(dtype = next(self.x_embedder.parameters()).dtype)
-        x_embedded = self.x_embedder(x)
-
-        combined = torch.cat([time_embedded, x_embedded], dim=1)
-
-        def block_wrap(args):
-            return block(
-                args["x"],
-                args["t"],
-                args["cond"],
-                skip_tensor=args.get("skip"),)
-
-        skip_stack = []
-        patches_replace = transformer_options.get("patches_replace", {})
-        blocks_replace = patches_replace.get("dit", {})
-        for idx, block in enumerate(self.blocks):
-            if idx <= self.depth // 2:
-                skip_input = None
-            else:
-                skip_input = skip_stack.pop()
-
-            if ("block", idx) in blocks_replace:
-
-                combined = blocks_replace[("block", idx)](
-                    {
-                        "x": combined,
-                        "t": time_embedded,
-                        "cond": main_condition,
-                        "skip": skip_input,
-                    },
-                    {"original_block": block_wrap},
-                )
-            else:
-                combined = block(combined, time_embedded, main_condition, skip_tensor=skip_input)
-
-            if idx < self.depth // 2:
-                skip_stack.append(combined)
-
-        output = self.final_layer(combined)
-        output =  output.movedim(-2, -1) * (-1.0)
-
-        cond_emb, uncond_emb = output.chunk(2, dim = 0)
-        return torch.cat([uncond_emb, cond_emb])

comfy/ldm/hunyuan_video/model.py

@@ -40,8 +40,6 @@ class HunyuanVideoParams:
     patch_size: list
     qkv_bias: bool
     guidance_embed: bool
-    byt5: bool
-    meanflow: bool
 
 
 class SelfAttentionRef(nn.Module):
@@ -164,30 +162,6 @@ class TokenRefiner(nn.Module):
         x = self.individual_token_refiner(x, c, mask, transformer_options=transformer_options)
         return x
 
-
-class ByT5Mapper(nn.Module):
-    def __init__(self, in_dim, out_dim, hidden_dim, out_dim1, use_res=False, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.layernorm = operations.LayerNorm(in_dim, dtype=dtype, device=device)
-        self.fc1 = operations.Linear(in_dim, hidden_dim, dtype=dtype, device=device)
-        self.fc2 = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device)
-        self.fc3 = operations.Linear(out_dim, out_dim1, dtype=dtype, device=device)
-        self.use_res = use_res
-        self.act_fn = nn.GELU()
-
-    def forward(self, x):
-        if self.use_res:
-            res = x
-        x = self.layernorm(x)
-        x = self.fc1(x)
-        x = self.act_fn(x)
-        x = self.fc2(x)
-        x2 = self.act_fn(x)
-        x2 = self.fc3(x2)
-        if self.use_res:
-            x2 = x2 + res
-        return x2
-
 class HunyuanVideo(nn.Module):
     """
     Transformer model for flow matching on sequences.
@@ -212,13 +186,9 @@ class HunyuanVideo(nn.Module):
         self.num_heads = params.num_heads
         self.pe_embedder = EmbedND(dim=pe_dim, theta=params.theta, axes_dim=params.axes_dim)
 
-        self.img_in = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed(None, self.patch_size, self.in_channels, self.hidden_size, conv3d=len(self.patch_size) == 3, dtype=dtype, device=device, operations=operations)
+        self.img_in = comfy.ldm.modules.diffusionmodules.mmdit.PatchEmbed(None, self.patch_size, self.in_channels, self.hidden_size, conv3d=True, dtype=dtype, device=device, operations=operations)
         self.time_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations)
-        if params.vec_in_dim is not None:
-            self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
-        else:
-            self.vector_in = None
-
+        self.vector_in = MLPEmbedder(params.vec_in_dim, self.hidden_size, dtype=dtype, device=device, operations=operations)
         self.guidance_in = (
             MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations) if params.guidance_embed else nn.Identity()
         )
@@ -246,23 +216,6 @@ class HunyuanVideo(nn.Module):
             ]
         )
 
-        if params.byt5:
-            self.byt5_in = ByT5Mapper(
-                in_dim=1472,
-                out_dim=2048,
-                hidden_dim=2048,
-                out_dim1=self.hidden_size,
-                use_res=False,
-                dtype=dtype, device=device, operations=operations
-            )
-        else:
-            self.byt5_in = None
-
-        if params.meanflow:
-            self.time_r_in = MLPEmbedder(in_dim=256, hidden_dim=self.hidden_size, dtype=dtype, device=device, operations=operations)
-        else:
-            self.time_r_in = None
-
         if final_layer:
             self.final_layer = LastLayer(self.hidden_size, self.patch_size[-1], self.out_channels, dtype=dtype, device=device, operations=operations)
 
@@ -274,12 +227,10 @@ class HunyuanVideo(nn.Module):
         txt_ids: Tensor,
         txt_mask: Tensor,
         timesteps: Tensor,
-        y: Tensor = None,
-        txt_byt5=None,
+        y: Tensor,
         guidance: Tensor = None,
         guiding_frame_index=None,
         ref_latent=None,
-        disable_time_r=False,
         control=None,
         transformer_options={},
     ) -> Tensor:
@@ -290,14 +241,6 @@ class HunyuanVideo(nn.Module):
         img = self.img_in(img)
         vec = self.time_in(timestep_embedding(timesteps, 256, time_factor=1.0).to(img.dtype))
 
-        if (self.time_r_in is not None) and (not disable_time_r):
-            w = torch.where(transformer_options['sigmas'][0] == transformer_options['sample_sigmas'])[0]  # This most likely could be improved
-            if len(w) > 0:
-                timesteps_r = transformer_options['sample_sigmas'][w[0] + 1]
-                timesteps_r = timesteps_r.unsqueeze(0).to(device=timesteps.device, dtype=timesteps.dtype)
-                vec_r = self.time_r_in(timestep_embedding(timesteps_r, 256, time_factor=1000.0).to(img.dtype))
-                vec = (vec + vec_r) / 2
-
         if ref_latent is not None:
             ref_latent_ids = self.img_ids(ref_latent)
             ref_latent = self.img_in(ref_latent)
@@ -308,17 +251,13 @@ class HunyuanVideo(nn.Module):
 
         if guiding_frame_index is not None:
             token_replace_vec = self.time_in(timestep_embedding(guiding_frame_index, 256, time_factor=1.0))
-            if self.vector_in is not None:
-                vec_ = self.vector_in(y[:, :self.params.vec_in_dim])
-                vec = torch.cat([(vec_ + token_replace_vec).unsqueeze(1), (vec_ + vec).unsqueeze(1)], dim=1)
-            else:
-                vec = torch.cat([(token_replace_vec).unsqueeze(1), (vec).unsqueeze(1)], dim=1)
+            vec_ = self.vector_in(y[:, :self.params.vec_in_dim])
+            vec = torch.cat([(vec_ + token_replace_vec).unsqueeze(1), (vec_ + vec).unsqueeze(1)], dim=1)
             frame_tokens = (initial_shape[-1] // self.patch_size[-1]) * (initial_shape[-2] // self.patch_size[-2])
             modulation_dims = [(0, frame_tokens, 0), (frame_tokens, None, 1)]
             modulation_dims_txt = [(0, None, 1)]
         else:
-            if self.vector_in is not None:
-                vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
+            vec = vec + self.vector_in(y[:, :self.params.vec_in_dim])
             modulation_dims = None
             modulation_dims_txt = None
 
@@ -331,12 +270,6 @@ class HunyuanVideo(nn.Module):
 
         txt = self.txt_in(txt, timesteps, txt_mask, transformer_options=transformer_options)
 
-        if self.byt5_in is not None and txt_byt5 is not None:
-            txt_byt5 = self.byt5_in(txt_byt5)
-            txt_byt5_ids = torch.zeros((txt_ids.shape[0], txt_byt5.shape[1], txt_ids.shape[-1]), device=txt_ids.device, dtype=txt_ids.dtype)
-            txt = torch.cat((txt, txt_byt5), dim=1)
-            txt_ids = torch.cat((txt_ids, txt_byt5_ids), dim=1)
-
         ids = torch.cat((img_ids, txt_ids), dim=1)
         pe = self.pe_embedder(ids)
 
@@ -396,16 +329,12 @@ class HunyuanVideo(nn.Module):
 
         img = self.final_layer(img, vec, modulation_dims=modulation_dims)  # (N, T, patch_size ** 2 * out_channels)
 
-        shape = initial_shape[-len(self.patch_size):]
+        shape = initial_shape[-3:]
         for i in range(len(shape)):
             shape[i] = shape[i] // self.patch_size[i]
         img = img.reshape([img.shape[0]] + shape + [self.out_channels] + self.patch_size)
-        if img.ndim == 8:
-            img = img.permute(0, 4, 1, 5, 2, 6, 3, 7)
-            img = img.reshape(initial_shape[0], self.out_channels, initial_shape[2], initial_shape[3], initial_shape[4])
-        else:
-            img = img.permute(0, 3, 1, 4, 2, 5)
-            img = img.reshape(initial_shape[0], self.out_channels, initial_shape[2], initial_shape[3])
+        img = img.permute(0, 4, 1, 5, 2, 6, 3, 7)
+        img = img.reshape(initial_shape[0], self.out_channels, initial_shape[2], initial_shape[3], initial_shape[4])
         return img
 
     def img_ids(self, x):
@@ -420,30 +349,16 @@ class HunyuanVideo(nn.Module):
         img_ids[:, :, :, 2] = img_ids[:, :, :, 2] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).reshape(1, 1, -1)
         return repeat(img_ids, "t h w c -> b (t h w) c", b=bs)
 
-    def img_ids_2d(self, x):
-        bs, c, h, w = x.shape
-        patch_size = self.patch_size
-        h_len = ((h + (patch_size[0] // 2)) // patch_size[0])
-        w_len = ((w + (patch_size[1] // 2)) // patch_size[1])
-        img_ids = torch.zeros((h_len, w_len, 2), device=x.device, dtype=x.dtype)
-        img_ids[:, :, 0] = img_ids[:, :, 0] + torch.linspace(0, h_len - 1, steps=h_len, device=x.device, dtype=x.dtype).unsqueeze(1)
-        img_ids[:, :, 1] = img_ids[:, :, 1] + torch.linspace(0, w_len - 1, steps=w_len, device=x.device, dtype=x.dtype).unsqueeze(0)
-        return repeat(img_ids, "h w c -> b (h w) c", b=bs)
-
-    def forward(self, x, timestep, context, y=None, txt_byt5=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
+    def forward(self, x, timestep, context, y, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, control=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, y, txt_byt5, guidance, attention_mask, guiding_frame_index, ref_latent, disable_time_r, control, transformer_options, **kwargs)
+        ).execute(x, timestep, context, y, guidance, attention_mask, guiding_frame_index, ref_latent, control, transformer_options, **kwargs)
 
-    def _forward(self, x, timestep, context, y=None, txt_byt5=None, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, disable_time_r=False, control=None, transformer_options={}, **kwargs):
-        bs = x.shape[0]
-        if len(self.patch_size) == 3:
-            img_ids = self.img_ids(x)
-            txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
-        else:
-            img_ids = self.img_ids_2d(x)
-            txt_ids = torch.zeros((bs, context.shape[1], 2), device=x.device, dtype=x.dtype)
-        out = self.forward_orig(x, img_ids, context, txt_ids, attention_mask, timestep, y, txt_byt5, guidance, guiding_frame_index, ref_latent, disable_time_r=disable_time_r, control=control, transformer_options=transformer_options)
+    def _forward(self, x, timestep, context, y, guidance=None, attention_mask=None, guiding_frame_index=None, ref_latent=None, control=None, transformer_options={}, **kwargs):
+        bs, c, t, h, w = x.shape
+        img_ids = self.img_ids(x)
+        txt_ids = torch.zeros((bs, context.shape[1], 3), device=x.device, dtype=x.dtype)
+        out = self.forward_orig(x, img_ids, context, txt_ids, attention_mask, timestep, y, guidance, guiding_frame_index, ref_latent, control=control, transformer_options=transformer_options)
         return out

comfy/ldm/hunyuan_video/vae.py

@@ -1,136 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock
-import comfy.ops
-ops = comfy.ops.disable_weight_init
-
-
-class PixelShuffle2D(nn.Module):
-    def __init__(self, in_dim, out_dim, op=ops.Conv2d):
-        super().__init__()
-        self.conv = op(in_dim, out_dim >> 2, 3, 1, 1)
-        self.ratio = (in_dim << 2) // out_dim
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        h2, w2 = h >> 1, w >> 1
-        y = self.conv(x).view(b, -1, h2, 2, w2, 2).permute(0, 3, 5, 1, 2, 4).reshape(b, -1, h2, w2)
-        r = x.view(b, c, h2, 2, w2, 2).permute(0, 3, 5, 1, 2, 4).reshape(b, c << 2, h2, w2)
-        return y + r.view(b, y.shape[1], self.ratio, h2, w2).mean(2)
-
-
-class PixelUnshuffle2D(nn.Module):
-    def __init__(self, in_dim, out_dim, op=ops.Conv2d):
-        super().__init__()
-        self.conv = op(in_dim, out_dim << 2, 3, 1, 1)
-        self.scale = (out_dim << 2) // in_dim
-
-    def forward(self, x):
-        b, c, h, w = x.shape
-        h2, w2 = h << 1, w << 1
-        y = self.conv(x).view(b, 2, 2, -1, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, -1, h2, w2)
-        r = x.repeat_interleave(self.scale, 1).view(b, 2, 2, -1, h, w).permute(0, 3, 4, 1, 5, 2).reshape(b, -1, h2, w2)
-        return y + r
-
-
-class Encoder(nn.Module):
-    def __init__(self, in_channels, z_channels, block_out_channels, num_res_blocks,
-                 ffactor_spatial, downsample_match_channel=True, **_):
-        super().__init__()
-        self.z_channels = z_channels
-        self.block_out_channels = block_out_channels
-        self.num_res_blocks = num_res_blocks
-        self.conv_in = ops.Conv2d(in_channels, block_out_channels[0], 3, 1, 1)
-
-        self.down = nn.ModuleList()
-        ch = block_out_channels[0]
-        depth = (ffactor_spatial >> 1).bit_length()
-
-        for i, tgt in enumerate(block_out_channels):
-            stage = nn.Module()
-            stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
-                                                     out_channels=tgt,
-                                                     temb_channels=0,
-                                                     conv_op=ops.Conv2d)
-                                        for j in range(num_res_blocks)])
-            ch = tgt
-            if i < depth:
-                nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and downsample_match_channel else ch
-                stage.downsample = PixelShuffle2D(ch, nxt, ops.Conv2d)
-                ch = nxt
-            self.down.append(stage)
-
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
-        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv2d)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
-
-        self.norm_out = ops.GroupNorm(32, ch, 1e-6, True)
-        self.conv_out = ops.Conv2d(ch, z_channels << 1, 3, 1, 1)
-
-    def forward(self, x):
-        x = self.conv_in(x)
-
-        for stage in self.down:
-            for blk in stage.block:
-                x = blk(x)
-            if hasattr(stage, 'downsample'):
-                x = stage.downsample(x)
-
-        x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
-
-        b, c, h, w = x.shape
-        grp = c // (self.z_channels << 1)
-        skip = x.view(b, c // grp, grp, h, w).mean(2)
-
-        return self.conv_out(F.silu(self.norm_out(x))) + skip
-
-
-class Decoder(nn.Module):
-    def __init__(self, z_channels, out_channels, block_out_channels, num_res_blocks,
-                 ffactor_spatial, upsample_match_channel=True, **_):
-        super().__init__()
-        block_out_channels = block_out_channels[::-1]
-        self.z_channels = z_channels
-        self.block_out_channels = block_out_channels
-        self.num_res_blocks = num_res_blocks
-
-        ch = block_out_channels[0]
-        self.conv_in = ops.Conv2d(z_channels, ch, 3, 1, 1)
-
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
-        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv2d)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=ops.Conv2d)
-
-        self.up = nn.ModuleList()
-        depth = (ffactor_spatial >> 1).bit_length()
-
-        for i, tgt in enumerate(block_out_channels):
-            stage = nn.Module()
-            stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
-                                                     out_channels=tgt,
-                                                     temb_channels=0,
-                                                     conv_op=ops.Conv2d)
-                                        for j in range(num_res_blocks + 1)])
-            ch = tgt
-            if i < depth:
-                nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and upsample_match_channel else ch
-                stage.upsample = PixelUnshuffle2D(ch, nxt, ops.Conv2d)
-                ch = nxt
-            self.up.append(stage)
-
-        self.norm_out = ops.GroupNorm(32, ch, 1e-6, True)
-        self.conv_out = ops.Conv2d(ch, out_channels, 3, 1, 1)
-
-    def forward(self, z):
-        x = self.conv_in(z) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
-        x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
-
-        for stage in self.up:
-            for blk in stage.block:
-                x = blk(x)
-            if hasattr(stage, 'upsample'):
-                x = stage.upsample(x)
-
-        return self.conv_out(F.silu(self.norm_out(x)))

comfy/ldm/hunyuan_video/vae_refiner.py

@@ -1,267 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from comfy.ldm.modules.diffusionmodules.model import ResnetBlock, AttnBlock, VideoConv3d
-import comfy.ops
-import comfy.ldm.models.autoencoder
-ops = comfy.ops.disable_weight_init
-
-class RMS_norm(nn.Module):
-    def __init__(self, dim):
-        super().__init__()
-        shape = (dim, 1, 1, 1)
-        self.scale = dim**0.5
-        self.gamma = nn.Parameter(torch.empty(shape))
-
-    def forward(self, x):
-        return F.normalize(x, dim=1) * self.scale * self.gamma
-
-class DnSmpl(nn.Module):
-    def __init__(self, ic, oc, tds=True):
-        super().__init__()
-        fct = 2 * 2 * 2 if tds else 1 * 2 * 2
-        assert oc % fct == 0
-        self.conv = VideoConv3d(ic, oc // fct, kernel_size=3)
-
-        self.tds = tds
-        self.gs = fct * ic // oc
-
-    def forward(self, x):
-        r1 = 2 if self.tds else 1
-        h = self.conv(x)
-
-        if self.tds:
-            hf = h[:, :, :1, :, :]
-            b, c, f, ht, wd = hf.shape
-            hf = hf.reshape(b, c, f, ht // 2, 2, wd // 2, 2)
-            hf = hf.permute(0, 4, 6, 1, 2, 3, 5)
-            hf = hf.reshape(b, 2 * 2 * c, f, ht // 2, wd // 2)
-            hf = torch.cat([hf, hf], dim=1)
-
-            hn = h[:, :, 1:, :, :]
-            b, c, frms, ht, wd = hn.shape
-            nf = frms // r1
-            hn = hn.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
-            hn = hn.permute(0, 3, 5, 7, 1, 2, 4, 6)
-            hn = hn.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
-
-            h = torch.cat([hf, hn], dim=2)
-
-            xf = x[:, :, :1, :, :]
-            b, ci, f, ht, wd = xf.shape
-            xf = xf.reshape(b, ci, f, ht // 2, 2, wd // 2, 2)
-            xf = xf.permute(0, 4, 6, 1, 2, 3, 5)
-            xf = xf.reshape(b, 2 * 2 * ci, f, ht // 2, wd // 2)
-            B, C, T, H, W = xf.shape
-            xf = xf.view(B, h.shape[1], self.gs // 2, T, H, W).mean(dim=2)
-
-            xn = x[:, :, 1:, :, :]
-            b, ci, frms, ht, wd = xn.shape
-            nf = frms // r1
-            xn = xn.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
-            xn = xn.permute(0, 3, 5, 7, 1, 2, 4, 6)
-            xn = xn.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
-            B, C, T, H, W = xn.shape
-            xn = xn.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
-            sc = torch.cat([xf, xn], dim=2)
-        else:
-            b, c, frms, ht, wd = h.shape
-            nf = frms // r1
-            h = h.reshape(b, c, nf, r1, ht // 2, 2, wd // 2, 2)
-            h = h.permute(0, 3, 5, 7, 1, 2, 4, 6)
-            h = h.reshape(b, r1 * 2 * 2 * c, nf, ht // 2, wd // 2)
-
-            b, ci, frms, ht, wd = x.shape
-            nf = frms // r1
-            sc = x.reshape(b, ci, nf, r1, ht // 2, 2, wd // 2, 2)
-            sc = sc.permute(0, 3, 5, 7, 1, 2, 4, 6)
-            sc = sc.reshape(b, r1 * 2 * 2 * ci, nf, ht // 2, wd // 2)
-            B, C, T, H, W = sc.shape
-            sc = sc.view(B, h.shape[1], self.gs, T, H, W).mean(dim=2)
-
-        return h + sc
-
-
-class UpSmpl(nn.Module):
-    def __init__(self, ic, oc, tus=True):
-        super().__init__()
-        fct = 2 * 2 * 2 if tus else 1 * 2 * 2
-        self.conv = VideoConv3d(ic, oc * fct, kernel_size=3)
-
-        self.tus = tus
-        self.rp = fct * oc // ic
-
-    def forward(self, x):
-        r1 = 2 if self.tus else 1
-        h = self.conv(x)
-
-        if self.tus:
-            hf = h[:, :, :1, :, :]
-            b, c, f, ht, wd = hf.shape
-            nc = c // (2 * 2)
-            hf = hf.reshape(b, 2, 2, nc, f, ht, wd)
-            hf = hf.permute(0, 3, 4, 5, 1, 6, 2)
-            hf = hf.reshape(b, nc, f, ht * 2, wd * 2)
-            hf = hf[:, : hf.shape[1] // 2]
-
-            hn = h[:, :, 1:, :, :]
-            b, c, frms, ht, wd = hn.shape
-            nc = c // (r1 * 2 * 2)
-            hn = hn.reshape(b, r1, 2, 2, nc, frms, ht, wd)
-            hn = hn.permute(0, 4, 5, 1, 6, 2, 7, 3)
-            hn = hn.reshape(b, nc, frms * r1, ht * 2, wd * 2)
-
-            h = torch.cat([hf, hn], dim=2)
-
-            xf = x[:, :, :1, :, :]
-            b, ci, f, ht, wd = xf.shape
-            xf = xf.repeat_interleave(repeats=self.rp // 2, dim=1)
-            b, c, f, ht, wd = xf.shape
-            nc = c // (2 * 2)
-            xf = xf.reshape(b, 2, 2, nc, f, ht, wd)
-            xf = xf.permute(0, 3, 4, 5, 1, 6, 2)
-            xf = xf.reshape(b, nc, f, ht * 2, wd * 2)
-
-            xn = x[:, :, 1:, :, :]
-            xn = xn.repeat_interleave(repeats=self.rp, dim=1)
-            b, c, frms, ht, wd = xn.shape
-            nc = c // (r1 * 2 * 2)
-            xn = xn.reshape(b, r1, 2, 2, nc, frms, ht, wd)
-            xn = xn.permute(0, 4, 5, 1, 6, 2, 7, 3)
-            xn = xn.reshape(b, nc, frms * r1, ht * 2, wd * 2)
-            sc = torch.cat([xf, xn], dim=2)
-        else:
-            b, c, frms, ht, wd = h.shape
-            nc = c // (r1 * 2 * 2)
-            h = h.reshape(b, r1, 2, 2, nc, frms, ht, wd)
-            h = h.permute(0, 4, 5, 1, 6, 2, 7, 3)
-            h = h.reshape(b, nc, frms * r1, ht * 2, wd * 2)
-
-            sc = x.repeat_interleave(repeats=self.rp, dim=1)
-            b, c, frms, ht, wd = sc.shape
-            nc = c // (r1 * 2 * 2)
-            sc = sc.reshape(b, r1, 2, 2, nc, frms, ht, wd)
-            sc = sc.permute(0, 4, 5, 1, 6, 2, 7, 3)
-            sc = sc.reshape(b, nc, frms * r1, ht * 2, wd * 2)
-
-        return h + sc
-
-class Encoder(nn.Module):
-    def __init__(self, in_channels, z_channels, block_out_channels, num_res_blocks,
-                 ffactor_spatial, ffactor_temporal, downsample_match_channel=True, **_):
-        super().__init__()
-        self.z_channels = z_channels
-        self.block_out_channels = block_out_channels
-        self.num_res_blocks = num_res_blocks
-        self.conv_in = VideoConv3d(in_channels, block_out_channels[0], 3, 1, 1)
-
-        self.down = nn.ModuleList()
-        ch = block_out_channels[0]
-        depth = (ffactor_spatial >> 1).bit_length()
-        depth_temporal = ((ffactor_spatial // ffactor_temporal) >> 1).bit_length()
-
-        for i, tgt in enumerate(block_out_channels):
-            stage = nn.Module()
-            stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
-                                                     out_channels=tgt,
-                                                     temb_channels=0,
-                                                     conv_op=VideoConv3d, norm_op=RMS_norm)
-                                        for j in range(num_res_blocks)])
-            ch = tgt
-            if i < depth:
-                nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and downsample_match_channel else ch
-                stage.downsample = DnSmpl(ch, nxt, tds=i >= depth_temporal)
-                ch = nxt
-            self.down.append(stage)
-
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
-        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=RMS_norm)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
-
-        self.norm_out = RMS_norm(ch)
-        self.conv_out = VideoConv3d(ch, z_channels << 1, 3, 1, 1)
-
-        self.regul = comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer()
-
-    def forward(self, x):
-        x = self.conv_in(x)
-
-        for stage in self.down:
-            for blk in stage.block:
-                x = blk(x)
-            if hasattr(stage, 'downsample'):
-                x = stage.downsample(x)
-
-        x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
-
-        b, c, t, h, w = x.shape
-        grp = c // (self.z_channels << 1)
-        skip = x.view(b, c // grp, grp, t, h, w).mean(2)
-
-        out = self.conv_out(F.silu(self.norm_out(x))) + skip
-        out = self.regul(out)[0]
-
-        out = torch.cat((out[:, :, :1], out), dim=2)
-        out = out.permute(0, 2, 1, 3, 4)
-        b, f_times_2, c, h, w = out.shape
-        out = out.reshape(b, f_times_2 // 2, 2 * c, h, w)
-        out = out.permute(0, 2, 1, 3, 4).contiguous()
-        return out
-
-class Decoder(nn.Module):
-    def __init__(self, z_channels, out_channels, block_out_channels, num_res_blocks,
-                 ffactor_spatial, ffactor_temporal, upsample_match_channel=True, **_):
-        super().__init__()
-        block_out_channels = block_out_channels[::-1]
-        self.z_channels = z_channels
-        self.block_out_channels = block_out_channels
-        self.num_res_blocks = num_res_blocks
-
-        ch = block_out_channels[0]
-        self.conv_in = VideoConv3d(z_channels, ch, 3)
-
-        self.mid = nn.Module()
-        self.mid.block_1 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
-        self.mid.attn_1 = AttnBlock(ch, conv_op=ops.Conv3d, norm_op=RMS_norm)
-        self.mid.block_2 = ResnetBlock(in_channels=ch, out_channels=ch, temb_channels=0, conv_op=VideoConv3d, norm_op=RMS_norm)
-
-        self.up = nn.ModuleList()
-        depth = (ffactor_spatial >> 1).bit_length()
-        depth_temporal = (ffactor_temporal >> 1).bit_length()
-
-        for i, tgt in enumerate(block_out_channels):
-            stage = nn.Module()
-            stage.block = nn.ModuleList([ResnetBlock(in_channels=ch if j == 0 else tgt,
-                                                     out_channels=tgt,
-                                                     temb_channels=0,
-                                                     conv_op=VideoConv3d, norm_op=RMS_norm)
-                                        for j in range(num_res_blocks + 1)])
-            ch = tgt
-            if i < depth:
-                nxt = block_out_channels[i + 1] if i + 1 < len(block_out_channels) and upsample_match_channel else ch
-                stage.upsample = UpSmpl(ch, nxt, tus=i < depth_temporal)
-                ch = nxt
-            self.up.append(stage)
-
-        self.norm_out = RMS_norm(ch)
-        self.conv_out = VideoConv3d(ch, out_channels, 3)
-
-    def forward(self, z):
-        z = z.permute(0, 2, 1, 3, 4)
-        b, f, c, h, w = z.shape
-        z = z.reshape(b, f, 2, c // 2, h, w)
-        z = z.permute(0, 1, 2, 3, 4, 5).reshape(b, f * 2, c // 2, h, w)
-        z = z.permute(0, 2, 1, 3, 4)
-        z = z[:, :, 1:]
-
-        x = self.conv_in(z) + z.repeat_interleave(self.block_out_channels[0] // self.z_channels, 1)
-        x = self.mid.block_2(self.mid.attn_1(self.mid.block_1(x)))
-
-        for stage in self.up:
-            for blk in stage.block:
-                x = blk(x)
-            if hasattr(stage, 'upsample'):
-                x = stage.upsample(x)
-
-        return self.conv_out(F.silu(self.norm_out(x)))

comfy/ldm/models/autoencoder.py

@@ -26,12 +26,6 @@ class DiagonalGaussianRegularizer(torch.nn.Module):
             z = posterior.mode()
         return z, None
 
-class EmptyRegularizer(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, z: torch.Tensor) -> Tuple[torch.Tensor, dict]:
-        return z, None
 
 class AbstractAutoencoder(torch.nn.Module):
     """

comfy/ldm/modules/attention.py

@@ -22,7 +22,7 @@ SAGE_ATTENTION_IS_AVAILABLE = False
 try:
     from sageattention import sageattn
     SAGE_ATTENTION_IS_AVAILABLE = True
-except ImportError as e:
+except ModuleNotFoundError as e:
     if model_management.sage_attention_enabled():
         if e.name == "sageattention":
             logging.error(f"\n\nTo use the `--use-sage-attention` feature, the `sageattention` package must be installed first.\ncommand:\n\t{sys.executable} -m pip install sageattention")
@@ -34,7 +34,7 @@ FLASH_ATTENTION_IS_AVAILABLE = False
 try:
     from flash_attn import flash_attn_func
     FLASH_ATTENTION_IS_AVAILABLE = True
-except ImportError:
+except ModuleNotFoundError:
     if model_management.flash_attention_enabled():
         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)
@@ -600,8 +600,7 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape
             mask = mask.unsqueeze(1)
 
     try:
-        if mask is not None:
-            raise RuntimeError("Mask must not be set for Flash attention")
+        assert mask is None
         out = flash_attn_wrapper(
             q.transpose(1, 2),
             k.transpose(1, 2),

comfy/ldm/modules/diffusionmodules/model.py

@@ -145,7 +145,7 @@ class Downsample(nn.Module):
 
 class ResnetBlock(nn.Module):
     def __init__(self, *, in_channels, out_channels=None, conv_shortcut=False,
-                 dropout=0.0, temb_channels=512, conv_op=ops.Conv2d, norm_op=Normalize):
+                 dropout, temb_channels=512, conv_op=ops.Conv2d):
         super().__init__()
         self.in_channels = in_channels
         out_channels = in_channels if out_channels is None else out_channels
@@ -153,7 +153,7 @@ class ResnetBlock(nn.Module):
         self.use_conv_shortcut = conv_shortcut
 
         self.swish = torch.nn.SiLU(inplace=True)
-        self.norm1 = norm_op(in_channels)
+        self.norm1 = Normalize(in_channels)
         self.conv1 = conv_op(in_channels,
                                      out_channels,
                                      kernel_size=3,
@@ -162,7 +162,7 @@ class ResnetBlock(nn.Module):
         if temb_channels > 0:
             self.temb_proj = ops.Linear(temb_channels,
                                              out_channels)
-        self.norm2 = norm_op(out_channels)
+        self.norm2 = Normalize(out_channels)
         self.dropout = torch.nn.Dropout(dropout, inplace=True)
         self.conv2 = conv_op(out_channels,
                                      out_channels,
@@ -183,7 +183,7 @@ class ResnetBlock(nn.Module):
                                                     stride=1,
                                                     padding=0)
 
-    def forward(self, x, temb=None):
+    def forward(self, x, temb):
         h = x
         h = self.norm1(h)
         h = self.swish(h)
@@ -305,11 +305,11 @@ def vae_attention():
         return normal_attention
 
 class AttnBlock(nn.Module):
-    def __init__(self, in_channels, conv_op=ops.Conv2d, norm_op=Normalize):
+    def __init__(self, in_channels, conv_op=ops.Conv2d):
         super().__init__()
         self.in_channels = in_channels
 
-        self.norm = norm_op(in_channels)
+        self.norm = Normalize(in_channels)
         self.q = conv_op(in_channels,
                                  in_channels,
                                  kernel_size=1,

comfy/ldm/wan/model.py

@@ -8,7 +8,7 @@ from einops import rearrange
 
 from comfy.ldm.modules.attention import optimized_attention
 from comfy.ldm.flux.layers import EmbedND
-from comfy.ldm.flux.math import apply_rope1
+from comfy.ldm.flux.math import apply_rope
 import comfy.ldm.common_dit
 import comfy.model_management
 import comfy.patcher_extension
@@ -34,9 +34,7 @@ class WanSelfAttention(nn.Module):
                  num_heads,
                  window_size=(-1, -1),
                  qk_norm=True,
-                 eps=1e-6,
-                 kv_dim=None,
-                 operation_settings={}):
+                 eps=1e-6, operation_settings={}):
         assert dim % num_heads == 0
         super().__init__()
         self.dim = dim
@@ -45,13 +43,11 @@ class WanSelfAttention(nn.Module):
         self.window_size = window_size
         self.qk_norm = qk_norm
         self.eps = eps
-        if kv_dim is None:
-            kv_dim = dim
 
         # layers
         self.q = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.k = operation_settings.get("operations").Linear(kv_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-        self.v = operation_settings.get("operations").Linear(kv_dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
+        self.k = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
+        self.v = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
         self.o = operation_settings.get("operations").Linear(dim, dim, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
         self.norm_q = operation_settings.get("operations").RMSNorm(dim, eps=eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) if qk_norm else nn.Identity()
         self.norm_k = operation_settings.get("operations").RMSNorm(dim, eps=eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype")) if qk_norm else nn.Identity()
@@ -64,24 +60,20 @@ class WanSelfAttention(nn.Module):
         """
         b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
 
-        def qkv_fn_q(x):
+        # query, key, value function
+        def qkv_fn(x):
             q = self.norm_q(self.q(x)).view(b, s, n, d)
-            return apply_rope1(q, freqs)
-
-        def qkv_fn_k(x):
             k = self.norm_k(self.k(x)).view(b, s, n, d)
-            return apply_rope1(k, freqs)
+            v = self.v(x).view(b, s, n * d)
+            return q, k, v
 
-        #These two are VRAM hogs, so we want to do all of q computation and
-        #have pytorch garbage collect the intermediates on the sub function
-        #return before we touch k
-        q = qkv_fn_q(x)
-        k = qkv_fn_k(x)
+        q, k, v = qkv_fn(x)
+        q, k = apply_rope(q, k, freqs)
 
         x = optimized_attention(
             q.view(b, s, n * d),
             k.view(b, s, n * d),
-            self.v(x).view(b, s, n * d),
+            v,
             heads=self.num_heads,
             transformer_options=transformer_options,
         )
@@ -406,7 +398,6 @@ class WanModel(torch.nn.Module):
                  eps=1e-6,
                  flf_pos_embed_token_number=None,
                  in_dim_ref_conv=None,
-                 wan_attn_block_class=WanAttentionBlock,
                  image_model=None,
                  device=None,
                  dtype=None,
@@ -484,8 +475,8 @@ class WanModel(torch.nn.Module):
         # blocks
         cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
         self.blocks = nn.ModuleList([
-            wan_attn_block_class(cross_attn_type, dim, ffn_dim, num_heads,
-                                 window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
+            WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+                              window_size, qk_norm, cross_attn_norm, eps, operation_settings=operation_settings)
             for _ in range(num_layers)
         ])
 
@@ -1330,250 +1321,3 @@ class WanModel_S2V(WanModel):
         # unpatchify
         x = self.unpatchify(x, grid_sizes)
         return x
-
-
-class WanT2VCrossAttentionGather(WanSelfAttention):
-
-    def forward(self, x, context, transformer_options={}, **kwargs):
-        r"""
-        Args:
-            x(Tensor): Shape [B, L1, C] - video tokens
-            context(Tensor): Shape [B, L2, C] - audio tokens with shape [B, frames*16, 1536]
-        """
-        b, n, d = x.size(0), self.num_heads, self.head_dim
-
-        q = self.norm_q(self.q(x))
-        k = self.norm_k(self.k(context))
-        v = self.v(context)
-
-        # Handle audio temporal structure (16 tokens per frame)
-        k = k.reshape(-1, 16, n, d).transpose(1, 2)
-        v = v.reshape(-1, 16, n, d).transpose(1, 2)
-
-        # Handle video spatial structure
-        q = q.reshape(k.shape[0], -1, n, d).transpose(1, 2)
-
-        x = optimized_attention(q, k, v, heads=self.num_heads, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
-
-        x = x.transpose(1, 2).reshape(b, -1, n * d)
-        x = self.o(x)
-        return x
-
-
-class AudioCrossAttentionWrapper(nn.Module):
-    def __init__(self, dim, kv_dim, num_heads, qk_norm=True, eps=1e-6, operation_settings={}):
-        super().__init__()
-
-        self.audio_cross_attn = WanT2VCrossAttentionGather(dim, num_heads, qk_norm=qk_norm, kv_dim=kv_dim, eps=eps, operation_settings=operation_settings)
-        self.norm1_audio = operation_settings.get("operations").LayerNorm(dim, eps, elementwise_affine=True, device=operation_settings.get("device"), dtype=operation_settings.get("dtype"))
-
-    def forward(self, x, audio, transformer_options={}):
-        x = x + self.audio_cross_attn(self.norm1_audio(x), audio, transformer_options=transformer_options)
-        return x
-
-
-class WanAttentionBlockAudio(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)
-        self.audio_cross_attn_wrapper = AudioCrossAttentionWrapper(dim, 1536, num_heads, qk_norm, eps, operation_settings=operation_settings)
-
-    def forward(
-        self,
-        x,
-        e,
-        freqs,
-        context,
-        context_img_len=257,
-        audio=None,
-        transformer_options={},
-    ):
-        r"""
-        Args:
-            x(Tensor): Shape [B, L, C]
-            e(Tensor): Shape [B, 6, C]
-            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
-        """
-        # assert e.dtype == torch.float32
-
-        if e.ndim < 4:
-            e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device) + e).chunk(6, dim=1)
-        else:
-            e = (comfy.model_management.cast_to(self.modulation, dtype=x.dtype, device=x.device).unsqueeze(0) + e).unbind(2)
-        # assert e[0].dtype == torch.float32
-
-        # self-attention
-        y = self.self_attn(
-            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))
-
-        # cross-attention & ffn
-        x = x + self.cross_attn(self.norm3(x), context, context_img_len=context_img_len, transformer_options=transformer_options)
-        if audio is not None:
-            x = self.audio_cross_attn_wrapper(x, audio, transformer_options=transformer_options)
-        y = self.ffn(torch.addcmul(repeat_e(e[3], x), self.norm2(x), 1 + repeat_e(e[4], x)))
-        x = torch.addcmul(x, y, repeat_e(e[5], x))
-        return x
-
-class DummyAdapterLayer(nn.Module):
-    def __init__(self, layer):
-        super().__init__()
-        self.layer = layer
-
-    def forward(self, *args, **kwargs):
-        return self.layer(*args, **kwargs)
-
-
-class AudioProjModel(nn.Module):
-    def __init__(
-        self,
-        seq_len=5,
-        blocks=13,  # add a new parameter blocks
-        channels=768,  # add a new parameter channels
-        intermediate_dim=512,
-        output_dim=1536,
-        context_tokens=16,
-        device=None,
-        dtype=None,
-        operations=None,
-    ):
-        super().__init__()
-
-        self.seq_len = seq_len
-        self.blocks = blocks
-        self.channels = channels
-        self.input_dim = seq_len * blocks * channels  # update input_dim to be the product of blocks and channels.
-        self.intermediate_dim = intermediate_dim
-        self.context_tokens = context_tokens
-        self.output_dim = output_dim
-
-        # define multiple linear layers
-        self.audio_proj_glob_1 = DummyAdapterLayer(operations.Linear(self.input_dim, intermediate_dim, dtype=dtype, device=device))
-        self.audio_proj_glob_2 = DummyAdapterLayer(operations.Linear(intermediate_dim, intermediate_dim, dtype=dtype, device=device))
-        self.audio_proj_glob_3 = DummyAdapterLayer(operations.Linear(intermediate_dim, context_tokens * output_dim, dtype=dtype, device=device))
-
-        self.audio_proj_glob_norm = DummyAdapterLayer(operations.LayerNorm(output_dim, dtype=dtype, device=device))
-
-    def forward(self, audio_embeds):
-        video_length = audio_embeds.shape[1]
-        audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c")
-        batch_size, window_size, blocks, channels = audio_embeds.shape
-        audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels)
-
-        audio_embeds = torch.relu(self.audio_proj_glob_1(audio_embeds))
-        audio_embeds = torch.relu(self.audio_proj_glob_2(audio_embeds))
-
-        context_tokens = self.audio_proj_glob_3(audio_embeds).reshape(batch_size, self.context_tokens, self.output_dim)
-
-        context_tokens = self.audio_proj_glob_norm(context_tokens)
-        context_tokens = rearrange(context_tokens, "(bz f) m c -> bz f m c", f=video_length)
-
-        return context_tokens
-
-
-class HumoWanModel(WanModel):
-    r"""
-    Wan diffusion backbone supporting both text-to-video and image-to-video.
-    """
-
-    def __init__(self,
-                 model_type='humo',
-                 patch_size=(1, 2, 2),
-                 text_len=512,
-                 in_dim=16,
-                 dim=2048,
-                 ffn_dim=8192,
-                 freq_dim=256,
-                 text_dim=4096,
-                 out_dim=16,
-                 num_heads=16,
-                 num_layers=32,
-                 window_size=(-1, -1),
-                 qk_norm=True,
-                 cross_attn_norm=True,
-                 eps=1e-6,
-                 flf_pos_embed_token_number=None,
-                 image_model=None,
-                 audio_token_num=16,
-                 device=None,
-                 dtype=None,
-                 operations=None,
-                 ):
-
-        super().__init__(model_type='t2v', 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, wan_attn_block_class=WanAttentionBlockAudio, image_model=image_model, device=device, dtype=dtype, operations=operations)
-
-        self.audio_proj = AudioProjModel(seq_len=8, blocks=5, channels=1280, intermediate_dim=512, output_dim=1536, context_tokens=audio_token_num, dtype=dtype, device=device, operations=operations)
-
-    def forward_orig(
-        self,
-        x,
-        t,
-        context,
-        freqs=None,
-        audio_embed=None,
-        reference_latent=None,
-        transformer_options={},
-        **kwargs,
-    ):
-        bs, _, time, height, width = x.shape
-
-        # embeddings
-        x = self.patch_embedding(x.float()).to(x.dtype)
-        grid_sizes = x.shape[2:]
-        x = x.flatten(2).transpose(1, 2)
-
-        # time embeddings
-        e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
-        e = e.reshape(t.shape[0], -1, e.shape[-1])
-        e0 = self.time_projection(e).unflatten(2, (6, self.dim))
-
-        if reference_latent is not None:
-            ref = self.patch_embedding(reference_latent.float()).to(x.dtype)
-            ref = ref.flatten(2).transpose(1, 2)
-            freqs_ref = self.rope_encode(reference_latent.shape[-3], reference_latent.shape[-2], reference_latent.shape[-1], t_start=time, device=x.device, dtype=x.dtype)
-            x = torch.cat([x, ref], dim=1)
-            freqs = torch.cat([freqs, freqs_ref], dim=1)
-            del ref, freqs_ref
-
-        # context
-        context = self.text_embedding(context)
-        context_img_len = None
-
-        if audio_embed is not None:
-            if reference_latent is not None:
-                zero_audio_pad = torch.zeros(audio_embed.shape[0], reference_latent.shape[-3], *audio_embed.shape[2:], device=audio_embed.device, dtype=audio_embed.dtype)
-                audio_embed = torch.cat([audio_embed, zero_audio_pad], dim=1)
-            audio = self.audio_proj(audio_embed).permute(0, 3, 1, 2).flatten(2).transpose(1, 2)
-        else:
-            audio = None
-
-        patches_replace = transformer_options.get("patches_replace", {})
-        blocks_replace = patches_replace.get("dit", {})
-        for i, block in enumerate(self.blocks):
-            if ("double_block", i) in blocks_replace:
-                def block_wrap(args):
-                    out = {}
-                    out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, audio=audio, transformer_options=args["transformer_options"])
-                    return out
-                out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
-                x = out["img"]
-            else:
-                x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, audio=audio, transformer_options=transformer_options)
-
-        # head
-        x = self.head(x, e)
-
-        # unpatchify
-        x = self.unpatchify(x, grid_sizes)
-        return x

comfy/ldm/wan/model_animate.py

@@ -1,548 +0,0 @@
-from torch import nn
-import torch
-from typing import Tuple, Optional
-from einops import rearrange
-import torch.nn.functional as F
-import math
-from .model import WanModel, sinusoidal_embedding_1d
-from comfy.ldm.modules.attention import optimized_attention
-import comfy.model_management
-
-class CausalConv1d(nn.Module):
-
-    def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", operations=None, **kwargs):
-        super().__init__()
-
-        self.pad_mode = pad_mode
-        padding = (kernel_size - 1, 0)  # T
-        self.time_causal_padding = padding
-
-        self.conv = operations.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
-
-    def forward(self, x):
-        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
-        return self.conv(x)
-
-
-class FaceEncoder(nn.Module):
-    def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, dtype=None, device=None, operations=None):
-        factory_kwargs = {"dtype": dtype, "device": device}
-        super().__init__()
-
-        self.num_heads = num_heads
-        self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1, operations=operations, **factory_kwargs)
-        self.norm1 = operations.LayerNorm(hidden_dim // 8, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-        self.act = nn.SiLU()
-        self.conv2 = CausalConv1d(1024, 1024, 3, stride=2, operations=operations, **factory_kwargs)
-        self.conv3 = CausalConv1d(1024, 1024, 3, stride=2, operations=operations, **factory_kwargs)
-
-        self.out_proj = operations.Linear(1024, hidden_dim, **factory_kwargs)
-        self.norm1 = operations.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.norm2 = operations.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.norm3 = operations.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.padding_tokens = nn.Parameter(torch.empty(1, 1, 1, hidden_dim, **factory_kwargs))
-
-    def forward(self, x):
-
-        x = rearrange(x, "b t c -> b c t")
-        b, c, t = x.shape
-
-        x = self.conv1_local(x)
-        x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads)
-
-        x = self.norm1(x)
-        x = self.act(x)
-        x = rearrange(x, "b t c -> b c t")
-        x = self.conv2(x)
-        x = rearrange(x, "b c t -> b t c")
-        x = self.norm2(x)
-        x = self.act(x)
-        x = rearrange(x, "b t c -> b c t")
-        x = self.conv3(x)
-        x = rearrange(x, "b c t -> b t c")
-        x = self.norm3(x)
-        x = self.act(x)
-        x = self.out_proj(x)
-        x = rearrange(x, "(b n) t c -> b t n c", b=b)
-        padding = comfy.model_management.cast_to(self.padding_tokens, dtype=x.dtype, device=x.device).repeat(b, x.shape[1], 1, 1)
-        x = torch.cat([x, padding], dim=-2)
-        x_local = x.clone()
-
-        return x_local
-
-
-def get_norm_layer(norm_layer, operations=None):
-    """
-    Get the normalization layer.
-
-    Args:
-        norm_layer (str): The type of normalization layer.
-
-    Returns:
-        norm_layer (nn.Module): The normalization layer.
-    """
-    if norm_layer == "layer":
-        return operations.LayerNorm
-    elif norm_layer == "rms":
-        return operations.RMSNorm
-    else:
-        raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
-
-
-class FaceAdapter(nn.Module):
-    def __init__(
-        self,
-        hidden_dim: int,
-        heads_num: int,
-        qk_norm: bool = True,
-        qk_norm_type: str = "rms",
-        num_adapter_layers: int = 1,
-        dtype=None, device=None, operations=None
-    ):
-
-        factory_kwargs = {"dtype": dtype, "device": device}
-        super().__init__()
-        self.hidden_size = hidden_dim
-        self.heads_num = heads_num
-        self.fuser_blocks = nn.ModuleList(
-            [
-                FaceBlock(
-                    self.hidden_size,
-                    self.heads_num,
-                    qk_norm=qk_norm,
-                    qk_norm_type=qk_norm_type,
-                    operations=operations,
-                    **factory_kwargs,
-                )
-                for _ in range(num_adapter_layers)
-            ]
-        )
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        motion_embed: torch.Tensor,
-        idx: int,
-        freqs_cis_q: Tuple[torch.Tensor, torch.Tensor] = None,
-        freqs_cis_k: Tuple[torch.Tensor, torch.Tensor] = None,
-    ) -> torch.Tensor:
-
-        return self.fuser_blocks[idx](x, motion_embed, freqs_cis_q, freqs_cis_k)
-
-
-
-class FaceBlock(nn.Module):
-    def __init__(
-        self,
-        hidden_size: int,
-        heads_num: int,
-        qk_norm: bool = True,
-        qk_norm_type: str = "rms",
-        qk_scale: float = None,
-        dtype: Optional[torch.dtype] = None,
-        device: Optional[torch.device] = None,
-        operations=None
-    ):
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super().__init__()
-
-        self.deterministic = False
-        self.hidden_size = hidden_size
-        self.heads_num = heads_num
-        head_dim = hidden_size // heads_num
-        self.scale = qk_scale or head_dim**-0.5
-
-        self.linear1_kv = operations.Linear(hidden_size, hidden_size * 2, **factory_kwargs)
-        self.linear1_q = operations.Linear(hidden_size, hidden_size, **factory_kwargs)
-
-        self.linear2 = operations.Linear(hidden_size, hidden_size, **factory_kwargs)
-
-        qk_norm_layer = get_norm_layer(qk_norm_type, operations=operations)
-        self.q_norm = (
-            qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
-        )
-        self.k_norm = (
-            qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
-        )
-
-        self.pre_norm_feat = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-        self.pre_norm_motion = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
-
-    def forward(
-        self,
-        x: torch.Tensor,
-        motion_vec: torch.Tensor,
-        motion_mask: Optional[torch.Tensor] = None,
-        # use_context_parallel=False,
-    ) -> torch.Tensor:
-
-        B, T, N, C = motion_vec.shape
-        T_comp = T
-
-        x_motion = self.pre_norm_motion(motion_vec)
-        x_feat = self.pre_norm_feat(x)
-
-        kv = self.linear1_kv(x_motion)
-        q = self.linear1_q(x_feat)
-
-        k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num)
-        q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num)
-
-        # Apply QK-Norm if needed.
-        q = self.q_norm(q).to(v)
-        k = self.k_norm(k).to(v)
-
-        k = rearrange(k, "B L N H D -> (B L) N H D")
-        v = rearrange(v, "B L N H D -> (B L) N H D")
-
-        q = rearrange(q, "B (L S) H D -> (B L) S (H D)", L=T_comp)
-
-        attn = optimized_attention(q, k, v, heads=self.heads_num)
-
-        attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T_comp)
-
-        output = self.linear2(attn)
-
-        if motion_mask is not None:
-            output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1)
-
-        return output
-
-# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/ops/upfirdn2d/upfirdn2d.py#L162
-def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
-    _, minor, in_h, in_w = input.shape
-    kernel_h, kernel_w = kernel.shape
-
-    out = input.view(-1, minor, in_h, 1, in_w, 1)
-    out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
-    out = out.view(-1, minor, in_h * up_y, in_w * up_x)
-
-    out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
-    out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0), max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0)]
-
-    out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
-    w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
-    out = F.conv2d(out, w)
-    out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1, in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1)
-    return out[:, :, ::down_y, ::down_x]
-
-def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
-    return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
-
-# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/ops/fused_act/fused_act.py#L81
-class FusedLeakyReLU(torch.nn.Module):
-    def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5, dtype=None, device=None):
-        super().__init__()
-        self.bias = torch.nn.Parameter(torch.empty(1, channel, 1, 1, dtype=dtype, device=device))
-        self.negative_slope = negative_slope
-        self.scale = scale
-
-    def forward(self, input):
-        return fused_leaky_relu(input, comfy.model_management.cast_to(self.bias, device=input.device, dtype=input.dtype), self.negative_slope, self.scale)
-
-def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
-    return F.leaky_relu(input + bias, negative_slope) * scale
-
-class Blur(torch.nn.Module):
-    def __init__(self, kernel, pad, dtype=None, device=None):
-        super().__init__()
-        kernel = torch.tensor(kernel, dtype=dtype, device=device)
-        kernel = kernel[None, :] * kernel[:, None]
-        kernel = kernel / kernel.sum()
-        self.register_buffer('kernel', kernel)
-        self.pad = pad
-
-    def forward(self, input):
-        return upfirdn2d(input, comfy.model_management.cast_to(self.kernel, dtype=input.dtype, device=input.device), pad=self.pad)
-
-#https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L590
-class ScaledLeakyReLU(torch.nn.Module):
-    def __init__(self, negative_slope=0.2):
-        super().__init__()
-        self.negative_slope = negative_slope
-
-    def forward(self, input):
-        return F.leaky_relu(input, negative_slope=self.negative_slope)
-
-# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L605
-class EqualConv2d(torch.nn.Module):
-    def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.weight = torch.nn.Parameter(torch.empty(out_channel, in_channel, kernel_size, kernel_size, device=device, dtype=dtype))
-        self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
-        self.stride = stride
-        self.padding = padding
-        self.bias = torch.nn.Parameter(torch.empty(out_channel, device=device, dtype=dtype)) if bias else None
-
-    def forward(self, input):
-        if self.bias is None:
-            bias = None
-        else:
-            bias = comfy.model_management.cast_to(self.bias, device=input.device, dtype=input.dtype)
-
-        return F.conv2d(input, comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) * self.scale, bias=bias, stride=self.stride, padding=self.padding)
-
-# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L134
-class EqualLinear(torch.nn.Module):
-    def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.weight = torch.nn.Parameter(torch.empty(out_dim, in_dim, device=device, dtype=dtype))
-        self.bias = torch.nn.Parameter(torch.empty(out_dim, device=device, dtype=dtype)) if bias else None
-        self.activation = activation
-        self.scale = (1 / math.sqrt(in_dim)) * lr_mul
-        self.lr_mul = lr_mul
-
-    def forward(self, input):
-        if self.bias is None:
-            bias = None
-        else:
-            bias = comfy.model_management.cast_to(self.bias, device=input.device, dtype=input.dtype) * self.lr_mul
-
-        if self.activation:
-            out = F.linear(input, comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) * self.scale)
-            return fused_leaky_relu(out, bias)
-        return F.linear(input, comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) * self.scale, bias=bias)
-
-# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L654
-class ConvLayer(torch.nn.Sequential):
-    def __init__(self, in_channel, out_channel, kernel_size, downsample=False, blur_kernel=[1, 3, 3, 1], bias=True, activate=True, dtype=None, device=None, operations=None):
-        layers = []
-
-        if downsample:
-            factor = 2
-            p = (len(blur_kernel) - factor) + (kernel_size - 1)
-            layers.append(Blur(blur_kernel, pad=((p + 1) // 2, p // 2)))
-            stride, padding = 2, 0
-        else:
-            stride, padding = 1, kernel_size // 2
-
-        layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=padding, stride=stride, bias=bias and not activate, dtype=dtype, device=device, operations=operations))
-
-        if activate:
-            layers.append(FusedLeakyReLU(out_channel) if bias else ScaledLeakyReLU(0.2))
-
-        super().__init__(*layers)
-
-# https://github.com/XPixelGroup/BasicSR/blob/8d56e3a045f9fb3e1d8872f92ee4a4f07f886b0a/basicsr/archs/stylegan2_arch.py#L704
-class ResBlock(torch.nn.Module):
-    def __init__(self, in_channel, out_channel, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.conv1 = ConvLayer(in_channel, in_channel, 3, dtype=dtype, device=device, operations=operations)
-        self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True, dtype=dtype, device=device, operations=operations)
-        self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False, dtype=dtype, device=device, operations=operations)
-
-    def forward(self, input):
-        out = self.conv2(self.conv1(input))
-        skip = self.skip(input)
-        return (out + skip) / math.sqrt(2)
-
-
-class EncoderApp(torch.nn.Module):
-    def __init__(self, w_dim=512, dtype=None, device=None, operations=None):
-        super().__init__()
-        kwargs = {"device": device, "dtype": dtype, "operations": operations}
-
-        self.convs = torch.nn.ModuleList([
-            ConvLayer(3, 32, 1, **kwargs), ResBlock(32, 64, **kwargs),
-            ResBlock(64, 128, **kwargs), ResBlock(128, 256, **kwargs),
-            ResBlock(256, 512, **kwargs), ResBlock(512, 512, **kwargs),
-            ResBlock(512, 512, **kwargs), ResBlock(512, 512, **kwargs),
-            EqualConv2d(512, w_dim, 4, padding=0, bias=False, **kwargs)
-        ])
-
-    def forward(self, x):
-        h = x
-        for conv in self.convs:
-            h = conv(h)
-        return h.squeeze(-1).squeeze(-1)
-
-class Encoder(torch.nn.Module):
-    def __init__(self, dim=512, motion_dim=20, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.net_app = EncoderApp(dim, dtype=dtype, device=device, operations=operations)
-        self.fc = torch.nn.Sequential(*[EqualLinear(dim, dim, dtype=dtype, device=device, operations=operations) for _ in range(4)] + [EqualLinear(dim, motion_dim, dtype=dtype, device=device, operations=operations)])
-
-    def encode_motion(self, x):
-        return self.fc(self.net_app(x))
-
-class Direction(torch.nn.Module):
-    def __init__(self, motion_dim, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.weight = torch.nn.Parameter(torch.empty(512, motion_dim, device=device, dtype=dtype))
-        self.motion_dim = motion_dim
-
-    def forward(self, input):
-        stabilized_weight = comfy.model_management.cast_to(self.weight, device=input.device, dtype=input.dtype) + 1e-8 * torch.eye(512, self.motion_dim, device=input.device, dtype=input.dtype)
-        Q, _ = torch.linalg.qr(stabilized_weight.float())
-        if input is None:
-            return Q
-        return torch.sum(input.unsqueeze(-1) * Q.T.to(input.dtype), dim=1)
-
-class Synthesis(torch.nn.Module):
-    def __init__(self, motion_dim, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.direction = Direction(motion_dim, dtype=dtype, device=device, operations=operations)
-
-class Generator(torch.nn.Module):
-    def __init__(self, style_dim=512, motion_dim=20, dtype=None, device=None, operations=None):
-        super().__init__()
-        self.enc = Encoder(style_dim, motion_dim, dtype=dtype, device=device, operations=operations)
-        self.dec = Synthesis(motion_dim, dtype=dtype, device=device, operations=operations)
-
-    def get_motion(self, img):
-        motion_feat = self.enc.encode_motion(img)
-        return self.dec.direction(motion_feat)
-
-class AnimateWanModel(WanModel):
-    r"""
-    Wan diffusion backbone supporting both text-to-video and image-to-video.
-    """
-
-    def __init__(self,
-                 model_type='animate',
-                 patch_size=(1, 2, 2),
-                 text_len=512,
-                 in_dim=16,
-                 dim=2048,
-                 ffn_dim=8192,
-                 freq_dim=256,
-                 text_dim=4096,
-                 out_dim=16,
-                 num_heads=16,
-                 num_layers=32,
-                 window_size=(-1, -1),
-                 qk_norm=True,
-                 cross_attn_norm=True,
-                 eps=1e-6,
-                 flf_pos_embed_token_number=None,
-                 motion_encoder_dim=512,
-                 image_model=None,
-                 device=None,
-                 dtype=None,
-                 operations=None,
-                 ):
-
-        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, image_model=image_model, device=device, dtype=dtype, operations=operations)
-
-        self.pose_patch_embedding = operations.Conv3d(
-            16, dim, kernel_size=patch_size, stride=patch_size, device=device, dtype=dtype
-        )
-
-        self.motion_encoder = Generator(style_dim=512, motion_dim=20, device=device, dtype=dtype, operations=operations)
-
-        self.face_adapter = FaceAdapter(
-            heads_num=self.num_heads,
-            hidden_dim=self.dim,
-            num_adapter_layers=self.num_layers // 5,
-            device=device, dtype=dtype, operations=operations
-        )
-
-        self.face_encoder = FaceEncoder(
-            in_dim=motion_encoder_dim,
-            hidden_dim=self.dim,
-            num_heads=4,
-            device=device, dtype=dtype, operations=operations
-        )
-
-    def after_patch_embedding(self, x, pose_latents, face_pixel_values):
-        if pose_latents is not None:
-            pose_latents = self.pose_patch_embedding(pose_latents)
-            x[:, :, 1:pose_latents.shape[2] + 1] += pose_latents[:, :, :x.shape[2] - 1]
-
-        if face_pixel_values is None:
-            return x, None
-
-        b, c, T, h, w = face_pixel_values.shape
-        face_pixel_values = rearrange(face_pixel_values, "b c t h w -> (b t) c h w")
-        encode_bs = 8
-        face_pixel_values_tmp = []
-        for i in range(math.ceil(face_pixel_values.shape[0] / encode_bs)):
-            face_pixel_values_tmp.append(self.motion_encoder.get_motion(face_pixel_values[i * encode_bs: (i + 1) * encode_bs]))
-
-        motion_vec = torch.cat(face_pixel_values_tmp)
-
-        motion_vec = rearrange(motion_vec, "(b t) c -> b t c", t=T)
-        motion_vec = self.face_encoder(motion_vec)
-
-        B, L, H, C = motion_vec.shape
-        pad_face = torch.zeros(B, 1, H, C).type_as(motion_vec)
-        motion_vec = torch.cat([pad_face, motion_vec], dim=1)
-
-        if motion_vec.shape[1] < x.shape[2]:
-            B, L, H, C = motion_vec.shape
-            pad = torch.zeros(B, x.shape[2] - motion_vec.shape[1], H, C).type_as(motion_vec)
-            motion_vec = torch.cat([motion_vec, pad], dim=1)
-        else:
-            motion_vec = motion_vec[:, :x.shape[2]]
-        return x, motion_vec
-
-    def forward_orig(
-        self,
-        x,
-        t,
-        context,
-        clip_fea=None,
-        pose_latents=None,
-        face_pixel_values=None,
-        freqs=None,
-        transformer_options={},
-        **kwargs,
-    ):
-        # embeddings
-        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:]
-        x = x.flatten(2).transpose(1, 2)
-
-        # time embeddings
-        e = self.time_embedding(
-            sinusoidal_embedding_1d(self.freq_dim, t.flatten()).to(dtype=x[0].dtype))
-        e = e.reshape(t.shape[0], -1, e.shape[-1])
-        e0 = self.time_projection(e).unflatten(2, (6, self.dim))
-
-        full_ref = None
-        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)
-
-        # context
-        context = self.text_embedding(context)
-
-        context_img_len = None
-        if clip_fea is not None:
-            if self.img_emb is not None:
-                context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
-                context = torch.concat([context_clip, context], dim=1)
-            context_img_len = clip_fea.shape[-2]
-
-        patches_replace = transformer_options.get("patches_replace", {})
-        blocks_replace = patches_replace.get("dit", {})
-        for i, block in enumerate(self.blocks):
-            if ("double_block", i) in blocks_replace:
-                def block_wrap(args):
-                    out = {}
-                    out["img"] = block(args["img"], context=args["txt"], e=args["vec"], freqs=args["pe"], context_img_len=context_img_len, transformer_options=args["transformer_options"])
-                    return out
-                out = blocks_replace[("double_block", i)]({"img": x, "txt": context, "vec": e0, "pe": freqs, "transformer_options": transformer_options}, {"original_block": block_wrap})
-                x = out["img"]
-            else:
-                x = block(x, e=e0, freqs=freqs, context=context, context_img_len=context_img_len, transformer_options=transformer_options)
-
-            if i % 5 == 0 and motion_vec is not None:
-                x = x + self.face_adapter.fuser_blocks[i // 5](x, motion_vec)
-
-        # head
-        x = self.head(x, e)
-
-        if full_ref is not None:
-            x = x[:, full_ref.shape[1]:]
-
-        # unpatchify
-        x = self.unpatchify(x, grid_sizes)
-        return x

comfy/lora.py

@@ -260,10 +260,6 @@ def model_lora_keys_unet(model, key_map={}):
                 key_map["transformer.{}".format(k[:-len(".weight")])] = to #simpletrainer and probably regular diffusers flux lora format
                 key_map["lycoris_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #simpletrainer lycoris
                 key_map["lora_transformer_{}".format(k[:-len(".weight")].replace(".", "_"))] = to #onetrainer
-        for k in sdk:
-            hidden_size = model.model_config.unet_config.get("hidden_size", 0)
-            if k.endswith(".weight") and ".linear1." in k:
-                key_map["{}".format(k.replace(".linear1.weight", ".linear1_qkv"))] = (k, (0, 0, hidden_size * 3))
 
     if isinstance(model, comfy.model_base.GenmoMochi):
         for k in sdk:
@@ -297,12 +293,6 @@ def model_lora_keys_unet(model, key_map={}):
                 key_lora = k[len("diffusion_model."):-len(".weight")]
                 key_map["{}".format(key_lora)] = k
 
-    if isinstance(model, comfy.model_base.Omnigen2):
-        for k in sdk:
-            if k.startswith("diffusion_model.") and k.endswith(".weight"):
-                key_lora = k[len("diffusion_model."):-len(".weight")]
-                key_map["{}".format(key_lora)] = k
-
     if isinstance(model, comfy.model_base.QwenImage):
         for k in sdk:
             if k.startswith("diffusion_model.") and k.endswith(".weight"): #QwenImage lora format

comfy/lora_convert.py

@@ -15,29 +15,10 @@ def convert_lora_bfl_control(sd): #BFL loras for Flux
 def convert_lora_wan_fun(sd): #Wan Fun loras
     return comfy.utils.state_dict_prefix_replace(sd, {"lora_unet__": "lora_unet_"})
 
-def convert_uso_lora(sd):
-    sd_out = {}
-    for k in sd:
-        tensor = sd[k]
-        k_to = "diffusion_model.{}".format(k.replace(".down.weight", ".lora_down.weight")
-                                           .replace(".up.weight", ".lora_up.weight")
-                                           .replace(".qkv_lora2.", ".txt_attn.qkv.")
-                                           .replace(".qkv_lora1.", ".img_attn.qkv.")
-                                           .replace(".proj_lora1.", ".img_attn.proj.")
-                                           .replace(".proj_lora2.", ".txt_attn.proj.")
-                                           .replace(".qkv_lora.", ".linear1_qkv.")
-                                           .replace(".proj_lora.", ".linear2.")
-                                           .replace(".processor.", ".")
-                                           )
-        sd_out[k_to] = tensor
-    return sd_out
-
 
 def convert_lora(sd):
     if "img_in.lora_A.weight" in sd and "single_blocks.0.norm.key_norm.scale" in sd:
         return convert_lora_bfl_control(sd)
     if "lora_unet__blocks_0_cross_attn_k.lora_down.weight" in sd:
         return convert_lora_wan_fun(sd)
-    if "single_blocks.37.processor.qkv_lora.up.weight" in sd and "double_blocks.18.processor.qkv_lora2.up.weight" in sd:
-        return convert_uso_lora(sd)
     return sd

comfy/model_base.py

@@ -16,8 +16,6 @@
     along with this program.  If not, see <https://www.gnu.org/licenses/>.
 """
 
-import comfy.ldm.hunyuan3dv2_1
-import comfy.ldm.hunyuan3dv2_1.hunyuandit
 import torch
 import logging
 from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
@@ -39,11 +37,9 @@ import comfy.ldm.cosmos.model
 import comfy.ldm.cosmos.predict2
 import comfy.ldm.lumina.model
 import comfy.ldm.wan.model
-import comfy.ldm.wan.model_animate
 import comfy.ldm.hunyuan3d.model
 import comfy.ldm.hidream.model
 import comfy.ldm.chroma.model
-import comfy.ldm.chroma_radiance.model
 import comfy.ldm.ace.model
 import comfy.ldm.omnigen.omnigen2
 import comfy.ldm.qwen_image.model
@@ -1214,63 +1210,6 @@ class WAN21_Camera(WAN21):
             out['camera_conditions'] = comfy.conds.CONDRegular(camera_conditions)
         return out
 
-class WAN21_HuMo(WAN21):
-    def __init__(self, model_config, model_type=ModelType.FLOW, image_to_video=False, device=None):
-        super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model.HumoWanModel)
-        self.image_to_video = image_to_video
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        noise = kwargs.get("noise", None)
-
-        audio_embed = kwargs.get("audio_embed", None)
-        if audio_embed is not None:
-            out['audio_embed'] = comfy.conds.CONDRegular(audio_embed)
-
-        if "c_concat" not in out:  # 1.7B model
-            reference_latents = kwargs.get("reference_latents", None)
-            if reference_latents is not None:
-                out['reference_latent'] = comfy.conds.CONDRegular(self.process_latent_in(reference_latents[-1]))
-        else:
-            noise_shape = list(noise.shape)
-            noise_shape[1] += 4
-            concat_latent = torch.zeros(noise_shape, device=noise.device, dtype=noise.dtype)
-            zero_vae_values_first = torch.tensor([0.8660, -0.4326, -0.0017, -0.4884, -0.5283, 0.9207, -0.9896, 0.4433, -0.5543, -0.0113, 0.5753, -0.6000, -0.8346, -0.3497, -0.1926, -0.6938]).view(1, 16, 1, 1, 1)
-            zero_vae_values_second = torch.tensor([1.0869, -1.2370, 0.0206, -0.4357, -0.6411, 2.0307, -1.5972, 1.2659, -0.8595, -0.4654, 0.9638, -1.6330, -1.4310, -0.1098, -0.3856, -1.4583]).view(1, 16, 1, 1, 1)
-            zero_vae_values = torch.tensor([0.8642, -1.8583, 0.1577, 0.1350, -0.3641, 2.5863, -1.9670, 1.6065, -1.0475, -0.8678, 1.1734, -1.8138, -1.5933, -0.7721, -0.3289, -1.3745]).view(1, 16, 1, 1, 1)
-            concat_latent[:, 4:] = zero_vae_values
-            concat_latent[:, 4:, :1] = zero_vae_values_first
-            concat_latent[:, 4:, 1:2] = zero_vae_values_second
-            out['c_concat'] = comfy.conds.CONDNoiseShape(concat_latent)
-            reference_latents = kwargs.get("reference_latents", None)
-            if reference_latents is not None:
-                ref_latent = self.process_latent_in(reference_latents[-1])
-                ref_latent_shape = list(ref_latent.shape)
-                ref_latent_shape[1] += 4 + ref_latent_shape[1]
-                ref_latent_full = torch.zeros(ref_latent_shape, device=ref_latent.device, dtype=ref_latent.dtype)
-                ref_latent_full[:, 20:] = ref_latent
-                ref_latent_full[:, 16:20] = 1.0
-                out['reference_latent'] = comfy.conds.CONDRegular(ref_latent_full)
-
-        return out
-
-class WAN22_Animate(WAN21):
-    def __init__(self, model_config, model_type=ModelType.FLOW, image_to_video=False, device=None):
-        super(WAN21, self).__init__(model_config, model_type, device=device, unet_model=comfy.ldm.wan.model_animate.AnimateWanModel)
-        self.image_to_video = image_to_video
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-
-        face_video_pixels = kwargs.get("face_video_pixels", None)
-        if face_video_pixels is not None:
-            out['face_pixel_values'] = comfy.conds.CONDRegular(face_video_pixels)
-
-        pose_latents = kwargs.get("pose_video_latent", None)
-        if pose_latents is not None:
-            out['pose_latents'] = comfy.conds.CONDRegular(self.process_latent_in(pose_latents))
-        return out
-
 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)
@@ -1343,21 +1282,6 @@ class Hunyuan3Dv2(BaseModel):
             out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
         return out
 
-class Hunyuan3Dv2_1(BaseModel):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan3dv2_1.hunyuandit.HunYuanDiTPlain)
-
-    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)
-
-        guidance = kwargs.get("guidance", 5.0)
-        if guidance is not None:
-            out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
-        return out
-
 class HiDream(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hidream.model.HiDreamImageTransformer2DModel)
@@ -1379,8 +1303,8 @@ class HiDream(BaseModel):
         return out
 
 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)
+    def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.chroma.model.Chroma)
 
     def extra_conds(self, **kwargs):
         out = super().extra_conds(**kwargs)
@@ -1390,10 +1314,6 @@ class Chroma(Flux):
             out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
         return out
 
-class ChromaRadiance(Chroma):
-    def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.chroma_radiance.model.ChromaRadiance)
-
 class ACEStep(BaseModel):
     def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ace.model.ACEStepTransformer2DModel)
@@ -1471,55 +1391,3 @@ class QwenImage(BaseModel):
         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):
-    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
-        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan_video.model.HunyuanVideo)
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        attention_mask = kwargs.get("attention_mask", None)
-        if attention_mask is not None:
-            if torch.numel(attention_mask) != attention_mask.sum():
-                out['attention_mask'] = comfy.conds.CONDRegular(attention_mask)
-        cross_attn = kwargs.get("cross_attn", None)
-        if cross_attn is not None:
-            out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
-
-        conditioning_byt5small = kwargs.get("conditioning_byt5small", None)
-        if conditioning_byt5small is not None:
-            out['txt_byt5'] = comfy.conds.CONDRegular(conditioning_byt5small)
-
-        guidance = kwargs.get("guidance", 6.0)
-        if guidance is not None:
-            out['guidance'] = comfy.conds.CONDRegular(torch.FloatTensor([guidance]))
-
-        return out
-
-class HunyuanImage21Refiner(HunyuanImage21):
-    def concat_cond(self, **kwargs):
-        noise = kwargs.get("noise", None)
-        image = kwargs.get("concat_latent_image", None)
-        noise_augmentation = kwargs.get("noise_augmentation", 0.0)
-        device = kwargs["device"]
-
-        if image is None:
-            shape_image = list(noise.shape)
-            image = torch.zeros(shape_image, dtype=noise.dtype, layout=noise.layout, device=noise.device)
-        else:
-            image = utils.common_upscale(image.to(device), noise.shape[-1], noise.shape[-2], "bilinear", "center")
-            image = self.process_latent_in(image)
-            image = utils.resize_to_batch_size(image, noise.shape[0])
-            if noise_augmentation > 0:
-                generator = torch.Generator(device="cpu")
-                generator.manual_seed(kwargs.get("seed", 0) - 10)
-                noise = torch.randn(image.shape, generator=generator, dtype=image.dtype, device="cpu").to(image.device)
-                image = noise_augmentation * noise + min(1.0 - noise_augmentation, 0.75) * image
-            else:
-                image = 0.75 * image
-        return image
-
-    def extra_conds(self, **kwargs):
-        out = super().extra_conds(**kwargs)
-        out['disable_time_r'] = comfy.conds.CONDConstant(True)
-        return out

comfy/model_detection.py

@@ -136,45 +136,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
 
     if '{}txt_in.individual_token_refiner.blocks.0.norm1.weight'.format(key_prefix) in state_dict_keys: #Hunyuan Video
         dit_config = {}
-        in_w = state_dict['{}img_in.proj.weight'.format(key_prefix)]
-        out_w = state_dict['{}final_layer.linear.weight'.format(key_prefix)]
         dit_config["image_model"] = "hunyuan_video"
-        dit_config["in_channels"] = in_w.shape[1] #SkyReels img2video has 32 input channels
-        dit_config["patch_size"] = list(in_w.shape[2:])
-        dit_config["out_channels"] = out_w.shape[0] // math.prod(dit_config["patch_size"])
-        if any(s.startswith('{}vector_in.'.format(key_prefix)) for s in state_dict_keys):
-            dit_config["vec_in_dim"] = 768
-        else:
-            dit_config["vec_in_dim"] = None
-
-        if len(dit_config["patch_size"]) == 2:
-            dit_config["axes_dim"] = [64, 64]
-        else:
-            dit_config["axes_dim"] = [16, 56, 56]
-
-        if any(s.startswith('{}time_r_in.'.format(key_prefix)) for s in state_dict_keys):
-            dit_config["meanflow"] = True
-        else:
-            dit_config["meanflow"] = False
-
-        dit_config["context_in_dim"] = state_dict['{}txt_in.input_embedder.weight'.format(key_prefix)].shape[1]
-        dit_config["hidden_size"] = in_w.shape[0]
+        dit_config["in_channels"] = state_dict['{}img_in.proj.weight'.format(key_prefix)].shape[1] #SkyReels img2video has 32 input channels
+        dit_config["patch_size"] = [1, 2, 2]
+        dit_config["out_channels"] = 16
+        dit_config["vec_in_dim"] = 768
+        dit_config["context_in_dim"] = 4096
+        dit_config["hidden_size"] = 3072
         dit_config["mlp_ratio"] = 4.0
-        dit_config["num_heads"] = in_w.shape[0] // 128
+        dit_config["num_heads"] = 24
         dit_config["depth"] = count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.')
         dit_config["depth_single_blocks"] = count_blocks(state_dict_keys, '{}single_blocks.'.format(key_prefix) + '{}.')
+        dit_config["axes_dim"] = [16, 56, 56]
         dit_config["theta"] = 256
         dit_config["qkv_bias"] = True
-        if '{}byt5_in.fc1.weight'.format(key_prefix) in state_dict:
-            dit_config["byt5"] = True
-        else:
-            dit_config["byt5"] = False
-
         guidance_keys = list(filter(lambda a: a.startswith("{}guidance_in.".format(key_prefix)), state_dict_keys))
         dit_config["guidance_embed"] = len(guidance_keys) > 0
         return dit_config
 
-    if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys and ('{}img_in.weight'.format(key_prefix) in state_dict_keys or f"{key_prefix}distilled_guidance_layer.norms.0.scale" in state_dict_keys): #Flux, Chroma or Chroma Radiance (has no img_in.weight)
+    if '{}double_blocks.0.img_attn.norm.key_norm.scale'.format(key_prefix) in state_dict_keys and '{}img_in.weight'.format(key_prefix) in state_dict_keys: #Flux
         dit_config = {}
         dit_config["image_model"] = "flux"
         dit_config["in_channels"] = 16
@@ -204,18 +184,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
             dit_config["out_dim"] = 3072
             dit_config["hidden_dim"] = 5120
             dit_config["n_layers"] = 5
-            if f"{key_prefix}nerf_blocks.0.norm.scale" in state_dict_keys: #Chroma Radiance
-                dit_config["image_model"] = "chroma_radiance"
-                dit_config["in_channels"] = 3
-                dit_config["out_channels"] = 3
-                dit_config["patch_size"] = 16
-                dit_config["nerf_hidden_size"] = 64
-                dit_config["nerf_mlp_ratio"] = 4
-                dit_config["nerf_depth"] = 4
-                dit_config["nerf_max_freqs"] = 8
-                dit_config["nerf_tile_size"] = 32
-                dit_config["nerf_final_head_type"] = "conv" if f"{key_prefix}nerf_final_layer_conv.norm.scale" in state_dict_keys else "linear"
-                dit_config["nerf_embedder_dtype"] = torch.float32
         else:
             dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
         return dit_config
@@ -402,10 +370,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
                 dit_config["model_type"] = "camera_2.2"
         elif '{}casual_audio_encoder.encoder.final_linear.weight'.format(key_prefix) in state_dict_keys:
             dit_config["model_type"] = "s2v"
-        elif '{}audio_proj.audio_proj_glob_1.layer.bias'.format(key_prefix) in state_dict_keys:
-            dit_config["model_type"] = "humo"
-        elif '{}face_adapter.fuser_blocks.0.k_norm.weight'.format(key_prefix) in state_dict_keys:
-            dit_config["model_type"] = "animate"
         else:
             if '{}img_emb.proj.0.bias'.format(key_prefix) in state_dict_keys:
                 dit_config["model_type"] = "i2v"
@@ -436,20 +400,6 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
         dit_config["guidance_embed"] = "{}guidance_in.in_layer.weight".format(key_prefix) in state_dict_keys
         return dit_config
 
-    if f"{key_prefix}t_embedder.mlp.2.weight" in state_dict_keys:  # Hunyuan 3D 2.1
-
-        dit_config = {}
-        dit_config["image_model"] = "hunyuan3d2_1"
-        dit_config["in_channels"] = state_dict[f"{key_prefix}x_embedder.weight"].shape[1]
-        dit_config["context_dim"] = 1024
-        dit_config["hidden_size"] = state_dict[f"{key_prefix}x_embedder.weight"].shape[0]
-        dit_config["mlp_ratio"] = 4.0
-        dit_config["num_heads"] = 16
-        dit_config["depth"] = count_blocks(state_dict_keys, f"{key_prefix}blocks.{{}}")
-        dit_config["qkv_bias"] = False
-        dit_config["guidance_cond_proj_dim"] = None#f"{key_prefix}t_embedder.cond_proj.weight" in state_dict_keys
-        return dit_config
-
     if '{}caption_projection.0.linear.weight'.format(key_prefix) in state_dict_keys:  # HiDream
         dit_config = {}
         dit_config["image_model"] = "hidream"

comfy/model_management.py

@@ -22,7 +22,6 @@ from enum import Enum
 from comfy.cli_args import args, PerformanceFeature
 import torch
 import sys
-import importlib
 import platform
 import weakref
 import gc
@@ -290,24 +289,6 @@ def is_amd():
             return True
     return False
 
-def amd_min_version(device=None, min_rdna_version=0):
-    if not is_amd():
-        return False
-
-    if is_device_cpu(device):
-        return False
-
-    arch = torch.cuda.get_device_properties(device).gcnArchName
-    if arch.startswith('gfx') and len(arch) == 7:
-        try:
-            cmp_rdna_version = int(arch[4]) + 2
-        except:
-            cmp_rdna_version = 0
-        if cmp_rdna_version >= min_rdna_version:
-            return True
-
-    return False
-
 MIN_WEIGHT_MEMORY_RATIO = 0.4
 if is_nvidia():
     MIN_WEIGHT_MEMORY_RATIO = 0.0
@@ -340,15 +321,14 @@ try:
         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 importlib.util.find_spec('triton') is not None:  # AMD efficient attention implementation depends on triton. TODO: better way of detecting if it's compiled in or not.
-                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", "gfx1100", "gfx1101", "gfx1151"]):  # TODO: more arches, TODO: gfx950
-                        ENABLE_PYTORCH_ATTENTION = True
-#                if torch_version_numeric >= (2, 8):
-#                    if any((a in arch) for a in ["gfx1201"]):
-#                        ENABLE_PYTORCH_ATTENTION = True
+            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", "gfx1100", "gfx1101", "gfx1151"]):  # TODO: more arches, TODO: gfx950
+                    ENABLE_PYTORCH_ATTENTION = True
+#            if torch_version_numeric >= (2, 8):
+#                if any((a in arch) for a in ["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", "gfx942", "gfx950"]):  # TODO: more arches
+            if any((a in arch) for a in ["gfx1201", "gfx942", "gfx950"]):  # TODO: more arches
                 SUPPORT_FP8_OPS = True
 
 except:
@@ -645,9 +625,7 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu
             if loaded_model.model.is_clone(current_loaded_models[i].model):
                 to_unload = [i] + to_unload
         for i in to_unload:
-            model_to_unload = current_loaded_models.pop(i)
-            model_to_unload.model.detach(unpatch_all=False)
-            model_to_unload.model_finalizer.detach()
+            current_loaded_models.pop(i).model.detach(unpatch_all=False)
 
     total_memory_required = {}
     for loaded_model in models_to_load:
@@ -927,9 +905,7 @@ def vae_dtype(device=None, allowed_dtypes=[]):
 
         # NOTE: bfloat16 seems to work on AMD for the VAE but is extremely slow in some cases compared to fp32
         # slowness still a problem on pytorch nightly 2.9.0.dev20250720+rocm6.4 tested on RDNA3
-        # also a problem on RDNA4 except fp32 is also slow there.
-        # This is due to large bf16 convolutions being extremely slow.
-        if d == torch.bfloat16 and ((not is_amd()) or amd_min_version(device, min_rdna_version=4)) and should_use_bf16(device):
+        if d == torch.bfloat16 and (not is_amd()) and should_use_bf16(device):
             return d
 
     return torch.float32

comfy/model_patcher.py

@@ -433,9 +433,6 @@ class ModelPatcher:
     def set_model_double_block_patch(self, patch):
         self.set_model_patch(patch, "double_block")
 
-    def set_model_post_input_patch(self, patch):
-        self.set_model_patch(patch, "post_input")
-
     def add_object_patch(self, name, obj):
         self.object_patches[name] = obj
 

comfy/ops.py

@@ -52,9 +52,6 @@ except (ModuleNotFoundError, TypeError):
 
 cast_to = comfy.model_management.cast_to #TODO: remove once no more references
 
-if torch.cuda.is_available() and torch.backends.cudnn.is_available() and PerformanceFeature.AutoTune in args.fast:
-    torch.backends.cudnn.benchmark = True
-
 def cast_to_input(weight, input, non_blocking=False, copy=True):
     return comfy.model_management.cast_to(weight, input.dtype, input.device, non_blocking=non_blocking, copy=copy)
 
@@ -365,13 +362,12 @@ class fp8_ops(manual_cast):
             return None
 
         def forward_comfy_cast_weights(self, input):
-            if not self.training:
-                try:
-                    out = fp8_linear(self, input)
-                    if out is not None:
-                        return out
-                except Exception as e:
-                    logging.info("Exception during fp8 op: {}".format(e))
+            try:
+                out = fp8_linear(self, input)
+                if out is not None:
+                    return out
+            except Exception as e:
+                logging.info("Exception during fp8 op: {}".format(e))
 
             weight, bias = cast_bias_weight(self, input)
             return torch.nn.functional.linear(input, weight, bias)

comfy/pixel_space_convert.py

@@ -1,16 +0,0 @@
-import torch
-
-
-# "Fake" VAE that converts from IMAGE B, H, W, C and values on the scale of 0..1
-# to LATENT B, C, H, W and values on the scale of -1..1.
-class PixelspaceConversionVAE(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.pixel_space_vae = torch.nn.Parameter(torch.tensor(1.0))
-
-    def encode(self, pixels: torch.Tensor, *_args, **_kwargs) -> torch.Tensor:
-        return pixels
-
-    def decode(self, samples: torch.Tensor, *_args, **_kwargs) -> torch.Tensor:
-        return samples
-

comfy/samplers.py

@@ -360,7 +360,7 @@ def calc_cond_uncond_batch(model, cond, uncond, x_in, timestep, model_options):
 def cfg_function(model, cond_pred, uncond_pred, cond_scale, x, timestep, model_options={}, cond=None, uncond=None):
     if "sampler_cfg_function" in model_options:
         args = {"cond": x - cond_pred, "uncond": x - uncond_pred, "cond_scale": cond_scale, "timestep": timestep, "input": x, "sigma": timestep,
-                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options, "input_cond": cond, "input_uncond": uncond}
+                "cond_denoised": cond_pred, "uncond_denoised": uncond_pred, "model": model, "model_options": model_options}
         cfg_result = x - model_options["sampler_cfg_function"](args)
     else:
         cfg_result = uncond_pred + (cond_pred - uncond_pred) * cond_scale
@@ -390,7 +390,7 @@ def sampling_function(model, x, timestep, uncond, cond, cond_scale, model_option
     for fn in model_options.get("sampler_pre_cfg_function", []):
         args = {"conds":conds, "conds_out": out, "cond_scale": cond_scale, "timestep": timestep,
                 "input": x, "sigma": timestep, "model": model, "model_options": model_options}
-        out = fn(args)
+        out  = fn(args)
 
     return cfg_function(model, out[0], out[1], cond_scale, x, timestep, model_options=model_options, cond=cond, uncond=uncond_)
 

comfy/sd.py

@@ -17,8 +17,6 @@ import comfy.ldm.wan.vae
 import comfy.ldm.wan.vae2_2
 import comfy.ldm.hunyuan3d.vae
 import comfy.ldm.ace.vae.music_dcae_pipeline
-import comfy.ldm.hunyuan_video.vae
-import comfy.pixel_space_convert
 import yaml
 import math
 import os
@@ -50,7 +48,6 @@ import comfy.text_encoders.hidream
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
-import comfy.text_encoders.hunyuan_image
 
 import comfy.model_patcher
 import comfy.lora
@@ -286,7 +283,6 @@ class VAE:
         self.process_output = lambda image: torch.clamp((image + 1.0) / 2.0, min=0.0, max=1.0)
         self.working_dtypes = [torch.bfloat16, torch.float32]
         self.disable_offload = False
-        self.not_video = False
 
         self.downscale_index_formula = None
         self.upscale_index_formula = None
@@ -332,19 +328,6 @@ class VAE:
                 self.first_stage_model = StageC_coder()
                 self.downscale_ratio = 32
                 self.latent_channels = 16
-            elif "decoder.conv_in.weight" in sd and 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}
-                self.latent_channels = ddconfig['z_channels'] = sd["decoder.conv_in.weight"].shape[1]
-                self.downscale_ratio = 32
-                self.upscale_ratio = 32
-                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-                self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.DiagonalGaussianRegularizer"},
-                                                            encoder_config={'target': "comfy.ldm.hunyuan_video.vae.Encoder", 'params': ddconfig},
-                                                            decoder_config={'target': "comfy.ldm.hunyuan_video.vae.Decoder", 'params': ddconfig})
-
-                self.memory_used_encode = lambda shape, dtype: (700 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
-                self.memory_used_decode = lambda shape, dtype: (700 * shape[2] * shape[3] * 32 * 32) * model_management.dtype_size(dtype)
-
             elif "decoder.conv_in.weight" in sd:
                 #default SD1.x/SD2.x VAE parameters
                 ddconfig = {'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}
@@ -411,23 +394,6 @@ class VAE:
                 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.conv.weight" in sd and sd['decoder.conv_in.conv.weight'].shape[1] == 32:
-                ddconfig = {"block_out_channels": [128, 256, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 16, "ffactor_temporal": 4, "downsample_match_channel": True, "upsample_match_channel": True}
-                ddconfig['z_channels'] = sd["decoder.conv_in.conv.weight"].shape[1]
-                self.latent_channels = 64
-                self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 16, 16)
-                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)
-                self.latent_dim = 3
-                self.not_video = True
-                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-                self.first_stage_model = AutoencodingEngine(regularizer_config={'target': "comfy.ldm.models.autoencoder.EmptyRegularizer"},
-                                                            encoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Encoder", 'params': ddconfig},
-                                                            decoder_config={'target': "comfy.ldm.hunyuan_video.vae_refiner.Decoder", 'params': ddconfig})
-
-                self.memory_used_encode = lambda shape, dtype: (1400 * shape[-2] * shape[-1]) * model_management.dtype_size(dtype)
-                self.memory_used_decode = lambda shape, dtype: (1400 * shape[-3] * shape[-2] * shape[-1] * 16 * 16) * model_management.dtype_size(dtype)
             elif "decoder.conv_in.conv.weight" in sd:
                 ddconfig = {'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}
                 ddconfig["conv3d"] = True
@@ -480,29 +446,17 @@ class VAE:
                     self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
                     self.memory_used_encode = lambda shape, dtype: 6000 * shape[3] * shape[4] * model_management.dtype_size(dtype)
                     self.memory_used_decode = lambda shape, dtype: 7000 * shape[3] * shape[4] * (8 * 8) * model_management.dtype_size(dtype)
-            # Hunyuan 3d v2 2.0 & 2.1
             elif "geo_decoder.cross_attn_decoder.ln_1.bias" in sd:
-
                 self.latent_dim = 1
-
-                def estimate_memory(shape, dtype, num_layers = 16, kv_cache_multiplier = 2):
-                    batch, num_tokens, hidden_dim = shape
-                    dtype_size = model_management.dtype_size(dtype)
-
-                    total_mem = batch * num_tokens * hidden_dim * dtype_size * (1 + kv_cache_multiplier * num_layers)
-                    return total_mem
-
-                # better memory estimations
-                self.memory_used_encode = lambda shape, dtype, num_layers = 8, kv_cache_multiplier = 0:\
-                    estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
-
-                self.memory_used_decode = lambda shape, dtype, num_layers = 16, kv_cache_multiplier = 2: \
-                    estimate_memory(shape, dtype, num_layers, kv_cache_multiplier)
-
-                self.first_stage_model = comfy.ldm.hunyuan3d.vae.ShapeVAE()
+                ln_post = "geo_decoder.ln_post.weight" in sd
+                inner_size = sd["geo_decoder.output_proj.weight"].shape[1]
+                downsample_ratio = sd["post_kl.weight"].shape[0] // inner_size
+                mlp_expand = sd["geo_decoder.cross_attn_decoder.mlp.c_fc.weight"].shape[0] // inner_size
+                self.memory_used_encode = lambda shape, dtype: (1000 * shape[2]) * model_management.dtype_size(dtype)  # TODO
+                self.memory_used_decode = lambda shape, dtype: (1024 * 1024 * 1024 * 2.0) * model_management.dtype_size(dtype)  # TODO
+                ddconfig = {"embed_dim": 64, "num_freqs": 8, "include_pi": False, "heads": 16, "width": 1024, "num_decoder_layers": 16, "qkv_bias": False, "qk_norm": True, "geo_decoder_mlp_expand_ratio": mlp_expand, "geo_decoder_downsample_ratio": downsample_ratio, "geo_decoder_ln_post": ln_post}
+                self.first_stage_model = comfy.ldm.hunyuan3d.vae.ShapeVAE(**ddconfig)
                 self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
-
-
             elif "vocoder.backbone.channel_layers.0.0.bias" in sd: #Ace Step Audio
                 self.first_stage_model = comfy.ldm.ace.vae.music_dcae_pipeline.MusicDCAE(source_sample_rate=44100)
                 self.memory_used_encode = lambda shape, dtype: (shape[2] * 330) * model_management.dtype_size(dtype)
@@ -517,15 +471,6 @@ class VAE:
                 self.working_dtypes = [torch.bfloat16, torch.float16, torch.float32]
                 self.disable_offload = True
                 self.extra_1d_channel = 16
-            elif "pixel_space_vae" in sd:
-                self.first_stage_model = comfy.pixel_space_convert.PixelspaceConversionVAE()
-                self.memory_used_encode = lambda shape, dtype: (1 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
-                self.memory_used_decode = lambda shape, dtype: (1 * shape[2] * shape[3]) * model_management.dtype_size(dtype)
-                self.downscale_ratio = 1
-                self.upscale_ratio = 1
-                self.latent_channels = 3
-                self.latent_dim = 2
-                self.output_channels = 3
             else:
                 logging.warning("WARNING: No VAE weights detected, VAE not initalized.")
                 self.first_stage_model = None
@@ -698,10 +643,7 @@ class VAE:
         pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
         pixel_samples = pixel_samples.movedim(-1, 1)
         if self.latent_dim == 3 and pixel_samples.ndim < 5:
-            if not self.not_video:
-                pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
-            else:
-                pixel_samples = pixel_samples.unsqueeze(2)
+            pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
         try:
             memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)
             model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -735,10 +677,7 @@ class VAE:
         dims = self.latent_dim
         pixel_samples = pixel_samples.movedim(-1, 1)
         if dims == 3:
-            if not self.not_video:
-                pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
-            else:
-                pixel_samples = pixel_samples.unsqueeze(2)
+            pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
 
         memory_used = self.memory_used_encode(pixel_samples.shape, self.vae_dtype)  # TODO: calculate mem required for tile
         model_management.load_models_gpu([self.patcher], memory_required=memory_used, force_full_load=self.disable_offload)
@@ -795,7 +734,6 @@ class VAE:
         except:
             return None
 
-
 class StyleModel:
     def __init__(self, model, device="cpu"):
         self.model = model
@@ -835,7 +773,6 @@ class CLIPType(Enum):
     ACE = 16
     OMNIGEN2 = 17
     QWEN_IMAGE = 18
-    HUNYUAN_IMAGE = 19
 
 
 def load_clip(ckpt_paths, embedding_directory=None, clip_type=CLIPType.STABLE_DIFFUSION, model_options={}):
@@ -857,7 +794,6 @@ class TEModel(Enum):
     GEMMA_2_2B = 9
     QWEN25_3B = 10
     QWEN25_7B = 11
-    BYT5_SMALL_GLYPH = 12
 
 def detect_te_model(sd):
     if "text_model.encoder.layers.30.mlp.fc1.weight" in sd:
@@ -875,9 +811,6 @@ def detect_te_model(sd):
     if 'encoder.block.23.layer.1.DenseReluDense.wi.weight' in sd:
         return TEModel.T5_XXL_OLD
     if "encoder.block.0.layer.0.SelfAttention.k.weight" in sd:
-        weight = sd['encoder.block.0.layer.0.SelfAttention.k.weight']
-        if weight.shape[0] == 384:
-            return TEModel.BYT5_SMALL_GLYPH
         return TEModel.T5_BASE
     if 'model.layers.0.post_feedforward_layernorm.weight' in sd:
         return TEModel.GEMMA_2_2B
@@ -992,12 +925,8 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
             clip_target.clip = comfy.text_encoders.omnigen2.te(**llama_detect(clip_data))
             clip_target.tokenizer = comfy.text_encoders.omnigen2.Omnigen2Tokenizer
         elif te_model == TEModel.QWEN25_7B:
-            if clip_type == CLIPType.HUNYUAN_IMAGE:
-                clip_target.clip = comfy.text_encoders.hunyuan_image.te(byt5=False, **llama_detect(clip_data))
-                clip_target.tokenizer = comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer
-            else:
-                clip_target.clip = comfy.text_encoders.qwen_image.te(**llama_detect(clip_data))
-                clip_target.tokenizer = comfy.text_encoders.qwen_image.QwenImageTokenizer
+            clip_target.clip = comfy.text_encoders.qwen_image.te(**llama_detect(clip_data))
+            clip_target.tokenizer = comfy.text_encoders.qwen_image.QwenImageTokenizer
         else:
             # clip_l
             if clip_type == CLIPType.SD3:
@@ -1041,9 +970,6 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
 
             clip_target.clip = comfy.text_encoders.hidream.hidream_clip(clip_l=clip_l, clip_g=clip_g, t5=t5, llama=llama, **t5_kwargs, **llama_kwargs)
             clip_target.tokenizer = comfy.text_encoders.hidream.HiDreamTokenizer
-        elif clip_type == CLIPType.HUNYUAN_IMAGE:
-            clip_target.clip = comfy.text_encoders.hunyuan_image.te(**llama_detect(clip_data))
-            clip_target.tokenizer = comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer
         else:
             clip_target.clip = sdxl_clip.SDXLClipModel
             clip_target.tokenizer = sdxl_clip.SDXLTokenizer

comfy/supported_models.py

@@ -20,7 +20,6 @@ import comfy.text_encoders.wan
 import comfy.text_encoders.ace
 import comfy.text_encoders.omnigen2
 import comfy.text_encoders.qwen_image
-import comfy.text_encoders.hunyuan_image
 
 from . import supported_models_base
 from . import latent_formats
@@ -995,7 +994,7 @@ class WAN21_T2V(supported_models_base.BASE):
     unet_extra_config = {}
     latent_format = latent_formats.Wan21
 
-    memory_usage_factor = 0.9
+    memory_usage_factor = 1.0
 
     supported_inference_dtypes = [torch.float16, torch.bfloat16, torch.float32]
 
@@ -1004,7 +1003,7 @@ class WAN21_T2V(supported_models_base.BASE):
 
     def __init__(self, unet_config):
         super().__init__(unet_config)
-        self.memory_usage_factor = self.unet_config.get("dim", 2000) / 2222
+        self.memory_usage_factor = self.unet_config.get("dim", 2000) / 2000
 
     def get_model(self, state_dict, prefix="", device=None):
         out = model_base.WAN21(self, device=device)
@@ -1073,16 +1072,6 @@ class WAN21_Vace(WAN21_T2V):
         out = model_base.WAN21_Vace(self, image_to_video=False, device=device)
         return out
 
-class WAN21_HuMo(WAN21_T2V):
-    unet_config = {
-        "image_model": "wan2.1",
-        "model_type": "humo",
-    }
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.WAN21_HuMo(self, image_to_video=False, device=device)
-        return out
-
 class WAN22_S2V(WAN21_T2V):
     unet_config = {
         "image_model": "wan2.1",
@@ -1096,19 +1085,6 @@ class WAN22_S2V(WAN21_T2V):
         out = model_base.WAN22_S2V(self, device=device)
         return out
 
-class WAN22_Animate(WAN21_T2V):
-    unet_config = {
-        "image_model": "wan2.1",
-        "model_type": "animate",
-    }
-
-    def __init__(self, unet_config):
-        super().__init__(unet_config)
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.WAN22_Animate(self, device=device)
-        return out
-
 class WAN22_T2V(WAN21_T2V):
     unet_config = {
         "image_model": "wan2.1",
@@ -1152,17 +1128,6 @@ class Hunyuan3Dv2(supported_models_base.BASE):
     def clip_target(self, state_dict={}):
         return None
 
-class Hunyuan3Dv2_1(Hunyuan3Dv2):
-    unet_config = {
-        "image_model": "hunyuan3d2_1",
-    }
-
-    latent_format = latent_formats.Hunyuan3Dv2_1
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.Hunyuan3Dv2_1(self, device = device)
-        return out
-
 class Hunyuan3Dv2mini(Hunyuan3Dv2):
     unet_config = {
         "image_model": "hunyuan3d2",
@@ -1228,19 +1193,6 @@ class Chroma(supported_models_base.BASE):
         t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.pixart_te(**t5_detect))
 
-class ChromaRadiance(Chroma):
-    unet_config = {
-        "image_model": "chroma_radiance",
-    }
-
-    latent_format = comfy.latent_formats.ChromaRadiance
-
-    # Pixel-space model, no spatial compression for model input.
-    memory_usage_factor = 0.038
-
-    def get_model(self, state_dict, prefix="", device=None):
-        return model_base.ChromaRadiance(self, device=device)
-
 class ACEStep(supported_models_base.BASE):
     unet_config = {
         "audio_model": "ace",
@@ -1332,48 +1284,7 @@ class QwenImage(supported_models_base.BASE):
         hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
 
-class HunyuanImage21(HunyuanVideo):
-    unet_config = {
-        "image_model": "hunyuan_video",
-        "vec_in_dim": None,
-    }
 
-    sampling_settings = {
-        "shift": 5.0,
-    }
-
-    latent_format = latent_formats.HunyuanImage21
-
-    memory_usage_factor = 7.7
-
-    supported_inference_dtypes = [torch.bfloat16, torch.float32]
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.HunyuanImage21(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, "{}qwen25_7b.transformer.".format(pref))
-        return supported_models_base.ClipTarget(comfy.text_encoders.hunyuan_image.HunyuanImageTokenizer, comfy.text_encoders.hunyuan_image.te(**hunyuan_detect))
-
-class HunyuanImage21Refiner(HunyuanVideo):
-    unet_config = {
-        "image_model": "hunyuan_video",
-        "patch_size": [1, 1, 1],
-        "vec_in_dim": None,
-    }
-
-    sampling_settings = {
-        "shift": 4.0,
-    }
-
-    latent_format = latent_formats.HunyuanImage21Refiner
-
-    def get_model(self, state_dict, prefix="", device=None):
-        out = model_base.HunyuanImage21Refiner(self, device=device)
-        return out
-
-models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanImage21Refiner, HunyuanImage21, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, WAN21_HuMo, WAN22_Animate, Hunyuan3Dv2mini, Hunyuan3Dv2, Hunyuan3Dv2_1, HiDream, Chroma, ChromaRadiance, ACEStep, Omnigen2, QwenImage]
+models = [LotusD, Stable_Zero123, SD15_instructpix2pix, SD15, SD20, SD21UnclipL, SD21UnclipH, SDXL_instructpix2pix, SDXLRefiner, SDXL, SSD1B, KOALA_700M, KOALA_1B, Segmind_Vega, SD_X4Upscaler, Stable_Cascade_C, Stable_Cascade_B, SV3D_u, SV3D_p, SD3, StableAudio, AuraFlow, PixArtAlpha, PixArtSigma, HunyuanDiT, HunyuanDiT1, FluxInpaint, Flux, FluxSchnell, GenmoMochi, LTXV, HunyuanVideoSkyreelsI2V, HunyuanVideoI2V, HunyuanVideo, CosmosT2V, CosmosI2V, CosmosT2IPredict2, CosmosI2VPredict2, Lumina2, WAN22_T2V, WAN21_T2V, WAN21_I2V, WAN21_FunControl2V, WAN21_Vace, WAN21_Camera, WAN22_Camera, WAN22_S2V, Hunyuan3Dv2mini, Hunyuan3Dv2, HiDream, Chroma, ACEStep, Omnigen2, QwenImage]
 
 models += [SVD_img2vid]

comfy/text_encoders/byt5_config_small_glyph.json

@@ -1,22 +0,0 @@
-{
-  "d_ff": 3584,
-  "d_kv": 64,
-  "d_model": 1472,
-  "decoder_start_token_id": 0,
-  "dropout_rate": 0.1,
-  "eos_token_id": 1,
-  "dense_act_fn": "gelu_pytorch_tanh",
-  "initializer_factor": 1.0,
-  "is_encoder_decoder": true,
-  "is_gated_act": true,
-  "layer_norm_epsilon": 1e-06,
-  "model_type": "t5",
-  "num_decoder_layers": 4,
-  "num_heads": 6,
-  "num_layers": 12,
-  "output_past": true,
-  "pad_token_id": 0,
-  "relative_attention_num_buckets": 32,
-  "tie_word_embeddings": false,
-  "vocab_size": 1510
-}

comfy/text_encoders/byt5_tokenizer/added_tokens.json

@@ -1,127 +0,0 @@
-{
-  "<extra_id_0>": 259,
-  "<extra_id_100>": 359,
-  "<extra_id_101>": 360,
-  "<extra_id_102>": 361,
-  "<extra_id_103>": 362,
-  "<extra_id_104>": 363,
-  "<extra_id_105>": 364,
-  "<extra_id_106>": 365,
-  "<extra_id_107>": 366,
-  "<extra_id_108>": 367,
-  "<extra_id_109>": 368,
-  "<extra_id_10>": 269,
-  "<extra_id_110>": 369,
-  "<extra_id_111>": 370,
-  "<extra_id_112>": 371,
-  "<extra_id_113>": 372,
-  "<extra_id_114>": 373,
-  "<extra_id_115>": 374,
-  "<extra_id_116>": 375,
-  "<extra_id_117>": 376,
-  "<extra_id_118>": 377,
-  "<extra_id_119>": 378,
-  "<extra_id_11>": 270,
-  "<extra_id_120>": 379,
-  "<extra_id_121>": 380,
-  "<extra_id_122>": 381,
-  "<extra_id_123>": 382,
-  "<extra_id_124>": 383,
-  "<extra_id_12>": 271,
-  "<extra_id_13>": 272,
-  "<extra_id_14>": 273,
-  "<extra_id_15>": 274,
-  "<extra_id_16>": 275,
-  "<extra_id_17>": 276,
-  "<extra_id_18>": 277,
-  "<extra_id_19>": 278,
-  "<extra_id_1>": 260,
-  "<extra_id_20>": 279,
-  "<extra_id_21>": 280,
-  "<extra_id_22>": 281,
-  "<extra_id_23>": 282,
-  "<extra_id_24>": 283,
-  "<extra_id_25>": 284,
-  "<extra_id_26>": 285,
-  "<extra_id_27>": 286,
-  "<extra_id_28>": 287,
-  "<extra_id_29>": 288,
-  "<extra_id_2>": 261,
-  "<extra_id_30>": 289,
-  "<extra_id_31>": 290,
-  "<extra_id_32>": 291,
-  "<extra_id_33>": 292,
-  "<extra_id_34>": 293,
-  "<extra_id_35>": 294,
-  "<extra_id_36>": 295,
-  "<extra_id_37>": 296,
-  "<extra_id_38>": 297,
-  "<extra_id_39>": 298,
-  "<extra_id_3>": 262,
-  "<extra_id_40>": 299,
-  "<extra_id_41>": 300,
-  "<extra_id_42>": 301,
-  "<extra_id_43>": 302,
-  "<extra_id_44>": 303,
-  "<extra_id_45>": 304,
-  "<extra_id_46>": 305,
-  "<extra_id_47>": 306,
-  "<extra_id_48>": 307,
-  "<extra_id_49>": 308,
-  "<extra_id_4>": 263,
-  "<extra_id_50>": 309,
-  "<extra_id_51>": 310,
-  "<extra_id_52>": 311,
-  "<extra_id_53>": 312,
-  "<extra_id_54>": 313,
-  "<extra_id_55>": 314,
-  "<extra_id_56>": 315,
-  "<extra_id_57>": 316,
-  "<extra_id_58>": 317,
-  "<extra_id_59>": 318,
-  "<extra_id_5>": 264,
-  "<extra_id_60>": 319,
-  "<extra_id_61>": 320,
-  "<extra_id_62>": 321,
-  "<extra_id_63>": 322,
-  "<extra_id_64>": 323,
-  "<extra_id_65>": 324,
-  "<extra_id_66>": 325,
-  "<extra_id_67>": 326,
-  "<extra_id_68>": 327,
-  "<extra_id_69>": 328,
-  "<extra_id_6>": 265,
-  "<extra_id_70>": 329,
-  "<extra_id_71>": 330,
-  "<extra_id_72>": 331,
-  "<extra_id_73>": 332,
-  "<extra_id_74>": 333,
-  "<extra_id_75>": 334,
-  "<extra_id_76>": 335,
-  "<extra_id_77>": 336,
-  "<extra_id_78>": 337,
-  "<extra_id_79>": 338,
-  "<extra_id_7>": 266,
-  "<extra_id_80>": 339,
-  "<extra_id_81>": 340,
-  "<extra_id_82>": 341,
-  "<extra_id_83>": 342,
-  "<extra_id_84>": 343,
-  "<extra_id_85>": 344,
-  "<extra_id_86>": 345,
-  "<extra_id_87>": 346,
-  "<extra_id_88>": 347,
-  "<extra_id_89>": 348,
-  "<extra_id_8>": 267,
-  "<extra_id_90>": 349,
-  "<extra_id_91>": 350,
-  "<extra_id_92>": 351,
-  "<extra_id_93>": 352,
-  "<extra_id_94>": 353,
-  "<extra_id_95>": 354,
-  "<extra_id_96>": 355,
-  "<extra_id_97>": 356,
-  "<extra_id_98>": 357,
-  "<extra_id_99>": 358,
-  "<extra_id_9>": 268
-}

comfy/text_encoders/byt5_tokenizer/special_tokens_map.json

@@ -1,150 +0,0 @@
-{
-  "additional_special_tokens": [
-    "<extra_id_0>",
-    "<extra_id_1>",
-    "<extra_id_2>",
-    "<extra_id_3>",
-    "<extra_id_4>",
-    "<extra_id_5>",
-    "<extra_id_6>",
-    "<extra_id_7>",
-    "<extra_id_8>",
-    "<extra_id_9>",
-    "<extra_id_10>",
-    "<extra_id_11>",
-    "<extra_id_12>",
-    "<extra_id_13>",
-    "<extra_id_14>",
-    "<extra_id_15>",
-    "<extra_id_16>",
-    "<extra_id_17>",
-    "<extra_id_18>",
-    "<extra_id_19>",
-    "<extra_id_20>",
-    "<extra_id_21>",
-    "<extra_id_22>",
-    "<extra_id_23>",
-    "<extra_id_24>",
-    "<extra_id_25>",
-    "<extra_id_26>",
-    "<extra_id_27>",
-    "<extra_id_28>",
-    "<extra_id_29>",
-    "<extra_id_30>",
-    "<extra_id_31>",
-    "<extra_id_32>",
-    "<extra_id_33>",
-    "<extra_id_34>",
-    "<extra_id_35>",
-    "<extra_id_36>",
-    "<extra_id_37>",
-    "<extra_id_38>",
-    "<extra_id_39>",
-    "<extra_id_40>",
-    "<extra_id_41>",
-    "<extra_id_42>",
-    "<extra_id_43>",
-    "<extra_id_44>",
-    "<extra_id_45>",
-    "<extra_id_46>",
-    "<extra_id_47>",
-    "<extra_id_48>",
-    "<extra_id_49>",
-    "<extra_id_50>",
-    "<extra_id_51>",
-    "<extra_id_52>",
-    "<extra_id_53>",
-    "<extra_id_54>",
-    "<extra_id_55>",
-    "<extra_id_56>",
-    "<extra_id_57>",
-    "<extra_id_58>",
-    "<extra_id_59>",
-    "<extra_id_60>",
-    "<extra_id_61>",
-    "<extra_id_62>",
-    "<extra_id_63>",
-    "<extra_id_64>",
-    "<extra_id_65>",
-    "<extra_id_66>",
-    "<extra_id_67>",
-    "<extra_id_68>",
-    "<extra_id_69>",
-    "<extra_id_70>",
-    "<extra_id_71>",
-    "<extra_id_72>",
-    "<extra_id_73>",
-    "<extra_id_74>",
-    "<extra_id_75>",
-    "<extra_id_76>",
-    "<extra_id_77>",
-    "<extra_id_78>",
-    "<extra_id_79>",
-    "<extra_id_80>",
-    "<extra_id_81>",
-    "<extra_id_82>",
-    "<extra_id_83>",
-    "<extra_id_84>",
-    "<extra_id_85>",
-    "<extra_id_86>",
-    "<extra_id_87>",
-    "<extra_id_88>",
-    "<extra_id_89>",
-    "<extra_id_90>",
-    "<extra_id_91>",
-    "<extra_id_92>",
-    "<extra_id_93>",
-    "<extra_id_94>",
-    "<extra_id_95>",
-    "<extra_id_96>",
-    "<extra_id_97>",
-    "<extra_id_98>",
-    "<extra_id_99>",
-    "<extra_id_100>",
-    "<extra_id_101>",
-    "<extra_id_102>",
-    "<extra_id_103>",
-    "<extra_id_104>",
-    "<extra_id_105>",
-    "<extra_id_106>",
-    "<extra_id_107>",
-    "<extra_id_108>",
-    "<extra_id_109>",
-    "<extra_id_110>",
-    "<extra_id_111>",
-    "<extra_id_112>",
-    "<extra_id_113>",
-    "<extra_id_114>",
-    "<extra_id_115>",
-    "<extra_id_116>",
-    "<extra_id_117>",
-    "<extra_id_118>",
-    "<extra_id_119>",
-    "<extra_id_120>",
-    "<extra_id_121>",
-    "<extra_id_122>",
-    "<extra_id_123>",
-    "<extra_id_124>"
-  ],
-  "eos_token": {
-    "content": "</s>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": {
-    "content": "<pad>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  },
-  "unk_token": {
-    "content": "<unk>",
-    "lstrip": false,
-    "normalized": true,
-    "rstrip": false,
-    "single_word": false
-  }
-}

comfy/text_encoders/hunyuan_image.py

@@ -1,103 +0,0 @@
-from comfy import sd1_clip
-import comfy.text_encoders.llama
-from .qwen_image import QwenImageTokenizer, QwenImageTEModel
-from transformers import ByT5Tokenizer
-import os
-import re
-
-class ByT5SmallTokenizer(sd1_clip.SDTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        tokenizer_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), "byt5_tokenizer")
-        super().__init__(tokenizer_path, pad_with_end=False, embedding_size=1472, embedding_key='byt5_small', tokenizer_class=ByT5Tokenizer, has_start_token=False, pad_to_max_length=False, max_length=99999999, min_length=1, tokenizer_data=tokenizer_data)
-
-class HunyuanImageTokenizer(QwenImageTokenizer):
-    def __init__(self, embedding_directory=None, tokenizer_data={}):
-        super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-        self.llama_template = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>"
-        # self.llama_template_images = "{}"
-        self.byt5 = ByT5SmallTokenizer(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data)
-
-    def tokenize_with_weights(self, text:str, return_word_ids=False, **kwargs):
-        out = super().tokenize_with_weights(text, return_word_ids, **kwargs)
-
-        # ByT5 processing for HunyuanImage
-        text_prompt_texts = []
-        pattern_quote_double = r'\"(.*?)\"'
-        pattern_quote_chinese_single = r'‘(.*?)’'
-        pattern_quote_chinese_double = r'“(.*?)”'
-
-        matches_quote_double = re.findall(pattern_quote_double, text)
-        matches_quote_chinese_single = re.findall(pattern_quote_chinese_single, text)
-        matches_quote_chinese_double = re.findall(pattern_quote_chinese_double, text)
-
-        text_prompt_texts.extend(matches_quote_double)
-        text_prompt_texts.extend(matches_quote_chinese_single)
-        text_prompt_texts.extend(matches_quote_chinese_double)
-
-        if len(text_prompt_texts) > 0:
-            out['byt5'] = self.byt5.tokenize_with_weights(''.join(map(lambda a: 'Text "{}". '.format(a), text_prompt_texts)), return_word_ids, **kwargs)
-        return out
-
-class Qwen25_7BVLIModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="hidden", layer_idx=-3, dtype=None, attention_mask=True, model_options={}):
-        llama_scaled_fp8 = model_options.get("qwen_scaled_fp8", None)
-        if llama_scaled_fp8 is not None:
-            model_options = model_options.copy()
-            model_options["scaled_fp8"] = llama_scaled_fp8
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={}, dtype=dtype, special_tokens={"pad": 151643}, layer_norm_hidden_state=False, model_class=comfy.text_encoders.llama.Qwen25_7BVLI, enable_attention_masks=attention_mask, return_attention_masks=attention_mask, model_options=model_options)
-
-
-class ByT5SmallModel(sd1_clip.SDClipModel):
-    def __init__(self, device="cpu", layer="last", layer_idx=None, dtype=None, model_options={}):
-        textmodel_json_config = os.path.join(os.path.dirname(os.path.realpath(__file__)), "byt5_config_small_glyph.json")
-        super().__init__(device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config=textmodel_json_config, dtype=dtype, model_options=model_options, special_tokens={"end": 1, "pad": 0}, model_class=comfy.text_encoders.t5.T5, enable_attention_masks=True, zero_out_masked=True)
-
-
-class HunyuanImageTEModel(QwenImageTEModel):
-    def __init__(self, byt5=True, device="cpu", dtype=None, model_options={}):
-        super(QwenImageTEModel, self).__init__(device=device, dtype=dtype, name="qwen25_7b", clip_model=Qwen25_7BVLIModel, model_options=model_options)
-
-        if byt5:
-            self.byt5_small = ByT5SmallModel(device=device, dtype=dtype, model_options=model_options)
-        else:
-            self.byt5_small = None
-
-    def encode_token_weights(self, token_weight_pairs):
-        tok_pairs = token_weight_pairs["qwen25_7b"][0]
-        template_end = -1
-        if tok_pairs[0][0] == 27:
-            if len(tok_pairs) > 36:  # refiner prompt uses a fixed 36 template_end
-                template_end = 36
-
-        cond, p, extra = super().encode_token_weights(token_weight_pairs, template_end=template_end)
-        if self.byt5_small is not None and "byt5" in token_weight_pairs:
-            out = self.byt5_small.encode_token_weights(token_weight_pairs["byt5"])
-            extra["conditioning_byt5small"] = out[0]
-        return cond, p, extra
-
-    def set_clip_options(self, options):
-        super().set_clip_options(options)
-        if self.byt5_small is not None:
-            self.byt5_small.set_clip_options(options)
-
-    def reset_clip_options(self):
-        super().reset_clip_options()
-        if self.byt5_small is not None:
-            self.byt5_small.reset_clip_options()
-
-    def load_sd(self, sd):
-        if "encoder.block.0.layer.0.SelfAttention.o.weight" in sd:
-            return self.byt5_small.load_sd(sd)
-        else:
-            return super().load_sd(sd)
-
-def te(byt5=True, dtype_llama=None, llama_scaled_fp8=None):
-    class QwenImageTEModel_(HunyuanImageTEModel):
-        def __init__(self, device="cpu", dtype=None, model_options={}):
-            if llama_scaled_fp8 is not None and "scaled_fp8" not in model_options:
-                model_options = model_options.copy()
-                model_options["qwen_scaled_fp8"] = llama_scaled_fp8
-            if dtype_llama is not None:
-                dtype = dtype_llama
-            super().__init__(byt5=byt5, device=device, dtype=dtype, model_options=model_options)
-    return QwenImageTEModel_

comfy/text_encoders/llama.py

@@ -128,12 +128,11 @@ def precompute_freqs_cis(head_dim, position_ids, theta, rope_dims=None, device=N
 
 
 def apply_rope(xq, xk, freqs_cis):
-    org_dtype = xq.dtype
     cos = freqs_cis[0]
     sin = freqs_cis[1]
     q_embed = (xq * cos) + (rotate_half(xq) * sin)
     k_embed = (xk * cos) + (rotate_half(xk) * sin)
-    return q_embed.to(org_dtype), k_embed.to(org_dtype)
+    return q_embed, k_embed
 
 
 class Attention(nn.Module):
@@ -400,25 +399,21 @@ class Qwen25_7BVLI(BaseLlama, torch.nn.Module):
 
     def forward(self, x, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[]):
         grid = None
-        position_ids = None
-        offset = 0
         for e in embeds_info:
             if e.get("type") == "image":
                 grid = e.get("extra", None)
+                position_ids = torch.zeros((3, embeds.shape[1]), device=embeds.device)
                 start = e.get("index")
-                if position_ids is None:
-                    position_ids = torch.zeros((3, embeds.shape[1]), device=embeds.device)
-                    position_ids[:, :start] = torch.arange(0, start, device=embeds.device)
+                position_ids[:, :start] = torch.arange(0, start, device=embeds.device)
                 end = e.get("size") + start
                 len_max = int(grid.max()) // 2
                 start_next = len_max + start
-                position_ids[:, end:] = torch.arange(start_next + offset, start_next + (embeds.shape[1] - end) + offset, device=embeds.device)
-                position_ids[0, start:end] = start + offset
+                position_ids[:, end:] = torch.arange(start_next, start_next + (embeds.shape[1] - end), device=embeds.device)
+                position_ids[0, start:end] = start
                 max_d = int(grid[0][1]) // 2
-                position_ids[1, start:end] = torch.arange(start + offset, start + max_d + offset, device=embeds.device).unsqueeze(1).repeat(1, math.ceil((end - start) / max_d)).flatten(0)[:end - start]
+                position_ids[1, start:end] = torch.arange(start, start + max_d, device=embeds.device).unsqueeze(1).repeat(1, math.ceil((end - start) / max_d)).flatten(0)[:end - start]
                 max_d = int(grid[0][2]) // 2
-                position_ids[2, start:end] = torch.arange(start + offset, start + max_d + offset, device=embeds.device).unsqueeze(0).repeat(math.ceil((end - start) / max_d), 1).flatten(0)[:end - start]
-                offset += len_max - (end - start)
+                position_ids[2, start:end] = torch.arange(start, start + max_d, device=embeds.device).unsqueeze(0).repeat(math.ceil((end - start) / max_d), 1).flatten(0)[:end - start]
 
         if grid is None:
             position_ids = None

comfy/text_encoders/qwen_image.py

@@ -18,22 +18,13 @@ class QwenImageTokenizer(sd1_clip.SD1Tokenizer):
         self.llama_template_images = "<|im_start|>system\nDescribe the key features of the input image (color, shape, size, texture, objects, background), then explain how the user's text instruction should alter or modify the image. Generate a new image that meets the user's requirements while maintaining consistency with the original input where appropriate.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{}<|im_end|>\n<|im_start|>assistant\n"
 
     def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], **kwargs):
-        skip_template = False
-        if text.startswith('<|im_start|>'):
-            skip_template = True
-        if text.startswith('<|start_header_id|>'):
-            skip_template = True
-
-        if skip_template:
-            llama_text = text
-        else:
-            if llama_template is None:
-                if len(images) > 0:
-                    llama_text = self.llama_template_images.format(text)
-                else:
-                    llama_text = self.llama_template.format(text)
+        if llama_template is None:
+            if len(images) > 0:
+                llama_text = self.llama_template_images.format(text)
             else:
-                llama_text = llama_template.format(text)
+                llama_text = self.llama_template.format(text)
+        else:
+            llama_text = llama_template.format(text)
         tokens = super().tokenize_with_weights(llama_text, return_word_ids=return_word_ids, disable_weights=True, **kwargs)
         key_name = next(iter(tokens))
         embed_count = 0
@@ -56,23 +47,22 @@ class QwenImageTEModel(sd1_clip.SD1ClipModel):
     def __init__(self, device="cpu", dtype=None, model_options={}):
         super().__init__(device=device, dtype=dtype, name="qwen25_7b", clip_model=Qwen25_7BVLIModel, model_options=model_options)
 
-    def encode_token_weights(self, token_weight_pairs, template_end=-1):
+    def encode_token_weights(self, token_weight_pairs):
         out, pooled, extra = super().encode_token_weights(token_weight_pairs)
         tok_pairs = token_weight_pairs["qwen25_7b"][0]
         count_im_start = 0
-        if template_end == -1:
-            for i, v in enumerate(tok_pairs):
-                elem = v[0]
-                if not torch.is_tensor(elem):
-                    if isinstance(elem, numbers.Integral):
-                        if elem == 151644 and count_im_start < 2:
-                            template_end = i
-                            count_im_start += 1
+        for i, v in enumerate(tok_pairs):
+            elem = v[0]
+            if not torch.is_tensor(elem):
+                if isinstance(elem, numbers.Integral):
+                    if elem == 151644 and count_im_start < 2:
+                        template_end = i
+                        count_im_start += 1
 
-            if out.shape[1] > (template_end + 3):
-                if tok_pairs[template_end + 1][0] == 872:
-                    if tok_pairs[template_end + 2][0] == 198:
-                        template_end += 3
+        if out.shape[1] > (template_end + 3):
+            if tok_pairs[template_end + 1][0] == 872:
+                if tok_pairs[template_end + 2][0] == 198:
+                    template_end += 3
 
         out = out[:, template_end:]
 

comfy/utils.py

@@ -50,16 +50,10 @@ if hasattr(torch.serialization, "add_safe_globals"):  # TODO: this was added in
 else:
     logging.info("Warning, you are using an old pytorch version and some ckpt/pt files might be loaded unsafely. Upgrading to 2.4 or above is recommended.")
 
-def is_html_file(file_path):
-    with open(file_path, "rb") as f:
-        content = f.read(100)
-        return b"<!DOCTYPE html>" in content or b"<html" in content
-
 def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
     if device is None:
         device = torch.device("cpu")
     metadata = None
-
     if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"):
         try:
             with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f:
@@ -72,8 +66,6 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
                 if return_metadata:
                     metadata = f.metadata()
         except Exception as e:
-            if is_html_file(ckpt):
-                raise ValueError("{}\n\nFile path: {}\n\nThe requested file is an HTML document not a safetensors file. Please re-download the file, not the web page.".format(e, ckpt))
             if len(e.args) > 0:
                 message = e.args[0]
                 if "HeaderTooLarge" in message:
@@ -101,8 +93,6 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False):
                     sd = pl_sd
             else:
                 sd = pl_sd
-
-    # populate_db_with_asset(ckpt) # surprise tool that can help us later - performs hashing on model file
     return (sd, metadata) if return_metadata else sd
 
 def save_torch_file(sd, ckpt, metadata=None):

comfy/weight_adapter/loha.py

@@ -130,12 +130,12 @@ class LoHaAdapter(WeightAdapterBase):
     def create_train(cls, weight, rank=1, alpha=1.0):
         out_dim = weight.shape[0]
         in_dim = weight.shape[1:].numel()
-        mat1 = torch.empty(out_dim, rank, device=weight.device, dtype=torch.float32)
-        mat2 = torch.empty(rank, in_dim, device=weight.device, dtype=torch.float32)
+        mat1 = torch.empty(out_dim, rank, device=weight.device, dtype=weight.dtype)
+        mat2 = torch.empty(rank, in_dim, device=weight.device, dtype=weight.dtype)
         torch.nn.init.normal_(mat1, 0.1)
         torch.nn.init.constant_(mat2, 0.0)
-        mat3 = torch.empty(out_dim, rank, device=weight.device, dtype=torch.float32)
-        mat4 = torch.empty(rank, in_dim, device=weight.device, dtype=torch.float32)
+        mat3 = torch.empty(out_dim, rank, device=weight.device, dtype=weight.dtype)
+        mat4 = torch.empty(rank, in_dim, device=weight.device, dtype=weight.dtype)
         torch.nn.init.normal_(mat3, 0.1)
         torch.nn.init.normal_(mat4, 0.01)
         return LohaDiff(

comfy/weight_adapter/lokr.py

@@ -89,8 +89,8 @@ class LoKrAdapter(WeightAdapterBase):
         in_dim = weight.shape[1:].numel()
         out1, out2 = factorization(out_dim, rank)
         in1, in2 = factorization(in_dim, rank)
-        mat1 = torch.empty(out1, in1, device=weight.device, dtype=torch.float32)
-        mat2 = torch.empty(out2, in2, device=weight.device, dtype=torch.float32)
+        mat1 = torch.empty(out1, in1, device=weight.device, dtype=weight.dtype)
+        mat2 = torch.empty(out2, in2, device=weight.device, dtype=weight.dtype)
         torch.nn.init.kaiming_uniform_(mat2, a=5**0.5)
         torch.nn.init.constant_(mat1, 0.0)
         return LokrDiff(

comfy/weight_adapter/lora.py

@@ -66,8 +66,8 @@ class LoRAAdapter(WeightAdapterBase):
     def create_train(cls, weight, rank=1, alpha=1.0):
         out_dim = weight.shape[0]
         in_dim = weight.shape[1:].numel()
-        mat1 = torch.empty(out_dim, rank, device=weight.device, dtype=torch.float32)
-        mat2 = torch.empty(rank, in_dim, device=weight.device, dtype=torch.float32)
+        mat1 = torch.empty(out_dim, rank, device=weight.device, dtype=weight.dtype)
+        mat2 = torch.empty(rank, in_dim, device=weight.device, dtype=weight.dtype)
         torch.nn.init.kaiming_uniform_(mat1, a=5**0.5)
         torch.nn.init.constant_(mat2, 0.0)
         return LoraDiff(

comfy/weight_adapter/oft.py

@@ -68,7 +68,7 @@ class OFTAdapter(WeightAdapterBase):
     def create_train(cls, weight, rank=1, alpha=1.0):
         out_dim = weight.shape[0]
         block_size, block_num = factorization(out_dim, rank)
-        block = torch.zeros(block_num, block_size, block_size, device=weight.device, dtype=torch.float32)
+        block = torch.zeros(block_num, block_size, block_size, device=weight.device, dtype=weight.dtype)
         return OFTDiff(
             (block, None, alpha, None)
         )

comfy_api/latest/_io.py

@@ -331,7 +331,7 @@ class String(ComfyTypeIO):
             })
 
 @comfytype(io_type="COMBO")
-class Combo(ComfyTypeIO):
+class Combo(ComfyTypeI):
     Type = str
     class Input(WidgetInput):
         """Combo input (dropdown)."""
@@ -360,14 +360,6 @@ class Combo(ComfyTypeIO):
                 "remote": self.remote.as_dict() if self.remote else None,
             })
 
-    class Output(Output):
-        def __init__(self, id: str=None, display_name: str=None, options: list[str]=None, tooltip: str=None, is_output_list=False):
-            super().__init__(id, display_name, tooltip, is_output_list)
-            self.options = options if options is not None else []
-
-        @property
-        def io_type(self):
-            return self.options
 
 @comfytype(io_type="COMBO")
 class MultiCombo(ComfyTypeI):
@@ -392,20 +384,6 @@ class MultiCombo(ComfyTypeI):
             })
             return to_return
 
-@comfytype(io_type="ASSET")
-class Asset(ComfyTypeI):
-    class Input(WidgetInput):
-        def __init__(self, id: str, query_tags: list[str], display_name: str=None, optional=False, tooltip: str=None, lazy: bool=None,
-                    default: str=None, socketless: bool=None):
-            super().__init__(id, display_name, optional, tooltip, lazy, default, socketless)
-            self.query_tags = query_tags
-
-        def as_dict(self):
-            to_return = super().as_dict() | prune_dict({
-                "query_tags": self.query_tags
-            })
-            return to_return
-
 @comfytype(io_type="IMAGE")
 class Image(ComfyTypeIO):
     Type = torch.Tensor
@@ -1212,18 +1190,13 @@ class _ComfyNodeBaseInternal(_ComfyNodeInternal):
         raise NotImplementedError
 
     @classmethod
-    def validate_inputs(cls, **kwargs) -> bool | str:
-        """Optionally, define this function to validate inputs; equivalent to V1's VALIDATE_INPUTS.
-
-        If the function returns a string, it will be used as the validation error message for the node.
-        """
+    def validate_inputs(cls, **kwargs) -> bool:
+        """Optionally, define this function to validate inputs; equivalent to V1's VALIDATE_INPUTS."""
         raise NotImplementedError
 
     @classmethod
     def fingerprint_inputs(cls, **kwargs) -> Any:
-        """Optionally, define this function to fingerprint inputs; equivalent to V1's IS_CHANGED.
-
-        If this function returns the same value as last run, the node will not be executed."""
+        """Optionally, define this function to fingerprint inputs; equivalent to V1's IS_CHANGED."""
         raise NotImplementedError
 
     @classmethod

comfy_api_nodes/apinode_utils.py

@@ -518,71 +518,6 @@ async def upload_audio_to_comfyapi(
     return await upload_file_to_comfyapi(audio_bytes_io, filename, mime_type, auth_kwargs)
 
 
-def f32_pcm(wav: torch.Tensor) -> torch.Tensor:
-    """Convert audio to float 32 bits PCM format. Copy-paste from nodes_audio.py file."""
-    if wav.dtype.is_floating_point:
-        return wav
-    elif wav.dtype == torch.int16:
-        return wav.float() / (2 ** 15)
-    elif wav.dtype == torch.int32:
-        return wav.float() / (2 ** 31)
-    raise ValueError(f"Unsupported wav dtype: {wav.dtype}")
-
-
-def audio_bytes_to_audio_input(audio_bytes: bytes,) -> dict:
-    """
-    Decode any common audio container from bytes using PyAV and return
-    a Comfy AUDIO dict: {"waveform": [1, C, T] float32, "sample_rate": int}.
-    """
-    with av.open(io.BytesIO(audio_bytes)) as af:
-        if not af.streams.audio:
-            raise ValueError("No audio stream found in response.")
-        stream = af.streams.audio[0]
-
-        in_sr = int(stream.codec_context.sample_rate)
-        out_sr = in_sr
-
-        frames: list[torch.Tensor] = []
-        n_channels = stream.channels or 1
-
-        for frame in af.decode(streams=stream.index):
-            arr = frame.to_ndarray()  # shape can be [C, T] or [T, C] or [T]
-            buf = torch.from_numpy(arr)
-            if buf.ndim == 1:
-                buf = buf.unsqueeze(0)  # [T] -> [1, T]
-            elif buf.shape[0] != n_channels and buf.shape[-1] == n_channels:
-                buf = buf.transpose(0, 1).contiguous()  # [T, C] -> [C, T]
-            elif buf.shape[0] != n_channels:
-                buf = buf.reshape(-1, n_channels).t().contiguous()  # fallback to [C, T]
-            frames.append(buf)
-
-    if not frames:
-        raise ValueError("Decoded zero audio frames.")
-
-    wav = torch.cat(frames, dim=1)  # [C, T]
-    wav = f32_pcm(wav)
-    return {"waveform": wav.unsqueeze(0).contiguous(), "sample_rate": out_sr}
-
-
-def audio_input_to_mp3(audio: AudioInput) -> io.BytesIO:
-    waveform = audio["waveform"].cpu()
-
-    output_buffer = io.BytesIO()
-    output_container = av.open(output_buffer, mode='w', format="mp3")
-
-    out_stream = output_container.add_stream("libmp3lame", rate=audio["sample_rate"])
-    out_stream.bit_rate = 320000
-
-    frame = av.AudioFrame.from_ndarray(waveform.movedim(0, 1).reshape(1, -1).float().numpy(), format='flt', layout='mono' if waveform.shape[0] == 1 else 'stereo')
-    frame.sample_rate = audio["sample_rate"]
-    frame.pts = 0
-    output_container.mux(out_stream.encode(frame))
-    output_container.mux(out_stream.encode(None))
-    output_container.close()
-    output_buffer.seek(0)
-    return output_buffer
-
-
 def audio_to_base64_string(
     audio: AudioInput, container_format: str = "mp4", codec_name: str = "aac"
 ) -> str:

comfy_api_nodes/apis/__init__.py

@@ -951,11 +951,7 @@ class MagicPrompt2(str, Enum):
 
 
 class StyleType1(str, Enum):
-    AUTO = 'AUTO'
     GENERAL = 'GENERAL'
-    REALISTIC = 'REALISTIC'
-    DESIGN = 'DESIGN'
-    FICTION = 'FICTION'
 
 
 class ImagenImageGenerationInstance(BaseModel):
@@ -2680,7 +2676,7 @@ class ReleaseNote(BaseModel):
 
 
 class RenderingSpeed(str, Enum):
-    DEFAULT = 'DEFAULT'
+    BALANCED = 'BALANCED'
     TURBO = 'TURBO'
     QUALITY = 'QUALITY'
 
@@ -4922,14 +4918,6 @@ class IdeogramV3EditRequest(BaseModel):
         None,
         description='A set of images to use as style references (maximum total size 10MB across all style references). The images should be in JPEG, PNG or WebP format.',
     )
-    character_reference_images: Optional[List[str]] = Field(
-        None,
-        description='Generations with character reference are subject to the character reference pricing. A set of images to use as character references (maximum total size 10MB across all character references), currently only supports 1 character reference image. The images should be in JPEG, PNG or WebP format.'
-    )
-    character_reference_images_mask: Optional[List[str]] = Field(
-        None,
-        description='Optional masks for character reference images. When provided, must match the number of character_reference_images. Each mask should be a grayscale image of the same dimensions as the corresponding character reference image. The images should be in JPEG, PNG or WebP format.'
-    )
 
 
 class IdeogramV3Request(BaseModel):
@@ -4963,14 +4951,6 @@ class IdeogramV3Request(BaseModel):
     style_type: Optional[StyleType1] = Field(
         None, description='The type of style to apply'
     )
-    character_reference_images: Optional[List[str]] = Field(
-        None,
-        description='Generations with character reference are subject to the character reference pricing. A set of images to use as character references (maximum total size 10MB across all character references), currently only supports 1 character reference image. The images should be in JPEG, PNG or WebP format.'
-    )
-    character_reference_images_mask: Optional[List[str]] = Field(
-        None,
-        description='Optional masks for character reference images. When provided, must match the number of character_reference_images. Each mask should be a grayscale image of the same dimensions as the corresponding character reference image. The images should be in JPEG, PNG or WebP format.'
-    )
 
 
 class ImagenGenerateImageResponse(BaseModel):

comfy_api_nodes/apis/client.py

@@ -683,7 +683,7 @@ class SynchronousOperation(Generic[T, R]):
         auth_token: Optional[str] = None,
         comfy_api_key: Optional[str] = None,
         auth_kwargs: Optional[Dict[str, str]] = None,
-        timeout: float = 7200.0,
+        timeout: float = 604800.0,
         verify_ssl: bool = True,
         content_type: str = "application/json",
         multipart_parser: Callable | None = None,

comfy_api_nodes/apis/rodin_api.py

@@ -9,9 +9,8 @@ class Rodin3DGenerateRequest(BaseModel):
     seed: int = Field(..., description="seed_")
     tier: str = Field(..., description="Tier of generation.")
     material: str = Field(..., description="The material type.")
-    quality_override: int = Field(..., description="The poly count of the mesh.")
+    quality: str = Field(..., description="The generation quality of the mesh.")
     mesh_mode: str = Field(..., description="It controls the type of faces of generated models.")
-    TAPose: Optional[bool] = Field(None, description="")
 
 class GenerateJobsData(BaseModel):
     uuids: List[str] = Field(..., description="str LIST")

comfy_api_nodes/apis/stability_api.py

@@ -125,25 +125,3 @@ class StabilityResultsGetResponse(BaseModel):
 
 class StabilityAsyncResponse(BaseModel):
     id: Optional[str] = Field(None)
-
-
-class StabilityTextToAudioRequest(BaseModel):
-    model: str = Field(...)
-    prompt: str = Field(...)
-    duration: int = Field(190, ge=1, le=190)
-    seed: int = Field(0, ge=0, le=4294967294)
-    steps: int = Field(8, ge=4, le=8)
-    output_format: str = Field("wav")
-
-
-class StabilityAudioToAudioRequest(StabilityTextToAudioRequest):
-    strength: float = Field(0.01, ge=0.01, le=1.0)
-
-
-class StabilityAudioInpaintRequest(StabilityTextToAudioRequest):
-    mask_start: int = Field(30, ge=0, le=190)
-    mask_end: int = Field(190, ge=0, le=190)
-
-
-class StabilityAudioResponse(BaseModel):
-    audio: Optional[str] = Field(None)

comfy_api_nodes/nodes_ideogram.py

@@ -255,7 +255,6 @@ class IdeogramV1(comfy_io.ComfyNode):
             display_name="Ideogram V1",
             category="api node/image/Ideogram",
             description="Generates images using the Ideogram V1 model.",
-            is_api_node=True,
             inputs=[
                 comfy_io.String.Input(
                     "prompt",
@@ -384,7 +383,6 @@ class IdeogramV2(comfy_io.ComfyNode):
             display_name="Ideogram V2",
             category="api node/image/Ideogram",
             description="Generates images using the Ideogram V2 model.",
-            is_api_node=True,
             inputs=[
                 comfy_io.String.Input(
                     "prompt",
@@ -554,7 +552,6 @@ class IdeogramV3(comfy_io.ComfyNode):
             category="api node/image/Ideogram",
             description="Generates images using the Ideogram V3 model. "
                         "Supports both regular image generation from text prompts and image editing with mask.",
-            is_api_node=True,
             inputs=[
                 comfy_io.String.Input(
                     "prompt",
@@ -615,21 +612,11 @@ class IdeogramV3(comfy_io.ComfyNode):
                 ),
                 comfy_io.Combo.Input(
                     "rendering_speed",
-                    options=["DEFAULT", "TURBO", "QUALITY"],
-                    default="DEFAULT",
+                    options=["BALANCED", "TURBO", "QUALITY"],
+                    default="BALANCED",
                     tooltip="Controls the trade-off between generation speed and quality",
                     optional=True,
                 ),
-                comfy_io.Image.Input(
-                    "character_image",
-                    tooltip="Image to use as character reference.",
-                    optional=True,
-                ),
-                comfy_io.Mask.Input(
-                    "character_mask",
-                    tooltip="Optional mask for character reference image.",
-                    optional=True,
-                ),
             ],
             outputs=[
                 comfy_io.Image.Output(),
@@ -652,46 +639,12 @@ class IdeogramV3(comfy_io.ComfyNode):
         magic_prompt_option="AUTO",
         seed=0,
         num_images=1,
-        rendering_speed="DEFAULT",
-        character_image=None,
-        character_mask=None,
+        rendering_speed="BALANCED",
     ):
         auth = {
             "auth_token": cls.hidden.auth_token_comfy_org,
             "comfy_api_key": cls.hidden.api_key_comfy_org,
         }
-        if rendering_speed == "BALANCED":  # for backward compatibility
-            rendering_speed = "DEFAULT"
-
-        character_img_binary = None
-        character_mask_binary = None
-
-        if character_image is not None:
-            input_tensor = character_image.squeeze().cpu()
-            if character_mask is not None:
-                character_mask = resize_mask_to_image(character_mask, character_image, allow_gradient=False)
-                character_mask = 1.0 - character_mask
-                if character_mask.shape[1:] != character_image.shape[1:-1]:
-                    raise Exception("Character mask and image must be the same size")
-
-                mask_np = (character_mask.squeeze().cpu().numpy() * 255).astype(np.uint8)
-                mask_img = Image.fromarray(mask_np)
-                mask_byte_arr = BytesIO()
-                mask_img.save(mask_byte_arr, format="PNG")
-                mask_byte_arr.seek(0)
-                character_mask_binary = mask_byte_arr
-                character_mask_binary.name = "mask.png"
-
-            img_np = (input_tensor.numpy() * 255).astype(np.uint8)
-            img = Image.fromarray(img_np)
-            img_byte_arr = BytesIO()
-            img.save(img_byte_arr, format="PNG")
-            img_byte_arr.seek(0)
-            character_img_binary = img_byte_arr
-            character_img_binary.name = "image.png"
-        elif character_mask is not None:
-            raise Exception("Character mask requires character image to be present")
-
         # Check if both image and mask are provided for editing mode
         if image is not None and mask is not None:
             # Edit mode
@@ -740,15 +693,6 @@ class IdeogramV3(comfy_io.ComfyNode):
             if num_images > 1:
                 edit_request.num_images = num_images
 
-            files = {
-                "image": img_binary,
-                "mask": mask_binary,
-            }
-            if character_img_binary:
-                files["character_reference_images"] = character_img_binary
-            if character_mask_binary:
-                files["character_mask_binary"] = character_mask_binary
-
             # Execute the operation for edit mode
             operation = SynchronousOperation(
                 endpoint=ApiEndpoint(
@@ -758,7 +702,10 @@ class IdeogramV3(comfy_io.ComfyNode):
                     response_model=IdeogramGenerateResponse,
                 ),
                 request=edit_request,
-                files=files,
+                files={
+                    "image": img_binary,
+                    "mask": mask_binary,
+                },
                 content_type="multipart/form-data",
                 auth_kwargs=auth,
             )
@@ -792,14 +739,6 @@ class IdeogramV3(comfy_io.ComfyNode):
             if num_images > 1:
                 gen_request.num_images = num_images
 
-            files = {}
-            if character_img_binary:
-                files["character_reference_images"] = character_img_binary
-            if character_mask_binary:
-                files["character_mask_binary"] = character_mask_binary
-            if files:
-                gen_request.style_type = "AUTO"
-
             # Execute the operation for generation mode
             operation = SynchronousOperation(
                 endpoint=ApiEndpoint(
@@ -809,8 +748,6 @@ class IdeogramV3(comfy_io.ComfyNode):
                     response_model=IdeogramGenerateResponse,
                 ),
                 request=gen_request,
-                files=files if files else None,
-                content_type="multipart/form-data",
                 auth_kwargs=auth,
             )
 

comfy_api_nodes/nodes_kling.py

@@ -846,8 +846,6 @@ class KlingStartEndFrameNode(KlingImage2VideoNode):
             "pro mode / 10s duration / kling-v1-5": ("pro", "10", "kling-v1-5"),
             "pro mode / 5s duration / kling-v1-6": ("pro", "5", "kling-v1-6"),
             "pro mode / 10s duration / kling-v1-6": ("pro", "10", "kling-v1-6"),
-            "pro mode / 5s duration / kling-v2-1": ("pro", "5", "kling-v2-1"),
-            "pro mode / 10s duration / kling-v2-1": ("pro", "10", "kling-v2-1"),
         }
 
     @classmethod

comfy_api_nodes/nodes_minimax.py

@@ -1,10 +1,9 @@
 from inspect import cleandoc
-from typing import Optional
+from typing import Union
 import logging
 import torch
 
-from typing_extensions import override
-from comfy_api.latest import ComfyExtension, io as comfy_io
+from comfy.comfy_types.node_typing import IO
 from comfy_api.input_impl.video_types import VideoFromFile
 from comfy_api_nodes.apis import (
     MinimaxVideoGenerationRequest,
@@ -12,7 +11,7 @@ from comfy_api_nodes.apis import (
     MinimaxFileRetrieveResponse,
     MinimaxTaskResultResponse,
     SubjectReferenceItem,
-    MiniMaxModel,
+    MiniMaxModel
 )
 from comfy_api_nodes.apis.client import (
     ApiEndpoint,
@@ -32,398 +31,372 @@ from server import PromptServer
 I2V_AVERAGE_DURATION = 114
 T2V_AVERAGE_DURATION = 234
 
-
-async def _generate_mm_video(
-    *,
-    auth: dict[str, str],
-    node_id: str,
-    prompt_text: str,
-    seed: int,
-    model: str,
-    image: Optional[torch.Tensor] = None,   # used for ImageToVideo
-    subject: Optional[torch.Tensor] = None, # used for SubjectToVideo
-    average_duration: Optional[int] = None,
-) -> comfy_io.NodeOutput:
-    if image is None:
-        validate_string(prompt_text, field_name="prompt_text")
-    # upload image, if passed in
-    image_url = None
-    if image is not None:
-        image_url = (await upload_images_to_comfyapi(image, max_images=1, auth_kwargs=auth))[0]
-
-    # TODO: figure out how to deal with subject properly, API returns invalid params when using S2V-01 model
-    subject_reference = None
-    if subject is not None:
-        subject_url = (await upload_images_to_comfyapi(subject, max_images=1, auth_kwargs=auth))[0]
-        subject_reference = [SubjectReferenceItem(image=subject_url)]
-
-
-    video_generate_operation = SynchronousOperation(
-        endpoint=ApiEndpoint(
-            path="/proxy/minimax/video_generation",
-            method=HttpMethod.POST,
-            request_model=MinimaxVideoGenerationRequest,
-            response_model=MinimaxVideoGenerationResponse,
-        ),
-        request=MinimaxVideoGenerationRequest(
-            model=MiniMaxModel(model),
-            prompt=prompt_text,
-            callback_url=None,
-            first_frame_image=image_url,
-            subject_reference=subject_reference,
-            prompt_optimizer=None,
-        ),
-        auth_kwargs=auth,
-    )
-    response = await video_generate_operation.execute()
-
-    task_id = response.task_id
-    if not task_id:
-        raise Exception(f"MiniMax generation failed: {response.base_resp}")
-
-    video_generate_operation = PollingOperation(
-        poll_endpoint=ApiEndpoint(
-            path="/proxy/minimax/query/video_generation",
-            method=HttpMethod.GET,
-            request_model=EmptyRequest,
-            response_model=MinimaxTaskResultResponse,
-            query_params={"task_id": task_id},
-        ),
-        completed_statuses=["Success"],
-        failed_statuses=["Fail"],
-        status_extractor=lambda x: x.status.value,
-        estimated_duration=average_duration,
-        node_id=node_id,
-        auth_kwargs=auth,
-    )
-    task_result = await video_generate_operation.execute()
-
-    file_id = task_result.file_id
-    if file_id is None:
-        raise Exception("Request was not successful. Missing file ID.")
-    file_retrieve_operation = SynchronousOperation(
-        endpoint=ApiEndpoint(
-            path="/proxy/minimax/files/retrieve",
-            method=HttpMethod.GET,
-            request_model=EmptyRequest,
-            response_model=MinimaxFileRetrieveResponse,
-            query_params={"file_id": int(file_id)},
-        ),
-        request=EmptyRequest(),
-        auth_kwargs=auth,
-    )
-    file_result = await file_retrieve_operation.execute()
-
-    file_url = file_result.file.download_url
-    if file_url is None:
-        raise Exception(
-            f"No video was found in the response. Full response: {file_result.model_dump()}"
-        )
-    logging.info("Generated video URL: %s", file_url)
-    if node_id:
-        if hasattr(file_result.file, "backup_download_url"):
-            message = f"Result URL: {file_url}\nBackup URL: {file_result.file.backup_download_url}"
-        else:
-            message = f"Result URL: {file_url}"
-        PromptServer.instance.send_progress_text(message, node_id)
-
-    # Download and return as VideoFromFile
-    video_io = await download_url_to_bytesio(file_url)
-    if video_io is None:
-        error_msg = f"Failed to download video from {file_url}"
-        logging.error(error_msg)
-        raise Exception(error_msg)
-    return comfy_io.NodeOutput(VideoFromFile(video_io))
-
-
-class MinimaxTextToVideoNode(comfy_io.ComfyNode):
+class MinimaxTextToVideoNode:
     """
     Generates videos synchronously based on a prompt, and optional parameters using MiniMax's API.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MinimaxTextToVideoNode",
-            display_name="MiniMax Text to Video",
-            category="api node/video/MiniMax",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt_text",
-                    multiline=True,
-                    default="",
-                    tooltip="Text prompt to guide the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=["T2V-01", "T2V-01-Director"],
-                    default="T2V-01",
-                    tooltip="Model to use for video generation",
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    step=1,
-                    control_after_generate=True,
-                    tooltip="The random seed used for creating the noise.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    AVERAGE_DURATION = T2V_AVERAGE_DURATION
 
     @classmethod
-    async def execute(
-        cls,
-        prompt_text: str,
-        model: str = "T2V-01",
-        seed: int = 0,
-    ) -> comfy_io.NodeOutput:
-        return await _generate_mm_video(
-            auth={
-                "auth_token": cls.hidden.auth_token_comfy_org,
-                "comfy_api_key": cls.hidden.api_key_comfy_org,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt_text": (
+                    "STRING",
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Text prompt to guide the video generation",
+                    },
+                ),
+                "model": (
+                    [
+                        "T2V-01",
+                        "T2V-01-Director",
+                    ],
+                    {
+                        "default": "T2V-01",
+                        "tooltip": "Model to use for video generation",
+                    },
+                ),
             },
-            node_id=cls.hidden.unique_id,
-            prompt_text=prompt_text,
-            seed=seed,
-            model=model,
-            image=None,
-            subject=None,
-            average_duration=T2V_AVERAGE_DURATION,
+            "optional": {
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "The random seed used for creating the noise.",
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    DESCRIPTION = "Generates videos from prompts using MiniMax's API"
+    FUNCTION = "generate_video"
+    CATEGORY = "api node/video/MiniMax"
+    API_NODE = True
+
+    async def generate_video(
+        self,
+        prompt_text,
+        seed=0,
+        model="T2V-01",
+        image: torch.Tensor=None, # used for ImageToVideo
+        subject: torch.Tensor=None, # used for SubjectToVideo
+        unique_id: Union[str, None]=None,
+        **kwargs,
+    ):
+        '''
+        Function used between MiniMax nodes - supports T2V, I2V, and S2V, based on provided arguments.
+        '''
+        if image is None:
+            validate_string(prompt_text, field_name="prompt_text")
+        # upload image, if passed in
+        image_url = None
+        if image is not None:
+            image_url = (await upload_images_to_comfyapi(image, max_images=1, auth_kwargs=kwargs))[0]
+
+        # TODO: figure out how to deal with subject properly, API returns invalid params when using S2V-01 model
+        subject_reference = None
+        if subject is not None:
+            subject_url = (await upload_images_to_comfyapi(subject, max_images=1, auth_kwargs=kwargs))[0]
+            subject_reference = [SubjectReferenceItem(image=subject_url)]
+
+
+        video_generate_operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/minimax/video_generation",
+                method=HttpMethod.POST,
+                request_model=MinimaxVideoGenerationRequest,
+                response_model=MinimaxVideoGenerationResponse,
+            ),
+            request=MinimaxVideoGenerationRequest(
+                model=MiniMaxModel(model),
+                prompt=prompt_text,
+                callback_url=None,
+                first_frame_image=image_url,
+                subject_reference=subject_reference,
+                prompt_optimizer=None,
+            ),
+            auth_kwargs=kwargs,
         )
+        response = await video_generate_operation.execute()
+
+        task_id = response.task_id
+        if not task_id:
+            raise Exception(f"MiniMax generation failed: {response.base_resp}")
+
+        video_generate_operation = PollingOperation(
+            poll_endpoint=ApiEndpoint(
+                path="/proxy/minimax/query/video_generation",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=MinimaxTaskResultResponse,
+                query_params={"task_id": task_id},
+            ),
+            completed_statuses=["Success"],
+            failed_statuses=["Fail"],
+            status_extractor=lambda x: x.status.value,
+            estimated_duration=self.AVERAGE_DURATION,
+            node_id=unique_id,
+            auth_kwargs=kwargs,
+        )
+        task_result = await video_generate_operation.execute()
+
+        file_id = task_result.file_id
+        if file_id is None:
+            raise Exception("Request was not successful. Missing file ID.")
+        file_retrieve_operation = SynchronousOperation(
+            endpoint=ApiEndpoint(
+                path="/proxy/minimax/files/retrieve",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=MinimaxFileRetrieveResponse,
+                query_params={"file_id": int(file_id)},
+            ),
+            request=EmptyRequest(),
+            auth_kwargs=kwargs,
+        )
+        file_result = await file_retrieve_operation.execute()
+
+        file_url = file_result.file.download_url
+        if file_url is None:
+            raise Exception(
+                f"No video was found in the response. Full response: {file_result.model_dump()}"
+            )
+        logging.info(f"Generated video URL: {file_url}")
+        if unique_id:
+            if hasattr(file_result.file, "backup_download_url"):
+                message = f"Result URL: {file_url}\nBackup URL: {file_result.file.backup_download_url}"
+            else:
+                message = f"Result URL: {file_url}"
+            PromptServer.instance.send_progress_text(message, unique_id)
+
+        video_io = await download_url_to_bytesio(file_url)
+        if video_io is None:
+            error_msg = f"Failed to download video from {file_url}"
+            logging.error(error_msg)
+            raise Exception(error_msg)
+        return (VideoFromFile(video_io),)
 
 
-class MinimaxImageToVideoNode(comfy_io.ComfyNode):
+class MinimaxImageToVideoNode(MinimaxTextToVideoNode):
     """
     Generates videos synchronously based on an image and prompt, and optional parameters using MiniMax's API.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MinimaxImageToVideoNode",
-            display_name="MiniMax Image to Video",
-            category="api node/video/MiniMax",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input(
-                    "image",
-                    tooltip="Image to use as first frame of video generation",
-                ),
-                comfy_io.String.Input(
-                    "prompt_text",
-                    multiline=True,
-                    default="",
-                    tooltip="Text prompt to guide the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=["I2V-01-Director", "I2V-01", "I2V-01-live"],
-                    default="I2V-01",
-                    tooltip="Model to use for video generation",
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    step=1,
-                    control_after_generate=True,
-                    tooltip="The random seed used for creating the noise.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    AVERAGE_DURATION = I2V_AVERAGE_DURATION
 
     @classmethod
-    async def execute(
-        cls,
-        image: torch.Tensor,
-        prompt_text: str,
-        model: str = "I2V-01",
-        seed: int = 0,
-    ) -> comfy_io.NodeOutput:
-        return await _generate_mm_video(
-            auth={
-                "auth_token": cls.hidden.auth_token_comfy_org,
-                "comfy_api_key": cls.hidden.api_key_comfy_org,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "image": (
+                    IO.IMAGE,
+                    {
+                        "tooltip": "Image to use as first frame of video generation"
+                    },
+                ),
+                "prompt_text": (
+                    "STRING",
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Text prompt to guide the video generation",
+                    },
+                ),
+                "model": (
+                    [
+                        "I2V-01-Director",
+                        "I2V-01",
+                        "I2V-01-live",
+                    ],
+                    {
+                        "default": "I2V-01",
+                        "tooltip": "Model to use for video generation",
+                    },
+                ),
             },
-            node_id=cls.hidden.unique_id,
-            prompt_text=prompt_text,
-            seed=seed,
-            model=model,
-            image=image,
-            subject=None,
-            average_duration=I2V_AVERAGE_DURATION,
-        )
+            "optional": {
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "The random seed used for creating the noise.",
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    DESCRIPTION = "Generates videos from an image and prompts using MiniMax's API"
+    FUNCTION = "generate_video"
+    CATEGORY = "api node/video/MiniMax"
+    API_NODE = True
 
 
-class MinimaxSubjectToVideoNode(comfy_io.ComfyNode):
+class MinimaxSubjectToVideoNode(MinimaxTextToVideoNode):
     """
     Generates videos synchronously based on an image and prompt, and optional parameters using MiniMax's API.
     """
 
-    @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MinimaxSubjectToVideoNode",
-            display_name="MiniMax Subject to Video",
-            category="api node/video/MiniMax",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.Image.Input(
-                    "subject",
-                    tooltip="Image of subject to reference for video generation",
-                ),
-                comfy_io.String.Input(
-                    "prompt_text",
-                    multiline=True,
-                    default="",
-                    tooltip="Text prompt to guide the video generation",
-                ),
-                comfy_io.Combo.Input(
-                    "model",
-                    options=["S2V-01"],
-                    default="S2V-01",
-                    tooltip="Model to use for video generation",
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    step=1,
-                    control_after_generate=True,
-                    tooltip="The random seed used for creating the noise.",
-                    optional=True,
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    AVERAGE_DURATION = T2V_AVERAGE_DURATION
 
     @classmethod
-    async def execute(
-        cls,
-        subject: torch.Tensor,
-        prompt_text: str,
-        model: str = "S2V-01",
-        seed: int = 0,
-    ) -> comfy_io.NodeOutput:
-        return await _generate_mm_video(
-            auth={
-                "auth_token": cls.hidden.auth_token_comfy_org,
-                "comfy_api_key": cls.hidden.api_key_comfy_org,
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "subject": (
+                    IO.IMAGE,
+                    {
+                        "tooltip": "Image of subject to reference video generation"
+                    },
+                ),
+                "prompt_text": (
+                    "STRING",
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Text prompt to guide the video generation",
+                    },
+                ),
+                "model": (
+                    [
+                        "S2V-01",
+                    ],
+                    {
+                        "default": "S2V-01",
+                        "tooltip": "Model to use for video generation",
+                    },
+                ),
             },
-            node_id=cls.hidden.unique_id,
-            prompt_text=prompt_text,
-            seed=seed,
-            model=model,
-            image=None,
-            subject=subject,
-            average_duration=T2V_AVERAGE_DURATION,
-        )
+            "optional": {
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "The random seed used for creating the noise.",
+                    },
+                ),
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    DESCRIPTION = "Generates videos from an image and prompts using MiniMax's API"
+    FUNCTION = "generate_video"
+    CATEGORY = "api node/video/MiniMax"
+    API_NODE = True
 
 
-class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
+class MinimaxHailuoVideoNode:
     """Generates videos from prompt, with optional start frame using the new MiniMax Hailuo-02 model."""
 
     @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MinimaxHailuoVideoNode",
-            display_name="MiniMax Hailuo Video",
-            category="api node/video/MiniMax",
-            description=cleandoc(cls.__doc__ or ""),
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt_text",
-                    multiline=True,
-                    default="",
-                    tooltip="Text prompt to guide the video generation.",
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "prompt_text": (
+                    "STRING",
+                    {
+                        "multiline": True,
+                        "default": "",
+                        "tooltip": "Text prompt to guide the video generation.",
+                    },
                 ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=0,
-                    min=0,
-                    max=0xFFFFFFFFFFFFFFFF,
-                    step=1,
-                    control_after_generate=True,
-                    tooltip="The random seed used for creating the noise.",
-                    optional=True,
+            },
+            "optional": {
+                "seed": (
+                    IO.INT,
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "max": 0xFFFFFFFFFFFFFFFF,
+                        "control_after_generate": True,
+                        "tooltip": "The random seed used for creating the noise.",
+                    },
                 ),
-                comfy_io.Image.Input(
-                    "first_frame_image",
-                    tooltip="Optional image to use as the first frame to generate a video.",
-                    optional=True,
+                "first_frame_image": (
+                    IO.IMAGE,
+                    {
+                        "tooltip": "Optional image to use as the first frame to generate a video."
+                    },
                 ),
-                comfy_io.Boolean.Input(
-                    "prompt_optimizer",
-                    default=True,
-                    tooltip="Optimize prompt to improve generation quality when needed.",
-                    optional=True,
+                "prompt_optimizer": (
+                    IO.BOOLEAN,
+                    {
+                        "tooltip": "Optimize prompt to improve generation quality when needed.",
+                        "default": True,
+                    },
                 ),
-                comfy_io.Combo.Input(
-                    "duration",
-                    options=[6, 10],
-                    default=6,
-                    tooltip="The length of the output video in seconds.",
-                    optional=True,
+                "duration": (
+                    IO.COMBO,
+                    {
+                        "tooltip": "The length of the output video in seconds.",
+                        "default": 6,
+                        "options": [6, 10],
+                    },
                 ),
-                comfy_io.Combo.Input(
-                    "resolution",
-                    options=["768P", "1080P"],
-                    default="768P",
-                    tooltip="The dimensions of the video display. 1080p is 1920x1080, 768p is 1366x768.",
-                    optional=True,
+                "resolution": (
+                    IO.COMBO,
+                    {
+                        "tooltip": "The dimensions of the video display. "
+                                   "1080p corresponds to 1920 x 1080 pixels, 768p corresponds to 1366 x 768 pixels.",
+                        "default": "768P",
+                        "options": ["768P", "1080P"],
+                    },
                 ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        prompt_text: str,
-        seed: int = 0,
-        first_frame_image: Optional[torch.Tensor] = None,  # used for ImageToVideo
-        prompt_optimizer: bool = True,
-        duration: int = 6,
-        resolution: str = "768P",
-        model: str = "MiniMax-Hailuo-02",
-    ) -> comfy_io.NodeOutput:
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
         }
+
+    RETURN_TYPES = ("VIDEO",)
+    DESCRIPTION = cleandoc(__doc__ or "")
+    FUNCTION = "generate_video"
+    CATEGORY = "api node/video/MiniMax"
+    API_NODE = True
+
+    async def generate_video(
+        self,
+        prompt_text,
+        seed=0,
+        first_frame_image: torch.Tensor=None, # used for ImageToVideo
+        prompt_optimizer=True,
+        duration=6,
+        resolution="768P",
+        model="MiniMax-Hailuo-02",
+        unique_id: Union[str, None]=None,
+        **kwargs,
+    ):
         if first_frame_image is None:
             validate_string(prompt_text, field_name="prompt_text")
 
@@ -435,7 +408,7 @@ class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
         # upload image, if passed in
         image_url = None
         if first_frame_image is not None:
-            image_url = (await upload_images_to_comfyapi(first_frame_image, max_images=1, auth_kwargs=auth))[0]
+            image_url = (await upload_images_to_comfyapi(first_frame_image, max_images=1, auth_kwargs=kwargs))[0]
 
         video_generate_operation = SynchronousOperation(
             endpoint=ApiEndpoint(
@@ -453,7 +426,7 @@ class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
                 duration=duration,
                 resolution=resolution,
             ),
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         response = await video_generate_operation.execute()
 
@@ -474,8 +447,8 @@ class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
             failed_statuses=["Fail"],
             status_extractor=lambda x: x.status.value,
             estimated_duration=average_duration,
-            node_id=cls.hidden.unique_id,
-            auth_kwargs=auth,
+            node_id=unique_id,
+            auth_kwargs=kwargs,
         )
         task_result = await video_generate_operation.execute()
 
@@ -491,7 +464,7 @@ class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
                 query_params={"file_id": int(file_id)},
             ),
             request=EmptyRequest(),
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         file_result = await file_retrieve_operation.execute()
 
@@ -501,31 +474,34 @@ class MinimaxHailuoVideoNode(comfy_io.ComfyNode):
                 f"No video was found in the response. Full response: {file_result.model_dump()}"
             )
         logging.info(f"Generated video URL: {file_url}")
-        if cls.hidden.unique_id:
+        if unique_id:
             if hasattr(file_result.file, "backup_download_url"):
                 message = f"Result URL: {file_url}\nBackup URL: {file_result.file.backup_download_url}"
             else:
                 message = f"Result URL: {file_url}"
-            PromptServer.instance.send_progress_text(message, cls.hidden.unique_id)
+            PromptServer.instance.send_progress_text(message, unique_id)
 
         video_io = await download_url_to_bytesio(file_url)
         if video_io is None:
             error_msg = f"Failed to download video from {file_url}"
             logging.error(error_msg)
             raise Exception(error_msg)
-        return comfy_io.NodeOutput(VideoFromFile(video_io))
+        return (VideoFromFile(video_io),)
 
 
-class MinimaxExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
-        return [
-            MinimaxTextToVideoNode,
-            MinimaxImageToVideoNode,
-            # MinimaxSubjectToVideoNode,
-            MinimaxHailuoVideoNode,
-        ]
+# A dictionary that contains all nodes you want to export with their names
+# NOTE: names should be globally unique
+NODE_CLASS_MAPPINGS = {
+    "MinimaxTextToVideoNode": MinimaxTextToVideoNode,
+    "MinimaxImageToVideoNode": MinimaxImageToVideoNode,
+    # "MinimaxSubjectToVideoNode": MinimaxSubjectToVideoNode,
+    "MinimaxHailuoVideoNode": MinimaxHailuoVideoNode,
+}
 
-
-async def comfy_entrypoint() -> MinimaxExtension:
-    return MinimaxExtension()
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MinimaxTextToVideoNode": "MiniMax Text to Video",
+    "MinimaxImageToVideoNode": "MiniMax Image to Video",
+    "MinimaxSubjectToVideoNode": "MiniMax Subject to Video",
+    "MinimaxHailuoVideoNode": "MiniMax Hailuo Video",
+}

comfy_api_nodes/nodes_moonvalley.py

@@ -1,7 +1,6 @@
 import logging
 from typing import Any, Callable, Optional, TypeVar
 import torch
-from typing_extensions import override
 from comfy_api_nodes.util.validation_utils import (
     get_image_dimensions,
     validate_image_dimensions,
@@ -27,9 +26,11 @@ from comfy_api_nodes.apinode_utils import (
     upload_images_to_comfyapi,
     upload_video_to_comfyapi,
 )
+from comfy_api_nodes.mapper_utils import model_field_to_node_input
 
-from comfy_api.input import VideoInput
-from comfy_api.latest import ComfyExtension, InputImpl, io as comfy_io
+from comfy_api.input.video_types import VideoInput
+from comfy.comfy_types.node_typing import IO
+from comfy_api.input_impl import VideoFromFile
 import av
 import io
 
@@ -361,7 +362,7 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
 
         # Return as VideoFromFile using the buffer
         output_buffer.seek(0)
-        return InputImpl.VideoFromFile(output_buffer)
+        return VideoFromFile(output_buffer)
 
     except Exception as e:
         # Clean up on error
@@ -372,150 +373,166 @@ def trim_video(video: VideoInput, duration_sec: float) -> VideoInput:
         raise RuntimeError(f"Failed to trim video: {str(e)}") from e
 
 
-def parse_width_height_from_res(resolution: str):
-    # Accepts a string like "16:9 (1920 x 1080)" and returns width, height as a dict
-    res_map = {
-        "16:9 (1920 x 1080)": {"width": 1920, "height": 1080},
-        "9:16 (1080 x 1920)": {"width": 1080, "height": 1920},
-        "1:1 (1152 x 1152)": {"width": 1152, "height": 1152},
-        "4:3 (1536 x 1152)": {"width": 1536, "height": 1152},
-        "3:4 (1152 x 1536)": {"width": 1152, "height": 1536},
-        "21:9 (2560 x 1080)": {"width": 2560, "height": 1080},
-    }
-    return res_map.get(resolution, {"width": 1920, "height": 1080})
+# --- BaseMoonvalleyVideoNode ---
+class BaseMoonvalleyVideoNode:
+    def parseWidthHeightFromRes(self, resolution: str):
+        # Accepts a string like "16:9 (1920 x 1080)" and returns width, height as a dict
+        res_map = {
+            "16:9 (1920 x 1080)": {"width": 1920, "height": 1080},
+            "9:16 (1080 x 1920)": {"width": 1080, "height": 1920},
+            "1:1 (1152 x 1152)": {"width": 1152, "height": 1152},
+            "4:3 (1536 x 1152)": {"width": 1536, "height": 1152},
+            "3:4 (1152 x 1536)": {"width": 1152, "height": 1536},
+            "21:9 (2560 x 1080)": {"width": 2560, "height": 1080},
+        }
+        if resolution in res_map:
+            return res_map[resolution]
+        else:
+            # Default to 1920x1080 if unknown
+            return {"width": 1920, "height": 1080}
 
+    def parseControlParameter(self, value):
+        control_map = {
+            "Motion Transfer": "motion_control",
+            "Canny": "canny_control",
+            "Pose Transfer": "pose_control",
+            "Depth": "depth_control",
+        }
+        if value in control_map:
+            return control_map[value]
+        else:
+            return control_map["Motion Transfer"]
 
-def parse_control_parameter(value):
-    control_map = {
-        "Motion Transfer": "motion_control",
-        "Canny": "canny_control",
-        "Pose Transfer": "pose_control",
-        "Depth": "depth_control",
-    }
-    return control_map.get(value, control_map["Motion Transfer"])
-
-
-async def get_response(
-    task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None
-) -> MoonvalleyPromptResponse:
-    return await poll_until_finished(
-        auth_kwargs,
-        ApiEndpoint(
-            path=f"{API_PROMPTS_ENDPOINT}/{task_id}",
-            method=HttpMethod.GET,
-            request_model=EmptyRequest,
-            response_model=MoonvalleyPromptResponse,
-        ),
-        result_url_extractor=get_video_url_from_response,
-        node_id=node_id,
-    )
-
-
-class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
+    async def get_response(
+        self, task_id: str, auth_kwargs: dict[str, str], node_id: Optional[str] = None
+    ) -> MoonvalleyPromptResponse:
+        return await poll_until_finished(
+            auth_kwargs,
+            ApiEndpoint(
+                path=f"{API_PROMPTS_ENDPOINT}/{task_id}",
+                method=HttpMethod.GET,
+                request_model=EmptyRequest,
+                response_model=MoonvalleyPromptResponse,
+            ),
+            result_url_extractor=get_video_url_from_response,
+            node_id=node_id,
+        )
 
     @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MoonvalleyImg2VideoNode",
-            display_name="Moonvalley Marey Image to Video",
-            category="api node/video/Moonvalley Marey",
-            description="Moonvalley Marey Image to Video Node",
-            inputs=[
-                comfy_io.Image.Input(
-                    "image",
-                    tooltip="The reference image used to generate the video",
-                ),
-                comfy_io.String.Input(
-                    "prompt",
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "prompt": model_field_to_node_input(
+                    IO.STRING,
+                    MoonvalleyTextToVideoRequest,
+                    "prompt_text",
                     multiline=True,
                 ),
-                comfy_io.String.Input(
+                "negative_prompt": model_field_to_node_input(
+                    IO.STRING,
+                    MoonvalleyTextToVideoInferenceParams,
                     "negative_prompt",
                     multiline=True,
-                    default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
-                            "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
-                            "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
-                            "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
-                            "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
-                            "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
-                    tooltip="Negative prompt text",
+                    default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, wobbly, weird, low quality, plastic, stock footage, video camera, boring",
                 ),
-                comfy_io.Combo.Input(
-                    "resolution",
-                    options=[
-                        "16:9 (1920 x 1080)",
-                        "9:16 (1080 x 1920)",
-                        "1:1 (1152 x 1152)",
-                        "4:3 (1536 x 1152)",
-                        "3:4 (1152 x 1536)",
-                        "21:9 (2560 x 1080)",
-                    ],
-                    default="16:9 (1920 x 1080)",
-                    tooltip="Resolution of the output video",
+                "resolution": (
+                    IO.COMBO,
+                    {
+                        "options": [
+                            "16:9 (1920 x 1080)",
+                            "9:16 (1080 x 1920)",
+                            "1:1 (1152 x 1152)",
+                            "4:3 (1440 x 1080)",
+                            "3:4 (1080 x 1440)",
+                            "21:9 (2560 x 1080)",
+                        ],
+                        "default": "16:9 (1920 x 1080)",
+                        "tooltip": "Resolution of the output video",
+                    },
                 ),
-                comfy_io.Float.Input(
-                    "prompt_adherence",
+                "prompt_adherence": model_field_to_node_input(
+                    IO.FLOAT,
+                    MoonvalleyTextToVideoInferenceParams,
+                    "guidance_scale",
                     default=10.0,
-                    min=1.0,
-                    max=20.0,
-                    step=1.0,
-                    tooltip="Guidance scale for generation control",
+                    step=1,
+                    min=1,
+                    max=20,
                 ),
-                comfy_io.Int.Input(
+                "seed": model_field_to_node_input(
+                    IO.INT,
+                    MoonvalleyTextToVideoInferenceParams,
                     "seed",
                     default=9,
                     min=0,
                     max=4294967295,
                     step=1,
-                    display_mode=comfy_io.NumberDisplay.number,
+                    display="number",
                     tooltip="Random seed value",
                 ),
-                comfy_io.Int.Input(
+                "steps": model_field_to_node_input(
+                    IO.INT,
+                    MoonvalleyTextToVideoInferenceParams,
                     "steps",
                     default=100,
                     min=1,
                     max=100,
-                    step=1,
-                    tooltip="Number of denoising steps",
                 ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+            "optional": {
+                "image": model_field_to_node_input(
+                    IO.IMAGE,
+                    MoonvalleyTextToVideoRequest,
+                    "image_url",
+                    tooltip="The reference image used to generate the video",
+                ),
+            },
+        }
+
+    RETURN_TYPES = ("STRING",)
+    FUNCTION = "generate"
+    CATEGORY = "api node/video/Moonvalley Marey"
+    API_NODE = True
+
+    def generate(self, **kwargs):
+        return None
+
+
+# --- MoonvalleyImg2VideoNode ---
+class MoonvalleyImg2VideoNode(BaseMoonvalleyVideoNode):
 
     @classmethod
-    async def execute(
-        cls,
-        image: torch.Tensor,
-        prompt: str,
-        negative_prompt: str,
-        resolution: str,
-        prompt_adherence: float,
-        seed: int,
-        steps: int,
-    ) -> comfy_io.NodeOutput:
+    def INPUT_TYPES(cls):
+        return super().INPUT_TYPES()
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+    DESCRIPTION = "Moonvalley Marey Image to Video Node"
+
+    async def generate(
+        self, prompt, negative_prompt, unique_id: Optional[str] = None, **kwargs
+    ):
+        image = kwargs.get("image", None)
+        if image is None:
+            raise MoonvalleyApiError("image is required")
+
         validate_input_image(image, True)
         validate_prompts(prompt, negative_prompt, MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
-        width_height = parse_width_height_from_res(resolution)
-
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
+        width_height = self.parseWidthHeightFromRes(kwargs.get("resolution"))
 
         inference_params = MoonvalleyTextToVideoInferenceParams(
             negative_prompt=negative_prompt,
-            steps=steps,
-            seed=seed,
-            guidance_scale=prompt_adherence,
+            steps=kwargs.get("steps"),
+            seed=kwargs.get("seed"),
+            guidance_scale=kwargs.get("prompt_adherence"),
             num_frames=128,
-            width=width_height["width"],
-            height=width_height["height"],
+            width=width_height.get("width"),
+            height=width_height.get("height"),
             use_negative_prompts=True,
         )
         """Upload image to comfy backend to have a URL available for further processing"""
@@ -524,7 +541,7 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
 
         image_url = (
             await upload_images_to_comfyapi(
-                image, max_images=1, auth_kwargs=auth, mime_type=mime_type
+                image, max_images=1, auth_kwargs=kwargs, mime_type=mime_type
             )
         )[0]
 
@@ -539,102 +556,127 @@ class MoonvalleyImg2VideoNode(comfy_io.ComfyNode):
                 response_model=MoonvalleyPromptResponse,
             ),
             request=request,
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         task_creation_response = await initial_operation.execute()
         validate_task_creation_response(task_creation_response)
         task_id = task_creation_response.id
 
-        final_response = await get_response(
-            task_id, auth_kwargs=auth, node_id=cls.hidden.unique_id
+        final_response = await self.get_response(
+            task_id, auth_kwargs=kwargs, node_id=unique_id
         )
         video = await download_url_to_video_output(final_response.output_url)
-        return comfy_io.NodeOutput(video)
+        return (video,)
 
 
-class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
+# --- MoonvalleyVid2VidNode ---
+class MoonvalleyVideo2VideoNode(BaseMoonvalleyVideoNode):
+    def __init__(self):
+        super().__init__()
 
     @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MoonvalleyVideo2VideoNode",
-            display_name="Moonvalley Marey Video to Video",
-            category="api node/video/Moonvalley Marey",
-            description="",
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt",
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "prompt": model_field_to_node_input(
+                    IO.STRING,
+                    MoonvalleyVideoToVideoRequest,
+                    "prompt_text",
                     multiline=True,
-                    tooltip="Describes the video to generate",
                 ),
-                comfy_io.String.Input(
+                "negative_prompt": model_field_to_node_input(
+                    IO.STRING,
+                    MoonvalleyVideoToVideoInferenceParams,
                     "negative_prompt",
                     multiline=True,
-                    default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
-                            "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
-                            "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
-                            "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
-                            "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
-                            "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
-                    tooltip="Negative prompt text",
+                    default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, wobbly, weird, low quality, plastic, stock footage, video camera, boring",
                 ),
-                comfy_io.Int.Input(
+                "seed": model_field_to_node_input(
+                    IO.INT,
+                    MoonvalleyVideoToVideoInferenceParams,
                     "seed",
                     default=9,
                     min=0,
                     max=4294967295,
                     step=1,
-                    display_mode=comfy_io.NumberDisplay.number,
+                    display="number",
                     tooltip="Random seed value",
                     control_after_generate=False,
                 ),
-                comfy_io.Video.Input(
-                    "video",
-                    tooltip="The reference video used to generate the output video. Must be at least 5 seconds long. "
-                            "Videos longer than 5s will be automatically trimmed. Only MP4 format supported.",
-                ),
-                comfy_io.Combo.Input(
-                    "control_type",
-                    options=["Motion Transfer", "Pose Transfer"],
-                    default="Motion Transfer",
-                    optional=True,
-                ),
-                comfy_io.Int.Input(
-                    "motion_intensity",
-                    default=100,
-                    min=0,
-                    max=100,
+                "prompt_adherence": model_field_to_node_input(
+                    IO.FLOAT,
+                    MoonvalleyVideoToVideoInferenceParams,
+                    "guidance_scale",
+                    default=10.0,
                     step=1,
-                    tooltip="Only used if control_type is 'Motion Transfer'",
-                    optional=True,
+                    min=1,
+                    max=20,
                 ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
-
-    @classmethod
-    async def execute(
-        cls,
-        prompt: str,
-        negative_prompt: str,
-        seed: int,
-        video: Optional[VideoInput] = None,
-        control_type: str = "Motion Transfer",
-        motion_intensity: Optional[int] = 100,
-    ) -> comfy_io.NodeOutput:
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
+            },
+            "hidden": {
+                "auth_token": "AUTH_TOKEN_COMFY_ORG",
+                "comfy_api_key": "API_KEY_COMFY_ORG",
+                "unique_id": "UNIQUE_ID",
+            },
+            "optional": {
+                "video": (
+                    IO.VIDEO,
+                    {
+                        "default": "",
+                        "multiline": False,
+                        "tooltip": "The reference video used to generate the output video. Must be at least 5 seconds long. Videos longer than 5s will be automatically trimmed. Only MP4 format supported.",
+                    },
+                ),
+                "control_type": (
+                    ["Motion Transfer", "Pose Transfer"],
+                    {"default": "Motion Transfer"},
+                ),
+                "motion_intensity": (
+                    "INT",
+                    {
+                        "default": 100,
+                        "step": 1,
+                        "min": 0,
+                        "max": 100,
+                        "tooltip": "Only used if control_type is 'Motion Transfer'",
+                    },
+                ),
+                "image": model_field_to_node_input(
+                    IO.IMAGE,
+                    MoonvalleyTextToVideoRequest,
+                    "image_url",
+                    tooltip="The reference image used to generate the video",
+                ),
+            },
         }
 
-        validated_video = validate_video_to_video_input(video)
-        video_url = await upload_video_to_comfyapi(validated_video, auth_kwargs=auth)
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+
+    async def generate(
+        self, prompt, negative_prompt, unique_id: Optional[str] = None, **kwargs
+    ):
+        video = kwargs.get("video")
+        image = kwargs.get("image", None)
+
+        if not video:
+            raise MoonvalleyApiError("video is required")
+
+        video_url = ""
+        if video:
+            validated_video = validate_video_to_video_input(video)
+            video_url = await upload_video_to_comfyapi(
+                validated_video, auth_kwargs=kwargs
+            )
+        mime_type = "image/png"
+
+        if not image is None:
+            validate_input_image(image, with_frame_conditioning=True)
+            image_url = await upload_images_to_comfyapi(
+                image=image, auth_kwargs=kwargs, max_images=1, mime_type=mime_type
+            )
+        control_type = kwargs.get("control_type")
+        motion_intensity = kwargs.get("motion_intensity")
 
         """Validate prompts and inference input"""
         validate_prompts(prompt, negative_prompt)
@@ -646,11 +688,11 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
 
         inference_params = MoonvalleyVideoToVideoInferenceParams(
             negative_prompt=negative_prompt,
-            seed=seed,
+            seed=kwargs.get("seed"),
             control_params=control_params,
         )
 
-        control = parse_control_parameter(control_type)
+        control = self.parseControlParameter(control_type)
 
         request = MoonvalleyVideoToVideoRequest(
             control_type=control,
@@ -658,6 +700,7 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
             prompt_text=prompt,
             inference_params=inference_params,
         )
+        request.image_url = image_url if not image is None else None
 
         initial_operation = SynchronousOperation(
             endpoint=ApiEndpoint(
@@ -667,125 +710,58 @@ class MoonvalleyVideo2VideoNode(comfy_io.ComfyNode):
                 response_model=MoonvalleyPromptResponse,
             ),
             request=request,
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
         task_creation_response = await initial_operation.execute()
         validate_task_creation_response(task_creation_response)
         task_id = task_creation_response.id
 
-        final_response = await get_response(
-            task_id, auth_kwargs=auth, node_id=cls.hidden.unique_id
+        final_response = await self.get_response(
+            task_id, auth_kwargs=kwargs, node_id=unique_id
         )
 
         video = await download_url_to_video_output(final_response.output_url)
-        return comfy_io.NodeOutput(video)
+
+        return (video,)
 
 
-class MoonvalleyTxt2VideoNode(comfy_io.ComfyNode):
+# --- MoonvalleyTxt2VideoNode ---
+class MoonvalleyTxt2VideoNode(BaseMoonvalleyVideoNode):
+    def __init__(self):
+        super().__init__()
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
 
     @classmethod
-    def define_schema(cls) -> comfy_io.Schema:
-        return comfy_io.Schema(
-            node_id="MoonvalleyTxt2VideoNode",
-            display_name="Moonvalley Marey Text to Video",
-            category="api node/video/Moonvalley Marey",
-            description="",
-            inputs=[
-                comfy_io.String.Input(
-                    "prompt",
-                    multiline=True,
-                ),
-                comfy_io.String.Input(
-                    "negative_prompt",
-                    multiline=True,
-                    default="<synthetic> <scene cut> gopro, bright, contrast, static, overexposed, vignette, "
-                            "artifacts, still, noise, texture, scanlines, videogame, 360 camera, VR, transition, "
-                            "flare, saturation, distorted, warped, wide angle, saturated, vibrant, glowing, "
-                            "cross dissolve, cheesy, ugly hands, mutated hands, mutant, disfigured, extra fingers, "
-                            "blown out, horrible, blurry, worst quality, bad, dissolve, melt, fade in, fade out, "
-                            "wobbly, weird, low quality, plastic, stock footage, video camera, boring",
-                    tooltip="Negative prompt text",
-                ),
-                comfy_io.Combo.Input(
-                    "resolution",
-                    options=[
-                        "16:9 (1920 x 1080)",
-                        "9:16 (1080 x 1920)",
-                        "1:1 (1152 x 1152)",
-                        "4:3 (1536 x 1152)",
-                        "3:4 (1152 x 1536)",
-                        "21:9 (2560 x 1080)",
-                    ],
-                    default="16:9 (1920 x 1080)",
-                    tooltip="Resolution of the output video",
-                ),
-                comfy_io.Float.Input(
-                    "prompt_adherence",
-                    default=10.0,
-                    min=1.0,
-                    max=20.0,
-                    step=1.0,
-                    tooltip="Guidance scale for generation control",
-                ),
-                comfy_io.Int.Input(
-                    "seed",
-                    default=9,
-                    min=0,
-                    max=4294967295,
-                    step=1,
-                    display_mode=comfy_io.NumberDisplay.number,
-                    tooltip="Random seed value",
-                ),
-                comfy_io.Int.Input(
-                    "steps",
-                    default=100,
-                    min=1,
-                    max=100,
-                    step=1,
-                    tooltip="Inference steps",
-                ),
-            ],
-            outputs=[comfy_io.Video.Output()],
-            hidden=[
-                comfy_io.Hidden.auth_token_comfy_org,
-                comfy_io.Hidden.api_key_comfy_org,
-                comfy_io.Hidden.unique_id,
-            ],
-            is_api_node=True,
-        )
+    def INPUT_TYPES(cls):
+        input_types = super().INPUT_TYPES()
+        # Remove image-specific parameters
+        for param in ["image"]:
+            if param in input_types["optional"]:
+                del input_types["optional"][param]
+        return input_types
 
-    @classmethod
-    async def execute(
-        cls,
-        prompt: str,
-        negative_prompt: str,
-        resolution: str,
-        prompt_adherence: float,
-        seed: int,
-        steps: int,
-    ) -> comfy_io.NodeOutput:
+    async def generate(
+        self, prompt, negative_prompt, unique_id: Optional[str] = None, **kwargs
+    ):
         validate_prompts(prompt, negative_prompt, MOONVALLEY_MAREY_MAX_PROMPT_LENGTH)
-        width_height = parse_width_height_from_res(resolution)
-
-        auth = {
-            "auth_token": cls.hidden.auth_token_comfy_org,
-            "comfy_api_key": cls.hidden.api_key_comfy_org,
-        }
+        width_height = self.parseWidthHeightFromRes(kwargs.get("resolution"))
 
         inference_params = MoonvalleyTextToVideoInferenceParams(
             negative_prompt=negative_prompt,
-            steps=steps,
-            seed=seed,
-            guidance_scale=prompt_adherence,
+            steps=kwargs.get("steps"),
+            seed=kwargs.get("seed"),
+            guidance_scale=kwargs.get("prompt_adherence"),
             num_frames=128,
-            width=width_height["width"],
-            height=width_height["height"],
+            width=width_height.get("width"),
+            height=width_height.get("height"),
         )
         request = MoonvalleyTextToVideoRequest(
             prompt_text=prompt, inference_params=inference_params
         )
 
-        init_op = SynchronousOperation(
+        initial_operation = SynchronousOperation(
             endpoint=ApiEndpoint(
                 path=API_TXT2VIDEO_ENDPOINT,
                 method=HttpMethod.POST,
@@ -793,29 +769,29 @@ class MoonvalleyTxt2VideoNode(comfy_io.ComfyNode):
                 response_model=MoonvalleyPromptResponse,
             ),
             request=request,
-            auth_kwargs=auth,
+            auth_kwargs=kwargs,
         )
-        task_creation_response = await init_op.execute()
+        task_creation_response = await initial_operation.execute()
         validate_task_creation_response(task_creation_response)
         task_id = task_creation_response.id
 
-        final_response = await get_response(
-            task_id, auth_kwargs=auth, node_id=cls.hidden.unique_id
+        final_response = await self.get_response(
+            task_id, auth_kwargs=kwargs, node_id=unique_id
         )
 
         video = await download_url_to_video_output(final_response.output_url)
-        return comfy_io.NodeOutput(video)
+        return (video,)
 
 
-class MoonvalleyExtension(ComfyExtension):
-    @override
-    async def get_node_list(self) -> list[type[comfy_io.ComfyNode]]:
-        return [
-            MoonvalleyImg2VideoNode,
-            MoonvalleyTxt2VideoNode,
-            MoonvalleyVideo2VideoNode,
-        ]
+NODE_CLASS_MAPPINGS = {
+    "MoonvalleyImg2VideoNode": MoonvalleyImg2VideoNode,
+    "MoonvalleyTxt2VideoNode": MoonvalleyTxt2VideoNode,
+    "MoonvalleyVideo2VideoNode": MoonvalleyVideo2VideoNode,
+}
 
 
-async def comfy_entrypoint() -> MoonvalleyExtension:
-    return MoonvalleyExtension()
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MoonvalleyImg2VideoNode": "Moonvalley Marey Image to Video",
+    "MoonvalleyTxt2VideoNode": "Moonvalley Marey Text to Video",
+    "MoonvalleyVideo2VideoNode": "Moonvalley Marey Video to Video",
+}