Improve logging for failed download links.

I get freezing issues on my test machine.

.github/workflows/cla.yml

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

alembic_db/versions/0005_allow_case_sensitive_tags.py

@@ -0,0 +1,107 @@
+"""
+Allow case-sensitive tag names.
+
+Revision ID: 0005_allow_case_sensitive_tags
+Revises: 0004_drop_tag_type
+Create Date: 2026-06-16
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+revision = "0005_allow_case_sensitive_tags"
+down_revision = "0004_drop_tag_type"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    bind = op.get_bind()
+    if bind.dialect.name == "sqlite":
+        # SQLite cannot ALTER/DROP CHECK constraints. Recreate the small tag
+        # vocabulary table without the lowercase constraint while preserving
+        # existing tag names.
+        op.execute("PRAGMA foreign_keys=OFF")
+        try:
+            op.execute(
+                "CREATE TABLE tags_new ("
+                "name VARCHAR(512) NOT NULL, "
+                "CONSTRAINT pk_tags PRIMARY KEY (name)"
+                ")"
+            )
+            op.execute("INSERT INTO tags_new(name) SELECT name FROM tags")
+            op.execute("DROP TABLE tags")
+            op.execute("ALTER TABLE tags_new RENAME TO tags")
+        finally:
+            op.execute("PRAGMA foreign_keys=ON")
+        return
+
+    op.drop_constraint("ck_tags_ck_tags_lowercase", "tags", type_="check")
+
+
+def downgrade() -> None:
+    # Existing mixed-case tags cannot satisfy the old constraint. Lowercase them
+    # before restoring it, merging duplicate vocabulary/link rows that collide.
+    bind = op.get_bind()
+
+    tag_names = [row[0] for row in bind.execute(sa.text("SELECT name FROM tags"))]
+    existing_names = set(tag_names)
+    lowercase_names = sorted({name.lower() for name in tag_names})
+    missing_lowercase_rows = [
+        {"name": name} for name in lowercase_names if name not in existing_names
+    ]
+    if missing_lowercase_rows:
+        bind.execute(sa.text("INSERT INTO tags(name) VALUES (:name)"), missing_lowercase_rows)
+
+    link_rows = bind.execute(
+        sa.text(
+            "SELECT asset_reference_id, tag_name, origin, added_at "
+            "FROM asset_reference_tags "
+            "ORDER BY asset_reference_id, tag_name"
+        )
+    ).mappings()
+    deduped_links = {}
+    for row in link_rows:
+        key = (row["asset_reference_id"], row["tag_name"].lower())
+        deduped_links.setdefault(
+            key,
+            {
+                "asset_reference_id": row["asset_reference_id"],
+                "tag_name": row["tag_name"].lower(),
+                "origin": row["origin"],
+                "added_at": row["added_at"],
+            },
+        )
+
+    op.execute("DELETE FROM asset_reference_tags")
+    if deduped_links:
+        bind.execute(
+            sa.text(
+                "INSERT INTO asset_reference_tags "
+                "(asset_reference_id, tag_name, origin, added_at) "
+                "VALUES (:asset_reference_id, :tag_name, :origin, :added_at)"
+            ),
+            list(deduped_links.values()),
+        )
+    op.execute("DELETE FROM tags WHERE name != lower(name)")
+
+    if bind.dialect.name == "sqlite":
+        op.execute("PRAGMA foreign_keys=OFF")
+        try:
+            op.execute(
+                "CREATE TABLE tags_new ("
+                "name VARCHAR(512) NOT NULL, "
+                "CONSTRAINT pk_tags PRIMARY KEY (name), "
+                "CONSTRAINT ck_tags_lowercase CHECK (name = lower(name))"
+                ")"
+            )
+            op.execute("INSERT INTO tags_new(name) SELECT name FROM tags")
+            op.execute("DROP TABLE tags")
+            op.execute("ALTER TABLE tags_new RENAME TO tags")
+        finally:
+            op.execute("PRAGMA foreign_keys=ON")
+        return
+
+    op.create_check_constraint(
+        "ck_tags_ck_tags_lowercase", "tags", "name = lower(name)"
+    )

alembic_db/versions/0005_download_manager.py

@@ -0,0 +1,84 @@
+"""
+Download manager schema.
+
+Adds the two tables that back the server-side model download manager:
+transient job/queue state (``downloads`` + per-segment ``download_segments``).
+
+Revision ID: 0005_download_manager
+Revises: 0004_drop_tag_type
+Create Date: 2026-06-27
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+revision = "0005_download_manager"
+down_revision = "0004_drop_tag_type"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.create_table(
+        "downloads",
+        sa.Column("id", sa.String(length=36), primary_key=True),
+        sa.Column("url", sa.Text(), nullable=False),
+        sa.Column("final_url", sa.Text(), nullable=True),
+        sa.Column("model_id", sa.String(length=1024), nullable=False),
+        sa.Column("dest_path", sa.Text(), nullable=False),
+        sa.Column("temp_path", sa.Text(), nullable=False),
+        sa.Column("status", sa.String(length=16), nullable=False),
+        sa.Column("priority", sa.Integer(), nullable=False, server_default="0"),
+        sa.Column("total_bytes", sa.BigInteger(), nullable=True),
+        sa.Column("bytes_done", sa.BigInteger(), nullable=False, server_default="0"),
+        sa.Column("etag", sa.String(length=512), nullable=True),
+        sa.Column("last_modified", sa.String(length=128), nullable=True),
+        sa.Column(
+            "accept_ranges", sa.Boolean(), nullable=False, server_default=sa.text("false")
+        ),
+        sa.Column("expected_sha256", sa.String(length=64), nullable=True),
+        sa.Column(
+            "allow_any_extension",
+            sa.Boolean(),
+            nullable=False,
+            server_default=sa.text("false"),
+        ),
+        sa.Column("attempts", sa.Integer(), nullable=False, server_default="0"),
+        sa.Column("error", sa.Text(), nullable=True),
+        sa.Column("created_at", sa.BigInteger(), nullable=False),
+        sa.Column("updated_at", sa.BigInteger(), nullable=False),
+        sa.CheckConstraint("bytes_done >= 0", name="ck_downloads_bytes_done_nonneg"),
+        sa.CheckConstraint(
+            "total_bytes IS NULL OR total_bytes >= 0",
+            name="ck_downloads_total_bytes_nonneg",
+        ),
+    )
+    op.create_index("ix_downloads_status", "downloads", ["status"])
+    op.create_index("ix_downloads_priority", "downloads", ["priority"])
+    op.create_index("ix_downloads_model_id", "downloads", ["model_id"])
+
+    op.create_table(
+        "download_segments",
+        sa.Column(
+            "download_id",
+            sa.String(length=36),
+            sa.ForeignKey("downloads.id", ondelete="CASCADE"),
+            nullable=False,
+        ),
+        sa.Column("idx", sa.Integer(), nullable=False),
+        sa.Column("start_offset", sa.BigInteger(), nullable=False),
+        sa.Column("end_offset", sa.BigInteger(), nullable=False),
+        sa.Column("bytes_done", sa.BigInteger(), nullable=False, server_default="0"),
+        sa.PrimaryKeyConstraint("download_id", "idx", name="pk_download_segments"),
+        sa.CheckConstraint("bytes_done >= 0", name="ck_segments_bytes_done_nonneg"),
+        sa.CheckConstraint("end_offset >= start_offset", name="ck_segments_range"),
+    )
+
+
+def downgrade() -> None:
+    op.drop_table("download_segments")
+
+    op.drop_index("ix_downloads_model_id", table_name="downloads")
+    op.drop_index("ix_downloads_priority", table_name="downloads")
+    op.drop_index("ix_downloads_status", table_name="downloads")
+    op.drop_table("downloads")

alembic_db/versions/0006_add_loader_path.py

@@ -0,0 +1,30 @@
+"""
+Add loader_path column to asset_references.
+
+Stores the in-root loader path (path relative to the storage root with the
+top-level model category dropped) derived from file_path at scan/ingest time,
+so the assets API can return it without re-resolving against every registered
+model-folder base on every request.
+
+Revision ID: 0006_add_loader_path
+Revises: 0005_allow_case_sensitive_tags
+Create Date: 2026-07-02
+"""
+
+from alembic import op
+import sqlalchemy as sa
+
+revision = "0006_add_loader_path"
+down_revision = "0005_allow_case_sensitive_tags"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    with op.batch_alter_table("asset_references") as batch_op:
+        batch_op.add_column(sa.Column("loader_path", sa.Text(), nullable=True))
+
+
+def downgrade() -> None:
+    with op.batch_alter_table("asset_references") as batch_op:
+        batch_op.drop_column("loader_path")

app/assets/api/routes.py

@@ -40,6 +40,7 @@ from app.assets.services import (
     upload_from_temp_path,
 )
 from app.assets.services.cursor import InvalidCursorError
+from app.assets.services.path_utils import compute_display_name
 from app.assets.services.tagging import list_tag_histogram
 
 ROUTES = web.RouteTableDef()
@@ -161,11 +162,19 @@ def _build_asset_response(result: schemas.AssetDetailResult | schemas.UploadResu
             preview_url = None
     else:
         preview_url = _build_preview_url_from_view(result.tags, result.ref.user_metadata)
+    if result.ref.file_path:
+        display_name = compute_display_name(result.ref.file_path)
+        # In-root loader path (model category dropped): what model loaders consume.
+        loader_path = result.ref.loader_path
+    else:
+        display_name, loader_path = None, None
     asset_content_hash = result.asset.hash if result.asset else None
     return schemas_out.Asset(
         id=result.ref.id,
         name=result.ref.name,
         hash=asset_content_hash,
+        loader_path=loader_path,
+        display_name=display_name,
         asset_hash=asset_content_hash,
         size=int(result.asset.size_bytes) if result.asset else None,
         mime_type=result.asset.mime_type if result.asset else None,
@@ -419,17 +428,6 @@ async def upload_asset(request: web.Request) -> web.Response:
             400, "INVALID_BODY", f"Validation failed: {ve.json()}"
         )
 
-    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
-        ):
-            delete_temp_file_if_exists(parsed.tmp_path)
-            category = spec.tags[1] if len(spec.tags) >= 2 else ""
-            return _build_error_response(
-                400, "INVALID_BODY", f"unknown models category '{category}'"
-            )
-
     try:
         # Fast path: hash exists, create AssetReference without writing anything
         if spec.hash and parsed.provided_hash_exists is True:
@@ -473,7 +471,7 @@ async def upload_asset(request: web.Request) -> web.Response:
         return _build_error_response(400, e.code, str(e))
     except ValueError as e:
         delete_temp_file_if_exists(parsed.tmp_path)
-        return _build_error_response(400, "BAD_REQUEST", str(e))
+        return _build_error_response(400, "INVALID_BODY", str(e))
     except HashMismatchError as e:
         delete_temp_file_if_exists(parsed.tmp_path)
         return _build_error_response(400, "HASH_MISMATCH", str(e))

app/assets/api/schemas_in.py

@@ -140,7 +140,7 @@ class CreateFromHashBody(BaseModel):
         if v is None:
             return []
         if isinstance(v, list):
-            out = [str(t).strip().lower() for t in v if str(t).strip()]
+            out = [str(t).strip() for t in v if str(t).strip()]
             seen = set()
             dedup = []
             for t in out:
@@ -149,7 +149,7 @@ class CreateFromHashBody(BaseModel):
                     dedup.append(t)
             return dedup
         if isinstance(v, str):
-            return [t.strip().lower() for t in v.split(",") if t.strip()]
+            return list(dict.fromkeys(t.strip() for t in v.split(",") if t.strip()))
         return []
 
 
@@ -206,7 +206,7 @@ class TagsListQuery(BaseModel):
         if v is None:
             return v
         v = v.strip()
-        return v.lower() or None
+        return v or None
 
 
 class TagsAdd(BaseModel):
@@ -220,7 +220,7 @@ class TagsAdd(BaseModel):
         for t in v:
             if not isinstance(t, str):
                 raise TypeError("tags must be strings")
-            tnorm = t.strip().lower()
+            tnorm = t.strip()
             if tnorm:
                 out.append(tnorm)
         seen = set()
@@ -239,8 +239,8 @@ class TagsRemove(TagsAdd):
 class UploadAssetSpec(BaseModel):
     """Upload Asset operation.
 
-    - tags: optional list; if provided, first is root ('models'|'input'|'output');
-            if root == 'models', second must be a valid category
+    - tags: labels plus one destination role ('models'|'input'|'output') for new bytes;
+            if role == 'models', exactly one model_type:<folder_name> tag is required
     - name: display name
     - user_metadata: arbitrary JSON object (optional)
     - hash: optional canonical 'blake3:<hex>' for validation / fast-path
@@ -309,7 +309,7 @@ class UploadAssetSpec(BaseModel):
         norm = []
         seen = set()
         for t in items:
-            tnorm = str(t).strip().lower()
+            tnorm = str(t).strip()
             if tnorm and tnorm not in seen:
                 seen.add(tnorm)
                 norm.append(tnorm)
@@ -335,14 +335,4 @@ class UploadAssetSpec(BaseModel):
 
     @model_validator(mode="after")
     def _validate_order(self):
-        if not self.tags:
-            raise ValueError("at least one tag is required for uploads")
-        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

app/assets/api/schemas_out.py

@@ -9,8 +9,20 @@ class Asset(BaseModel):
     ``id`` here is the AssetReference id, not the content-addressed Asset id."""
 
     id: str
-    name: str
+    name: str = Field(
+        ...,
+        deprecated=True,
+        description="Reference label, often caller-provided or derived from the filename. Deprecated for storage path/display semantics; use `loader_path` and `display_name` when present.",
+    )
     hash: str | None = None
+    loader_path: str | None = Field(
+        default=None,
+        description="The value a loader consumes to load this asset. `None` when no loader can resolve the file.",
+    )
+    display_name: str | None = Field(
+        default=None,
+        description="Human-facing label for the asset. Not unique.",
+    )
     asset_hash: str | None = None
     size: int | None = None
     mime_type: str | None = None

app/assets/api/upload.py

