Merge remote-tracking branch 'origin/master' into temp_pr
This commit is contained in:
commit
cff62dd32e
|
|
@ -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' ||
|
||||
|
|
|
|||
41
AGENTS.md
41
AGENTS.md
|
|
@ -19,6 +19,9 @@
|
|||
better to remove a broken feature path than keep a complicated partial fix.
|
||||
- Preserve existing APIs, node names, model-loading behavior, file layout, and
|
||||
workflow compatibility unless the change is explicitly about replacing them.
|
||||
- When compatibility is explicitly out of scope, remove compatibility-only
|
||||
aliases, duplicate nodes, legacy entry points, and preset wrappers instead of
|
||||
retaining parallel ways to perform the same operation.
|
||||
- Code must look hand-written for this repository. Changes that read like
|
||||
generic AI-generated code will be rejected automatically: unnecessary helper
|
||||
layers, vague names, boilerplate comments, defensive branches without a real
|
||||
|
|
@ -96,6 +99,13 @@
|
|||
unless they are read by current code and change current behavior. Remove
|
||||
pass-through or stored-but-unused values instead of preserving upstream or
|
||||
deprecated API baggage.
|
||||
- Do not add a model-specific option to a shared helper when only one caller
|
||||
needs it. Keep one-off behavior at the model integration boundary, or extend
|
||||
the shared helper only when the option is a coherent reusable capability.
|
||||
- Implementations of shared model interfaces should accept the standard caller
|
||||
contract without model-specific rejection branches for optional capabilities
|
||||
they do not consume. Let supported behavior be determined by implementation
|
||||
paths that actually use those inputs.
|
||||
- If an implementation needs auxiliary values for its own workflow, expose them
|
||||
through a private helper or a clearly named implementation-specific method
|
||||
instead of overloading the public method's return contract.
|
||||
|
|
@ -154,6 +164,10 @@
|
|||
`comfy-kitchen` helpers where they already solve the problem.
|
||||
- Use optimized comfy-kitchen ops in places where they improve performance
|
||||
without changing the expected dtype, device, memory, or interface behavior.
|
||||
- Prefer ComfyUI's shared optimized kernels and backend dispatchers over
|
||||
handwritten implementations of the same operation. Remove duplicate local
|
||||
kernels and adapt inputs to the shared operation's documented layout while
|
||||
preserving the model's original math and output contract.
|
||||
- All models should use the optimized attention function selected by ComfyUI.
|
||||
Treat optimized backend functions, dispatch helpers, and capability-selected
|
||||
callables as opaque. Higher-level code must not inspect function identity,
|
||||
|
|
@ -176,6 +190,12 @@
|
|||
- Model detection code that inspects linear weight shapes should only use the
|
||||
first dimension. The second dimension may be half the original size for
|
||||
NVFP4 or other 4-bit quantized models.
|
||||
- A model-detection signature must guard every state-dict key it dereferences.
|
||||
Do not partially match a format and then raise an incidental `KeyError` while
|
||||
extracting its configuration.
|
||||
- Order model-detection checks from established or more-specific signatures to
|
||||
newer or broader signatures. Put a broad new detector near the generic
|
||||
fallback when giving it higher precedence could steal another model family.
|
||||
- Avoid adding `einops` usage in core inference code. Use native torch tensor
|
||||
ops such as `reshape`, `view`, `permute`, `transpose`, `flatten`, `unflatten`,
|
||||
`unsqueeze`, and `squeeze` instead.
|
||||
|
|
@ -192,11 +212,23 @@
|
|||
methods for scalar or structural calculations.
|
||||
- Avoid unnecessary casts and transfers. Preserve the intended compute dtype,
|
||||
storage dtype, bias dtype, and original tensor shape metadata.
|
||||
- Do not cast the result of an optimized backend operation back to its input
|
||||
dtype unless that backend's documented result contract requires normalization.
|
||||
In particular, trust the selected optimized-attention implementation to honor
|
||||
its dtype contract.
|
||||
- Keep model-native latent layout handling inside the model or latent-format
|
||||
owner, not in helper nodes. Do not collapse, expand, pack, or unpack latent
|
||||
dimensions in nodes or other caller-side adapters just to satisfy a model
|
||||
forward; the model path should consume and return the native latent shape for
|
||||
that model family.
|
||||
- DiT models should accept latent dimensions that are not exact patch-size
|
||||
multiples. Use `comfy.ldm.common_dit.pad_to_patch_size` on every patchified
|
||||
target or reference input, then crop only the target output back to its
|
||||
original dimensions.
|
||||
- Avoid defensive shape and configuration checks that merely replace the clear
|
||||
failure from the tensor operation immediately below them. Add explicit
|
||||
validation only when it provides materially better context at a real boundary
|
||||
or prevents silent incorrect output.
|
||||
- Assume inputs to the main model forward are already in the compute dtype by
|
||||
default, except integer inputs such as some model timestep tensors. Do not add
|
||||
defensive or convenience casts in model code; it is better for invalid dtype
|
||||
|
|
@ -260,6 +292,15 @@
|
|||
- Model implementations should add the minimal number of ComfyUI nodes required
|
||||
to run the model. Reuse existing nodes as much as possible; adapting the model
|
||||
to work with existing nodes is strongly preferred over creating new nodes.
|
||||
- Use `io.Autogrow` for a variable number of repeated inputs instead of a fixed
|
||||
series of numbered optional sockets. Set its minimum to zero when the model
|
||||
has a valid no-item path, and cap it only when the model has a real limit.
|
||||
- Mark inputs optional when execution has a valid path that does not read them.
|
||||
If one optional input is needed only to process another optional input, do not
|
||||
force users on the path that supplies neither to connect it.
|
||||
- Conditioning nodes should normally output conditioning only. Do not expose
|
||||
input or intermediate images as convenience outputs for downstream sizing or
|
||||
routing; use the existing image path or a dedicated image operation instead.
|
||||
- Nodes should output only values they own. Do not add pass-through outputs for
|
||||
workflow convenience unless the node is explicitly an output node. Existing
|
||||
models, latents, conditioning, or other inputs should flow directly to the
|
||||
|
|
|
|||
|
|
@ -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)"
|
||||
)
|
||||
|
|
@ -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")
|
||||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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'."
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}`
|
||||
|
|
|
|||
|
|
@ -92,6 +92,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"
|
||||
|
|
|
|||
|
|
@ -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 "")
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,278 @@
|
|||
import re
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
import comfy.ops
|
||||
import comfy.utils
|
||||
|
||||
|
||||
MODULE_PATTERN = re.compile(r"lllite_dit_blocks_(\d+)_(self_attn_[qkv]_proj|cross_attn_q_proj|mlp_layer1)$")
|
||||
|
||||
|
||||
def _group_norm(channels, device=None, dtype=None, operations=None):
|
||||
groups = 8
|
||||
while groups > 1 and channels % groups != 0:
|
||||
groups //= 2
|
||||
return operations.GroupNorm(groups, channels, device=device, dtype=dtype)
|
||||
|
||||
|
||||
class AnimaLLLiteResBlock(nn.Module):
|
||||
def __init__(self, channels, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.norm1 = _group_norm(channels, device=device, dtype=dtype, operations=operations)
|
||||
self.conv1 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype)
|
||||
self.norm2 = _group_norm(channels, device=device, dtype=dtype, operations=operations)
|
||||
self.conv2 = operations.Conv2d(channels, channels, kernel_size=3, padding=1, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
h = self.conv1(F.silu(self.norm1(x)))
|
||||
h = self.conv2(F.silu(self.norm2(h)))
|
||||
return x + h
|
||||
|
||||
|
||||
class AnimaLLLiteASPP(nn.Module):
|
||||
def __init__(self, channels, dilations, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
branches = []
|
||||
for dilation in dilations:
|
||||
if dilation == 1:
|
||||
conv = operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype)
|
||||
else:
|
||||
conv = operations.Conv2d(channels, channels, kernel_size=3, padding=dilation, dilation=dilation, device=device, dtype=dtype)
|
||||
branches.append(nn.Sequential(conv, _group_norm(channels, device=device, dtype=dtype, operations=operations), nn.SiLU()))
|
||||
self.branches = nn.ModuleList(branches)
|
||||
self.global_pool = nn.AdaptiveAvgPool2d(1)
|
||||
self.global_conv = nn.Sequential(
|
||||
operations.Conv2d(channels, channels, kernel_size=1, device=device, dtype=dtype),
|
||||
_group_norm(channels, device=device, dtype=dtype, operations=operations),
|
||||
nn.SiLU(),
|
||||
)
|
||||
self.proj = nn.Sequential(
|
||||
operations.Conv2d(channels * (len(dilations) + 1), channels, kernel_size=1, device=device, dtype=dtype),
|
||||
_group_norm(channels, device=device, dtype=dtype, operations=operations),
|
||||
nn.SiLU(),
|
||||
)
|
||||
|
||||
def forward(self, x):
|
||||
height, width = x.shape[-2:]
|
||||
outputs = [branch(x) for branch in self.branches]
|
||||
pooled = self.global_conv(self.global_pool(x))
|
||||
outputs.append(F.interpolate(pooled, size=(height, width), mode="bilinear", align_corners=False))
|
||||
return self.proj(torch.cat(outputs, dim=1))
|
||||
|
||||
|
||||
class AnimaLLLiteConditioning(nn.Module):
|
||||
def __init__(self, cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
half_dim = cond_dim // 2
|
||||
self.conv1 = operations.Conv2d(cond_in_channels, half_dim, kernel_size=4, stride=4, device=device, dtype=dtype)
|
||||
self.norm1 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations)
|
||||
self.conv2 = operations.Conv2d(half_dim, half_dim, kernel_size=3, padding=1, device=device, dtype=dtype)
|
||||
self.norm2 = _group_norm(half_dim, device=device, dtype=dtype, operations=operations)
|
||||
self.conv3 = operations.Conv2d(half_dim, cond_dim, kernel_size=4, stride=4, device=device, dtype=dtype)
|
||||
self.norm3 = _group_norm(cond_dim, device=device, dtype=dtype, operations=operations)
|
||||
self.resblocks = nn.ModuleList([
|
||||
AnimaLLLiteResBlock(cond_dim, device=device, dtype=dtype, operations=operations)
|
||||
for _ in range(cond_resblocks)
|
||||
])
|
||||
self.aspp = AnimaLLLiteASPP(cond_dim, aspp_dilations, device=device, dtype=dtype, operations=operations) if aspp_dilations else None
|
||||
self.proj = operations.Conv2d(cond_dim, cond_emb_dim, kernel_size=1, device=device, dtype=dtype)
|
||||
self.out_norm = operations.LayerNorm(cond_emb_dim, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
x = F.silu(self.norm1(self.conv1(x)))
|
||||
x = F.silu(self.norm2(self.conv2(x)))
|
||||
x = F.silu(self.norm3(self.conv3(x)))
|
||||
for block in self.resblocks:
|
||||
x = block(x)
|
||||
if self.aspp is not None:
|
||||
x = self.aspp(x)
|
||||
x = self.proj(x).flatten(2).transpose(1, 2).contiguous()
|
||||
return self.out_norm(x)
|
||||
|
||||
|
||||
class AnimaLLLiteModule(nn.Module):
|
||||
def __init__(self, in_dim, cond_emb_dim, mlp_dim, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.down = operations.Linear(in_dim, mlp_dim, device=device, dtype=dtype)
|
||||
self.mid = operations.Linear(mlp_dim + cond_emb_dim, mlp_dim, device=device, dtype=dtype)
|
||||
self.cond_to_film = operations.Linear(cond_emb_dim, 2 * mlp_dim, device=device, dtype=dtype)
|
||||
self.up = operations.Linear(mlp_dim, in_dim, device=device, dtype=dtype)
|
||||
self.depth_embed = nn.Parameter(torch.empty(cond_emb_dim, device=device, dtype=dtype), requires_grad=False)
|
||||
|
||||
def forward(self, x, cond_emb, strength):
|
||||
original_shape = x.shape
|
||||
if x.ndim == 5:
|
||||
x = x.flatten(1, 3)
|
||||
|
||||
if x.shape[0] != cond_emb.shape[0]:
|
||||
if x.shape[0] % cond_emb.shape[0] != 0:
|
||||
raise ValueError(f"Anima LLLite batch mismatch: model input batch {x.shape[0]}, control batch {cond_emb.shape[0]}")
|
||||
cond_emb = cond_emb.repeat(x.shape[0] // cond_emb.shape[0], 1, 1)
|
||||
if x.shape[1] != cond_emb.shape[1]:
|
||||
raise ValueError(f"Anima LLLite sequence mismatch: model input has {x.shape[1]} tokens, control has {cond_emb.shape[1]}")
|
||||
|
||||
cond_local = cond_emb + comfy.ops.cast_to_input(self.depth_embed, cond_emb)
|
||||
hidden = F.silu(self.down(x))
|
||||
gamma, beta = self.cond_to_film(cond_local).chunk(2, dim=-1)
|
||||
hidden = self.mid(torch.cat((cond_local, hidden), dim=-1))
|
||||
hidden = F.silu(hidden * (1 + gamma) + beta)
|
||||
x = x + self.up(hidden) * strength
|
||||
|
||||
if len(original_shape) == 5:
|
||||
x = x.reshape(original_shape)
|
||||
return x
|
||||
|
||||
|
||||
class AnimaLLLite(nn.Module):
|
||||
def __init__(self, state_dict, metadata, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
metadata = metadata or {}
|
||||
version = metadata.get("lllite.version", "2")
|
||||
if version != "2":
|
||||
raise ValueError(f"Unsupported Anima LLLite version {version!r}; only named-key v2 checkpoints are supported")
|
||||
|
||||
module_names = sorted({key.split(".", 1)[0] for key in state_dict if key.startswith("lllite_dit_blocks_")})
|
||||
if not module_names:
|
||||
raise ValueError("Anima LLLite checkpoint has no lllite_dit_blocks_* modules")
|
||||
|
||||
cond_in_channels = state_dict["lllite_conditioning1.conv1.weight"].shape[1]
|
||||
cond_dim = state_dict["lllite_conditioning1.conv3.weight"].shape[0]
|
||||
cond_emb_dim = state_dict["lllite_conditioning1.proj.weight"].shape[0]
|
||||
resblock_ids = {int(key.split(".")[2]) for key in state_dict if key.startswith("lllite_conditioning1.resblocks.")}
|
||||
cond_resblocks = max(resblock_ids) + 1 if resblock_ids else 0
|
||||
use_aspp = any(key.startswith("lllite_conditioning1.aspp.") for key in state_dict)
|
||||
dilation_string = metadata.get("lllite.aspp_dilations", "1,2,4,8")
|
||||
aspp_dilations = tuple(int(value) for value in dilation_string.split(",") if value.strip()) if use_aspp else ()
|
||||
|
||||
self.cond_in_channels = cond_in_channels
|
||||
self.inpaint_masked_input = metadata.get("lllite.inpaint_masked_input", "false").lower() == "true"
|
||||
self.lllite_conditioning1 = AnimaLLLiteConditioning(
|
||||
cond_in_channels, cond_dim, cond_emb_dim, cond_resblocks, aspp_dilations,
|
||||
device=device, dtype=dtype, operations=operations,
|
||||
)
|
||||
|
||||
self.module_names = set()
|
||||
self.block_count = 0
|
||||
self.model_dim = None
|
||||
for name in module_names:
|
||||
match = MODULE_PATTERN.fullmatch(name)
|
||||
if match is None:
|
||||
raise ValueError(f"Unsupported Anima LLLite module name: {name}")
|
||||
down_shape = state_dict[f"{name}.down.weight"].shape
|
||||
mlp_dim, in_dim = down_shape
|
||||
module_cond_dim = state_dict[f"{name}.cond_to_film.weight"].shape[1]
|
||||
if module_cond_dim != cond_emb_dim:
|
||||
raise ValueError(f"Anima LLLite conditioning dimension mismatch in {name}: {module_cond_dim} != {cond_emb_dim}")
|
||||
if self.model_dim is None:
|
||||
self.model_dim = in_dim
|
||||
elif self.model_dim != in_dim:
|
||||
raise ValueError(f"Anima LLLite model dimension mismatch in {name}: {in_dim} != {self.model_dim}")
|
||||
self.add_module(name, AnimaLLLiteModule(in_dim, cond_emb_dim, mlp_dim, device=device, dtype=dtype, operations=operations))
|
||||
self.module_names.add(name)
|
||||
self.block_count = max(self.block_count, int(match.group(1)) + 1)
|
||||
|
||||
def encode_conditioning(self, image):
|
||||
return self.lllite_conditioning1(image)
|
||||
|
||||
def apply(self, x, cond_emb, block_index, target, strength):
|
||||
name = f"lllite_dit_blocks_{block_index}_{target}"
|
||||
if name not in self.module_names:
|
||||
return x
|
||||
return self.get_submodule(name)(x, cond_emb, strength)
|
||||
|
||||
|
||||
class AnimaLLLitePatch:
|
||||
def __init__(self, model_patch, image, mask, strength, sigma_start, sigma_end):
|
||||
self.model_patch = model_patch
|
||||
self.image = image
|
||||
self.mask = mask
|
||||
self.strength = strength
|
||||
self.sigma_start = sigma_start
|
||||
self.sigma_end = sigma_end
|
||||
|
||||
def __call__(self, args):
|
||||
x = args["x"]
|
||||
transformer_options = args["transformer_options"]
|
||||
if self.strength == 0.0:
|
||||
return args
|
||||
sigmas = transformer_options.get("sigmas")
|
||||
if sigmas is not None:
|
||||
sigma = float(sigmas.max().item())
|
||||
if not self.sigma_end <= sigma <= self.sigma_start:
|
||||
return args
|
||||
if x.shape[2] != 1:
|
||||
raise ValueError(f"Anima LLLite only supports T=1, got T={x.shape[2]}")
|
||||
|
||||
target_height = x.shape[-2] * 8
|
||||
target_width = x.shape[-1] * 8
|
||||
image = comfy.utils.common_upscale(
|
||||
self.image.movedim(-1, 1), target_width, target_height, "bicubic", crop="center"
|
||||
).clamp(0.0, 1.0)
|
||||
image = image.to(device=x.device, dtype=x.dtype) * 2.0 - 1.0
|
||||
|
||||
if self.model_patch.model.cond_in_channels == 4:
|
||||
mask = self.mask
|
||||
if mask.ndim == 3:
|
||||
mask = mask.unsqueeze(1)
|
||||
if mask.ndim != 4 or mask.shape[1] != 1:
|
||||
raise ValueError(f"Anima LLLite mask must have one channel, got shape {tuple(mask.shape)}")
|
||||
mask = comfy.utils.common_upscale(
|
||||
mask.float(), target_width, target_height, "nearest-exact", crop="center"
|
||||
)
|
||||
if mask.shape[0] != image.shape[0]:
|
||||
if image.shape[0] % mask.shape[0] != 0:
|
||||
raise ValueError(
|
||||
f"Anima LLLite mask batch {mask.shape[0]} cannot be broadcast to image batch {image.shape[0]}"
|
||||
)
|
||||
mask = mask.repeat(image.shape[0] // mask.shape[0], 1, 1, 1)
|
||||
mask = (mask >= 0.5).to(device=x.device, dtype=x.dtype)
|
||||
if self.model_patch.model.inpaint_masked_input:
|
||||
image = image * (mask < 0.5).to(image.dtype)
|
||||
image = torch.cat((image, mask * 2.0 - 1.0), dim=1)
|
||||
|
||||
cond_emb = self.model_patch.model.encode_conditioning(image)
|
||||
transformer_options["model_patch_data"][self] = cond_emb
|
||||
return args
|
||||
|
||||
def to(self, device_or_dtype):
|
||||
return self
|
||||
|
||||
def models(self):
|
||||
return [self.model_patch]
|
||||
|
||||
|
||||
class AnimaLLLiteAttentionPatch:
|
||||
def __init__(self, patch, targets):
|
||||
self.patch = patch
|
||||
self.targets = targets
|
||||
|
||||
def __call__(self, q, k, v, pe=None, attn_mask=None, extra_options=None):
|
||||
cond_emb = extra_options["model_patch_data"].get(self.patch)
|
||||
if cond_emb is None:
|
||||
return {"q": q, "k": k, "v": v, "pe": pe, "attn_mask": attn_mask}
|
||||
|
||||
block_index = extra_options["block_index"]
|
||||
values = {"q": q, "k": k, "v": v}
|
||||
for value_name, target in self.targets.items():
|
||||
values[value_name] = self.patch.model_patch.model.apply(
|
||||
values[value_name], cond_emb, block_index, target, self.patch.strength
|
||||
)
|
||||
|
||||
return {"q": values["q"], "k": values["k"], "v": values["v"], "pe": pe, "attn_mask": attn_mask}
|
||||
|
||||
|
||||
class AnimaLLLiteMLPPatch:
|
||||
def __init__(self, patch):
|
||||
self.patch = patch
|
||||
|
||||
def __call__(self, args):
|
||||
cond_emb = args["transformer_options"]["model_patch_data"].get(self.patch)
|
||||
if cond_emb is None:
|
||||
return args
|
||||
args["x"] = self.patch.model_patch.model.apply(
|
||||
args["x"], cond_emb, args["transformer_options"]["block_index"], "mlp_layer1", self.patch.strength
|
||||
)
|
||||
return args
|
||||
|
|
@ -14,6 +14,7 @@ from torchvision import transforms
|
|||
import comfy.patcher_extension
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.ops
|
||||
import comfy.quant_ops
|
||||
|
||||
|
||||
|
|
@ -148,11 +149,29 @@ class Attention(nn.Module):
|
|||
x: torch.Tensor,
|
||||
context: Optional[torch.Tensor] = None,
|
||||
rope_emb: Optional[torch.Tensor] = None,
|
||||
transformer_options: Optional[dict] = {},
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
q = self.q_proj(x)
|
||||
context = x if context is None else context
|
||||
k = self.k_proj(context)
|
||||
v = self.v_proj(context)
|
||||
q_input = x
|
||||
k_input = context
|
||||
v_input = context
|
||||
|
||||
transformer_patches = transformer_options.get("patches", {})
|
||||
patch_name = "attn1_patch" if self.is_selfattn else "attn2_patch"
|
||||
if patch_name in transformer_patches:
|
||||
extra_options = transformer_options.copy()
|
||||
extra_options["n_heads"] = self.n_heads
|
||||
extra_options["dim_head"] = self.head_dim
|
||||
for patch in transformer_patches[patch_name]:
|
||||
out = patch(q_input, k_input, v_input, pe=rope_emb, attn_mask=None, extra_options=extra_options)
|
||||
q_input = out.get("q", q_input)
|
||||
k_input = out.get("k", k_input)
|
||||
v_input = out.get("v", v_input)
|
||||
rope_emb = out.get("pe", rope_emb)
|
||||
|
||||
q = self.q_proj(q_input)
|
||||
k = self.k_proj(k_input)
|
||||
v = self.v_proj(v_input)
|
||||
q, k, v = map(
|
||||
lambda t: rearrange(t, "b ... (h d) -> b ... h d", h=self.n_heads, d=self.head_dim),
|
||||
(q, k, v),
|
||||
|
|
@ -161,11 +180,16 @@ class Attention(nn.Module):
|
|||
def apply_norm_and_rotary_pos_emb(
|
||||
q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, rope_emb: Optional[torch.Tensor]
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
v = self.v_norm(v)
|
||||
if self.is_selfattn and rope_emb is not None: # only apply to self-attention!
|
||||
q, k = comfy.quant_ops.ck.apply_rope_split_half(q, k, rope_emb)
|
||||
q_scale, _, q_offload_stream = comfy.ops.cast_bias_weight(self.q_norm, q, offloadable=True)
|
||||
k_scale, _, k_offload_stream = comfy.ops.cast_bias_weight(self.k_norm, k, offloadable=True)
|
||||
q, k = comfy.quant_ops.ck.rms_rope_split_half(q, k, rope_emb, q_scale, k_scale, self.q_norm.eps)
|
||||
comfy.ops.uncast_bias_weight(self.q_norm, q_scale, None, q_offload_stream)
|
||||
comfy.ops.uncast_bias_weight(self.k_norm, k_scale, None, k_offload_stream)
|
||||
else:
|
||||
q = self.q_norm(q)
|
||||
k = self.k_norm(k)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = apply_norm_and_rotary_pos_emb(q, k, v, rope_emb)
|
||||
|
|
@ -188,7 +212,7 @@ class Attention(nn.Module):
|
|||
x (Tensor): The query tensor of shape [B, Mq, K]
|
||||
context (Optional[Tensor]): The key tensor of shape [B, Mk, K] or use x as context [self attention] if None
|
||||
"""
|
||||
q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb)
|
||||
q, k, v = self.compute_qkv(x, context, rope_emb=rope_emb, transformer_options=transformer_options)
|
||||
return self.compute_attention(q, k, v, transformer_options=transformer_options)
|
||||
|
||||
|
||||
|
|
@ -555,8 +579,14 @@ class Block(nn.Module):
|
|||
self.layer_norm_mlp,
|
||||
scale_mlp_B_T_1_1_D,
|
||||
shift_mlp_B_T_1_1_D,
|
||||
)
|
||||
result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D.to(compute_dtype))
|
||||
).to(compute_dtype)
|
||||
patches = transformer_options.get("patches", {})
|
||||
if "mlp_patch" in patches:
|
||||
args = {"x": normalized_x_B_T_H_W_D, "transformer_options": transformer_options}
|
||||
for patch in patches["mlp_patch"]:
|
||||
args = patch(args)
|
||||
normalized_x_B_T_H_W_D = args["x"]
|
||||
result_B_T_H_W_D = self.mlp(normalized_x_B_T_H_W_D)
|
||||
x_B_T_H_W_D = torch.addcmul(x_B_T_H_W_D, gate_mlp_B_T_1_1_D.to(residual_dtype), result_B_T_H_W_D.to(residual_dtype))
|
||||
return x_B_T_H_W_D
|
||||
|
||||
|
|
@ -863,11 +893,22 @@ class MiniTrainDIT(nn.Module):
|
|||
x_B_T_H_W_D.shape == extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape
|
||||
), f"{x_B_T_H_W_D.shape} != {extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D.shape}"
|
||||
|
||||
transformer_options = kwargs.get("transformer_options", {})
|
||||
patches = transformer_options.get("patches", {})
|
||||
if "post_input" in patches:
|
||||
transformer_options = transformer_options.copy()
|
||||
transformer_options["model_patch_data"] = {}
|
||||
|
||||
if "post_input" in patches:
|
||||
for patch in patches["post_input"]:
|
||||
out = patch({"img": x_B_T_H_W_D, "x": x_B_C_T_H_W, "transformer_options": transformer_options})
|
||||
x_B_T_H_W_D = out["img"]
|
||||
|
||||
block_kwargs = {
|
||||
"rope_emb_L_1_1_D": rope_emb_L_1_1_D.unsqueeze(1).unsqueeze(0),
|
||||
"adaln_lora_B_T_3D": adaln_lora_B_T_3D,
|
||||
"extra_per_block_pos_emb": extra_pos_emb_B_T_H_W_D_or_T_H_W_B_D,
|
||||
"transformer_options": kwargs.get("transformer_options", {}),
|
||||
"transformer_options": transformer_options,
|
||||
}
|
||||
|
||||
# The residual stream for this model has large values. To make fp16 compute_dtype work, we keep the residual stream
|
||||
|
|
@ -877,7 +918,8 @@ class MiniTrainDIT(nn.Module):
|
|||
if x_B_T_H_W_D.dtype == torch.float16:
|
||||
x_B_T_H_W_D = x_B_T_H_W_D.float()
|
||||
|
||||
for block in self.blocks:
|
||||
for block_index, block in enumerate(self.blocks):
|
||||
transformer_options["block_index"] = block_index
|
||||
x_B_T_H_W_D = block(
|
||||
x_B_T_H_W_D,
|
||||
t_embedding_B_T_D,
|
||||
|
|
|
|||
|
|
@ -15,24 +15,24 @@ def make_two_pass_attention(ar_len: int, transformer_options=None):
|
|||
The AR pass goes through SDPA directand bypasses wrappers, it is only ~1% of T at typical edit sizes.
|
||||
"""
|
||||
|
||||
def two_pass_attention(q, k, v, heads, **kwargs):
|
||||
def two_pass_attention(q, k, v, heads, enable_gqa=False, **kwargs):
|
||||
B, H, T, D = q.shape
|
||||
|
||||
if T < k.shape[2]: # KV-cache hot path: Q is shorter than K/V (cached AR prefix is in K/V only), all fresh Q positions are in the gen region, single full-attention call
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
elif ar_len >= T:
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True)
|
||||
out = comfy.ops.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa)
|
||||
elif ar_len <= 0:
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options)
|
||||
out = optimized_attention(q, k, v, heads, mask=None, skip_reshape=True, skip_output_reshape=True, transformer_options=transformer_options, enable_gqa=enable_gqa)
|
||||
else:
|
||||
out_ar = comfy.ops.scaled_dot_product_attention(
|
||||
q[:, :, :ar_len], k[:, :, :ar_len], v[:, :, :ar_len],
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True,
|
||||
attn_mask=None, dropout_p=0.0, is_causal=True, enable_gqa=enable_gqa,
|
||||
)
|
||||
out_gen = optimized_attention(
|
||||
q[:, :, ar_len:], k, v, heads,
|
||||
mask=None, skip_reshape=True, skip_output_reshape=True,
|
||||
transformer_options=transformer_options,
|
||||
transformer_options=transformer_options, enable_gqa=enable_gqa,
|
||||
)
|
||||
out = torch.cat([out_ar, out_gen], dim=2)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,445 @@
|
|||
# https://github.com/jdopensource/JoyAI-Image-Edit (Apache 2.0)
|
||||
import math
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import comfy_kitchen
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
import comfy.ldm.common_dit
|
||||
import comfy.ops
|
||||
import comfy.patcher_extension
|
||||
from comfy.ldm.lightricks.model import GELU_approx, PixArtAlphaTextProjection, TimestepEmbedding, Timesteps
|
||||
from comfy.ldm.modules.attention import optimized_attention
|
||||
|
||||
|
||||
class JoyImageModulate(nn.Module):
|
||||
def __init__(self, hidden_size: int, factor: int, dtype=None, device=None):
|
||||
super().__init__()
|
||||
self.factor = factor
|
||||
self.modulate_table = nn.Parameter(
|
||||
torch.empty(1, factor, hidden_size, dtype=dtype, device=device)
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> list:
|
||||
if x.ndim != 3:
|
||||
x = x.unsqueeze(1)
|
||||
table = comfy.ops.cast_to_input(self.modulate_table, x)
|
||||
return [o.squeeze(1) for o in (table + x).chunk(self.factor, dim=1)]
|
||||
|
||||
|
||||
class JoyImageFeedForward(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
inner_dim: int,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.net = nn.ModuleList([
|
||||
GELU_approx(dim, inner_dim, dtype=dtype, device=device, operations=operations),
|
||||
nn.Identity(),
|
||||
operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device),
|
||||
])
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
for module in self.net:
|
||||
x = module(x)
|
||||
return x
|
||||
|
||||
|
||||
class JoyImageAttention(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
eps: float = 1e-6,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.num_attention_heads = num_attention_heads
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
self.img_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device)
|
||||
self.img_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||
self.img_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||
self.img_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device)
|
||||
|
||||
self.txt_attn_qkv = operations.Linear(dim, inner_dim * 3, bias=True, dtype=dtype, device=device)
|
||||
self.txt_attn_q_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||
self.txt_attn_k_norm = operations.RMSNorm(attention_head_dim, eps=eps, dtype=dtype, device=device)
|
||||
self.txt_attn_proj = operations.Linear(inner_dim, dim, bias=True, dtype=dtype, device=device)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
img: torch.Tensor,
|
||||
txt: torch.Tensor,
|
||||
image_rotary_emb: torch.Tensor,
|
||||
transformer_options=None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
heads = self.num_attention_heads
|
||||
|
||||
img_q, img_k, img_v = self.img_attn_qkv(img).chunk(3, dim=-1)
|
||||
txt_q, txt_k, txt_v = self.txt_attn_qkv(txt).chunk(3, dim=-1)
|
||||
|
||||
img_q = img_q.unflatten(-1, (heads, -1))
|
||||
img_k = img_k.unflatten(-1, (heads, -1))
|
||||
img_v = img_v.unflatten(-1, (heads, -1))
|
||||
txt_q = txt_q.unflatten(-1, (heads, -1))
|
||||
txt_k = txt_k.unflatten(-1, (heads, -1))
|
||||
txt_v = txt_v.unflatten(-1, (heads, -1))
|
||||
|
||||
img_q = self.img_attn_q_norm(img_q)
|
||||
img_k = self.img_attn_k_norm(img_k)
|
||||
txt_q = self.txt_attn_q_norm(txt_q)
|
||||
txt_k = self.txt_attn_k_norm(txt_k)
|
||||
|
||||
img_q, img_k = comfy_kitchen.apply_rope(img_q, img_k, image_rotary_emb)
|
||||
|
||||
joint_q = torch.cat([img_q, txt_q], dim=1)
|
||||
joint_k = torch.cat([img_k, txt_k], dim=1)
|
||||
joint_v = torch.cat([img_v, txt_v], dim=1)
|
||||
|
||||
joint_q = joint_q.flatten(2, 3)
|
||||
joint_k = joint_k.flatten(2, 3)
|
||||
joint_v = joint_v.flatten(2, 3)
|
||||
|
||||
joint_out = optimized_attention(joint_q, joint_k, joint_v, heads=heads, transformer_options=transformer_options)
|
||||
|
||||
seq_img = img.shape[1]
|
||||
img_out = joint_out[:, :seq_img, :]
|
||||
txt_out = joint_out[:, seq_img:, :]
|
||||
|
||||
img_out = self.img_attn_proj(img_out)
|
||||
txt_out = self.txt_attn_proj(txt_out)
|
||||
return img_out, txt_out
|
||||
|
||||
|
||||
class JoyImageTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
eps: float = 1e-6,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
mlp_hidden_dim = int(dim * mlp_width_ratio)
|
||||
|
||||
self.img_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device)
|
||||
self.img_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||
self.img_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||
self.img_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
|
||||
|
||||
self.txt_mod = JoyImageModulate(dim, factor=6, dtype=dtype, device=device)
|
||||
self.txt_norm1 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||
self.txt_norm2 = operations.LayerNorm(dim, elementwise_affine=False, eps=eps, dtype=dtype, device=device)
|
||||
self.txt_mlp = JoyImageFeedForward(dim, inner_dim=mlp_hidden_dim, dtype=dtype, device=device, operations=operations)
|
||||
|
||||
self.attn = JoyImageAttention(
|
||||
dim=dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
eps=eps,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
temb: torch.Tensor,
|
||||
image_rotary_emb: torch.Tensor,
|
||||
transformer_options=None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
(
|
||||
img_mod1_shift,
|
||||
img_mod1_scale,
|
||||
img_mod1_gate,
|
||||
img_mod2_shift,
|
||||
img_mod2_scale,
|
||||
img_mod2_gate,
|
||||
) = self.img_mod(temb)
|
||||
(
|
||||
txt_mod1_shift,
|
||||
txt_mod1_scale,
|
||||
txt_mod1_gate,
|
||||
txt_mod2_shift,
|
||||
txt_mod2_scale,
|
||||
txt_mod2_gate,
|
||||
) = self.txt_mod(temb)
|
||||
|
||||
img_normed = self.img_norm1(hidden_states)
|
||||
txt_normed = self.txt_norm1(encoder_hidden_states)
|
||||
img_modulated = img_normed * (1 + img_mod1_scale.unsqueeze(1)) + img_mod1_shift.unsqueeze(1)
|
||||
txt_modulated = txt_normed * (1 + txt_mod1_scale.unsqueeze(1)) + txt_mod1_shift.unsqueeze(1)
|
||||
|
||||
img_attn, txt_attn = self.attn(img_modulated, txt_modulated, image_rotary_emb, transformer_options=transformer_options)
|
||||
|
||||
hidden_states = hidden_states + img_attn * img_mod1_gate.unsqueeze(1)
|
||||
encoder_hidden_states = encoder_hidden_states + txt_attn * txt_mod1_gate.unsqueeze(1)
|
||||
|
||||
img_ffn_normed = self.img_norm2(hidden_states)
|
||||
txt_ffn_normed = self.txt_norm2(encoder_hidden_states)
|
||||
img_ffn_input = img_ffn_normed * (1 + img_mod2_scale.unsqueeze(1)) + img_mod2_shift.unsqueeze(1)
|
||||
txt_ffn_input = txt_ffn_normed * (1 + txt_mod2_scale.unsqueeze(1)) + txt_mod2_shift.unsqueeze(1)
|
||||
hidden_states = hidden_states + self.img_mlp(img_ffn_input) * img_mod2_gate.unsqueeze(1)
|
||||
encoder_hidden_states = encoder_hidden_states + self.txt_mlp(txt_ffn_input) * txt_mod2_gate.unsqueeze(1)
|
||||
|
||||
return hidden_states, encoder_hidden_states
|
||||
|
||||
|
||||
class JoyImageTimeTextImageEmbedding(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
time_freq_dim: int,
|
||||
time_proj_dim: int,
|
||||
text_embed_dim: int,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.timesteps_proj = Timesteps(num_channels=time_freq_dim, flip_sin_to_cos=True, downscale_freq_shift=0)
|
||||
self.time_embedder = TimestepEmbedding(
|
||||
in_channels=time_freq_dim,
|
||||
time_embed_dim=dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
self.act_fn = nn.SiLU()
|
||||
self.time_proj = operations.Linear(dim, time_proj_dim, bias=True, dtype=dtype, device=device)
|
||||
self.text_embedder = PixArtAlphaTextProjection(
|
||||
text_embed_dim, dim, act_fn="gelu_tanh", dtype=dtype, device=device, operations=operations,
|
||||
)
|
||||
|
||||
def forward(self, timestep: torch.Tensor, encoder_hidden_states: torch.Tensor):
|
||||
timestep = self.timesteps_proj(timestep)
|
||||
temb = self.time_embedder(timestep.to(dtype=encoder_hidden_states.dtype)).type_as(encoder_hidden_states)
|
||||
timestep_proj = self.time_proj(self.act_fn(temb))
|
||||
encoder_hidden_states = self.text_embedder(encoder_hidden_states)
|
||||
return temb, timestep_proj, encoder_hidden_states
|
||||
|
||||
|
||||
class JoyImageTransformer3DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
patch_size: list = [1, 2, 2],
|
||||
in_channels: int = 16,
|
||||
out_channels: Optional[int] = None,
|
||||
hidden_size: int = 3072,
|
||||
num_attention_heads: int = 24,
|
||||
text_dim: int = 4096,
|
||||
mlp_width_ratio: float = 4.0,
|
||||
num_layers: int = 20,
|
||||
rope_dim_list: list = [16, 56, 56],
|
||||
theta: int = 256,
|
||||
image_model=None,
|
||||
dtype=None,
|
||||
device=None,
|
||||
operations=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
self.out_channels = out_channels or in_channels
|
||||
self.patch_size = list(patch_size)
|
||||
self.rope_dim_list = list(rope_dim_list)
|
||||
self.theta = theta
|
||||
|
||||
attention_head_dim = hidden_size // num_attention_heads
|
||||
|
||||
self.img_in = operations.Conv3d(
|
||||
in_channels,
|
||||
hidden_size,
|
||||
kernel_size=tuple(self.patch_size),
|
||||
stride=tuple(self.patch_size),
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
self.condition_embedder = JoyImageTimeTextImageEmbedding(
|
||||
dim=hidden_size,
|
||||
time_freq_dim=256,
|
||||
time_proj_dim=hidden_size * 6,
|
||||
text_embed_dim=text_dim,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
|
||||
self.double_blocks = nn.ModuleList([
|
||||
JoyImageTransformerBlock(
|
||||
dim=hidden_size,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
mlp_width_ratio=mlp_width_ratio,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
operations=operations,
|
||||
)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
|
||||
self.norm_out = operations.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, dtype=dtype, device=device)
|
||||
self.proj_out = operations.Linear(
|
||||
hidden_size,
|
||||
self.out_channels * math.prod(self.patch_size),
|
||||
bias=True,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
def _get_rotary_pos_embed_for_range(
|
||||
self,
|
||||
start: Tuple[int, int, int],
|
||||
stop: Tuple[int, int, int],
|
||||
device=None,
|
||||
) -> torch.Tensor:
|
||||
# 3D RoPE for the patch grid range [start, stop) over (t, h, w). Token order after
|
||||
# reshape(-1) is (t, h, w), matching the img_in Conv3d flatten.
|
||||
rope_dim_list = self.rope_dim_list
|
||||
|
||||
grids = [torch.arange(start[i], stop[i], dtype=torch.float32, device=device) for i in range(3)]
|
||||
mesh = torch.stack(torch.meshgrid(*grids, indexing="ij"), dim=0)
|
||||
|
||||
angles_parts = []
|
||||
for i, dim in enumerate(rope_dim_list):
|
||||
pos = mesh[i].reshape(-1)
|
||||
freqs = 1.0 / (self.theta ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device)[: (dim // 2)] / dim))
|
||||
angles_parts.append(torch.outer(pos, freqs))
|
||||
|
||||
angles = torch.cat(angles_parts, dim=1)
|
||||
cos = angles.cos()
|
||||
sin = angles.sin()
|
||||
return torch.stack((cos, -sin, sin, cos), dim=-1).unflatten(-1, (2, 2))
|
||||
|
||||
def get_rotary_pos_embed_for_components(
|
||||
self,
|
||||
component_sizes,
|
||||
device=None,
|
||||
) -> torch.Tensor:
|
||||
# Per-component 3D RoPE. component_sizes is a list of (t, h, w) patch grid sizes in
|
||||
# sequence order [target, ref0, ref1, ...]; h/w restart at 0 for each component while t
|
||||
# continues from the running offset, giving every image its own temporal position band.
|
||||
freqs_parts = []
|
||||
t_offset = 0
|
||||
for (t, h, w) in component_sizes:
|
||||
freqs = self._get_rotary_pos_embed_for_range(
|
||||
start=(t_offset, 0, 0),
|
||||
stop=(t_offset + t, h, w),
|
||||
device=device,
|
||||
)
|
||||
freqs_parts.append(freqs)
|
||||
t_offset += t
|
||||
return torch.cat(freqs_parts, dim=0).unsqueeze(0).unsqueeze(2)
|
||||
|
||||
def unpatchify(self, x: torch.Tensor, t: int, h: int, w: int) -> torch.Tensor:
|
||||
c = self.out_channels
|
||||
pt, ph, pw = self.patch_size
|
||||
x = x.reshape(x.shape[0], t, h, w, pt, ph, pw, c)
|
||||
x = x.permute(0, 7, 1, 4, 2, 5, 3, 6)
|
||||
return x.reshape(x.shape[0], c, t * pt, h * ph, w * pw)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
context: torch.Tensor = None,
|
||||
ref_latents=None,
|
||||
control=None,
|
||||
transformer_options=None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
transformer_options = {} if transformer_options is None else transformer_options.copy()
|
||||
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
self._forward,
|
||||
self,
|
||||
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options)
|
||||
).execute(hidden_states, timestep, context, ref_latents, transformer_options, **kwargs)
|
||||
|
||||
def _forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
timestep: torch.Tensor,
|
||||
context: torch.Tensor,
|
||||
ref_latents=None,
|
||||
transformer_options=None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
pt, ph, pw = self.patch_size
|
||||
_, _, ot, oh, ow = hidden_states.shape
|
||||
|
||||
components = [hidden_states, *(ref_latents or [])]
|
||||
component_sizes = []
|
||||
img_tokens = []
|
||||
for comp in components:
|
||||
comp = comfy.ldm.common_dit.pad_to_patch_size(comp, self.patch_size)
|
||||
_, _, ct, ch, cw = comp.shape
|
||||
component_sizes.append((ct // pt, ch // ph, cw // pw))
|
||||
tokens = self.img_in(comp).flatten(2).transpose(1, 2) # (B, n_i, D)
|
||||
img_tokens.append(tokens)
|
||||
|
||||
img = torch.cat(img_tokens, dim=1)
|
||||
|
||||
_, vec, txt = self.condition_embedder(timestep, context)
|
||||
vec = vec.unflatten(1, (6, -1))
|
||||
|
||||
image_rotary_emb = self.get_rotary_pos_embed_for_components(
|
||||
component_sizes,
|
||||
device=hidden_states.device,
|
||||
)
|
||||
|
||||
patches_replace = transformer_options.get("patches_replace", {})
|
||||
blocks_replace = patches_replace.get("dit", {})
|
||||
transformer_options["total_blocks"] = len(self.double_blocks)
|
||||
transformer_options["block_type"] = "double"
|
||||
for i, block in enumerate(self.double_blocks):
|
||||
transformer_options["block_index"] = i
|
||||
if ("double_block", i) in blocks_replace:
|
||||
def block_wrap(args):
|
||||
out = {}
|
||||
out["img"], out["txt"] = block(
|
||||
hidden_states=args["img"],
|
||||
encoder_hidden_states=args["txt"],
|
||||
temb=args["vec"],
|
||||
image_rotary_emb=args["pe"],
|
||||
transformer_options=args.get("transformer_options"),
|
||||
)
|
||||
return out
|
||||
|
||||
out = blocks_replace[("double_block", i)]({"img": img,
|
||||
"txt": txt,
|
||||
"vec": vec,
|
||||
"pe": image_rotary_emb,
|
||||
"transformer_options": transformer_options},
|
||||
{"original_block": block_wrap})
|
||||
txt = out["txt"]
|
||||
img = out["img"]
|
||||
else:
|
||||
img, txt = block(
|
||||
hidden_states=img,
|
||||
encoder_hidden_states=txt,
|
||||
temb=vec,
|
||||
image_rotary_emb=image_rotary_emb,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
|
||||
tt, th, tw = component_sizes[0]
|
||||
target_tokens = tt * th * tw
|
||||
img = img[:, :target_tokens, :]
|
||||
img = self.proj_out(self.norm_out(img))
|
||||
img = self.unpatchify(img, tt, th, tw)
|
||||
return img[:, :, :ot, :oh, :ow]
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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).
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
|
|
@ -55,8 +55,10 @@ 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.joyimage.model
|
||||
import comfy.ldm.ideogram4.model
|
||||
import comfy.ldm.krea2.model
|
||||
import comfy.ldm.kandinsky5.model
|
||||
|
|
@ -932,6 +934,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)
|
||||
|
|
@ -2261,6 +2274,28 @@ class QwenImage(BaseModel):
|
|||
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||
return out
|
||||
|
||||
class JoyImage(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.joyimage.model.JoyImageTransformer3DModel)
|
||||
self.memory_usage_factor_conds = ("ref_latents",)
|
||||
|
||||
def extra_conds(self, **kwargs):
|
||||
out = super().extra_conds(**kwargs)
|
||||
cross_attn = kwargs.get("cross_attn", None)
|
||||
if cross_attn is not None:
|
||||
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
|
||||
ref_latents = kwargs.get("reference_latents", None)
|
||||
if ref_latents is not None:
|
||||
out['ref_latents'] = comfy.conds.CONDList([self.process_latent_in(lat) for lat in ref_latents])
|
||||
return out
|
||||
|
||||
def extra_conds_shapes(self, **kwargs):
|
||||
out = {}
|
||||
ref_latents = kwargs.get("reference_latents", None)
|
||||
if ref_latents is not None:
|
||||
out['ref_latents'] = list([1, 16, sum(map(lambda a: math.prod(a.size()), ref_latents)) // 16])
|
||||
return out
|
||||
|
||||
class Ideogram4(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.ideogram4.model.Ideogram4Transformer2DModel)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
@ -989,6 +1058,25 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
|||
dit_config["image_model"] = "SAM31"
|
||||
return dit_config
|
||||
|
||||
if (
|
||||
'{}double_blocks.0.attn.img_attn_qkv.weight'.format(key_prefix) in state_dict_keys
|
||||
and '{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix) in state_dict_keys
|
||||
and '{}condition_embedder.time_embedder.linear_1.weight'.format(key_prefix) in state_dict_keys
|
||||
and '{}img_in.weight'.format(key_prefix) in state_dict_keys
|
||||
and len(state_dict['{}img_in.weight'.format(key_prefix)].shape) == 5
|
||||
):
|
||||
img_in = state_dict['{}img_in.weight'.format(key_prefix)]
|
||||
head_dim = state_dict['{}double_blocks.0.attn.img_attn_q_norm.weight'.format(key_prefix)].shape[0]
|
||||
return {
|
||||
"image_model": "joyimage",
|
||||
"in_channels": img_in.shape[1],
|
||||
"hidden_size": img_in.shape[0],
|
||||
"patch_size": list(img_in.shape[2:]),
|
||||
"num_layers": count_blocks(state_dict_keys, '{}double_blocks.'.format(key_prefix) + '{}.'),
|
||||
"num_attention_heads": img_in.shape[0] // head_dim,
|
||||
"text_dim": 4096,
|
||||
}
|
||||
|
||||
if '{}input_blocks.0.0.weight'.format(key_prefix) not in state_dict_keys:
|
||||
return None
|
||||
|
||||
|
|
@ -1119,9 +1207,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 +1220,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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
40
comfy/ops.py
40
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):
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
108
comfy/sd.py
108
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
|
||||
|
|
@ -75,6 +76,7 @@ import comfy.text_encoders.gemma4
|
|||
import comfy.text_encoders.cogvideo
|
||||
import comfy.text_encoders.sa3
|
||||
import comfy.text_encoders.gpt_oss
|
||||
import comfy.text_encoders.joyimage
|
||||
|
||||
import comfy.model_patcher
|
||||
import comfy.lora
|
||||
|
|
@ -473,7 +475,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 +503,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 +551,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 +1033,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 +1073,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 +1139,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 +1170,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 +1233,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 +1251,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 +1275,27 @@ class VAE:
|
|||
elif dims == 2:
|
||||
samples = self.encode_tiled_(pixel_samples, **args)
|
||||
elif dims == 3:
|
||||
if tile_t is not None:
|
||||
tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
|
||||
if self.handles_tiling:
|
||||
samples = self._encode_tiled_owned(pixel_samples, **self._owned_tiled_args(tile_x, tile_y, overlap, tile_t, overlap_t))
|
||||
else:
|
||||
tile_t_latent = 9999
|
||||
args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
|
||||
if tile_t is not None:
|
||||
tile_t_latent = max(2, self.downscale_ratio[0](tile_t))
|
||||
else:
|
||||
tile_t_latent = 9999
|
||||
args["tile_t"] = self.upscale_ratio[0](tile_t_latent)
|
||||
|
||||
if overlap_t is None:
|
||||
args["overlap"] = (1, overlap, overlap)
|
||||
else:
|
||||
args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), overlap, overlap)
|
||||
maximum = pixel_samples.shape[2]
|
||||
maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
|
||||
spatial_overlap = overlap if overlap is not None else 64
|
||||
if overlap_t is None:
|
||||
args["overlap"] = (1, spatial_overlap, spatial_overlap)
|
||||
else:
|
||||
args["overlap"] = (self.upscale_ratio[0](max(1, min(tile_t_latent // 2, self.downscale_ratio[0](overlap_t)))), spatial_overlap, spatial_overlap)
|
||||
maximum = pixel_samples.shape[2]
|
||||
maximum = self.upscale_ratio[0](self.downscale_ratio[0](maximum))
|
||||
|
||||
samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
|
||||
samples = self.encode_tiled_3d(pixel_samples[:,:,:maximum], **args)
|
||||
|
||||
if self.format_encoded is not None:
|
||||
samples = self.format_encoded(samples)
|
||||
return samples
|
||||
|
||||
def get_sd(self):
|
||||
|
|
@ -1313,6 +1378,7 @@ class CLIPType(Enum):
|
|||
IDEOGRAM4 = 30
|
||||
BOOGU = 31
|
||||
KREA2 = 32
|
||||
JOYIMAGE = 33
|
||||
|
||||
|
||||
|
||||
|
|
@ -1642,6 +1708,10 @@ def load_text_encoder_state_dicts(state_dicts=[], embedding_directory=None, clip
|
|||
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
|
||||
clip_target.clip = comfy.text_encoders.krea2.te(**llama_detect(clip_data))
|
||||
clip_target.tokenizer = comfy.text_encoders.krea2.Krea2Tokenizer
|
||||
elif clip_type == CLIPType.JOYIMAGE and te_model == TEModel.QWEN3VL_8B: # JoyImageEdit: full Qwen3-VL-8B, edit-conditioning template + drop_idx.
|
||||
clip_data[0] = comfy.utils.state_dict_prefix_replace(clip_data[0], {"model.language_model.": "model.", "model.visual.": "visual.", "lm_head.": "model.lm_head."})
|
||||
clip_target.clip = comfy.text_encoders.joyimage.te(**llama_detect(clip_data))
|
||||
clip_target.tokenizer = comfy.text_encoders.joyimage.JoyImageTokenizer
|
||||
elif clip_type in (CLIPType.FLUX, CLIPType.FLUX2): # Flux2 Klein reuses the Qwen3-VL LM (3-layer tap -> 12288); visual unused.
|
||||
klein_model_type = "qwen3_8b" if te_model == TEModel.QWEN3VL_8B else "qwen3_4b"
|
||||
clip_target.clip = comfy.text_encoders.flux.klein_te(**llama_detect(clip_data), model_type=klein_model_type)
|
||||
|
|
@ -1898,7 +1968,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 +2109,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
|
||||
|
|
|
|||
|
|
@ -27,6 +27,7 @@ import comfy.text_encoders.z_image
|
|||
import comfy.text_encoders.ideogram4
|
||||
import comfy.text_encoders.boogu
|
||||
import comfy.text_encoders.krea2
|
||||
import comfy.text_encoders.joyimage
|
||||
import comfy.text_encoders.anima
|
||||
import comfy.text_encoders.ace15
|
||||
import comfy.text_encoders.longcat_image
|
||||
|
|
@ -1685,6 +1686,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",
|
||||
|
|
@ -1877,6 +1912,38 @@ class QwenImage(supported_models_base.BASE):
|
|||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen25_7b.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.qwen_image.QwenImageTokenizer, comfy.text_encoders.qwen_image.te(**hunyuan_detect))
|
||||
|
||||
class JoyImage(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "joyimage",
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
"multiplier": 1000,
|
||||
"shift": 1.5,
|
||||
}
|
||||
|
||||
memory_usage_factor = 1.8
|
||||
|
||||
unet_extra_config = {
|
||||
"theta": 10000,
|
||||
"rope_dim_list": [16, 56, 56],
|
||||
}
|
||||
|
||||
latent_format = latent_formats.Wan21
|
||||
|
||||
supported_inference_dtypes = [torch.bfloat16, torch.float32]
|
||||
|
||||
vae_key_prefix = ["vae."]
|
||||
text_encoder_key_prefix = ["text_encoders."]
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
return model_base.JoyImage(self, device=device)
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
pref = self.text_encoder_key_prefix[0]
|
||||
qwen3vl_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.joyimage.JoyImageTokenizer, comfy.text_encoders.joyimage.te(**qwen3vl_detect))
|
||||
|
||||
class HunyuanImage21(HunyuanVideo):
|
||||
unet_config = {
|
||||
"image_model": "hunyuan_video",
|
||||
|
|
@ -2348,12 +2415,14 @@ models = [
|
|||
HiDream,
|
||||
HiDreamO1,
|
||||
Chroma,
|
||||
SeedVR2,
|
||||
ChromaRadiance,
|
||||
ACEStep,
|
||||
ACEStep15,
|
||||
Omnigen2,
|
||||
Boogu,
|
||||
QwenImage,
|
||||
JoyImage,
|
||||
Ideogram4,
|
||||
Krea2,
|
||||
Flux2,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -0,0 +1,97 @@
|
|||
import torch
|
||||
|
||||
from comfy import sd1_clip
|
||||
import comfy.text_encoders.qwen_vl
|
||||
from comfy.text_encoders.qwen3vl import Qwen3VL, Qwen3VLTokenizer
|
||||
|
||||
JOYIMAGE_VISION_BLOCK = "<|vision_start|><|image_pad|><|vision_end|>"
|
||||
JOYIMAGE_TEMPLATE_TEXT = (
|
||||
"<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
|
||||
"quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
|
||||
"<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
JOYIMAGE_TEMPLATE_IMAGE = (
|
||||
"<|im_start|>system\n \\nDescribe the image by detailing the color, shape, size, texture, "
|
||||
"quantity, text, spatial relationships of the objects and background:<|im_end|>\n"
|
||||
f"<|im_start|>user\n{JOYIMAGE_VISION_BLOCK}{{}}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
# The DiT was trained without the leading system-prompt tokens.
|
||||
JOYIMAGE_DROP_IDX = 34
|
||||
PAD_TOKEN = 151643
|
||||
|
||||
|
||||
class Qwen3VL8B_JoyImage(Qwen3VL):
|
||||
model_type = "qwen3vl_8b"
|
||||
|
||||
def preprocess_embed(self, embed, device):
|
||||
if embed["type"] == "image":
|
||||
image, grid = comfy.text_encoders.qwen_vl.process_qwen2vl_images(
|
||||
embed["data"], min_pixels=65536, max_pixels=16777216, patch_size=16,
|
||||
image_mean=[0.5, 0.5, 0.5], image_std=[0.5, 0.5, 0.5],
|
||||
interpolation="bicubic",
|
||||
)
|
||||
merged, deepstack = self.visual(image.to(device, dtype=torch.float32), grid)
|
||||
return merged, {"grid": grid, "deepstack": deepstack}
|
||||
return None, None
|
||||
|
||||
|
||||
class JoyImageTokenizer(Qwen3VLTokenizer):
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
super().__init__(
|
||||
embedding_directory=embedding_directory, tokenizer_data=tokenizer_data,
|
||||
model_type="qwen3vl_8b",
|
||||
)
|
||||
self.llama_template = JOYIMAGE_TEMPLATE_TEXT
|
||||
self.llama_template_images = JOYIMAGE_TEMPLATE_IMAGE
|
||||
|
||||
def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=None, **kwargs):
|
||||
kwargs.pop("thinking", None)
|
||||
return super().tokenize_with_weights(
|
||||
text, return_word_ids=return_word_ids, llama_template=llama_template,
|
||||
images=images or [], thinking=True, **kwargs,
|
||||
)
|
||||
|
||||
|
||||
class _JoyImageClipModel(sd1_clip.SDClipModel):
|
||||
def __init__(self, device="cpu", layer="hidden", layer_idx=-1, dtype=None,
|
||||
attention_mask=True, model_options={}):
|
||||
super().__init__(
|
||||
device=device, layer=layer, layer_idx=layer_idx, textmodel_json_config={},
|
||||
# JoyImage conditions on the pre-final-norm output of the last decoder layer.
|
||||
dtype=dtype, special_tokens={"pad": PAD_TOKEN}, layer_norm_hidden_state=False,
|
||||
model_class=Qwen3VL8B_JoyImage, enable_attention_masks=attention_mask,
|
||||
return_attention_masks=attention_mask, model_options=model_options,
|
||||
)
|
||||
|
||||
|
||||
class JoyImageTEModel(sd1_clip.SD1ClipModel):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||
super().__init__(
|
||||
device=device, dtype=dtype, name="qwen3vl_8b",
|
||||
clip_model=_JoyImageClipModel, model_options=model_options,
|
||||
)
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs):
|
||||
out, pooled, extra = super().encode_token_weights(token_weight_pairs)
|
||||
if out.shape[1] <= JOYIMAGE_DROP_IDX:
|
||||
raise ValueError(
|
||||
f"JoyImageTEModel: encoded sequence length {out.shape[1]} is shorter "
|
||||
f"than drop_idx={JOYIMAGE_DROP_IDX}; the prompt did not include the "
|
||||
f"template prefix."
|
||||
)
|
||||
out = out[:, JOYIMAGE_DROP_IDX:]
|
||||
if "attention_mask" in extra:
|
||||
extra["attention_mask"] = extra["attention_mask"][:, JOYIMAGE_DROP_IDX:]
|
||||
return out, pooled, extra
|
||||
|
||||
|
||||
def te(dtype_llama=None, llama_quantization_metadata=None):
|
||||
class JoyImageTEModel_(JoyImageTEModel):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||
if llama_quantization_metadata is not None:
|
||||
model_options = model_options.copy()
|
||||
model_options["quantization_metadata"] = llama_quantization_metadata
|
||||
if dtype_llama is not None:
|
||||
dtype = dtype_llama
|
||||
super().__init__(device=device, dtype=dtype, model_options=model_options)
|
||||
return JoyImageTEModel_
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ def process_qwen2vl_images(
|
|||
merge_size: int = 2,
|
||||
image_mean: list = None,
|
||||
image_std: list = None,
|
||||
interpolation: str = "bilinear",
|
||||
):
|
||||
if image_mean is None:
|
||||
image_mean = [0.48145466, 0.4578275, 0.40821073]
|
||||
|
|
@ -47,10 +48,9 @@ def process_qwen2vl_images(
|
|||
img_resized = F.interpolate(
|
||||
img.unsqueeze(0),
|
||||
size=(h_bar, w_bar),
|
||||
mode='bilinear',
|
||||
mode=interpolation,
|
||||
align_corners=False
|
||||
).squeeze(0)
|
||||
|
||||
normalized = img_resized.clone()
|
||||
for c in range(3):
|
||||
normalized[c] = (img_resized[c] - image_mean[c]) / image_std[c]
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ 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,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
from av.container import InputContainer
|
||||
from av.subtitles.stream import SubtitleStream
|
||||
from av.video.reformatter import ColorRange
|
||||
from fractions import Fraction
|
||||
from typing import Optional
|
||||
from .._input import AudioInput, VideoInput
|
||||
|
|
@ -9,6 +10,7 @@ import itertools
|
|||
import json
|
||||
import numpy as np
|
||||
import math
|
||||
import os
|
||||
import torch
|
||||
from .._util import VideoContainer, VideoCodec, VideoComponents
|
||||
import logging
|
||||
|
|
@ -58,6 +60,57 @@ def video_stream_bit_depth(stream) -> int:
|
|||
return max(component.bits for component in stream.format.components)
|
||||
|
||||
|
||||
def last_decodable_audio_stream(container: InputContainer):
|
||||
"""Streams FFmpeg has no decoder for have no codec context, and decoding their
|
||||
packets crashes the process (e.g. APAC spatial-audio track in iPhone)."""
|
||||
stream = next(
|
||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||
None,
|
||||
)
|
||||
if stream is None and len(container.streams.audio):
|
||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||
return stream
|
||||
|
||||
|
||||
def probe_audio_params(container: InputContainer, audio_stream, max_packets: int = 200):
|
||||
"""Containers probed only up to a window (mpegts) leave audio codec parameters unset when
|
||||
audio starts beyond it; learn them by decoding ahead. The caller must seek back afterwards.
|
||||
Returns (sample_rate, channels), zeros when the stream never yields a decodable frame."""
|
||||
for i, packet in enumerate(container.demux(audio_stream)):
|
||||
try:
|
||||
frames = packet.decode()
|
||||
except av.error.FFmpegError:
|
||||
frames = ()
|
||||
if frames:
|
||||
return frames[0].sample_rate, frames[0].layout.nb_channels
|
||||
if i >= max_packets:
|
||||
break
|
||||
return 0, 0
|
||||
|
||||
|
||||
def write_output_metadata(container: InputContainer, output, metadata: dict | None):
|
||||
"""Copy the source container's metadata, then overlay the caller's tags."""
|
||||
for key, value in container.metadata.items():
|
||||
if metadata is None or key not in metadata:
|
||||
output.metadata[key] = value
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
output.metadata[key] = value if isinstance(value, str) else json.dumps(value)
|
||||
|
||||
|
||||
def mp4_output_open_kwargs(path: str | io.BytesIO, format: VideoContainer, codec: VideoCodec) -> dict:
|
||||
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
||||
raise ValueError("Only MP4 format is supported for now")
|
||||
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
||||
raise ValueError("Only H264 codec is supported for now")
|
||||
open_kwargs = {"mode": "w", "options": {"movflags": "use_metadata_tags"}}
|
||||
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
||||
open_kwargs["format"] = format.value
|
||||
elif isinstance(path, io.BytesIO):
|
||||
open_kwargs["format"] = "mp4" # no file extension to infer the format from
|
||||
return open_kwargs
|
||||
|
||||
|
||||
class VideoFromFile(VideoInput):
|
||||
"""
|
||||
Class representing video input from a file.
|
||||
|
|
@ -192,13 +245,10 @@ class VideoFromFile(VideoInput):
|
|||
return estimated_frames
|
||||
|
||||
# 3. Last resort: decode frames and count them (streaming)
|
||||
if self.__start_time < 0:
|
||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
||||
else:
|
||||
start_time = self.__start_time
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
frame_count = 1
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
frame_iterator = (
|
||||
container.decode(video_stream)
|
||||
|
|
@ -253,17 +303,14 @@ class VideoFromFile(VideoInput):
|
|||
|
||||
def get_components_internal(self, container: InputContainer) -> VideoComponents:
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
if self.__start_time < 0:
|
||||
start_time = max(self._get_raw_duration() + self.__start_time, 0)
|
||||
else:
|
||||
start_time = self.__start_time
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
|
||||
# Get video frames
|
||||
frames = []
|
||||
audio_frames = []
|
||||
alphas = None
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + self.__duration) / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base)
|
||||
|
||||
if start_pts != 0:
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
|
|
@ -281,18 +328,11 @@ class VideoFromFile(VideoInput):
|
|||
video_done = False
|
||||
audio_done = True
|
||||
|
||||
# Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context,
|
||||
# and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone)
|
||||
audio_stream = next(
|
||||
(s for s in reversed(container.streams.audio) if s.codec_context is not None),
|
||||
None,
|
||||
)
|
||||
audio_stream = last_decodable_audio_stream(container)
|
||||
if audio_stream is not None:
|
||||
streams += [audio_stream]
|
||||
resampler = av.audio.resampler.AudioResampler(format='fltp')
|
||||
audio_done = False
|
||||
elif len(container.streams.audio):
|
||||
logging.warning("No decodable audio stream found in video; ignoring audio.")
|
||||
|
||||
for packet in container.demux(*streams):
|
||||
if video_done and audio_done:
|
||||
|
|
@ -305,7 +345,7 @@ class VideoFromFile(VideoInput):
|
|||
for frame in packet.decode():
|
||||
if frame.pts < start_pts:
|
||||
continue
|
||||
if self.__duration and frame.pts >= end_pts:
|
||||
if duration and frame.pts >= end_pts:
|
||||
video_done = True
|
||||
break
|
||||
|
||||
|
|
@ -372,7 +412,7 @@ class VideoFromFile(VideoInput):
|
|||
map(resampler.resample, packet.decode())
|
||||
)
|
||||
for frame in aframes:
|
||||
if self.__duration and frame.time > start_time + self.__duration:
|
||||
if duration and frame.time > start_time + duration:
|
||||
audio_done = True
|
||||
break
|
||||
|
||||
|
|
@ -394,8 +434,8 @@ class VideoFromFile(VideoInput):
|
|||
|
||||
if len(audio_frames) > 0:
|
||||
audio_data = np.concatenate(audio_frames, axis=1) # shape: (channels, total_samples)
|
||||
if self.__duration:
|
||||
audio_data = audio_data[..., :int(self.__duration * audio_stream.sample_rate)]
|
||||
if duration:
|
||||
audio_data = audio_data[..., :int(duration * audio_stream.sample_rate)]
|
||||
|
||||
audio_tensor = torch.from_numpy(audio_data).unsqueeze(0) # shape: (1, channels, total_samples)
|
||||
audio = AudioInput({
|
||||
|
|
@ -441,28 +481,14 @@ class VideoFromFile(VideoInput):
|
|||
if not reuse_streams:
|
||||
if bit_depth is None:
|
||||
bit_depth = source_bit_depth
|
||||
components = self.get_components_internal(container)
|
||||
video = VideoFromComponents(components)
|
||||
return video.save_to(
|
||||
path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth,
|
||||
)
|
||||
return self._save_transcoded(container, path, format=format, codec=codec, metadata=metadata, bit_depth=bit_depth)
|
||||
|
||||
streams = container.streams
|
||||
|
||||
open_kwargs = get_open_write_kwargs(path, container_format, format)
|
||||
with av.open(path, **open_kwargs) as output_container:
|
||||
# Copy over the original metadata
|
||||
for key, value in container.metadata.items():
|
||||
if metadata is None or key not in metadata:
|
||||
output_container.metadata[key] = value
|
||||
|
||||
# Add our new metadata
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
if isinstance(value, str):
|
||||
output_container.metadata[key] = value
|
||||
else:
|
||||
output_container.metadata[key] = json.dumps(value)
|
||||
# Add metadata before writing any streams
|
||||
write_output_metadata(container, output_container, metadata)
|
||||
|
||||
# Add streams to the new container. Streams with no codec context cannot be used as an output template.
|
||||
stream_map = {}
|
||||
|
|
@ -480,6 +506,282 @@ class VideoFromFile(VideoInput):
|
|||
packet.stream = stream_map[packet.stream]
|
||||
output_container.mux(packet)
|
||||
|
||||
def _save_transcoded(
|
||||
self,
|
||||
container: InputContainer,
|
||||
path: str | io.BytesIO,
|
||||
format: VideoContainer,
|
||||
codec: VideoCodec,
|
||||
metadata: dict | None,
|
||||
bit_depth: int,
|
||||
):
|
||||
"""Re-encode to H.264/AAC one frame at a time; peak memory does not scale with video length."""
|
||||
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||
video_stream = self._get_first_video_stream(container)
|
||||
start_time, duration = self.get_active_trim_window()
|
||||
start_pts = int(start_time / video_stream.time_base)
|
||||
end_pts = int((start_time + duration) / video_stream.time_base) if duration else None
|
||||
stream_end_pts = None
|
||||
if video_stream.duration is not None:
|
||||
stream_end_pts = (video_stream.start_time or 0) + video_stream.duration
|
||||
output_end_pts = end_pts
|
||||
if stream_end_pts is not None and (output_end_pts is None or stream_end_pts < output_end_pts):
|
||||
output_end_pts = stream_end_pts
|
||||
if start_pts != 0:
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
|
||||
audio_stream = last_decodable_audio_stream(container)
|
||||
pix_fmt = "yuv420p10le" if bit_depth >= 10 else "yuv420p"
|
||||
rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||
|
||||
resampler = None
|
||||
sample_rate = 0
|
||||
audio_time_base = None
|
||||
duration_cap = None
|
||||
if audio_stream is not None:
|
||||
sample_rate = audio_stream.codec_context.sample_rate
|
||||
channels = audio_stream.codec_context.channels
|
||||
if not sample_rate:
|
||||
sample_rate, channels = probe_audio_params(container, audio_stream)
|
||||
container.seek(start_pts, stream=video_stream)
|
||||
if sample_rate:
|
||||
audio_stream.codec_context.flush_buffers()
|
||||
else:
|
||||
logging.warning("Audio stream parameters could not be determined; ignoring audio.")
|
||||
audio_stream = None
|
||||
if audio_stream is not None:
|
||||
audio_time_base = Fraction(1, sample_rate)
|
||||
layout = {1: "mono", 2: "stereo", 6: "5.1"}.get(channels, "stereo")
|
||||
resampler = av.audio.resampler.AudioResampler(format="fltp", layout=layout, rate=sample_rate)
|
||||
if duration:
|
||||
duration_cap = math.ceil(duration * sample_rate)
|
||||
|
||||
streams = [video_stream] if audio_stream is None else [video_stream, audio_stream]
|
||||
pts_step = max(1, int(round((1 / rate) / video_stream.time_base)))
|
||||
video_done = False
|
||||
audio_done = audio_stream is None
|
||||
video_pts_offset = None
|
||||
last_video_pts = None
|
||||
last_video_end = None
|
||||
# rebased pts -> true display duration: the mp4 muxer pads the last sample with 1/rate otherwise
|
||||
video_frame_durations = {}
|
||||
source_size = None
|
||||
rotation_k = 0
|
||||
rotation_filter = None
|
||||
audio_started = False
|
||||
samples_written = 0
|
||||
pending_audio = []
|
||||
# The output opens lazily on the first kept frame: it decides the geometry (90/270 rotation swaps dims),
|
||||
# and never seeking back keeps webm/mkv leading audio intact.
|
||||
output = None
|
||||
out_video = None
|
||||
out_audio = None
|
||||
|
||||
def audio_frame_from_ndarray(nd_planar):
|
||||
frame = av.AudioFrame.from_ndarray(np.ascontiguousarray(nd_planar), format="fltp", layout=layout)
|
||||
frame.sample_rate = sample_rate
|
||||
return frame
|
||||
|
||||
def drain_audio(final=False):
|
||||
# Audio may cover the pts span of the video written so far, capped by the requested duration
|
||||
nonlocal samples_written, audio_done
|
||||
if last_video_end is None:
|
||||
cap = 0
|
||||
else:
|
||||
cap = math.ceil(last_video_end * video_stream.time_base * sample_rate)
|
||||
if duration_cap is not None:
|
||||
cap = min(cap, duration_cap)
|
||||
while pending_audio and not audio_done:
|
||||
frame = pending_audio[0]
|
||||
if samples_written + frame.samples <= cap:
|
||||
frame.pts = samples_written
|
||||
frame.time_base = audio_time_base
|
||||
output.mux(out_audio.encode(frame))
|
||||
samples_written += frame.samples
|
||||
pending_audio.pop(0)
|
||||
continue
|
||||
if final:
|
||||
keep = frame.to_ndarray()[..., :cap - samples_written]
|
||||
if keep.shape[-1] > 0:
|
||||
tail = audio_frame_from_ndarray(keep)
|
||||
tail.pts = samples_written
|
||||
tail.time_base = audio_time_base
|
||||
output.mux(out_audio.encode(tail))
|
||||
samples_written += keep.shape[-1]
|
||||
pending_audio.clear()
|
||||
break
|
||||
if duration_cap is not None and samples_written >= duration_cap:
|
||||
audio_done = True
|
||||
return cap
|
||||
|
||||
try:
|
||||
for packet in container.demux(*streams):
|
||||
if video_done and audio_done:
|
||||
break
|
||||
|
||||
if packet.stream == video_stream and not video_done:
|
||||
try:
|
||||
frames = packet.decode()
|
||||
except av.error.InvalidDataError:
|
||||
logging.info("pyav decode error")
|
||||
continue
|
||||
for frame in frames:
|
||||
if frame.pts is not None and frame.pts < start_pts:
|
||||
continue
|
||||
if end_pts is not None and frame.pts is not None and frame.pts >= end_pts:
|
||||
video_done = True
|
||||
if last_video_pts is not None:
|
||||
# the source continues past the window: hold the last kept frame to the window end
|
||||
end_offset = video_pts_offset if video_pts_offset is not None else start_pts
|
||||
last_video_end = max(last_video_end, end_pts - end_offset)
|
||||
break
|
||||
# the source's true display duration of this frame; average_rate is not a
|
||||
# frame duration (sparse/VFR sources), so it is only the fallback
|
||||
frame_duration = frame.duration if frame.duration else pts_step
|
||||
if end_pts is not None and frame.pts is not None:
|
||||
frame_duration = min(frame_duration, end_pts - frame.pts)
|
||||
if output is None:
|
||||
rotation_k = int(round(frame.rotation // 90)) % 4 if frame.rotation else 0
|
||||
if rotation_k % 2:
|
||||
out_width, out_height = frame.height, frame.width
|
||||
else:
|
||||
out_width, out_height = frame.width, frame.height
|
||||
if out_width % 2 or out_height % 2:
|
||||
raise ValueError(f"H.264 output requires even dimensions, got {out_width}x{out_height}")
|
||||
source_size = (frame.width, frame.height)
|
||||
output = av.open(path, **open_kwargs)
|
||||
# Add metadata before writing any streams
|
||||
write_output_metadata(container, output, metadata)
|
||||
out_video = output.add_stream("h264", rate=rate)
|
||||
# no B-frames: reordering makes mp4 sample durations follow decode order,
|
||||
# so irregular-VFR spans and trim windows land wrong
|
||||
out_video.codec_context.max_b_frames = 0
|
||||
out_video.width = out_width
|
||||
out_video.height = out_height
|
||||
out_video.pix_fmt = pix_fmt
|
||||
# source pts pass through (rebased to 0), so variable frame rate survives
|
||||
out_video.codec_context.time_base = video_stream.time_base
|
||||
if audio_stream is not None:
|
||||
out_audio = output.add_stream("aac", rate=sample_rate, layout=layout)
|
||||
if (frame.width, frame.height) != source_size:
|
||||
# encoding would silently rescale the new geometry into the old one
|
||||
raise ValueError(
|
||||
f"Video resolution changes mid-stream "
|
||||
f"({source_size[0]}x{source_size[1]} -> {frame.width}x{frame.height}); cannot transcode"
|
||||
)
|
||||
if rotation_k:
|
||||
if rotation_filter is None:
|
||||
g = av.filter.Graph()
|
||||
g_src = g.add_buffer(width=frame.width, height=frame.height,
|
||||
format=frame.format.name, time_base=video_stream.time_base)
|
||||
tail = g_src
|
||||
for filter_name, filter_args in {1: [("transpose", "cclock")],
|
||||
2: [("hflip", None), ("vflip", None)],
|
||||
3: [("transpose", "clock")]}[rotation_k]:
|
||||
step = g.add(filter_name, filter_args)
|
||||
tail.link_to(step)
|
||||
tail = step
|
||||
g_sink = g.add("buffersink")
|
||||
tail.link_to(g_sink)
|
||||
g.configure()
|
||||
rotation_filter = (g_src, g_sink)
|
||||
rotation_filter[0].push(frame)
|
||||
frame = rotation_filter[1].pull()
|
||||
if frame.color_range == ColorRange.JPEG:
|
||||
# compress full-range sources (yuvj/MJPEG) to limited range
|
||||
frame = frame.reformat(format=pix_fmt, src_color_range="JPEG", dst_color_range="MPEG")
|
||||
else:
|
||||
frame = frame.reformat(format=pix_fmt)
|
||||
frame_output_end = None
|
||||
if frame.pts is not None:
|
||||
if video_pts_offset is None:
|
||||
video_pts_offset = frame.pts
|
||||
frame.pts -= video_pts_offset
|
||||
if output_end_pts is not None:
|
||||
frame_output_end = output_end_pts - video_pts_offset
|
||||
if frame.pts + frame_duration > frame_output_end:
|
||||
clamped_pts = frame_output_end - frame_duration
|
||||
if clamped_pts >= 0 and (last_video_pts is None or clamped_pts > last_video_pts):
|
||||
frame.pts = min(frame.pts, clamped_pts)
|
||||
elif frame.pts < frame_output_end:
|
||||
frame_duration = frame_output_end - frame.pts
|
||||
else:
|
||||
continue
|
||||
if frame.pts is None or (last_video_pts is not None and frame.pts <= last_video_pts):
|
||||
# broken sources emit missing/backward timestamps mid-stream, which the
|
||||
# muxer rejects; nudge them forward by one nominal frame interval
|
||||
frame.pts = 0 if last_video_pts is None else last_video_pts + pts_step
|
||||
if frame_output_end is not None and frame.pts + frame_duration > frame_output_end:
|
||||
if frame.pts >= frame_output_end:
|
||||
continue
|
||||
frame_duration = frame_output_end - frame.pts
|
||||
last_video_pts = frame.pts
|
||||
last_video_end = frame.pts + frame_duration
|
||||
video_frame_durations[frame.pts] = frame_duration
|
||||
# the decoded pict_type would force x264's frame types (intra-only
|
||||
# sources like MJPEG/ProRes would come out all-keyframe)
|
||||
frame.pict_type = 0
|
||||
for out_packet in out_video.encode(frame):
|
||||
out_packet.duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||
output.mux(out_packet)
|
||||
drain_audio()
|
||||
|
||||
elif packet.stream == audio_stream and not audio_done:
|
||||
for resampled in itertools.chain.from_iterable(map(resampler.resample, packet.decode())):
|
||||
frame_start = None
|
||||
if resampled.pts is not None:
|
||||
# passthrough frames keep the source stream's time base
|
||||
tb = resampled.time_base if resampled.time_base else audio_time_base
|
||||
frame_start = float(resampled.pts * tb)
|
||||
if duration and not audio_started and frame_start >= start_time + duration:
|
||||
audio_done = True
|
||||
break
|
||||
if not audio_started:
|
||||
if frame_start is None:
|
||||
frame_start = 0.0
|
||||
to_skip = max(0, int((start_time - frame_start) * sample_rate))
|
||||
if to_skip >= resampled.samples:
|
||||
continue
|
||||
audio_started = True
|
||||
if duration and frame_start > start_time:
|
||||
duration_cap = min(duration_cap, math.ceil((start_time + duration - frame_start) * sample_rate))
|
||||
if to_skip:
|
||||
pending_audio.append(audio_frame_from_ndarray(resampled.to_ndarray()[..., to_skip:]))
|
||||
continue
|
||||
pending_audio.append(resampled)
|
||||
if video_done:
|
||||
# the video window is complete so the cap is final, but containers
|
||||
# that interleave audio behind video (fragmented mp4) still owe most
|
||||
# of it: stop only once the demuxed audio covers the cap
|
||||
cap = drain_audio()
|
||||
if pending_audio or samples_written >= cap:
|
||||
drain_audio(final=True)
|
||||
audio_done = True
|
||||
break
|
||||
|
||||
if output is None:
|
||||
raise ValueError(f"No decodable video frames found in file '{self.__file}'")
|
||||
if out_audio is not None and not audio_done:
|
||||
drain_audio(final=True)
|
||||
window_fill = last_video_end - last_video_pts if video_done and last_video_pts is not None else 0
|
||||
for out_packet in out_video.encode(None):
|
||||
duration = video_frame_durations.pop(out_packet.pts, 0)
|
||||
if out_packet.pts == last_video_pts:
|
||||
duration = max(duration, window_fill)
|
||||
out_packet.duration = duration
|
||||
output.mux(out_packet)
|
||||
if out_audio is not None:
|
||||
output.mux(out_audio.encode(None))
|
||||
except BaseException:
|
||||
if output is not None:
|
||||
output.close()
|
||||
if isinstance(path, (str, os.PathLike)) and os.path.exists(path):
|
||||
os.remove(path)
|
||||
raise
|
||||
else:
|
||||
if output is not None:
|
||||
output.close()
|
||||
|
||||
def _get_first_video_stream(self, container: InputContainer):
|
||||
if len(container.streams.video):
|
||||
return container.streams.video[0]
|
||||
|
|
@ -527,22 +829,12 @@ class VideoFromComponents(VideoInput):
|
|||
bit_depth: int | None = None,
|
||||
):
|
||||
"""Save the video to a file path or BytesIO buffer."""
|
||||
if format != VideoContainer.AUTO and format != VideoContainer.MP4:
|
||||
raise ValueError("Only MP4 format is supported for now")
|
||||
if codec != VideoCodec.AUTO and codec != VideoCodec.H264:
|
||||
raise ValueError("Only H264 codec is supported for now")
|
||||
open_kwargs = mp4_output_open_kwargs(path, format, codec)
|
||||
# None means "use the depth this video was created with" (CreateVideo's choice).
|
||||
if bit_depth is None:
|
||||
bit_depth = self.__bit_depth
|
||||
is_10bit = bit_depth >= 10
|
||||
extra_kwargs = {}
|
||||
if isinstance(format, VideoContainer) and format != VideoContainer.AUTO:
|
||||
extra_kwargs["format"] = format.value
|
||||
elif isinstance(path, io.BytesIO):
|
||||
# BytesIO has no file extension, so av.open can't infer the format.
|
||||
# Default to mp4 since that's the only supported format anyway.
|
||||
extra_kwargs["format"] = "mp4"
|
||||
with av.open(path, mode='w', options={'movflags': 'use_metadata_tags'}, **extra_kwargs) as output:
|
||||
with av.open(path, **open_kwargs) as output:
|
||||
# Add metadata before writing any streams
|
||||
if metadata is not None:
|
||||
for key, value in metadata.items():
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
from datetime import date
|
||||
from enum import Enum
|
||||
from typing import Any
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
|
@ -242,3 +242,60 @@ class GeminiGenerateContentResponse(BaseModel):
|
|||
promptFeedback: GeminiPromptFeedback | None = Field(None)
|
||||
usageMetadata: GeminiUsageMetadata | None = Field(None)
|
||||
modelVersion: str | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionTextPart(BaseModel):
|
||||
type: Literal["text"] = "text"
|
||||
text: str = Field(...)
|
||||
|
||||
|
||||
class GeminiInteractionMediaPart(BaseModel):
|
||||
type: str = Field(..., description="One of: image, video, audio, document.")
|
||||
data: str | None = Field(None, description="Base64-encoded media bytes.")
|
||||
uri: str | None = Field(None, description="URI of the media, as an alternative to inline data.")
|
||||
mime_type: str | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionGenerationConfig(BaseModel):
|
||||
temperature: float | None = Field(None, ge=0.0, le=2.0)
|
||||
top_p: float | None = Field(None, ge=0.0, le=1.0)
|
||||
|
||||
|
||||
class GeminiInteractionRequest(BaseModel):
|
||||
model: str = Field(...)
|
||||
input: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = Field(...)
|
||||
generation_config: GeminiInteractionGenerationConfig | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionModalityTokens(BaseModel):
|
||||
modality: str | None = Field(None, description="One of: text, image, audio, video, document.")
|
||||
tokens: int | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionUsage(BaseModel):
|
||||
input_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
|
||||
output_tokens_by_modality: list[GeminiInteractionModalityTokens] | None = Field(None)
|
||||
total_thought_tokens: int | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionContent(BaseModel):
|
||||
type: str | None = Field(None)
|
||||
text: str | None = Field(None)
|
||||
data: str | None = Field(None)
|
||||
uri: str | None = Field(None)
|
||||
mime_type: str | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteractionStep(BaseModel):
|
||||
type: str | None = Field(None)
|
||||
content: list[GeminiInteractionContent] | None = Field(None)
|
||||
|
||||
|
||||
class GeminiInteraction(BaseModel):
|
||||
id: str | None = Field(None)
|
||||
status: str | None = Field(
|
||||
None,
|
||||
description="One of: in_progress, requires_action, completed, failed, cancelled, incomplete.",
|
||||
)
|
||||
steps: list[GeminiInteractionStep] | None = Field(None)
|
||||
usage: GeminiInteractionUsage | None = Field(None)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,452 @@
|
|||
# (label, avatar_id, avatar_type, supported engines)
|
||||
HEYGEN_AVATAR_LOOKS: list[tuple[str, str, str, tuple[str, ...]]] = [
|
||||
(
|
||||
"Annie Lounge Standing Side",
|
||||
"Annie_Lounge_Standing_Side_public",
|
||||
"studio_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Yara Modern Lecture Hall",
|
||||
"fd6814ecc5e143cd899e615a80eaa2dc",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Brandon Business Sitting Front",
|
||||
"Brandon_Business_Sitting_Front_public",
|
||||
"studio_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Caroline Business Sitting Side",
|
||||
"Caroline_Business_Sitting_Side_public",
|
||||
"studio_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Ursula Lawyer Angle 4",
|
||||
"f7173d2bb8584c00bfec6905c5e9a492",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Sofia Corporate Presenter 01 Angle 3",
|
||||
"fe563971fd2d438e957372dac9e2be8c",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Seoyeon Health Nutrition Coach Angle 3",
|
||||
"fe3c5d5028d941398d064b8fc64a2dea",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Sanne Fitness Coach Angle 4",
|
||||
"d967f935a8bf4a0c8f0bccfd66c501d2",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
("Sander", "f5cd7b94056f495ca0610602d64a9aa3", "photo_avatar", ("avatar_v", "avatar_iv", "avatar_iii")),
|
||||
(
|
||||
"Rupert Personal Development Coach Angle 4",
|
||||
"f57b3e626adb4bc997b38f64884adce4",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Olivier Professor Angle 2",
|
||||
"f6659bbb094b459c87c967edbb9ee481",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Obi Health Nutrition Coach Angle 5",
|
||||
"f3dc2c38201d414382f506d2d8e8d029",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Matilda Modern Office Setting",
|
||||
"fda889ac354a440da8dbecc410981273",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Mateo Traditional Law Office",
|
||||
"ff172d6c499c4e47ba6fcc5de631e9fc",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Marlon Inviting Armchair Setting",
|
||||
"f5a57db099ab462daa3e7c604a05dacc",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Margaret Professor Angle 1",
|
||||
"fb472bc29ab04bcca576e3703978fecb",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Marek Therapy Coach Angle 3",
|
||||
"e197768703f1463a93dc25ada1f421fb",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Maeve Warm, Professional Setting",
|
||||
"faf66681d8cc48dc82c4283200b3e782",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
(
|
||||
"Lorenzo Professor Angle 5",
|
||||
"fc268dc244bb40d7a554663ce723dcf0",
|
||||
"photo_avatar",
|
||||
("avatar_v", "avatar_iv", "avatar_iii"),
|
||||
),
|
||||
("Luca", "Luca_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Bruce", "Bruce_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Nico", "Nico_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Lisa", "Lisa_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Sophie", "Sophie_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Aiko", "Aiko_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Rebecca (portrait)", "Rebecca_public", "studio_avatar", ("avatar_iii",)),
|
||||
("Daphne in Grey blazer (portrait)", "Daphne_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Bryce in Black t-shirt", "Bryce_public_5", "studio_avatar", ("avatar_iii",)),
|
||||
("Diora in White shirt", "Diora_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Freja in White blazer", "Freja_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Albert in Blue blazer", "Albert_public_2", "studio_avatar", ("avatar_iii",)),
|
||||
("Emery in Red blazer", "Emery_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Minho in Blue shirt", "Minho_public_6", "studio_avatar", ("avatar_iii",)),
|
||||
("Aditya in Brown blazer", "Aditya_public_4", "studio_avatar", ("avatar_iii",)),
|
||||
("Nadim in Blue blazer", "Nadim_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Iker in Black blazer", "Iker_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Nour in Black blazer", "Nour_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Saskia in Blue blazer", "Saskia_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Lucien in Blue blazer", "Lucien_public_1", "studio_avatar", ("avatar_iii",)),
|
||||
("Esmond in Blue suit", "Esmond_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Jinwoo in Blue suit", "Jinwoo_public_5", "studio_avatar", ("avatar_iii",)),
|
||||
("Annelore in Red sweater (portrait)", "Annelore_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Bastien in Blue shirt", "Bastien_public_4", "studio_avatar", ("avatar_iii",)),
|
||||
("Zosia in Khaki blazer", "Zosia_public_3", "studio_avatar", ("avatar_iii",)),
|
||||
("Tahlia in Dark blue suit", "Tahlia_public_4", "studio_avatar", ("avatar_iii",)),
|
||||
]
|
||||
HEYGEN_AVATAR_OPTIONS = [x[0] for x in HEYGEN_AVATAR_LOOKS]
|
||||
HEYGEN_AVATAR_MAP = {x[0]: (x[1], x[2], x[3]) for x in HEYGEN_AVATAR_LOOKS}
|
||||
|
||||
# (label, voice_id) — Starfish-compatible voices for the TTS endpoint
|
||||
HEYGEN_VOICE_TTS: list[tuple[str, str]] = [
|
||||
("Chill Brian (English, male)", "d2f4f24783d04e22ab49ee8fdc3715e0"),
|
||||
("Zain (English, female)", "0047732240584155b1588455313e78ec"),
|
||||
("Narrator Mateo - Excited 🤩 (Spanish, male)", "0077225a877e457db4572ccaf245910b"),
|
||||
("Aria (English, female)", "007e1378fc454a9f976db570ba6164a7"),
|
||||
("Caryns (English, female)", "0082e70326864107823605db0d77f5e0"),
|
||||
("Klara (English, female)", "01209fdcd1c24a109c86dc24ee0f71c0"),
|
||||
("Bold Kasia - Excited 🤩 (Polish, female)", "015482a78b9a46ebae74bd0beb17765b"),
|
||||
("Shaun (English, male)", "01c42cddcfdc4665a57b8d89cba8ffc1"),
|
||||
("Senthil (English, male)", "01d674cfd32b4728a3fddd21b7e7d543"),
|
||||
("Cody (English, male)", "01f98ed43e6140349f47dbd37a416827"),
|
||||
("Saffron (English, female)", "0258bbc2cd8648cfa357adfb833f6d7b"),
|
||||
("Blanka - Lifelike (English, female)", "02880d1c6fd94b7799d91135581ed810"),
|
||||
("Rami (English, male)", "02d5366a90af4c7a87157808ff352e33"),
|
||||
("Rhodes (English, female)", "02dce0a169b3460084b6c914d18fb2c8"),
|
||||
("Michelle - Voice 1 (English, female)", "02X8sHnuxFpsq1caYWN0"),
|
||||
("Autumn - UGC 3 (English, female)", "03dca9ebfca441dba55fb14afa0791b7"),
|
||||
("Reassuring Rupert (English, male)", "03fcf8ecb0a94b6b94e9007edb7c35f8"),
|
||||
("Rose - UGC -2 (English, female)", "0495e14c2bd74eb3aeeef03583e0bce5"),
|
||||
("Derya - Lifelike - Broadcaster 🎙️ (English, female)", "04d0ae1d0af2489ca7d3bb402a39a890"),
|
||||
("Dynamic Derek (English, male)", "0516c2d857eb425c94e90b068241914e"),
|
||||
("Lotte (English, female)", "052fcfb83d1a4c2f8d0368c226fea4b9"),
|
||||
("Thanos - Broadcaster 🎙️ (English, male)", "054af44a167344d0af2722fdfef08d17"),
|
||||
("Marcia (English, female)", "05f19352e8f74b0392a8f411eba40de1"),
|
||||
("Camden (English, male)", "06468055edd4458aa131a1dfd813c1e9"),
|
||||
("Rumi (English, female)", "06672207805f41a9ad0af6797f8aa14b"),
|
||||
("Pippa (English, female)", "06b68c4dbb544935b9af984e80efa4fb"),
|
||||
("William Prescott - Broadcaster 🎙️ (English, male)", "06c816b952f14fa9b3a6c42aa151f731"),
|
||||
("Sammy (English, female)", "06e6facd99654b9dbb9308f67bf3a31c"),
|
||||
("Breezy Bagus (Indonesian, male)", "06e81a5d7c8b41818d3f0b38f7cf15a1"),
|
||||
("Ben (English, male)", "07ca39b243184dbcb82e7e0f0e524b21"),
|
||||
("Smooth Dev (English, male)", "07d2ba65847541feb97abc9b60181555"),
|
||||
("Daran inside booth (English, male)", "080d9383c0314056aef392892e009806"),
|
||||
("Peppy Stella (English, female)", "084760b4922a44599575c770070ec2d7"),
|
||||
("Silas (English, male)", "08f561403ec846dbbd8c691cc448f45a"),
|
||||
("Aditya (English, male)", "09c3d65e44e247dd8b78a97a903feb58"),
|
||||
("Christy (English, female)", "09d88c036bf449fa905900c08b235a37"),
|
||||
("Elio (English, male)", "0a0b38624ac64ec6afcd5842a977ca10"),
|
||||
("Luminous Laksh (Hindi, male)", "0adc547b76a5401c856274c379904eb7"),
|
||||
("Jeff (English, male)", "0add542e349f4ccaba6ecb3b7ced6034"),
|
||||
("Tahlia Brooks - Excited 🤩 (English, female)", "0b440d1ac2454d69a73302fc806522b1"),
|
||||
("Riya Mehta (Hindi, female)", "0b464b2f4e2249a4b5a05e60eaf41e7e"),
|
||||
("Ben Hart (English, male)", "0b47b5a637e944f9bfd49913999b344b"),
|
||||
("Skylar (English, female)", "0bbfbda5aa924a68a9d1da7b8496052a"),
|
||||
("Relaxed Reece (English, male)", "0c2151d538844c70a8b096de533f2828"),
|
||||
("Daniel (English, male)", "0c23804af39a4946ac6fda42bfff2738"),
|
||||
("Melani (English, female)", "0c54c6399ad64551a304e1a346677723"),
|
||||
("Clover (English, female)", "0ccb0bea067d4449ad367baeed7ea2e9"),
|
||||
("Pedro Lima - Serious 😐 (Portuguese, male)", "0d0e23e8170446e38b18a7380b2d30a8"),
|
||||
("Ana Carvalho (Portuguese, female)", "0d23c5b2f6004e909802a2e8bfcd52c2"),
|
||||
("Confident Connor - Excited 🤩 (English, male)", "0dd34c3eb79247238219eea35aeb58cd"),
|
||||
("Vibrant Victor (Spanish, male)", "1062976ea8bf42f4adc27c7e868b8fde"),
|
||||
("Young Olivier (French, male)", "1c5dc9a8f8cf4de0932f91d75f43a15d"),
|
||||
("Émile Noir (French, male)", "25a6a67280574d3da78e97b1935ebfc7"),
|
||||
("Steadfast Stefan (German, male)", "0eb85e6e8710473b82f7e88609ba3053"),
|
||||
("Deep Dieter (German, male)", "118949676b0a46629d1ad52981c3ef84"),
|
||||
("Serene Marco (Italian, male)", "72e922488a614041b5ab5f6ee07e3deb"),
|
||||
("Murmuring Matteo (Italian, male)", "755902b751654f30a6ef49e8bbcacfec"),
|
||||
("Gail in car (Multilingual, female)", "0214ac51f93e420f8711d568dcfbc50e"),
|
||||
("Daran outside walking (Multilingual, male)", "0ac81e725f4948dfa9638ceca216bcfa"),
|
||||
("BOB - Voice 1 (Chinese, unknown)", "dMkR1XwIkarpNqWUJLnX"),
|
||||
("Hakeem Hassan (Arabic, male)", "61a4359785664d01a59664ceb87ce6d4"),
|
||||
("Rami Idris (Arabic, male)", "a0bd2e5d41a74643be47ac75ca9171a2"),
|
||||
("Bold Kasia - Friendly 😊 (Polish, female)", "331624aec8b24a6c9287b8e16bdf54e8"),
|
||||
("Tranquil Tulin (Turkish, female)", "61646c861eb64e2d9036d8db51385356"),
|
||||
("Dynamic Derya (Turkish, female)", "664b73058b784aa89ddb2924c141d441"),
|
||||
("Quiet Dewa (Indonesian, male)", "1fa1193cf1d74f27ba58531c07ef9862"),
|
||||
("Cuong (Vietnamese, male)", "8af68d7ea38f4e7ca05cf46c3f7a590b"),
|
||||
]
|
||||
HEYGEN_VOICE_TTS_OPTIONS = [x[0] for x in HEYGEN_VOICE_TTS]
|
||||
HEYGEN_VOICE_TTS_MAP = dict(HEYGEN_VOICE_TTS)
|
||||
|
||||
# (label, voice_id) — top-ranked voices for video narration (any engine)
|
||||
HEYGEN_VOICE_GENERAL: list[tuple[str, str]] = [
|
||||
("Cassidy (English, female)", "16a09e4706f74997ba4ed05ea11470f6"),
|
||||
("Hope (English, female)", "42d00d4aac5441279d8536cd6b52c53c"),
|
||||
("Archer (English, male)", "453c20e1525a429080e2ad9e4b26f2cd"),
|
||||
("Brittney (English, female)", "4754e1ec667544b0bd18cdf4bec7d6a7"),
|
||||
("Mark (English, male)", "5d8c378ba8c3434586081a52ac368738"),
|
||||
("Andrew (English, male)", "6be73833ef9a4eb0aeee399b8fe9d62b"),
|
||||
("Spuds Oxley (English, male)", "76940a9adcd0490a9ce2cfe9a64a2664"),
|
||||
("Patrick (English, male)", "7e157ec62c9c45f1adca12faae72c86f"),
|
||||
("David Castlemore (English, male)", "828b59f834fd4c7188da322b6d9b6c75"),
|
||||
("Michael C (English, male)", "8661cd40d6c44c709e2d0031c0186ada"),
|
||||
("Adam Stone (English, male)", "88bb9ee1c81b466eb2a08fdde86d3619"),
|
||||
("Alex (English, male)", "897d6a9b2c844f56aa077238768fe10a"),
|
||||
("Monika Sogam (English, female)", "97dd67ab8ce242b6a9e7689cb00c6414"),
|
||||
("Jessica Anne Bogart (English, female)", "b966c31caf124c2a99f19ff1479c964f"),
|
||||
("John Doe (English, male)", "c4a8ceb7a2954500bc047fb092bcff3f"),
|
||||
("Ivy (English, female)", "cef3bc4e0a84424cafcde6f2cf466c97"),
|
||||
("Chill Brian (English, male)", "d2f4f24783d04e22ab49ee8fdc3715e0"),
|
||||
("Allison (English, female)", "f8c69e517f424cafaecde32dde57096b"),
|
||||
("Mia Starset (Norwegian, female)", "000466f8ac6d47a49f5743d50b3778de"),
|
||||
("William Shanks (Spanish, male)", "001248bb63f847888d37b766ee8b3a47"),
|
||||
("Zain (English, female)", "0047732240584155b1588455313e78ec"),
|
||||
("Jora Slobod (Romanian, male)", "00631519159a402ab5d8f719e51532bb"),
|
||||
("Narrator Mateo - Excited 🤩 (Spanish, male)", "0077225a877e457db4572ccaf245910b"),
|
||||
("Aria (English, female)", "007e1378fc454a9f976db570ba6164a7"),
|
||||
("Caryns (English, female)", "0082e70326864107823605db0d77f5e0"),
|
||||
("Klara (English, female)", "01209fdcd1c24a109c86dc24ee0f71c0"),
|
||||
("Son Tran (Vietnamese, male)", "0132f85950a94d11ba180f885101bf84"),
|
||||
("Bold Kasia - Excited 🤩 (Polish, female)", "015482a78b9a46ebae74bd0beb17765b"),
|
||||
("Marc Aurèle (French, male)", "018a94cf15574491a0bab7f6799ac15b"),
|
||||
("Shaun (English, male)", "01c42cddcfdc4665a57b8d89cba8ffc1"),
|
||||
("Senthil (English, male)", "01d674cfd32b4728a3fddd21b7e7d543"),
|
||||
("Cody (English, male)", "01f98ed43e6140349f47dbd37a416827"),
|
||||
("Saffron (English, female)", "0258bbc2cd8648cfa357adfb833f6d7b"),
|
||||
("Blanka - Lifelike (English, female)", "02880d1c6fd94b7799d91135581ed810"),
|
||||
("Rami (English, male)", "02d5366a90af4c7a87157808ff352e33"),
|
||||
("Rhodes (English, female)", "02dce0a169b3460084b6c914d18fb2c8"),
|
||||
("Michelle - Voice 1 (English, female)", "02X8sHnuxFpsq1caYWN0"),
|
||||
("Tuba (, female)", "034ca0c32b6542028748d6d365d90d6a"),
|
||||
("Autumn - UGC 3 (English, female)", "03dca9ebfca441dba55fb14afa0791b7"),
|
||||
("Reassuring Rupert (English, male)", "03fcf8ecb0a94b6b94e9007edb7c35f8"),
|
||||
]
|
||||
HEYGEN_VOICE_GENERAL_OPTIONS = [x[0] for x in HEYGEN_VOICE_GENERAL]
|
||||
HEYGEN_VOICE_GENERAL_MAP = dict(HEYGEN_VOICE_GENERAL)
|
||||
|
||||
HEYGEN_TRANSLATE_LANGUAGES = [
|
||||
"English",
|
||||
"Spanish",
|
||||
"Spanish (Spain)",
|
||||
"Spanish (Mexico)",
|
||||
"French",
|
||||
"French (France)",
|
||||
"German",
|
||||
"German (Germany)",
|
||||
"Portuguese",
|
||||
"Portuguese (Brazil)",
|
||||
"Italian",
|
||||
"Italian (Italy)",
|
||||
"Japanese",
|
||||
"Japanese (Japan)",
|
||||
"Korean",
|
||||
"Chinese (Mandarin, Simplified)",
|
||||
"Arabic",
|
||||
"Hindi",
|
||||
"Hindi (India)",
|
||||
"Russian",
|
||||
"Russian (Russia)",
|
||||
"Dutch",
|
||||
"Polish",
|
||||
"Turkish",
|
||||
"Indonesian",
|
||||
"Vietnamese",
|
||||
"Ukrainian",
|
||||
"Afrikaans (South Africa)",
|
||||
"Albanian (Albania)",
|
||||
"Amharic (Ethiopia)",
|
||||
"Arabic (Algeria)",
|
||||
"Arabic (Bahrain)",
|
||||
"Arabic (Egypt)",
|
||||
"Arabic (Iraq)",
|
||||
"Arabic (Jordan)",
|
||||
"Arabic (Kuwait)",
|
||||
"Arabic (Lebanon)",
|
||||
"Arabic (Libya)",
|
||||
"Arabic (Morocco)",
|
||||
"Arabic (Oman)",
|
||||
"Arabic (Qatar)",
|
||||
"Arabic (Saudi Arabia)",
|
||||
"Arabic (Syria)",
|
||||
"Arabic (Tunisia)",
|
||||
"Arabic (United Arab Emirates)",
|
||||
"Arabic (World)",
|
||||
"Arabic (Yemen)",
|
||||
"Armenian (Armenia)",
|
||||
"Azerbaijani (Latin, Azerbaijan)",
|
||||
"Bangla (Bangladesh)",
|
||||
"Basque",
|
||||
"Belarusian (Belarus)",
|
||||
"Bengali (India)",
|
||||
"Bosnian (Bosnia and Herzegovina)",
|
||||
"Bulgarian",
|
||||
"Bulgarian (Bulgaria)",
|
||||
"Burmese (Myanmar)",
|
||||
"Catalan",
|
||||
"Chinese (Cantonese, Traditional)",
|
||||
"Chinese (Jilu Mandarin, Simplified)",
|
||||
"Chinese (Northeastern Mandarin, Simplified)",
|
||||
"Chinese (Southwestern Mandarin, Simplified)",
|
||||
"Chinese (Taiwanese Mandarin, Traditional)",
|
||||
"Chinese (Wu, Simplified)",
|
||||
"Chinese (Zhongyuan Mandarin Henan, Simplified)",
|
||||
"Chinese (Zhongyuan Mandarin Shaanxi, Simplified)",
|
||||
"Croatian",
|
||||
"Croatian (Croatia)",
|
||||
"Czech",
|
||||
"Czech (Czechia)",
|
||||
"Danish",
|
||||
"Danish (Denmark)",
|
||||
"Dutch (Belgium)",
|
||||
"Dutch (Netherlands)",
|
||||
"English (Australia)",
|
||||
"English (Canada)",
|
||||
"English (Hong Kong SAR)",
|
||||
"English (India)",
|
||||
"English (Ireland)",
|
||||
"English (Kenya)",
|
||||
"English (New Zealand)",
|
||||
"English (Nigeria)",
|
||||
"English (Philippines)",
|
||||
"English (Singapore)",
|
||||
"English (South Africa)",
|
||||
"English (Tanzania)",
|
||||
"English (UK)",
|
||||
"English (United States)",
|
||||
"Estonian (Estonia)",
|
||||
"Filipino",
|
||||
"Filipino (Cebuano)",
|
||||
"Filipino (Philippines)",
|
||||
"Finnish",
|
||||
"Finnish (Finland)",
|
||||
"French (Belgium)",
|
||||
"French (Canada)",
|
||||
"French (Switzerland)",
|
||||
"Galician",
|
||||
"Georgian (Georgia)",
|
||||
"German (Austria)",
|
||||
"German (Switzerland)",
|
||||
"Greek",
|
||||
"Greek (Greece)",
|
||||
"Gujarati (India)",
|
||||
"Haitian Creole (Haiti)",
|
||||
"Hebrew (Israel)",
|
||||
"Hungarian (Hungary)",
|
||||
"Icelandic (Iceland)",
|
||||
"Indonesian (Indonesia)",
|
||||
"Irish (Ireland)",
|
||||
"Javanese (Latin, Indonesia)",
|
||||
"Kannada (India)",
|
||||
"Kazakh (Kazakhstan)",
|
||||
"Khmer (Cambodia)",
|
||||
"Konkani (India)",
|
||||
"Korean (Korea)",
|
||||
"Lao (Laos)",
|
||||
"Latin (Vatican City)",
|
||||
"Latvian (Latvia)",
|
||||
"Lithuanian (Lithuania)",
|
||||
"Luxembourgish (Luxembourg)",
|
||||
"Macedonian (North Macedonia)",
|
||||
"Maithili (India)",
|
||||
"Malagasy (Madagascar)",
|
||||
"Malay",
|
||||
"Malay (Malaysia)",
|
||||
"Malayalam (India)",
|
||||
"Maltese (Malta)",
|
||||
"Mandarin",
|
||||
"Marathi (India)",
|
||||
"Mongolian (Mongolia)",
|
||||
"Nepali (Nepal)",
|
||||
"Norwegian Bokmål (Norway)",
|
||||
"Norwegian Nynorsk (Norway)",
|
||||
"Odia (India)",
|
||||
"Pashto (Afghanistan)",
|
||||
"Persian (Iran)",
|
||||
"Polish (Poland)",
|
||||
"Portuguese (Portugal)",
|
||||
"Punjabi (India)",
|
||||
"Romanian",
|
||||
"Romanian (Romania)",
|
||||
"Serbian (Latin, Serbia)",
|
||||
"Sindhi (India)",
|
||||
"Sinhala (Sri Lanka)",
|
||||
"Slovak",
|
||||
"Slovak (Slovakia)",
|
||||
"Slovenian (Slovenia)",
|
||||
"Somali (Somalia)",
|
||||
"Spanish (Argentina)",
|
||||
"Spanish (Bolivia)",
|
||||
"Spanish (Chile)",
|
||||
"Spanish (Colombia)",
|
||||
"Spanish (Costa Rica)",
|
||||
"Spanish (Cuba)",
|
||||
"Spanish (Dominican Republic)",
|
||||
"Spanish (Ecuador)",
|
||||
"Spanish (El Salvador)",
|
||||
"Spanish (Equatorial Guinea)",
|
||||
"Spanish (Guatemala)",
|
||||
"Spanish (Honduras)",
|
||||
"Spanish (Latin America)",
|
||||
"Spanish (Nicaragua)",
|
||||
"Spanish (Panama)",
|
||||
"Spanish (Paraguay)",
|
||||
"Spanish (Peru)",
|
||||
"Spanish (Puerto Rico)",
|
||||
"Spanish (United States)",
|
||||
"Spanish (Uruguay)",
|
||||
"Spanish (Venezuela)",
|
||||
"Sundanese (Indonesia)",
|
||||
"Swahili (Kenya)",
|
||||
"Swahili (Tanzania)",
|
||||
"Swedish",
|
||||
"Swedish (Sweden)",
|
||||
"Tamil",
|
||||
"Tamil (India)",
|
||||
"Tamil (Malaysia)",
|
||||
"Tamil (Singapore)",
|
||||
"Tamil (Sri Lanka)",
|
||||
"Telugu (India)",
|
||||
"Thai (Thailand)",
|
||||
"Turkish (Türkiye)",
|
||||
"Ukrainian (Ukraine)",
|
||||
"Urdu (India)",
|
||||
"Urdu (Pakistan)",
|
||||
"Uzbek (Latin, Uzbekistan)",
|
||||
"Vietnamese (Vietnam)",
|
||||
"Welsh (United Kingdom)",
|
||||
"Zulu (South Africa)",
|
||||
]
|
||||
|
|
@ -77,6 +77,7 @@ class To3DUVTaskRequest(BaseModel):
|
|||
|
||||
class To3DPartTaskRequest(BaseModel):
|
||||
File: TaskFile3DInput = Field(...)
|
||||
EnableStagedGeneration: bool | None = Field(None)
|
||||
|
||||
|
||||
class TextureEditImageInfo(BaseModel):
|
||||
|
|
|
|||
|
|
@ -128,7 +128,7 @@ class OpenAIResponse(ModelResponseProperties, ResponseProperties):
|
|||
parallel_tool_calls: bool | None = Field(True)
|
||||
status: str | None = Field(
|
||||
None,
|
||||
description="One of `completed`, `failed`, `in_progress`, or `incomplete`.",
|
||||
description="One of `completed`, `failed`, `in_progress`, `incomplete`, `queued`, or `cancelled`.",
|
||||
)
|
||||
usage: ResponseUsage | None = Field(None)
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,49 @@
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class SyncInputItem(BaseModel):
|
||||
type: str = Field(..., description="Input kind: 'video', 'image' or 'audio'.")
|
||||
url: str = Field(...)
|
||||
|
||||
|
||||
class SyncActiveSpeakerDetection(BaseModel):
|
||||
auto_detect: bool | None = Field(
|
||||
None, description="Detect the active speaker automatically. Video input only; rejected for images."
|
||||
)
|
||||
frame_number: int | None = Field(
|
||||
None, description="Frame used for manual speaker selection. Must be 0 for image inputs."
|
||||
)
|
||||
coordinates: list[int] | None = Field(
|
||||
None, description="Pixel [x, y] of the speaker's face in the frame selected by frame_number."
|
||||
)
|
||||
|
||||
|
||||
class SyncGenerationOptions(BaseModel):
|
||||
sync_mode: str | None = Field(
|
||||
None,
|
||||
description="How to resolve an audio/video duration mismatch: "
|
||||
"cut_off, bounce, loop, silence or remap. Ignored for image inputs.",
|
||||
)
|
||||
i2v_prompt: str | None = Field(
|
||||
None, description="Motion prompt for image-to-video generation. Image input only."
|
||||
)
|
||||
active_speaker_detection: SyncActiveSpeakerDetection | None = Field(None)
|
||||
|
||||
|
||||
class SyncGenerationRequest(BaseModel):
|
||||
model: str = Field(..., description="Generation model, e.g. 'sync-3'.")
|
||||
input: list[SyncInputItem] = Field(
|
||||
..., description="Exactly one visual input (video or image) plus one audio input."
|
||||
)
|
||||
options: SyncGenerationOptions | None = Field(None)
|
||||
|
||||
|
||||
class SyncGeneration(BaseModel):
|
||||
"""Subset of the Generation object returned by POST /v2/generate and GET /v2/generate/{id}."""
|
||||
|
||||
id: str = Field(...)
|
||||
status: str = Field(..., description="PENDING | PROCESSING | COMPLETED | FAILED | REJECTED")
|
||||
outputUrl: str | None = Field(None)
|
||||
outputDuration: float | None = Field(None)
|
||||
error: str | None = Field(None, description="Human-readable failure message.")
|
||||
errorCode: str | None = Field(None, description="Stable machine-readable code from the GET /v2/errors catalog.")
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -24,6 +24,11 @@ from comfy_api_nodes.apis.gemini import (
|
|||
GeminiImageGenerateContentRequest,
|
||||
GeminiImageGenerationConfig,
|
||||
GeminiInlineData,
|
||||
GeminiInteraction,
|
||||
GeminiInteractionGenerationConfig,
|
||||
GeminiInteractionMediaPart,
|
||||
GeminiInteractionRequest,
|
||||
GeminiInteractionTextPart,
|
||||
GeminiMimeType,
|
||||
GeminiPart,
|
||||
GeminiRole,
|
||||
|
|
@ -51,6 +56,7 @@ from comfy_api_nodes.util import (
|
|||
)
|
||||
|
||||
GEMINI_BASE_ENDPOINT = "/proxy/vertexai/gemini"
|
||||
GEMINI_INTERACTIONS_ENDPOINT = "/proxy/gemini-interactions"
|
||||
GEMINI_MAX_INPUT_FILE_SIZE = 20 * 1024 * 1024 # 20 MB
|
||||
GEMINI_URL_INPUT_BUDGET = 10
|
||||
GEMINI_MAX_INLINE_BYTES = 18 * 1024 * 1024
|
||||
|
|
@ -231,29 +237,10 @@ async def get_image_from_response(response: GeminiGenerateContentResponse, thoug
|
|||
return torch.cat(image_tensors, dim=0)
|
||||
|
||||
|
||||
async def get_video_from_response(
|
||||
response: GeminiGenerateContentResponse, cls: type[IO.ComfyNode] | None = None
|
||||
) -> InputImpl.VideoFromFile:
|
||||
parts = get_parts_by_type(response, "video/*")
|
||||
for part in parts:
|
||||
if part.inlineData and part.inlineData.data:
|
||||
return InputImpl.VideoFromFile(BytesIO(base64.b64decode(part.inlineData.data)))
|
||||
if part.fileData and part.fileData.fileUri:
|
||||
return await download_url_to_video_output(part.fileData.fileUri, cls=cls)
|
||||
model_message = get_text_from_response(response).strip()
|
||||
if model_message:
|
||||
raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
|
||||
raise ValueError(
|
||||
"Gemini did not generate a video. Try rephrasing your prompt, "
|
||||
"shortening the requested duration, or reducing the number of input images/videos."
|
||||
)
|
||||
|
||||
|
||||
def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | None:
|
||||
if not response.modelVersion:
|
||||
return None
|
||||
# Define prices (Cost per 1,000,000 tokens), see https://cloud.google.com/vertex-ai/generative-ai/pricing
|
||||
output_video_tokens_price = 0.0
|
||||
if response.modelVersion == "gemini-2.5-pro":
|
||||
input_tokens_price = 1.25
|
||||
output_text_tokens_price = 10.0
|
||||
|
|
@ -274,6 +261,10 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
input_tokens_price = 0.25
|
||||
output_text_tokens_price = 1.50
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion == "gemini-3.5-flash":
|
||||
input_tokens_price = 1.50
|
||||
output_text_tokens_price = 9.0
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion in ("gemini-3-pro-image-preview", "gemini-3-pro-image"):
|
||||
input_tokens_price = 2
|
||||
output_text_tokens_price = 12.0
|
||||
|
|
@ -286,11 +277,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
input_tokens_price = 0.25
|
||||
output_text_tokens_price = 1.50
|
||||
output_image_tokens_price = 30.0
|
||||
elif response.modelVersion == "gemini-omni-flash-preview":
|
||||
input_tokens_price = 2.145
|
||||
output_text_tokens_price = 12.87
|
||||
output_image_tokens_price = 0.0
|
||||
output_video_tokens_price = 25.025
|
||||
else:
|
||||
return None
|
||||
final_price = response.usageMetadata.promptTokenCount * input_tokens_price
|
||||
|
|
@ -298,8 +284,6 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
for i in response.usageMetadata.candidatesTokensDetails:
|
||||
if i.modality == Modality.IMAGE:
|
||||
final_price += output_image_tokens_price * i.tokenCount # for Nano Banana models
|
||||
elif i.modality == Modality.VIDEO:
|
||||
final_price += output_video_tokens_price * i.tokenCount # for Omni Flash
|
||||
else:
|
||||
final_price += output_text_tokens_price * i.tokenCount
|
||||
if response.usageMetadata.thoughtsTokenCount:
|
||||
|
|
@ -307,6 +291,58 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
|||
return final_price / 1_000_000.0
|
||||
|
||||
|
||||
def get_text_from_interaction(interaction: GeminiInteraction) -> str:
|
||||
"""Extract and concatenate all model output text from an Interactions API response."""
|
||||
texts = []
|
||||
for step in interaction.steps or []:
|
||||
if step.type != "model_output":
|
||||
continue
|
||||
for content in step.content or []:
|
||||
if content.type == "text" and content.text:
|
||||
texts.append(content.text)
|
||||
return "\n".join(texts)
|
||||
|
||||
|
||||
async def get_video_from_interaction(
|
||||
interaction: GeminiInteraction, cls: type[IO.ComfyNode] | None = None
|
||||
) -> InputImpl.VideoFromFile:
|
||||
for step in interaction.steps or []:
|
||||
if step.type != "model_output":
|
||||
continue
|
||||
for content in step.content or []:
|
||||
if content.type != "video":
|
||||
continue
|
||||
if content.data:
|
||||
return InputImpl.VideoFromFile(BytesIO(base64.b64decode(content.data)))
|
||||
if content.uri:
|
||||
return await download_url_to_video_output(content.uri, cls=cls)
|
||||
model_message = get_text_from_interaction(interaction).strip()
|
||||
if model_message:
|
||||
raise ValueError(f"Gemini did not generate a video. Model response: {model_message}")
|
||||
raise ValueError(
|
||||
"Gemini did not generate a video. Try rephrasing your prompt, "
|
||||
"shortening the requested duration, or reducing the number of input images/videos."
|
||||
)
|
||||
|
||||
|
||||
def calculate_interaction_tokens_price(interaction: GeminiInteraction) -> float | None:
|
||||
if interaction.usage is None:
|
||||
return None
|
||||
input_tokens_price = 1.5
|
||||
output_tokens_prices = {"text": 9.0, "video": 17.5}
|
||||
thoughts_tokens_price = 9.0
|
||||
final_price = 0.0
|
||||
for i in interaction.usage.input_tokens_by_modality or []:
|
||||
if i.tokens:
|
||||
final_price += input_tokens_price * i.tokens
|
||||
for i in interaction.usage.output_tokens_by_modality or []:
|
||||
if i.tokens and i.modality in output_tokens_prices:
|
||||
final_price += output_tokens_prices[i.modality] * i.tokens
|
||||
if interaction.usage.total_thought_tokens:
|
||||
final_price += thoughts_tokens_price * interaction.usage.total_thought_tokens
|
||||
return final_price / 1_000_000.0
|
||||
|
||||
|
||||
def create_video_parts(video_input: Input.Video) -> list[GeminiPart]:
|
||||
"""Convert a single video input to Gemini API compatible parts (inline MP4/H.264)."""
|
||||
base_64_string = video_to_base64_string(
|
||||
|
|
@ -441,6 +477,15 @@ async def build_gemini_media_parts(
|
|||
return parts
|
||||
|
||||
|
||||
def to_interaction_media_part(part: GeminiPart) -> GeminiInteractionMediaPart:
|
||||
"""Convert a fileData/inlineData GeminiPart into an Interactions API media part."""
|
||||
if part.fileData:
|
||||
mime = part.fileData.mimeType.value
|
||||
return GeminiInteractionMediaPart(type=mime.split("/")[0], uri=part.fileData.fileUri, mime_type=mime)
|
||||
mime = part.inlineData.mimeType.value
|
||||
return GeminiInteractionMediaPart(type=mime.split("/")[0], data=part.inlineData.data, mime_type=mime)
|
||||
|
||||
|
||||
class GeminiNode(IO.ComfyNode):
|
||||
"""
|
||||
Node to generate text responses from a Gemini model.
|
||||
|
|
@ -619,11 +664,12 @@ class GeminiNode(IO.ComfyNode):
|
|||
|
||||
GEMINI_V2_MODELS: dict[str, str] = {
|
||||
"Gemini 3.1 Pro": "gemini-3.1-pro-preview",
|
||||
"Gemini 3.5 Flash": "gemini-3.5-flash",
|
||||
"Gemini 3.1 Flash-Lite": "gemini-3.1-flash-lite-preview",
|
||||
}
|
||||
|
||||
|
||||
def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
|
||||
def _gemini_text_model_inputs(thinking_default: str, thinking_options: list[str] | None = None) -> list[Input]:
|
||||
"""Per-model inputs revealed by the model DynamicCombo (shared media + sampling controls)."""
|
||||
return [
|
||||
IO.Autogrow.Input(
|
||||
|
|
@ -661,7 +707,7 @@ def _gemini_text_model_inputs(thinking_default: str) -> list[Input]:
|
|||
),
|
||||
IO.Combo.Input(
|
||||
"thinking_level",
|
||||
options=["LOW", "HIGH"],
|
||||
options=thinking_options or ["LOW", "HIGH"],
|
||||
default=thinking_default,
|
||||
tooltip="How hard the model reasons internally before answering. "
|
||||
"HIGH improves quality on difficult tasks but costs more (thinking) tokens and is slower.",
|
||||
|
|
@ -719,6 +765,10 @@ class GeminiNodeV2(IO.ComfyNode):
|
|||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Gemini 3.5 Flash",
|
||||
_gemini_text_model_inputs("MEDIUM", ["MINIMAL", "LOW", "MEDIUM", "HIGH"]),
|
||||
),
|
||||
IO.DynamicCombo.Option("Gemini 3.1 Pro", _gemini_text_model_inputs("HIGH")),
|
||||
IO.DynamicCombo.Option("Gemini 3.1 Flash-Lite", _gemini_text_model_inputs("LOW")),
|
||||
],
|
||||
|
|
@ -759,7 +809,13 @@ class GeminiNodeV2(IO.ComfyNode):
|
|||
"type": "list_usd",
|
||||
"usd": [0.00025, 0.0015],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
} : {
|
||||
}
|
||||
: $contains($m, "3.5 flash") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.0015, 0.009],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
}
|
||||
: {
|
||||
"type": "list_usd",
|
||||
"usd": [0.002, 0.012],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
|
|
@ -1133,7 +1189,9 @@ class GeminiImage2(IO.ComfyNode):
|
|||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||
if model == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||
model = "gemini-3.1-flash-image-preview"
|
||||
model = "gemini-3.1-flash-image"
|
||||
elif model == "gemini-3-pro-image-preview":
|
||||
model = "gemini-3-pro-image"
|
||||
|
||||
parts: list[GeminiPart] = [GeminiPart(text=prompt)]
|
||||
if images is not None:
|
||||
|
|
@ -1507,7 +1565,7 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
|||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||
model_choice = model["model"]
|
||||
if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||
model_id = "gemini-3.1-flash-image-preview"
|
||||
model_id = "gemini-3.1-flash-image"
|
||||
elif model_choice == "Nano Banana 2 Lite":
|
||||
model_id = "gemini-3.1-flash-lite-image"
|
||||
else:
|
||||
|
|
@ -1659,7 +1717,7 @@ class GeminiVideoOmni(IO.ComfyNode):
|
|||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr='{"type":"usd","usd":0.146,"format":{"suffix":"/second","approximate":true}}'
|
||||
expr='{"type":"usd","usd":0.101,"format":{"suffix":"/second","approximate":true}}'
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -1678,27 +1736,34 @@ class GeminiVideoOmni(IO.ComfyNode):
|
|||
for video in videos:
|
||||
validate_video_duration(video, max_duration=10)
|
||||
|
||||
parts: list[GeminiPart] = []
|
||||
parts: list[GeminiInteractionTextPart | GeminiInteractionMediaPart] = []
|
||||
if images or videos:
|
||||
parts.extend(await build_gemini_media_parts(cls, images, [], videos))
|
||||
parts.append(GeminiPart(text=prompt))
|
||||
response = await sync_op(
|
||||
media_parts = await build_gemini_media_parts(cls, images, [], videos)
|
||||
parts.extend(to_interaction_media_part(p) for p in media_parts)
|
||||
parts.append(GeminiInteractionTextPart(text=prompt))
|
||||
interaction = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{GEMINI_BASE_ENDPOINT}/{model_id}", method="POST"),
|
||||
data=GeminiGenerateContentRequest(
|
||||
contents=[GeminiContent(role=GeminiRole.user, parts=parts)],
|
||||
generationConfig=GeminiGenerationConfig(
|
||||
responseModalities=["TEXT", "VIDEO"],
|
||||
ApiEndpoint(path=GEMINI_INTERACTIONS_ENDPOINT, method="POST"),
|
||||
data=GeminiInteractionRequest(
|
||||
model=model_id,
|
||||
input=parts,
|
||||
generation_config=GeminiInteractionGenerationConfig(
|
||||
temperature=model.get("temperature", 1.0),
|
||||
topP=model.get("top_p", 0.95),
|
||||
top_p=model.get("top_p", 0.95),
|
||||
),
|
||||
),
|
||||
response_model=GeminiGenerateContentResponse,
|
||||
price_extractor=calculate_tokens_price,
|
||||
response_model=GeminiInteraction,
|
||||
price_extractor=calculate_interaction_tokens_price,
|
||||
)
|
||||
if interaction.status != "completed":
|
||||
model_message = get_text_from_interaction(interaction).strip()
|
||||
raise ValueError(
|
||||
f"Gemini interaction did not complete (status: {interaction.status})."
|
||||
+ (f" Model response: {model_message}" if model_message else "")
|
||||
)
|
||||
return IO.NodeOutput(
|
||||
await get_video_from_response(response, cls=cls),
|
||||
get_text_from_response(response),
|
||||
await get_video_from_interaction(interaction, cls=cls),
|
||||
get_text_from_interaction(interaction),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,799 @@
|
|||
import uuid
|
||||
|
||||
import torch
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.heygen import (
|
||||
HEYGEN_AVATAR_MAP,
|
||||
HEYGEN_AVATAR_OPTIONS,
|
||||
HEYGEN_TRANSLATE_LANGUAGES,
|
||||
HEYGEN_VOICE_GENERAL_MAP,
|
||||
HEYGEN_VOICE_GENERAL_OPTIONS,
|
||||
HEYGEN_VOICE_TTS_MAP,
|
||||
HEYGEN_VOICE_TTS_OPTIONS,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
audio_bytes_to_audio_input,
|
||||
download_url_as_bytesio,
|
||||
download_url_to_image_tensor,
|
||||
download_url_to_video_output,
|
||||
downscale_image_tensor_by_max_side,
|
||||
get_number_of_images,
|
||||
poll_op_raw,
|
||||
sync_op_raw,
|
||||
upload_audio_to_comfyapi,
|
||||
upload_image_to_comfyapi,
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_string,
|
||||
)
|
||||
from server import PromptServer
|
||||
|
||||
_VIDEOS_PATH = "/proxy/heygen/v3/videos"
|
||||
_TRANSLATIONS_PATH = "/proxy/heygen/v3/video-translations"
|
||||
_SPEECH_PATH = "/proxy/heygen/v3/voices/speech"
|
||||
_AVATARS_PATH = "/proxy/heygen/v3/avatars"
|
||||
_LOOKS_PATH = "/proxy/heygen/v3/avatars/looks"
|
||||
|
||||
_DEFAULT_VOICE_OPTION = "(avatar's default voice)"
|
||||
|
||||
_AVATARS_BY_ENGINE = {
|
||||
e: [label for label, (_aid, _atype, engines) in HEYGEN_AVATAR_MAP.items() if e in engines]
|
||||
for e in ("avatar_iv", "avatar_iii", "avatar_v")
|
||||
}
|
||||
|
||||
|
||||
async def _apply_speech_source(cls: type[IO.ComfyNode], payload: dict, speech: dict, require_voice: bool) -> None:
|
||||
"""Fill script/audio speech fields of a /v3/videos payload from the DynamicCombo dict."""
|
||||
if speech["speech"] == "audio":
|
||||
payload["audio_url"] = await upload_audio_to_comfyapi(
|
||||
cls, speech["audio"], container_format="mp3", codec_name="libmp3lame", mime_type="audio/mpeg"
|
||||
)
|
||||
elif speech["speech"] == "script":
|
||||
validate_string(speech["text"], strip_whitespace=True, min_length=1, max_length=5000)
|
||||
payload["script"] = speech["text"]
|
||||
voice_id = speech.get("custom_voice_id", "").strip()
|
||||
if not voice_id and speech["voice"] != _DEFAULT_VOICE_OPTION:
|
||||
voice_id = HEYGEN_VOICE_GENERAL_MAP[speech["voice"]]
|
||||
if voice_id:
|
||||
payload["voice_id"] = voice_id
|
||||
elif require_voice:
|
||||
raise ValueError("A voice is required when driving the video with a text script.")
|
||||
speed = speech.get("voice_speed", 1.0)
|
||||
if speed != 1.0:
|
||||
payload["voice_settings"] = {"speed": round(speed, 2)}
|
||||
|
||||
|
||||
async def _create_and_poll_video(cls: type[IO.ComfyNode], payload: dict) -> dict:
|
||||
"""POST a /v3/videos payload, poll until terminal, and return the final video data."""
|
||||
created = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_VIDEOS_PATH, method="POST", headers={"Idempotency-Key": uuid.uuid4().hex}),
|
||||
data=payload,
|
||||
)
|
||||
video_id = (created.get("data") or {}).get("video_id")
|
||||
if not video_id:
|
||||
raise ValueError(f"HeyGen did not return a video_id: {created}")
|
||||
final = await poll_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_VIDEOS_PATH}/{video_id}"),
|
||||
status_extractor=lambda r: (r.get("data") or {}).get("status"),
|
||||
queued_statuses=["pending", "waiting"],
|
||||
poll_interval=5.0,
|
||||
)
|
||||
data = final["data"]
|
||||
if not data.get("video_url"):
|
||||
raise ValueError(f"HeyGen returned no video_url for video {video_id}.")
|
||||
return data
|
||||
|
||||
|
||||
async def _resolve_avatar(
|
||||
cls: type[IO.ComfyNode], avatar_label: str, custom_avatar_id: str, engine_choice: str
|
||||
) -> tuple[str, str | None]:
|
||||
"""Resolve (avatar_id, engine_type) from the combo/override + engine widgets."""
|
||||
custom_avatar_id = custom_avatar_id.strip()
|
||||
if custom_avatar_id:
|
||||
look = (
|
||||
await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_LOOKS_PATH}/{custom_avatar_id}"),
|
||||
final_label_on_success=None,
|
||||
)
|
||||
).get("data") or {}
|
||||
avatar_id = custom_avatar_id
|
||||
avatar_label = look.get("name") or custom_avatar_id
|
||||
supported = look.get("supported_api_engines") or []
|
||||
else:
|
||||
avatar_id, avatar_type, supported = HEYGEN_AVATAR_MAP[avatar_label]
|
||||
|
||||
if engine_choice == "auto":
|
||||
engine = next((e for e in ("avatar_iv", "avatar_iii", "avatar_v") if e in supported), None)
|
||||
else:
|
||||
engine = engine_choice
|
||||
if supported and engine not in supported:
|
||||
raise ValueError(
|
||||
f"Avatar '{avatar_label}' does not support the {engine} engine "
|
||||
f"(supported: {', '.join(supported)}). Set engine to 'auto' to pick "
|
||||
"a compatible engine automatically."
|
||||
)
|
||||
return avatar_id, engine
|
||||
|
||||
|
||||
class HeyGenTalkingPhotoNode(IO.ComfyNode):
|
||||
"""Animate a still image of a person into a lip-synced talking video."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenTalkingPhotoNode",
|
||||
display_name="HeyGen Talking Photo",
|
||||
category="partner/video/HeyGen",
|
||||
description="Animate any image of a person into a lip-synced talking video "
|
||||
"(HeyGen Avatar IV). Drive it with a text script or your own audio.",
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
tooltip="Image of a person to animate. Downscaled automatically if larger than 2K.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"speech",
|
||||
display_name="speech source",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"script",
|
||||
[
|
||||
IO.String.Input(
|
||||
"text",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Text for the avatar to speak (up to 5000 characters). "
|
||||
"The generated speech must be at least 1 second long.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
options=HEYGEN_VOICE_GENERAL_OPTIONS,
|
||||
tooltip="Voice for the script (HeyGen's most popular voices).",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_voice_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
|
||||
"Any voice from HeyGen's library (2000+) can be used.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"voice_speed",
|
||||
default=1.0,
|
||||
min=0.5,
|
||||
max=1.5,
|
||||
step=0.05,
|
||||
optional=True,
|
||||
tooltip="Speech speed multiplier.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"audio",
|
||||
[
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Audio for the avatar to lip-sync, up to 10 minutes.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Drive the avatar with a text script (HeyGen text-to-speech) or your own audio.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720p", "1080p"],
|
||||
default="1080p",
|
||||
optional=True,
|
||||
tooltip="Output video resolution.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=["auto", "16:9", "9:16", "1:1", "4:5", "5:4"],
|
||||
default="auto",
|
||||
optional=True,
|
||||
tooltip="Output aspect ratio. 'auto' follows the input image.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"expressiveness",
|
||||
options=["low", "medium", "high"],
|
||||
default="low",
|
||||
optional=True,
|
||||
tooltip="How expressive the animated face and gestures are.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.0715,"format":{"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
speech: dict,
|
||||
resolution: str = "1080p",
|
||||
aspect_ratio: str = "auto",
|
||||
expressiveness: str = "low",
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
image = downscale_image_tensor_by_max_side(image, max_side=2000)
|
||||
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||
payload = {
|
||||
"type": "image",
|
||||
"image": {"type": "url", "url": image_url},
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"expressiveness": expressiveness,
|
||||
"title": "ComfyUI Talking Photo",
|
||||
}
|
||||
await _apply_speech_source(cls, payload, speech, require_voice=True)
|
||||
video = await _create_and_poll_video(cls, payload)
|
||||
return IO.NodeOutput(await download_url_to_video_output(video["video_url"]))
|
||||
|
||||
|
||||
class HeyGenAvatarVideoNode(IO.ComfyNode):
|
||||
"""Generate a presenter video from a HeyGen avatar look."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenAvatarVideoNode",
|
||||
display_name="HeyGen Avatar Video",
|
||||
category="partner/video/HeyGen",
|
||||
description="Generate a talking-presenter video from a HeyGen avatar. "
|
||||
"Includes HeyGen's most popular public avatars; any look ID can be supplied as an override.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"engine",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"auto",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=HEYGEN_AVATAR_OPTIONS,
|
||||
tooltip="Avatar look to present the video (curated from HeyGen's "
|
||||
"public library). The best engine the look supports is chosen "
|
||||
"automatically.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"avatar_iv",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=_AVATARS_BY_ENGINE["avatar_iv"],
|
||||
tooltip="Avatar looks that support the Avatar IV engine.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"avatar_iii",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=_AVATARS_BY_ENGINE["avatar_iii"],
|
||||
tooltip="Avatar looks that support the Avatar III engine.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"avatar_v",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"avatar",
|
||||
options=_AVATARS_BY_ENGINE["avatar_v"],
|
||||
tooltip="Avatar looks that support the Avatar V engine.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Rendering engine; each choice lists only the avatars that support it. "
|
||||
"'auto' offers every avatar and picks its best engine (Avatar IV preferred). "
|
||||
"Avatar V is highest fidelity, Avatar III is the most affordable.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_avatar_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen avatar look ID. When set, overrides the avatar selected above. "
|
||||
"Any of HeyGen's 3000+ public looks (or your private avatars) can be used.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"speech",
|
||||
display_name="speech source",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"script",
|
||||
[
|
||||
IO.String.Input(
|
||||
"text",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Text for the avatar to speak (up to 5000 characters). "
|
||||
"The generated speech must be at least 1 second long.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
options=[_DEFAULT_VOICE_OPTION] + HEYGEN_VOICE_GENERAL_OPTIONS,
|
||||
tooltip="Voice for the script. The default option uses the voice HeyGen assigned to the avatar.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_voice_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
|
||||
"Any voice from HeyGen's library (2000+) can be used.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"voice_speed",
|
||||
default=1.0,
|
||||
min=0.5,
|
||||
max=1.5,
|
||||
step=0.05,
|
||||
optional=True,
|
||||
tooltip="Speech speed multiplier.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"audio",
|
||||
[
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Audio for the avatar to lip-sync, up to 10 minutes.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Drive the avatar with a text script (HeyGen text-to-speech) or your own audio.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720p", "1080p"],
|
||||
default="1080p",
|
||||
optional=True,
|
||||
tooltip="Output video resolution.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=["auto", "16:9", "9:16", "1:1", "4:5", "5:4"],
|
||||
default="auto",
|
||||
optional=True,
|
||||
tooltip="Output aspect ratio. 'auto' follows the avatar's source footage.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"background_color",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional solid background color as a hex code (e.g. '#00ff00'). "
|
||||
"Leave empty for the avatar's own background.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["engine"]),
|
||||
expr="""
|
||||
widgets.engine = "avatar_iii"
|
||||
? {"type":"range_usd","min_usd":0.023881,"max_usd":0.061919,"format":{"suffix":"/second"}}
|
||||
: widgets.engine = "avatar_v"
|
||||
? {"type":"usd","usd":0.095381,"format":{"suffix":"/second"}}
|
||||
: widgets.engine = "avatar_iv"
|
||||
? {"type":"range_usd","min_usd":0.0715,"max_usd":0.095381,"format":{"suffix":"/second"}}
|
||||
: {"type":"range_usd","min_usd":0.023881,"max_usd":0.095381,"format":{"suffix":"/second"}}
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
engine: dict,
|
||||
speech: dict,
|
||||
custom_avatar_id: str = "",
|
||||
resolution: str = "1080p",
|
||||
aspect_ratio: str = "auto",
|
||||
background_color: str = "",
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
avatar_id, engine_type = await _resolve_avatar(cls, engine["avatar"], custom_avatar_id, engine["engine"])
|
||||
payload = {
|
||||
"type": "avatar",
|
||||
"avatar_id": avatar_id,
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"title": "ComfyUI Avatar Video",
|
||||
}
|
||||
if engine_type:
|
||||
payload["engine"] = {"type": engine_type}
|
||||
background_color = background_color.strip()
|
||||
if background_color:
|
||||
if not background_color.startswith("#"):
|
||||
raise ValueError("background_color must be a hex color code like '#00ff00'.")
|
||||
payload["background"] = {"type": "color", "value": background_color}
|
||||
await _apply_speech_source(cls, payload, speech, require_voice=False)
|
||||
video = await _create_and_poll_video(cls, payload)
|
||||
return IO.NodeOutput(await download_url_to_video_output(video["video_url"]))
|
||||
|
||||
|
||||
class HeyGenCreateAvatarNode(IO.ComfyNode):
|
||||
"""Create a reusable HeyGen avatar from a photo or a text prompt."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenCreateAvatarNode",
|
||||
display_name="HeyGen Create Avatar",
|
||||
category="partner/video/HeyGen",
|
||||
description="Create your own reusable HeyGen avatar from a photo of a person or "
|
||||
"from a text prompt (a generated character). Feed the resulting avatar_id into "
|
||||
"HeyGen Avatar Video's custom_avatar_id — and save the ID somewhere to reuse the "
|
||||
"avatar in future workflows.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"source",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"prompt",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Description of the avatar to generate (up to 1000 characters).",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_images",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Image.Input("ref_image"),
|
||||
names=[f"ref_image_{i}" for i in range(1, 4)],
|
||||
min=0,
|
||||
),
|
||||
tooltip="Up to 3 reference images guiding the generated look.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"photo",
|
||||
[
|
||||
IO.Image.Input(
|
||||
"identity_photo",
|
||||
tooltip="Photo of the person to turn into an avatar. "
|
||||
"Downscaled automatically if larger than 2K.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="Generate a new character from a text prompt, or create the avatar "
|
||||
"from a connected photo of a person.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.String.Output(
|
||||
display_name="avatar_id",
|
||||
tooltip="Avatar look ID. Pass it to HeyGen Avatar Video's custom_avatar_id; "
|
||||
"save it to reuse the avatar later.",
|
||||
),
|
||||
IO.Image.Output(display_name="preview"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":1.43}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
source: dict,
|
||||
) -> IO.NodeOutput:
|
||||
payload: dict = {"name": "ComfyUI Avatar"}
|
||||
if source["source"] == "photo":
|
||||
image = downscale_image_tensor_by_max_side(source["identity_photo"], max_side=2000)
|
||||
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||
payload["type"] = "photo"
|
||||
payload["file"] = {"type": "url", "url": image_url}
|
||||
else:
|
||||
validate_string(source["prompt"], strip_whitespace=True, min_length=1, max_length=1000)
|
||||
payload["type"] = "prompt"
|
||||
payload["prompt"] = source["prompt"]
|
||||
ref_tensors = [t for t in (source.get("reference_images") or {}).values() if t is not None]
|
||||
if ref_tensors:
|
||||
n_images = sum(get_number_of_images(t) for t in ref_tensors)
|
||||
if n_images > 3:
|
||||
raise ValueError(f"HeyGen accepts at most 3 reference images; got {n_images}.")
|
||||
scaled = [downscale_image_tensor_by_max_side(t, max_side=2000) for t in ref_tensors]
|
||||
ref_urls = await upload_images_to_comfyapi(
|
||||
cls, scaled, max_images=3, mime_type="image/png", total_pixels=None
|
||||
)
|
||||
payload["reference_images"] = [{"type": "url", "url": u} for u in ref_urls]
|
||||
created = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_AVATARS_PATH, method="POST"),
|
||||
data=payload,
|
||||
)
|
||||
look_id = ((created.get("data") or {}).get("avatar_item") or {}).get("id")
|
||||
if not look_id:
|
||||
raise ValueError(f"HeyGen did not return an avatar: {created}")
|
||||
final = await poll_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_LOOKS_PATH}/{look_id}"),
|
||||
# A missing status means the look needed no training and is ready.
|
||||
status_extractor=lambda r: (r.get("data") or {}).get("status") or "completed",
|
||||
failed_statuses=["failed", "pending_consent"],
|
||||
poll_interval=5.0,
|
||||
)
|
||||
data = final["data"]
|
||||
if data.get("preview_image_url"):
|
||||
preview = await download_url_to_image_tensor(data["preview_image_url"])
|
||||
else:
|
||||
preview = torch.zeros(1, 64, 64, 3)
|
||||
PromptServer.instance.send_progress_text(
|
||||
f"Please save the avatar_id for reuse.\n\navatar_id: {look_id}",
|
||||
cls.hidden.unique_id,
|
||||
)
|
||||
return IO.NodeOutput(look_id, preview)
|
||||
|
||||
|
||||
class HeyGenVideoTranslateNode(IO.ComfyNode):
|
||||
"""Translate a spoken video into another language with voice cloning and lip sync."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenVideoTranslateNode",
|
||||
display_name="HeyGen Video Translate",
|
||||
category="partner/video/HeyGen",
|
||||
description="Translate a spoken video into another language. Clones the original "
|
||||
"speaker's voice and re-animates the mouth to match the translated speech.",
|
||||
inputs=[
|
||||
IO.Video.Input(
|
||||
"video",
|
||||
tooltip="Video with speech to translate.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"output_language",
|
||||
options=HEYGEN_TRANSLATE_LANGUAGES,
|
||||
tooltip="Target language for the translated video.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"mode",
|
||||
options=["speed", "precision"],
|
||||
default="speed",
|
||||
tooltip="'speed' is faster; 'precision' produces higher-quality lip sync at twice the price.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"translate_audio_only",
|
||||
default=False,
|
||||
optional=True,
|
||||
tooltip="Only swap the audio track, keeping the original mouth movements (no lip sync).",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_count",
|
||||
default=0,
|
||||
min=0,
|
||||
max=10,
|
||||
optional=True,
|
||||
tooltip="Number of speakers in the video. 0 = detect automatically.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["mode"]),
|
||||
expr="""{"type":"usd","usd": widgets.mode = "precision" ? 0.095381 : 0.047619,"""
|
||||
""""format":{"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
video: Input.Video,
|
||||
output_language: str,
|
||||
mode: str,
|
||||
translate_audio_only: bool = False,
|
||||
speaker_count: int = 0,
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
video_url = await upload_video_to_comfyapi(cls, video)
|
||||
payload = {
|
||||
"video": {"type": "url", "url": video_url},
|
||||
"output_languages": [output_language],
|
||||
"mode": mode,
|
||||
"translate_audio_only": translate_audio_only,
|
||||
"title": "ComfyUI Video Translate",
|
||||
}
|
||||
if speaker_count > 0:
|
||||
payload["speaker_num"] = speaker_count
|
||||
created = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_TRANSLATIONS_PATH, method="POST"),
|
||||
data=payload,
|
||||
)
|
||||
translation_ids = (created.get("data") or {}).get("video_translation_ids") or []
|
||||
if not translation_ids:
|
||||
raise ValueError(f"HeyGen did not return a translation ID: {created}")
|
||||
final = await poll_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=f"{_TRANSLATIONS_PATH}/{translation_ids[0]}"),
|
||||
status_extractor=lambda r: (r.get("data") or {}).get("status"),
|
||||
queued_statuses=["pending"],
|
||||
poll_interval=5.0,
|
||||
)
|
||||
data = final["data"]
|
||||
if not data.get("video_url"):
|
||||
raise ValueError(f"HeyGen returned no video_url for translation {translation_ids[0]}.")
|
||||
return IO.NodeOutput(await download_url_to_video_output(data["video_url"]))
|
||||
|
||||
|
||||
class HeyGenTextToSpeechNode(IO.ComfyNode):
|
||||
"""Synthesize speech audio from text with HeyGen's Starfish TTS engine."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="HeyGenTextToSpeechNode",
|
||||
display_name="HeyGen Text to Speech",
|
||||
category="partner/audio/HeyGen",
|
||||
description="Generate speech audio from text using HeyGen's Starfish TTS engine. "
|
||||
"Includes HeyGen's most popular voices across 17 languages.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"text",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Text to synthesize (up to 5000 characters). The generated speech "
|
||||
"must be at least 1 second long.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"voice",
|
||||
options=HEYGEN_VOICE_TTS_OPTIONS,
|
||||
tooltip="Voice to use (curated from HeyGen's most popular Starfish-compatible voices).",
|
||||
),
|
||||
IO.String.Input(
|
||||
"custom_voice_id",
|
||||
default="",
|
||||
optional=True,
|
||||
tooltip="Optional HeyGen voice ID. When set, overrides the voice selected above. "
|
||||
"The voice must support the Starfish engine.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"speed",
|
||||
default=1.0,
|
||||
min=0.5,
|
||||
max=2.0,
|
||||
step=0.05,
|
||||
optional=True,
|
||||
tooltip="Speech speed multiplier.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"ssml",
|
||||
default=False,
|
||||
optional=True,
|
||||
tooltip="Treat the text as SSML markup (for pauses, emphasis, and pronunciation control).",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
optional=True,
|
||||
tooltip="Not sent to HeyGen; change it to force a re-run.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Audio.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.00095381,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
text: str,
|
||||
voice: str,
|
||||
custom_voice_id: str = "",
|
||||
speed: float = 1.0,
|
||||
ssml: bool = False,
|
||||
seed: int = 0,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(text, strip_whitespace=True, min_length=1, max_length=5000)
|
||||
payload = {
|
||||
"text": text,
|
||||
"voice_id": custom_voice_id.strip() or HEYGEN_VOICE_TTS_MAP[voice],
|
||||
"speed": round(speed, 2),
|
||||
}
|
||||
if ssml:
|
||||
payload["input_type"] = "ssml"
|
||||
response = await sync_op_raw(
|
||||
cls,
|
||||
ApiEndpoint(path=_SPEECH_PATH, method="POST"),
|
||||
data=payload,
|
||||
)
|
||||
audio_url = (response.get("data") or {}).get("audio_url")
|
||||
if not audio_url:
|
||||
raise ValueError(f"HeyGen did not return an audio_url: {response}")
|
||||
audio_bytes = await download_url_as_bytesio(audio_url)
|
||||
return IO.NodeOutput(audio_bytes_to_audio_input(audio_bytes.getvalue()))
|
||||
|
||||
|
||||
class HeyGenExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
HeyGenTalkingPhotoNode,
|
||||
HeyGenAvatarVideoNode,
|
||||
HeyGenCreateAvatarNode,
|
||||
HeyGenVideoTranslateNode,
|
||||
HeyGenTextToSpeechNode,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> HeyGenExtension:
|
||||
return HeyGenExtension()
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -41,6 +41,9 @@ STARTING_POINT_ID_PATTERN = r"<starting_point_id:(.*)>"
|
|||
|
||||
|
||||
class SupportedOpenAIModel(str, Enum):
|
||||
gpt_5_6_sol = "gpt-5.6-sol"
|
||||
gpt_5_6_terra = "gpt-5.6-terra"
|
||||
gpt_5_6_luna = "gpt-5.6-luna"
|
||||
gpt_5_5_pro = "gpt-5.5-pro"
|
||||
gpt_5_5 = "gpt-5.5"
|
||||
gpt_5 = "gpt-5"
|
||||
|
|
@ -1063,6 +1066,21 @@ class OpenAIChatNode(IO.ComfyNode):
|
|||
"usd": [0.002, 0.008],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
}
|
||||
: $contains($m, "gpt-5.6-terra") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.0025, 0.015],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
}
|
||||
: $contains($m, "gpt-5.6-luna") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.001, 0.006],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
}
|
||||
: $contains($m, "gpt-5.6") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.005, 0.03],
|
||||
"format": { "approximate": true, "separator": "-", "suffix": " per 1K tokens" }
|
||||
}
|
||||
: $contains($m, "gpt-5.5-pro") ? {
|
||||
"type": "list_usd",
|
||||
"usd": [0.03, 0.18],
|
||||
|
|
|
|||
|
|
@ -0,0 +1,391 @@
|
|||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.sync_so import (
|
||||
SyncActiveSpeakerDetection,
|
||||
SyncGeneration,
|
||||
SyncGenerationOptions,
|
||||
SyncGenerationRequest,
|
||||
SyncInputItem,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
download_url_to_video_output,
|
||||
downscale_image_tensor,
|
||||
downscale_image_tensor_by_max_side,
|
||||
get_image_dimensions,
|
||||
get_number_of_images,
|
||||
poll_op,
|
||||
sync_op,
|
||||
upload_audio_to_comfyapi,
|
||||
upload_image_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_audio_duration,
|
||||
)
|
||||
|
||||
|
||||
class SyncLipSyncNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="SyncLipSyncNode",
|
||||
display_name="sync.so Lip Sync",
|
||||
category="partner/video/sync.so",
|
||||
description=(
|
||||
"Re-sync mouth movement in a video to new speech audio using sync.so. "
|
||||
"Handles close-ups, profiles and obstructions automatically while preserving "
|
||||
"the speaker's expression. Cost scales with output duration."
|
||||
),
|
||||
inputs=[
|
||||
IO.Video.Input(
|
||||
"video",
|
||||
tooltip="Footage of the speaker to re-sync. Up to 4K (4096x2160); "
|
||||
"a constant frame rate of 24/25/30 fps works best.",
|
||||
),
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Speech audio to sync the mouth to.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=42,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed controls whether the node should re-run; "
|
||||
"results are non-deterministic regardless of seed.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"sync-3",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"sync_mode",
|
||||
options=["bounce", "cut_off", "loop", "silence", "remap"],
|
||||
default="bounce",
|
||||
tooltip=(
|
||||
"How to handle a duration mismatch between video and audio; "
|
||||
"this also sets the output length. "
|
||||
"bounce: video plays forward then backward until the audio ends "
|
||||
"(output = audio length). "
|
||||
"loop: video restarts until the audio ends (output = audio length). "
|
||||
"remap: video is time-stretched to match the audio (output = audio length). "
|
||||
"cut_off: the longer track is trimmed (output = shorter length). "
|
||||
"silence: nothing is trimmed; the shorter track is padded "
|
||||
"(output = longer length)."
|
||||
),
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"speaker_selection",
|
||||
options=["default", "auto-detect", "coordinates"],
|
||||
default="default",
|
||||
tooltip=(
|
||||
"Which face to lipsync when several people are visible. "
|
||||
"default: let the model decide. "
|
||||
"auto-detect: detect and follow the active speaker. "
|
||||
"coordinates: target the face at pixel (speaker_x, speaker_y) "
|
||||
"in the frame chosen by speaker_frame."
|
||||
),
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_frame",
|
||||
default=0,
|
||||
min=0,
|
||||
max=1_000_000,
|
||||
advanced=True,
|
||||
tooltip="Video frame used to locate the speaker. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_x",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="X pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_y",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="Y pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
tooltip="sync.so generation model.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
video: Input.Video,
|
||||
audio: Input.Audio,
|
||||
seed: int,
|
||||
model: dict,
|
||||
) -> IO.NodeOutput:
|
||||
try:
|
||||
width, height = video.get_dimensions()
|
||||
except Exception:
|
||||
width = height = None
|
||||
if width and height and (max(width, height) > 4096 or width * height > 4096 * 2160):
|
||||
raise ValueError(
|
||||
f"sync.so rejects videos above 4K (4096x2160); got {width}x{height}. Downscale the video first."
|
||||
)
|
||||
validate_audio_duration(audio, max_duration=600)
|
||||
|
||||
if model["speaker_selection"] == "auto-detect":
|
||||
speaker_detection = SyncActiveSpeakerDetection(auto_detect=True)
|
||||
elif model["speaker_selection"] == "coordinates":
|
||||
speaker_detection = SyncActiveSpeakerDetection(
|
||||
frame_number=model["speaker_frame"],
|
||||
coordinates=[model["speaker_x"], model["speaker_y"]],
|
||||
)
|
||||
else:
|
||||
speaker_detection = None
|
||||
|
||||
video_url = await upload_video_to_comfyapi(cls, video, max_duration=600)
|
||||
audio_url = await upload_audio_to_comfyapi(cls, audio)
|
||||
|
||||
generation = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
|
||||
response_model=SyncGeneration,
|
||||
data=SyncGenerationRequest(
|
||||
model=model["model"],
|
||||
input=[
|
||||
SyncInputItem(type="video", url=video_url),
|
||||
SyncInputItem(type="audio", url=audio_url),
|
||||
],
|
||||
options=SyncGenerationOptions(
|
||||
sync_mode=model["sync_mode"],
|
||||
active_speaker_detection=speaker_detection,
|
||||
),
|
||||
),
|
||||
)
|
||||
generation = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
|
||||
response_model=SyncGeneration,
|
||||
status_extractor=lambda g: g.status,
|
||||
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
|
||||
failed_statuses=[],
|
||||
queued_statuses=["PENDING"],
|
||||
poll_interval=10.0,
|
||||
)
|
||||
if generation.status != "COMPLETED":
|
||||
code = f" [{generation.errorCode}]" if generation.errorCode else ""
|
||||
raise ValueError(
|
||||
f"sync.so generation {generation.status.lower()}{code}: "
|
||||
f"{generation.error or 'no error details provided'}"
|
||||
)
|
||||
if not generation.outputUrl:
|
||||
raise ValueError("sync.so generation completed but no output URL was returned.")
|
||||
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
|
||||
|
||||
|
||||
class SyncTalkingImageNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="SyncTalkingImageNode",
|
||||
display_name="sync.so Talking Image",
|
||||
category="partner/video/sync.so",
|
||||
description=(
|
||||
"Animate a still portrait into a talking video driven by speech audio, "
|
||||
"using sync.so's sync-3 model. The output duration matches the audio. "
|
||||
"Cost scales with output duration."
|
||||
),
|
||||
inputs=[
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
tooltip="A single image with a clearly visible face, up to 4K (4096x2160).",
|
||||
),
|
||||
IO.Audio.Input(
|
||||
"audio",
|
||||
tooltip="Speech audio driving the talking video; the output duration matches it. "
|
||||
"Chain any TTS node here to drive the animation from text.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Optional guidance for how the portrait comes to life, e.g. "
|
||||
"'make the subject smile and look at the camera'. "
|
||||
"Leave empty for natural talking motion.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed controls whether the node should re-run; "
|
||||
"results are non-deterministic regardless of seed.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"sync-3",
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"speaker_selection",
|
||||
options=["default", "coordinates"],
|
||||
default="default",
|
||||
tooltip=(
|
||||
"Which face to animate when several people are visible. "
|
||||
"default: let the model decide. "
|
||||
"coordinates: target the face at pixel (speaker_x, speaker_y) "
|
||||
"in the image. Auto-detection is not supported for images."
|
||||
),
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_x",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="X pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speaker_y",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4096,
|
||||
advanced=True,
|
||||
tooltip="Y pixel coordinate of the speaker's face. "
|
||||
"Only used when speaker_selection is 'coordinates'.",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"auto_downscale",
|
||||
default=True,
|
||||
advanced=True,
|
||||
tooltip="Automatically downscale the image if it exceeds the 4K "
|
||||
"(4096x2160) input limit; speaker coordinates are scaled to match. "
|
||||
"When disabled, an oversized image raises an error instead.",
|
||||
),
|
||||
],
|
||||
)
|
||||
],
|
||||
tooltip="sync.so generation model. Image input is exclusive to sync-3.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output()],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.19019,"format":{"approximate":true,"suffix":"/second"}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: Input.Image,
|
||||
audio: Input.Audio,
|
||||
prompt: str,
|
||||
seed: int,
|
||||
model: dict,
|
||||
) -> IO.NodeOutput:
|
||||
if get_number_of_images(image) != 1:
|
||||
raise ValueError("Exactly one image is required; got a batch. Pick one frame first.")
|
||||
validate_audio_duration(audio, max_duration=600)
|
||||
|
||||
height, width = get_image_dimensions(image)
|
||||
speaker_x, speaker_y = model["speaker_x"], model["speaker_y"]
|
||||
if max(width, height) > 4096 or width * height > 4096 * 2160:
|
||||
if not model["auto_downscale"]:
|
||||
raise ValueError(
|
||||
f"sync.so rejects images above 4K (4096x2160); got {width}x{height}. "
|
||||
"Downscale the image first or enable auto_downscale."
|
||||
)
|
||||
image = downscale_image_tensor(image, total_pixels=4096 * 2160)
|
||||
image = downscale_image_tensor_by_max_side(image, max_side=4096)
|
||||
new_height, new_width = get_image_dimensions(image)
|
||||
# speaker coordinates are given in the original image's pixel space
|
||||
speaker_x = min(new_width - 1, round(speaker_x * new_width / width))
|
||||
speaker_y = min(new_height - 1, round(speaker_y * new_height / height))
|
||||
|
||||
if model["speaker_selection"] == "coordinates":
|
||||
speaker_detection = SyncActiveSpeakerDetection(
|
||||
frame_number=0, # images have a single frame; auto_detect is rejected by the API
|
||||
coordinates=[speaker_x, speaker_y],
|
||||
)
|
||||
else:
|
||||
speaker_detection = None
|
||||
|
||||
image_url = await upload_image_to_comfyapi(cls, image, mime_type="image/png", total_pixels=None)
|
||||
audio_url = await upload_audio_to_comfyapi(cls, audio)
|
||||
|
||||
generation = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/synclabs/v2/generate", method="POST"),
|
||||
response_model=SyncGeneration,
|
||||
data=SyncGenerationRequest(
|
||||
model=model["model"],
|
||||
input=[
|
||||
SyncInputItem(type="image", url=image_url),
|
||||
SyncInputItem(type="audio", url=audio_url),
|
||||
],
|
||||
options=SyncGenerationOptions(
|
||||
i2v_prompt=prompt.strip() or None,
|
||||
active_speaker_detection=speaker_detection,
|
||||
),
|
||||
),
|
||||
)
|
||||
generation = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/synclabs/v2/generate/{generation.id}"),
|
||||
response_model=SyncGeneration,
|
||||
status_extractor=lambda g: g.status,
|
||||
completed_statuses=["COMPLETED", "FAILED", "REJECTED"],
|
||||
failed_statuses=[],
|
||||
queued_statuses=["PENDING"],
|
||||
poll_interval=10.0,
|
||||
)
|
||||
if generation.status != "COMPLETED":
|
||||
code = f" [{generation.errorCode}]" if generation.errorCode else ""
|
||||
raise ValueError(
|
||||
f"sync.so generation {generation.status.lower()}{code}: "
|
||||
f"{generation.error or 'no error details provided'}"
|
||||
)
|
||||
if not generation.outputUrl:
|
||||
raise ValueError("sync.so generation completed but no output URL was returned.")
|
||||
return IO.NodeOutput(await download_url_to_video_output(generation.outputUrl))
|
||||
|
||||
|
||||
class SyncExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
SyncLipSyncNode,
|
||||
SyncTalkingImageNode,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> SyncExtension:
|
||||
return SyncExtension()
|
||||
|
|
@ -11,9 +11,12 @@ from io import BytesIO
|
|||
from yarl import URL
|
||||
|
||||
from comfy.cli_args import args
|
||||
from comfy.comfy_api_env import normalize_comfy_api_base
|
||||
from comfy.deploy_environment import get_deploy_environment
|
||||
from comfy.model_management import processing_interrupted
|
||||
from comfy_api.latest import IO
|
||||
from comfy_execution.utils import get_executing_context
|
||||
from comfyui_version import __version__ as comfyui_version
|
||||
|
||||
from .common_exceptions import ProcessingInterrupted
|
||||
|
||||
|
|
@ -55,15 +58,20 @@ def get_comfy_api_headers(node_cls: type[IO.ComfyNode]) -> dict[str, str]:
|
|||
relative/cloud URLs resolved against ``default_base_url()``; because the result
|
||||
includes auth, callers must not attach it to arbitrary absolute/presigned URLs.
|
||||
"""
|
||||
return {
|
||||
headers = {
|
||||
**get_auth_header(node_cls),
|
||||
"Comfy-Env": get_deploy_environment(),
|
||||
"Comfy-Usage-Source": get_usage_source(node_cls),
|
||||
"Comfy-Core-Version": comfyui_version,
|
||||
}
|
||||
ctx = get_executing_context()
|
||||
if ctx is not None:
|
||||
headers["Comfy-Job-Id"] = ctx.prompt_id
|
||||
return headers
|
||||
|
||||
|
||||
def default_base_url() -> str:
|
||||
return getattr(args, "comfy_api_base", "https://api.comfy.org")
|
||||
return normalize_comfy_api_base(getattr(args, "comfy_api_base", "https://api.comfy.org"))
|
||||
|
||||
|
||||
async def sleep_with_interrupt(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ''
|
||||
|
|
|
|||
|
|
@ -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"),
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -844,15 +844,18 @@ class ImageMergeTileList(IO.ComfyNode):
|
|||
# Format specifications
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# Maps (file_format, bit_depth, has_alpha) -> (numpy dtype scale, av pixel format,
|
||||
# stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
||||
# Maps (file_format, bit_depth, num_channels) -> (quantization scale, numpy dtype,
|
||||
# av frame pix_fmt, stream pix_fmt). Keeps the encode path declarative instead of branchy.
|
||||
_FORMAT_SPECS = {
|
||||
("png", "8-bit", False): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
|
||||
("png", "8-bit", True): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
|
||||
("png", "16-bit", False): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
|
||||
("png", "16-bit", True): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"},
|
||||
("exr", "32-bit float", False): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"},
|
||||
("exr", "32-bit float", True): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
||||
("png", "8-bit", 1): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "gray", "stream_fmt": "gray"},
|
||||
("png", "8-bit", 3): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgb24", "stream_fmt": "rgb24"},
|
||||
("png", "8-bit", 4): {"scale": 255.0, "dtype": np.uint8, "frame_fmt": "rgba", "stream_fmt": "rgba"},
|
||||
("png", "16-bit", 1): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "gray16le", "stream_fmt": "gray16be"},
|
||||
("png", "16-bit", 3): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgb48le", "stream_fmt": "rgb48be"},
|
||||
("png", "16-bit", 4): {"scale": 65535.0, "dtype": np.uint16, "frame_fmt": "rgba64le", "stream_fmt": "rgba64be"},
|
||||
("exr", "32-bit float", 1): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "grayf32le", "stream_fmt": "grayf32le"},
|
||||
("exr", "32-bit float", 3): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrpf32le", "stream_fmt": "gbrpf32le"},
|
||||
("exr", "32-bit float", 4): {"scale": 1.0, "dtype": np.float32, "frame_fmt": "gbrapf32le", "stream_fmt": "gbrapf32le"},
|
||||
}
|
||||
|
||||
|
||||
|
|
@ -891,10 +894,11 @@ def hlg_to_linear(t: torch.Tensor) -> torch.Tensor:
|
|||
return torch.cat([hlg_to_linear(rgb), alpha], dim=-1)
|
||||
|
||||
# Piecewise: sqrt branch below 0.5, log branch above.
|
||||
# Clamp inside the log branch so negative / out-of-range values don't blow up;
|
||||
# Clamp the log branch at the 0.5 branch point (not above it) so the
|
||||
# unselected lane stays finite in exp() without altering selected values;
|
||||
# values above 1.0 are allowed and extrapolate naturally.
|
||||
low = (t ** 2) / 3.0
|
||||
high = (torch.exp((t.clamp(min=_HLG_C) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
high = (torch.exp((t.clamp(min=0.5) - _HLG_C) / _HLG_A) + _HLG_B) / 12.0
|
||||
return torch.where(t <= 0.5, low, high)
|
||||
|
||||
|
||||
|
|
@ -1087,7 +1091,8 @@ def _encode_image(
|
|||
bit_depth: str,
|
||||
colorspace: str,
|
||||
) -> bytes:
|
||||
"""Encode a single HxWxC tensor to PNG or EXR bytes in memory.
|
||||
"""Encode a single HxWxC (or channel-less HxW grayscale) tensor to PNG or
|
||||
EXR bytes in memory. Grayscale is written as single-channel PNG / Y-only EXR.
|
||||
|
||||
For EXR the input is interpreted according to `colorspace` and converted
|
||||
to scene-linear (EXR's convention) before writing:
|
||||
|
|
@ -1101,10 +1106,16 @@ def _encode_image(
|
|||
For PNG, colorspace selection does not modify pixels — PNG is delivered
|
||||
sRGB-encoded and there is no PNG path for wide-gamut HDR in this node.
|
||||
"""
|
||||
if img_tensor.ndim == 2:
|
||||
img_tensor = img_tensor.unsqueeze(-1) # Some nodes emit grayscale as (H, W) with no channel dim, mask-style.
|
||||
height, width, num_channels = img_tensor.shape
|
||||
has_alpha = num_channels == 4
|
||||
|
||||
spec = _FORMAT_SPECS[(file_format, bit_depth, has_alpha)]
|
||||
spec = _FORMAT_SPECS.get((file_format, bit_depth, num_channels))
|
||||
if spec is None:
|
||||
raise ValueError(
|
||||
f"No {file_format}/{bit_depth} encoder for {num_channels}-channel images: "
|
||||
"supported channel counts are 1 (grayscale), 3 (RGB) and 4 (RGBA)."
|
||||
)
|
||||
|
||||
if spec["dtype"] == np.float32:
|
||||
# EXR path: preserve full range, no clamp.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,102 @@
|
|||
from typing_extensions import override
|
||||
|
||||
import comfy.utils
|
||||
import node_helpers
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
|
||||
# fmt: off
|
||||
BUCKETS_1024 = [
|
||||
(512, 1792), (512, 1856), (512, 1920), (512, 1984), (512, 2048),
|
||||
(576, 1600), (576, 1664), (576, 1728), (576, 1792),
|
||||
(640, 1472), (640, 1536), (640, 1600),
|
||||
(704, 1344), (704, 1408), (704, 1472),
|
||||
(768, 1216), (768, 1280), (768, 1344),
|
||||
(832, 1152), (832, 1216),
|
||||
(896, 1088), (896, 1152),
|
||||
(960, 1024), (960, 1088),
|
||||
(1024, 960), (1024, 1024),
|
||||
(1088, 896), (1088, 960),
|
||||
(1152, 832), (1152, 896),
|
||||
(1216, 768), (1216, 832),
|
||||
(1280, 768),
|
||||
(1344, 704), (1344, 768),
|
||||
(1408, 704),
|
||||
(1472, 640), (1472, 704),
|
||||
(1536, 640),
|
||||
(1600, 576), (1600, 640),
|
||||
(1664, 576),
|
||||
(1728, 576),
|
||||
(1792, 512), (1792, 576),
|
||||
(1856, 512),
|
||||
(1920, 512),
|
||||
(1984, 512),
|
||||
(2048, 512),
|
||||
]
|
||||
# fmt: on
|
||||
|
||||
|
||||
def _find_best_bucket(height: int, width: int) -> tuple[int, int]:
|
||||
target_ratio = height / width
|
||||
return min(BUCKETS_1024, key=lambda hw: abs(hw[0] / hw[1] - target_ratio))
|
||||
|
||||
|
||||
def _resize_reference(image):
|
||||
if image.shape[0] != 1:
|
||||
raise ValueError("JoyImage reference inputs must contain one image each")
|
||||
samples = image.movedim(-1, 1)
|
||||
bucket_h, bucket_w = _find_best_bucket(samples.shape[2], samples.shape[3])
|
||||
resized = comfy.utils.common_upscale(samples, bucket_w, bucket_h, "bilinear", "center")
|
||||
return resized.movedim(1, -1)[:, :, :, :3]
|
||||
|
||||
|
||||
def _encode(clip, prompt, vae, images):
|
||||
resized_images = [_resize_reference(image) for image in images]
|
||||
conditioning = clip.encode_from_tokens_scheduled(clip.tokenize(prompt, images=resized_images))
|
||||
if vae is not None and resized_images:
|
||||
ref_latents = [vae.encode(image) for image in resized_images]
|
||||
conditioning = node_helpers.conditioning_set_values(
|
||||
conditioning, {"reference_latents": ref_latents}, append=True,
|
||||
)
|
||||
return conditioning
|
||||
|
||||
|
||||
class TextEncodeJoyImageEdit(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
image_template = io.Autogrow.TemplatePrefix(
|
||||
io.Image.Input("image"),
|
||||
prefix="image",
|
||||
min=0,
|
||||
max=6,
|
||||
)
|
||||
return io.Schema(
|
||||
node_id="TextEncodeJoyImageEdit",
|
||||
category="model/conditioning/joyimage",
|
||||
inputs=[
|
||||
io.Clip.Input("clip"),
|
||||
io.String.Input("prompt", multiline=True, dynamic_prompts=True),
|
||||
io.Vae.Input("vae", optional=True),
|
||||
io.Autogrow.Input("images", template=image_template, optional=True),
|
||||
],
|
||||
outputs=[
|
||||
io.Conditioning.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, clip, prompt, vae=None, images: io.Autogrow.Type = None) -> io.NodeOutput:
|
||||
images = images or {}
|
||||
return io.NodeOutput(_encode(clip, prompt, vae, list(images.values())))
|
||||
|
||||
|
||||
class JoyImageExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [
|
||||
TextEncodeJoyImageEdit,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> JoyImageExtension:
|
||||
return JoyImageExtension()
|
||||
|
|
@ -61,14 +61,10 @@ class Load3D(IO.ComfyNode):
|
|||
|
||||
@classmethod
|
||||
def execute(cls, model_file, image, **kwargs) -> IO.NodeOutput:
|
||||
image_path = folder_paths.get_annotated_filepath(image['image'])
|
||||
mask_path = folder_paths.get_annotated_filepath(image['mask'])
|
||||
normal_path = folder_paths.get_annotated_filepath(image['normal'])
|
||||
|
||||
load_image_node = nodes.LoadImage()
|
||||
output_image, ignore_mask = load_image_node.load_image(image=image_path)
|
||||
ignore_image, output_mask = load_image_node.load_image(image=mask_path)
|
||||
normal_image, ignore_mask2 = load_image_node.load_image(image=normal_path)
|
||||
output_image, ignore_mask = load_image_node.load_image(image=image['image'])
|
||||
ignore_image, output_mask = load_image_node.load_image(image=image['mask'])
|
||||
normal_image, ignore_mask2 = load_image_node.load_image(image=image['normal'])
|
||||
|
||||
video = None
|
||||
|
||||
|
|
@ -96,6 +92,7 @@ class Preview3D(IO.ComfyNode):
|
|||
search_aliases=["view mesh", "3d viewer"],
|
||||
display_name="Preview 3D & Animation",
|
||||
category="3d",
|
||||
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
|
|
@ -140,6 +137,7 @@ class Preview3DAdvanced(IO.ComfyNode):
|
|||
display_name="Preview 3D (Advanced)",
|
||||
search_aliases=["preview 3d", "3d viewer", "view mesh", "frame 3d", "3d camera output"],
|
||||
category="3d",
|
||||
description="Preview a 3D model file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
|
|
@ -176,8 +174,9 @@ class Preview3DAdvanced(IO.ComfyNode):
|
|||
filename = f"preview3d_advanced_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
|
|
@ -197,6 +196,7 @@ class PreviewGaussianSplat(IO.ComfyNode):
|
|||
node_id="PreviewGaussianSplat",
|
||||
display_name="Preview Splat",
|
||||
category="3d",
|
||||
description="Preview a gaussian splat 3D file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
search_aliases=[
|
||||
|
|
@ -244,8 +244,9 @@ class PreviewGaussianSplat(IO.ComfyNode):
|
|||
filename = f"preview_splat_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
|
|
@ -265,6 +266,7 @@ class PreviewPointCloud(IO.ComfyNode):
|
|||
node_id="PreviewPointCloud",
|
||||
display_name="Preview Point Cloud",
|
||||
category="3d",
|
||||
description="Preview a point cloud 3D file without saving it to the ComfyUI output directory.",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
search_aliases=[
|
||||
|
|
@ -303,8 +305,9 @@ class PreviewPointCloud(IO.ComfyNode):
|
|||
filename = f"preview_pointcloud_{uuid.uuid4().hex}.{model_3d.format}"
|
||||
model_3d.save_to(os.path.join(folder_paths.get_temp_directory(), filename))
|
||||
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state['camera_info']
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
|
|
@ -375,8 +378,9 @@ class Load3DAdvanced(IO.ComfyNode):
|
|||
file_3d = None
|
||||
if model_file and model_file != "none":
|
||||
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
model_3d_info = viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height)
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state.get('camera_info'), width, height)
|
||||
|
||||
|
||||
class Load3DExtension(ComfyExtension):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ import comfy.ldm.common_dit
|
|||
import comfy.latent_formats
|
||||
import comfy.ldm.lumina.controlnet
|
||||
import comfy.ldm.supir.supir_modules
|
||||
import comfy.ldm.anima.lllite
|
||||
from comfy.ldm.wan.model_multitalk import WanMultiTalkAttentionBlock, MultiTalkAudioProjModel
|
||||
from comfy_api.latest import io
|
||||
from comfy.ldm.supir.supir_patch import SUPIRPatch
|
||||
|
|
@ -236,10 +237,12 @@ class ModelPatchLoader:
|
|||
|
||||
def load_model_patch(self, name):
|
||||
model_patch_path = folder_paths.get_full_path_or_raise("model_patches", name)
|
||||
sd = comfy.utils.load_torch_file(model_patch_path, safe_load=True)
|
||||
sd, metadata = comfy.utils.load_torch_file(model_patch_path, safe_load=True, return_metadata=True)
|
||||
dtype = comfy.utils.weight_dtype(sd)
|
||||
|
||||
if 'controlnet_blocks.0.y_rms.weight' in sd:
|
||||
if 'lllite_conditioning1.conv1.weight' in sd:
|
||||
model = comfy.ldm.anima.lllite.AnimaLLLite(sd, metadata, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
|
||||
elif 'controlnet_blocks.0.y_rms.weight' in sd:
|
||||
additional_in_dim = sd["img_in.weight"].shape[1] - 64
|
||||
model = QwenImageBlockWiseControlNet(additional_in_dim=additional_in_dim, device=comfy.model_management.unet_offload_device(), dtype=dtype, operations=comfy.ops.manual_cast)
|
||||
elif 'feature_embedder.mid_layer_norm.bias' in sd:
|
||||
|
|
@ -296,6 +299,50 @@ class ModelPatchLoader:
|
|||
return (model_patcher,)
|
||||
|
||||
|
||||
class AnimaLLLiteApply:
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": {"model": ("MODEL",),
|
||||
"model_patch": ("MODEL_PATCH",),
|
||||
"image": ("IMAGE",),
|
||||
"strength": ("FLOAT", {"default": 1.0, "min": -10.0, "max": 10.0, "step": 0.01}),
|
||||
"start_percent": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"end_percent": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
},
|
||||
"optional": {"mask": ("MASK",),
|
||||
}}
|
||||
RETURN_TYPES = ("MODEL",)
|
||||
FUNCTION = "apply_patch"
|
||||
EXPERIMENTAL = True
|
||||
|
||||
CATEGORY = "model_patches/anima"
|
||||
|
||||
def apply_patch(self, model, model_patch, image, strength, start_percent, end_percent, mask=None):
|
||||
image = image[..., :3]
|
||||
|
||||
if model_patch.model.cond_in_channels == 4 and mask is None:
|
||||
mask = torch.zeros_like(image[..., 0])
|
||||
elif model_patch.model.cond_in_channels != 4:
|
||||
mask = None
|
||||
|
||||
model_sampling = model.get_model_object("model_sampling")
|
||||
sigma_start = float(model_sampling.percent_to_sigma(start_percent))
|
||||
sigma_end = float(model_sampling.percent_to_sigma(end_percent))
|
||||
patch = comfy.ldm.anima.lllite.AnimaLLLitePatch(model_patch, image, mask, strength, sigma_start, sigma_end)
|
||||
model_patched = model.clone()
|
||||
model_patched.set_model_post_input_patch(patch)
|
||||
model_patched.set_model_attn1_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(
|
||||
patch,
|
||||
{"q": "self_attn_q_proj", "k": "self_attn_k_proj", "v": "self_attn_v_proj"},
|
||||
))
|
||||
model_patched.set_model_attn2_patch(comfy.ldm.anima.lllite.AnimaLLLiteAttentionPatch(
|
||||
patch,
|
||||
{"q": "cross_attn_q_proj"},
|
||||
))
|
||||
model_patched.set_model_patch(comfy.ldm.anima.lllite.AnimaLLLiteMLPPatch(patch), "mlp_patch")
|
||||
return (model_patched,)
|
||||
|
||||
|
||||
class DiffSynthCnetPatch:
|
||||
def __init__(self, model_patch, vae, image, strength, mask=None):
|
||||
self.model_patch = model_patch
|
||||
|
|
@ -674,6 +721,7 @@ NODE_CLASS_MAPPINGS = {
|
|||
"ZImageFunControlnet": ZImageFunControlnet,
|
||||
"USOStyleReference": USOStyleReference,
|
||||
"SUPIRApply": SUPIRApply,
|
||||
"AnimaLLLiteApply": AnimaLLLiteApply,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
|
|
@ -682,4 +730,5 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
|||
"ZImageFunControlnet": "Apply Z-Image Fun ControlNet",
|
||||
"USOStyleReference": "Apply USO Style Reference",
|
||||
"SUPIRApply": "Apply SUPIR Patch",
|
||||
"AnimaLLLiteApply": "Apply Anima LLLite",
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)],
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ from typing_extensions import override
|
|||
|
||||
import folder_paths
|
||||
from comfy.cli_args import args
|
||||
from comfy_api.latest import ComfyExtension, IO, Types
|
||||
from comfy_api.latest import ComfyExtension, IO, Types, UI
|
||||
|
||||
|
||||
def pack_variable_mesh_batch(vertices, faces, colors=None, uvs=None, texture=None, unlit=False):
|
||||
|
|
@ -406,10 +406,165 @@ class SaveGLB(IO.ComfyNode):
|
|||
return IO.NodeOutput(ui={"3d": results})
|
||||
|
||||
|
||||
def _save_file3d_to_output(model_3d: Types.File3D, filename_prefix: str) -> str:
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(
|
||||
filename_prefix, folder_paths.get_output_directory()
|
||||
)
|
||||
ext = model_3d.format or "glb"
|
||||
saved_filename = f"{filename}_{counter:05}.{ext}"
|
||||
model_3d.save_to(os.path.join(full_output_folder, saved_filename))
|
||||
return f"{subfolder}/{saved_filename}" if subfolder else saved_filename
|
||||
|
||||
|
||||
def execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs) -> IO.NodeOutput:
|
||||
model_file = _save_file3d_to_output(model_3d, filename_prefix)
|
||||
viewport_state = viewport_state if isinstance(viewport_state, dict) else {}
|
||||
camera_info_input = kwargs.get("camera_info", None)
|
||||
camera_info = camera_info_input if camera_info_input is not None else viewport_state.get('camera_info')
|
||||
model_3d_info_input = kwargs.get("model_3d_info", None)
|
||||
model_3d_info = model_3d_info_input if model_3d_info_input is not None else viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(
|
||||
model_3d,
|
||||
model_3d_info,
|
||||
camera_info,
|
||||
width,
|
||||
height,
|
||||
ui=UI.PreviewUI3DAdvanced(model_file, camera_info, model_3d_info),
|
||||
)
|
||||
|
||||
|
||||
class Save3DAdvanced(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="Save3DAdvanced",
|
||||
display_name="Save 3D (Advanced)",
|
||||
search_aliases=["save 3d", "export 3d model", "save mesh advanced"],
|
||||
category="3d",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[
|
||||
IO.File3DGLB,
|
||||
IO.File3DGLTF,
|
||||
IO.File3DFBX,
|
||||
IO.File3DOBJ,
|
||||
IO.File3DSTL,
|
||||
IO.File3DUSDZ,
|
||||
IO.File3DAny,
|
||||
],
|
||||
tooltip="3D model file from an upstream 3D node.",
|
||||
),
|
||||
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
|
||||
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
|
||||
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
|
||||
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
|
||||
],
|
||||
outputs=[
|
||||
IO.File3DAny.Output(display_name="model_3d"),
|
||||
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
|
||||
IO.Load3DCamera.Output(display_name="camera_info"),
|
||||
IO.Int.Output(display_name="width"),
|
||||
IO.Int.Output(display_name="height"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
|
||||
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
|
||||
|
||||
|
||||
class SaveGaussianSplat(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SaveGaussianSplat",
|
||||
display_name="Save Splat",
|
||||
search_aliases=["save splat", "save gaussian splat", "export gaussian", "export splat"],
|
||||
category="3d",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[
|
||||
IO.File3DSplatAny,
|
||||
IO.File3DPLY,
|
||||
IO.File3DSPLAT,
|
||||
IO.File3DSPZ,
|
||||
IO.File3DKSPLAT,
|
||||
],
|
||||
tooltip="A gaussian splat 3D file.",
|
||||
),
|
||||
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
|
||||
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
|
||||
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
|
||||
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
|
||||
],
|
||||
outputs=[
|
||||
IO.File3DSplatAny.Output(display_name="model_3d"),
|
||||
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
|
||||
IO.Load3DCamera.Output(display_name="camera_info"),
|
||||
IO.Int.Output(display_name="width"),
|
||||
IO.Int.Output(display_name="height"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
|
||||
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
|
||||
|
||||
|
||||
class SavePointCloud(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="SavePointCloud",
|
||||
display_name="Save Point Cloud",
|
||||
search_aliases=["save point cloud", "save pointcloud", "export point cloud"],
|
||||
category="3d",
|
||||
is_experimental=True,
|
||||
is_output_node=True,
|
||||
inputs=[
|
||||
IO.MultiType.Input(
|
||||
"model_3d",
|
||||
types=[
|
||||
IO.File3DPointCloudAny,
|
||||
IO.File3DPLY,
|
||||
],
|
||||
tooltip="Point cloud file (.ply)",
|
||||
),
|
||||
IO.String.Input("filename_prefix", default="3d/ComfyUI"),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
IO.Load3DModelInfo.Input("model_3d_info", optional=True, advanced=True),
|
||||
IO.Load3DCamera.Input("camera_info", optional=True, advanced=True),
|
||||
IO.Int.Input("width", default=1024, min=1, max=4096, step=1),
|
||||
IO.Int.Input("height", default=1024, min=1, max=4096, step=1),
|
||||
],
|
||||
outputs=[
|
||||
IO.File3DPointCloudAny.Output(display_name="model_3d"),
|
||||
IO.Load3DModelInfo.Output(display_name="model_3d_info"),
|
||||
IO.Load3DCamera.Output(display_name="camera_info"),
|
||||
IO.Int.Output(display_name="width"),
|
||||
IO.Int.Output(display_name="height"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_3d: Types.File3D, viewport_state, width: int, height: int, filename_prefix: str, **kwargs) -> IO.NodeOutput:
|
||||
return execute_save_3d_advanced(model_3d, viewport_state, width, height, filename_prefix, kwargs)
|
||||
|
||||
|
||||
class Save3DExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [SaveGLB]
|
||||
return [SaveGLB, Save3DAdvanced, SaveGaussianSplat, SavePointCloud]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> Save3DExtension:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
|
@ -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."),
|
||||
|
|
|
|||
|
|
@ -1,3 +1,3 @@
|
|||
# This file is automatically generated by the build process when version is
|
||||
# updated in pyproject.toml.
|
||||
__version__ = "0.27.0"
|
||||
__version__ = "0.28.0"
|
||||
|
|
|
|||
21
execution.py
21
execution.py
|
|
@ -29,6 +29,7 @@ from comfy_execution.caching import (
|
|||
HierarchicalCache,
|
||||
LRUCache,
|
||||
RAMPressureCache,
|
||||
RAM_CACHE_LARGE_INTERMEDIATE,
|
||||
)
|
||||
from comfy_execution.graph import (
|
||||
DynamicPrompt,
|
||||
|
|
@ -425,12 +426,12 @@ def _is_intermediate_output(dynprompt, node_id):
|
|||
|
||||
|
||||
def _send_cached_ui(server, node_id, display_node_id, cached, prompt_id, ui_outputs):
|
||||
if cached.ui is not None:
|
||||
ui_outputs[node_id] = cached.ui
|
||||
if server.client_id is None:
|
||||
return
|
||||
cached_ui = cached.ui or {}
|
||||
server.send_sync("executed", { "node": node_id, "display_node": display_node_id, "output": cached_ui.get("output", None), "prompt_id": prompt_id }, server.client_id)
|
||||
if cached.ui is not None:
|
||||
ui_outputs[node_id] = cached.ui
|
||||
|
||||
async def execute(server, dynprompt, caches, current_item, extra_data, executed, prompt_id, execution_list, pending_subgraph_results, pending_async_nodes, ui_outputs):
|
||||
unique_id = current_item
|
||||
|
|
@ -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
|
||||
|
|
|
|||
8
nodes.py
8
nodes.py
|
|
@ -992,7 +992,7 @@ class CLIPLoader:
|
|||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "clip_name": (folder_paths.get_filename_list("text_encoders"), ),
|
||||
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2"], ),
|
||||
"type": (["stable_diffusion", "stable_cascade", "sd3", "stable_audio", "mochi", "ltxv", "pixart", "cosmos", "lumina2", "wan", "hidream", "chroma", "ace", "omnigen2", "qwen_image", "hunyuan_image", "flux2", "ovis", "longcat_image", "cogvideox", "lens", "pixeldit", "ideogram4", "boogu", "krea2", "joyimage"], ),
|
||||
},
|
||||
"optional": {
|
||||
"device": (["default", "cpu"], {"advanced": True}),
|
||||
|
|
@ -1002,7 +1002,7 @@ class CLIPLoader:
|
|||
|
||||
CATEGORY = "model/loaders"
|
||||
|
||||
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
||||
DESCRIPTION = "Recipes:\nsd: clip-l\nstable cascade: clip-g\nsd3: t5 xxl / clip-g / clip-l\nstable audio: t5 base\nmochi: t5 xxl\ncogvideox: t5 xxl (226-token padding)\ncosmos: old t5 xxl\nlumina2: gemma 2 2B\nwan: umt5 xxl\nhidream: llama-3.1 (Recommend) or t5\nomnigen2: qwen vl 2.5 3B\njoyimage: qwen3-vl 8B\nlens: gpt-oss-20b\npixeldit: gemma 2 2B elm"
|
||||
|
||||
def load_clip(self, clip_name, type="stable_diffusion", device="default"):
|
||||
clip_type = getattr(comfy.sd.CLIPType, type.upper(), comfy.sd.CLIPType.STABLE_DIFFUSION)
|
||||
|
|
@ -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,8 +2459,10 @@ 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_joyimage.py",
|
||||
"nodes_boogu.py",
|
||||
"nodes_chroma_radiance.py",
|
||||
"nodes_pid.py",
|
||||
|
|
@ -2503,6 +2506,7 @@ async def init_builtin_extra_nodes():
|
|||
"nodes_triposplat.py",
|
||||
"nodes_depth_anything_3.py",
|
||||
"nodes_seed.py",
|
||||
"nodes_text.py",
|
||||
]
|
||||
|
||||
import_failed = []
|
||||
|
|
|
|||
57
openapi.yaml
57
openapi.yaml
|
|
@ -530,6 +530,10 @@ components:
|
|||
description: Job creation timestamp (Unix timestamp in milliseconds)
|
||||
format: int64
|
||||
type: integer
|
||||
execution_end_time:
|
||||
description: Workflow execution completion timestamp (Unix milliseconds, only present for terminal states)
|
||||
format: int64
|
||||
type: integer
|
||||
execution_error:
|
||||
allOf:
|
||||
- $ref: '#/components/schemas/ExecutionError'
|
||||
|
|
@ -538,6 +542,10 @@ components:
|
|||
additionalProperties: true
|
||||
description: Node-level execution metadata (only for terminal states)
|
||||
type: object
|
||||
execution_start_time:
|
||||
description: Workflow execution start timestamp (Unix milliseconds, only present once execution has started)
|
||||
format: int64
|
||||
type: integer
|
||||
execution_status:
|
||||
additionalProperties: true
|
||||
description: ComfyUI execution status and timeline (only for terminal states)
|
||||
|
|
@ -570,6 +578,12 @@ components:
|
|||
description: Last update timestamp (Unix timestamp in milliseconds)
|
||||
format: int64
|
||||
type: integer
|
||||
user_id:
|
||||
description: |
|
||||
ID of the user that owns this job (see the `workspace_id`
|
||||
description above for why this is always the caller's own id
|
||||
on a successful response).
|
||||
type: string
|
||||
workflow:
|
||||
additionalProperties: true
|
||||
description: |
|
||||
|
|
@ -583,6 +597,18 @@ components:
|
|||
workflow_id:
|
||||
description: UUID identifying the workflow graph definition
|
||||
type: string
|
||||
workspace_id:
|
||||
description: |
|
||||
ID of the workspace that owns this job. A successful (200)
|
||||
response from this operation is only ever returned for the
|
||||
caller's own job (see this operation's ownership-scoped
|
||||
query), so this is always the caller's own workspace —
|
||||
consumers that also need to correlate this job to its
|
||||
live-progress broadcast channel (workspace+user scoped; see
|
||||
the internal common/gateways/broadcast package) can use this
|
||||
value directly rather than resolving their own identity a
|
||||
second way.
|
||||
type: string
|
||||
required:
|
||||
- id
|
||||
- status
|
||||
|
|
@ -1565,7 +1591,13 @@ paths:
|
|||
schema:
|
||||
default: true
|
||||
type: boolean
|
||||
- description: Filter assets by exact content hash.
|
||||
- description: |
|
||||
Filter assets by content hash, in the canonical `blake3:<hex>`
|
||||
form. Matches regardless of which of this asset store's two
|
||||
internal hash storage formats the matching row was written
|
||||
under (the canonical form used by from-hash-created references,
|
||||
or the raw `<hex>.<ext>`/bare `<hex>` storage key used by direct
|
||||
uploads) — both represent the same content hash.
|
||||
in: query
|
||||
name: hash
|
||||
schema:
|
||||
|
|
@ -2464,6 +2496,23 @@ paths:
|
|||
schema:
|
||||
additionalProperties: true
|
||||
properties:
|
||||
free_tier_balance:
|
||||
description: Free-tier job allowance for an authenticated non-paid (FREE-tier) user in the rollout. Absent for paid users and unauthenticated requests. Synthesized from config before a grant row exists so a brand-new user still sees their full allowance.
|
||||
properties:
|
||||
allowance:
|
||||
description: Total free jobs granted for the current period
|
||||
type: integer
|
||||
remaining:
|
||||
description: Free jobs remaining (allowance - used, floored at 0)
|
||||
type: integer
|
||||
used:
|
||||
description: Free jobs consumed so far
|
||||
type: integer
|
||||
required:
|
||||
- allowance
|
||||
- used
|
||||
- remaining
|
||||
type: object
|
||||
max_upload_size:
|
||||
description: Maximum upload size in bytes
|
||||
type: integer
|
||||
|
|
@ -3297,6 +3346,12 @@ paths:
|
|||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Invalid request parameters
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unauthorized - Authentication required
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "ComfyUI"
|
||||
version = "0.27.0"
|
||||
version = "0.28.0"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
comfyui-frontend-package==1.45.20
|
||||
comfyui-workflow-templates==0.11.6
|
||||
comfyui-embedded-docs==0.5.7
|
||||
comfyui-frontend-package==1.45.21
|
||||
comfyui-workflow-templates==0.11.12
|
||||
comfyui-embedded-docs==0.5.8
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
|
|
@ -22,7 +22,7 @@ alembic
|
|||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.16
|
||||
comfy-kitchen==0.2.22
|
||||
comfy-aimdo==0.4.10
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
|
|
|
|||
12
server.py
12
server.py
|
|
@ -39,6 +39,7 @@ 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
|
||||
|
|
@ -46,6 +47,7 @@ from comfy_api.internal import _ComfyNodeInternal
|
|||
from app.assets.seeder import asset_seeder
|
||||
from app.assets.api.routes import register_assets_routes
|
||||
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
|
||||
|
|
@ -441,7 +443,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 +728,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):
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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",))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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}"]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -19,7 +19,8 @@ def test_seed_asset_removed_when_file_is_deleted(
|
|||
"""Asset without hash (seed) whose file disappears:
|
||||
after triggering sync_seed_assets, Asset + AssetInfo disappear.
|
||||
"""
|
||||
# Create a file directly under input/unit-tests/<case> so tags include "unit-tests"
|
||||
# Create a file directly under input/unit-tests/<case>. Backend tags only
|
||||
# classify the root; nested path components are not exposed as tags.
|
||||
case_dir = comfy_tmp_base_dir / root / "unit-tests" / "syncseed"
|
||||
case_dir.mkdir(parents=True, exist_ok=True)
|
||||
name = f"seed_{uuid.uuid4().hex[:8]}.bin"
|
||||
|
|
@ -32,7 +33,7 @@ def test_seed_asset_removed_when_file_is_deleted(
|
|||
# Verify it is visible via API and carries no hash (seed)
|
||||
r1 = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": "unit-tests,syncseed", "name_contains": name},
|
||||
params={"include_tags": root, "name_contains": name},
|
||||
timeout=120,
|
||||
)
|
||||
body1 = r1.json()
|
||||
|
|
@ -54,7 +55,7 @@ def test_seed_asset_removed_when_file_is_deleted(
|
|||
# It should disappear (AssetInfo and seed Asset gone)
|
||||
r2 = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": "unit-tests,syncseed", "name_contains": name},
|
||||
params={"include_tags": root, "name_contains": name},
|
||||
timeout=120,
|
||||
)
|
||||
body2 = r2.json()
|
||||
|
|
@ -132,7 +133,7 @@ def test_hashed_asset_two_asset_infos_both_get_missing(
|
|||
second_id = b2["id"]
|
||||
|
||||
# Remove the single underlying file
|
||||
p = comfy_tmp_base_dir / "input" / "unit-tests" / "multiinfo" / get_asset_filename(b2["asset_hash"], ".png")
|
||||
p = comfy_tmp_base_dir / "input" / get_asset_filename(created["asset_hash"], ".png")
|
||||
assert p.exists()
|
||||
p.unlink()
|
||||
|
||||
|
|
@ -250,8 +251,7 @@ def test_missing_tag_clears_on_fastpass_when_mtime_and_size_match(
|
|||
|
||||
a = asset_factory(name, [root, "unit-tests", scope], {}, data)
|
||||
aid = a["id"]
|
||||
base = comfy_tmp_base_dir / root / "unit-tests" / scope
|
||||
p = base / get_asset_filename(a["asset_hash"], ".bin")
|
||||
p = comfy_tmp_base_dir / root / get_asset_filename(a["asset_hash"], ".bin")
|
||||
st0 = p.stat()
|
||||
orig_mtime_ns = getattr(st0, "st_mtime_ns", int(st0.st_mtime * 1_000_000_000))
|
||||
|
||||
|
|
|
|||
|
|
@ -290,7 +290,7 @@ def test_metadata_filename_is_set_for_seed_asset_without_hash(
|
|||
|
||||
r1 = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": f"unit-tests,{scope}", "name_contains": name},
|
||||
params={"include_tags": root, "name_contains": name},
|
||||
timeout=120,
|
||||
)
|
||||
body = r1.json()
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ def test_download_svg_forced_to_attachment(http: requests.Session, api_base: str
|
|||
svg = b'<svg xmlns="http://www.w3.org/2000/svg"><script>alert(1)</script></svg>'
|
||||
files = {"file": ("evil.svg", svg, "image/svg+xml")}
|
||||
form_data = {
|
||||
"tags": json.dumps(["models", "checkpoints", "unit-tests", "svgxss"]),
|
||||
"tags": json.dumps(["models", "model_type:checkpoints", "unit-tests", "svgxss"]),
|
||||
"name": "evil.svg",
|
||||
}
|
||||
up = http.post(api_base + "/api/assets", files=files, data=form_data, timeout=120)
|
||||
|
|
@ -131,7 +131,7 @@ def test_download_chooses_existing_state_and_updates_access_time(
|
|||
assert t1 > t0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seeded_asset", [{"tags": ["models", "checkpoints"]}], indirect=True)
|
||||
@pytest.mark.parametrize("seeded_asset", [{"tags": ["models", "model_type:checkpoints"]}], indirect=True)
|
||||
def test_download_missing_file_returns_404(
|
||||
http: requests.Session, api_base: str, comfy_tmp_base_dir: Path, seeded_asset: dict
|
||||
):
|
||||
|
|
|
|||
|
|
@ -13,7 +13,7 @@ def _seed(asset_factory, make_asset_bytes, count: int, tag: str) -> list[str]:
|
|||
for n in names:
|
||||
asset_factory(
|
||||
n,
|
||||
["models", "checkpoints", "unit-tests", tag],
|
||||
["models", "model_type:checkpoints", "unit-tests", tag],
|
||||
{},
|
||||
make_asset_bytes(n, size=2048),
|
||||
)
|
||||
|
|
@ -208,7 +208,7 @@ def test_cursor_walks_for_non_name_sorts(sort_field, http: requests.Session, api
|
|||
names = []
|
||||
for i in range(4):
|
||||
n = f"cursor_{sort_field}_{i:02d}.safetensors"
|
||||
asset_factory(n, ["models", "checkpoints", "unit-tests", f"cursor-{sort_field}"], {}, make_asset_bytes(n, size=2048 + i))
|
||||
asset_factory(n, ["models", "model_type:checkpoints", "unit-tests", f"cursor-{sort_field}"], {}, make_asset_bytes(n, size=2048 + i))
|
||||
names.append(n)
|
||||
|
||||
params = {
|
||||
|
|
|
|||
|
|
@ -11,7 +11,7 @@ def test_list_assets_paging_and_sort(http: requests.Session, api_base: str, asse
|
|||
for n in names:
|
||||
asset_factory(
|
||||
n,
|
||||
["models", "checkpoints", "unit-tests", "paging"],
|
||||
["models", "model_type:checkpoints", "unit-tests", "paging"],
|
||||
{"epoch": 1},
|
||||
make_asset_bytes(n, size=2048),
|
||||
)
|
||||
|
|
@ -45,8 +45,8 @@ def test_list_assets_paging_and_sort(http: requests.Session, api_base: str, asse
|
|||
|
||||
|
||||
def test_list_assets_include_exclude_and_name_contains(http: requests.Session, api_base: str, asset_factory):
|
||||
a = asset_factory("inc_a.safetensors", ["models", "checkpoints", "unit-tests", "alpha"], {}, b"X" * 1024)
|
||||
b = asset_factory("inc_b.safetensors", ["models", "checkpoints", "unit-tests", "beta"], {}, b"Y" * 1024)
|
||||
a = asset_factory("inc_a.safetensors", ["models", "model_type:checkpoints", "unit-tests", "alpha"], {}, b"X" * 1024)
|
||||
b = asset_factory("inc_b.safetensors", ["models", "model_type:checkpoints", "unit-tests", "beta"], {}, b"Y" * 1024)
|
||||
|
||||
r = http.get(
|
||||
api_base + "/api/assets",
|
||||
|
|
@ -81,7 +81,7 @@ def test_list_assets_include_exclude_and_name_contains(http: requests.Session, a
|
|||
|
||||
|
||||
def test_list_assets_sort_by_size_both_orders(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-size"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-size"]
|
||||
n1, n2, n3 = "sz1.safetensors", "sz2.safetensors", "sz3.safetensors"
|
||||
asset_factory(n1, t, {}, make_asset_bytes(n1, 1024))
|
||||
asset_factory(n2, t, {}, make_asset_bytes(n2, 2048))
|
||||
|
|
@ -108,7 +108,7 @@ def test_list_assets_sort_by_size_both_orders(http, api_base, asset_factory, mak
|
|||
|
||||
|
||||
def test_list_assets_sort_by_updated_at_desc(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-upd"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-upd"]
|
||||
a1 = asset_factory("upd_a.safetensors", t, {}, make_asset_bytes("upd_a", 1200))
|
||||
a2 = asset_factory("upd_b.safetensors", t, {}, make_asset_bytes("upd_b", 1200))
|
||||
|
||||
|
|
@ -131,7 +131,7 @@ def test_list_assets_sort_by_updated_at_desc(http, api_base, asset_factory, make
|
|||
|
||||
|
||||
def test_list_assets_sort_by_last_access_time_desc(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-access"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-access"]
|
||||
asset_factory("acc_a.safetensors", t, {}, make_asset_bytes("acc_a", 1100))
|
||||
time.sleep(0.02)
|
||||
a2 = asset_factory("acc_b.safetensors", t, {}, make_asset_bytes("acc_b", 1100))
|
||||
|
|
@ -154,14 +154,14 @@ def test_list_assets_sort_by_last_access_time_desc(http, api_base, asset_factory
|
|||
|
||||
|
||||
def test_list_assets_include_tags_variants_and_case(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-include"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-include"]
|
||||
a = asset_factory("incvar_alpha.safetensors", [*t, "alpha"], {}, make_asset_bytes("iva"))
|
||||
asset_factory("incvar_beta.safetensors", [*t, "beta"], {}, make_asset_bytes("ivb"))
|
||||
|
||||
# CSV + case-insensitive
|
||||
# CSV tag filters are whitespace-trimmed and case-sensitive.
|
||||
r1 = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": "UNIT-TESTS,LF-INCLUDE,alpha"},
|
||||
params={"include_tags": "unit-tests,lf-include,alpha"},
|
||||
timeout=120,
|
||||
)
|
||||
b1 = r1.json()
|
||||
|
|
@ -196,14 +196,14 @@ def test_list_assets_include_tags_variants_and_case(http, api_base, asset_factor
|
|||
|
||||
|
||||
def test_list_assets_exclude_tags_dedup_and_case(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-exclude"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-exclude"]
|
||||
a = asset_factory("ex_a_alpha.safetensors", [*t, "alpha"], {}, make_asset_bytes("exa", 900))
|
||||
asset_factory("ex_b_beta.safetensors", [*t, "beta"], {}, make_asset_bytes("exb", 900))
|
||||
|
||||
# Exclude uppercase should work
|
||||
# Exclude filters are case-sensitive.
|
||||
r1 = http.get(
|
||||
api_base + "/api/assets",
|
||||
params={"include_tags": "unit-tests,lf-exclude", "exclude_tags": "BETA"},
|
||||
params={"include_tags": "unit-tests,lf-exclude", "exclude_tags": "beta"},
|
||||
timeout=120,
|
||||
)
|
||||
b1 = r1.json()
|
||||
|
|
@ -225,7 +225,7 @@ def test_list_assets_exclude_tags_dedup_and_case(http, api_base, asset_factory,
|
|||
|
||||
|
||||
def test_list_assets_name_contains_case_and_specials(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-name"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-name"]
|
||||
a1 = asset_factory("CaseMix.SAFE", t, {}, make_asset_bytes("cm", 800))
|
||||
a2 = asset_factory("case-other.safetensors", t, {}, make_asset_bytes("co", 800))
|
||||
|
||||
|
|
@ -261,7 +261,7 @@ def test_list_assets_name_contains_case_and_specials(http, api_base, asset_facto
|
|||
|
||||
|
||||
def test_list_assets_offset_beyond_total_and_limit_boundary(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "lf-pagelimits"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "lf-pagelimits"]
|
||||
asset_factory("pl1.safetensors", t, {}, make_asset_bytes("pl1", 600))
|
||||
asset_factory("pl2.safetensors", t, {}, make_asset_bytes("pl2", 600))
|
||||
asset_factory("pl3.safetensors", t, {}, make_asset_bytes("pl3", 600))
|
||||
|
|
@ -319,7 +319,7 @@ def test_list_assets_name_contains_literal_underscore(
|
|||
- foobar.safetensors (must NOT match)
|
||||
"""
|
||||
scope = f"lf-underscore-{uuid.uuid4().hex[:6]}"
|
||||
tags = ["models", "checkpoints", "unit-tests", scope]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", scope]
|
||||
|
||||
a = asset_factory("foo_bar.safetensors", tags, {}, make_asset_bytes("a", 700))
|
||||
b = asset_factory("fooxbar.safetensors", tags, {}, make_asset_bytes("b", 700))
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ def test_meta_and_across_keys_and_types(
|
|||
http, api_base: str, asset_factory, make_asset_bytes
|
||||
):
|
||||
name = "mf_and_mix.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-and"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-and"]
|
||||
meta = {"purpose": "mix", "epoch": 1, "active": True, "score": 1.23}
|
||||
asset_factory(name, tags, meta, make_asset_bytes(name, 4096))
|
||||
|
||||
|
|
@ -41,7 +41,7 @@ def test_meta_and_across_keys_and_types(
|
|||
|
||||
def test_meta_type_strictness_int_vs_str_and_bool(http, api_base, asset_factory, make_asset_bytes):
|
||||
name = "mf_types.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-types"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-types"]
|
||||
meta = {"epoch": 1, "active": True}
|
||||
asset_factory(name, tags, meta, make_asset_bytes(name))
|
||||
|
||||
|
|
@ -95,7 +95,7 @@ def test_meta_type_strictness_int_vs_str_and_bool(http, api_base, asset_factory,
|
|||
|
||||
def test_meta_any_of_list_of_scalars(http, api_base, asset_factory, make_asset_bytes):
|
||||
name = "mf_list_scalars.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-list"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-list"]
|
||||
meta = {"flags": ["red", "green"]}
|
||||
asset_factory(name, tags, meta, make_asset_bytes(name, 3000))
|
||||
|
||||
|
|
@ -134,7 +134,7 @@ def test_meta_none_semantics_missing_or_null_and_any_of_with_none(
|
|||
http, api_base, asset_factory, make_asset_bytes
|
||||
):
|
||||
# a1: key missing; a2: explicit null; a3: concrete value
|
||||
t = ["models", "checkpoints", "unit-tests", "mf-none"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "mf-none"]
|
||||
a1 = asset_factory("mf_none_missing.safetensors", t, {"x": 1}, make_asset_bytes("a1"))
|
||||
a2 = asset_factory("mf_none_null.safetensors", t, {"maybe": None}, make_asset_bytes("a2"))
|
||||
a3 = asset_factory("mf_none_value.safetensors", t, {"maybe": "x"}, make_asset_bytes("a3"))
|
||||
|
|
@ -166,7 +166,7 @@ def test_meta_none_semantics_missing_or_null_and_any_of_with_none(
|
|||
|
||||
def test_meta_nested_json_object_equality(http, api_base, asset_factory, make_asset_bytes):
|
||||
name = "mf_nested_json.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-nested"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-nested"]
|
||||
cfg = {"optimizer": "adam", "lr": 0.001, "schedule": {"type": "cosine", "warmup": 100}}
|
||||
asset_factory(name, tags, {"config": cfg}, make_asset_bytes(name, 2200))
|
||||
|
||||
|
|
@ -197,7 +197,7 @@ def test_meta_nested_json_object_equality(http, api_base, asset_factory, make_as
|
|||
|
||||
def test_meta_list_of_objects_any_of(http, api_base, asset_factory, make_asset_bytes):
|
||||
name = "mf_list_objects.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-objlist"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-objlist"]
|
||||
transforms = [{"type": "crop", "size": 128}, {"type": "flip", "p": 0.5}]
|
||||
asset_factory(name, tags, {"transforms": transforms}, make_asset_bytes(name, 2048))
|
||||
|
||||
|
|
@ -228,7 +228,7 @@ def test_meta_list_of_objects_any_of(http, api_base, asset_factory, make_asset_b
|
|||
|
||||
def test_meta_with_special_and_unicode_keys(http, api_base, asset_factory, make_asset_bytes):
|
||||
name = "mf_keys_unicode.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-keys"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-keys"]
|
||||
meta = {
|
||||
"weird.key": "v1",
|
||||
"path/like": 7,
|
||||
|
|
@ -259,7 +259,7 @@ def test_meta_with_special_and_unicode_keys(http, api_base, asset_factory, make_
|
|||
|
||||
|
||||
def test_meta_with_zero_and_boolean_lists(http, api_base, asset_factory, make_asset_bytes):
|
||||
t = ["models", "checkpoints", "unit-tests", "mf-zero-bool"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "mf-zero-bool"]
|
||||
a0 = asset_factory("mf_zero_count.safetensors", t, {"count": 0}, make_asset_bytes("z", 1025))
|
||||
a1 = asset_factory("mf_bool_list.safetensors", t, {"choices": [True, False]}, make_asset_bytes("b", 1026))
|
||||
|
||||
|
|
@ -286,7 +286,7 @@ def test_meta_with_zero_and_boolean_lists(http, api_base, asset_factory, make_as
|
|||
|
||||
def test_meta_mixed_list_types_and_strictness(http, api_base, asset_factory, make_asset_bytes):
|
||||
name = "mf_mixed_list.safetensors"
|
||||
tags = ["models", "checkpoints", "unit-tests", "mf-mixed"]
|
||||
tags = ["models", "model_type:checkpoints", "unit-tests", "mf-mixed"]
|
||||
meta = {"mix": ["1", 1, True, None]}
|
||||
asset_factory(name, tags, meta, make_asset_bytes(name, 1999))
|
||||
|
||||
|
|
@ -311,7 +311,7 @@ def test_meta_mixed_list_types_and_strictness(http, api_base, asset_factory, mak
|
|||
|
||||
def test_meta_unknown_key_and_none_behavior_with_scope_tags(http, api_base, asset_factory, make_asset_bytes):
|
||||
# Use a unique scope tag to avoid interference
|
||||
t = ["models", "checkpoints", "unit-tests", "mf-unknown-scope"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "mf-unknown-scope"]
|
||||
x = asset_factory("mf_unknown_a.safetensors", t, {"k1": 1}, make_asset_bytes("ua"))
|
||||
y = asset_factory("mf_unknown_b.safetensors", t, {"k2": 2}, make_asset_bytes("ub"))
|
||||
|
||||
|
|
@ -340,13 +340,13 @@ def test_meta_with_tags_include_exclude_and_name_contains(http, api_base, asset_
|
|||
# alpha matches epoch=1; beta has epoch=2
|
||||
a = asset_factory(
|
||||
"mf_tag_alpha.safetensors",
|
||||
["models", "checkpoints", "unit-tests", "mf-tag", "alpha"],
|
||||
["models", "model_type:checkpoints", "unit-tests", "mf-tag", "alpha"],
|
||||
{"epoch": 1},
|
||||
make_asset_bytes("alpha"),
|
||||
)
|
||||
b = asset_factory(
|
||||
"mf_tag_beta.safetensors",
|
||||
["models", "checkpoints", "unit-tests", "mf-tag", "beta"],
|
||||
["models", "model_type:checkpoints", "unit-tests", "mf-tag", "beta"],
|
||||
{"epoch": 2},
|
||||
make_asset_bytes("beta"),
|
||||
)
|
||||
|
|
@ -367,7 +367,7 @@ def test_meta_with_tags_include_exclude_and_name_contains(http, api_base, asset_
|
|||
|
||||
def test_meta_sort_and_paging_under_filter(http, api_base, asset_factory, make_asset_bytes):
|
||||
# Three assets in same scope with different sizes and a common filter key
|
||||
t = ["models", "checkpoints", "unit-tests", "mf-sort"]
|
||||
t = ["models", "model_type:checkpoints", "unit-tests", "mf-sort"]
|
||||
n1, n2, n3 = "mf_sort_1.safetensors", "mf_sort_2.safetensors", "mf_sort_3.safetensors"
|
||||
asset_factory(n1, t, {"group": "g"}, make_asset_bytes(n1, 1024))
|
||||
asset_factory(n2, t, {"group": "g"}, make_asset_bytes(n2, 2048))
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in New Issue