@@ -140,7 +140,6 @@ async def parse_multipart_upload(
             provided_mime_type = ((await field.text()) or "").strip() or None
         elif fname == "preview_id":
             provided_preview_id = ((await field.text()) or "").strip() or None
-
     if not file_present and not (provided_hash and provided_hash_exists):
         raise UploadError(
             400, "MISSING_FILE", "Form must include a 'file' part or a known 'hash'."

app/assets/database/models.py

@@ -76,6 +76,8 @@ class AssetReference(Base):
 
     # Cache state fields (from former AssetCacheState)
     file_path: Mapped[str | None] = mapped_column(Text, nullable=True)
+    # In-root loader path derived from file_path at scan/ingest time.
+    loader_path: Mapped[str | None] = mapped_column(Text, nullable=True)
     mtime_ns: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
     needs_verify: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
     is_missing: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)

app/assets/database/queries/asset_reference.py

@@ -650,6 +650,7 @@ def upsert_reference(
     name: str,
     mtime_ns: int,
     owner_id: str = "",
+    loader_path: str | None = None,
 ) -> tuple[bool, bool]:
     """Upsert a reference by file_path. Returns (created, updated).
 
@@ -659,6 +660,7 @@ def upsert_reference(
     vals = {
         "asset_id": asset_id,
         "file_path": file_path,
+        "loader_path": loader_path,
         "name": name,
         "owner_id": owner_id,
         "mtime_ns": int(mtime_ns),
@@ -686,13 +688,14 @@ def upsert_reference(
                 AssetReference.asset_id != asset_id,
                 AssetReference.mtime_ns.is_(None),
                 AssetReference.mtime_ns != int(mtime_ns),
+                AssetReference.loader_path.is_distinct_from(loader_path),
                 AssetReference.is_missing == True,  # noqa: E712
                 AssetReference.deleted_at.isnot(None),
             )
         )
         .values(
-            asset_id=asset_id, mtime_ns=int(mtime_ns), is_missing=False,
-            deleted_at=None, updated_at=now,
+            asset_id=asset_id, mtime_ns=int(mtime_ns), loader_path=loader_path,
+            is_missing=False, deleted_at=None, updated_at=now,
         )
     )
     res2 = session.execute(upd)

app/assets/database/queries/tags.py

@@ -265,6 +265,8 @@ def list_tags_with_usage(
     order: str = "count_desc",
     owner_id: str = "",
 ) -> tuple[list[tuple[str, str, int]], int]:
+    prefix_filter = prefix.strip() if prefix else ""
+
     counts_sq = (
         select(
             AssetReferenceTag.tag_name.label("tag_name"),
@@ -293,9 +295,8 @@ def list_tags_with_usage(
         .join(counts_sq, counts_sq.c.tag_name == Tag.name, isouter=True)
     )
 
-    if prefix:
-        escaped, esc = escape_sql_like_string(prefix.strip().lower())
-        q = q.where(Tag.name.like(escaped + "%", escape=esc))
+    if prefix_filter:
+        q = q.where(func.substr(Tag.name, 1, len(prefix_filter)) == prefix_filter)
 
     if not include_zero:
         q = q.where(func.coalesce(counts_sq.c.cnt, 0) > 0)
@@ -306,9 +307,8 @@ def list_tags_with_usage(
         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_sql_like_string(prefix.strip().lower())
-        total_q = total_q.where(Tag.name.like(escaped + "%", escape=esc))
+    if prefix_filter:
+        total_q = total_q.where(func.substr(Tag.name, 1, len(prefix_filter)) == prefix_filter)
     if not include_zero:
         visible_tags_sq = (
             select(AssetReferenceTag.tag_name)

app/assets/helpers.py

@@ -41,10 +41,10 @@ def get_utc_now() -> datetime:
 def normalize_tags(tags: list[str] | None) -> list[str]:
     """
     Normalize a list of tags by:
-      - Stripping whitespace and converting to lowercase.
-      - Removing duplicates.
+      - Stripping whitespace.
+      - Removing exact duplicates while preserving order and case.
     """
-    return list(dict.fromkeys(t.strip().lower() for t in (tags or []) if (t or "").strip()))
+    return list(dict.fromkeys(t.strip() for t in (tags or []) if (t or "").strip()))
 
 
 def validate_blake3_hash(s: str) -> str:

app/assets/scanner.py

@@ -36,7 +36,7 @@ from app.assets.services.hashing import HashCheckpoint, compute_blake3_hash
 from app.assets.services.image_dimensions import extract_image_dimensions
 from app.assets.services.metadata_extract import extract_file_metadata
 from app.assets.services.path_utils import (
-    compute_relative_filename,
+    compute_loader_path,
     get_comfy_models_folders,
     get_name_and_tags_from_asset_path,
 )
@@ -63,7 +63,7 @@ RootType = Literal["models", "input", "output"]
 def get_prefixes_for_root(root: RootType) -> list[str]:
     if root == "models":
         bases: list[str] = []
-        for _bucket, paths in get_comfy_models_folders():
+        for _bucket, paths, _exts in get_comfy_models_folders():
             bases.extend(paths)
         return [os.path.abspath(p) for p in bases]
     if root == "input":
@@ -81,7 +81,7 @@ def get_all_known_prefixes() -> list[str]:
 
 def collect_models_files() -> list[str]:
     out: list[str] = []
-    for folder_name, bases in get_comfy_models_folders():
+    for folder_name, bases, _exts in get_comfy_models_folders():
         rel_files = folder_paths.get_filename_list(folder_name) or []
         for rel_path in rel_files:
             if not all(is_visible(part) for part in Path(rel_path).parts):
@@ -308,7 +308,7 @@ def build_asset_specs(
         if not stat_p.st_size:
             continue
         name, tags = get_name_and_tags_from_asset_path(abs_p)
-        rel_fname = compute_relative_filename(abs_p)
+        rel_fname = compute_loader_path(abs_p)
 
         # Extract metadata (tier 1: filesystem, tier 2: safetensors header)
         metadata = None
@@ -430,7 +430,7 @@ def enrich_asset(
         return new_level
 
     initial_mtime_ns = get_mtime_ns(stat_p)
-    rel_fname = compute_relative_filename(file_path)
+    rel_fname = compute_loader_path(file_path)
     mime_type: str | None = None
     metadata = None
 

app/assets/services/asset_management.py

@@ -38,7 +38,7 @@ from app.assets.database.queries import (
     update_reference_updated_at,
 )
 from app.assets.helpers import select_best_live_path
-from app.assets.services.path_utils import compute_relative_filename
+from app.assets.services.path_utils import compute_loader_path
 from app.assets.services.schemas import (
     AssetData,
     AssetDetailResult,
@@ -91,7 +91,7 @@ def update_asset_metadata(
             update_reference_name(session, reference_id=reference_id, name=name)
             touched = True
 
-        computed_filename = compute_relative_filename(ref.file_path) if ref.file_path else None
+        computed_filename = compute_loader_path(ref.file_path) if ref.file_path else None
 
         new_meta: dict | None = None
         if user_metadata is not None:

app/assets/services/bulk_ingest.py

@@ -56,6 +56,7 @@ class ReferenceRow(TypedDict):
     id: str
     asset_id: str
     file_path: str
+    loader_path: str | None
     mtime_ns: int
     owner_id: str
     name: str
@@ -134,6 +135,14 @@ def batch_insert_seed_assets(
 
     for spec in specs:
         absolute_path = os.path.abspath(spec["abs_path"])
+        existing_asset_id = path_to_asset_id.get(absolute_path)
+        if existing_asset_id is not None:
+            existing_tags = asset_id_to_ref_data[existing_asset_id]["tags"]
+            asset_id_to_ref_data[existing_asset_id]["tags"] = list(
+                dict.fromkeys([*existing_tags, *spec["tags"]])
+            )
+            continue
+
         asset_id = str(uuid.uuid4())
         reference_id = str(uuid.uuid4())
         absolute_path_list.append(absolute_path)
@@ -164,6 +173,8 @@ def batch_insert_seed_assets(
                 "id": reference_id,
                 "asset_id": asset_id,
                 "file_path": absolute_path,
+                # spec["fname"] is compute_loader_path(abs_path) from build_asset_specs.
+                "loader_path": spec["fname"],
                 "mtime_ns": spec["mtime_ns"],
                 "owner_id": owner_id,
                 "name": spec["info_name"],

app/assets/services/ingest.py

@@ -33,8 +33,9 @@ from app.assets.services.bulk_ingest import batch_insert_seed_assets
 from app.assets.services.file_utils import get_size_and_mtime_ns
 from app.assets.services.image_dimensions import extract_image_dimensions
 from app.assets.services.path_utils import (
-    compute_relative_filename,
+    compute_loader_path,
     get_name_and_tags_from_asset_path,
+    get_path_derived_tags_from_path,
     resolve_destination_from_tags,
     validate_path_within_base,
 )
@@ -91,6 +92,7 @@ def _ingest_file_from_path(
             name=info_name or os.path.basename(locator),
             mtime_ns=mtime_ns,
             owner_id=owner_id,
+            loader_path=compute_loader_path(locator),
         )
 
         # Get the reference we just created/updated
@@ -101,17 +103,32 @@ def _ingest_file_from_path(
             if preview_id and ref.preview_id != preview_id:
                 ref.preview_id = preview_id
 
-            norm = normalize_tags(list(tags))
-            if norm:
+            try:
+                backend_tags = get_path_derived_tags_from_path(locator)
+            except ValueError:
+                backend_tags = []
+            caller_tags = normalize_tags(tags)
+            backend_tags = normalize_tags(backend_tags)
+            all_tags = normalize_tags([*caller_tags, *backend_tags])
+            if all_tags:
                 if require_existing_tags:
-                    validate_tags_exist(session, norm)
-                add_tags_to_reference(
-                    session,
-                    reference_id=reference_id,
-                    tags=norm,
-                    origin=tag_origin,
-                    create_if_missing=not require_existing_tags,
-                )
+                    validate_tags_exist(session, all_tags)
+                if backend_tags:
+                    add_tags_to_reference(
+                        session,
+                        reference_id=reference_id,
+                        tags=backend_tags,
+                        origin="automatic",
+                        create_if_missing=not require_existing_tags,
+                    )
+                if caller_tags:
+                    add_tags_to_reference(
+                        session,
+                        reference_id=reference_id,
+                        tags=caller_tags,
+                        origin=tag_origin,
+                        create_if_missing=not require_existing_tags,
+                    )
 
             _update_metadata_with_filename(
                 session,
@@ -228,7 +245,7 @@ def ingest_existing_file(
             "mtime_ns": mtime_ns,
             "info_name": name,
             "tags": tags,
-            "fname": os.path.basename(abs_path),
+            "fname": compute_loader_path(abs_path),
             "metadata": None,
             "hash": None,
             "mime_type": mime_type,
@@ -288,7 +305,7 @@ def _register_existing_asset(
             return result
 
         new_meta = dict(user_metadata)
-        computed_filename = compute_relative_filename(ref.file_path) if ref.file_path else None
+        computed_filename = compute_loader_path(ref.file_path) if ref.file_path else None
         if computed_filename:
             new_meta["filename"] = computed_filename
 
@@ -335,7 +352,7 @@ def _update_metadata_with_filename(
     current_metadata: dict | None,
     user_metadata: dict[str, Any],
 ) -> None:
-    computed_filename = compute_relative_filename(file_path) if file_path else None
+    computed_filename = compute_loader_path(file_path) if file_path else None
 
     current_meta = current_metadata or {}
     new_meta = dict(current_meta)
@@ -474,6 +491,10 @@ def upload_from_temp_path(
         existing = get_asset_by_hash(session, asset_hash=asset_hash)
 
     if existing is not None:
+        # Once content is already known, duplicate byte uploads are treated as
+        # reference-only creation. Request tags are labels only here: do not
+        # require upload destination tags, do not move bytes, and do not
+        # synthesize path-derived classification or uploaded provenance.
         with contextlib.suppress(Exception):
             if temp_path and os.path.exists(temp_path):
                 os.remove(temp_path)
@@ -535,7 +556,7 @@ def upload_from_temp_path(
         owner_id=owner_id,
         preview_id=preview_id,
         user_metadata=user_metadata or {},
-        tags=tags,
+        tags=[*(tags or []), "uploaded"],
         tag_origin="manual",
         require_existing_tags=False,
     )
@@ -569,15 +590,19 @@ def register_file_in_place(
 ) -> UploadResult:
     """Register an already-saved file in the asset database without moving it.
 
-    Tags are derived from the filesystem path (root category + subfolder names),
-    merged with any caller-provided tags, matching the behavior of the scanner.
+    This helper is used by upload paths that have already written bytes before
+    registering the file, so it records the same ``uploaded`` tag as the
+    multipart byte-upload path.
+
+    Tags are derived from trusted filesystem classification and merged with any
+    caller-provided tags, matching the behavior of the scanner.
     If the path is not under a known root, only the caller-provided tags are used.
     """
     try:
         _, path_tags = get_name_and_tags_from_asset_path(abs_path)
     except ValueError:
         path_tags = []
-    merged_tags = normalize_tags([*path_tags, *tags])
+    merged_tags = normalize_tags([*path_tags, *tags, "uploaded"])
 
     try:
         digest, _ = hashing.compute_blake3_hash(abs_path)

app/assets/services/path_utils.py

@@ -3,59 +3,66 @@ from pathlib import Path
 from typing import Literal
 
 import folder_paths
-from app.assets.helpers import normalize_tags
 
 
-_NON_MODEL_FOLDER_NAMES = frozenset({"custom_nodes"})
+_NON_MODEL_FOLDER_NAMES = frozenset({"configs", "custom_nodes"})
+_KNOWN_SUBFOLDER_TAGS = frozenset({"3d", "pasted", "painter", "threed", "webcam"})
 
 
-def get_comfy_models_folders() -> list[tuple[str, list[str]]]:
-    """Build list of (folder_name, base_paths[]) for all model locations.
+def get_comfy_models_folders() -> list[tuple[str, list[str], set[str]]]:
+    """Build list of (folder_name, base_paths[], extensions) for all model locations.
 
     Includes every category registered in folder_names_and_paths,
     regardless of whether its paths are under the main models_dir,
-    but excludes non-model entries like custom_nodes.
+    but excludes non-model entries like configs and custom_nodes.
+
+    An empty extensions set means the category accepts any extension,
+    matching folder_paths.filter_files_extensions semantics.
     """
-    targets: list[tuple[str, list[str]]] = []
+    targets: list[tuple[str, list[str], set[str]]] = []
     for name, values in folder_paths.folder_names_and_paths.items():
         if name in _NON_MODEL_FOLDER_NAMES:
             continue
-        paths, _exts = values[0], values[1]
+        paths, exts = values[0], values[1]
         if paths:
-            targets.append((name, paths))
+            targets.append((name, paths, set(exts)))
     return targets
 
 
 def resolve_destination_from_tags(tags: list[str]) -> tuple[str, list[str]]:
-    """Validates and maps tags -> (base_dir, subdirs_for_fs)"""
-    if not tags:
-        raise ValueError("tags must not be empty")
-    root = tags[0].lower()
+    """Validates and maps upload routing tags -> (base_dir, subdirs_for_fs).
+
+    The request tags are only used to choose the write destination. Extra tags
+    remain labels; they do not become path components or trusted classification.
+    """
+    destination_roles = [t for t in tags if t in {"input", "models", "output"}]
+    if len(destination_roles) != 1:
+        raise ValueError("uploads require exactly one destination role: input, models, or output")
+
+    root = destination_roles[0]
     if root == "models":
-        if len(tags) < 2:
-            raise ValueError("at least two tags required for model asset")
+        model_type_tags = [t for t in tags if t.startswith("model_type:")]
+        if len(model_type_tags) != 1:
+            raise ValueError("models uploads require exactly one model_type:<folder_name> tag")
+        folder_name = model_type_tags[0].split(":", 1)[1]
+        if not folder_name:
+            raise ValueError("models uploads require exactly one model_type:<folder_name> tag")
+        model_folder_paths = {
+            name: paths for name, paths, _exts in get_comfy_models_folders()
+        }
         try:
-            bases = folder_paths.folder_names_and_paths[tags[1]][0]
+            bases = model_folder_paths[folder_name]
         except KeyError:
-            raise ValueError(f"unknown model category '{tags[1]}'")
+            raise ValueError(f"unknown model category '{folder_name}'")
         if not bases:
-            raise ValueError(f"no base path configured for category '{tags[1]}'")
+            raise ValueError(f"no base path configured for category '{folder_name}'")
         base_dir = os.path.abspath(bases[0])
-        raw_subdirs = tags[2:]
     elif root == "input":
         base_dir = os.path.abspath(folder_paths.get_input_directory())
-        raw_subdirs = tags[1:]
-    elif root == "output":
-        base_dir = os.path.abspath(folder_paths.get_output_directory())
-        raw_subdirs = tags[1:]
     else:
-        raise ValueError(f"unknown root tag '{tags[0]}'; expected 'models', 'input', or 'output'")
-    _sep_chars = frozenset(("/", "\\", os.sep))
-    for i in raw_subdirs:
-        if i in (".", "..") or _sep_chars & set(i):
-            raise ValueError("invalid path component in tags")
+        base_dir = os.path.abspath(folder_paths.get_output_directory())
 
-    return base_dir, raw_subdirs if raw_subdirs else []
+    return base_dir, []
 
 
 def validate_path_within_base(candidate: str, base: str) -> None:
@@ -65,14 +72,79 @@ def validate_path_within_base(candidate: str, base: str) -> None:
         raise ValueError("destination escapes base directory")
 
 
-def compute_relative_filename(file_path: str) -> str | None:
+def _compute_relative_path(child: str, parent: str) -> str:
+    rel = os.path.relpath(os.path.abspath(child), os.path.abspath(parent))
+    if rel == ".":
+        return ""
+    return rel.replace(os.sep, "/")
+
+
+def _is_relative_to(child: str, parent: str) -> bool:
+    return Path(os.path.abspath(child)).is_relative_to(os.path.abspath(parent))
+
+
+def compute_asset_response_paths(file_path: str) -> tuple[str, str | None] | None:
+    """Return (logical_path, display_name) for a file path.
+
+    ``logical_path`` is the internal namespaced storage locator (e.g.
+    ``models/checkpoints/foo/bar.safetensors``); ``display_name`` is the
+    human-facing label below that namespace, served on Asset responses. These
+    are storage locators, not model-loader namespaces. Registered model-folder
+    membership is represented by backend tags such as
+    ``model_type:<folder_name>``; these paths only use known storage roots.
     """
-    Return the model's path relative to the last well-known folder (the model category),
-    using forward slashes, eg:
+    fp_abs = os.path.abspath(file_path)
+    candidates: list[tuple[int, int, str, str]] = []
+
+    for order, (namespace, base) in enumerate(
+        (
+            ("input", folder_paths.get_input_directory()),
+            ("output", folder_paths.get_output_directory()),
+            ("temp", folder_paths.get_temp_directory()),
+            ("models", getattr(folder_paths, "models_dir", "")),
+        )
+    ):
+        if not base:
+            continue
+        base_abs = os.path.abspath(base)
+        if _is_relative_to(fp_abs, base_abs):
+            candidates.append((len(base_abs), -order, namespace, base_abs))
+
+    if not candidates:
+        return None
+
+    _base_len, _order, namespace, base = max(candidates)
+    rel = _compute_relative_path(fp_abs, base)
+    public_path = f"{namespace}/{rel}" if rel else namespace
+    return public_path, rel or None
+
+
+def compute_display_name(file_path: str) -> str | None:
+    """Return the asset's `display_name`, or None for unknown paths."""
+    result = compute_asset_response_paths(file_path)
+    return result[1] if result else None
+
+
+def compute_logical_path(file_path: str) -> str | None:
+    """Return the internal namespaced storage locator, or None for unknown paths."""
+    result = compute_asset_response_paths(file_path)
+    return result[0] if result else None
+
+
+def compute_loader_path(file_path: str) -> str | None:
+    """
+    Return the asset's in-root loader path: the 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.
+    This is the value model loaders consume (the model category is dropped). It
+    is persisted as ``AssetReference.loader_path`` and served as the public
+    Asset response `loader_path` field. The human-facing `display_name` comes
+    from compute_asset_response_paths().
+
+    For input/output/temp paths the full path relative to that root is returned.
+    For paths outside any known root, returns None.
     """
     try:
         root_category, rel_path = get_asset_category_and_relative_path(file_path)
@@ -116,9 +188,10 @@ def get_asset_category_and_relative_path(
     def _compute_relative(child: str, parent: str) -> str:
         # Normalize relative path, stripping any leading ".." components
         # by anchoring to root (os.sep) then computing relpath back from it.
-        return os.path.relpath(
+        rel = os.path.relpath(
             os.path.join(os.sep, os.path.relpath(child, parent)), os.sep
         )
+        return "" if rel == "." else rel.replace(os.sep, "/")
 
     # 1) input
     input_base = os.path.abspath(folder_paths.get_input_directory())
@@ -136,8 +209,14 @@ def get_asset_category_and_relative_path(
         return "temp", _compute_relative(fp_abs, temp_base)
 
     # 4) models (check deepest matching base to avoid ambiguity)
+    ext = os.path.splitext(fp_abs)[1].lower()
     best: tuple[int, str, str] | None = None  # (base_len, bucket, rel_inside_bucket)
-    for bucket, bases in get_comfy_models_folders():
+    for bucket, bases, extensions in get_comfy_models_folders():
+        # A bucket only lists files within its extension set (empty set
+        # accepts any extension), so a bucket that cannot load the file
+        # must not contribute a loader path.
+        if extensions and ext not in extensions:
+            continue
         for b in bases:
             base_abs = os.path.abspath(b)
             if not _check_is_within(fp_abs, base_abs):
@@ -149,25 +228,111 @@ def get_asset_category_and_relative_path(
     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)
+        normalized = os.path.relpath(os.path.join(os.sep, combined), os.sep)
+        return "models", normalized.replace(os.sep, "/")
 
     raise ValueError(
         f"Path is not within input, output, temp, or configured model bases: {file_path}"
     )
 
 
+def get_backend_system_tags_from_path(path: str) -> list[str]:
+    """Return trusted backend tags derived from current filesystem facts.
+
+    The returned tags are only the backend-generated system tags: ``models``,
+    ``model_type:<folder_name>``, ``input``, ``output``, and ``temp``. Model
+    type tags are based on registered folder names, not path components.
+
+    A ``model_type:<folder_name>`` tag is only emitted when the file's
+    extension is accepted by that folder's registered extension set, so
+    categories sharing a base directory tag only the files they can
+    actually load. Files under a model base whose extension matches no
+    category still get the ``models`` tag.
+    """
+    fp_abs = os.path.abspath(path)
+    fp_path = Path(fp_abs)
+    tags: list[str] = []
+
+    def _add(tag: str) -> None:
+        if tag not in tags:
+            tags.append(tag)
+
+    for role, base in (
+        ("input", folder_paths.get_input_directory()),
+        ("output", folder_paths.get_output_directory()),
+        ("temp", folder_paths.get_temp_directory()),
+    ):
+        if fp_path.is_relative_to(os.path.abspath(base)):
+            _add(role)
+
+    ext = os.path.splitext(fp_abs)[1].lower()
+    model_types: list[str] = []
+    under_models_base = False
+    for folder_name, bases, extensions in get_comfy_models_folders():
+        for base in bases:
+            if fp_path.is_relative_to(os.path.abspath(base)):
+                under_models_base = True
+                # Empty set accepts any extension, matching
+                # folder_paths.filter_files_extensions semantics.
+                if not extensions or ext in extensions:
+                    model_types.append(folder_name)
+                break
+
+    if under_models_base:
+        _add("models")
+    for folder_name in model_types:
+        _add(f"model_type:{folder_name}")
+
+    if not tags:
+        raise ValueError(
+            f"Path is not within input, output, temp, or configured model bases: {path}"
+        )
+    return tags
+
+
+def get_known_subfolder_tags(subfolder: str | None) -> list[str]:
+    """Return tags for known UI/input subfolder names."""
+    if subfolder in _KNOWN_SUBFOLDER_TAGS:
+        return [subfolder]
+    return []
+
+
+def get_known_input_subfolder_tags_from_path(path: str) -> list[str]:
+    """Return known input-layout tags for files in canonical input subfolders.
+
+    These are compatibility tags for current UI-origin input directories such as
+    ``pasted`` and ``webcam``. They are intentionally narrow: only files directly
+    inside a known top-level input directory receive the matching tag.
+    """
+    fp_abs = os.path.abspath(path)
+    input_base = os.path.abspath(folder_paths.get_input_directory())
+    if not Path(fp_abs).is_relative_to(input_base):
+        return []
+
+    rel = os.path.relpath(fp_abs, input_base)
+    parts = Path(rel).parts
+    if len(parts) == 2:
+        return get_known_subfolder_tags(parts[0])
+    return []
+
+
+def get_path_derived_tags_from_path(path: str) -> list[str]:
+    """Return all backend-derived tags for an asset path."""
+    tags = get_backend_system_tags_from_path(path)
+    for tag in get_known_input_subfolder_tags_from_path(path):
+        if tag not in tags:
+            tags.append(tag)
+    return tags
+
+
 def get_name_and_tags_from_asset_path(file_path: str) -> tuple[str, list[str]]:
     """Return (name, tags) derived from a filesystem path.
 
     - name: base filename with extension
-    - tags: [root_category] + parent folder names in order
+    - tags: backend-derived tags from root/model classification and known input
+      subfolder layout conventions
 
     Raises:
         ValueError: path does not belong to any known root.
     """
-    root_category, some_path = get_asset_category_and_relative_path(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])))
+    return Path(file_path).name, get_path_derived_tags_from_path(file_path)

app/assets/services/schemas.py

@@ -25,6 +25,7 @@ class ReferenceData:
     preview_id: str | None
     created_at: datetime
     updated_at: datetime
+    loader_path: str | None = None
     system_metadata: dict[str, Any] | None = None
     job_id: str | None = None
     last_access_time: datetime | None = None
@@ -93,6 +94,7 @@ def extract_reference_data(ref: AssetReference) -> ReferenceData:
         id=ref.id,
         name=ref.name,
         file_path=ref.file_path,
+        loader_path=ref.loader_path,
         user_metadata=ref.user_metadata,
         preview_id=ref.preview_id,
         system_metadata=ref.system_metadata,

app/database/db.py

@@ -4,7 +4,8 @@ import shutil
 from app.logger import log_startup_warning
 from utils.install_util import get_missing_requirements_message
 from filelock import FileLock, Timeout
-from comfy.cli_args import args
+# Import the module so tests that reload comfy.cli_args see the live object.
+import comfy.cli_args
 
 _DB_AVAILABLE = False
 Session = None
@@ -21,6 +22,7 @@ try:
 
     from app.database.models import Base
     import app.assets.database.models  # noqa: F401 — register models with Base.metadata
+    import app.model_downloader.database.models  # noqa: F401 — register models with Base.metadata
 
     _DB_AVAILABLE = True
 except ImportError as e:
@@ -57,13 +59,13 @@ def get_alembic_config():
 
     config = Config(config_path)
     config.set_main_option("script_location", scripts_path)
-    config.set_main_option("sqlalchemy.url", args.database_url)
+    config.set_main_option("sqlalchemy.url", comfy.cli_args.args.database_url)
 
     return config
 
 
 def get_db_path():
-    url = args.database_url
+    url = comfy.cli_args.args.database_url
     if url.startswith("sqlite:///"):
         return url.split("///")[1]
     else:
@@ -97,7 +99,7 @@ def _is_memory_db(db_url):
 
 
 def init_db():
-    db_url = args.database_url
+    db_url = comfy.cli_args.args.database_url
     logging.debug(f"Database URL: {db_url}")
 
     if _is_memory_db(db_url):

app/model_downloader/api/routes.py

@@ -0,0 +1,202 @@
+"""aiohttp routes for the download manager.
+
+Endpoint surface (all under ``/api/download``), mirroring the response
+envelope used by ``app/assets/api/routes.py``:
+
+  POST   /api/download/enqueue
+  GET    /api/download
+  POST   /api/download/availability
+  POST   /api/download/clear
+  GET    /api/download/auth
+  POST   /api/download/auth/{provider}/login
+  POST   /api/download/auth/{provider}/logout
+  GET    /api/download/{id}
+  DELETE /api/download/{id}
+  POST   /api/download/{id}/pause
+  POST   /api/download/{id}/resume
+  POST   /api/download/{id}/cancel
+  POST   /api/download/{id}/priority
+
+Note on ordering: the static ``auth`` routes are registered before the dynamic
+``/api/download/{id}`` route so a request to ``.../auth`` is not captured as
+``id == "auth"``.
+"""
+
+from __future__ import annotations
+
+import json
+
+from aiohttp import web
+from pydantic import BaseModel, ValidationError
+
+from app.model_downloader.api import schemas_in, schemas_out
+from app.model_downloader.auth.oauth import LoginInProgress, OAuthNotConfigured
+from app.model_downloader.auth.providers import PROVIDERS
+from app.model_downloader.auth.store import AUTH_STORE
+from app.model_downloader.manager import DOWNLOAD_MANAGER, DownloadError
+
+ROUTES = web.RouteTableDef()
+
+
+def register_routes(app: web.Application) -> None:
+    """Wire the download-manager routes into the running aiohttp app."""
+    app.add_routes(ROUTES)
+
+
+# ----- envelope helpers (same shape as app/assets/api/routes.py) -----
+
+
+def _error(status: int, code: str, message: str, details: dict | None = None) -> web.Response:
+    return web.json_response(
+        {"error": {"code": code, "message": message, "details": details or {}}},
+        status=status,
+    )
+
+
+def _ok(payload, status: int = 200) -> web.Response:
+    return web.json_response(payload, status=status)
+
+
+async def _parse(request: web.Request, model: type[BaseModel]):
+    try:
+        raw = await request.json()
+    except json.JSONDecodeError:
+        return _error(400, "INVALID_JSON", "Request body must be valid JSON.")
+    try:
+        return model.model_validate(raw)
+    except ValidationError as ve:
+        return _error(400, "INVALID_BODY", "Validation failed.", {"errors": json.loads(ve.json())})
+
+
+def _from_download_error(e: DownloadError) -> web.Response:
+    return _error(e.http_status, e.code, e.message)
+
+
+# ----- downloads: collection + enqueue + availability -----
+
+
+@ROUTES.post("/api/download/enqueue")
+async def enqueue(request: web.Request) -> web.Response:
+    parsed = await _parse(request, schemas_in.EnqueueRequest)
+    if isinstance(parsed, web.Response):
+        return parsed
+    try:
+        download_id = await DOWNLOAD_MANAGER.enqueue(
+            parsed.url,
+            parsed.model_id,
+            priority=parsed.priority,
+            expected_sha256=parsed.expected_sha256,
+            allow_any_extension=parsed.allow_any_extension,
+        )
+    except DownloadError as e:
+        return _from_download_error(e)
+    return _ok({"download_id": download_id, "accepted": True}, status=202)
+
+
+@ROUTES.get("/api/download")
+async def list_downloads(request: web.Request) -> web.Response:
+    return _ok({"downloads": await DOWNLOAD_MANAGER.list()})
+
+
+@ROUTES.post("/api/download/availability")
+async def availability(request: web.Request) -> web.Response:
+    parsed = await _parse(request, schemas_in.AvailabilityRequest)
+    if isinstance(parsed, web.Response):
+        return parsed
+    return _ok({"models": await DOWNLOAD_MANAGER.availability(parsed.models)})
+
+
+@ROUTES.post("/api/download/clear")
+async def clear(request: web.Request) -> web.Response:
+    deleted = await DOWNLOAD_MANAGER.clear()
+    return _ok({"deleted": deleted})
+
+
+# ----- auth (OAuth login + env-key status) — must precede /{id} -----
+
+
+@ROUTES.get("/api/download/auth")
+async def auth_status(request: web.Request) -> web.Response:
+    return _ok({"providers": schemas_out.auth_status()})
+
+
+@ROUTES.post("/api/download/auth/{provider}/login")
+async def auth_login(request: web.Request) -> web.Response:
+    provider = PROVIDERS.get(request.match_info["provider"])
+    if provider is None:
+        return _error(400, "UNKNOWN_PROVIDER", "No such auth provider.")
+    try:
+        authorize_url = await AUTH_STORE.begin_login(provider)
+    except OAuthNotConfigured as e:
+        return _error(400, "OAUTH_NOT_CONFIGURED", str(e))
+    except LoginInProgress as e:
+        return _error(409, "LOGIN_IN_PROGRESS", str(e))
+    return _ok({"authorize_url": authorize_url})
+
+
+@ROUTES.post("/api/download/auth/{provider}/logout")
+async def auth_logout(request: web.Request) -> web.Response:
+    provider = PROVIDERS.get(request.match_info["provider"])
+    if provider is None:
+        return _error(400, "UNKNOWN_PROVIDER", "No such auth provider.")
+    AUTH_STORE.clear(provider.name)
+    return _ok({"logged_out": True})
+
+
+# ----- single download by id (dynamic; registered last) -----
+
+
+@ROUTES.get("/api/download/{id}")
+async def get_download(request: web.Request) -> web.Response:
+    view = await DOWNLOAD_MANAGER.status(request.match_info["id"])
+    if view is None:
+        return _error(404, "NOT_FOUND", "No such download.")
+    return _ok(view)
+
+
+@ROUTES.delete("/api/download/{id}")
+async def delete_download(request: web.Request) -> web.Response:
+    try:
+        await DOWNLOAD_MANAGER.delete(request.match_info["id"])
+    except DownloadError as e:
+        return _from_download_error(e)
+    return _ok({"deleted": True})
+
+
+@ROUTES.post("/api/download/{id}/pause")
+async def pause(request: web.Request) -> web.Response:
+    try:
+        await DOWNLOAD_MANAGER.pause(request.match_info["id"])
+    except DownloadError as e:
+        return _from_download_error(e)
+    return _ok({"ok": True})
+
+
+@ROUTES.post("/api/download/{id}/resume")
+async def resume(request: web.Request) -> web.Response:
+    try:
+        await DOWNLOAD_MANAGER.resume(request.match_info["id"])
+    except DownloadError as e:
+        return _from_download_error(e)
+    return _ok({"ok": True})
+
+
+@ROUTES.post("/api/download/{id}/cancel")
+async def cancel(request: web.Request) -> web.Response:
+    try:
+        await DOWNLOAD_MANAGER.cancel(request.match_info["id"])
+    except DownloadError as e:
+        return _from_download_error(e)
+    return _ok({"ok": True})
+
+
+@ROUTES.post("/api/download/{id}/priority")
+async def set_priority(request: web.Request) -> web.Response:
+    parsed = await _parse(request, schemas_in.PriorityRequest)
+    if isinstance(parsed, web.Response):
+        return parsed
+    try:
+        await DOWNLOAD_MANAGER.set_priority(request.match_info["id"], parsed.priority)
+    except DownloadError as e:
+        return _from_download_error(e)
+    return _ok({"ok": True})

app/model_downloader/api/schemas_in.py

@@ -0,0 +1,46 @@
+"""Request schemas for the download manager API.
+
+Pydantic enforces shape at the boundary; handlers operate only on validated
+values past that point.
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+from pydantic import BaseModel, Field, field_validator
+
+
+class EnqueueRequest(BaseModel):
+    url: str
+    model_id: str
+    priority: int = 0
+    expected_sha256: Optional[str] = None
+    allow_any_extension: bool = False
+
+    @field_validator("url")
+    @classmethod
+    def _strip_url(cls, v: str) -> str:
+        return v.strip()
+
+
+class PriorityRequest(BaseModel):
+    priority: int
+
+
+class AvailabilityRequest(BaseModel):
+    """``{model_id: url}`` — the URLs declared in the workflow JSON."""
+
+    models: dict[str, str] = Field(default_factory=dict)
+
+    @field_validator("models")
+    @classmethod
+    def _strip_urls(cls, v: dict[str, str]) -> dict[str, str]:
+        return {k: url.strip() for k, url in v.items()}
+
+
+__all__ = [
+    "EnqueueRequest",
+    "PriorityRequest",
+    "AvailabilityRequest",
+]

app/model_downloader/api/schemas_out.py

@@ -0,0 +1,15 @@
+"""Response helpers for the download manager API.
+
+The download/status read models are plain dicts produced by the manager. This
+module serializes the per-provider auth status (never a token) for the API.
+"""
+
+from __future__ import annotations
+
+from app.model_downloader.auth.providers import PROVIDERS
+from app.model_downloader.auth.store import AUTH_STORE
+
+
+def auth_status() -> list[dict]:
+    """Per-provider auth status — never includes a token."""
+    return [AUTH_STORE.status(p) for p in PROVIDERS.values()]

app/model_downloader/auth/oauth.py

@@ -0,0 +1,269 @@
+"""Generic OAuth 2.0 PKCE engine + transient loopback callback server.
+
+The flow, per provider:
+
+1. :func:`start_login_flow` builds a PKCE challenge, binds a loopback callback
+   server on ``127.0.0.1:<CALLBACK_PORT>`` at ``/callback/<provider>``, and
+   returns the provider's authorize URL for the user to open.
+2. The provider redirects the browser back to the loopback URL with a ``code``
+   and the ``state`` we generated. The server validates ``state``, exchanges the
+   code for a :class:`Token`, hands it to the ``deliver`` sink, and tears down.
+3. If no callback arrives within :data:`_LOGIN_TIMEOUT`, the server tears down.
+
+The callback runs on its own bare server, not ComfyUI's main server: the main
+server rejects cross-site navigations (``Sec-Fetch-Site: cross-site`` → 403),
+and an OAuth redirect from the provider is exactly such a navigation. The port
+is fixed because HuggingFace and Civitai require an exact registered
+``redirect_uri`` (port included); only one login runs at a time so the port
+never contends with itself.
+
+Only public PKCE clients are supported (no client secret). All outbound calls
+go to the provider's own authorize/token endpoints, strictly user-initiated.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import base64
+import hashlib
+import logging
+import os
+import secrets
+import time
+from typing import Callable
+from urllib.parse import urlencode
+
+from aiohttp import web
+
+from app.model_downloader.auth.providers import Provider
+from app.model_downloader.auth.token_store import Token
+from app.model_downloader.net.session import get_session, ssl_context
+
+CALLBACK_HOST = "127.0.0.1"
+# Fixed loopback port for the OAuth redirect. Must match the redirect URI
+# registered on the provider's OAuth app; override in lockstep if you change it.
+CALLBACK_PORT = int(os.environ.get("COMFY_OAUTH_CALLBACK_PORT", "41954"))
+_LOGIN_TIMEOUT = 300.0  # seconds to wait for the browser callback
+
+# The auth tab is opened by the frontend via window.open, so window.close() is
+# allowed here; the visible text is the fallback when the browser blocks it.
+_SUCCESS_HTML = (
+    "<!doctype html><meta charset=utf-8><title>ComfyUI</title>"
+    "<p>Login successful. You can close this window and return to ComfyUI.</p>"
+    "<script>window.close()</script>"
+)
+
+# Token sink: called with the provider name and the exchanged Token.
+TokenSink = Callable[[str, Token], None]
+
+
+class OAuthError(Exception):
+    """A user-facing OAuth failure."""
+
+
+class OAuthNotConfigured(OAuthError):
+    """The provider has no public client id configured."""
+
+
+class LoginInProgress(OAuthError):
+    """A login flow for this provider is already running."""
+
+
+def _b64url(data: bytes) -> str:
+    return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
+
+
+def _make_pkce() -> tuple[str, str]:
+    """Return ``(verifier, challenge)`` for the S256 PKCE method."""
+    verifier = _b64url(secrets.token_bytes(32))
+    challenge = _b64url(hashlib.sha256(verifier.encode("ascii")).digest())
+    return verifier, challenge
+
+
+def build_authorize_url(
+    provider: Provider, challenge: str, state: str, redirect_uri: str
+) -> str:
+    params = {
+        "response_type": "code",
+        "client_id": provider.client_id,
+        "redirect_uri": redirect_uri,
+        "scope": provider.scope,
+        "state": state,
+        "code_challenge": challenge,
+        "code_challenge_method": "S256",
+    }
+    return f"{provider.authorize_url}?{urlencode(params)}"
+
+
+def _token_from_payload(payload: dict) -> Token:
+    expires_in = payload.get("expires_in")
+    expires_at = int(time.time()) + int(expires_in) if expires_in else 0
+    return Token(
+        access_token=payload["access_token"],
+        refresh_token=payload.get("refresh_token"),
+        expires_at=expires_at,
+        token_type=payload.get("token_type", "Bearer"),
+        scope=payload.get("scope"),
+    )
+
+
+async def _post_token(provider: Provider, data: dict) -> Token:
+    session = await get_session()
+    resp = await session.post(
+        provider.token_url,
+        data=data,
+        headers={"Accept": "application/json"},
+        ssl=ssl_context(),
+    )
+    try:
+        if resp.status != 200:
+            body = await resp.text()
+            raise OAuthError(
+                f"{provider.name} token endpoint returned HTTP {resp.status}: {body[:200]}"
+            )
+        payload = await resp.json()
+    finally:
+        await resp.release()
+    if "access_token" not in payload:
+        raise OAuthError(f"{provider.name} token response missing access_token")
+    return _token_from_payload(payload)
+
+
+async def exchange_code(
+    provider: Provider, code: str, verifier: str, redirect_uri: str
+) -> Token:
+    return await _post_token(
+        provider,
+        {
+            "grant_type": "authorization_code",
+            "code": code,
+            "redirect_uri": redirect_uri,
+            "client_id": provider.client_id,
+            "code_verifier": verifier,
+        },
+    )
+
+
+async def refresh_access_token(provider: Provider, token: Token) -> Token:
+    if not token.refresh_token:
+        raise OAuthError(f"{provider.name} token is not refreshable")
+    refreshed = await _post_token(
+        provider,
+        {
+            "grant_type": "refresh_token",
+            "refresh_token": token.refresh_token,
+            "client_id": provider.client_id,
+        },
+    )
+    # Some providers omit a new refresh token on refresh; keep the old one.
+    if refreshed.refresh_token is None:
+        refreshed.refresh_token = token.refresh_token
+    return refreshed
+
+
+class _LoginFlow:
+    """A single in-flight login: owns the loopback server and PKCE state."""
+
+    def __init__(self, provider: Provider, deliver: TokenSink) -> None:
+        self.provider = provider
+        self.deliver = deliver
+        self.verifier, self.challenge = _make_pkce()
+        self.state = secrets.token_urlsafe(24)
+        self.redirect_uri = f"http://{CALLBACK_HOST}:{CALLBACK_PORT}/callback/{provider.name}"
+        self._runner: web.AppRunner | None = None
+        self._timeout_handle: asyncio.TimerHandle | None = None
+
+    async def start(self) -> str:
+        app = web.Application()
+        app.router.add_get("/callback/{provider}", self._handle_callback)
+        self._runner = web.AppRunner(app)
+        await self._runner.setup()
+        site = web.TCPSite(self._runner, CALLBACK_HOST, CALLBACK_PORT, reuse_address=True)
+        try:
+            await site.start()
+        except OSError as e:
+            await self._runner.cleanup()
+            self._runner = None
+            raise OAuthError(f"could not bind callback port {CALLBACK_PORT}: {e}")
+        loop = asyncio.get_running_loop()
+        self._timeout_handle = loop.call_later(
+            _LOGIN_TIMEOUT, lambda: asyncio.ensure_future(self._teardown())
+        )
+        return build_authorize_url(
+            self.provider, self.challenge, self.state, self.redirect_uri
+        )
+
+    async def _handle_callback(self, request: web.Request) -> web.Response:
+        if request.match_info.get("provider") != self.provider.name:
+            return web.Response(text="Unknown login.", content_type="text/plain", status=404)
+        error = request.query.get("error")
+        if error:
+            asyncio.ensure_future(self._teardown())
+            return web.Response(
+                text=f"Login failed: {error}", content_type="text/plain", status=400
+            )
+        if request.query.get("state") != self.state:
+            return web.Response(
+                text="Login failed: state mismatch.",
+                content_type="text/plain",
+                status=400,
+            )
+        code = request.query.get("code")
+        if not code:
+            return web.Response(
+                text="Login failed: no authorization code.",
+                content_type="text/plain",
+                status=400,
+            )
+        try:
+            token = await exchange_code(
+                self.provider, code, self.verifier, self.redirect_uri
+            )
+        except OAuthError as e:
+            logging.warning("[model_downloader] %s login failed: %s", self.provider.name, e)
+            asyncio.ensure_future(self._teardown())
+            return web.Response(
+                text=f"Login failed: {e}", content_type="text/plain", status=502
+            )
+        self.deliver(self.provider.name, token)
+        asyncio.ensure_future(self._teardown())
+        return web.Response(text=_SUCCESS_HTML, content_type="text/html")
+
+    async def _teardown(self) -> None:
+        _ACTIVE.pop(self.provider.name, None)
+        if self._timeout_handle is not None:
+            self._timeout_handle.cancel()
+            self._timeout_handle = None
+        if self._runner is not None:
+            try:
+                await self._runner.cleanup()
+            except Exception:
+                logging.debug("[model_downloader] callback server cleanup error", exc_info=True)
+            self._runner = None
+
+
+_ACTIVE: dict[str, _LoginFlow] = {}
+
+
+def login_in_progress(provider_name: str) -> bool:
+    return provider_name in _ACTIVE
+
+
+async def start_login_flow(provider: Provider, deliver: TokenSink) -> str:
+    """Begin a login flow and return the authorize URL to open in a browser.
+
+    Binds the fixed-port loopback callback server; only one login may run at a
+    time since that port is shared.
+    """
+    if not provider.client_id:
+        raise OAuthNotConfigured(
+            f"OAuth app not configured for {provider.name}; set "
+            f"{provider.client_id_env} or use an env API key."
+        )
+    if _ACTIVE:
+        active = next(iter(_ACTIVE))
+        raise LoginInProgress(f"A login for {active} is already in progress.")
+    flow = _LoginFlow(provider, deliver)
+    authorize_url = await flow.start()
+    _ACTIVE[provider.name] = flow
+    return authorize_url

app/model_downloader/auth/providers.py

@@ -0,0 +1,87 @@
+"""Provider registry for download authentication.
+
+A :class:`Provider` describes a hub that can authenticate downloads either from
+an environment API key or from an OAuth 2.0 access token. Both HuggingFace and
+Civitai are public PKCE clients, so no client secret is ever stored; the public
+``client_id`` is a placeholder overridable via env.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from urllib.parse import urlsplit
+
+
+def normalize_host(host: str) -> str:
+    """Lowercase, strip port, IDNA-encode."""
+    if not host:
+        return ""
+    host = host.strip()
+    if "://" in host:  # a full URL was pasted — extract just the host
+        host = urlsplit(host).hostname or ""
+    host = host.lower()
+    if host.startswith("[") and "]" in host:  # bracketed IPv6 literal
+        host = host[1 : host.index("]")]
+    elif host.count(":") == 1:  # host:port (not IPv6)
+        host = host.split(":", 1)[0]
+    try:
+        host = host.encode("idna").decode("ascii")
+    except (UnicodeError, ValueError):
+        pass
+    return host
+
+
+@dataclass(frozen=True)
+class Provider:
+    name: str
+    host: str
+    authorize_url: str
+    token_url: str
+    scope: str
+    # Env vars to try, in order, for a plain API key.
+    env_keys: tuple[str, ...]
+    # Env var overriding the public OAuth client id.
+    client_id_env: str
+    # Public PKCE client id. Empty means "not configured" until the env sets it.
+    default_client_id: str = ""
+
+    @property
+    def client_id(self) -> str:
+        return os.environ.get(self.client_id_env, self.default_client_id) or ""
+
+    def env_token(self) -> str | None:
+        for var in self.env_keys:
+            token = os.environ.get(var)
+            if token:
+                return token
+        return None
+
+
+PROVIDERS: dict[str, Provider] = {
+    "huggingface": Provider(
+        name="huggingface",
+        host="huggingface.co",
+        authorize_url="https://huggingface.co/oauth/authorize",
+        token_url="https://huggingface.co/oauth/token",
+        scope="openid read-repos gated-repos",
+        env_keys=("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"),
+        client_id_env="COMFY_HF_OAUTH_CLIENT_ID",
+    ),
+    "civitai": Provider(
+        name="civitai",
+        host="civitai.com",
+        authorize_url="https://auth.civitai.com/api/auth/oauth/authorize",
+        token_url="https://auth.civitai.com/api/auth/oauth/token",
+        scope="4",  # ModelsRead; UserRead is auto-granted
+        env_keys=("CIVITAI_API_TOKEN", "CIVITAI_API_KEY"),
+        client_id_env="COMFY_CIVITAI_OAUTH_CLIENT_ID",
+    ),
+}
+
+_HOST_TO_PROVIDER = {p.host: p for p in PROVIDERS.values()}
+
+
+def provider_for_host(host: str) -> Provider | None:
+    """Return the provider whose host exactly matches ``host`` (normalized)."""
+    return _HOST_TO_PROVIDER.get(normalize_host(host))

app/model_downloader/auth/resolver.py

@@ -0,0 +1,42 @@
+"""Per-hop auth resolution (https only).
+
+Recomputed from scratch on every redirect hop: a hop only gets a bearer token
+when *its own host* matches a configured provider, so a token bound to
+``huggingface.co`` is silently dropped when the request is redirected to a
+presigned CDN host — which is exactly what these hubs expect.
+
+For a matching hop: env API key first, then the provider's OAuth access token
+(refreshed if expired), else no auth.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+from app.model_downloader.auth.providers import provider_for_host
+from app.model_downloader.auth.store import AUTH_STORE
+
+
+@dataclass
+class RequestAuth:
+    """How to modify a single request to carry a bearer token."""
+
+    headers: dict[str, str] = field(default_factory=dict)
+
+
+async def resolve_auth_for_hop(host: str, scheme: str) -> RequestAuth | None:
+    """Resolve the bearer token (if any) to attach for one request hop."""
+    if scheme.lower() != "https":
+        return None
+    provider = provider_for_host(host)
+    if provider is None:
+        return None
+
+    token = provider.env_token()
+    if token:
+        return RequestAuth(headers={"Authorization": f"Bearer {token}"})
+
+    access = await AUTH_STORE.get_valid_token(provider)
+    if access:
+        return RequestAuth(headers={"Authorization": f"Bearer {access}"})
+    return None

app/model_downloader/auth/store.py

@@ -0,0 +1,64 @@
+"""In-memory OAuth token cache over the on-disk token store.
+
+:data:`AUTH_STORE` is the process singleton the resolver and API talk to. It
+lazily loads each provider's token from disk, refreshes an expired access token
+via its refresh token, and orchestrates the login flow (delegating the loopback
+callback server to :mod:`oauth`).
+"""
+
+from __future__ import annotations
+
+from app.model_downloader.auth import oauth, token_store
+from app.model_downloader.auth.providers import Provider
+from app.model_downloader.auth.token_store import Token
+
+
+class AuthStore:
+    def __init__(self) -> None:
+        # provider name -> Token, or None when known to be absent. A missing key
+        # means "not yet loaded from disk".
+        self._cache: dict[str, Token | None] = {}
+
+    def _load(self, name: str) -> Token | None:
+        if name not in self._cache:
+            self._cache[name] = token_store.load(name)
+        return self._cache[name]
+
+    def set_token(self, name: str, token: Token) -> None:
+        self._cache[name] = token
+        token_store.save(name, token)
+
+    def clear(self, name: str) -> None:
+        self._cache[name] = None
+        token_store.delete(name)
+
+    async def get_valid_token(self, provider: Provider) -> str | None:
+        """Return a valid access token string for ``provider``, or ``None``.
+
+        Refreshes an expired token when a refresh token is available.
+        """
+        token = self._load(provider.name)
+        if token is None or not token.access_token:
+            return None
+        if token.is_expired():
+            if not token.refresh_token:
+                return None
+            token = await oauth.refresh_access_token(provider, token)
+            self.set_token(provider.name, token)
+        return token.access_token
+
+    async def begin_login(self, provider: Provider) -> str:
+        """Start a login flow; returns the authorize URL to open in a browser."""
+        return await oauth.start_login_flow(provider, self.set_token)
+
+    def status(self, provider: Provider) -> dict:
+        token = self._load(provider.name)
+        return {
+            "provider": provider.name,
+            "logged_in": token is not None and bool(token.access_token),
+            "login_in_progress": oauth.login_in_progress(provider.name),
+            "env_key_present": provider.env_token() is not None,
+        }
+
+
+AUTH_STORE = AuthStore()

app/model_downloader/auth/token_store.py

@@ -0,0 +1,186 @@
+"""On-disk OAuth token persistence — one machine-bound blob per provider.
+
+Tokens live under ``folder_paths.get_system_user_directory("download_auth")``,
+never in the SQLite DB. Each provider file is written ``0600`` and holds an
+opaque blob, not readable JSON: the token JSON is XORed with an HMAC-SHA256
+keystream whose key is derived from stable machine/install attributes plus a
+per-install random salt.
+
+This is obfuscation, not confidentiality. It stops a token from being read by a
+human browsing files, grepped out of a backup, or lifted from a folder copied to
+another machine (the blob won't decrypt off its origin machine). It does not
+protect against code running inside this process (custom nodes) or an attacker
+who reads this source and recomputes the key.
+"""
+
+from __future__ import annotations
+
+import base64
+import getpass
+import hashlib
+import hmac
+import json
+import os
+import platform
+import secrets
+import time
+from dataclasses import asdict, dataclass
+
+import folder_paths
+
+_SALT_FILE = ".salt"
+_SALT_LEN = 32
+_NONCE_LEN = 16
+_PBKDF2_ITERS = 200_000
+
+
+@dataclass
+class Token:
+    access_token: str
+    refresh_token: str | None = None
+    # Epoch seconds when the access token expires; 0 means "unknown / no expiry".
+    expires_at: int = 0
+    token_type: str = "Bearer"
+    scope: str | None = None
+
+    def is_expired(self, skew: int = 60) -> bool:
+        if not self.expires_at:
+            return False
+        return time.time() + skew >= self.expires_at
+
+
+def _auth_dir() -> str:
+    path = folder_paths.get_system_user_directory("download_auth")
+    os.makedirs(path, exist_ok=True)
+    return path
+
+
+def _token_path(provider: str) -> str:
+    return os.path.join(_auth_dir(), f"{provider}.bin")
+
+
+def _machine_id() -> bytes:
+    """A stable per-machine identifier, best-effort across platforms."""
+    for path in ("/etc/machine-id", "/var/lib/dbus/machine-id"):
+        try:
+            with open(path, "rb") as f:
+                return f.read().strip()
+        except OSError:
+            pass
+    if os.name == "nt":
+        import winreg
+
+        try:
+            key = winreg.OpenKey(
+                winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Cryptography"
+            )
+            try:
+                guid, _ = winreg.QueryValueEx(key, "MachineGuid")
+                return str(guid).encode("utf-8")
+            finally:
+                winreg.CloseKey(key)
+        except OSError:
+            pass
+    return platform.node().encode("utf-8")
+
+
+def _machine_material(auth_dir: str) -> bytes:
+    try:
+        user = getpass.getuser()
+    except Exception:
+        user = ""
+    parts = (_machine_id(), platform.node().encode("utf-8"), user.encode("utf-8"), auth_dir.encode("utf-8"))
+    return b"\x00".join(parts)
+
+
+def _load_or_create_salt(auth_dir: str) -> bytes | None:
+    path = os.path.join(auth_dir, _SALT_FILE)
+    try:
+        with open(path, "rb") as f:
+            salt = f.read()
+        if len(salt) == _SALT_LEN:
+            return salt
+    except FileNotFoundError:
+        pass
+    except OSError:
+        return None
+    salt = secrets.token_bytes(_SALT_LEN)
+    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "wb") as f:
+        f.write(salt)
+    os.chmod(path, 0o600)
+    return salt
+
+
+def _derive_key(salt: bytes, auth_dir: str) -> bytes:
+    return hashlib.pbkdf2_hmac(
+        "sha256", _machine_material(auth_dir), salt, _PBKDF2_ITERS, dklen=32
+    )
+
+
+def _keystream(key: bytes, nonce: bytes, n: int) -> bytes:
+    out = bytearray()
+    counter = 0
+    while len(out) < n:
+        out.extend(hmac.new(key, nonce + counter.to_bytes(8, "big"), hashlib.sha256).digest())
+        counter += 1
+    return bytes(out[:n])
+
+
+def _xor(data: bytes, stream: bytes) -> bytes:
+    return bytes(a ^ b for a, b in zip(data, stream))
+
+
+def load(provider: str) -> Token | None:
+    auth_dir = _auth_dir()
+    try:
+        with open(_token_path(provider), "rb") as f:
+            blob = base64.b64decode(f.read())
+    except FileNotFoundError:
+        return None
+    except (ValueError, OSError):
+        return None
+    salt = _load_or_create_salt(auth_dir)
+    if salt is None or len(blob) <= _NONCE_LEN:
+        return None
+    nonce, ciphertext = blob[:_NONCE_LEN], blob[_NONCE_LEN:]
+    key = _derive_key(salt, auth_dir)
+    plaintext = _xor(ciphertext, _keystream(key, nonce, len(ciphertext)))
+    # A wrong machine / corrupt file decrypts to garbage; treat as logged out.
+    try:
+        data = json.loads(plaintext)
+    except ValueError:
+        return None
+    if not isinstance(data, dict) or "access_token" not in data:
+        return None
+    return Token(
+        access_token=data.get("access_token", ""),
+        refresh_token=data.get("refresh_token"),
+        expires_at=int(data.get("expires_at", 0) or 0),
+        token_type=data.get("token_type", "Bearer"),
+        scope=data.get("scope"),
+    )
+
+
+def save(provider: str, token: Token) -> None:
+    auth_dir = _auth_dir()
+    salt = _load_or_create_salt(auth_dir)
+    if salt is None:
+        return
+    key = _derive_key(salt, auth_dir)
+    nonce = secrets.token_bytes(_NONCE_LEN)
+    plaintext = json.dumps(asdict(token)).encode("utf-8")
+    ciphertext = _xor(plaintext, _keystream(key, nonce, len(plaintext)))
+    blob = base64.b64encode(nonce + ciphertext)
+    path = _token_path(provider)
+    fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
+    with os.fdopen(fd, "wb") as f:
+        f.write(blob)
+    os.chmod(path, 0o600)
+
+
+def delete(provider: str) -> None:
+    try:
+        os.remove(_token_path(provider))
+    except FileNotFoundError:
+        pass

app/model_downloader/constants.py

@@ -0,0 +1,32 @@
+"""Shared constants for the download manager.
+
+Status values are persisted as TEXT in the ``downloads`` table; keep them
+stable. The lifecycle is:
+
+    queued -> active -> verifying -> completed
+       |        |-> paused -> (resume) -> active
+       |        |-> failed (network, retryable) -> queued (backoff)
+       |-> cancelled
+"""
+
+from __future__ import annotations
+
+
+class DownloadStatus:
+    QUEUED = "queued"
+    ACTIVE = "active"
+    PAUSED = "paused"
+    VERIFYING = "verifying"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+    #: States from which a worker is doing (or about to do) network I/O.
+    LIVE = (QUEUED, ACTIVE, VERIFYING)
+    #: Terminal states — the job will not transition again on its own.
+    TERMINAL = (COMPLETED, FAILED, CANCELLED)
+
+
+# Default temp-file suffix. Distinctive so the startup orphan sweep only
+# removes files THIS subsystem created, never unrelated *.tmp files.
+TMP_SUFFIX = ".comfy-download.part"

app/model_downloader/database/models.py

@@ -0,0 +1,125 @@
+"""SQLAlchemy models for the download manager.
+
+Two tables:
+
+- ``downloads``         one row per requested file (job + queue state).
+- ``download_segments`` per-segment byte progress, for segmented resume.
+
+On completion a finished file is registered into the assets catalog;
+``downloads`` is kept only as job history.
+"""
+
+from __future__ import annotations
+
+import time
+import uuid
+
+from sqlalchemy import (
+    BigInteger,
+    Boolean,
+    CheckConstraint,
+    ForeignKey,
+    Index,
+    Integer,
+    String,
+    Text,
+)
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from app.database.models import Base
+
+
+def _uuid() -> str:
+    return str(uuid.uuid4())
+
+
+def _now() -> int:
+    return int(time.time())
+
+
+class Download(Base):
+    __tablename__ = "downloads"
+
+    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=_uuid)
+    # Original requested URL and the final URL after validated redirects.
+    url: Mapped[str] = mapped_column(Text, nullable=False)
+    final_url: Mapped[str | None] = mapped_column(Text, nullable=True)
+    # Canonical "<directory>/<filename>" identifier (resolved via folder_paths).
+    model_id: Mapped[str] = mapped_column(String(1024), nullable=False)
+    # Final on-disk location and the .part write target.
+    dest_path: Mapped[str] = mapped_column(Text, nullable=False)
+    temp_path: Mapped[str] = mapped_column(Text, nullable=False)
+
+    status: Mapped[str] = mapped_column(String(16), nullable=False)
+    priority: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+
+    total_bytes: Mapped[int | None] = mapped_column(BigInteger, nullable=True)
+    bytes_done: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
+
+    etag: Mapped[str | None] = mapped_column(String(512), nullable=True)
+    last_modified: Mapped[str | None] = mapped_column(String(128), nullable=True)
+    accept_ranges: Mapped[bool] = mapped_column(Boolean, nullable=False, default=False)
+
+    # Optional hub-provided checksum to verify against (NOT the dedup key).
+    expected_sha256: Mapped[str | None] = mapped_column(String(64), nullable=True)
+
+    allow_any_extension: Mapped[bool] = mapped_column(
+        Boolean, nullable=False, default=False
+    )
+    # How many retryable failures we have seen (for backoff capping).
+    attempts: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
+
+    error: Mapped[str | None] = mapped_column(Text, nullable=True)
+    created_at: Mapped[int] = mapped_column(BigInteger, nullable=False, default=_now)
+    updated_at: Mapped[int] = mapped_column(
+        BigInteger, nullable=False, default=_now, onupdate=_now
+    )
+
+    segments: Mapped[list[DownloadSegment]] = relationship(
+        "DownloadSegment",
+        back_populates="download",
+        cascade="all,delete-orphan",
+        passive_deletes=True,
+        order_by="DownloadSegment.idx",
+    )
+
+    __table_args__ = (
+        Index("ix_downloads_status", "status"),
+        Index("ix_downloads_priority", "priority"),
+        Index("ix_downloads_model_id", "model_id"),
+        CheckConstraint("bytes_done >= 0", name="ck_downloads_bytes_done_nonneg"),
+        CheckConstraint(
+            "total_bytes IS NULL OR total_bytes >= 0",
+            name="ck_downloads_total_bytes_nonneg",
+        ),
+    )
+
+    def __repr__(self) -> str:
+        return f"<Download id={self.id} model_id={self.model_id!r} status={self.status}>"
+
+
+class DownloadSegment(Base):
+    __tablename__ = "download_segments"
+
+    download_id: Mapped[str] = mapped_column(
+        String(36),
+        ForeignKey("downloads.id", ondelete="CASCADE"),
+        primary_key=True,
+    )
+    idx: Mapped[int] = mapped_column(Integer, primary_key=True)
+    start_offset: Mapped[int] = mapped_column(BigInteger, nullable=False)
+    end_offset: Mapped[int] = mapped_column(BigInteger, nullable=False)
+    bytes_done: Mapped[int] = mapped_column(BigInteger, nullable=False, default=0)
+
+    download: Mapped[Download] = relationship("Download", back_populates="segments")
+
+    __table_args__ = (
+        CheckConstraint("bytes_done >= 0", name="ck_segments_bytes_done_nonneg"),
+        CheckConstraint("end_offset >= start_offset", name="ck_segments_range"),
+    )
+
+    def __repr__(self) -> str:
+        return (
+            f"<DownloadSegment {self.download_id}#{self.idx} "
+            f"{self.start_offset}-{self.end_offset} done={self.bytes_done}>"
+        )

app/model_downloader/database/queries.py

@@ -0,0 +1,177 @@
+"""Synchronous DB access for the download manager.
+
+All functions open their own short-lived session via ``create_session`` and
+commit before returning, mirroring ``app/assets`` usage. They are blocking
+(SQLite) and should be called from async code through ``asyncio.to_thread``.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Optional
+
+from sqlalchemy import delete, select
+
+from app.database.db import create_session
+from app.model_downloader.constants import DownloadStatus
+from app.model_downloader.database.models import Download, DownloadSegment
+
+
+# ----- downloads -----
+
+
+def insert_download(values: dict) -> None:
+    with create_session() as session:
+        session.add(Download(**values))
+        session.commit()
+
+
+def get_download(download_id: str) -> Optional[Download]:
+    with create_session() as session:
+        row = session.get(Download, download_id)
+        if row is not None:
+            session.expunge_all()
+        return row
+
+
+def list_downloads() -> list[Download]:
+    with create_session() as session:
+        rows = list(
+            session.execute(
+                select(Download).order_by(Download.created_at.desc())
+            ).scalars()
+        )
+        session.expunge_all()
+        return rows
+
+
+def has_live_download_for_model(
+    model_id: str, live_statuses: tuple[str, ...], exclude_id: Optional[str] = None
+) -> bool:
+    with create_session() as session:
+        stmt = select(Download.id).where(
+            Download.model_id == model_id,
+            Download.status.in_(live_statuses),
+        ).limit(1)
+        if exclude_id is not None:
+            stmt = stmt.where(Download.id != exclude_id)
+        return session.execute(stmt).first() is not None
+
+
+def list_segments(download_id: str) -> list[DownloadSegment]:
+    with create_session() as session:
+        rows = list(
+            session.execute(
+                select(DownloadSegment)
+                .where(DownloadSegment.download_id == download_id)
+                .order_by(DownloadSegment.idx)
+            ).scalars()
+        )
+        session.expunge_all()
+        return rows
+
+
+def update_download(download_id: str, **fields) -> None:
+    if not fields:
+        return
+    fields.setdefault("updated_at", int(time.time()))
+    with create_session() as session:
+        row = session.get(Download, download_id)
+        if row is None:
+            return
+        for key, value in fields.items():
+            setattr(row, key, value)
+        session.commit()
+
+
+def delete_download(download_id: str) -> None:
+    with create_session() as session:
+        row = session.get(Download, download_id)
+        if row is not None:
+            session.delete(row)
+            session.commit()
+
+
+def delete_downloads(download_ids: list[str]) -> int:
+    """Delete many downloads in one transaction; returns the number removed.
+
+    Uses a bulk ``DELETE ... WHERE id IN (...)``. Segment rows are removed by
+    the ``ON DELETE CASCADE`` foreign key (SQLite ``PRAGMA foreign_keys=ON`` is
+    set in ``app/database/db.py``), so this stays consistent without loading the
+    ORM relationship.
+    """
+    if not download_ids:
+        return 0
+    with create_session() as session:
+        result = session.execute(
+            delete(Download).where(Download.id.in_(download_ids))
+        )
+        session.commit()
+        return result.rowcount or 0
+
+
+def replace_segments(download_id: str, segments: list[dict]) -> None:
+    """Atomically replace the segment plan for a download."""
+    with create_session() as session:
+        session.query(DownloadSegment).filter(
+            DownloadSegment.download_id == download_id
+        ).delete()
+        for seg in segments:
+            session.add(DownloadSegment(download_id=download_id, **seg))
+        session.commit()
+
+
+def update_segment_progress(download_id: str, idx: int, bytes_done: int) -> None:
+    with create_session() as session:
+        row = session.get(DownloadSegment, {"download_id": download_id, "idx": idx})
+        if row is None:
+            return
+        row.bytes_done = bytes_done
+        session.commit()
+
+
+def list_queued_downloads() -> list[Download]:
+    """Queued rows ordered for admission (priority desc, then FIFO)."""
+    with create_session() as session:
+        rows = list(
+            session.execute(
+                select(Download)
+                .where(Download.status == DownloadStatus.QUEUED)
+                .order_by(Download.priority.desc(), Download.created_at.asc())
+            ).scalars()
+        )
+        session.expunge_all()
+        return rows
+
+
+def reconcile_live_downloads() -> list[Download]:
+    """Reset any ``active``/``verifying`` rows left by a previous run.
+
+    On a clean restart there can be no live worker, so anything still marked
+    live is stale. Move it back to ``queued`` (offsets are preserved on the
+    segment rows) so the scheduler re-admits it. Returns the rows that should
+    be re-queued by the scheduler (queued + paused).
+    """
+    with create_session() as session:
+        stale = list(
+            session.execute(
+                select(Download).where(
+                    Download.status.in_([DownloadStatus.ACTIVE, DownloadStatus.VERIFYING])
+                )
+            ).scalars()
+        )
+        now = int(time.time())
+        for row in stale:
+            row.status = DownloadStatus.QUEUED
+            row.updated_at = now
+        session.commit()
+
+        resumable = list(
+            session.execute(
+                select(Download)
+                .where(Download.status == DownloadStatus.QUEUED)
+                .order_by(Download.priority.desc(), Download.created_at.asc())
+            ).scalars()
+        )
+        session.expunge_all()
+        return resumable

app/model_downloader/engine/job.py

@@ -0,0 +1,611 @@
+"""The per-download worker.
+
+One :class:`DownloadJob` drives a single file from probe to verified, cataloged
+completion. It supports cooperative pause / resume / cancel, segmented
+multi-connection transfer with positioned writes, and a verification gate
+(size + structural + optional sha256) before the atomic rename into place.
+
+Control is cooperative: external callers flip ``_control`` via
+:meth:`request_pause` / :meth:`request_cancel`; segment loops observe it between
+chunks and raise, which unwinds cleanly and persists resume offsets.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Callable, Optional
+
+from comfy.cli_args import args
+from app.model_downloader.constants import DownloadStatus
+from app.model_downloader.database import queries
+from app.model_downloader.engine.planner import (
+    effective_segment_count,
+    plan_segments,
+)
+from app.model_downloader.engine.writer import FileWriter
+from app.model_downloader.net.http import open_validated, redact_url
+from app.model_downloader.net.probe import gated_error_message, probe
+from app.model_downloader.verify import checksum, dedup, structural
+
+_RETRYABLE_STATUSES = {408, 429, 500, 502, 503, 504}
+_PERSIST_INTERVAL = 2.0  # seconds between throttled progress persists
+
+
+class Paused(Exception):
+    pass
+
+
+class Cancelled(Exception):
+    pass
+
+
+class RemoteChanged(Exception):
+    """The remote file changed under a resume (got 200 where 206 expected)."""
+
+
+class RetryableError(Exception):
+    pass
+
+
+class FatalError(Exception):
+    """Non-retryable: 4xx, checksum mismatch, structural failure, gated, etc."""
+
+
+@dataclass
+class SegmentRuntime:
+    idx: int
+    start: int
+    end: int  # inclusive; may be -1 for unknown-size single stream
+    bytes_done: int = 0
+
+    @property
+    def length(self) -> int:
+        return self.end - self.start + 1
+
+
+@dataclass
+class RuntimeState:
+    download_id: str
+    model_id: str
+    url: str
+    priority: int
+    status: str
+    total_bytes: Optional[int] = None
+    bytes_done: int = 0
+    error: Optional[str] = None
+    segments: list[SegmentRuntime] = field(default_factory=list)
+    started_at: float = field(default_factory=time.monotonic)
+    _last_bytes: int = 0
+    _last_time: float = field(default_factory=time.monotonic)
+    speed_bps: float = 0.0
+
+    @property
+    def progress(self) -> Optional[float]:
+        if not self.total_bytes:
+            return None
+        return min(1.0, self.bytes_done / self.total_bytes)
+
+    @property
+    def eta_seconds(self) -> Optional[float]:
+        if not self.total_bytes or self.speed_bps <= 0:
+            return None
+        remaining = max(0, self.total_bytes - self.bytes_done)
+        return remaining / self.speed_bps
+
+
+@dataclass
+class JobSpec:
+    download_id: str
+    url: str
+    model_id: str
+    dest_path: str
+    temp_path: str
+    priority: int = 0
+    expected_sha256: Optional[str] = None
+    allow_any_extension: bool = False
+    etag: Optional[str] = None
+    attempts: int = 0
+
+
+class DownloadJob:
+    def __init__(
+        self, spec: JobSpec, notify_cb: Optional[Callable[[str], None]] = None
+    ) -> None:
+        self.spec = spec
+        self._notify = notify_cb
+        self._control = "run"  # run | pause | cancel
+        self.state = RuntimeState(
+            download_id=spec.download_id,
+            model_id=spec.model_id,
+            url=spec.url,
+            priority=spec.priority,
+            status=DownloadStatus.QUEUED,
+        )
+        self._writer: Optional[FileWriter] = None
+        self._etag: Optional[str] = spec.etag
+        self._last_persist = 0.0
+
+    # ----- external control -----
+
+    def request_pause(self) -> None:
+        if self._control == "run":
+            self._control = "pause"
+
+    def request_cancel(self) -> None:
+        self._control = "cancel"
+
+    def _check_control(self) -> None:
+        if self._control == "cancel":
+            raise Cancelled()
+        if self._control == "pause":
+            raise Paused()
+
+    # ----- lifecycle -----
+
+    async def run(self) -> str:
+        """Run to a terminal/paused state; returns the final status string."""
+        await self._set_status(DownloadStatus.ACTIVE, error=None)
+        try:
+            pr = await self._probe_and_plan()
+            await self._transfer(pr)
+            await self._finalize()
+            await self._set_status(DownloadStatus.COMPLETED)
+        except Paused:
+            await self._persist_progress(force=True)
+            await self._set_status(DownloadStatus.PAUSED)
+        except Cancelled:
+            await self._close_writer()
+            self._remove_temp()
+            await self._set_status(DownloadStatus.CANCELLED)
+        except RemoteChanged:
+            await self._reset_for_restart()
+            await self._set_status(
+                DownloadStatus.QUEUED, error="remote file changed; restarting"
+            )
+        except RetryableError as e:
+            await self._persist_progress(force=True)
+            await self._set_status(DownloadStatus.QUEUED, error=str(e))
+        except FatalError as e:
+            await self._close_writer()
+            self._remove_temp()
+            await self._set_status(DownloadStatus.FAILED, error=str(e))
+        except Exception as e:  # unexpected -> treat as retryable
+            logging.warning(
+                "[model_downloader] %s unexpected error: %s",
+                self.spec.model_id, e, exc_info=True,
+            )
+            await self._persist_progress(force=True)
+            await self._set_status(DownloadStatus.QUEUED, error=f"{type(e).__name__}: {e}")
+        finally:
+            await self._close_writer()
+        return self.state.status
+
+    # ----- probe + plan -----
+
+    async def _probe_and_plan(self):
+        pr = await probe(self.spec.url)
+        if not pr.ok:
+            if pr.gated:
+                raise FatalError(gated_error_message(self.spec.url, pr))
+            if pr.status == 0 or pr.status in _RETRYABLE_STATUSES:
+                raise RetryableError(pr.error or "probe failed")
+            raise FatalError(pr.error or f"probe returned HTTP {pr.status}")
+
+        max_bytes = self._max_download_bytes()
+        if max_bytes is not None and pr.total_bytes is not None and pr.total_bytes > max_bytes:
+            raise FatalError(
+                f"file size {pr.total_bytes} exceeds the maximum allowed "
+                f"download size {max_bytes} (--download-max-bytes)"
+            )
+
+        self._etag = pr.etag or self._etag
+        self.state.total_bytes = pr.total_bytes
+        await asyncio.to_thread(
+            queries.update_download,
+            self.spec.download_id,
+            final_url=pr.final_url,
+            total_bytes=pr.total_bytes,
+            accept_ranges=pr.accept_ranges,
+            etag=pr.etag,
+            last_modified=pr.last_modified,
+        )
+
+        seg_count = effective_segment_count(
+            pr.total_bytes, pr.accept_ranges, max(1, args.download_segments)
+        )
+        existing = await asyncio.to_thread(queries.list_segments, self.spec.download_id)
+        can_resume_segmented = (
+            seg_count > 1
+            and existing
+            and pr.total_bytes is not None
+            and existing[-1].end_offset == pr.total_bytes - 1
+        )
+        if can_resume_segmented and not self._segmented_part_valid(pr.total_bytes):
+            # The persisted per-segment offsets describe bytes in a preallocated
+            # .part that is now gone or the wrong size (e.g. the partial of a
+            # failed download was swept on restart, or removed by a fatal
+            # error). Trusting them would skip already-"complete" segments and
+            # leave zero-filled holes. Discard the offsets and re-plan fresh.
+            logging.info(
+                "[model_downloader] %s discarding segmented resume offsets "
+                "(preallocated .part missing or wrong size); restarting",
+                self.spec.model_id,
+            )
+            self._remove_temp()
+            await asyncio.to_thread(
+                queries.replace_segments, self.spec.download_id, []
+            )
+            await asyncio.to_thread(
+                queries.update_download, self.spec.download_id, bytes_done=0
+            )
+            existing = []
+            can_resume_segmented = False
+
+        if can_resume_segmented:
+            # Resume an existing segmented plan.
+            self.state.segments = [
+                SegmentRuntime(s.idx, s.start_offset, s.end_offset, s.bytes_done)
+                for s in existing
+            ]
+        elif seg_count > 1 and pr.total_bytes is not None:
+            plans = plan_segments(pr.total_bytes, seg_count)
+            await asyncio.to_thread(
+                queries.replace_segments,
+                self.spec.download_id,
+                [
+                    {"idx": p.idx, "start_offset": p.start, "end_offset": p.end, "bytes_done": 0}
+                    for p in plans
+                ],
+            )
+            self.state.segments = [SegmentRuntime(p.idx, p.start, p.end, 0) for p in plans]
+        else:
+            # Single-stream: one logical segment; bytes_done tracked on the row.
+            row = await asyncio.to_thread(queries.get_download, self.spec.download_id)
+            resume_from = row.bytes_done if row else 0
+            end = (pr.total_bytes - 1) if pr.total_bytes else -1
+            # ``row.bytes_done`` may be the SUM of per-segment offsets from a
+            # prior segmented run (a preallocated, non-contiguous .part). A
+            # single-stream resume writes a contiguous prefix, so the offset is
+            # only trustworthy when the on-disk file is exactly that many
+            # contiguous bytes. This guards the case where a download that ran
+            # segmented now resolves to one segment (server dropped
+            # Accept-Ranges, or --download-segments was lowered between runs):
+            # resuming over non-contiguous data would corrupt the output.
+            if resume_from > 0 and not self._contiguous_prefix_valid(resume_from):
+                logging.info(
+                    "[model_downloader] %s discarding untrusted resume offset "
+                    "%d (on-disk .part not a contiguous prefix); restarting",
+                    self.spec.model_id, resume_from,
+                )
+                resume_from = 0
+                self._remove_temp()
+                if await asyncio.to_thread(queries.list_segments, self.spec.download_id):
+                    await asyncio.to_thread(
+                        queries.replace_segments, self.spec.download_id, []
+                    )
+                await asyncio.to_thread(
+                    queries.update_download, self.spec.download_id, bytes_done=0
+                )
+            self.state.segments = [SegmentRuntime(0, 0, end, resume_from)]
+        self._recompute_bytes_done()
+        return pr
+
+    # ----- transfer -----
+
+    async def _transfer(self, pr) -> None:
+        self._writer = FileWriter(self.spec.temp_path)
+        await self._writer.open()
+
+        segmented = len(self.state.segments) > 1
+        if segmented and self.state.total_bytes:
+            await self._writer.preallocate(self.state.total_bytes)
+            await self._run_segmented()
+        else:
+            await self._run_single()
+
+        await self._writer.flush()
+
+    async def _run_segmented(self) -> None:
+        pending = [
+            asyncio.ensure_future(self._run_segment(seg))
+            for seg in self.state.segments
+            if seg.bytes_done < seg.length
+        ]
+        if not pending:
+            return
+        done, not_done = await asyncio.wait(
+            pending, return_when=asyncio.FIRST_EXCEPTION
+        )
+        first_exc: Optional[BaseException] = None
+        for task in done:
+            exc = task.exception()
+            if exc is not None and first_exc is None:
+                first_exc = exc
+        if first_exc is not None:
+            for task in not_done:
+                task.cancel()
+            await asyncio.gather(*not_done, return_exceptions=True)
+            raise first_exc
+
+    async def _run_segment(self, seg: SegmentRuntime) -> None:
+        offset = seg.start + seg.bytes_done
+        headers = {
+            "Range": f"bytes={offset}-{seg.end}",
+            "Accept-Encoding": "identity",
+        }
+        if self._etag:
+            headers["If-Range"] = self._etag
+        async with open_validated(
+            "GET", self.spec.url, headers=headers
+        ) as (resp, _final):
+            if resp.status == 200:
+                # Server ignored the range -> remote changed / no resume support.
+                raise RemoteChanged()
+            if resp.status not in (206,):
+                self._raise_for_status(resp.status)
+            async for chunk in resp.content.iter_chunked(args.download_chunk_size):
+                self._check_control()
+                # Never write past this segment's planned range: a
+                # non-conforming 206 that returns more than the requested
+                # bytes would otherwise overrun adjacent segments and the
+                # preallocated file. Cap the write and abort on overflow.
+                remaining = seg.length - seg.bytes_done
+                if remaining <= 0:
+                    raise FatalError(
+                        f"segment {seg.idx}: server returned more than the "
+                        f"requested {seg.length} bytes"
+                    )
+                overflow = len(chunk) > remaining
+                if overflow:
+                    chunk = chunk[:remaining]
+                await self._writer.write_at(offset, chunk)
+                offset += len(chunk)
+                seg.bytes_done += len(chunk)
+                self._recompute_bytes_done()
+                await self._persist_progress()
+                if overflow:
+                    raise FatalError(
+                        f"segment {seg.idx}: server returned more than the "
+                        f"requested {seg.length} bytes"
+                    )
+
+    async def _run_single(self) -> None:
+        seg = self.state.segments[0]
+        offset = seg.bytes_done  # resume from here for single-stream
+        headers = {"Accept-Encoding": "identity"}
+        if offset > 0:
+            headers["Range"] = f"bytes={offset}-"
+            if self._etag:
+                headers["If-Range"] = self._etag
+        async with open_validated(
+            "GET", self.spec.url, headers=headers
+        ) as (resp, _final):
+            if offset > 0 and resp.status == 200:
+                # Resume not honoured -> start over from the beginning. Truncate
+                # the existing partial so stale trailing bytes from the prior
+                # attempt cannot survive past the new (possibly shorter) end.
+                offset = 0
+                seg.bytes_done = 0
+                self.state.bytes_done = 0
+                await self._writer.truncate(0)
+            elif offset > 0 and resp.status != 206:
+                self._raise_for_status(resp.status)
+            elif offset == 0 and resp.status != 200:
+                self._raise_for_status(resp.status)
+            # Byte ceiling for this stream: the known total when the server
+            # reported a size, otherwise the configured maximum download size.
+            # Without a bound, a non-conforming response or an unknown-length
+            # stream (end == -1) that never closes could fill the disk (DoS).
+            limit = (seg.end + 1) if seg.end >= 0 else self._max_download_bytes()
+            async for chunk in resp.content.iter_chunked(args.download_chunk_size):
+                self._check_control()
+                overflow = False
+                if limit is not None:
+                    remaining = limit - offset
+                    if remaining <= 0:
+                        raise FatalError(
+                            f"download exceeded the maximum size {limit} bytes"
+                        )
+                    if len(chunk) > remaining:
+                        chunk = chunk[:remaining]
+                        overflow = True
+                await self._writer.write_at(offset, chunk)
+                offset += len(chunk)
+                seg.bytes_done = offset
+                self.state.bytes_done = offset
+                await self._persist_progress()
+                if overflow:
+                    raise FatalError(
+                        f"download exceeded the maximum size {limit} bytes"
+                    )
+
+    def _max_download_bytes(self) -> Optional[int]:
+        """Configured maximum download size in bytes, or ``None`` if disabled."""
+        cap = getattr(args, "download_max_bytes", 0)
+        return cap if cap and cap > 0 else None
+
+    def _raise_for_status(self, status: int) -> None:
+        if status in (401, 403):
+            raise FatalError(
+                f"{redact_url(self.spec.url)} returned {status}; authenticate this "
+                f"host via /api/download/auth or set its API key env var."
+            )
+        if status in _RETRYABLE_STATUSES:
+            raise RetryableError(f"HTTP {status}")
+        raise FatalError(f"unexpected HTTP {status}")
+
+    # ----- finalize / verify (PRD section 8.4) -----
+
+    async def _finalize(self) -> None:
+        self._check_control()
+        await self._close_writer()
+        await self._set_status(DownloadStatus.VERIFYING)
+
+        total = self.state.total_bytes
+        segmented = len(self.state.segments) > 1
+        if segmented:
+            # The .part was preallocated to total_bytes, so its on-disk size is
+            # not evidence of completeness: a segment that ends short (truncated
+            # 206 / server closes mid-range) leaves a zero-filled hole while the
+            # file size still equals total. Verify each segment wrote its full
+            # planned range, and trust the byte counter (== sum of segments)
+            # rather than os.path.getsize for the total check.
+            for seg in self.state.segments:
+                if seg.bytes_done != seg.length:
+                    raise FatalError(
+                        f"segment {seg.idx} incomplete: wrote {seg.bytes_done} "
+                        f"of {seg.length} bytes"
+                    )
+            observed = self.state.bytes_done
+        else:
+            # Single-stream writes a contiguous prefix, so the on-disk size is
+            # an independent witness of how much actually landed.
+            observed = os.path.getsize(self.spec.temp_path)
+        if total is not None and observed != total:
+            raise FatalError(
+                f"size mismatch: wrote {observed} of {total} bytes"
+            )
+
+        # Structural gate (cheap, no full read) then optional sha256 (full read).
+        # Both failures are non-retryable (a truncated/corrupt or mismatched file
+        # will not heal on retry), so surface them as FatalError rather than
+        # letting the plain Exceptions fall through to the retryable handler.
+        # ``temp_path`` carries the ``.part`` suffix; pass ``dest_path`` so the
+        # structural check detects the real file format instead of skipping it.
+        try:
+            await asyncio.to_thread(
+                structural.validate, self.spec.temp_path, self.spec.dest_path
+            )
+            if self.spec.expected_sha256:
+                await asyncio.to_thread(
+                    checksum.verify_sha256,
+                    self.spec.temp_path,
+                    self.spec.expected_sha256,
+                )
+        except (structural.StructuralError, checksum.ChecksumError) as e:
+            raise FatalError(str(e)) from e
+
+        os.makedirs(os.path.dirname(self.spec.dest_path), exist_ok=True)
+        os.replace(self.spec.temp_path, self.spec.dest_path)
+        logging.info(
+            "[model_downloader] completed %s (%d bytes)",
+            self.spec.model_id, observed,
+        )
+        # Catalog into the assets system (blake3 dedup identity). Best-effort.
+        await dedup.register_completed(self.spec.dest_path)
+
+    # ----- helpers -----
+
+    def _recompute_bytes_done(self) -> None:
+        self.state.bytes_done = sum(s.bytes_done for s in self.state.segments)
+        now = time.monotonic()
+        dt = now - self.state._last_time
+        if dt >= 0.5:
+            self.state.speed_bps = (self.state.bytes_done - self.state._last_bytes) / dt
+            self.state._last_bytes = self.state.bytes_done
+            self.state._last_time = now
+
+    async def _persist_progress(self, force: bool = False) -> None:
+        # Both the DB write and the websocket notify are gated by the same
+        # throttle: persisting hits SQLite, and notifying broadcasts to every
+        # client, so doing either per-chunk (small --download-chunk-size or
+        # many concurrent segments) would overwhelm both. Skip entirely inside
+        # the window; the next persist (or a forced one) ships the latest bytes.
+        now = time.monotonic()
+        if not force and now - self._last_persist < _PERSIST_INTERVAL:
+            return
+        self._last_persist = now
+        # SQLite is blocking; run it off the event loop per the queries module
+        # contract so progress persists don't stall the web server.
+        await asyncio.to_thread(self._write_progress)
+        if self._notify:
+            self._notify(self.spec.download_id)
+
+    def _write_progress(self) -> None:
+        queries.update_download(self.spec.download_id, bytes_done=self.state.bytes_done)
+        for seg in self.state.segments:
+            if seg.end >= seg.start:  # skip unknown-size sentinel
+                queries.update_segment_progress(
+                    self.spec.download_id, seg.idx, seg.bytes_done
+                )
+
+    async def _reset_for_restart(self) -> None:
+        await self._close_writer()
+        self._remove_temp()
+        for seg in self.state.segments:
+            seg.bytes_done = 0
+        self.state.bytes_done = 0
+        await asyncio.to_thread(
+            queries.update_download, self.spec.download_id, bytes_done=0
+        )
+        if await asyncio.to_thread(queries.list_segments, self.spec.download_id):
+            await asyncio.to_thread(
+                queries.replace_segments, self.spec.download_id, []
+            )
+
+    async def _close_writer(self) -> None:
+        if self._writer is not None:
+            try:
+                await self._writer.close()
+            except Exception:
+                logging.debug("[model_downloader] writer close error", exc_info=True)
+            self._writer = None
+
+    def _segmented_part_valid(self, total_bytes: int) -> bool:
+        """True when the temp file is the preallocated segmented ``.part``.
+
+        A segmented transfer preallocates the .part to ``total_bytes`` up front
+        and tracks how much of each range landed via per-segment offsets. Those
+        offsets are only trustworthy when the file they describe is still on
+        disk at its full preallocated size. A missing file (swept after a
+        failure, removed on a fatal error, deleted by hand) or a wrong-sized one
+        means the persisted offsets no longer correspond to real bytes and must
+        not be resumed over. Doing so would skip "complete" segments and leave
+        zero-filled holes that pass the size-only verification gate.
+        """
+        try:
+            return os.path.getsize(self.spec.temp_path) == total_bytes
+        except OSError:
+            return False
+
+    def _contiguous_prefix_valid(self, prefix_len: int) -> bool:
+        """True when the temp file is exactly ``prefix_len`` contiguous bytes.
+
+        Single-stream resume appends sequentially, so a valid resume point
+        implies the .part size equals the persisted offset. A larger file (e.g.
+        one preallocated to ``total_bytes`` by a previous segmented run) or a
+        missing/short file means the persisted offset is not a trustworthy
+        contiguous prefix and must not be resumed over.
+        """
+        try:
+            return os.path.getsize(self.spec.temp_path) == prefix_len
+        except OSError:
+            return False
+
+    def _remove_temp(self) -> None:
+        try:
+            os.remove(self.spec.temp_path)
+        except FileNotFoundError:
+            pass
+        except OSError as e:
+            logging.warning(
+                "[model_downloader] could not remove %s: %s", self.spec.temp_path, e
+            )
+
+    async def _set_status(self, status: str, error: Optional[str] = None) -> None:
+        # ``error`` is authoritative: passing None clears any prior failure
+        # text so transitions out of a failure state (retry/success) don't
+        # leave stale messages on RuntimeState or in the persisted row.
+        self.state.status = status
+        self.state.error = error
+        fields = {"status": status, "bytes_done": self.state.bytes_done, "error": error}
+        if status == DownloadStatus.QUEUED:
+            fields["attempts"] = self.spec.attempts + 1
+            self.spec.attempts += 1
+        await asyncio.to_thread(queries.update_download, self.spec.download_id, **fields)
+        if self._notify:
+            self._notify(self.spec.download_id)

app/model_downloader/engine/planner.py

@@ -0,0 +1,51 @@
+"""Segment planning.
+
+Split a known byte range into S roughly-equal segments, each fetched by its
+own coroutine with ``Range: bytes=start-end``. Falls back to a single segment
+when the server doesn't support ranges or the size is unknown/too small for
+segmentation to be worthwhile.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Below this size, the per-connection setup cost outweighs any parallelism.
+_MIN_SEGMENT_BYTES = 1 * 1024 * 1024
+
+
+@dataclass(frozen=True)
+class SegmentPlan:
+    idx: int
+    start: int
+    end: int  # inclusive
+
+    @property
+    def length(self) -> int:
+        return self.end - self.start + 1
+
+
+def effective_segment_count(
+    total_bytes: int | None, accept_ranges: bool, configured: int
+) -> int:
+    """How many segments to actually use for this file."""
+    if not accept_ranges or total_bytes is None or total_bytes <= 0:
+        return 1
+    by_size = max(1, total_bytes // _MIN_SEGMENT_BYTES)
+    return max(1, min(configured, by_size))
+
+
+def plan_segments(total_bytes: int, num_segments: int) -> list[SegmentPlan]:
+    """Return ``num_segments`` contiguous, inclusive byte ranges covering [0, total)."""
+    if total_bytes <= 0 or num_segments <= 1:
+        return [SegmentPlan(idx=0, start=0, end=max(0, total_bytes - 1))]
+    base = total_bytes // num_segments
+    plans: list[SegmentPlan] = []
+    start = 0
+    for i in range(num_segments):
+        # Last segment soaks up the remainder.
+        length = base if i < num_segments - 1 else total_bytes - start
+        end = start + length - 1
+        plans.append(SegmentPlan(idx=i, start=start, end=end))
+        start = end + 1
+    return plans

app/model_downloader/engine/writer.py

@@ -0,0 +1,110 @@
+"""Positioned, off-loop file writes.
+
+Network I/O stays on the event loop; every blocking disk op (preallocate,
+positioned write, fsync) is run in a bounded thread pool via
+``run_in_executor`` so downloads never stall inference or the web server.
+
+A single file descriptor is opened for the whole download. Segments write to
+their own offsets with ``os.pwrite`` — which is offset-addressed and atomic
+per call, so concurrent segment writers need no extra locking. Per-chunk
+fsync is avoided; we fsync once at completion.
+
+``os.pwrite`` is unavailable on Windows, so there we fall back to
+``os.lseek`` + ``os.write`` guarded by a per-writer lock (the seek/write pair
+is not atomic, so concurrent segment writers must be serialized).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import os
+import threading
+from concurrent.futures import ThreadPoolExecutor
+from typing import Optional
+
+# One shared, bounded pool for all download disk I/O.
+_EXECUTOR = ThreadPoolExecutor(max_workers=8, thread_name_prefix="dl-writer")
+
+_HAS_PWRITE = hasattr(os, "pwrite")
+
+# On Windows ``os.open`` defaults to text mode, which translates every ``\n``
+# byte into ``\r\n`` on write and corrupts binary payloads (the file grows by
+# one byte per 0x0A). ``O_BINARY`` disables that translation; it does not exist
+# on POSIX, where the default is already binary.
+_O_BINARY = getattr(os, "O_BINARY", 0)
+
+
+class FileWriter:
+    """Owns the ``.part`` file descriptor for one download."""
+
+    def __init__(self, path: str) -> None:
+        self.path = path
+        self._fd: Optional[int] = None
+        # Serializes lseek+write on platforms without os.pwrite (Windows).
+        self._seek_lock = threading.Lock()
+
+    def _open(self) -> None:
+        os.makedirs(os.path.dirname(self.path), exist_ok=True)
+        self._fd = os.open(self.path, os.O_RDWR | os.O_CREAT | _O_BINARY, 0o644)
+
+    async def open(self) -> None:
+        await asyncio.get_running_loop().run_in_executor(_EXECUTOR, self._open)
+
+    async def preallocate(self, size: int) -> None:
+        """Grow the file to ``size`` so segments write to their offsets."""
+        if self._fd is None or size <= 0:
+            return
+        await asyncio.get_running_loop().run_in_executor(
+            _EXECUTOR, os.ftruncate, self._fd, size
+        )
+
+    async def truncate(self, size: int = 0) -> None:
+        """Truncate the file to ``size`` bytes (default: empty it)."""
+        if self._fd is None:
+            return
+        await asyncio.get_running_loop().run_in_executor(
+            _EXECUTOR, os.ftruncate, self._fd, size
+        )
+
+    def _pwrite_all(self, data: bytes, offset: int) -> None:
+        """A positioned write may write fewer bytes than requested (signal
+        interruption, near-ENOSPC); loop until every byte lands so we never
+        leave a gap while the caller advances by the full chunk length.
+
+        Uses ``os.pwrite`` where available (offset-addressed, atomic per call).
+        On Windows it falls back to ``os.lseek`` + ``os.write`` under a lock,
+        since that pair is not atomic across concurrent segment writers."""
+        assert self._fd is not None, "writer not opened"
+        view = memoryview(data)
+        written = 0
+        total = len(view)
+        while written < total:
+            if _HAS_PWRITE:
+                n = os.pwrite(self._fd, view[written:], offset + written)
+            else:
+                with self._seek_lock:
+                    os.lseek(self._fd, offset + written, os.SEEK_SET)
+                    n = os.write(self._fd, view[written:])
+            if n == 0:
+                raise OSError(
+                    f"positioned write wrote 0 bytes at offset {offset + written} "
+                    f"({written}/{total} bytes written)"
+                )
+            written += n
+
+    async def write_at(self, offset: int, data: bytes) -> None:
+        assert self._fd is not None, "writer not opened"
+        await asyncio.get_running_loop().run_in_executor(
+            _EXECUTOR, self._pwrite_all, data, offset
+        )
+
+    async def flush(self) -> None:
+        if self._fd is None:
+            return
+        await asyncio.get_running_loop().run_in_executor(_EXECUTOR, os.fsync, self._fd)
+
+    async def close(self) -> None:
+        if self._fd is None:
+            return
+        fd, self._fd = self._fd, None
+        await asyncio.get_running_loop().run_in_executor(_EXECUTOR, os.close, fd)

app/model_downloader/manager.py

@@ -0,0 +1,444 @@
+"""Public facade for the download manager.
+
+This is the only object the server imports. It validates requests, owns the
+:class:`Scheduler`, and exposes a small async API plus read models for status.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import uuid
+from typing import Callable, Optional
+
+from app.model_downloader.constants import DownloadStatus
+from app.model_downloader.database import queries
+from app.model_downloader.net.probe import gated_error_message, probe
+from app.model_downloader.scheduler import SCHEDULER
+from app.model_downloader.security import paths
+from app.model_downloader.net.http import redact_url
+from app.model_downloader.security.allowlist import (
+    ALLOWED_MODEL_EXTENSIONS,
+    filename_extension,
+    is_host_allowed_url,
+    is_url_downloadable,
+    url_path_extension,
+)
+from app.model_downloader.security.paths import InvalidModelId
+
+# Non-terminal statuses: an existing row in one of these blocks a re-enqueue.
+_LIVE_STATUSES = (
+    DownloadStatus.QUEUED,
+    DownloadStatus.ACTIVE,
+    DownloadStatus.PAUSED,
+    DownloadStatus.VERIFYING,
+)
+
+
+class DownloadError(Exception):
+    """A user-facing error with a stable machine-readable code."""
+
+    def __init__(self, code: str, message: str, status: int = 400) -> None:
+        super().__init__(message)
+        self.code = code
+        self.message = message
+        self.http_status = status
+
+
+class DownloadManager:
+    def __init__(self) -> None:
+        self._scheduler = SCHEDULER
+        self._notify_cb: Optional[Callable[[str], None]] = None
+        # Serializes the "check for a live download, then write" critical section
+        # per model_id. ``downloads`` has no uniqueness constraint on model_id
+        # (history rows are kept), so without this two concurrent enqueue/resume
+        # calls could both pass the live check and admit two jobs sharing one
+        # temp/dest path. The manager is a process singleton over a local SQLite
+        # DB, so an in-process lock is sufficient (and avoids a migration).
+        self._model_locks: dict[str, asyncio.Lock] = {}
+
+    def set_notify(self, cb: Optional[Callable[[str], None]]) -> None:
+        self._notify_cb = cb
+        self._scheduler.set_notify(cb)
+
+    async def start(self) -> None:
+        await self._scheduler.start()
+
+    # ----- enqueue -----
+
+    async def enqueue(
+        self,
+        url: str,
+        model_id: str,
+        *,
+        priority: int = 0,
+        expected_sha256: Optional[str] = None,
+        allow_any_extension: bool = False,
+    ) -> str:
+        # Coarse gate first: host/scheme must be allowlisted, and any extension
+        # present in the URL path must be a known model type. A URL whose path
+        # carries NO extension (e.g. Civitai's ``/api/download/models/<id>``) is
+        # admitted here and its real extension is resolved from the network
+        # below before the download is finally accepted.
+        if allow_any_extension:
+            if not is_host_allowed_url(url):
+                raise DownloadError(
+                    "URL_NOT_ALLOWED",
+                    "URL is not on the download allowlist (host/scheme).",
+                )
+        elif not is_url_downloadable(url):
+            raise DownloadError(
+                "URL_NOT_ALLOWED",
+                "URL is not on the download allowlist (host/scheme/extension).",
+            )
+
+        # When the URL path has no extension, follow it to where it resolves and
+        # adopt the real extension from the response, forcing the stored
+        # filename to match. Skipped when the caller opted into any extension.
+        if not allow_any_extension and url_path_extension(url) == "":
+            resolved_ext = await self._resolve_extension(url)
+            model_id = paths.apply_extension(model_id, resolved_ext)
+
+        try:
+            paths.parse_model_id(model_id, allow_any_extension)
+            dest_path, temp_path = paths.resolve_destination(model_id, allow_any_extension)
+        except InvalidModelId as e:
+            raise DownloadError("INVALID_MODEL_ID", str(e))
+
+        if await asyncio.to_thread(
+            paths.resolve_existing, model_id, allow_any_extension
+        ):
+            raise DownloadError(
+                "ALREADY_AVAILABLE",
+                f"Model already exists on disk: {model_id}",
+                status=409,
+            )
+        download_id = str(uuid.uuid4())
+        # Hold the per-model lock across the live check and the insert so a
+        # concurrent enqueue/resume for the same model_id cannot interleave
+        # between them and create a second job against the same temp/dest path.
+        async with self._model_lock(model_id):
+            if await self._has_live_download(model_id):
+                raise DownloadError(
+                    "ALREADY_DOWNLOADING",
+                    f"A download for {model_id} is already in progress.",
+                    status=409,
+                )
+            await asyncio.to_thread(
+                queries.insert_download,
+                {
+                    "id": download_id,
+                    "url": url,
+                    "model_id": model_id,
+                    "dest_path": dest_path,
+                    "temp_path": temp_path,
+                    "status": DownloadStatus.QUEUED,
+                    "priority": priority,
+                    "expected_sha256": expected_sha256,
+                    "allow_any_extension": allow_any_extension,
+                },
+            )
+        logging.info("[model_downloader] enqueued %s -> %s", redact_url(url), model_id)
+        await self._scheduler.pump()
+        return download_id
+
+    async def _resolve_extension(self, url: str) -> str:
+        """Follow ``url`` to its final response and return the real extension.
+
+        Used for allowlisted URLs whose path has no extension (e.g. Civitai
+        download endpoints): the filename lives in the ``Content-Disposition``
+        header or the post-redirect URL. Raises :class:`DownloadError` when the
+        URL can't be resolved, needs authentication, or resolves to something
+        that is not a known model file — so we never persist a bogus destination.
+        """
+        pr = await probe(url)
+        if not pr.ok:
+            if pr.gated:
+                raise DownloadError(
+                    "GATED_REPO" if pr.is_gated_repo else "CREDENTIALS_REQUIRED",
+                    gated_error_message(url, pr),
+                    status=401,
+                )
+            raise DownloadError(
+                "URL_RESOLVE_FAILED",
+                f"Could not resolve {redact_url(url)}: {pr.error or 'unknown error'}",
+                status=502,
+            )
+        ext = filename_extension(pr.filename) if pr.filename else ""
+        if ext not in ALLOWED_MODEL_EXTENSIONS:
+            raise DownloadError(
+                "URL_NOT_ALLOWED",
+                f"URL resolves to {pr.filename or '<unknown>'!r}, which is not a "
+                f"known model file type {ALLOWED_MODEL_EXTENSIONS}.",
+            )
+        return ext
+
+    def _model_lock(self, model_id: str) -> asyncio.Lock:
+        # Lazily create one lock per model_id. There is no ``await`` between the
+        # lookup and the insert, so under the single asyncio thread this is
+        # atomic and cannot hand out two different locks for the same model_id.
+        lock = self._model_locks.get(model_id)
+        if lock is None:
+            lock = asyncio.Lock()
+            self._model_locks[model_id] = lock
+        return lock
+
+    async def _has_live_download(
+        self, model_id: str, *, exclude_id: Optional[str] = None
+    ) -> bool:
+        return await asyncio.to_thread(
+            queries.has_live_download_for_model, model_id, _LIVE_STATUSES, exclude_id
+        )
+
+    # ----- control -----
+
+    async def pause(self, download_id: str) -> None:
+        job = self._scheduler.get_job(download_id)
+        if job is not None:
+            job.request_pause()
+            return
+        row = await asyncio.to_thread(queries.get_download, download_id)
+        if row is None:
+            raise DownloadError("NOT_FOUND", "No such download.", status=404)
+        if row.status == DownloadStatus.QUEUED:
+            await asyncio.to_thread(
+                queries.update_download, download_id, status=DownloadStatus.PAUSED
+            )
+
+    async def resume(self, download_id: str) -> None:
+        row = await asyncio.to_thread(queries.get_download, download_id)
+        if row is None:
+            raise DownloadError("NOT_FOUND", "No such download.", status=404)
+        if row.status not in (DownloadStatus.PAUSED, DownloadStatus.FAILED):
+            return
+        # Re-queueing a paused/failed row must respect the single-live-per-model
+        # invariant: another download (e.g. a newer enqueue) may already be live
+        # for this model_id and would share this row's temp/dest path. Hold the
+        # per-model lock across the check and the status flip, and exclude this
+        # row itself (a paused row is already a "live" status).
+        async with self._model_lock(row.model_id):
+            if await self._has_live_download(row.model_id, exclude_id=download_id):
+                raise DownloadError(
+                    "ALREADY_DOWNLOADING",
+                    f"A download for {row.model_id} is already in progress.",
+                    status=409,
+                )
+            await asyncio.to_thread(
+                queries.update_download,
+                download_id,
+                status=DownloadStatus.QUEUED,
+                error=None,
+            )
+        await self._scheduler.pump()
+
+    async def cancel(self, download_id: str) -> None:
+        job = self._scheduler.get_job(download_id)
+        if job is not None:
+            job.request_cancel()
+            return
+        row = await asyncio.to_thread(queries.get_download, download_id)
+        if row is None:
+            raise DownloadError("NOT_FOUND", "No such download.", status=404)
+        if row.status in _LIVE_STATUSES:
+            try:
+                os.remove(row.temp_path)
+            except OSError:
+                pass
+            await asyncio.to_thread(
+                queries.update_download, download_id, status=DownloadStatus.CANCELLED
+            )
+
+    async def set_priority(self, download_id: str, priority: int) -> None:
+        row = await asyncio.to_thread(queries.get_download, download_id)
+        if row is None:
+            raise DownloadError("NOT_FOUND", "No such download.", status=404)
+        await asyncio.to_thread(
+            queries.update_download, download_id, priority=priority
+        )
+        # Admission-order only; a higher priority is
+        # picked up the next time a slot frees. Pump in case a slot is free now.
+        await self._scheduler.pump()
+
+    async def delete(self, download_id: str) -> None:
+        """Delete a terminal download so it stays gone from history.
+
+        Refuses to delete a live download so a record is never removed out from
+        under a running worker; cancel it first. Any leftover ``.part`` temp
+        file (e.g. from a failed transfer) is removed, but the finished model
+        file on disk is never touched.
+        """
+        if self._scheduler.get_job(download_id) is not None:
+            raise DownloadError(
+                "DOWNLOAD_ACTIVE",
+                "Cannot delete a download that is still in progress.",
+                status=409,
+            )
+        row = await asyncio.to_thread(queries.get_download, download_id)
+        if row is None:
+            raise DownloadError("NOT_FOUND", "No such download.", status=404)
+        if row.status in _LIVE_STATUSES:
+            raise DownloadError(
+                "DOWNLOAD_ACTIVE",
+                "Cannot delete a download that is still in progress.",
+                status=409,
+            )
+
+        try:
+            os.remove(row.temp_path)
+        except OSError:
+            pass
+        await asyncio.to_thread(queries.delete_download, download_id)
+
+    async def clear(self) -> int:
+        """Delete all terminal downloads from history in one transaction.
+
+        Skips anything still live (queued/active/paused/verifying, or a running
+        job) so an in-flight download is never removed out from under a worker.
+        Finished model files on disk are never touched; only leftover ``.part``
+        temp files from failed/cancelled transfers are removed. Returns the
+        number of history rows deleted.
+        """
+
+        rows = await asyncio.to_thread(queries.list_downloads)
+        deletable = [
+            r
+            for r in rows
+            if r.status not in _LIVE_STATUSES
+            and self._scheduler.get_job(r.id) is None
+        ]
+        if not deletable:
+            return 0
+        for r in deletable:
+            try:
+                os.remove(r.temp_path)
+            except OSError:
+                pass
+        return await asyncio.to_thread(
+            queries.delete_downloads, [r.id for r in deletable]
+        )
+
+    # ----- read models -----
+
+    def _view(self, row) -> dict:
+        """Combine the persisted row with live in-memory progress, if running."""
+        job = self._scheduler.get_job(row.id)
+        bytes_done = row.bytes_done
+        total = row.total_bytes
+        speed = None
+        eta = None
+        segments = None
+        if job is not None:
+            st = job.state
+            bytes_done = st.bytes_done
+            total = st.total_bytes if st.total_bytes is not None else total
+            speed = st.speed_bps
+            eta = st.eta_seconds
+            segments = [
+                {"idx": s.idx, "bytes_done": s.bytes_done, "length": s.length}
+                for s in st.segments
+                if s.end >= s.start
+            ]
+        progress = (bytes_done / total) if total else None
+        return {
+            "download_id": row.id,
+            "model_id": row.model_id,
+            "url": redact_url(row.url),
+            "status": row.status,
+            "priority": row.priority,
+            "total_bytes": total,
+            "bytes_done": bytes_done,
+            "progress": progress,
+            "speed_bps": speed,
+            "eta_seconds": eta,
+            "segments": segments,
+            "error": row.error,
+            "created_at": row.created_at,
+            "updated_at": row.updated_at,
+        }
+
+    def _view_from_state(self, job) -> dict:
+        """Build a view purely from the live in-memory job state (no DB)."""
+        st = job.state
+        return {
+            "download_id": st.download_id,
+            "model_id": st.model_id,
+            "url": redact_url(st.url),
+            "status": st.status,
+            "priority": st.priority,
+            "total_bytes": st.total_bytes,
+            "bytes_done": st.bytes_done,
+            "progress": st.progress,
+            "speed_bps": st.speed_bps,
+            "eta_seconds": st.eta_seconds,
+            "segments": [
+                {"idx": s.idx, "bytes_done": s.bytes_done, "length": s.length}
+                for s in st.segments
+                if s.end >= s.start
+            ],
+            "error": st.error,
+        }
+
+    def status_sync(self, download_id: str) -> Optional[dict]:
+        """Synchronous status read for the websocket notify path.
+
+        Uses live in-memory state when the job is running (no DB round-trip on
+        the hot path); falls back to a quick DB read otherwise.
+        """
+        job = self._scheduler.get_job(download_id)
+        if job is not None:
+            return self._view_from_state(job)
+        row = queries.get_download(download_id)
+        return self._view(row) if row is not None else None
+
+    async def status(self, download_id: str) -> Optional[dict]:
+        row = await asyncio.to_thread(queries.get_download, download_id)
+        return self._view(row) if row is not None else None
+
+    async def list(self) -> list[dict]:
+        rows = await asyncio.to_thread(queries.list_downloads)
+        return [self._view(r) for r in rows]
+
+    async def availability(self, models: dict[str, str]) -> dict[str, dict]:
+        """Bulk per-id ``{state, progress, ...}`` for the frontend poll.
+
+        ``state`` is ``available`` (on disk), ``downloading`` (live row), or
+        ``missing``. Cheap: a path lookup plus an in-memory/DB status check.
+        """
+        rows = await asyncio.to_thread(queries.list_downloads)
+        by_model: dict[str, object] = {}
+        for r in rows:
+            if r.status in _LIVE_STATUSES or r.model_id not in by_model:
+                by_model[r.model_id] = r
+
+        # ``url_allowed`` mirrors the coarse enqueue gate (host/scheme + a
+        # non-disallowed extension); URLs whose extension is only known after a
+        # network resolve — e.g. Civitai download endpoints — report allowed.
+        out: dict[str, dict] = {}
+        for model_id, url in models.items():
+            try:
+                exists = await asyncio.to_thread(paths.resolve_existing, model_id)
+            except InvalidModelId:
+                out[model_id] = {"state": "missing", "url_allowed": is_url_downloadable(url)}
+                continue
+            if exists:
+                out[model_id] = {"state": "available", "url_allowed": is_url_downloadable(url)}
+                continue
+            row = by_model.get(model_id)
+            if row is not None and row.status in _LIVE_STATUSES:
+                view = self._view(row)
+                out[model_id] = {
+                    "state": "downloading",
+                    "url_allowed": is_url_downloadable(url),
+                    "download_id": view["download_id"],
+                    "progress": view["progress"],
+                    "bytes_done": view["bytes_done"],
+                    "total_bytes": view["total_bytes"],
+                    "speed_bps": view["speed_bps"],
+                }
+            else:
+                out[model_id] = {"state": "missing", "url_allowed": is_url_downloadable(url)}
+        return out
+
+
+DOWNLOAD_MANAGER = DownloadManager()

app/model_downloader/net/http.py

@@ -0,0 +1,142 @@
+"""Manual, validated redirect-following request opener.
+
+Automatic redirects are disabled. We follow hops ourselves
+so that on *every* hop we (a) re-validate scheme + reject credentials-in-URL,
+(b) recompute which auth — if any — applies to that hop's host, and (c) let the
+connector's resolver screen the IP. This is the single place that attaches a
+token, so it can never ride a redirect to a CDN host.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from contextlib import asynccontextmanager
+from typing import AsyncIterator, Optional
+from urllib.parse import unquote, urljoin, urlsplit, urlunsplit
+
+import aiohttp
+
+from app.model_downloader.auth.resolver import resolve_auth_for_hop
+from app.model_downloader.net.session import get_session
+from app.model_downloader.security.ssrf import (
+    MAX_REDIRECTS,
+    SSRFError,
+    check_redirect_hop,
+)
+
+_REDIRECT_CODES = {301, 302, 303, 307, 308}
+DEFAULT_TIMEOUT = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=120)
+
+
+def redact_url(url: str) -> str:
+    """Drop the query string so a query-scheme secret is never logged/stored."""
+    try:
+        parts = urlsplit(url)
+    except ValueError:
+        return "<unparseable-url>"
+    return urlunsplit(parts._replace(query=""))
+
+
+_CD_FILENAME_STAR = re.compile(
+    r"filename\*\s*=\s*[^']*'[^']*'([^;]+)", re.IGNORECASE
+)
+_CD_FILENAME_QUOTED = re.compile(r'filename\s*=\s*"([^"]+)"', re.IGNORECASE)
+_CD_FILENAME_BARE = re.compile(r"filename\s*=\s*([^;]+)", re.IGNORECASE)
+
+
+def filename_from_content_disposition(value: Optional[str]) -> Optional[str]:
+    """Extract the download filename from a ``Content-Disposition`` header.
+
+    Prefers the RFC 5987 ``filename*=`` form (percent-decoded) over the plain
+    ``filename=`` form. Any directory components in the value are stripped so a
+    hostile header can only influence the *name*, never the target directory.
+    Returns ``None`` when no filename is present.
+    """
+    if not value:
+        return None
+    for pat, decode in (
+        (_CD_FILENAME_STAR, True),
+        (_CD_FILENAME_QUOTED, False),
+        (_CD_FILENAME_BARE, False),
+    ):
+        m = pat.search(value)
+        if not m:
+            continue
+        raw = m.group(1).strip().strip('"')
+        if decode:
+            try:
+                raw = unquote(raw)
+            except Exception:
+                pass
+        name = raw.replace("\\", "/").rsplit("/", 1)[-1].strip()
+        if name:
+            return name
+    return None
+
+
+async def _resolve_final_response(
+    method: str,
+    url: str,
+    base_headers: dict[str, str],
+    timeout: aiohttp.ClientTimeout,
+) -> tuple[aiohttp.ClientResponse, str]:
+    """Follow redirects manually until a non-redirect response.
+
+    Each intermediate redirect response is released before the next hop.
+    Returns the final ``(response, final_url)``; the caller owns releasing it.
+    """
+    session = await get_session()
+    current = url
+    hops = 0
+    while True:
+        check_redirect_hop(current, is_initial_url=(hops == 0))
+        parts = urlsplit(current)
+        auth = await resolve_auth_for_hop(parts.hostname or "", parts.scheme)
+        req_headers = dict(base_headers)
+        if auth is not None:
+            req_headers.update(auth.headers)
+
+        resp = await session.request(
+            method,
+            current,
+            allow_redirects=False,
+            headers=req_headers,
+            timeout=timeout,
+        )
+        if resp.status in _REDIRECT_CODES and resp.headers.get("Location"):
+            next_url = urljoin(str(resp.url), resp.headers["Location"])
+            await resp.release()
+            hops += 1
+            if hops > MAX_REDIRECTS:
+                raise SSRFError(
+                    f"too many redirects (> {MAX_REDIRECTS}) for {redact_url(url)}"
+                )
+            current = next_url
+            continue
+        return resp, redact_url(str(resp.url))
+
+
+@asynccontextmanager
+async def open_validated(
+    method: str,
+    url: str,
+    *,
+    headers: Optional[dict[str, str]] = None,
+    timeout: aiohttp.ClientTimeout = DEFAULT_TIMEOUT,
+) -> AsyncIterator[tuple[aiohttp.ClientResponse, str]]:
+    """Open ``method url`` following redirects manually and validated.
+
+    Yields ``(response, final_url)`` where ``final_url`` is redacted of any
+    query string. The response is released automatically on exit.
+    """
+    resp, final_url = await _resolve_final_response(
+        method, url, dict(headers or {}), timeout
+    )
+    try:
+        yield resp, final_url
+    finally:
+        try:
+            await resp.release()
+        except Exception:  # pragma: no cover - best-effort cleanup
+            logging.debug("[model_downloader] response release error", exc_info=True)

app/model_downloader/net/probe.py

@@ -0,0 +1,192 @@
+"""Pre-download probe.
+
+Issues a tiny ranged GET (``Range: bytes=0-0``) — which doubles as a
+range-support test — to discover ``Content-Length``, ``Accept-Ranges``,
+``ETag``/``Last-Modified``, and the final post-redirect URL. For HuggingFace
+LFS files the true size also appears in the non-standard ``X-Linked-Size``
+header, which we read as a fallback.
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import Optional
+from urllib.parse import urlparse, urlsplit
+
+import aiohttp
+
+from app.model_downloader.net.http import (
+    filename_from_content_disposition,
+    open_validated,
+    redact_url,
+)
+from app.model_downloader.net.session import parse_int_header
+
+_PROBE_TIMEOUT = aiohttp.ClientTimeout(total=60, sock_connect=30, sock_read=30)
+
+
+@dataclass
+class ProbeResult:
+    ok: bool
+    status: int
+    final_url: Optional[str] = None
+    total_bytes: Optional[int] = None
+    accept_ranges: bool = False
+    etag: Optional[str] = None
+    last_modified: Optional[str] = None
+    gated: bool = False  # 401/403 — needs (or has wrong) credentials
+    error: Optional[str] = None
+    # HuggingFace's ``X-Error-Code`` header (e.g. ``GatedRepo``,
+    # ``RepoNotFound``) when the host reports one. Lets us tell "this repo is
+    # gated — request access" apart from "you just need a token".
+    error_code: Optional[str] = None
+    # Filename the server intends this response to be saved as: the
+    # ``Content-Disposition`` name if present, else the post-redirect URL's
+    # basename. Used to resolve the real extension for URLs (e.g. Civitai's
+    # ``/api/download`` endpoints) that carry no extension in their path.
+    filename: Optional[str] = None
+
+    @property
+    def is_gated_repo(self) -> bool:
+        """True when the host says the repo is gated (access must be granted).
+
+        Distinct from a plain missing/invalid token: even a valid credential
+        won't help until the user accepts the model's terms on its page.
+        """
+        return (self.error_code or "").lower() == "gatedrepo"
+
+
+def _error_detail(error_code: Optional[str], error_message: Optional[str]) -> str:
+    """Format the host's ``X-Error-Code``/``X-Error-Message`` for logs/messages."""
+    detail = ": ".join(p.strip() for p in (error_code, error_message) if p and p.strip())
+    return f" ({detail})" if detail else ""
+
+
+def _probe_failure_message(
+    status: int, error_code: Optional[str], error_message: Optional[str]
+) -> str:
+    msg = f"probe returned HTTP {status}{_error_detail(error_code, error_message)}"
+    if status == 404:
+        # HuggingFace returns 404 (not 403) for a private repo the current
+        # credentials cannot see, so it is indistinguishable from a missing
+        # file without the hint. Name both causes so the user can check the
+        # URL or their access/token scope.
+        msg += (
+            " — the file may not exist, or it is private/gated and the "
+            "credentials in use lack access to it"
+        )
+    return msg
+
+
+def _total_from_content_range(value: Optional[str]) -> Optional[int]:
+    # "bytes 0-0/12345" -> 12345 ; "bytes 0-0/*" -> None
+    if not value or "/" not in value:
+        return None
+    total = value.rsplit("/", 1)[1].strip()
+    return parse_int_header(total)
+
+
+def _filename_from_response(
+    content_disposition: Optional[str], final_url: Optional[str]
+) -> Optional[str]:
+    name = filename_from_content_disposition(content_disposition)
+    if name:
+        return name
+    if final_url:
+        base = urlsplit(final_url).path.rsplit("/", 1)[-1]
+        if base:
+            return base
+    return None
+
+
+async def probe(url: str) -> ProbeResult:
+    """Probe ``url`` and return discovered metadata, failing soft."""
+    try:
+        async with open_validated(
+            "GET",
+            url,
+            headers={"Range": "bytes=0-0", "Accept-Encoding": "identity"},
+            timeout=_PROBE_TIMEOUT,
+        ) as (resp, final_url):
+            # HuggingFace (and some others) report the real reason in these
+            # headers on any status, including 404 for a private/missing repo.
+            error_code = resp.headers.get("X-Error-Code")
+            error_message = resp.headers.get("X-Error-Message")
+            if resp.status in (401, 403):
+                logging.warning(
+                    "[model_downloader] probe %s -> HTTP %d%s",
+                    redact_url(final_url or url), resp.status,
+                    _error_detail(error_code, error_message),
+                )
+                return ProbeResult(
+                    ok=False, status=resp.status, final_url=final_url, gated=True,
+                    error_code=error_code,
+                    error=(
+                        error_message
+                        or f"host returned {resp.status} (authentication required)"
+                    ),
+                )
+            if resp.status not in (200, 206):
+                logging.warning(
+                    "[model_downloader] probe %s -> HTTP %d%s",
+                    redact_url(final_url or url), resp.status,
+                    _error_detail(error_code, error_message),
+                )
+                return ProbeResult(
+                    ok=False, status=resp.status, final_url=final_url,
+                    error_code=error_code,
+                    error=_probe_failure_message(resp.status, error_code, error_message),
+                )
+
+            headers = resp.headers
+            accept_ranges = False
+            total: Optional[int] = None
+            if resp.status == 206:
+                accept_ranges = True
+                total = _total_from_content_range(headers.get("Content-Range"))
+            else:  # 200: server ignored the range
+                accept_ranges = headers.get("Accept-Ranges", "").lower() == "bytes"
+                total = parse_int_header(headers.get("Content-Length"))
+
+            if total is None:
+                total = parse_int_header(headers.get("X-Linked-Size"))
+
+            return ProbeResult(
+                ok=True,
+                status=resp.status,
+                final_url=final_url,
+                total_bytes=total,
+                accept_ranges=accept_ranges,
+                etag=headers.get("ETag"),
+                last_modified=headers.get("Last-Modified"),
+                filename=_filename_from_response(
+                    headers.get("Content-Disposition"), final_url
+                ),
+            )
+    except Exception as e:  # network / SSRF / timeout
+        host = urlparse(url).netloc or "<unknown>"
+        logging.debug("[model_downloader] probe failed for %s: %s", host, type(e).__name__)
+        return ProbeResult(ok=False, status=0, error="probe failed: network error")
+
+
+def gated_error_message(url: str, pr: ProbeResult) -> str:
+    """Build a user-facing message for a gated/auth-required probe result.
+
+    Distinguishes a *gated* repo (access must be requested/granted on the model
+    page — a token alone is not enough) from a plain missing/invalid credential.
+    """
+    redacted = redact_url(url)
+    if pr.is_gated_repo:
+        detail = (pr.error or "access is restricted").rstrip()
+        if detail and not detail.endswith((".", "!", "?")):
+            detail += "."
+        return (
+            f"{redacted} is a gated model — {detail} Request access on the model's "
+            f"page, authenticate this host via /api/download/auth (or set its API "
+            f"key env var), and retry."
+        )
+    return (
+        f"{redacted} requires authentication. Authenticate this host via "
+        f"/api/download/auth or set its API key env var, and retry."
+    )

app/model_downloader/net/session.py

@@ -0,0 +1,72 @@
+"""Lazily-created shared :class:`aiohttp.ClientSession`.
+
+A single session reuses TLS handshakes and TCP connections across the probe
+and the many segment GETs to the same host (HuggingFace is the dominant
+case), which is a large speedup on cold connections and exactly the
+connection-reuse strategy that lets us match aria2c.
+
+The connector uses :class:`ValidatingResolver` so every connection — initial
+or post-redirect — is screened for private/special-use IPs at connect time.
+TLS is pinned to certifi's CA bundle because the OS trust store is not wired
+up on some Python installs (python.org macOS, slim containers).
+"""
+
+from __future__ import annotations
+
+import asyncio
+import ssl
+from typing import Optional
+
+import aiohttp
+
+try:
+    import certifi
+    _CA_FILE = certifi.where()
+except Exception:  # pragma: no cover - certifi is a transitive dep of aiohttp
+    _CA_FILE = None
+
+from comfy.cli_args import args
+from app.model_downloader.security.ssrf import ValidatingResolver
+
+_session: Optional[aiohttp.ClientSession] = None
+_lock = asyncio.Lock()
+
+
+def ssl_context() -> ssl.SSLContext:
+    if _CA_FILE is not None:
+        return ssl.create_default_context(cafile=_CA_FILE)
+    return ssl.create_default_context()
+
+
+async def get_session() -> aiohttp.ClientSession:
+    """Return the shared session, creating it on first use."""
+    global _session
+    if _session is not None and not _session.closed:
+        return _session
+    async with _lock:
+        if _session is None or _session.closed:
+            connector = aiohttp.TCPConnector(
+                limit_per_host=max(1, getattr(args, "download_max_connections_per_host", 16)),
+                ssl=ssl_context(),
+                resolver=ValidatingResolver(),
+            )
+            _session = aiohttp.ClientSession(connector=connector)
+    return _session
+
+
+async def close_session() -> None:
+    global _session
+    if _session is not None and not _session.closed:
+        await _session.close()
+    _session = None
+
+
+def parse_int_header(value: Optional[str]) -> Optional[int]:
+    """Parse a non-negative integer header value, or None if bad/absent."""
+    if not value:
+        return None
+    try:
+        n = int(value)
+    except (TypeError, ValueError):
+        return None
+    return n if n >= 0 else None

app/model_downloader/scheduler.py

@@ -0,0 +1,176 @@
+"""Priority scheduler + lifecycle.
+
+Owns the set of running jobs and admits queued downloads up to a global
+concurrency limit (K), highest priority first, FIFO within a priority. Runs
+entirely on the existing ComfyUI asyncio loop; blocking work (disk, hashing,
+DB) is offloaded by the job/writer layers.
+
+On startup it reconciles DB vs. disk: ``active``/``verifying`` rows left by a
+previous run are reset to ``queued`` and resumed from persisted offsets, and
+orphaned ``.part`` files with no live download row are swept.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+import random
+import time
+from typing import Callable, Optional
+
+from comfy.cli_args import args
+from app.model_downloader.constants import DownloadStatus
+from app.model_downloader.database import queries
+from app.model_downloader.engine.job import DownloadJob, JobSpec
+from app.model_downloader.security import paths
+
+# Backoff for retryable failures
+_BACKOFF_BASE = 2.0
+_BACKOFF_CAP = 300.0
+_MAX_ATTEMPTS = 6
+
+
+class Scheduler:
+    def __init__(self) -> None:
+        self._jobs: dict[str, DownloadJob] = {}
+        self._tasks: dict[str, asyncio.Task] = {}
+        self._backoff_until: dict[str, float] = {}
+        self._pump_lock = asyncio.Lock()
+        self._notify_cb: Optional[Callable[[str], None]] = None
+        self._started = False
+
+    @property
+    def max_active(self) -> int:
+        return max(1, getattr(args, "download_max_active", 3))
+
+    def set_notify(self, cb: Optional[Callable[[str], None]]) -> None:
+        self._notify_cb = cb
+
+    def get_job(self, download_id: str) -> Optional[DownloadJob]:
+        return self._jobs.get(download_id)
+
+    def is_active(self, download_id: str) -> bool:
+        return download_id in self._tasks
+
+    # ----- startup -----
+
+    async def start(self) -> None:
+        if self._started:
+            return
+        self._started = True
+        try:
+            await asyncio.to_thread(queries.reconcile_live_downloads)
+            await asyncio.to_thread(self._sweep_orphan_temp_files)
+        except Exception as e:
+            logging.warning("[model_downloader] startup reconcile failed: %s", e)
+        await self.pump()
+
+    @staticmethod
+    def _sweep_orphan_temp_files() -> None:
+        """Remove ``.part`` files not referenced by a resumable download row.
+
+        Resumable partials are preserved; only truly orphaned temp files from
+        crashed runs are deleted. ``FAILED`` is included because
+        :meth:`DownloadManager.resume` explicitly permits resuming a
+        retry-exhausted failed row: deleting its partial here while the
+        per-segment offsets survive in the DB would make the next resume
+        preallocate a fresh sparse file, skip every "complete" segment, and
+        leave zero-filled holes that pass the size-only verification gate.
+        """
+        live = {
+            row.temp_path
+            for row in queries.list_downloads()
+            if row.status
+            in (
+                DownloadStatus.QUEUED,
+                DownloadStatus.PAUSED,
+                DownloadStatus.FAILED,
+            )
+        }
+        for path in paths.iter_all_tmp_paths():
+            if path in live:
+                continue
+            try:
+                os.remove(path)
+                logging.info("[model_downloader] removed orphan temp file: %s", path)
+            except OSError as e:
+                logging.warning("[model_downloader] could not remove %s: %s", path, e)
+
+    # ----- admission -----
+
+    async def pump(self) -> None:
+        async with self._pump_lock:
+            slots = self.max_active - len(self._tasks)
+            if slots <= 0:
+                return
+            now = time.monotonic()
+            candidates = await asyncio.to_thread(queries.list_queued_downloads)
+            for row in candidates:
+                if slots <= 0:
+                    break
+                if row.id in self._tasks:
+                    continue
+                if self._backoff_until.get(row.id, 0.0) > now:
+                    continue
+                self._admit(row)
+                slots -= 1
+
+    def _admit(self, row) -> None:
+        spec = JobSpec(
+            download_id=row.id,
+            url=row.url,
+            model_id=row.model_id,
+            dest_path=row.dest_path,
+            temp_path=row.temp_path,
+            priority=row.priority,
+            expected_sha256=row.expected_sha256,
+            allow_any_extension=row.allow_any_extension,
+            etag=row.etag,
+            attempts=row.attempts,
+        )
+        job = DownloadJob(spec, notify_cb=self._notify_cb)
+        self._jobs[row.id] = job
+        self._tasks[row.id] = asyncio.ensure_future(self._run_job(job))
+
+    async def _run_job(self, job: DownloadJob) -> None:
+        download_id = job.spec.download_id
+        status = DownloadStatus.FAILED
+        try:
+            status = await job.run()
+        except Exception as e:  # run() is defensive, but never let a task die silently
+            logging.error("[model_downloader] job %s crashed: %s", download_id, e)
+            queries.update_download(
+                download_id,
+                status=DownloadStatus.FAILED,
+                error=f"internal error: {e}",
+            )
+            if self._notify_cb:
+                self._notify_cb(download_id)
+        finally:
+            self._tasks.pop(download_id, None)
+            self._jobs.pop(download_id, None)
+
+        if status == DownloadStatus.QUEUED:
+            if job.spec.attempts >= _MAX_ATTEMPTS:
+                queries.update_download(
+                    download_id,
+                    status=DownloadStatus.FAILED,
+                    error=f"giving up after {job.spec.attempts} attempts",
+                )
+                if self._notify_cb:
+                    self._notify_cb(download_id)
+            else:
+                delay = min(
+                    _BACKOFF_CAP, _BACKOFF_BASE ** job.spec.attempts
+                ) + random.uniform(0, 1.0)
+                self._backoff_until[download_id] = time.monotonic() + delay
+                asyncio.ensure_future(self._delayed_pump(delay))
+        await self.pump()
+
+    async def _delayed_pump(self, delay: float) -> None:
+        await asyncio.sleep(delay)
+        await self.pump()
+
+
+SCHEDULER = Scheduler()

app/model_downloader/security/allowlist.py

@@ -0,0 +1,140 @@
+"""URL allowlist for server-side model fetches.
+
+Default-deny. A URL is downloadable only when its parsed host + scheme are
+allowlisted AND (unless explicitly relaxed) its final filename ends in a
+known model extension.
+
+The built-in host defaults mirror the frontend's ``isModelDownloadable``
+allowlist so the two flows agree on what is eligible; ``--download-allowed-hosts``
+extends it for self-hosted mirrors. Matching is done on ``urlparse().hostname``
+(never a raw string prefix) so userinfo tricks like
+``http://127.0.0.1@169.254.169.254/x.safetensors`` — whose real host is the
+metadata IP — cannot slip past.
+"""
+
+from __future__ import annotations
+
+from urllib.parse import urlparse
+
+from comfy.cli_args import args
+
+# host -> set of allowed schemes. Frontend parity (HuggingFace / Civitai /
+# localhost). Extra hosts from --download-allowed-hosts are https-only.
+_DEFAULT_ALLOWED_HOSTS: dict[str, set[str]] = {
+    "huggingface.co": {"https"},
+    "civitai.com": {"https"},
+    "localhost": {"http", "https"},
+    "127.0.0.1": {"http", "https"},
+}
+
+# Hosts for which loopback addresses are intentionally permitted (the localhost
+# "download a local model" feature). Every other host's loopback resolution is
+# rejected by the SSRF resolver.
+LOOPBACK_HOSTS = frozenset({"localhost", "127.0.0.1", "::1"})
+
+# Known model file extensions (frontend parity). Checked on the final filename.
+ALLOWED_MODEL_EXTENSIONS = (
+    ".safetensors",
+    ".sft",
+    ".ckpt",
+    ".pth",
+    ".pt",
+    ".gguf",
+    ".bin",
+)
+
+
+def _allowed_hosts() -> dict[str, set[str]]:
+    hosts = {h: set(s) for h, s in _DEFAULT_ALLOWED_HOSTS.items()}
+    for extra in getattr(args, "download_allowed_hosts", []) or []:
+        host = extra.strip().lower()
+        if host:
+            hosts.setdefault(host, set()).add("https")
+    return hosts
+
+
+def is_host_allowed(host: str | None, scheme: str | None) -> bool:
+    """True iff ``host`` is allowlisted for ``scheme``.
+
+    Used both for the initial URL and re-checked on every redirect hop,
+    so a whitelisted URL cannot 30x into an off-list host.
+    """
+    if not host or not scheme:
+        return False
+    allowed = _allowed_hosts().get(host.lower())
+    return allowed is not None and scheme.lower() in allowed
+
+
+def has_allowed_extension(path: str, allow_any_extension: bool = False) -> bool:
+    if allow_any_extension:
+        return True
+    return path.lower().endswith(ALLOWED_MODEL_EXTENSIONS)
+
+
+def filename_extension(name: str) -> str:
+    """Lowercased extension (including the leading dot) of a bare filename.
+
+    Returns ``""`` when there is no extension. A leading-dot name
+    (``.safetensors``) is treated as having no extension (all stem), matching
+    ``os.path.splitext`` semantics so dotfiles aren't mistaken for typed files.
+    """
+    base = name.replace("\\", "/").rsplit("/", 1)[-1]
+    dot = base.rfind(".")
+    if dot <= 0:
+        return ""
+    return base[dot:].lower()
+
+
+def is_allowed_extension_name(name: str) -> bool:
+    """True iff ``name`` ends in one of the known model extensions."""
+    return name.lower().endswith(ALLOWED_MODEL_EXTENSIONS)
+
+
+def is_host_allowed_url(url: str) -> bool:
+    """True iff ``url`` parses and its host+scheme are allowlisted."""
+    if not isinstance(url, str) or not url:
+        return False
+    try:
+        parsed = urlparse(url)
+    except ValueError:
+        return False
+    return is_host_allowed(parsed.hostname, parsed.scheme)
+
+
+def url_path_extension(url: str) -> str:
+    """Extension of the URL *path* basename (query ignored), or ``""``."""
+    try:
+        parsed = urlparse(url)
+    except ValueError:
+        return ""
+    return filename_extension(parsed.path)
+
+
+def is_url_downloadable(url: str) -> bool:
+    """Coarse enqueue gate: host/scheme allowed and extension not disallowed.
+
+    Unlike :func:`is_url_allowed` (which demands a known extension *in the URL*),
+    this also admits URLs whose path carries no extension at all — e.g. a Civitai
+    ``/api/download/models/<id>`` endpoint whose real filename only shows up in
+    the redirect target / ``Content-Disposition``. The true extension is then
+    resolved from the network and re-validated before the download is admitted.
+    A path bearing an explicit *non-model* extension (``.zip``, ``.html``, ...)
+    is still rejected here.
+    """
+    if not is_host_allowed_url(url):
+        return False
+    ext = url_path_extension(url)
+    return ext == "" or ext in ALLOWED_MODEL_EXTENSIONS
+
+
+def is_url_allowed(url: str, allow_any_extension: bool = False) -> bool:
+    """Check whether ``url`` is permitted as a server-side download source."""
+    if not isinstance(url, str) or not url:
+        return False
+    try:
+        parsed = urlparse(url)
+    except ValueError:
+        return False
+    if not is_host_allowed(parsed.hostname, parsed.scheme):
+        return False
+    return has_allowed_extension(parsed.path, allow_any_extension)

app/model_downloader/security/paths.py

@@ -0,0 +1,132 @@
+"""Path resolution + traversal safety for downloads.
+
+A ``model_id`` is a *relative destination path* of the form
+``<directory>/<filename>`` (e.g. ``loras/my_lora.safetensors``). This module
+turns one into an absolute on-disk path under one of ComfyUI's registered
+model folders, rejecting unknown folders, path traversal, and symlink escape.
+This is the only thing that composes destination paths, so the engine never
+touches user-supplied path strings directly.
+"""
+
+from __future__ import annotations
+
+import os
+import re
+from typing import Iterator, Optional
+
+import folder_paths
+
+from app.model_downloader.constants import TMP_SUFFIX
+from app.model_downloader.security.allowlist import ALLOWED_MODEL_EXTENSIONS
+
+# A model_id component is a single path segment of safe characters — no slashes,
+# no "..", no leading dots that could escape the target directory.
+_SEGMENT_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+class InvalidModelId(ValueError):
+    """Raised when a model_id is malformed or names an unknown model folder."""
+
+
+def parse_model_id(model_id: str, allow_any_extension: bool = False) -> tuple[str, str]:
+    """Split ``<directory>/<filename>`` and validate both components.
+
+    Returns ``(directory, filename)``. Does not touch the filesystem.
+    """
+    if not isinstance(model_id, str) or "/" not in model_id:
+        raise InvalidModelId(
+            f"model_id must be '<directory>/<filename>', got {model_id!r}"
+        )
+    directory, _, filename = model_id.partition("/")
+    if "/" in filename or not directory or not filename:
+        raise InvalidModelId(
+            f"model_id must have exactly one '/' separator, got {model_id!r}"
+        )
+    if not _SEGMENT_RE.match(directory):
+        raise InvalidModelId(f"invalid directory segment {directory!r}")
+    if not _SEGMENT_RE.match(filename):
+        raise InvalidModelId(f"invalid filename segment {filename!r}")
+    if not allow_any_extension and not filename.lower().endswith(
+        ALLOWED_MODEL_EXTENSIONS
+    ):
+        raise InvalidModelId(
+            f"filename must end with a known model extension "
+            f"{ALLOWED_MODEL_EXTENSIONS}, got {filename!r}"
+        )
+    if directory not in folder_paths.folder_names_and_paths:
+        raise InvalidModelId(f"unknown model folder {directory!r}")
+    return directory, filename
+
+
+def apply_extension(model_id: str, ext: str) -> str:
+    """Return ``model_id`` with its filename forced to end in ``ext``.
+
+    ``ext`` includes the leading dot (e.g. ``".safetensors"``). If the filename
+    already ends in a *known model extension* it is replaced; otherwise ``ext``
+    is appended (so ``loras/mymodel`` -> ``loras/mymodel.safetensors`` and
+    ``loras/mymodel.ckpt`` -> ``loras/mymodel.safetensors``). A filename with a
+    non-model suffix (``my.model.v2``) is treated as an extensionless stem and
+    ``ext`` is appended. The directory part is left untouched; validation is
+    still the caller's job via :func:`parse_model_id`.
+    """
+    directory, sep, filename = model_id.partition("/")
+    if not sep:
+        return model_id  # malformed; parse_model_id will reject it
+    low = filename.lower()
+    for known in ALLOWED_MODEL_EXTENSIONS:
+        if low.endswith(known):
+            filename = filename[: -len(known)]
+            break
+    return f"{directory}{sep}{filename}{ext}"
+
+
+def resolve_existing(model_id: str, allow_any_extension: bool = False) -> Optional[str]:
+    """Return the absolute path of an installed model, or None if missing.
+
+    Honours ``extra_model_paths.yaml`` transparently via ``get_full_path``.
+    """
+    directory, filename = parse_model_id(model_id, allow_any_extension)
+    return folder_paths.get_full_path(directory, filename)
+
+
+def resolve_destination(
+    model_id: str, allow_any_extension: bool = False
+) -> tuple[str, str]:
+    """Return ``(final_path, temp_path)`` for a download.
+
+    Downloads land at the first registered path for the model's directory
+    (the "primary" location). ``temp_path`` is a sibling ``.part`` file that
+    is atomically renamed onto ``final_path`` on success. The result is
+    asserted to stay within the registered root (defence in depth on top of
+    the segment regex).
+    """
+    directory, filename = parse_model_id(model_id, allow_any_extension)
+    roots = folder_paths.get_folder_paths(directory)
+    if not roots:
+        raise InvalidModelId(f"no on-disk path registered for folder {directory!r}")
+    root = os.path.realpath(roots[0])
+    final_path = os.path.realpath(os.path.join(root, filename))
+    if final_path != root and not final_path.startswith(root + os.sep):
+        raise InvalidModelId(f"resolved path escapes model root: {model_id!r}")
+    temp_path = f"{final_path}{TMP_SUFFIX}"
+    return final_path, temp_path
+
+
+def iter_all_tmp_paths() -> Iterator[str]:
+    """Yield this subsystem's temp files under every registered model folder.
+
+    Matches only the distinctive ``TMP_SUFFIX`` so the startup orphan sweep
+    can never delete temp files created by other tools.
+    """
+    seen_roots: set[str] = set()
+    for directory in list(folder_paths.folder_names_and_paths.keys()):
+        for root in folder_paths.get_folder_paths(directory):
+            if root in seen_roots or not os.path.isdir(root):
+                continue
+            seen_roots.add(root)
+            try:
+                for entry in os.scandir(root):
+                    if entry.is_file() and entry.name.endswith(TMP_SUFFIX):
+                        yield entry.path
+            except OSError:
+                continue

app/model_downloader/security/ssrf.py

@@ -0,0 +1,163 @@
+"""SSRF / exfiltration defenses.
+
+Two cooperating layers:
+
+1. :class:`ValidatingResolver` is installed on the shared connector. Every
+   connection — the initial probe and every segment GET, including ones made
+   after a redirect — resolves its host through this resolver, which rejects
+   any address that lands on a private / special-use IP range. Because the
+   resolve and the connect happen together inside the connector, there is no
+   check-then-connect window for DNS rebinding to exploit.
+
+2. :func:`check_redirect_hop` re-validates every hop. The host allowlist gates
+   only the *initial* user-supplied URL (anti-SSRF for arbitrary input);
+   legitimate downloads from allowlisted origins redirect to presigned CDN
+   hosts that are deliberately NOT on the allowlist (HF ->
+   ``cdn-lfs*.huggingface.co``, Civitai -> signed Cloudflare/S3), so hops are
+   instead screened for scheme, embedded credentials, and — via the resolver
+   above — private IPs. Credentials are only ever attached when a hop's host
+   exactly matches a stored credential, so they are dropped on the CDN hop.
+   Loopback (the "download a local model" feature) is exempt from IP filtering
+   only for the initial URL: a *redirect* may never target a loopback host or
+   a blocked IP-literal, which the resolver alone can't enforce (it exempts
+   loopback literals and never sees IP literals through DNS).
+"""
+
+from __future__ import annotations
+
+import ipaddress
+import socket
+from urllib.parse import urlparse
+
+from aiohttp.abc import AbstractResolver
+from aiohttp.resolver import DefaultResolver
+
+from app.model_downloader.security.allowlist import LOOPBACK_HOSTS
+
+# Cap the redirect chain length a hop may use.
+MAX_REDIRECTS = 5
+
+
+class SSRFError(Exception):
+    """A hop failed an SSRF / allowlist check."""
+
+
+def is_scheme_allowed(scheme: str | None, host: str | None) -> bool:
+    """True iff ``scheme`` is permitted for ``host`` on a download hop.
+
+    https is always allowed; plain http only for loopback/approved dev hosts.
+    """
+    if not scheme:
+        return False
+    scheme = scheme.lower()
+    if scheme == "https":
+        return True
+    if scheme == "http":
+        return bool(host) and host.lower() in LOOPBACK_HOSTS
+    return False
+
+
+def is_blocked_ip(ip_str: str) -> bool:
+    """True for any address we refuse to connect to.
+
+    Covers loopback, link-local (incl. 169.254.169.254 cloud metadata),
+    RFC1918 private ranges, unique-local (ULA), unspecified (0.0.0.0/::),
+    multicast and other reserved ranges.
+    """
+    try:
+        ip = ipaddress.ip_address(ip_str)
+    except ValueError:
+        return True  # unparseable -> refuse
+    # On CPython before the gh-113171 fix (backported to 3.12.4/3.11.9/
+    # 3.10.14/3.9.19) the is_* properties don't see through IPv4-mapped IPv6
+    # (e.g. ::ffff:169.254.169.254), so resolve and re-check the embedded IPv4
+    # to keep mapped metadata/private addresses from slipping past the filter.
+    mapped = getattr(ip, "ipv4_mapped", None)
+    if mapped is not None:
+        ip = mapped
+    return (
+        ip.is_private
+        or ip.is_loopback
+        or ip.is_link_local
+        or ip.is_multicast
+        or ip.is_reserved
+        or ip.is_unspecified
+    )
+
+
+class ValidatingResolver(AbstractResolver):
+    """Delegating resolver that drops blocked IPs from every resolution.
+
+    If a hostname resolves only to blocked addresses, the connection fails
+    closed with an :class:`OSError`, which aiohttp surfaces as a connection
+    error to the caller.
+    """
+
+    def __init__(self) -> None:
+        self._inner = DefaultResolver()
+
+    async def resolve(self, host, port=0, family=socket.AF_INET):
+        infos = await self._inner.resolve(host, port, family)
+        # localhost/127.0.0.1 are an explicit, opt-in allowlist feature.
+        if isinstance(host, str) and host.lower() in LOOPBACK_HOSTS:
+            return infos
+        safe = [info for info in infos if not is_blocked_ip(info["host"])]
+        if not safe:
+            raise OSError(
+                f"refusing to connect to {host!r}: resolves only to "
+                f"private/special-use addresses"
+            )
+        return safe
+
+    async def close(self) -> None:
+        await self._inner.close()
+
+
+def check_redirect_hop(url: str, *, is_initial_url: bool = False) -> str:
+    """Validate one hop's URL.
+
+    Returns the URL unchanged on success; raises :class:`SSRFError` otherwise.
+    Requires https for external hosts (http only for loopback/approved dev
+    hosts) and forbids credentials-in-URL. The host is NOT re-checked against
+    the allowlist (CDN redirect targets are off-list by design); credential
+    leakage is prevented by exact host matching at attach time, and the landing
+    filename's extension is gated separately by the caller.
+
+    Loopback/blocked-IP screening: the connector's resolver filters resolvable
+    hostnames but exempts literal loopback hosts (``localhost``/``127.0.0.1``/
+    ``::1``) and never sees IP literals through DNS. That loopback exemption is
+    legitimate only for the *initial* user-supplied URL (``is_initial_url``);
+    on a redirect hop we reject loopback hosts and any blocked IP-literal here,
+    so a 30x can't steer a server-side GET at loopback/internal services.
+    """
+    try:
+        parsed = urlparse(url)
+    except ValueError as e:
+        raise SSRFError(f"unparseable redirect URL {url!r}: {e}") from e
+    host = parsed.hostname
+    if not host:
+        raise SSRFError(f"redirect URL has no host: {url!r}")
+    if not is_scheme_allowed(parsed.scheme, host):
+        raise SSRFError(
+            f"redirect to disallowed scheme {parsed.scheme!r} for host "
+            f"{host!r} (https required for external hosts)"
+        )
+    if parsed.username or parsed.password:
+        raise SSRFError("credentials-in-URL are not allowed")
+    host_is_loopback = host.lower() in LOOPBACK_HOSTS
+    if not is_initial_url and host_is_loopback:
+        raise SSRFError(f"redirect to loopback host {host!r} is not allowed")
+    # IP-literal targets never go through DNS, so the connector's resolver can't
+    # screen them — check them directly. The only blocked IP allowed through is
+    # a loopback literal on the initial URL (handled by the exemption above).
+    try:
+        ipaddress.ip_address(host)
+    except ValueError:
+        is_ip_literal = False
+    else:
+        is_ip_literal = True
+    if is_ip_literal and is_blocked_ip(host) and not (
+        is_initial_url and host_is_loopback
+    ):
+        raise SSRFError(f"redirect to blocked internal address {host!r}")
+    return url

app/model_downloader/verify/checksum.py

@@ -0,0 +1,49 @@
+"""Hub-checksum verification = SHA256.
+
+Only used to confirm a download matches a *provided* ``expected_sha256``. It
+is NOT the dedup key (that is blake3, owned by the assets system). The full
+sequential read happens at most once, here, only when a checksum was supplied.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from typing import Callable, Optional
+
+_CHUNK = 8 * 1024 * 1024
+
+InterruptCheck = Callable[[], bool]
+
+
+class ChecksumError(Exception):
+    """The computed SHA256 did not match the expected value."""
+
+
+def sha256_file(path: str, interrupt_check: Optional[InterruptCheck] = None) -> Optional[str]:
+    """Stream the file and return its lowercase hex SHA256.
+
+    Returns ``None`` if interrupted via ``interrupt_check``.
+    """
+    h = hashlib.sha256()
+    with open(path, "rb") as f:
+        while True:
+            if interrupt_check is not None and interrupt_check():
+                return None
+            chunk = f.read(_CHUNK)
+            if not chunk:
+                break
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def verify_sha256(
+    path: str, expected: str, interrupt_check: Optional[InterruptCheck] = None
+) -> None:
+    """Raise :class:`ChecksumError` unless the file's SHA256 matches ``expected``."""
+    actual = sha256_file(path, interrupt_check)
+    if actual is None:
+        return  # interrupted; caller will re-verify on resume
+    if actual.lower() != expected.lower():
+        raise ChecksumError(
+            f"sha256 mismatch: expected {expected.lower()}, got {actual.lower()}"
+        )

app/model_downloader/verify/dedup.py

@@ -0,0 +1,53 @@
+"""Dedup + catalog handoff — reuse the assets system.
+
+We do NOT build a parallel indexer. "Do I already have it?" is answered by
+``resolve_existing`` (path) at enqueue time and, where a hash is known, by the
+assets blake3 catalog. After a completed download we register the file
+through the assets ingest path so it is cataloged and (eventually) hashed by
+the existing enrichment worker.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from typing import Optional
+
+
+def _register_sync(abs_path: str) -> Optional[str]:
+    """Register a finished file into the assets catalog. Returns asset hash."""
+    try:
+        from app.assets.services.ingest import register_file_in_place
+    except Exception as e:  # assets package import failure — non-fatal
+        logging.debug("[model_downloader] assets ingest unavailable: %s", e)
+        return None
+    try:
+        result = register_file_in_place(abs_path, name=os.path.basename(abs_path), tags=[])
+        return result.asset.hash if result and result.asset else None
+    except Exception as e:
+        # The file is already safely on disk; cataloging is best-effort.
+        logging.warning(
+            "[model_downloader] could not register %s into assets catalog: %s",
+            abs_path, e,
+        )
+        return None
+
+
+async def register_completed(abs_path: str) -> Optional[str]:
+    """Catalog a completed download via the assets system (off the event loop)."""
+    return await asyncio.to_thread(_register_sync, abs_path)
+
+
+def _find_by_hash_sync(blake3_hex: str) -> Optional[str]:
+    try:
+        from app.assets.services.asset_management import get_asset_by_hash
+    except Exception:
+        return None
+    asset = get_asset_by_hash("blake3:" + blake3_hex)
+    return asset.hash if asset is not None else None
+
+
+async def find_existing_by_hash(blake3_hex: str) -> Optional[str]:
+    """Pure DB lookup — never triggers hashing on the hot path."""
+    return await asyncio.to_thread(_find_by_hash_sync, blake3_hex)

app/model_downloader/verify/structural.py

@@ -0,0 +1,86 @@
+"""Cheap structural validation, no full read.
+
+For ``.safetensors``/``.sft`` we parse the header (first few KB): it carries
+the tensor table and the byte length of the data region. We assert
+``file_size == 8 + header_len + data_region_len``. This detects truncation
+and most corruption for free, before any crypto hashing. Other extensions
+have no cheap structural check and pass through.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import struct
+from typing import Optional
+
+_SAFETENSORS_EXTS = (".safetensors", ".sft")
+# A sane upper bound so a corrupt header length can't make us read gigabytes.
+_MAX_HEADER_BYTES = 100 * 1024 * 1024
+
+
+class StructuralError(Exception):
+    """The file failed its structural integrity check."""
+
+
+def validate(path: str, name_hint: Optional[str] = None) -> None:
+    """Validate the file at ``path``. Raises :class:`StructuralError` on failure.
+
+    The file format is detected from ``name_hint`` when provided, otherwise from
+    ``path``. Callers that download into a temp file with an opaque suffix (e.g.
+    ``*.comfy-download.part``) must pass the final destination name as
+    ``name_hint`` so the format check is not silently skipped.
+    """
+    lower = (name_hint or path).lower()
+    if lower.endswith(_SAFETENSORS_EXTS):
+        _validate_safetensors(path)
+    # No structural check for other formats; the size + (optional) checksum
+    # gates in the engine cover those.
+
+
+def _validate_safetensors(path: str) -> None:
+    file_size = os.path.getsize(path)
+    if file_size < 8:
+        raise StructuralError(f"file too small to be safetensors ({file_size} bytes)")
+    with open(path, "rb") as f:
+        header_len = struct.unpack("<Q", f.read(8))[0]
+        if header_len <= 0 or header_len > _MAX_HEADER_BYTES:
+            raise StructuralError(f"implausible safetensors header length {header_len}")
+        if 8 + header_len > file_size:
+            raise StructuralError("safetensors header extends past end of file")
+        try:
+            header = json.loads(f.read(header_len).decode("utf-8"))
+        except (UnicodeDecodeError, json.JSONDecodeError) as e:
+            raise StructuralError(f"safetensors header is not valid JSON: {e}") from e
+
+    if not isinstance(header, dict):
+        raise StructuralError("safetensors header is not a JSON object")
+
+    data_len = 0
+    for name, entry in header.items():
+        if name == "__metadata__":
+            continue
+        if not isinstance(entry, dict) or "data_offsets" not in entry:
+            raise StructuralError(f"tensor {name!r} missing data_offsets")
+        offsets = entry["data_offsets"]
+        if not (isinstance(offsets, list) and len(offsets) == 2):
+            raise StructuralError(f"tensor {name!r} has malformed data_offsets")
+        begin, end = offsets
+        # bool is an int subclass; reject it explicitly to avoid True/False offsets.
+        if (
+            not isinstance(begin, int)
+            or not isinstance(end, int)
+            or isinstance(begin, bool)
+            or isinstance(end, bool)
+            or begin < 0
+            or end < begin
+        ):
+            raise StructuralError(f"tensor {name!r} has malformed data_offsets")
+        data_len = max(data_len, end)
+
+    expected = 8 + header_len + data_len
+    if file_size != expected:
+        raise StructuralError(
+            f"size mismatch: file is {file_size} bytes, header implies {expected} "
+            f"(8 + {header_len} header + {data_len} data)"
+        )

app/model_manager.py

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

comfy/cli_args.py

@@ -33,6 +33,28 @@ class EnumAction(argparse.Action):
         setattr(namespace, self.dest, value)
 
 
+def _positive_int(value: str) -> int:
+    """argparse type that rejects zero and negative integers."""
+    try:
+        ivalue = int(value)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"{value!r} is not an integer")
+    if ivalue <= 0:
+        raise argparse.ArgumentTypeError(f"{value!r} must be a positive integer (> 0)")
+    return ivalue
+
+
+def _non_negative_int(value: str) -> int:
+    """argparse type that rejects negatives but allows zero (a disable sentinel)."""
+    try:
+        ivalue = int(value)
+    except ValueError:
+        raise argparse.ArgumentTypeError(f"{value!r} is not an integer")
+    if ivalue < 0:
+        raise argparse.ArgumentTypeError(f"{value!r} must be a non-negative integer (>= 0)")
+    return ivalue
+
+
 parser = argparse.ArgumentParser()
 
 parser.add_argument("--listen", type=str, default="127.0.0.1", metavar="IP", nargs="?", const="0.0.0.0,::", help="Specify the IP address to listen on (default: 127.0.0.1). You can give a list of ip addresses by separating them with a comma like: 127.2.2.2,127.3.3.3 If --listen is provided without an argument, it defaults to 0.0.0.0,:: (listens on all ipv4 and ipv6)")
@@ -92,6 +114,7 @@ parser.add_argument("--directml", type=int, nargs="?", metavar="DIRECTML_DEVICE"
 parser.add_argument("--oneapi-device-selector", type=str, default=None, metavar="SELECTOR_STRING", help="Sets the oneAPI device(s) this instance will use.")
 parser.add_argument("--supports-fp8-compute", action="store_true", help="ComfyUI will act like if the device supports fp8 compute.")
 parser.add_argument("--enable-triton-backend", action="store_true", help="ComfyUI will enable the use of Triton backend in comfy-kitchen. Is disabled at launch by default.")
+parser.add_argument("--disable-triton-backend", action="store_true", help="Force-disable the comfy-kitchen Triton backend, overriding the automatic ROCm/AMD default and --enable-triton-backend.")
 
 class LatentPreviewMethod(enum.Enum):
     NoPreviews = "none"
@@ -245,6 +268,15 @@ parser.add_argument("--enable-asset-hashing", action="store_true", help="Compute
 parser.add_argument("--feature-flag", type=str, action='append', default=[], metavar="KEY[=VALUE]", help="Set a server feature flag. Use KEY=VALUE to set an explicit value, or bare KEY to set it to true. Can be specified multiple times. Boolean values (true/false) and numbers are auto-converted. Examples: --feature-flag show_signin_button=true  or  --feature-flag show_signin_button")
 parser.add_argument("--list-feature-flags", action="store_true", help="Print the registry of known CLI-settable feature flags as JSON and exit.")
 
+# ----- Model download manager (PRD: docs/prd-download-manager.md) -----
+parser.add_argument("--download-segments", type=_positive_int, default=8, metavar="N", help="Number of parallel HTTP range segments per file for the model download manager (default: 8).")
+parser.add_argument("--download-max-active", type=_positive_int, default=3, metavar="N", help="Maximum number of model downloads running concurrently (default: 3).")
+parser.add_argument("--download-max-connections-per-host", type=_positive_int, default=16, metavar="N", help="Maximum simultaneous connections to a single host for the download manager (default: 16).")
+parser.add_argument("--download-chunk-size", type=_positive_int, default=4 * 1024 * 1024, metavar="BYTES", help="Read chunk size in bytes for the download manager (default: 4 MiB).")
+parser.add_argument("--download-max-bytes", type=_non_negative_int, default=1024 * 1024 * 1024 * 1024, metavar="BYTES", help="Maximum size in bytes of a single download; aborts transfers that exceed it (guards against malicious/non-conforming hosts filling the disk). Set to 0 to disable (default: 1 TiB).")
+parser.add_argument("--download-allowed-hosts", type=str, nargs="*", default=[], metavar="HOST", help="Additional hostnames to add to the download manager allowlist (https only). The built-in defaults always include huggingface.co and civitai.com.")
+parser.add_argument("--download-allow-any-extension", action="store_true", help="Allow the download manager to fetch files with any extension (default: only known model extensions like .safetensors).")
+
 if comfy.options.args_parsing:
     args = parser.parse_args()
 else:

comfy/comfy_api_env.py

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

comfy/latent_formats.py

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

comfy/ldm/modules/attention.py

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

comfy/ldm/modules/diffusionmodules/model.py

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

comfy/ldm/pixeldit/model.py

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

comfy/ldm/pixeldit/pid.py

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

comfy/ldm/seedvr/attention.py

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

comfy/ldm/seedvr/color_fix.py

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

comfy/ldm/seedvr/constants.py

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

comfy/model_base.py

@@ -55,6 +55,7 @@ import comfy.ldm.pixeldit.model
 import comfy.ldm.pixeldit.pid
 import comfy.ldm.ace.model
 import comfy.ldm.omnigen.omnigen2
+import comfy.ldm.seedvr.model
 import comfy.ldm.boogu.model
 import comfy.ldm.qwen_image.model
 import comfy.ldm.ideogram4.model
@@ -932,6 +933,17 @@ class HunyuanDiT(BaseModel):
         out['image_meta_size'] = comfy.conds.CONDRegular(torch.FloatTensor([[height, width, target_height, target_width, 0, 0]]))
         return out
 
+class SeedVR2(BaseModel):
+    def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
+        super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.seedvr.model.NaDiT)
+
+    def extra_conds(self, **kwargs):
+        out = super().extra_conds(**kwargs)
+        condition = kwargs.get("condition", None)
+        if condition is not None:
+            out["condition"] = comfy.conds.CONDRegular(condition)
+        return out
+
 class PixArt(BaseModel):
     def __init__(self, model_config, model_type=ModelType.EPS, device=None):
         super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.pixart.pixartms.PixArtMS)

comfy/model_detection.py

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

comfy/model_management.py

@@ -616,6 +616,8 @@ PIN_PRESSURE_HYSTERESIS = 256 * 1024 * 1024
 #Freeing registerables on pressure does imply a GPU sync, so go big on
 #the hysteresis so each expensive sync gives us back a good chunk.
 REGISTERABLE_PIN_HYSTERESIS = 2048 * 1024 * 1024
+WINDOWS_PIN_EVICTION_SWAP_PERCENT = 5.0
+WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE = 512 * 1024 ** 2
 
 def module_size(module):
     module_mem = 0
@@ -642,6 +644,15 @@ def free_pins(size, evict_active=False):
             size -= freed
     return freed_total
 
+def should_free_pins_for_ram_pressure(shortfall):
+    if shortfall <= 0:
+        return False
+    if not WINDOWS:
+        return True
+    if psutil.virtual_memory().available < WINDOWS_PIN_EVICTION_EMERGENCY_AVAILABLE:
+        return True
+    return psutil.swap_memory().percent >= WINDOWS_PIN_EVICTION_SWAP_PERCENT
+
 def ensure_pin_budget(size, evict_active=False):
     if args.high_ram:
         return True

comfy/ops.py

@@ -1104,6 +1104,21 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
                 scales["convrot_groupsize"] = int(
                     layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
                 )
+        elif module.quant_format == "convrot_w4a4":
+            scale = pop_scale("weight_scale")
+            if scale is None:
+                raise ValueError(f"Missing ConvRot W4A4 weight scale for layer {layer_name}")
+            params_conf = layer_conf.get("params", {})
+            if not isinstance(params_conf, dict):
+                params_conf = {}
+            scales = {
+                "scale": scale,
+                "convrot_groupsize": int(
+                    layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
+                ),
+                "quant_group_size": 64,
+                "linear_dtype": layer_conf.get("linear_dtype", params_conf.get("linear_dtype", "int4")),
+            }
         else:
             raise ValueError(f"Unsupported quantization format: {module.quant_format}")
 
@@ -1150,6 +1165,11 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
         if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False):
             quant_conf["convrot"] = True
             quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
+        elif module.quant_format == "convrot_w4a4":
+            quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
+            linear_dtype = getattr(params, "linear_dtype", "int4")
+            if linear_dtype != "int4":
+                quant_conf["linear_dtype"] = linear_dtype
         if extra_quant_conf:
             quant_conf.update(extra_quant_conf)
         sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
@@ -1237,7 +1257,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 run_every_op()
 
                 input_shape = input.shape
-                reshaped_3d = False
+                reshaped_nd = False
                 #If cast needs to apply lora, it should be done in the compute dtype
                 compute_dtype = input.dtype
 
@@ -1274,12 +1294,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 # Inference path (unchanged)
                 if _use_quantized and quantize_input:
 
-                    # Reshape 3D tensors to 2D for quantization (needed for NVFP4 and others)
-                    input_reshaped = input.reshape(-1, input_shape[2]) if input.ndim == 3 else input
+                    # Reshape >=3D tensors to 2D for quantization (needed for NVFP4 and others)
+                    input_reshaped = input.reshape(-1, input_shape[-1]) if input.ndim >= 3 else input
 
                     # Fall back to non-quantized for non-2D tensors
                     if input_reshaped.ndim == 2:
-                        reshaped_3d = input.ndim == 3
+                        reshaped_nd = input.ndim >= 3
                         # dtype is now implicit in the layout class
                         scale = getattr(self, 'input_scale', None)
                         if scale is not None:
@@ -1294,9 +1314,9 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                     weight_only_quant=weight_only_quant,
                 )
 
-                # Reshape output back to 3D if input was 3D
-                if reshaped_3d:
-                    output = output.reshape((input_shape[0], input_shape[1], self.weight.shape[0]))
+                # Reshape output back to original rank if input was >2D
+                if reshaped_nd:
+                    output = output.reshape((*input_shape[:-1], self.weight.shape[0]))
 
                 return output
 
@@ -1430,6 +1450,12 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 }
                 if hasattr(params, "block_scale"): # NVFP4
                     kwargs["block_scale"] = params.block_scale[i]
+                if hasattr(params, "quant_group_size"):
+                    kwargs["quant_group_size"] = params.quant_group_size
+                if hasattr(params, "convrot_groupsize"):
+                    kwargs["convrot_groupsize"] = params.convrot_groupsize
+                if hasattr(params, "linear_dtype"):
+                    kwargs["linear_dtype"] = params.linear_dtype
                 return QuantizedTensor(weight._qdata[i], weight._layout_cls, type(params)(**kwargs))
 
             def state_dict(self, *args, destination=None, prefix="", **kwargs):

comfy/quant_ops.py

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

comfy/sd.py

@@ -16,6 +16,7 @@ import comfy.ldm.cosmos.vae
 import comfy.ldm.wan.vae
 import comfy.ldm.wan.vae2_2
 import comfy.ldm.hunyuan3d.vae
+import comfy.ldm.seedvr.vae
 import comfy.ldm.triposplat.vae
 import comfy.ldm.ace.vae.music_dcae_pipeline
 import comfy.ldm.cogvideo.vae
@@ -473,7 +474,8 @@ class CLIP:
 
 class VAE:
     def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None):
-        if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
+        is_seedvr2_vae = "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd
+        if not is_seedvr2_vae and 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format
             sd = diffusers_convert.convert_vae_state_dict(sd)
 
         if model_management.is_amd():
@@ -500,6 +502,8 @@ class VAE:
         self.upscale_index_formula = None
         self.extra_1d_channel = None
         self.crop_input = True
+        self.handles_tiling = False
+        self.format_encoded = None
 
         self.audio_sample_rate = 44100
 
@@ -546,6 +550,22 @@ class VAE:
                 self.first_stage_model = StageC_coder()
                 self.downscale_ratio = 32
                 self.latent_channels = 16
+            elif "decoder.up_blocks.2.upsamplers.0.upscale_conv.weight" in sd: # seedvr2
+                self.first_stage_model = comfy.ldm.seedvr.vae.VideoAutoencoderKLWrapper()
+                self.latent_channels = comfy.ldm.seedvr.vae.SEEDVR2_LATENT_CHANNELS
+                self.latent_dim = 3
+                self.disable_offload = True
+                self.memory_used_decode = lambda shape, dtype: self.first_stage_model.comfy_memory_used_decode(shape)
+                self.memory_used_encode = lambda shape, dtype: (max(shape[2], 5) * shape[3] * shape[4] * 64) * model_management.dtype_size(dtype)
+                self.working_dtypes = [torch.float16, torch.bfloat16, torch.float32]
+                self.handles_tiling = True
+                self.format_encoded = self.first_stage_model.comfy_format_encoded
+                self.downscale_ratio = (lambda a: max(0, math.floor((a + 3) / 4)), 8, 8)
+                self.downscale_index_formula = (4, 8, 8)
+                self.upscale_ratio = (lambda a: max(0, a * 4 - 3), 8, 8)
+                self.upscale_index_formula = (4, 8, 8)
+                self.process_input = lambda image: image * 2.0 - 1.0
+                self.crop_input = False
             elif "decoder.conv_in.weight" in sd:
                 if sd['decoder.conv_in.weight'].shape[1] == 64:
                     ddconfig = {"block_out_channels": [128, 256, 512, 512, 1024, 1024], "in_channels": 3, "out_channels": 3, "num_res_blocks": 2, "ffactor_spatial": 32, "downsample_match_channel": True, "upsample_match_channel": True}
@@ -1012,6 +1032,10 @@ class VAE:
         decode_fn = lambda a: self.first_stage_model.decode(a.to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
         return self.process_output(comfy.utils.tiled_scale_multidim(samples, decode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.upscale_ratio, out_channels=self.output_channels, index_formulas=self.upscale_index_formula, output_device=self.output_device))
 
+    def _decode_tiled_owned(self, samples, **kwargs):
+        out = self.first_stage_model.decode_tiled(samples.to(self.vae_dtype).to(self.device), **kwargs)
+        return self.process_output(out.to(device=self.output_device, dtype=self.vae_output_dtype(), copy=True))
+
     def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
         steps = pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x, tile_y, overlap)
         steps += pixel_samples.shape[0] * comfy.utils.get_tiled_scale_steps(pixel_samples.shape[3], pixel_samples.shape[2], tile_x // 2, tile_y * 2, overlap)
@@ -1048,6 +1072,25 @@ class VAE:
         encode_fn = lambda a: self.first_stage_model.encode((self.process_input(a)).to(self.vae_dtype).to(self.device)).to(dtype=self.vae_output_dtype())
         return comfy.utils.tiled_scale_multidim(samples, encode_fn, tile=(tile_t, tile_x, tile_y), overlap=overlap, upscale_amount=self.downscale_ratio, out_channels=self.latent_channels, downscale=True, index_formulas=self.downscale_index_formula, output_device=self.output_device)
 
+    def _encode_tiled_owned(self, pixel_samples, **kwargs):
+        x = self.process_input(pixel_samples).to(self.vae_dtype).to(self.device)
+        out = self.first_stage_model.encode_tiled(x, **kwargs)
+        return out.to(device=self.output_device, dtype=self.vae_output_dtype())
+
+    def _owned_tiled_args(self, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
+        args = {}
+        if tile_x is not None:
+            args["tile_x"] = tile_x
+        if tile_y is not None:
+            args["tile_y"] = tile_y
+        if overlap is not None:
+            args["overlap"] = overlap
+        if tile_t is not None:
+            args["tile_t"] = tile_t
+        if overlap_t is not None:
+            args["overlap_t"] = overlap_t
+        return args
+
     def decode(self, samples_in, vae_options={}):
         self.throw_exception_if_invalid()
         pixel_samples = None
@@ -1095,11 +1138,19 @@ class VAE:
                 if dims == 1 or self.extra_1d_channel is not None:
                     pixel_samples = self.decode_tiled_1d(samples_in)
                 elif dims == 2:
-                    pixel_samples = self.decode_tiled_(samples_in)
+                    if self.handles_tiling:
+                        tile = 256 // self.spacial_compression_decode()
+                        overlap = tile // 4
+                        pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
+                    else:
+                        pixel_samples = self.decode_tiled_(samples_in)
                 elif dims == 3:
                     tile = 256 // self.spacial_compression_decode()
                     overlap = tile // 4
-                    pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+                    if self.handles_tiling:
+                        pixel_samples = self._decode_tiled_owned(samples_in, tile_x=tile, tile_y=tile, overlap=overlap)
+                    else:
+                        pixel_samples = self.decode_tiled_3d(samples_in, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
 
         pixel_samples = pixel_samples.to(self.output_device).movedim(1,-1)
         return pixel_samples
@@ -1118,7 +1169,9 @@ class VAE:
             args["overlap"] = overlap
 
         with model_management.cuda_device_context(self.device):
-            if dims == 1 or self.extra_1d_channel is not None:
+            if self.handles_tiling and dims in (2, 3):
+                output = self._decode_tiled_owned(samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
+            elif dims == 1 or self.extra_1d_channel is not None:
                 args.pop("tile_y")
                 output = self.decode_tiled_1d(samples, **args)
             elif dims == 2:
@@ -1179,12 +1232,17 @@ class VAE:
                 if self.latent_dim == 3:
                     tile = 256
                     overlap = tile // 4
-                    samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
+                    if self.handles_tiling:
+                        samples = self._encode_tiled_owned(pixel_samples, tile_x=tile, tile_y=tile, overlap=overlap)
+                    else:
+                        samples = self.encode_tiled_3d(pixel_samples, tile_x=tile, tile_y=tile, overlap=(1, overlap, overlap))
                 elif self.latent_dim == 1 or self.extra_1d_channel is not None:
                     samples = self.encode_tiled_1d(pixel_samples)
                 else:
                     samples = self.encode_tiled_(pixel_samples)
 
+        if self.format_encoded is not None:
+            samples = self.format_encoded(samples)
         return samples
 
     def encode_tiled(self, pixel_samples, tile_x=None, tile_y=None, overlap=None, tile_t=None, overlap_t=None):
@@ -1192,7 +1250,7 @@ class VAE:
         pixel_samples = self.vae_encode_crop_pixels(pixel_samples)
         dims = self.latent_dim
         pixel_samples = pixel_samples.movedim(-1, 1)
-        if dims == 3:
+        if dims == 3 and pixel_samples.ndim < 5:
             if not self.not_video:
                 pixel_samples = pixel_samples.movedim(1, 0).unsqueeze(0)
             else:
@@ -1216,21 +1274,27 @@ class VAE:
             elif dims == 2:
                 samples = self.encode_tiled_(pixel_samples, **args)
             elif dims == 3:
-                if tile_t is not None:
-                    tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
+                if self.handles_tiling:
+                    samples = self._encode_tiled_owned(pixel_samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
                 else:
-                    tile_t_latent = 9999
-                args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
+                    if tile_t is not None:
+                        tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
+                    else:
+                        tile_t_latent = 9999
+                    args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
 
-                if overlap_t is None:
-                    args["overlap"] = (1, overlap, overlap)
-                else:
-                    args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap)
-                maximum = pixel_samples.shape[2]
-                maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
+                    spatial_overlap = overlap if overlap is not None else 64
+                    if overlap_t is None:
+                        args["overlap"] = (1, spatial_overlap, spatial_overlap)
+                    else:
+                        args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), spatial_overlap, spatial_overlap)
+                    maximum = pixel_samples.shape[2]
+                    maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
 
-                samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
+                    samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
 
+        if self.format_encoded is not None:
+            samples = self.format_encoded(samples)
         return samples
 
     def get_sd(self):
@@ -1898,7 +1962,7 @@ def load_state_dict_guess_config(sd, output_vae=True, output_clip=True, output_c
         manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
     else:
         manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if model_config.clip_vision_prefix is not None:
         if output_clipvision:
@@ -2039,7 +2103,7 @@ def load_diffusion_model_state_dict(sd, model_options={}, metadata=None, disable
         manual_cast_dtype = model_management.unet_manual_cast(None, load_device, model_config.supported_inference_dtypes)
     else:
         manual_cast_dtype = model_management.unet_manual_cast(unet_dtype, load_device, model_config.supported_inference_dtypes)
-    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype)
+    model_config.set_inference_dtype(unet_dtype, manual_cast_dtype, device=load_device)
 
     if custom_operations is not None:
         model_config.custom_operations = custom_operations

comfy/supported_models.py

@@ -1685,6 +1685,40 @@ class Chroma(supported_models_base.BASE):
         t5_detect = comfy.text_encoders.sd3_clip.t5_xxl_detect(state_dict, "{}t5xxl.transformer.".format(pref))
         return supported_models_base.ClipTarget(comfy.text_encoders.pixart_t5.PixArtTokenizer, comfy.text_encoders.pixart_t5.pixart_te(**t5_detect))
 
+class SeedVR2(supported_models_base.BASE):
+    unet_config = {
+        "image_model": "seedvr2"
+    }
+    unet_extra_config = {}
+    required_keys = {
+        "{}positive_conditioning",
+        "{}negative_conditioning",
+    }
+    latent_format = comfy.latent_formats.SeedVR2
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+    supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
+    sampling_settings = {
+        "shift": 1.0,
+    }
+
+    def set_inference_dtype(self, dtype, manual_cast_dtype, device=None):
+        if (
+            dtype == torch.float16
+            and manual_cast_dtype is None
+            and comfy.model_management.should_use_bf16(device)
+        ):
+            manual_cast_dtype = torch.bfloat16
+        super().set_inference_dtype(dtype, manual_cast_dtype, device=device)
+
+    def get_model(self, state_dict, prefix="", device=None):
+        out = model_base.SeedVR2(self, device=device)
+        return out
+
+    def clip_target(self, state_dict={}):
+        return None
+
 class ChromaRadiance(Chroma):
     unet_config = {
         "image_model": "chroma_radiance",
@@ -2348,6 +2382,7 @@ models = [
     HiDream,
     HiDreamO1,
     Chroma,
+    SeedVR2,
     ChromaRadiance,
     ACEStep,
     ACEStep15,

comfy/supported_models_base.py

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

comfy/text_encoders/gemma4.py

@@ -1088,7 +1088,7 @@ class Gemma4_Tokenizer():
             h, w = samples.shape[2], samples.shape[3]
             patch_size = 16
             pooling_k = 3
-            max_soft_tokens = 70 if is_video else 280  # video uses smaller token budget per frame
+            max_soft_tokens = kwargs.get("max_soft_tokens", 70 if is_video else 280)
             max_patches = max_soft_tokens * pooling_k * pooling_k
             target_px = max_patches * patch_size * patch_size
             factor = (target_px / (h * w)) ** 0.5

comfy/text_encoders/qwen3vl.py

@@ -90,6 +90,27 @@ class Qwen3VL(BaseLlama, BaseQwen3, BaseGenerate, torch.nn.Module):
                 deepstack = [torch.cat([deepstack[i], ds[i]], dim=0) for i in range(len(ds))]
         return position_ids, visual_pos_masks, deepstack
 
+    def forward(self, input_ids, attention_mask=None, embeds=None, num_tokens=None, intermediate_output=None, final_layer_norm_intermediate=True, dtype=None, embeds_info=[], **kwargs):
+        position_ids = kwargs.pop("position_ids", None)
+        visual_pos_masks = kwargs.pop("visual_pos_masks", None)
+        deepstack_embeds = kwargs.pop("deepstack_embeds", None)
+        if embeds is not None and position_ids is None:
+            position_ids, visual_pos_masks, deepstack_embeds = self.build_image_inputs(embeds, embeds_info)
+        return self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            embeds=embeds,
+            num_tokens=num_tokens,
+            intermediate_output=intermediate_output,
+            final_layer_norm_intermediate=final_layer_norm_intermediate,
+            dtype=dtype,
+            position_ids=position_ids,
+            embeds_info=embeds_info,
+            visual_pos_masks=visual_pos_masks,
+            deepstack_embeds=deepstack_embeds,
+            **kwargs,
+        )
+
 
 def _make_qwen3vl_model(model_type):
     class Qwen3VL_(Qwen3VL):

comfy_api/feature_flags.py

@@ -100,10 +100,12 @@ def _parse_cli_feature_flags() -> dict[str, Any]:
 # Default server capabilities
 _CORE_FEATURE_FLAGS: dict[str, Any] = {
     "supports_preview_metadata": True,
+    "supports_model_type_tags": True,
     "max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes
     "extension": {"manager": {"supports_v4": True}},
     "node_replacements": True,
     "assets": args.enable_assets,
+    "server_side_model_downloads": True,
 }
 
 # CLI-provided flags cannot overwrite core flags

comfy_api_nodes/apis/bytedance.py

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

comfy_api_nodes/apis/hunyuan3d.py

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

comfy_api_nodes/nodes_bytedance.py

@@ -34,6 +34,7 @@ from comfy_api_nodes.apis.bytedance import (
     SeedanceVirtualLibraryCreateAssetRequest,
     Seedream4Options,
     Seedream4TaskCreationRequest,
+    Seedream5OptimizePromptOptions,
     TaskAudioContent,
     TaskAudioContentUrl,
     TaskCreationResponse,
@@ -875,6 +876,17 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                     tooltip='Whether to add an "AI generated" watermark to the image.',
                     advanced=True,
                 ),
+                IO.Boolean.Input(
+                    "thinking",
+                    default=True,
+                    tooltip=(
+                        "Enable the model's prompt-optimization reasoning ('thinking') for better adherence. "
+                        "Can substantially increase generation time — notably on Seedream 5.0 Pro. "
+                        "Can only be disabled for text-to-image (not when reference images are provided)."
+                    ),
+                    optional=True,
+                    advanced=True,
+                ),
             ],
             outputs=[
                 IO.Image.Output(),
@@ -920,6 +932,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
         model: dict,
         seed: int = 0,
         watermark: bool = False,
+        thinking: bool = True,
     ) -> IO.NodeOutput:
         validate_string(prompt, strip_whitespace=True, min_length=1)
         model_id = SEEDREAM_MODELS[model["model"]]
@@ -979,6 +992,10 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
             raise ValueError(
                 "The maximum number of generated images plus the number of reference images cannot exceed 15."
             )
+        if not thinking and n_input_images > 0:
+            raise ValueError(
+                "'thinking' can only be disabled for text-to-image; enable it when using reference images."
+            )
 
         reference_images_urls: list[str] = []
         if image_tensors:
@@ -992,6 +1009,9 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                 wait_label="Uploading reference images",
             )
 
+        optimize_prompt_options = None
+        if n_input_images == 0:
+            optimize_prompt_options = Seedream5OptimizePromptOptions(thinking="enabled" if thinking else "disabled")
         response = await sync_op(
             cls,
             ApiEndpoint(path=BYTEPLUS_IMAGE_ENDPOINT, method="POST"),
@@ -1005,6 +1025,7 @@ class ByteDanceSeedreamNodeV2(IO.ComfyNode):
                 sequential_image_generation=None if is_pro else sequential_image_generation,
                 sequential_image_generation_options=None if is_pro else Seedream4Options(max_images=max_images),
                 watermark=watermark,
+                optimize_prompt_options=optimize_prompt_options,
             ),
         )
         if len(response.data) == 1:

comfy_api_nodes/nodes_hunyuan3d.py

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

comfy_api_nodes/util/_helpers.py

@@ -11,6 +11,7 @@ from io import BytesIO
 from yarl import URL
 
 from comfy.cli_args import args
+from comfy.comfy_api_env import normalize_comfy_api_base
 from comfy.deploy_environment import get_deploy_environment
 from comfy.model_management import processing_interrupted
 from comfy_api.latest import IO
@@ -63,7 +64,7 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
 
 
 def default_base_url() -> str:
-    return getattr(args, "comfy_api_base", "https://api.comfy.org")
+    return normalize_comfy_api_base(getattr(args, "comfy_api_base", "https://api.comfy.org"))
 
 
 async def sleep_with_interrupt(

comfy_execution/caching.py

@@ -503,6 +503,8 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05
 
 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3
 
+RAM_CACHE_LARGE_INTERMEDIATE = 512 * 1024 ** 2
+
 
 def all_outputs_dynamic(outputs):
     if outputs is None:
@@ -517,7 +519,6 @@ def all_outputs_dynamic(outputs):
 
     return True
 
-
 class RAMPressureCache(LRUCache):
 
     def __init__(self, key_class, enable_providers=False):
@@ -539,9 +540,9 @@ class RAMPressureCache(LRUCache):
         self.timestamps[self.cache_key_set.get_data_key(node_id)] = time.time()
         super().set_local(node_id, value)
 
-    def ram_release(self, target, free_active=False):
+    def ram_release(self, target, free_active=False, min_entry_size=0):
         if psutil.virtual_memory().available >= target:
-            return
+            return 0
 
         clean_list = []
 
@@ -555,8 +556,9 @@ class RAMPressureCache(LRUCache):
             oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key])
 
             ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE
+            oom_ram_usage = ram_usage
             def scan_list_for_ram_usage(outputs):
-                nonlocal ram_usage
+                nonlocal ram_usage, oom_ram_usage
                 if outputs is None:
                     return
                 for output in outputs:
@@ -564,19 +566,26 @@ class RAMPressureCache(LRUCache):
                         scan_list_for_ram_usage(output)
                     elif isinstance(output, torch.Tensor) and output.device.type == 'cpu':
                         ram_usage += output.numel() * output.element_size()
+                        oom_ram_usage += output.numel() * output.element_size()
                     elif isinstance(output, ModelPatcher) and self.used_generation[key] != self.generation:
                         #old ModelPatchers are the first to go
-                        ram_usage = 1e30
+                        oom_ram_usage = 1e30
             scan_list_for_ram_usage(cache_entry.outputs)
 
-            oom_score *= ram_usage
+            if ram_usage < min_entry_size:
+                continue
+
+            oom_score *= oom_ram_usage
             #In the case where we have no information on the node ram usage at all,
             #break OOM score ties on the last touch timestamp (pure LRU)
-            bisect.insort(clean_list, (oom_score, self.timestamps[key], key))
+            bisect.insort(clean_list, (oom_score, self.timestamps[key], key, ram_usage))
 
+        freed = 0
         while psutil.virtual_memory().available < target and clean_list:
-            _, _, key = clean_list.pop()
+            _, _, key, ram_usage = clean_list.pop()
             del self.cache[key]
             self.used_generation.pop(key, None)
             self.timestamps.pop(key, None)
             self.children.pop(key, None)
+            freed += ram_usage
+        return freed

comfy_execution/jobs.py

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

comfy_extras/nodes_audio.py

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

comfy_extras/nodes_bounding_boxes.py

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

comfy_extras/nodes_load_3d.py

@@ -61,14 +61,10 @@ class Load3D(IO.ComfyNode):
 
     @classmethod
     def execute(cls, model_file, image, **kwargs) -> IO.NodeOutput:
-        image_path = folder_paths.get_annotated_filepath(image['image'])
-        mask_path = folder_paths.get_annotated_filepath(image['mask'])
-        normal_path = folder_paths.get_annotated_filepath(image['normal'])
-
         load_image_node = nodes.LoadImage()
-        output_image, ignore_mask = load_image_node.load_image(image=image_path)
-        ignore_image, output_mask = load_image_node.load_image(image=mask_path)
-        normal_image, ignore_mask2 = load_image_node.load_image(image=normal_path)
+        output_image, ignore_mask = load_image_node.load_image(image=image['image'])
+        ignore_image, output_mask = load_image_node.load_image(image=image['mask'])
+        normal_image, ignore_mask2 = load_image_node.load_image(image=image['normal'])
 
         video = None
 
@@ -96,6 +92,7 @@ class Preview3D(IO.ComfyNode):
             search_aliases=["view mesh", "3d viewer"],
             display_name="Preview 3D & Animation",
             category="3d",
+            description="Preview a 3D model file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             inputs=[
@@ -140,6 +137,7 @@ class Preview3DAdvanced(IO.ComfyNode):
             display_name="Preview 3D (Advanced)",
             search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"],
             category="3d",
+            description="Preview a 3D model file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             inputs=[
@@ -197,6 +195,7 @@ class PreviewGaussianSplat(IO.ComfyNode):
             node_id="PreviewGaussianSplat",
             display_name="Preview Splat",
             category="3d",
+            description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             search_aliases=[
@@ -265,6 +264,7 @@ class PreviewPointCloud(IO.ComfyNode):
             node_id="PreviewPointCloud",
             display_name="Preview Point Cloud",
             category="3d",
+            description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.",
             is_experimental=True,
             is_output_node=True,
             search_aliases=[

comfy_extras/nodes_mask.py

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

comfy_extras/nodes_preview_any.py

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

comfy_extras/nodes_primitive.py

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

comfy_extras/nodes_save_3d.py

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

comfy_extras/nodes_seedvr.py

@@ -0,0 +1,614 @@
+import logging
+
+from typing_extensions import override
+from comfy_api.latest import ComfyExtension, io
+import torch
+
+import comfy.model_management
+from comfy.ldm.seedvr.color_fix import (
+    adain_color_transfer,
+    lab_color_transfer,
+    wavelet_color_transfer,
+)
+from comfy.ldm.seedvr.constants import (
+    BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE,
+    SEEDVR2_ADAIN_SCALE_MULTIPLIER,
+    SEEDVR2_CHUNK_GIB_PER_MPX_FRAME,
+    SEEDVR2_CHUNK_RESERVED_GIB,
+    SEEDVR2_CHUNK_SIGMA_GIB,
+    SEEDVR2_CHUNK_SIGMA_K,
+    SEEDVR2_COLOR_MEM_HEADROOM,
+    SEEDVR2_DTYPE_BYTES_FLOOR,
+    SEEDVR2_LAB_SCALE_MULTIPLIER,
+    SEEDVR2_LATENT_CHANNELS,
+    SEEDVR2_OOM_BACKOFF_DIVISOR,
+    SEEDVR2_WAVELET_SCALE_MULTIPLIER,
+)
+
+from torchvision.transforms import functional as TVF
+from torchvision.transforms.functional import InterpolationMode
+
+
+_SEEDVR2_INVALID_MODEL_MSG_PREFIX = "SeedVR2Conditioning: model object does not match expected SeedVR2 structure"
+_ATTR_MISSING = object()
+
+
+def _resolve_seedvr2_diffusion_model(model):
+    inner = getattr(model, "model", _ATTR_MISSING)
+    if inner is _ATTR_MISSING:
+        raise RuntimeError(
+            f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: input has no 'model' attribute "
+            f"(got type {type(model).__name__})."
+        )
+    if inner is None:
+        raise RuntimeError(
+            f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: input.model is None "
+            f"(input type {type(model).__name__})."
+        )
+    diffusion_model = getattr(inner, "diffusion_model", _ATTR_MISSING)
+    if diffusion_model is _ATTR_MISSING:
+        raise RuntimeError(
+            f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: 'model.model' has no "
+            f"'diffusion_model' attribute (got type {type(inner).__name__})."
+        )
+    if diffusion_model is None:
+        raise RuntimeError(
+            f"{_SEEDVR2_INVALID_MODEL_MSG_PREFIX}: 'model.model.diffusion_model' "
+            f"is None (model.model type {type(inner).__name__})."
+        )
+    return diffusion_model
+
+
+def div_pad(image, factor):
+    height_factor, width_factor = factor
+    height, width = image.shape[-2:]
+
+    pad_height = (height_factor - (height % height_factor)) % height_factor
+    pad_width = (width_factor - (width % width_factor)) % width_factor
+
+    if pad_height == 0 and pad_width == 0:
+        return image
+
+    padding = (0, pad_width, 0, pad_height)
+    return torch.nn.functional.pad(image, padding, mode='constant', value=0.0)
+
+def cut_videos(videos):
+    t = videos.size(1)
+    if t < 1:
+        raise ValueError("SeedVR2Preprocess expected at least one frame.")
+    if t == 1:
+        return videos
+    if t <= 4:
+        padding = videos[:, -1:].repeat(1, 4 - t + 1, 1, 1, 1)
+        return torch.cat([videos, padding], dim=1)
+    if (t - 1) % 4 == 0:
+        return videos
+    padding = videos[:, -1:].repeat(1, 4 - ((t - 1) % 4), 1, 1, 1)
+    videos = torch.cat([videos, padding], dim=1)
+    if (videos.size(1) - 1) % 4 != 0:
+        raise ValueError(f"SeedVR2Preprocess failed to pad video length to 4n+1; got {videos.size(1)} frames.")
+    return videos
+
+def _seedvr2_input_shorter_edge(images, node_name):
+    if images.dim() == 4:
+        return min(images.shape[1], images.shape[2])
+    if images.dim() == 5:
+        return min(images.shape[2], images.shape[3])
+    raise ValueError(
+        f"{node_name}: expected 4-D or 5-D IMAGE tensor, "
+        f"got shape {tuple(images.shape)}"
+    )
+
+
+def _seedvr2_pad(images, upscaled_shorter_edge, node_name):
+    if upscaled_shorter_edge < 2:
+        raise ValueError(
+            f"{node_name}: input shorter edge must be at least 2 pixels; "
+            f"got {upscaled_shorter_edge}."
+        )
+    if images.shape[-1] > 3:
+        images = images[..., :3]
+    if images.dim() == 4:
+        # Comfy video components arrive as a 4-D IMAGE frame sequence:
+        # (frames, H, W, C). SeedVR2 consumes that as one video.
+        images = images.unsqueeze(0)
+    elif images.dim() != 5:
+        raise ValueError(
+            f"{node_name}: expected 4-D or 5-D IMAGE tensor, "
+            f"got shape {tuple(images.shape)}"
+        )
+    images = images.permute(0, 1, 4, 2, 3)
+
+    b, t, c, h, w = images.shape
+    images = images.reshape(b * t, c, h, w)
+
+    images = torch.clamp(images, 0.0, 1.0)
+    images = div_pad(images, (16, 16))
+    _, _, new_h, new_w = images.shape
+
+    images = images.reshape(b, t, c, new_h, new_w)
+    images = cut_videos(images)
+    images_bthwc = images.permute(0, 1, 3, 4, 2).contiguous()
+
+    return io.NodeOutput(images_bthwc)
+
+
+class SeedVR2Preprocess(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SeedVR2Preprocess",
+            display_name="Pre-Process SeedVR2 Input",
+            category="image/pre-processors",
+            description="Pad a resized image for SeedVR2 model. Alpha channel is dropped. The node Post-Process SeedVR2 Output re-applies it from the original resized image.",
+            search_aliases=["seedvr2", "upscale", "video upscale", "pad", "preprocess"],
+            inputs=[
+                io.Image.Input("resized_images", tooltip="The resized image to process."),
+            ],
+            outputs=[
+                io.Image.Output("images", tooltip="The padded image for VAE encoding."),
+            ]
+        )
+
+    @classmethod
+    def execute(cls, resized_images):
+        upscaled_shorter_edge = _seedvr2_input_shorter_edge(resized_images, "SeedVR2Preprocess")
+        return _seedvr2_pad(
+            resized_images, upscaled_shorter_edge, "SeedVR2Preprocess",
+        )
+
+
+class SeedVR2PostProcessing(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SeedVR2PostProcessing",
+            display_name="Post-Process SeedVR2 Output",
+            category="image/post-processors",
+            description="Align the generated image with the original resized image and apply color correction.",
+            search_aliases=["seedvr2", "upscale", "color correction", "color match", "postprocess"],
+            inputs=[
+                io.Image.Input("images", tooltip="The generated image to process."),
+                io.Image.Input("original_resized_images", tooltip="The original resized image before pre-processing, used as reference."),
+                io.Combo.Input("color_correction_method", options=["lab", "wavelet", "adain", "none"], default="lab", tooltip="Method to match the generated image colors to the original image. lab: transfer color in CIELAB space, preserving detail (most faithful). wavelet: transfer low-frequency color, keeping upscaled high-frequency detail. adain: match per-channel mean/std (fastest, global tint). none: skip color transfer (geometry alignment only)."),
+            ],
+            outputs=[io.Image.Output(display_name="images", tooltip="The aligned, color-corrected image.")],
+        )
+
+    @classmethod
+    def execute(cls, images, original_resized_images, color_correction_method):
+        alpha_input = None
+        if original_resized_images.shape[-1] == 4:
+            alpha_input = original_resized_images[..., 3:4]
+            original_resized_images = original_resized_images[..., :3]
+        decoded_5d, decoded_was_4d = cls._as_bthwc(images)
+        reference_full, _ = cls._as_bthwc(original_resized_images)
+        decoded_5d = cls._restore_reference_batch_time(decoded_5d, reference_full)
+
+        b = min(decoded_5d.shape[0], reference_full.shape[0])
+        t = min(decoded_5d.shape[1], reference_full.shape[1])
+        reference_h = reference_full.shape[2]
+        reference_w = reference_full.shape[3]
+
+        decoded_5d = decoded_5d[:b, :t, :, :, :]
+        target_h = min(decoded_5d.shape[2], reference_h)
+        target_w = min(decoded_5d.shape[3], reference_w)
+        decoded_5d = decoded_5d[:, :, :target_h, :target_w, :]
+        if color_correction_method in ("lab", "wavelet", "adain"):
+            reference_5d = reference_full[:b, :t, :, :, :]
+            reference_5d = cls._resize_reference(reference_5d, target_h, target_w)
+            output_device = decoded_5d.device
+            decoded_raw = cls._to_seedvr2_raw(decoded_5d)
+            reference_raw = cls._to_seedvr2_raw(reference_5d)
+            decoded_flat = decoded_raw.permute(0, 1, 4, 2, 3).reshape(b * t, decoded_raw.shape[4], target_h, target_w)
+            reference_flat = reference_raw.permute(0, 1, 4, 2, 3).reshape(b * t, reference_raw.shape[4], target_h, target_w)
+            output = cls._color_transfer_chunked(
+                decoded_flat, reference_flat, output_device, color_correction_method,
+            )
+            output = output.reshape(b, t, output.shape[1], output.shape[2], output.shape[3]).permute(0, 1, 3, 4, 2)
+            output = output.add(1.0).div(2.0).clamp(0.0, 1.0)
+        elif color_correction_method == "none":
+            output = decoded_5d
+        else:
+            raise ValueError(f"SeedVR2PostProcessing: unknown color_correction_method {color_correction_method!r}")
+
+        if alpha_input is not None:
+            alpha_5d, _ = cls._as_bthwc(alpha_input)
+            alpha_5d = alpha_5d[:output.shape[0], :output.shape[1], :output.shape[2], :output.shape[3], :]
+            output = torch.cat([output, alpha_5d.to(dtype=output.dtype, device=output.device)], dim=-1)
+        h2 = output.shape[-3] - (output.shape[-3] % 2)
+        w2 = output.shape[-2] - (output.shape[-2] % 2)
+        output = output[:, :, :h2, :w2, :]
+        if decoded_was_4d:
+            output = output.reshape(-1, output.shape[-3], output.shape[-2], output.shape[-1])
+        return io.NodeOutput(output)
+
+    @staticmethod
+    def _as_bthwc(images):
+        if images.ndim == 4:
+            return images.unsqueeze(0), True
+        if images.ndim == 5:
+            return images, False
+        raise ValueError(
+            f"SeedVR2PostProcessing: expected 4-D or 5-D IMAGE tensor, got shape {tuple(images.shape)}"
+        )
+
+    @staticmethod
+    def _restore_reference_batch_time(decoded, reference):
+        if decoded.shape[0] != 1:
+            return decoded
+        ref_b, ref_t = reference.shape[:2]
+        if ref_b < 1 or decoded.shape[1] % ref_b != 0:
+            return decoded
+        decoded_t = decoded.shape[1] // ref_b
+        if decoded_t < ref_t:
+            return decoded
+        return decoded.reshape(ref_b, decoded_t, decoded.shape[2], decoded.shape[3], decoded.shape[4])
+
+    @staticmethod
+    def _to_seedvr2_raw(images):
+        return images.mul(2.0).sub(1.0)
+
+    @staticmethod
+    def _color_transfer_on_vae_device(decoded_flat, reference_flat, output_device, transfer_fn):
+        color_device = comfy.model_management.vae_device()
+        decoded_flat = decoded_flat.to(device=color_device)
+        reference_flat = reference_flat.to(device=color_device)
+        output = transfer_fn(decoded_flat, reference_flat)
+        return output.to(device=output_device)
+
+    @staticmethod
+    def _lab_color_transfer_on_vae_device(decoded_flat, reference_flat, output_device):
+        color_device = comfy.model_management.vae_device()
+        result = None
+        for start in range(decoded_flat.shape[0]):
+            decoded_frame = decoded_flat[start:start + 1].to(device=color_device).clone()
+            reference_frame = reference_flat[start:start + 1].to(device=color_device).clone()
+            output = lab_color_transfer(decoded_frame, reference_frame).to(device=output_device)
+            if result is None:
+                result = torch.empty(
+                    (decoded_flat.shape[0],) + tuple(output.shape[1:]),
+                    device=output_device,
+                    dtype=output.dtype,
+                )
+            result[start:start + 1].copy_(output)
+        if result is None:
+            raise ValueError("SeedVR2PostProcessing: LAB color correction requires at least one frame.")
+        return result
+
+    @classmethod
+    def _color_transfer_chunked(cls, decoded_flat, reference_flat, output_device, color_correction_method):
+        chunk_size = cls._estimate_color_correction_chunk_size(decoded_flat, color_correction_method)
+        while True:
+            try:
+                return cls._run_color_transfer_chunks(
+                    decoded_flat, reference_flat, output_device, color_correction_method, chunk_size,
+                )
+            except Exception as e:
+                comfy.model_management.raise_non_oom(e)
+                if chunk_size <= 1:
+                    raise RuntimeError(
+                        "SeedVR2PostProcessing: color correction OOM at one frame; "
+                        f"color_correction_method={color_correction_method}, shape={tuple(decoded_flat.shape)}."
+                    ) from e
+                chunk_size = max(1, chunk_size // SEEDVR2_OOM_BACKOFF_DIVISOR)
+
+    @classmethod
+    def _run_color_transfer_chunks(cls, decoded_flat, reference_flat, output_device, color_correction_method, chunk_size):
+        result = None
+        for start in range(0, decoded_flat.shape[0], chunk_size):
+            end = min(start + chunk_size, decoded_flat.shape[0])
+            decoded_chunk = decoded_flat[start:end]
+            reference_chunk = reference_flat[start:end]
+            if color_correction_method == "lab":
+                output = cls._lab_color_transfer_on_vae_device(decoded_chunk, reference_chunk, output_device)
+            elif color_correction_method == "wavelet":
+                output = cls._color_transfer_on_vae_device(
+                    decoded_chunk, reference_chunk, output_device, wavelet_color_transfer,
+                )
+            else:
+                output = cls._color_transfer_on_vae_device(
+                    decoded_chunk, reference_chunk, output_device, adain_color_transfer,
+                )
+            if result is None:
+                result = torch.empty(
+                    (decoded_flat.shape[0],) + tuple(output.shape[1:]),
+                    device=output_device,
+                    dtype=output.dtype,
+                )
+            result[start:end].copy_(output)
+        if result is None:
+            raise ValueError("SeedVR2PostProcessing: color correction requires at least one frame.")
+        return result
+
+    @classmethod
+    def _estimate_color_correction_chunk_size(cls, decoded_flat, color_correction_method):
+        multiplier = cls._color_correction_memory_multiplier(color_correction_method)
+        frames = decoded_flat.shape[0]
+        _, channels, height, width = decoded_flat.shape
+        dtype_bytes = max(decoded_flat.element_size(), SEEDVR2_DTYPE_BYTES_FLOOR)
+        bytes_per_frame = height * width * channels * dtype_bytes * multiplier
+        if bytes_per_frame <= 0:
+            return frames
+        color_device = comfy.model_management.vae_device()
+        free_memory = comfy.model_management.get_free_memory(color_device)
+        chunk_size = int((free_memory * SEEDVR2_COLOR_MEM_HEADROOM) // bytes_per_frame)
+        return max(1, min(frames, chunk_size))
+
+    @staticmethod
+    def _color_correction_memory_multiplier(color_correction_method):
+        if color_correction_method == "lab":
+            return SEEDVR2_LAB_SCALE_MULTIPLIER
+        if color_correction_method == "wavelet":
+            return SEEDVR2_WAVELET_SCALE_MULTIPLIER
+        if color_correction_method == "adain":
+            return SEEDVR2_ADAIN_SCALE_MULTIPLIER
+        raise ValueError(f"SeedVR2PostProcessing: unknown color_correction_method {color_correction_method!r}")
+
+    @staticmethod
+    def _resize_reference(reference, height, width):
+        if reference.shape[2] == height and reference.shape[3] == width:
+            return reference
+        b, t = reference.shape[:2]
+        reference_flat = reference.permute(0, 1, 4, 2, 3).reshape(b * t, reference.shape[4], reference.shape[2], reference.shape[3])
+        resized = TVF.resize(
+            reference_flat,
+            size=(height, width),
+            interpolation=InterpolationMode.BICUBIC,
+            antialias=not (isinstance(reference_flat, torch.Tensor) and reference_flat.device.type == "mps"),
+        )
+        return resized.reshape(b, t, resized.shape[1], height, width).permute(0, 1, 3, 4, 2)
+
+
+class SeedVR2Conditioning(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SeedVR2Conditioning",
+            display_name="Apply SeedVR2 Conditioning",
+            category="model/conditioning",
+            description="Build SeedVR2 positive/negative conditioning from a VAE latent.",
+            search_aliases=["seedvr2", "upscale", "conditioning"],
+            inputs=[
+                io.Model.Input("model", tooltip="The SeedVR2 model."),
+                io.Latent.Input("vae_conditioning", display_name="latent"),
+            ],
+            outputs=[
+                io.Conditioning.Output(display_name="positive", tooltip="The positive conditioning for sampling."),
+                io.Conditioning.Output(display_name="negative", tooltip="The negative conditioning for sampling."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, model, vae_conditioning) -> io.NodeOutput:
+
+        vae_conditioning = vae_conditioning["samples"]
+        if vae_conditioning.ndim != 5:
+            raise ValueError(
+                "SeedVR2Conditioning expects a 5-D VAE latent in Comfy "
+                f"channel-first layout; got shape {tuple(vae_conditioning.shape)}."
+            )
+        if vae_conditioning.shape[1] != SEEDVR2_LATENT_CHANNELS:
+            if vae_conditioning.shape[-1] == SEEDVR2_LATENT_CHANNELS:
+                raise ValueError(
+                    "SeedVR2Conditioning expects SeedVR2 VAE latents in Comfy "
+                    f"channel-first layout (B, {SEEDVR2_LATENT_CHANNELS}, T, H, W); "
+                    f"got channel-last shape {tuple(vae_conditioning.shape)}."
+                )
+            raise ValueError(
+                "SeedVR2Conditioning expects SeedVR2 VAE latents with "
+                f"{SEEDVR2_LATENT_CHANNELS} channels; got shape {tuple(vae_conditioning.shape)}."
+            )
+        vae_conditioning = vae_conditioning.movedim(1, -1).contiguous()
+        model = _resolve_seedvr2_diffusion_model(model)
+        pos_cond = model.positive_conditioning
+        neg_cond = model.negative_conditioning
+
+        mask = vae_conditioning.new_ones(vae_conditioning.shape[:-1] + (1,))
+        condition = torch.cat((vae_conditioning, mask), dim=-1)
+        condition = condition.movedim(-1, 1)
+
+        negative = [[neg_cond.unsqueeze(0), {"condition": condition}]]
+        positive = [[pos_cond.unsqueeze(0), {"condition": condition}]]
+
+        return io.NodeOutput(positive, negative)
+
+def _seedvr2_chunk_crossfade_weights(overlap, device, dtype):
+    """Descending previous-chunk weights across the overlap (next chunk gets ``1 - w``): a Hann fade over the middle third, flat shoulders on the outer thirds."""
+    ramp = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype)
+    ramp = ((ramp - 1.0 / 3.0) / (1.0 / 3.0)).clamp(0.0, 1.0)
+    return 0.5 + 0.5 * torch.cos(torch.pi * ramp)
+
+
+class SeedVR2TemporalChunk(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SeedVR2TemporalChunk",
+            display_name="Split SeedVR2 Latent",
+            category="model/latent/batch",
+            description="Split a SeedVR2 video latent into overlapping temporal chunks small enough to sample one at a time within VRAM, wiring latents outputs to both Apply SeedVR2 Conditioning and the sampler latent input before recombining with Merge SeedVR2 Latents.",
+            search_aliases=["seedvr2", "split", "chunk", "temporal", "video upscale", "rebatch"],
+            inputs=[
+                io.Latent.Input("latent", tooltip="The VAE-encoded SeedVR2 latent to split."),
+                io.Int.Input("temporal_overlap", default=0, min=0, max=16384,
+                             tooltip="Latent frames shared between adjacent chunks and crossfaded at merge; 0 = no overlap."),
+                io.DynamicCombo.Input("chunking_mode",
+                                      tooltip="manual = use frames_per_chunk exactly; auto = predict the largest chunk that fits free VRAM.",
+                                      options=[
+                                          io.DynamicCombo.Option("auto", []),
+                                          io.DynamicCombo.Option("manual", [
+                                              io.Int.Input("frames_per_chunk", default=21, min=1, max=16384, step=4,
+                                                           tooltip="Pixel frames per temporal chunk (4n+1: 1, 5, 9, 13, ...)."),
+                                          ]),
+                                      ]),
+            ],
+            outputs=[
+                io.Latent.Output(display_name="latents", is_output_list=True,
+                                 tooltip="The temporal chunks in sequence order."),
+                io.Int.Output(display_name="temporal_overlap",
+                              tooltip="The effective latent-frame overlap between adjacent chunks, for Merge SeedVR2 Latents."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, latent, temporal_overlap, chunking_mode) -> io.NodeOutput:
+        samples = latent["samples"]
+        if samples.ndim != 5:
+            raise ValueError(
+                f"SeedVR2TemporalChunk: expected a 5-D video latent (B, C, T, H, W); "
+                f"got shape {tuple(samples.shape)}."
+            )
+        if samples.shape[1] != SEEDVR2_LATENT_CHANNELS:
+            raise ValueError(
+                f"SeedVR2TemporalChunk: expected {SEEDVR2_LATENT_CHANNELS} latent channels; "
+                f"got shape {tuple(samples.shape)}."
+            )
+        if temporal_overlap < 0:
+            raise ValueError(
+                f"SeedVR2TemporalChunk: temporal_overlap must be >= 0; got {temporal_overlap}."
+            )
+        mode = chunking_mode["chunking_mode"]
+        if mode not in ("auto", "manual"):
+            raise ValueError(
+                f"SeedVR2TemporalChunk: chunking_mode must be 'auto' or 'manual'; "
+                f"got {mode!r}."
+            )
+        t_latent = samples.shape[2]
+        t_pixel = 4 * (t_latent - 1) + 1
+
+        if mode == "auto":
+            free_gb = comfy.model_management.get_free_memory(
+                comfy.model_management.get_torch_device()) / (1024 ** 3)
+            mpx_per_frame = (samples.shape[0] * samples.shape[3] * samples.shape[4]) * (BYTEDANCE_VAE_SPATIAL_DOWNSAMPLE ** 2) / 1e6
+            budget_gb = free_gb - SEEDVR2_CHUNK_RESERVED_GIB - SEEDVR2_CHUNK_SIGMA_K * SEEDVR2_CHUNK_SIGMA_GIB
+            chunk_latent_max = max(1, int(budget_gb / (SEEDVR2_CHUNK_GIB_PER_MPX_FRAME * mpx_per_frame)))
+            frames_per_chunk = min(4 * (chunk_latent_max - 1) + 1, t_pixel)
+            logging.info(
+                "SeedVR2TemporalChunk auto: free=%.2fGiB, %.2fMpx -> frames_per_chunk=%d (t_pixel=%d).",
+                free_gb, mpx_per_frame, frames_per_chunk, t_pixel,
+            )
+        else:
+            frames_per_chunk = chunking_mode["frames_per_chunk"]
+            if frames_per_chunk < 1 or (frames_per_chunk - 1) % 4 != 0:
+                raise ValueError(
+                    f"SeedVR2TemporalChunk: frames_per_chunk must be a 4n+1 pixel-frame count "
+                    f"(1, 5, 9, 13, 17, 21, ...); got {frames_per_chunk}."
+                )
+
+        if t_pixel <= frames_per_chunk:
+            return io.NodeOutput([latent], 0)
+
+        chunk_latent = (frames_per_chunk - 1) // 4 + 1
+        temporal_overlap = min(temporal_overlap, chunk_latent - 1)
+        step = chunk_latent - temporal_overlap
+
+        chunks = []
+        for start in range(0, t_latent, step):
+            end = min(start + chunk_latent, t_latent)
+            chunk = latent.copy()
+            chunk["samples"] = samples[:, :, start:end].contiguous()
+            chunks.append(chunk)
+            if end >= t_latent:
+                break
+        return io.NodeOutput(chunks, temporal_overlap)
+
+
+class SeedVR2TemporalMerge(io.ComfyNode):
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SeedVR2TemporalMerge",
+            display_name="Merge SeedVR2 Latents",
+            category="model/latent/batch",
+            is_input_list=True,
+            description="Recombine sampled SeedVR2 latent temporal chunks into one latent, crossfading each overlap with a Hann window sized by the temporal_overlap wired from Split SeedVR2 Latent.",
+            search_aliases=["seedvr2", "merge", "temporal", "hann", "crossfade"],
+            inputs=[
+                io.Latent.Input("latents", tooltip="The sampled temporal chunks in sequence order."),
+                io.Int.Input("temporal_overlap", default=0, min=0, max=16384, force_input=True,
+                             tooltip="The temporal_overlap output of Split SeedVR2 Latent. 0 = plain concatenation."),
+            ],
+            outputs=[
+                io.Latent.Output(display_name="latent", tooltip="The recombined full-length latent."),
+            ],
+        )
+
+    @classmethod
+    def execute(cls, latents, temporal_overlap) -> io.NodeOutput:
+        temporal_overlap = temporal_overlap[0]
+        if temporal_overlap < 0:
+            raise ValueError(
+                f"SeedVR2TemporalMerge: temporal_overlap must be >= 0; got {temporal_overlap}."
+            )
+        chunks = [entry["samples"] for entry in latents]
+        first = chunks[0]
+        if first.ndim != 5:
+            raise ValueError(
+                f"SeedVR2TemporalMerge: expected 5-D video latents (B, C, T, H, W); "
+                f"chunk 0 has shape {tuple(first.shape)}."
+            )
+        for i, chunk in enumerate(chunks[1:], start=1):
+            if chunk.shape[:2] != first.shape[:2] or chunk.shape[3:] != first.shape[3:]:
+                raise ValueError(
+                    f"SeedVR2TemporalMerge: chunk {i} shape {tuple(chunk.shape)} does not "
+                    f"match chunk 0 shape {tuple(first.shape)} outside the temporal axis."
+                )
+            if i < len(chunks) - 1 and chunk.shape[2] != first.shape[2]:
+                raise ValueError(
+                    f"SeedVR2TemporalMerge: chunk {i} has {chunk.shape[2]} latent frames but "
+                    f"chunk 0 has {first.shape[2]}; only the final chunk may be shorter."
+                )
+
+        out = latents[0].copy()
+        out.pop("noise_mask", None)
+
+        if len(chunks) == 1:
+            out["samples"] = first
+            return io.NodeOutput(out)
+        if temporal_overlap == 0:
+            out["samples"] = torch.cat(chunks, dim=2)
+            return io.NodeOutput(out)
+
+        chunk_latent = first.shape[2]
+        step = chunk_latent - min(temporal_overlap, chunk_latent - 1)
+        t_total = step * (len(chunks) - 1) + chunks[-1].shape[2]
+        b, c, _, h, w = first.shape
+        merged = torch.empty((b, c, t_total, h, w), device=first.device, dtype=first.dtype)
+
+        merged[:, :, :chunk_latent] = first
+        filled = chunk_latent
+        for i, chunk in enumerate(chunks[1:], start=1):
+            start = i * step
+            end = start + chunk.shape[2]
+            # Crossfade width is bounded by the previous fill frontier and by a runt
+            # final chunk shorter than the configured overlap.
+            fade = min(filled - start, chunk.shape[2])
+            if fade > 0:
+                w_prev = _seedvr2_chunk_crossfade_weights(
+                    fade, chunk.device, chunk.dtype).view(1, 1, fade, 1, 1)
+                merged[:, :, start:start + fade] = (
+                    merged[:, :, start:start + fade] * w_prev + chunk[:, :, :fade] * (1.0 - w_prev)
+                )
+                merged[:, :, start + fade:end] = chunk[:, :, fade:]
+            else:
+                merged[:, :, start:end] = chunk
+            filled = end
+
+        out["samples"] = merged
+        return io.NodeOutput(out)
+
+
+class SeedVRExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            SeedVR2Conditioning,
+            SeedVR2Preprocess,
+            SeedVR2PostProcessing,
+            SeedVR2TemporalChunk,
+            SeedVR2TemporalMerge,
+        ]
+
+async def comfy_entrypoint() -> SeedVRExtension:
+    return SeedVRExtension()

comfy_extras/nodes_text.py

@@ -0,0 +1,71 @@
+import os
+import json
+from typing_extensions import override
+from comfy_api.latest import io, ComfyExtension, ui
+import folder_paths
+
+
+class SaveTextNode(io.ComfyNode):
+    """Save text content to .txt, .md, or .json."""
+
+    @classmethod
+    def define_schema(cls):
+        return io.Schema(
+            node_id="SaveText",
+            search_aliases=["save text", "write text", "export text"],
+            display_name="Save Text",
+            category="text",
+            description="Save text content to a file in the output directory.",
+            inputs=[
+                io.String.Input("text", force_input=True),
+                io.String.Input("filename_prefix", default="ComfyUI"),
+                io.Combo.Input("format", options=["txt", "md", "json"], default="txt"),
+            ],
+            outputs=[io.String.Output(display_name="text")],
+            is_output_node=True,
+        )
+
+    @classmethod
+    def execute(cls, text, filename_prefix, format):
+        full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
+            filename_prefix,
+            folder_paths.get_output_directory(),
+            1,
+            1,
+        )
+
+        file = f"{filename}_{counter:05}.{format}"
+        filepath = os.path.join(full_output_folder, file)
+
+        if format == "json":
+            # tries to pretty print otherwise saves normally
+            try:
+                data = json.loads(text)
+                with open(filepath, "w", encoding="utf-8") as f:
+                    json.dump(data, f, indent=2, ensure_ascii=False)
+            except json.JSONDecodeError:
+                with open(filepath, "w", encoding="utf-8") as f:
+                    f.write(text)
+        else:
+            with open(filepath, "w", encoding="utf-8") as f:
+                f.write(text)
+
+        return io.NodeOutput(
+            text,
+            ui={
+                "text": (text,),
+                "files": [
+                    ui.SavedResult(file, subfolder, io.FolderType.output)
+                ]
+            }
+        )
+
+class TextExtension(ComfyExtension):
+    @override
+    async def get_node_list(self) -> list[type[io.ComfyNode]]:
+        return [
+            SaveTextNode
+        ]
+
+async def comfy_entrypoint() -> TextExtension:
+    return TextExtension()

comfy_extras/nodes_video.py

@@ -81,7 +81,7 @@ class SaveVideo(io.ComfyNode):
             display_name="Save Video",
             category="video",
             essentials_category="Basics",
-            description="Saves the input images to your ComfyUI output directory.",
+            description="Saves the input videos to your ComfyUI output directory.",
             inputs=[
                 io.Video.Input("video", tooltip="The video to save."),
                 io.String.Input("filename_prefix", default="video/ComfyUI", tooltip="The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes."),

execution.py

@@ -29,6 +29,7 @@ from comfy_execution.caching import (
     HierarchicalCache,
     LRUCache,
     RAMPressureCache,
+    RAM_CACHE_LARGE_INTERMEDIATE,
 )
 from comfy_execution.graph import (
     DynamicPrompt,
@@ -794,12 +795,16 @@ class PromptExecutor:
                     if self.cache_type == CacheType.RAM_PRESSURE:
                         ram_release_callback(ram_inactive_headroom)
                         ram_shortfall = ram_headroom - psutil.virtual_memory().available
-                        freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
-                        if freed < ram_shortfall:
-                            if freed > 64 * (1024 ** 2):
-                                # AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
-                                time.sleep(0.05)
-                            ram_release_callback(ram_headroom, free_active=True)
+                        if ram_shortfall > 0:
+                            freed = ram_release_callback(ram_headroom, free_active=True, min_entry_size=RAM_CACHE_LARGE_INTERMEDIATE)
+                            ram_shortfall -= freed
+                        if comfy.model_management.should_free_pins_for_ram_pressure(ram_shortfall):
+                            freed = comfy.model_management.free_pins(ram_shortfall + 512 * (1024 ** 2))
+                            if freed < ram_shortfall:
+                                if freed > 64 * (1024 ** 2):
+                                    # AIMDO MEM_DECOMMIT can outrun psutil.available catching up.
+                                    time.sleep(0.05)
+                                ram_release_callback(ram_headroom, free_active=True)
                 else:
                     # Only execute when the while-loop ends without break
                     # Send cached UI for intermediate output nodes that weren't executed

nodes.py

@@ -1709,6 +1709,7 @@ class PreviewImage(SaveImage):
         self.compress_level = 1
 
     SEARCH_ALIASES = ["preview", "preview image", "show image", "view image", "display image", "image viewer"]
+    DESCRIPTION = "Preview the images without saving them to the ComfyUI output directory."
 
     @classmethod
     def INPUT_TYPES(s):
@@ -2458,6 +2459,7 @@ async def init_builtin_extra_nodes():
         "nodes_camera_trajectory.py",
         "nodes_edit_model.py",
         "nodes_tcfg.py",
+        "nodes_seedvr.py",
         "nodes_context_windows.py",
         "nodes_qwen.py",
         "nodes_boogu.py",
@@ -2503,6 +2505,7 @@ async def init_builtin_extra_nodes():
         "nodes_triposplat.py",
         "nodes_depth_anything_3.py",
         "nodes_seed.py",
+        "nodes_text.py",
     ]
 
     import_failed = []

openapi.yaml

@@ -7,18 +7,18 @@ components:
                     description: Timestamp when the asset was created
                     format: date-time
                     type: string
-                display_name:
-                    description: Display name of the asset. Mirrors name for backwards compatibility.
-                    nullable: true
-                    type: string
-                file_path:
-                    description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
-                    nullable: true
-                    type: string
                 hash:
                     description: Blake3 hash of the asset content.
                     pattern: ^blake3:[a-f0-9]{64}$
                     type: string
+                loader_path:
+                    description: The value a loader consumes to load this asset. Null when no loader can resolve the file.
+                    nullable: true
+                    type: string
+                display_name:
+                    description: Human-facing label for the asset. Not unique.
+                    nullable: true
+                    type: string
                 id:
                     description: Unique identifier for the asset
                     format: uuid
@@ -144,14 +144,6 @@ components:
         AssetUpdated:
             description: Response returned when an existing asset is successfully updated.
             properties:
-                display_name:
-                    description: Display name of the asset. Mirrors name for backwards compatibility.
-                    nullable: true
-                    type: string
-                file_path:
-                    description: Relative path in global-namespace-root form (e.g. "models/checkpoints/flux.safetensors")
-                    nullable: true
-                    type: string
                 hash:
                     description: Blake3 hash of the asset content.
                     pattern: ^blake3:[a-f0-9]{64}$
@@ -230,6 +222,89 @@ components:
                 - base_version
                 - workflow_json
             type: object
+        DownloadEnqueueRequest:
+            description: Request body for enqueuing a server-side model download.
+            properties:
+                allow_any_extension:
+                    default: false
+                    description: Permit a non-model file extension (default only allows known model extensions).
+                    type: boolean
+                expected_sha256:
+                    description: Optional hub-provided SHA256 to verify the completed file against (fail-closed).
+                    nullable: true
+                    type: string
+                model_id:
+                    description: Destination as "<directory>/<filename>", resolving to a registered model folder (e.g. "loras/my_lora.safetensors").
+                    type: string
+                priority:
+                    default: 0
+                    description: Scheduling priority; higher is admitted first.
+                    type: integer
+                url:
+                    description: Source URL; must be on the allowlist (host + scheme + extension).
+                    type: string
+            required:
+                - url
+                - model_id
+            type: object
+        DownloadStatus:
+            description: Current state and live progress of a single download.
+            properties:
+                bytes_done:
+                    type: integer
+                created_at:
+                    type: integer
+                download_id:
+                    format: uuid
+                    type: string
+                error:
+                    nullable: true
+                    type: string
+                eta_seconds:
+                    nullable: true
+                    type: number
+                model_id:
+                    type: string
+                priority:
+                    type: integer
+                progress:
+                    description: Fraction in [0,1]; null until total size is known.
+                    nullable: true
+                    type: number
+                segments:
+                    description: Per-segment progress (segmented downloads only).
+                    items:
+                        properties:
+                            bytes_done:
+                                type: integer
+                            idx:
+                                type: integer
+                            length:
+                                type: integer
+                        type: object
+                    nullable: true
+                    type: array
+                speed_bps:
+                    nullable: true
+                    type: number
+                status:
+                    enum:
+                        - queued
+                        - active
+                        - paused
+                        - verifying
+                        - completed
+                        - failed
+                        - cancelled
+                    type: string
+                total_bytes:
+                    nullable: true
+                    type: integer
+                updated_at:
+                    type: integer
+                url:
+                    type: string
+            type: object
         ErrorResponse:
             description: Standard error response with a machine-readable code and human-readable message.
             properties:
@@ -511,6 +586,27 @@ components:
             required:
                 - history
             type: object
+        DownloadAuthProvider:
+            description: Per-provider download-auth status. Never includes a token.
+            properties:
+                env_key_present:
+                    description: Whether an API key for this provider is set via an environment variable.
+                    type: boolean
+                logged_in:
+                    description: Whether a stored OAuth token exists for this provider.
+                    type: boolean
+                login_in_progress:
+                    description: Whether an OAuth login flow is currently awaiting the browser callback.
+                    type: boolean
+                provider:
+                    description: Provider name (e.g. "huggingface", "civitai").
+                    type: string
+            required:
+                - provider
+                - logged_in
+                - login_in_progress
+                - env_key_present
+            type: object
         JobCancelResponse:
             description: Response for POST /api/jobs/{job_id}/cancel. Returned on both fresh cancels and idempotent no-ops.
             properties:
@@ -783,6 +879,14 @@ components:
         ModelFolder:
             description: Represents a folder containing models
             properties:
+                extensions:
+                    description: The folder's registered file-extension allowlist. An empty array means the folder accepts any extension (match-all).
+                    example:
+                        - .ckpt
+                        - .safetensors
+                    items:
+                        type: string
+                    type: array
                 folders:
                     description: List of paths where models of this type are stored
                     example:
@@ -1644,7 +1748,7 @@ paths:
                                     format: uuid
                                     type: string
                                 tags:
-                                    description: JSON-encoded array of freeform tag strings, e.g. '["models","checkpoint"]'. Common types include "models", "input", "output", and "temp", but any tag can be used in any order.
+                                    description: JSON-encoded array of tag strings. For new byte uploads, include exactly one destination role (`input`, `output`, or `models`); `models` uploads also require exactly one `model_type:<folder_name>` tag. Extra tags are stored as labels and do not create path components.
                                     type: string
                                 user_metadata:
                                     description: Custom JSON metadata as a string
@@ -1829,7 +1933,7 @@ paths:
                     content:
                         application/json:
                             schema:
-                                $ref: '#/components/schemas/AssetUpdated'
+                                $ref: '#/components/schemas/Asset'
                     description: Asset updated successfully
                 "400":
                     content:
@@ -2350,6 +2454,377 @@ paths:
             summary: Get tag histogram for filtered assets
             tags:
                 - file
+    /api/download:
+        get:
+            description: List all known downloads (queued, active, paused, and terminal) with live progress.
+            operationId: listDownloads
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    downloads:
+                                        items:
+                                            $ref: '#/components/schemas/DownloadStatus'
+                                        type: array
+                                type: object
+                    description: List of downloads
+            summary: List downloads
+            tags:
+                - download
+    /api/download/availability:
+        post:
+            description: |
+                Bulk per-id availability for a set of model_ids declared in a workflow.
+                Returns whether each model is available on disk, currently downloading
+                (with progress), or missing, plus whether its URL is on the allowlist.
+            operationId: getModelsAvailability
+            requestBody:
+                content:
+                    application/json:
+                        schema:
+                            properties:
+                                models:
+                                    additionalProperties:
+                                        type: string
+                                    description: Map of "<directory>/<filename>" model_id to its declared source URL.
+                                    type: object
+                            type: object
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    models:
+                                        additionalProperties: true
+                                        type: object
+                                type: object
+                    description: Per-id availability map
+            summary: Bulk model availability + status
+            tags:
+                - download
+    /api/download/clear:
+        post:
+            description: |
+                Delete all terminal downloads (completed, failed, cancelled) from history
+                in one transaction, so the cleared history persists across reloads. Live
+                downloads (queued, active, paused, verifying) are skipped. Finished model
+                files on disk are never removed; only leftover .part temp files are cleaned up.
+            operationId: clearDownloads
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    deleted:
+                                        description: Number of history rows removed.
+                                        type: integer
+                                type: object
+                    description: History cleared
+            summary: Clear terminal downloads from history
+            tags:
+                - download
+    /api/download/auth:
+        get:
+            description: Per-provider download-auth status (OAuth login state + env-key presence). Never returns a token.
+            operationId: getDownloadAuth
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    providers:
+                                        items:
+                                            $ref: '#/components/schemas/DownloadAuthProvider'
+                                        type: array
+                                type: object
+                    description: Per-provider auth status
+            summary: Get download auth status
+            tags:
+                - download
+    /api/download/auth/{provider}/login:
+        post:
+            description: |
+                Start an OAuth 2.0 PKCE login for a provider. Spawns a transient loopback
+                callback server and returns the authorize URL to open in a browser.
+            operationId: loginDownloadAuth
+            parameters:
+                - in: path
+                  name: provider
+                  required: true
+                  schema:
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    authorize_url:
+                                        type: string
+                                type: object
+                    description: Authorize URL to open in a browser
+                "400":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unknown provider or OAuth app not configured
+                "409":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: A login for this provider is already in progress
+            summary: Start OAuth login for a provider
+            tags:
+                - download
+    /api/download/auth/{provider}/logout:
+        post:
+            description: Clear the stored OAuth token (memory + on-disk file) for a provider.
+            operationId: logoutDownloadAuth
+            parameters:
+                - in: path
+                  name: provider
+                  required: true
+                  schema:
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    logged_out:
+                                        type: boolean
+                                type: object
+                    description: Token cleared
+                "400":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Unknown provider
+            summary: Log out a provider
+            tags:
+                - download
+    /api/download/enqueue:
+        post:
+            description: |
+                Enqueue a server-side model download. The URL must be on the allowlist
+                (host + scheme + extension) and the model_id must be "<directory>/<filename>"
+                resolving to a registered model folder. Returns immediately; track progress
+                via GET /api/download/{id} or the "download_progress" websocket event.
+            operationId: enqueueDownload
+            requestBody:
+                content:
+                    application/json:
+                        schema:
+                            $ref: '#/components/schemas/DownloadEnqueueRequest'
+            responses:
+                "202":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    accepted:
+                                        type: boolean
+                                    download_id:
+                                        format: uuid
+                                        type: string
+                                type: object
+                    description: Download accepted and queued
+                "400":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Invalid request (bad URL, model_id, or not allowlisted)
+                "409":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Already on disk or already downloading
+            summary: Enqueue a model download
+            tags:
+                - download
+    /api/download/{id}:
+        delete:
+            description: |
+                Delete a single terminal download from history so it stays gone across
+                reloads. Refuses (409) to delete a live download (queued, active, paused,
+                verifying) — cancel it first. The finished model file on disk is never
+                removed; only a leftover .part temp file is cleaned up.
+            operationId: deleteDownload
+            parameters:
+                - in: path
+                  name: id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    deleted:
+                                        type: boolean
+                                type: object
+                    description: Deleted
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: No such download
+                "409":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: Download is still in progress
+            summary: Delete a download from history
+            tags:
+                - download
+        get:
+            description: Get the current status + progress of a single download.
+            operationId: getDownloadStatus
+            parameters:
+                - in: path
+                  name: id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/DownloadStatus'
+                    description: Download status
+                "404":
+                    content:
+                        application/json:
+                            schema:
+                                $ref: '#/components/schemas/ErrorResponse'
+                    description: No such download
+            summary: Get download status
+            tags:
+                - download
+    /api/download/{id}/cancel:
+        post:
+            description: Cancel a download. The partial file is removed.
+            operationId: cancelDownload
+            parameters:
+                - in: path
+                  name: id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    ok:
+                                        type: boolean
+                                type: object
+                    description: Cancelled
+            summary: Cancel a download
+            tags:
+                - download
+    /api/download/{id}/pause:
+        post:
+            description: Pause a download. The partial file and per-segment offsets are retained for resume.
+            operationId: pauseDownload
+            parameters:
+                - in: path
+                  name: id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    ok:
+                                        type: boolean
+                                type: object
+                    description: Paused
+            summary: Pause a download
+            tags:
+                - download
+    /api/download/{id}/priority:
+        post:
+            description: Set a download's scheduling priority. Higher priority is admitted first when a slot frees.
+            operationId: setDownloadPriority
+            parameters:
+                - in: path
+                  name: id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+            requestBody:
+                content:
+                    application/json:
+                        schema:
+                            properties:
+                                priority:
+                                    type: integer
+                            required:
+                                - priority
+                            type: object
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    ok:
+                                        type: boolean
+                                type: object
+                    description: Priority updated
+            summary: Set download priority
+            tags:
+                - download
+    /api/download/{id}/resume:
+        post:
+            description: Resume a paused (or failed) download from its persisted offsets.
+            operationId: resumeDownload
+            parameters:
+                - in: path
+                  name: id
+                  required: true
+                  schema:
+                    format: uuid
+                    type: string
+            responses:
+                "200":
+                    content:
+                        application/json:
+                            schema:
+                                properties:
+                                    ok:
+                                        type: boolean
+                                type: object
+                    description: Resumed
+            summary: Resume a download
+            tags:
+                - download
     /api/embeddings:
         get:
             description: Returns the list of text-encoder embeddings available on disk.
@@ -2470,6 +2945,9 @@ paths:
                                     supports_preview_metadata:
                                         description: Whether the server supports preview metadata
                                         type: boolean
+                                    supports_model_type_tags:
+                                        description: Whether the server supports namespaced model type asset tags
+                                        type: boolean
                                 type: object
                     description: Success
                     headers:
@@ -5103,3 +5581,5 @@ tags:
       name: queue
     - description: Job lifecycle queries
       name: job
+    - description: Model download management
+      name: download

requirements.txt

@@ -22,7 +22,7 @@ alembic
 SQLAlchemy>=2.0.0
 filelock
 av>=16.0.0
-comfy-kitchen==0.2.16
+comfy-kitchen==0.2.18
 comfy-aimdo==0.4.10
 requests
 simpleeval>=1.0.0

server.py

@@ -39,13 +39,17 @@ from comfy.deploy_environment import get_deploy_environment
 import comfy.utils
 import comfy.model_management
 from comfy_api import feature_flags
+from comfy.comfy_api_env import get_environment_overrides
 import node_helpers
 from comfyui_version import __version__
 from app.frontend_management import FrontendManager, parse_version
 from comfy_api.internal import _ComfyNodeInternal
 from app.assets.seeder import asset_seeder
 from app.assets.api.routes import register_assets_routes
+from app.model_downloader.api.routes import register_routes as register_model_downloader_routes
+from app.model_downloader.manager import DOWNLOAD_MANAGER
 from app.assets.services.ingest import register_file_in_place
+from app.assets.services.path_utils import get_known_subfolder_tags
 from app.assets.services.asset_management import resolve_hash_to_path
 
 from app.user_manager import UserManager
@@ -257,6 +261,7 @@ class PromptServer():
         else:
             register_assets_routes(self.app)
             asset_seeder.disable()
+        register_model_downloader_routes(self.app)
         routes = web.RouteTableDef()
         self.routes = routes
         self.last_node_id = None
@@ -441,7 +446,9 @@ class PromptServer():
                 if args.enable_assets:
                     try:
                         tag = image_upload_type if image_upload_type in ("input", "output") else "input"
-                        result = register_file_in_place(abs_path=filepath, name=filename, tags=[tag])
+                        tags = [tag]
+                        tags.extend(get_known_subfolder_tags(subfolder))
+                        result = register_file_in_place(abs_path=filepath, name=filename, tags=tags)
                         resp["asset"] = {
                             "id": result.ref.id,
                             "name": result.ref.name,
@@ -724,7 +731,11 @@ class PromptServer():
 
         @routes.get("/features")
         async def get_features(request):
-            return web.json_response(feature_flags.get_server_features())
+            features = feature_flags.get_server_features()
+            overrides = get_environment_overrides()
+            if overrides:
+                features.update(overrides)
+            return web.json_response(features)
 
         @routes.get("/prompt")
         async def get_prompt(request):
@@ -1198,6 +1209,29 @@ class PromptServer():
     async def setup(self):
         timeout = aiohttp.ClientTimeout(total=None) # no timeout
         self.client_session = aiohttp.ClientSession(timeout=timeout)
+        await self._setup_model_downloader()
+
+    async def _setup_model_downloader(self):
+        """Start the download manager: push progress over the websocket and
+        resume any downloads interrupted by a previous run."""
+        def _notify(download_id: str) -> None:
+            try:
+                view = DOWNLOAD_MANAGER.status_sync(download_id)
+                if view is not None:
+                    # Drop the url field before broadcasting: the redacted URL
+                    # (scheme + host + path) should not leak to every connected
+                    # websocket client. download_id / model_id are sufficient to
+                    # correlate progress on the frontend.
+                    broadcast = {k: v for k, v in view.items() if k != "url"}
+                    self.send_sync("download_progress", broadcast)
+            except Exception:
+                logging.debug("download progress notify failed", exc_info=True)
+
+        DOWNLOAD_MANAGER.set_notify(_notify)
+        try:
+            await DOWNLOAD_MANAGER.start()
+        except Exception as e:
+            logging.warning("Failed to start model download manager: %s", e)
 
     def add_routes(self):
         self.user_manager.add_routes(self.routes)

tests-unit/app_test/model_manager_test.py

@@ -24,6 +24,28 @@ def app(model_manager):
     app.add_routes(routes)
     return app
 
+async def test_get_model_folders_includes_registered_extensions(aiohttp_client, app, tmp_path):
+    """Folders expose their registered extension set verbatim; an empty list
+    means match-all (filter_files_extensions semantics)."""
+    with patch('folder_paths.folder_names_and_paths', {
+        'test_checkpoints': ([str(tmp_path)], {'.safetensors', '.ckpt'}),
+        'test_configs': ([str(tmp_path)], ['.yaml']),
+        'test_match_all': ([str(tmp_path)], set()),
+        'configs': ([str(tmp_path)], ['.yaml']),
+    }):
+        client = await aiohttp_client(app)
+        response = await client.get('/experiment/models')
+
+        assert response.status == 200
+        folders = {f['name']: f for f in await response.json()}
+
+        assert 'configs' not in folders  # blocklisted
+        assert folders['test_checkpoints']['folders'] == [str(tmp_path)]
+        assert folders['test_checkpoints']['extensions'] == ['.ckpt', '.safetensors']
+        assert folders['test_configs']['extensions'] == ['.yaml']
+        # Match-all registrations are exposed honestly, not substituted.
+        assert folders['test_match_all']['extensions'] == []
+
 async def test_get_model_preview_safetensors(aiohttp_client, app, tmp_path):
     img = Image.new('RGB', (100, 100), 'white')
     img_byte_arr = BytesIO()

tests-unit/app_test/test_migrations.py

@@ -8,6 +8,7 @@ upgrade/downgrade for 0003+.
 """
 
 import os
+import sqlite3
 
 import pytest
 from alembic import command
@@ -30,6 +31,12 @@ def _make_config(db_path: str) -> Config:
     return cfg
 
 
+def _sqlite_path(cfg: Config) -> str:
+    url = cfg.get_main_option("sqlalchemy.url")
+    assert url is not None and url.startswith("sqlite:///")
+    return url.removeprefix("sqlite:///")
+
+
 @pytest.fixture
 def migration_db(tmp_path):
     """Yield an alembic Config pre-upgraded to the baseline revision."""
@@ -55,3 +62,26 @@ def test_upgrade_downgrade_cycle(migration_db):
     command.upgrade(migration_db, "head")
     command.downgrade(migration_db, _BASELINE)
     command.upgrade(migration_db, "head")
+
+
+def test_case_sensitive_tags_downgrade_normalizes_existing_tags(migration_db):
+    """Downgrading 0005 folds mixed-case tag vocabulary before restoring CHECK."""
+    command.upgrade(migration_db, "0005_allow_case_sensitive_tags")
+
+    db_path = _sqlite_path(migration_db)
+    with sqlite3.connect(db_path) as conn:
+        conn.execute("INSERT INTO tags(name) VALUES (?)", ("NewTag",))
+        conn.execute("INSERT INTO tags(name) VALUES (?)", ("newtag",))
+        conn.execute("INSERT INTO tags(name) VALUES (?)", ("model_type:LLM",))
+
+    command.downgrade(migration_db, "0004_drop_tag_type")
+
+    with sqlite3.connect(db_path) as conn:
+        tags = {row[0] for row in conn.execute("SELECT name FROM tags")}
+        assert "newtag" in tags
+        assert "model_type:llm" in tags
+        assert "NewTag" not in tags
+        assert "model_type:LLM" not in tags
+
+        with pytest.raises(sqlite3.IntegrityError):
+            conn.execute("INSERT INTO tags(name) VALUES (?)", ("Upper",))

tests-unit/assets_test/conftest.py

@@ -234,7 +234,7 @@ def seeded_asset(request: pytest.FixtureRequest, http: requests.Session, api_bas
     p = getattr(request, "param", {}) or {}
     tags: Optional[list[str]] = p.get("tags")
     if tags is None:
-        tags = ["models", "checkpoints", "unit-tests", "alpha"]
+        tags = ["models", "model_type:checkpoints", "unit-tests", "alpha"]
     meta = {"purpose": "test", "epoch": 1, "flags": ["x", "y"], "nullable": None}
     # Unique content per test so the seed always creates a fresh asset (201).
     # Delete is now always a soft delete, so content from a prior test survives

tests-unit/assets_test/queries/test_asset_info.py

@@ -133,6 +133,66 @@ class TestListReferencesPage:
         assert total == 1
         assert refs[0].name == "tagged"
 
+    def test_include_tags_filter_ands_persisted_model_tags(self, session: Session):
+        asset = _make_asset(session, "hash-model-tags")
+        checkpoint = _make_reference(session, asset, name="checkpoint")
+        lora = _make_reference(session, asset, name="lora")
+        input_ref = _make_reference(session, asset, name="input")
+        ensure_tags_exist(
+            session,
+            ["models", "model_type:checkpoints", "model_type:loras", "unit-tests"],
+        )
+        add_tags_to_reference(
+            session,
+            reference_id=checkpoint.id,
+            tags=["models", "model_type:checkpoints", "unit-tests"],
+            origin="automatic",
+        )
+        add_tags_to_reference(
+            session,
+            reference_id=lora.id,
+            tags=["models", "model_type:loras", "unit-tests"],
+            origin="automatic",
+        )
+        add_tags_to_reference(
+            session,
+            reference_id=input_ref.id,
+            tags=["unit-tests"],
+        )
+        session.commit()
+
+        refs, _, total = list_references_page(
+            session,
+            include_tags=["models", "model_type:checkpoints", "unit-tests"],
+        )
+
+        assert total == 1
+        assert refs[0].id == checkpoint.id
+
+    def test_include_tags_filter_preserves_model_type_case(self, session: Session):
+        asset = _make_asset(session, "hash-model-case")
+        ref = _make_reference(session, asset, name="llm")
+        ensure_tags_exist(session, ["models", "model_type:LLM"])
+        add_tags_to_reference(
+            session,
+            reference_id=ref.id,
+            tags=["models", "model_type:LLM"],
+            origin="automatic",
+        )
+        session.commit()
+
+        refs, _, total = list_references_page(
+            session, include_tags=["models", "model_type:LLM"]
+        )
+        refs_lower, _, total_lower = list_references_page(
+            session, include_tags=["models", "model_type:llm"]
+        )
+
+        assert total == 1
+        assert refs[0].id == ref.id
+        assert total_lower == 0
+        assert refs_lower == []
+
     def test_exclude_tags_filter(self, session: Session):
         asset = _make_asset(session, "hash1")
         _make_reference(session, asset, name="keep")

tests-unit/assets_test/queries/test_cache_state.py

@@ -176,6 +176,39 @@ class TestUpsertReference:
         ref = session.query(AssetReference).filter_by(file_path=file_path).one()
         assert ref.mtime_ns == final_mtime
 
+    def test_upsert_refreshes_loader_path_on_existing_reference(self, session: Session):
+        """Re-ingesting an existing reference writes the loader_path computed
+        by that ingest, healing NULL or stale values even when nothing else
+        about the row changed."""
+        asset = _make_asset(session, "hash1")
+        file_path = "/models/checkpoints/sub/model.safetensors"
+
+        upsert_reference(
+            session, asset_id=asset.id, file_path=file_path, name="model",
+            mtime_ns=100, loader_path=None,
+        )
+        session.commit()
+
+        created, updated = upsert_reference(
+            session, asset_id=asset.id, file_path=file_path, name="model",
+            mtime_ns=100, loader_path="sub/model.safetensors",
+        )
+        session.commit()
+
+        assert created is False
+        assert updated is True
+        ref = session.query(AssetReference).filter_by(file_path=file_path).one()
+        assert ref.loader_path == "sub/model.safetensors"
+
+        # Identical loader_path is a no-op, not a spurious update.
+        created, updated = upsert_reference(
+            session, asset_id=asset.id, file_path=file_path, name="model",
+            mtime_ns=100, loader_path="sub/model.safetensors",
+        )
+        session.commit()
+        assert created is False
+        assert updated is False
+
     def test_upsert_restores_missing_reference(self, session: Session):
         """Upserting a reference that was marked missing should restore it."""
         asset = _make_asset(session, "hash1")

tests-unit/assets_test/queries/test_tags.py

@@ -58,7 +58,7 @@ class TestEnsureTagsExist:
         session.commit()
 
         tags = session.query(Tag).all()
-        assert {t.name for t in tags} == {"alpha", "beta"}
+        assert {t.name for t in tags} == {"ALPHA", "Beta", "alpha"}
 
     def test_empty_list_is_noop(self, session: Session):
         ensure_tags_exist(session, [])
@@ -258,6 +258,16 @@ class TestListTagsWithUsage:
         tag_names = {name for name, _ in rows}
         assert tag_names == {"alpha", "alphabet"}
 
+    def test_prefix_filter_is_case_sensitive(self, session: Session):
+        ensure_tags_exist(session, ["model_type:LLM", "model_type:llm"])
+        session.commit()
+
+        rows, total = list_tags_with_usage(session, prefix="model_type:L")
+
+        tag_names = {name for name, _ in rows}
+        assert tag_names == {"model_type:LLM"}
+        assert total == 1
+
     def test_order_by_name(self, session: Session):
         ensure_tags_exist(session, ["zebra", "alpha", "middle"])
         session.commit()

tests-unit/assets_test/services/test_asset_response_loader_path.py

@@ -0,0 +1,83 @@
+"""Tests for how _build_asset_response derives the response `loader_path`.
+
+Guards the persist-and-read contract: the response reads the stored
+`loader_path` verbatim, with no read-time recomputation. Like tags, the
+value is a seed-time derivative healed by the scan lifecycle.
+"""
+
+from datetime import datetime
+from pathlib import Path
+from unittest.mock import patch
+
+from app.assets.api.routes import _build_asset_response
+from app.assets.services.schemas import AssetDetailResult, ReferenceData
+
+_TS = datetime(2024, 1, 1, 0, 0, 0)
+
+
+def _make_result(
+    *, file_path: str | None, loader_path: str | None
+) -> AssetDetailResult:
+    ref = ReferenceData(
+        id="ref-1",
+        name="model.safetensors",
+        file_path=file_path,
+        loader_path=loader_path,
+        user_metadata=None,
+        preview_id=None,
+        created_at=_TS,
+        updated_at=_TS,
+        last_access_time=_TS,
+    )
+    return AssetDetailResult(ref=ref, asset=None, tags=[])
+
+
+def test_uses_persisted_loader_path_without_recomputing():
+    """A stored loader_path is returned verbatim, not re-derived from file_path.
+
+    The sentinel value could never be produced by compute_loader_path for this
+    file_path, so seeing it in the response proves the stored column is read.
+    """
+    result = _make_result(
+        file_path="/unmatched/root/model.safetensors",
+        loader_path="SENTINEL/stored.safetensors",
+    )
+
+    resp = _build_asset_response(result)
+
+    assert resp.loader_path == "SENTINEL/stored.safetensors"
+
+
+def test_null_stored_loader_path_is_served_as_null(tmp_path: Path):
+    """No read-time recomputation: a NULL column is served as null even when
+    the path would resolve."""
+    models = tmp_path / "models"
+    ckpt = models / "checkpoints"
+    ckpt.mkdir(parents=True)
+    f = ckpt / "bar.safetensors"
+    f.touch()
+
+    with patch("app.assets.services.path_utils.folder_paths") as mock_fp, patch(
+        "app.assets.services.path_utils.get_comfy_models_folders",
+        return_value=[("checkpoints", [str(ckpt)], {".safetensors"})],
+    ):
+        mock_fp.get_input_directory.return_value = str(tmp_path / "in")
+        mock_fp.get_output_directory.return_value = str(tmp_path / "out")
+        mock_fp.get_temp_directory.return_value = str(tmp_path / "tmp")
+        mock_fp.models_dir = str(models)
+
+        result = _make_result(file_path=str(f), loader_path=None)
+        resp = _build_asset_response(result)
+
+        assert resp.loader_path is None
+        assert resp.display_name == "checkpoints/bar.safetensors"
+
+
+def test_all_path_fields_null_without_file_path():
+    """API-created / hash-only references (no file_path) expose no paths."""
+    result = _make_result(file_path=None, loader_path=None)
+
+    resp = _build_asset_response(result)
+
+    assert resp.loader_path is None
+    assert resp.display_name is None

tests-unit/assets_test/services/test_bulk_ingest.py

@@ -1,10 +1,14 @@
 """Tests for bulk ingest services."""
 
+import os
 from pathlib import Path
+from unittest.mock import patch
 
 from sqlalchemy.orm import Session
 
 from app.assets.database.models import Asset, AssetReference
+from app.assets.database.queries import get_reference_tags
+from app.assets.scanner import build_asset_specs
 from app.assets.services.bulk_ingest import SeedAssetSpec, batch_insert_seed_assets
 
 
@@ -101,6 +105,184 @@ class TestBatchInsertSeedAssets:
             asset = session.query(Asset).filter_by(id=ref.asset_id).first()
             assert asset.mime_type == expected_mime, f"Expected {expected_mime} for {filename}, got {asset.mime_type}"
 
+    def test_duplicate_paths_merge_tags_before_insert(
+        self, session: Session, temp_dir: Path
+    ):
+        """Overlapping model-folder registrations can emit the same path twice."""
+        file_path = temp_dir / "shared.safetensors"
+        file_path.write_bytes(b"shared model")
+
+        specs: list[SeedAssetSpec] = [
+            {
+                "abs_path": str(file_path),
+                "size_bytes": 12,
+                "mtime_ns": 1234567890000000000,
+                "info_name": "Shared Model",
+                "tags": ["models", "model_type:checkpoints"],
+                "fname": "shared.safetensors",
+                "metadata": None,
+                "hash": None,
+                "mime_type": "application/safetensors",
+            },
+            {
+                "abs_path": str(file_path),
+                "size_bytes": 12,
+                "mtime_ns": 1234567890000000000,
+                "info_name": "Shared Model",
+                "tags": ["models", "model_type:diffusion_models"],
+                "fname": "shared.safetensors",
+                "metadata": None,
+                "hash": None,
+                "mime_type": "application/safetensors",
+            },
+        ]
+
+        result = batch_insert_seed_assets(session, specs=specs, owner_id="")
+
+        assert result.inserted_refs == 1
+        assert result.won_paths == 1
+        refs = session.query(AssetReference).all()
+        assert len(refs) == 1
+        assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
+            "models",
+            "model_type:checkpoints",
+            "model_type:diffusion_models",
+        }
+
+    def test_duplicate_paths_are_merged_after_abspath_normalization(
+        self, session: Session, temp_dir: Path, monkeypatch
+    ):
+        """The scanner may emit equivalent paths with different spelling."""
+        file_path = temp_dir / "same-file.safetensors"
+        file_path.write_bytes(b"shared model")
+        monkeypatch.chdir(temp_dir)
+        relative_path = file_path.name
+        absolute_path = os.path.abspath(relative_path)
+
+        specs: list[SeedAssetSpec] = [
+            {
+                "abs_path": relative_path,
+                "size_bytes": 12,
+                "mtime_ns": 1234567890000000000,
+                "info_name": "Shared Model",
+                "tags": ["models", "model_type:checkpoints"],
+                "fname": "same-file.safetensors",
+                "metadata": None,
+                "hash": None,
+                "mime_type": "application/safetensors",
+            },
+            {
+                "abs_path": absolute_path,
+                "size_bytes": 12,
+                "mtime_ns": 1234567890000000000,
+                "info_name": "Shared Model",
+                "tags": ["models", "model_type:diffusion_models"],
+                "fname": "same-file.safetensors",
+                "metadata": None,
+                "hash": None,
+                "mime_type": "application/safetensors",
+            },
+        ]
+
+        result = batch_insert_seed_assets(session, specs=specs, owner_id="")
+
+        assert result.inserted_refs == 1
+        assert result.won_paths == 1
+        refs = session.query(AssetReference).all()
+        assert len(refs) == 1
+        assert refs[0].file_path == absolute_path
+        # loader_path is persisted from the spec's fname (compute_loader_path).
+        assert refs[0].loader_path == "same-file.safetensors"
+        assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
+            "models",
+            "model_type:checkpoints",
+            "model_type:diffusion_models",
+        }
+
+    def test_scanner_duplicate_shared_model_paths_keep_all_model_type_tags(
+        self, session: Session, temp_dir: Path
+    ):
+        """Shared extra model roots make scanner collection emit duplicate paths."""
+        shared_root = temp_dir / "shared"
+        input_dir = temp_dir / "input"
+        output_dir = temp_dir / "output"
+        temp_root = temp_dir / "temp"
+        for directory in (shared_root, input_dir, output_dir, temp_root):
+            directory.mkdir()
+        file_path = shared_root / "dual_use_model.safetensors"
+        file_path.write_bytes(b"shared model")
+
+        with (
+            patch("app.assets.services.path_utils.folder_paths") as mock_fp,
+            patch(
+                "app.assets.services.path_utils.get_comfy_models_folders",
+                return_value=[
+                    ("checkpoints", [str(shared_root)], {".safetensors"}),
+                    ("diffusion_models", [str(shared_root)], {".safetensors"}),
+                ],
+            ),
+        ):
+            mock_fp.get_input_directory.return_value = str(input_dir)
+            mock_fp.get_output_directory.return_value = str(output_dir)
+            mock_fp.get_temp_directory.return_value = str(temp_root)
+
+            specs, tag_pool, skipped = build_asset_specs(
+                paths=[str(file_path), str(file_path)],
+                existing_paths=set(),
+                enable_metadata_extraction=False,
+                compute_hashes=False,
+            )
+
+        assert skipped == 0
+        assert len(specs) == 2
+        assert tag_pool == {
+            "models",
+            "model_type:checkpoints",
+            "model_type:diffusion_models",
+        }
+
+        result = batch_insert_seed_assets(session, specs=specs, owner_id="")
+
+        assert result.inserted_refs == 1
+        assert result.won_paths == 1
+        refs = session.query(AssetReference).all()
+        assert len(refs) == 1
+        assert set(get_reference_tags(session, reference_id=refs[0].id)) == {
+            "models",
+            "model_type:checkpoints",
+            "model_type:diffusion_models",
+        }
+
+    def test_loader_path_persisted_as_null_when_fname_is_none(
+        self, session: Session, temp_dir: Path
+    ):
+        """A file with no in-root loader path (fname=None, e.g. an orphan under
+        models_root) persists loader_path as NULL rather than a synthesized value."""
+        file_path = temp_dir / "orphan.bin"
+        file_path.write_bytes(b"x")
+
+        specs: list[SeedAssetSpec] = [
+            {
+                "abs_path": str(file_path),
+                "size_bytes": 1,
+                "mtime_ns": 1234567890000000000,
+                "info_name": "orphan.bin",
+                "tags": [],
+                "fname": None,
+                "metadata": None,
+                "hash": None,
+                "mime_type": None,
+            }
+        ]
+
+        result = batch_insert_seed_assets(session, specs=specs, owner_id="")
+
+        assert result.inserted_refs == 1
+        refs = session.query(AssetReference).all()
+        assert len(refs) == 1
+        assert refs[0].file_path == str(file_path)
+        assert refs[0].loader_path is None
+
 
 class TestMetadataExtraction:
     def test_extracts_mime_type_for_model_files(self, temp_dir: Path):

tests-unit/assets_test/services/test_ingest.py

@@ -94,6 +94,47 @@ class TestIngestFileFromPath:
         ref_tags = get_reference_tags(session, reference_id=result.reference_id)
         assert set(ref_tags) == {"models", "checkpoints"}
 
+    def test_path_derived_tags_use_automatic_origin(
+        self, mock_create_session, temp_dir: Path, session: Session
+    ):
+        input_dir = temp_dir / "input"
+        output_dir = temp_dir / "output"
+        temp_root = temp_dir / "temp"
+        for directory in (input_dir, output_dir, temp_root):
+            directory.mkdir()
+        file_path = input_dir / "pasted" / "tagged.png"
+        file_path.parent.mkdir()
+        file_path.write_bytes(b"data")
+
+        with (
+            patch("app.assets.services.path_utils.folder_paths") as mock_fp,
+            patch(
+                "app.assets.services.path_utils.get_comfy_models_folders",
+                return_value=[],
+            ),
+        ):
+            mock_fp.get_input_directory.return_value = str(input_dir)
+            mock_fp.get_output_directory.return_value = str(output_dir)
+            mock_fp.get_temp_directory.return_value = str(temp_root)
+
+            result = _ingest_file_from_path(
+                abs_path=str(file_path),
+                asset_hash="blake3:pathorigin",
+                size_bytes=4,
+                mtime_ns=1234567890000000000,
+                info_name="Tagged Asset",
+                tags=["input", "manual-label"],
+            )
+
+        assert result.reference_id is not None
+        links = session.query(AssetReferenceTag).filter_by(
+            asset_reference_id=result.reference_id
+        )
+        origin_by_tag = {link.tag_name: link.origin for link in links}
+        assert origin_by_tag["input"] == "automatic"
+        assert origin_by_tag["pasted"] == "automatic"
+        assert origin_by_tag["manual-label"] == "manual"
+
     def test_idempotent_upsert(self, mock_create_session, temp_dir: Path, session: Session):
         file_path = temp_dir / "dup.bin"
         file_path.write_bytes(b"content")
@@ -288,6 +329,45 @@ class TestIngestExistingFileTagFK:
             assert "output" in ref_tag_names
 
 
+class TestIngestExistingFileLoaderPath:
+    """Outputs saved into a subfolder must persist the subfolder-qualified
+    loader path, not the bare basename (regression: spec["fname"] was
+    os.path.basename)."""
+
+    def test_subfoldered_output_persists_relative_loader_path(
+        self, mock_create_session, temp_dir: Path, session: Session
+    ):
+        input_dir = temp_dir / "input"
+        output_dir = temp_dir / "output"
+        temp_root = temp_dir / "temp"
+        for directory in (input_dir, output_dir, temp_root):
+            directory.mkdir()
+        file_path = output_dir / "sub" / "img_00001_.png"
+        file_path.parent.mkdir()
+        file_path.write_bytes(b"image data")
+
+        with (
+            patch("app.assets.services.path_utils.folder_paths") as mock_fp,
+            patch(
+                "app.assets.services.path_utils.get_comfy_models_folders",
+                return_value=[],
+            ),
+        ):
+            mock_fp.get_input_directory.return_value = str(input_dir)
+            mock_fp.get_output_directory.return_value = str(output_dir)
+            mock_fp.get_temp_directory.return_value = str(temp_root)
+
+            assert ingest_existing_file(abs_path=str(file_path)) is True
+
+        ref = (
+            session.query(AssetReference)
+            .filter_by(file_path=str(file_path))
+            .one()
+        )
+        assert ref.loader_path == "sub/img_00001_.png"
+        assert (ref.user_metadata or {}).get("filename") == "sub/img_00001_.png"
+
+
 class TestIngestImageDimensions:
     """system_metadata should carry {kind, width, height} for image assets."""
 

tests-unit/assets_test/services/test_path_utils.py

@@ -6,7 +6,16 @@ from unittest.mock import patch
 
 import pytest
 
-from app.assets.services.path_utils import get_asset_category_and_relative_path
+from app.assets.services.path_utils import (
+    compute_display_name,
+    compute_loader_path,
+    compute_logical_path,
+    get_asset_category_and_relative_path,
+    get_known_input_subfolder_tags_from_path,
+    get_known_subfolder_tags,
+    get_name_and_tags_from_asset_path,
+    resolve_destination_from_tags,
+)
 
 
 @pytest.fixture
@@ -17,7 +26,8 @@ def fake_dirs():
         input_dir = root_path / "input"
         output_dir = root_path / "output"
         temp_dir = root_path / "temp"
-        models_dir = root_path / "models" / "checkpoints"
+        models_root = root_path / "models"
+        models_dir = models_root / "checkpoints"
         for d in (input_dir, output_dir, temp_dir, models_dir):
             d.mkdir(parents=True)
 
@@ -25,15 +35,17 @@ def fake_dirs():
             mock_fp.get_input_directory.return_value = str(input_dir)
             mock_fp.get_output_directory.return_value = str(output_dir)
             mock_fp.get_temp_directory.return_value = str(temp_dir)
+            mock_fp.models_dir = str(models_root)
 
             with patch(
                 "app.assets.services.path_utils.get_comfy_models_folders",
-                return_value=[("checkpoints", [str(models_dir)])],
+                return_value=[("checkpoints", [str(models_dir)], {".safetensors"})],
             ):
                 yield {
                     "input": input_dir,
                     "output": output_dir,
                     "temp": temp_dir,
+                    "models_root": models_root,
                     "models": models_dir,
                 }
 
@@ -76,6 +88,538 @@ class TestGetAssetCategoryAndRelativePath:
         cat, rel = get_asset_category_and_relative_path(str(f))
         assert cat == "models"
 
+    def test_model_path_tags_include_registered_model_type_only(self, fake_dirs):
+        f = fake_dirs["models"] / "subdir" / "model.safetensors"
+        f.parent.mkdir()
+        f.touch()
+
+        _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert "model_type:checkpoints" in tags
+        assert "checkpoints" not in tags
+        assert "subdir" not in tags
+
+    def test_model_type_preserves_registered_folder_case(self, fake_dirs):
+        llm_dir = fake_dirs["models"].parent / "LLM"
+        llm_dir.mkdir()
+        f = llm_dir / "model.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("LLM", [str(llm_dir)], {".safetensors"})],
+        ):
+            _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert "model_type:LLM" in tags
+        assert "model_type:llm" not in tags
+
+    def test_path_components_do_not_create_model_type_tags(self, fake_dirs):
+        f = fake_dirs["models"] / "loras" / "model.safetensors"
+        f.parent.mkdir()
+        f.touch()
+
+        _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert "model_type:checkpoints" in tags
+        assert "loras" not in tags
+        assert "model_type:loras" not in tags
+
+    def test_shared_root_returns_all_matching_model_type_tags(self, fake_dirs):
+        shared_root = fake_dirs["models"].parent / "shared"
+        shared_root.mkdir()
+        f = shared_root / "foo.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[
+                ("checkpoints", [str(shared_root)], {".safetensors"}),
+                ("loras", [str(shared_root)], {".safetensors"}),
+            ],
+        ):
+            _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert "model_type:checkpoints" in tags
+        assert "model_type:loras" in tags
+
+    def test_shared_root_model_type_tags_respect_bucket_extensions(self, fake_dirs):
+        """Buckets sharing a base dir only tag files matching their extensions."""
+        shared_root = fake_dirs["models"].parent / "unet"
+        shared_root.mkdir()
+        safetensors_file = shared_root / "wan.safetensors"
+        gguf_file = shared_root / "wan.gguf"
+        safetensors_file.touch()
+        gguf_file.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[
+                ("diffusion_models", [str(shared_root)], {".safetensors"}),
+                ("unet_gguf", [str(shared_root)], {".gguf"}),
+            ],
+        ):
+            _name, safetensors_tags = get_name_and_tags_from_asset_path(str(safetensors_file))
+            _name, gguf_tags = get_name_and_tags_from_asset_path(str(gguf_file))
+
+        assert "model_type:diffusion_models" in safetensors_tags
+        assert "model_type:unet_gguf" not in safetensors_tags
+        assert "model_type:unet_gguf" in gguf_tags
+        assert "model_type:diffusion_models" not in gguf_tags
+
+    def test_empty_extension_set_tags_any_extension(self, fake_dirs):
+        """Custom buckets registered without extensions accept every file."""
+        custom_root = fake_dirs["models"].parent / "custom_bucket"
+        custom_root.mkdir()
+        f = custom_root / "weights.bin"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("custom_bucket", [str(custom_root)], set())],
+        ):
+            _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert "model_type:custom_bucket" in tags
+
+    def test_no_extension_match_keeps_models_tag_without_model_type(self, fake_dirs):
+        f = fake_dirs["models"] / "notes.txt"
+        f.touch()
+
+        _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert not any(tag.startswith("model_type:") for tag in tags)
+
+    def test_output_backed_registered_folder_gets_model_and_output_tags(self, fake_dirs):
+        output_checkpoints_dir = fake_dirs["output"] / "checkpoints"
+        output_checkpoints_dir.mkdir()
+        f = output_checkpoints_dir / "saved.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("checkpoints", [str(output_checkpoints_dir)], {".safetensors"})],
+        ):
+            _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "models" in tags
+        assert "model_type:checkpoints" in tags
+        assert "output" in tags
+
+    def test_temp_path_tags_include_temp_not_output_or_preview(self, fake_dirs):
+        f = fake_dirs["temp"] / "preview.png"
+        f.touch()
+
+        _name, tags = get_name_and_tags_from_asset_path(str(f))
+
+        assert "temp" in tags
+        assert "output" not in tags
+        assert "preview:true" not in tags
+
+    def test_known_subfolder_tags_are_centralized(self):
+        assert get_known_subfolder_tags("pasted") == ["pasted"]
+        assert get_known_subfolder_tags("arbitrary") == []
+
+    def test_known_input_subfolder_tags_are_path_derived_for_direct_children(self, fake_dirs):
+        f = fake_dirs["input"] / "pasted" / "image.png"
+        f.parent.mkdir()
+        f.touch()
+
+        assert get_known_input_subfolder_tags_from_path(str(f)) == ["pasted"]
+
+        _name, tags = get_name_and_tags_from_asset_path(str(f))
+        assert "input" in tags
+        assert "pasted" in tags
+
+    def test_known_input_subfolder_tags_do_not_apply_to_nested_or_other_roots(self, fake_dirs):
+        nested = fake_dirs["input"] / "pasted" / "session" / "image.png"
+        output = fake_dirs["output"] / "pasted" / "image.png"
+        for path in (nested, output):
+            path.parent.mkdir(parents=True)
+            path.touch()
+
+        assert get_known_input_subfolder_tags_from_path(str(nested)) == []
+        assert get_known_input_subfolder_tags_from_path(str(output)) == []
+
     def test_unknown_path_raises(self, fake_dirs):
         with pytest.raises(ValueError, match="not within"):
             get_asset_category_and_relative_path("/some/random/path.png")
+
+
+class TestResponseStoragePaths:
+    def test_input_file_path_and_display_name_include_subfolder(self, fake_dirs):
+        sub = fake_dirs["input"] / "some" / "folder"
+        sub.mkdir(parents=True)
+        f = sub / "image.png"
+        f.touch()
+
+        assert compute_logical_path(str(f)) == "input/some/folder/image.png"
+        assert compute_display_name(str(f)) == "some/folder/image.png"
+
+    def test_output_file_path_and_display_name_include_subfolder(self, fake_dirs):
+        sub = fake_dirs["output"] / "renders"
+        sub.mkdir()
+        f = sub / "ComfyUI_00001_.png"
+        f.touch()
+
+        assert compute_logical_path(str(f)) == "output/renders/ComfyUI_00001_.png"
+        assert compute_display_name(str(f)) == "renders/ComfyUI_00001_.png"
+
+    def test_temp_file_path_and_display_name(self, fake_dirs):
+        f = fake_dirs["temp"] / "preview.png"
+        f.touch()
+
+        assert compute_logical_path(str(f)) == "temp/preview.png"
+        assert compute_display_name(str(f)) == "preview.png"
+
+    def test_exact_storage_root_has_no_display_name(self, fake_dirs):
+        assert compute_logical_path(str(fake_dirs["input"])) == "input"
+        assert compute_display_name(str(fake_dirs["input"])) is None
+
+    def test_longest_matching_builtin_root_wins(self, fake_dirs, tmp_path: Path):
+        nested_output = fake_dirs["input"] / "nested-output"
+        nested_output.mkdir()
+        f = nested_output / "image.png"
+        f.touch()
+
+        with patch("app.assets.services.path_utils.folder_paths") as mock_fp:
+            mock_fp.get_input_directory.return_value = str(fake_dirs["input"])
+            mock_fp.get_output_directory.return_value = str(nested_output)
+            mock_fp.get_temp_directory.return_value = str(tmp_path / "temp")
+            mock_fp.models_dir = str(fake_dirs["models_root"])
+
+            assert compute_logical_path(str(f)) == "output/image.png"
+            assert compute_display_name(str(f)) == "image.png"
+
+    def test_model_file_path_is_relative_to_physical_models_root(self, fake_dirs):
+        sub = fake_dirs["models"] / "flux"
+        sub.mkdir()
+        f = sub / "model.safetensors"
+        f.touch()
+
+        assert compute_logical_path(str(f)) == "models/checkpoints/flux/model.safetensors"
+        assert compute_display_name(str(f)) == "checkpoints/flux/model.safetensors"
+
+        name, tags = get_name_and_tags_from_asset_path(str(f))
+        assert name == "model.safetensors"
+        assert "models" in tags
+        assert "model_type:checkpoints" in tags
+        assert "checkpoints" not in tags
+        assert "flux" not in tags
+
+    @pytest.mark.parametrize(
+        "folder_name",
+        ["checkpoints", "clip", "vae", "diffusion_models", "loras"],
+    )
+    def test_output_model_folder_uses_output_storage_file_path(self, fake_dirs, folder_name):
+        output_model_dir = fake_dirs["output"] / folder_name
+        output_model_dir.mkdir(exist_ok=True)
+        default_model_dir = fake_dirs["models_root"] / folder_name
+        default_model_dir.mkdir(exist_ok=True)
+        f = output_model_dir / "saved.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[
+                (folder_name, [str(default_model_dir), str(output_model_dir)], {".safetensors"})
+            ],
+        ):
+            assert compute_logical_path(str(f)) == f"output/{folder_name}/saved.safetensors"
+            assert compute_display_name(str(f)) == f"{folder_name}/saved.safetensors"
+
+            name, tags = get_name_and_tags_from_asset_path(str(f))
+            assert name == "saved.safetensors"
+            assert "output" in tags
+            assert "models" in tags
+            assert f"model_type:{folder_name}" in tags
+            assert folder_name not in tags
+
+    def test_output_model_subfolder_uses_output_storage_file_path(self, fake_dirs):
+        folder_name = "loras"
+        output_model_dir = fake_dirs["output"] / folder_name
+        subdir = output_model_dir / "experiments"
+        subdir.mkdir(parents=True)
+        f = subdir / "my_lora.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[(folder_name, [str(output_model_dir)], {".safetensors"})],
+        ):
+            assert (
+                compute_logical_path(str(f))
+                == "output/loras/experiments/my_lora.safetensors"
+            )
+            assert compute_display_name(str(f)) == "loras/experiments/my_lora.safetensors"
+
+            name, tags = get_name_and_tags_from_asset_path(str(f))
+            assert name == "my_lora.safetensors"
+            assert "output" in tags
+            assert "models" in tags
+            assert "model_type:loras" in tags
+            assert "loras" not in tags
+            assert "experiments" not in tags
+
+    def test_external_model_folder_without_provenance_has_no_file_path(self, tmp_path: Path):
+        external_checkpoints_dir = tmp_path / "external" / "not_named_like_category"
+        external_checkpoints_dir.mkdir(parents=True)
+        f = external_checkpoints_dir / "external.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("checkpoints", [str(external_checkpoints_dir)], {".safetensors"})],
+        ):
+            assert compute_logical_path(str(f)) is None
+            assert compute_display_name(str(f)) is None
+
+            name, tags = get_name_and_tags_from_asset_path(str(f))
+            assert name == "external.safetensors"
+            assert "models" in tags
+            assert "model_type:checkpoints" in tags
+
+    def test_same_relative_model_file_under_multiple_external_roots_has_no_storage_file_path(
+        self, tmp_path: Path
+    ):
+        foo_dir = tmp_path / "foo"
+        bar_dir = tmp_path / "bar"
+        foo_dir.mkdir()
+        bar_dir.mkdir()
+        foo_file = foo_dir / "baz.safetensors"
+        bar_file = bar_dir / "baz.safetensors"
+        foo_file.touch()
+        bar_file.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("checkpoints", [str(foo_dir), str(bar_dir)], {".safetensors"})],
+        ):
+            assert compute_logical_path(str(foo_file)) is None
+            assert compute_logical_path(str(bar_file)) is None
+            assert compute_display_name(str(foo_file)) is None
+            assert compute_display_name(str(bar_file)) is None
+
+    def test_output_clip_folder_uses_output_storage_and_text_encoder_tag(self, fake_dirs):
+        output_clip_dir = fake_dirs["output"] / "clip"
+        output_clip_dir.mkdir()
+        f = output_clip_dir / "clip_l.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("text_encoders", [str(output_clip_dir)], {".safetensors"})],
+        ):
+            assert compute_logical_path(str(f)) == "output/clip/clip_l.safetensors"
+            assert compute_display_name(str(f)) == "clip/clip_l.safetensors"
+
+            name, tags = get_name_and_tags_from_asset_path(str(f))
+            assert name == "clip_l.safetensors"
+            assert "output" in tags
+            assert "models" in tags
+            assert "model_type:text_encoders" in tags
+            assert "clip" not in tags
+
+    def test_physical_unet_folder_uses_storage_path_and_diffusion_models_tag(self, fake_dirs):
+        unet_dir = fake_dirs["models_root"] / "unet"
+        diffusion_models_dir = fake_dirs["models_root"] / "diffusion_models"
+        unet_dir.mkdir()
+        diffusion_models_dir.mkdir()
+        f = unet_dir / "wan.safetensors"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[
+                ("diffusion_models", [str(unet_dir), str(diffusion_models_dir)], {".safetensors"})
+            ],
+        ):
+            assert compute_logical_path(str(f)) == "models/unet/wan.safetensors"
+            assert compute_display_name(str(f)) == "unet/wan.safetensors"
+
+            name, tags = get_name_and_tags_from_asset_path(str(f))
+            assert name == "wan.safetensors"
+            assert "models" in tags
+            assert "model_type:diffusion_models" in tags
+            assert "unet" not in tags
+
+    def test_unregistered_file_under_physical_models_root_still_has_storage_file_path(self, fake_dirs):
+        f = fake_dirs["models_root"] / "not_registered" / "orphan.bin"
+        f.parent.mkdir()
+        f.touch()
+
+        assert compute_logical_path(str(f)) == "models/not_registered/orphan.bin"
+        assert compute_display_name(str(f)) == "not_registered/orphan.bin"
+
+    def test_output_checkpoint_folder_without_registration_has_only_output_tag(self, fake_dirs):
+        f = fake_dirs["output"] / "checkpoints" / "saved.safetensors"
+        f.parent.mkdir(exist_ok=True)
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[],
+        ):
+            assert compute_logical_path(str(f)) == "output/checkpoints/saved.safetensors"
+            assert compute_display_name(str(f)) == "checkpoints/saved.safetensors"
+
+            name, tags = get_name_and_tags_from_asset_path(str(f))
+            assert name == "saved.safetensors"
+            assert "output" in tags
+            assert "models" not in tags
+            assert not any(tag.startswith("model_type:") for tag in tags)
+
+    def test_unknown_path_returns_none(self):
+        assert compute_logical_path("/some/random/path.png") is None
+        assert compute_display_name("/some/random/path.png") is None
+
+
+class TestLoaderPath:
+    """In-root loader path: relative to the storage root, model category dropped."""
+
+    def test_model_loader_path_drops_category(self, fake_dirs):
+        sub = fake_dirs["models"] / "flux"
+        sub.mkdir()
+        f = sub / "model.safetensors"
+        f.touch()
+
+        # logical_path keeps the category, file_path (loader) drops it
+        assert compute_logical_path(str(f)) == "models/checkpoints/flux/model.safetensors"
+        assert compute_loader_path(str(f)) == "flux/model.safetensors"
+
+    def test_model_loader_path_flat_file(self, fake_dirs):
+        f = fake_dirs["models"] / "model.safetensors"
+        f.touch()
+
+        assert compute_loader_path(str(f)) == "model.safetensors"
+
+    def test_input_loader_path_keeps_subfolders(self, fake_dirs):
+        sub = fake_dirs["input"] / "some" / "folder"
+        sub.mkdir(parents=True)
+        f = sub / "image.png"
+        f.touch()
+
+        assert compute_loader_path(str(f)) == "some/folder/image.png"
+
+    def test_temp_loader_path(self, fake_dirs):
+        f = fake_dirs["temp"] / "preview.png"
+        f.touch()
+
+        assert compute_loader_path(str(f)) == "preview.png"
+
+    def test_unregistered_file_under_models_root_has_no_loader_path(self, fake_dirs):
+        # Under models_root but not within any registered category base.
+        f = fake_dirs["models_root"] / "not_registered" / "orphan.bin"
+        f.parent.mkdir()
+        f.touch()
+
+        # It still has a namespaced logical_path, but no loader path.
+        assert compute_logical_path(str(f)) == "models/not_registered/orphan.bin"
+        assert compute_loader_path(str(f)) is None
+
+    def test_extension_mismatch_in_registered_bucket_has_no_loader_path(self, fake_dirs):
+        # Inside a registered bucket, but the bucket's extension set cannot
+        # load it: no model_type tag, and no loader path either.
+        f = fake_dirs["models"] / "notes.txt"
+        f.touch()
+
+        assert compute_logical_path(str(f)) == "models/checkpoints/notes.txt"
+        assert compute_loader_path(str(f)) is None
+
+    def test_shared_base_loader_path_uses_extension_matching_bucket(self, fake_dirs):
+        shared_root = fake_dirs["models"].parent / "unet"
+        shared_root.mkdir()
+        f = shared_root / "wan.gguf"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[
+                ("diffusion_models", [str(shared_root)], {".safetensors"}),
+                ("unet_gguf", [str(shared_root)], {".gguf"}),
+            ],
+        ):
+            assert compute_loader_path(str(f)) == "wan.gguf"
+
+    def test_match_all_bucket_provides_loader_path_for_any_extension(self, fake_dirs):
+        custom_root = fake_dirs["models"].parent / "custom_bucket"
+        custom_root.mkdir()
+        f = custom_root / "weights.bin"
+        f.touch()
+
+        with patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("custom_bucket", [str(custom_root)], set())],
+        ):
+            assert compute_loader_path(str(f)) == "weights.bin"
+
+    def test_extra_path_model_has_loader_path_but_no_logical_path(self, tmp_path: Path):
+        """Registered category base outside models_dir (extra_model_paths style).
+
+        Loadable, so loader_path resolves; but it is not under any canonical
+        storage root, so logical_path/display_name are None. This asymmetry is
+        intentional: loader_path resolves every registered model-folder base,
+        logical_path only resolves the canonical storage roots.
+        """
+        extra = tmp_path / "extra_ckpts"
+        extra.mkdir()
+        f = extra / "foo.safetensors"
+        f.touch()
+
+        with patch("app.assets.services.path_utils.folder_paths") as mock_fp, patch(
+            "app.assets.services.path_utils.get_comfy_models_folders",
+            return_value=[("checkpoints", [str(extra)], {".safetensors"})],
+        ):
+            mock_fp.get_input_directory.return_value = str(tmp_path / "in")
+            mock_fp.get_output_directory.return_value = str(tmp_path / "out")
+            mock_fp.get_temp_directory.return_value = str(tmp_path / "tmp")
+            mock_fp.models_dir = str(tmp_path / "models")  # extra is NOT under this
+
+            assert compute_loader_path(str(f)) == "foo.safetensors"
+            assert compute_logical_path(str(f)) is None
+            assert compute_display_name(str(f)) is None
+
+    def test_unknown_path_returns_none(self):
+        assert compute_loader_path("/some/random/path.png") is None
+
+
+class TestResolveDestinationFromTags:
+    def test_extra_tags_are_not_path_components(self, fake_dirs):
+        base_dir, subdirs = resolve_destination_from_tags(["input", "unit-tests", "foo"])
+
+        assert base_dir == os.path.abspath(fake_dirs["input"])
+        assert subdirs == []
+
+    def test_model_upload_rejects_non_writable_registered_folders(self):
+        with tempfile.TemporaryDirectory() as root:
+            root_path = Path(root)
+            checkpoints_dir = root_path / "models" / "checkpoints"
+            configs_dir = root_path / "models" / "configs"
+            custom_nodes_dir = root_path / "custom_nodes"
+            for path in (checkpoints_dir, configs_dir, custom_nodes_dir):
+                path.mkdir(parents=True)
+
+            with patch("app.assets.services.path_utils.folder_paths") as mock_fp:
+                mock_fp.folder_names_and_paths = {
+                    "checkpoints": ([str(checkpoints_dir)], set()),
+                    "configs": ([str(configs_dir)], set()),
+                    "custom_nodes": ([str(custom_nodes_dir)], set()),
+                }
+
+                base_dir, subdirs = resolve_destination_from_tags(
+                    ["models", "model_type:checkpoints"]
+                )
+                assert base_dir == os.path.abspath(checkpoints_dir)
+                assert subdirs == []
+
+                for folder_name in ("configs", "custom_nodes"):
+                    with pytest.raises(ValueError, match="unknown model category"):
+                        resolve_destination_from_tags(
+                            ["models", f"model_type:{folder_name}"]
+                        )