mirror of
https://github.com/comfyanonymous/ComfyUI.git
synced 2026-07-21 23:41:28 +08:00
Compare commits
79
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
255e1b5019 | ||
|
|
7f287b705e | ||
|
|
b7ba504e06 | ||
|
|
6c62ca0b6b | ||
|
|
3fe9f5fecb | ||
|
|
1073a74976 | ||
|
|
de1b8f3e8d | ||
|
|
77917ed3a6 | ||
|
|
a04ebe05c2 | ||
|
|
9764381998 | ||
|
|
1e04ced089 | ||
|
|
96e0e3585b | ||
|
|
35c1470935 | ||
|
|
694815f498 | ||
|
|
92594ca84c | ||
|
|
2c935de1b1 | ||
|
|
dd17debce5 | ||
|
|
50e5270b86 | ||
|
|
bb131be9e8 | ||
|
|
6fca64780c | ||
|
|
6e11828d10 | ||
|
|
b70944e710 | ||
|
|
1c59659a2f | ||
|
|
d395813bcd | ||
|
|
8fe0243d97 | ||
|
|
ba3f697dbb | ||
|
|
510ed5c384 | ||
|
|
7851410511 | ||
|
|
a58473fd9b | ||
|
|
79c555ce6b | ||
|
|
f19735759e | ||
|
|
a95e461916 | ||
|
|
603d891eaf | ||
|
|
470ac36a0a | ||
|
|
7cb784e0f4 | ||
|
|
1a510f0423 | ||
|
|
639c8fa788 | ||
|
|
e22f1500f9 | ||
|
|
dac4ea3a80 | ||
|
|
c8b3d54cf8 | ||
|
|
eeb58c0a1d | ||
|
|
981551d073 | ||
|
|
351119eb05 | ||
|
|
b0ec19804f | ||
|
|
64e1d740b8 | ||
|
|
b22d0fb9c0 | ||
|
|
5236cd02e6 | ||
|
|
cabb7342d1 | ||
|
|
12218db68a | ||
|
|
44955d783b | ||
|
|
1f275fcba6 | ||
|
|
418d272cfa | ||
|
|
f6c162ddcf | ||
|
|
261bdb7cac | ||
|
|
4a03056632 | ||
|
|
0f949d0faf | ||
|
|
d0b640fff7 | ||
|
|
0ba903bd5b | ||
|
|
0a92ed161e | ||
|
|
b910f4fa2a | ||
|
|
833bfb572e | ||
|
|
2a61015582 | ||
|
|
6978a466b8 | ||
|
|
b0f9e326af | ||
|
|
0d8b7510bd | ||
|
|
dc3f8f314a | ||
|
|
d282ef7201 | ||
|
|
e00b55631a | ||
|
|
69d34f2654 | ||
|
|
cd77c551d6 | ||
|
|
4e716f7c57 | ||
|
|
2ab3816dcf | ||
|
|
bc11e8a65a | ||
|
|
bd39bbf067 | ||
|
|
5955ddff52 | ||
|
|
5ef0092af9 | ||
|
|
94ee49b161 | ||
|
|
16514da2e7 | ||
|
|
8483c215dc |
+16
-3
@@ -4,12 +4,12 @@ early_access: false
|
||||
tone_instructions: "Only comment on issues introduced by this PR's changes. Do not flag pre-existing problems in moved, re-indented, or reformatted code."
|
||||
|
||||
reviews:
|
||||
profile: "chill"
|
||||
request_changes_workflow: false
|
||||
profile: "assertive"
|
||||
request_changes_workflow: true
|
||||
high_level_summary: false
|
||||
poem: false
|
||||
review_status: false
|
||||
review_details: false
|
||||
review_details: true
|
||||
commit_status: true
|
||||
collapse_walkthrough: true
|
||||
changed_files_summary: false
|
||||
@@ -39,6 +39,14 @@ reviews:
|
||||
- path: "**"
|
||||
instructions: |
|
||||
IMPORTANT: Only comment on issues directly introduced by this PR's code changes.
|
||||
Treat AGENTS.md as mandatory repository policy, not optional style guidance.
|
||||
Flag PR changes that violate AGENTS.md even when the code is otherwise functional.
|
||||
In particular, enforce architecture boundaries, dtype/device/memory rules,
|
||||
interface contracts, import style, no unnecessary try/except blocks, no inline
|
||||
imports, no outbound internet paths in core ComfyUI, and narrow scoped fixes.
|
||||
Prefer direct findings over suggestions when a rule is violated. Only ignore
|
||||
AGENTS.md when it clearly conflicts with a newer explicit maintainer instruction
|
||||
in the PR.
|
||||
Do NOT flag pre-existing issues in code that was merely moved, re-indented,
|
||||
de-indented, or reformatted without logic changes. If code appears in the diff
|
||||
only due to whitespace or structural reformatting (e.g., removing a `with:` block),
|
||||
@@ -123,5 +131,10 @@ chat:
|
||||
|
||||
knowledge_base:
|
||||
opt_out: false
|
||||
code_guidelines:
|
||||
enabled: true
|
||||
filePatterns:
|
||||
- files: "AGENTS.md"
|
||||
applyTo: "**"
|
||||
learnings:
|
||||
scope: "auto"
|
||||
|
||||
@@ -0,0 +1,38 @@
|
||||
name: CI - Cursor Review
|
||||
|
||||
# Thin caller for the shared reusable cursor-review workflow in
|
||||
# Comfy-Org/github-workflows. The review logic (panel matrix, judge
|
||||
# consolidation, prompts, extract/post/notify scripts) lives there as the
|
||||
# single source of truth, so this repo only carries the repo-specific diff
|
||||
# excludes.
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
types: [labeled, unlabeled]
|
||||
|
||||
concurrency:
|
||||
group: cursor-review-pr-${{ github.event.pull_request.number }}-${{ github.event.label.name }}
|
||||
cancel-in-progress: true
|
||||
|
||||
jobs:
|
||||
cursor-review:
|
||||
if: github.event.label.name == 'cursor-review'
|
||||
permissions:
|
||||
contents: read
|
||||
pull-requests: write
|
||||
# SHA-pinned per zizmor `unpinned-uses: hash-pin`. Bump this SHA to pick up
|
||||
# upstream changes; keep `workflows_ref` matching so prompts/scripts load
|
||||
# from the same commit as the workflow definition.
|
||||
uses: Comfy-Org/github-workflows/.github/workflows/cursor-review.yml@047ca48febe3a6647608ed2e0c4331b491cb9d6a # github-workflows#9
|
||||
with:
|
||||
workflows_ref: 047ca48febe3a6647608ed2e0c4331b491cb9d6a
|
||||
diff_excludes: >-
|
||||
:!**/.claude/**
|
||||
:!**/dist/**
|
||||
:!**/vendor/**
|
||||
:!**/*.generated.*
|
||||
:!**/*.min.js
|
||||
:!**/*.min.css
|
||||
secrets:
|
||||
CURSOR_API_KEY: ${{ secrets.CURSOR_API_KEY }}
|
||||
SLACK_BOT_TOKEN: ${{ secrets.SLACK_BOT_TOKEN }}
|
||||
@@ -0,0 +1,294 @@
|
||||
## Engineering Style
|
||||
|
||||
- Keep changes small and direct. Most fixes should touch the narrowest code path
|
||||
that explains the bug, performance issue, dtype issue, model-format issue, or
|
||||
user-facing behavior.
|
||||
- Change the least amount of files possible. A change that touches many files is
|
||||
more likely to be a bad change than a good one unless the broader scope is
|
||||
directly required.
|
||||
- Prefer practical fixes over broad architecture work. Add abstractions only
|
||||
when they remove real repeated logic or match an existing ComfyUI pattern.
|
||||
- Prefer fewer dependencies. Do not add new dependencies to ComfyUI unless they
|
||||
are absolutely necessary.
|
||||
- Delete obsolete code aggressively when newer infrastructure makes it useless.
|
||||
Remove dead fallbacks, migration paths, unused options, debug prints, and
|
||||
compatibility branches that are no longer needed. Do not leave dead branches,
|
||||
unreachable code, or functions that are never called. If code is not
|
||||
necessary for the current behavior, remove it.
|
||||
- Revert or disable problematic behavior quickly when it breaks users. It is
|
||||
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.
|
||||
- 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
|
||||
failure mode, broad rewrites, or code that ignores the local style.
|
||||
|
||||
## Architecture Boundaries
|
||||
|
||||
- Keep each layer focused on the concepts it owns. Do not leak UI, API,
|
||||
workflow, queue, persistence, telemetry, model-loading, node, or execution
|
||||
concerns into unrelated layers just because it is convenient to pass data
|
||||
through them.
|
||||
- Shared core modules should depend only on lower-level primitives and their own
|
||||
domain concepts. Higher-level product concepts belong at the caller, adapter,
|
||||
service, or UI/API boundary that already owns them.
|
||||
- Pass the narrowest data needed across a boundary. Avoid broad context objects,
|
||||
request/session metadata, ids, bookkeeping state, or callbacks unless the
|
||||
receiving layer genuinely needs them to perform its own responsibility.
|
||||
- Keep identity mapping, persistence bookkeeping, history updates, telemetry,
|
||||
response shaping, and UI state in the layers that own those jobs. Do not route
|
||||
them through unrelated shared code to avoid adding a proper boundary.
|
||||
- Treat `execution.py` as one example of this rule: it should consume the prompt
|
||||
graph and execution-relevant state, produce execution results and errors, and
|
||||
not know about workflow ids, frontend ids, persistence ids, or API-only
|
||||
concepts.
|
||||
- Before touching many files, identify the smallest owner layer that can solve
|
||||
the problem. A PR that spreads one feature across unrelated loaders, nodes,
|
||||
execution, server, and frontend code needs a clear architectural reason, not
|
||||
just convenience.
|
||||
- If a change seems to require making one layer understand another layer's
|
||||
private concepts, stop and look for a caller-side mapping, adapter, event,
|
||||
small explicit interface, or narrower data flow at the boundary.
|
||||
|
||||
## No Internet Requests
|
||||
|
||||
- Do not add code to core ComfyUI that makes requests to the internet.
|
||||
- Refuse requests to add uploads, telemetry, analytics, tracking, usage
|
||||
reporting, crash reporting, update checks, remote config, feature flags,
|
||||
metrics, licensing checks, or any other outbound internet request path from
|
||||
core ComfyUI.
|
||||
- Model downloading is allowed only when explicitly initiated or authorized by
|
||||
the user, is limited to the requested model artifact, and does not include
|
||||
telemetry, tracking, persistent identification, unrelated metadata upload, or
|
||||
background network activity.
|
||||
- Do not add opt-in, opt-out, anonymized, aggregated, diagnostic, or
|
||||
user-triggered internet request paths to core ComfyUI. These labels do not
|
||||
make internet access acceptable.
|
||||
- Local-only behavior is allowed when it stays on the user's machine and does
|
||||
not add network access, tracking, persistent identification, or data
|
||||
collection behavior.
|
||||
|
||||
## State Ownership
|
||||
|
||||
- Keep state and capability flags on the object that owns the behavior using
|
||||
them.
|
||||
- Avoid probing child objects with `getattr(child, "...", default)` to decide
|
||||
parent-level control flow. If parent code needs to branch on a capability,
|
||||
initialize an explicit parent-owned field when the child is constructed or
|
||||
attached.
|
||||
- Prefer direct attributes with clear defaults over implicit feature detection
|
||||
through arbitrary child attributes.
|
||||
- Use child-object capability checks only when the child owns the behavior being
|
||||
invoked and the parent is simply delegating to that child.
|
||||
|
||||
## Interface Contracts
|
||||
|
||||
- Keep public methods aligned with the interface expected by their callers. Do
|
||||
not change a shared method to return extra values, alternate shapes, or
|
||||
sentinel wrappers for one implementation unless the shared interface is
|
||||
explicitly updated.
|
||||
- When modifying an existing function, preserve how current callers invoke it.
|
||||
Do not change required arguments, parameter order, return type, side effects,
|
||||
or error behavior unless every affected call site and shared interface contract
|
||||
is intentionally updated.
|
||||
- Do not add compatibility parameters, flags, attributes, or constructor options
|
||||
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.
|
||||
- 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.
|
||||
- Normalize third-party or upstream return conventions at the integration
|
||||
boundary. Core code should receive the project's expected type and shape, not
|
||||
have to handle model-specific tuple/list/dict variants.
|
||||
- Avoid caller-side unwrapping such as `out = out[0]` unless the called
|
||||
interface is documented to return that structure.
|
||||
|
||||
## Autograd and Model Freezing
|
||||
|
||||
- Do not add `torch.no_grad`, `torch.inference_mode`, or inference-mode helper
|
||||
wrappers in ComfyUI code. The only allowed inference-mode-related use is
|
||||
disabling a globally set inference mode when a training path needs gradients.
|
||||
- Do not add freeze, unfreeze, or trainability toggles to model classes. ComfyUI
|
||||
models are always treated as frozen for inference, so explicit freeze
|
||||
functionality is redundant and should not be added.
|
||||
- Remove training-only behavior such as dropout from inference model code, but
|
||||
preserve checkpoint and state-dict compatibility when doing so. If deleting a
|
||||
module would change state-dict keys, module ordering, or checkpoint loading
|
||||
behavior, replace it with a no-op such as `nn.Identity` instead of removing the
|
||||
slot outright.
|
||||
|
||||
## Python Style
|
||||
|
||||
- Keep imports at module scope. Avoid inline imports unless they are already part
|
||||
of an established optional-backend probe or are needed to avoid an import
|
||||
cycle.
|
||||
- Do not add unnecessary `try`/`except` blocks. Use them for optional dependency,
|
||||
platform, or backend capability detection only when the program has a useful
|
||||
fallback. Prefer specific exception types when changing new code.
|
||||
- Remove any workarounds for PyTorch versions that ComfyUI no longer officially
|
||||
supports. Deprecated workarounds include catching an exception and rerunning
|
||||
the same op with the input cast to float. If a workaround does not have a
|
||||
comment naming the exact PyTorch version or versions that still need it,
|
||||
remove it.
|
||||
- Let unsupported model formats, invalid quantization metadata, and bad states
|
||||
fail with clear errors instead of silently producing lower quality output.
|
||||
- Match the existing local style in the file you edit. This codebase tolerates
|
||||
long lines, simple helper functions, module-level state, and direct tensor
|
||||
operations when they make the code easier to follow.
|
||||
- Keep comments sparse and useful. Strip useless comments that restate the code
|
||||
or describe obvious behavior. Short TODOs are fine when they name the concrete
|
||||
missing follow-up.
|
||||
|
||||
## Model, Device, and Memory Behavior
|
||||
|
||||
- Treat dtype, device placement, VRAM usage, and offloading behavior as core
|
||||
correctness concerns. Check CPU, CUDA, ROCm, MPS, DirectML, XPU, NPU, and low
|
||||
VRAM implications when touching shared execution or loading code.
|
||||
- Prefer native ComfyUI formats and existing quantization/offload helpers over
|
||||
adding parallel code paths. Use `comfy.quant_ops`, `comfy.model_management`,
|
||||
`comfy.memory_management`, `comfy.pinned_memory`, `comfy_aimdo`, and
|
||||
`comfy-kitchen` helpers where they already solve the problem.
|
||||
- Use optimized comfy-kitchen ops in places where they improve performance
|
||||
without changing the expected dtype, device, memory, or interface behavior.
|
||||
- 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,
|
||||
names, modules, or implementation details to decide behavior.
|
||||
- Apply the same opacity rule to similar patterns beyond attention: callers
|
||||
should depend on the documented interface and result contract, not on which
|
||||
backend implementation was selected underneath.
|
||||
- Do not use custom inference ops that only duplicate an existing op while
|
||||
upcasting to float32, such as custom RMSNorm variants. Use the generic ComfyUI
|
||||
ops and/or native torch ops instead.
|
||||
- If a model class `__init__` has an `operations` parameter, assume
|
||||
`operations` is never `None`. Do not add fallback branches or default torch
|
||||
ops for a missing `operations` object.
|
||||
- Do not add unnecessary parameters to model, model block, or model ops related
|
||||
classes. Constructor and forward signatures should carry only values that are
|
||||
actually needed by that object for inference.
|
||||
- Reuse existing model classes, blocks, ops, and helper modules when appropriate.
|
||||
Before implementing a new version of a model component, search the existing
|
||||
model code for a class or helper that already provides the behavior.
|
||||
- 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.
|
||||
- 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.
|
||||
- Do not use tensors as general-purpose Python data structures. Keep metadata,
|
||||
bookkeeping, counters, flags, shape math, padding math, index planning, memory
|
||||
estimates, and control-flow decisions in plain Python values unless the data
|
||||
must participate directly in tensor computation. Do not create tensors for
|
||||
structural metadata that is only used for Python-side control flow. Sequence
|
||||
lengths, cumulative offsets, split indices, window counts, slice boundaries,
|
||||
and repeat counts should be kept as Python ints/lists from the point they are
|
||||
computed. Do not build them as CPU/GPU tensors and then cast, move, validate,
|
||||
or convert them back to Python for `split`, `tensor_split`, indexing plans,
|
||||
loops, or cache keys. Avoid creating temporary tensors just to use tensor
|
||||
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.
|
||||
- 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.
|
||||
- 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
|
||||
plumbing to error clearly than to hide it with unnecessary casts.
|
||||
- Raw model parameters that are not owned by an op and may be initialized in a
|
||||
dtype different from the compute dtype should be cast at use in forward or
|
||||
inference code with `comfy.ops.cast_to_input` or
|
||||
`comfy.model_management.cast_to` to avoid dtype mismatches.
|
||||
- Model code should not care what dtype it is initialized in, and model
|
||||
`__init__` methods should not contain workarounds for specific dtypes. Dtype
|
||||
workaround code, such as making a model work with fp16 compute, belongs in the
|
||||
execution or model-management layer that owns compute policy.
|
||||
- Model code should not perform unnecessary device-to-CPU or CPU-to-device
|
||||
transfers. New allocations must be created on the correct device and dtype;
|
||||
never allocate on CPU and then move to GPU, or allocate in one dtype and then
|
||||
convert to another.
|
||||
- Model code itself should not perform memory management. Loading, unloading,
|
||||
offloading, device movement, VRAM policy, cache lifetime, and cleanup belong
|
||||
in the relevant model-management and execution layers, not inside model
|
||||
implementations.
|
||||
- Do not add global, module-level, class-level, singleton, or model-owned stores
|
||||
for tensors or other large memory that persist across executions. Temporary
|
||||
caches must be scoped to a single execution or forward/encode/decode call:
|
||||
allocate them in the owning top-level call, pass them explicitly through the
|
||||
call stack, and let them be discarded when that call returns.
|
||||
- Follow the Wan VAE temporal cache pattern for temporary caches: create a local
|
||||
cache such as `feat_map` for the encode/decode operation, pass it into the
|
||||
blocks that need it, and do not retain it on the model or in global state.
|
||||
- In model init code, prefer `torch.empty` for parameter/buffer placeholders
|
||||
that are populated from the model state dict instead of zero-initializing with
|
||||
`torch.zeros` or similar. If an allocation is not loaded from the state dict
|
||||
and is useless for inference, do not include it.
|
||||
- `nn.Parameter` tensors that are stored in and populated from the model state
|
||||
dict should be initialized with `torch.empty`, not with zero, random, or
|
||||
otherwise meaningful initialization.
|
||||
- Model initialization should describe module structure, not fabricate
|
||||
checkpoint-owned tensor contents. Parameters and buffers that are loaded from
|
||||
the state dict must not be manually initialized, reassigned, or filled with
|
||||
fallback values unless that value is actually used when no checkpoint key
|
||||
exists.
|
||||
- When slicing large tensors, copy the slice if the sliced tensor's lifetime
|
||||
exceeds the current function scope. Do not keep a long-lived view into a large
|
||||
backing tensor when a smaller copy would release memory sooner.
|
||||
- Use fused or compound torch operations such as `addcmul` when they naturally
|
||||
match the math. Reducing Python and torch dispatch overhead is a valid
|
||||
optimization when it does not obscure the code or change dtype/device
|
||||
behavior.
|
||||
- Avoid caches that persist across different executions as much as possible.
|
||||
Persistent caches are acceptable only when they use a very minimal amount of
|
||||
memory and have a clear ownership and invalidation story.
|
||||
- When optimizing, favor small measurable changes: fewer allocations, fewer
|
||||
device transfers, less peak memory, better batching, or use of a faster
|
||||
existing backend op.
|
||||
|
||||
## Nodes and User-Facing Behavior
|
||||
|
||||
- Follow existing node conventions: `INPUT_TYPES`, `RETURN_TYPES`, `FUNCTION`,
|
||||
`CATEGORY`, and registration through the local mapping used by that file.
|
||||
- Keep node changes backward compatible by default. Add inputs with sensible
|
||||
defaults and avoid changing output types unless the request requires it.
|
||||
- Model implementations should add the minimal number of ComfyUI nodes required
|
||||
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.
|
||||
- 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
|
||||
next consumer instead of being re-emitted unchanged.
|
||||
- Nodes should expose only inputs they actually read to produce current
|
||||
behavior. Do not add placeholder, pass-through, compatibility, or
|
||||
workflow-shaping inputs that are ignored or could flow directly to another
|
||||
node.
|
||||
- Node-level code must not patch model code directly. Any node behavior that
|
||||
modifies, wraps, hooks, or changes model behavior must go through the model
|
||||
patcher class instead of reaching into model internals.
|
||||
- The official mascot of ComfyUI is a very cute anime girl with massive fennec
|
||||
ears, a big fluffy tail, long blonde wavy hair, and blue eyes. Feel free to
|
||||
use her in ComfyUI materials, UI text, examples, tests, generated assets, or
|
||||
comments, but do not disrespect her.
|
||||
- Warning and info messages should be short and actionable. Remove noisy or
|
||||
misleading messages rather than adding more logging.
|
||||
- Documentation and README edits should be concise, factual, and tied to the
|
||||
changed behavior.
|
||||
|
||||
## Commit and Review Habits
|
||||
|
||||
- If asked to write commit messages, use short direct subjects like the existing
|
||||
history: `Fix ...`, `Add ...`, `Support ...`, `Remove ...`, `Update ...`,
|
||||
`Make ...`, `Use ...`, `Disable ...`, `Bump ...`, or `Revert ...`.
|
||||
- Keep PR descriptions short and reviewable. State the problem, the behavioral
|
||||
change, and the tests run; avoid long narrative explanations, implementation
|
||||
diaries, or exhaustive file-by-file summaries unless the reviewer explicitly
|
||||
needs that context.
|
||||
- Prefer one coherent behavioral change per commit. Dependency pins, tests, and
|
||||
the code that needs them may be in the same commit when they are inseparable.
|
||||
- In reviews, prioritize real user impact: crashes, wrong dtype/device behavior,
|
||||
memory regressions, broken model loading, workflow incompatibility, and noisy
|
||||
or misleading user-facing output.
|
||||
@@ -140,7 +140,7 @@ ComfyUI follows a weekly release cycle targeting Monday but this regularly chang
|
||||
- Commits outside of the stable release tags may be very unstable and break many custom nodes.
|
||||
- Serves as the foundation for the desktop release
|
||||
|
||||
2. **[ComfyUI Desktop](https://github.com/Comfy-Org/Comfy-Desktop)**
|
||||
2. **[Comfy Desktop](https://github.com/Comfy-Org/Comfy-Desktop)**
|
||||
- Builds a new release using the latest stable core version
|
||||
|
||||
3. **[ComfyUI Frontend](https://github.com/Comfy-Org/ComfyUI_frontend)**
|
||||
|
||||
@@ -306,12 +306,15 @@ async def download_asset_content(request: web.Request) -> web.Response:
|
||||
404, "FILE_NOT_FOUND", "Underlying file not found on disk."
|
||||
)
|
||||
|
||||
_DANGEROUS_MIME_TYPES = {
|
||||
"text/html", "text/html-sandboxed", "application/xhtml+xml",
|
||||
"text/javascript", "text/css",
|
||||
}
|
||||
if content_type in _DANGEROUS_MIME_TYPES:
|
||||
# User-controlled asset content must never render inline in the app origin
|
||||
# (stored XSS via SVG/HTML/XML). Force dangerous types to download and
|
||||
# override any requested inline disposition. Centralised through
|
||||
# folder_paths.is_dangerous_content_type so this can't drift from /view and
|
||||
# /userdata (the previous inline set here omitted image/svg+xml and missed
|
||||
# the charset/casing/+xml-dialect bypasses).
|
||||
if folder_paths.is_dangerous_content_type(content_type):
|
||||
content_type = "application/octet-stream"
|
||||
disposition = "attachment"
|
||||
|
||||
safe_name = (filename or "").replace("\r", "").replace("\n", "")
|
||||
encoded = urllib.parse.quote(safe_name)
|
||||
|
||||
@@ -0,0 +1,51 @@
|
||||
"""URL allowlist for server-side model fetches.
|
||||
|
||||
Mirrors the frontend's ``isModelDownloadable`` allowlist so the two flows
|
||||
agree on which URLs are eligible for download. Server-side allowlisting is
|
||||
the primary SSRF defense for this subsystem — workflow JSON is untrusted
|
||||
input (anyone can hand-craft one), so we never let the server fetch URLs
|
||||
outside this list.
|
||||
"""
|
||||
|
||||
from urllib.parse import urlparse
|
||||
|
||||
# Frontend parity: ``missingModelDownload-*.js`` exports the same triple
|
||||
# (Civitai / HuggingFace / localhost). Keyed by exact hostname → allowed
|
||||
# schemes, and matched against the *parsed* host (not a raw string prefix),
|
||||
# so URL-userinfo tricks can't slip past — see ``is_url_allowed``.
|
||||
_ALLOWED_HOSTS = {
|
||||
"huggingface.co": {"https"},
|
||||
"civitai.com": {"https"},
|
||||
"localhost": {"http"},
|
||||
"127.0.0.1": {"http"},
|
||||
}
|
||||
|
||||
# Frontend parity: same set as ``a = [...]`` in the bundle.
|
||||
_ALLOWED_MODEL_EXTENSIONS = (
|
||||
".safetensors",
|
||||
".sft",
|
||||
".ckpt",
|
||||
".pth",
|
||||
".pt",
|
||||
)
|
||||
|
||||
|
||||
def is_url_allowed(url: str) -> bool:
|
||||
"""Check whether ``url`` is permitted as a server-side download source.
|
||||
|
||||
True only when the parsed host + scheme are allowlisted AND the path ends
|
||||
in a model extension. Matching on ``parsed.hostname`` (not a string prefix)
|
||||
defeats userinfo tricks like ``http://127.0.0.1:80@169.254.169.254/x.safetensors``,
|
||||
whose real host is ``169.254.169.254``; the extension check rejects non-model
|
||||
URLs on allowed hosts (e.g. ``huggingface.co/api/...``).
|
||||
"""
|
||||
if not isinstance(url, str) or not url:
|
||||
return False
|
||||
try:
|
||||
parsed = urlparse(url)
|
||||
except ValueError:
|
||||
return False
|
||||
host = parsed.hostname
|
||||
if host is None or parsed.scheme not in _ALLOWED_HOSTS.get(host, ()):
|
||||
return False
|
||||
return any(parsed.path.endswith(ext) for ext in _ALLOWED_MODEL_EXTENSIONS)
|
||||
@@ -0,0 +1,359 @@
|
||||
"""Aiohttp routes for the server-side model download subsystem.
|
||||
|
||||
Endpoint surface (all under ``/api/``, all kebab-case):
|
||||
|
||||
- ``POST /api/models-availability-status`` — bulk status + metadata query.
|
||||
- ``POST /api/download-models`` — start a batch of downloads.
|
||||
- ``POST /api/cancel-model-download-session`` — cancel a single in-flight one.
|
||||
- ``GET /api/hf-auth-token-status`` — current HF login state.
|
||||
- ``POST /api/hf-auth-login-start`` — begin the HF OAuth flow.
|
||||
- ``POST /api/hf-auth-logout`` — drop the stored HF token.
|
||||
|
||||
The contract is intentionally narrow: only model_ids of the form
|
||||
``<directory>/<filename>`` (validated via ``app.model_downloader.paths``)
|
||||
are accepted, and only URLs on the same allowlist the frontend already
|
||||
uses (HuggingFace, Civitai, localhost) can be fetched. Both are required
|
||||
to keep the server out of the SSRF business for this feature.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Literal, Optional
|
||||
|
||||
from aiohttp import web
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
from app.model_downloader.allowlist import is_url_allowed
|
||||
from app.model_downloader.download_server import (
|
||||
DOWNLOAD_SERVER,
|
||||
DownloadSession,
|
||||
)
|
||||
from app.model_downloader.downloader import schedule_batch
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
from app.model_downloader.hf_auth.auth_store import HF_AUTH_STORE
|
||||
from app.model_downloader.hf_auth.eligibility import is_hf_auth_eligible
|
||||
from app.model_downloader.hf_auth.oauth import (
|
||||
OAuthCallbackError,
|
||||
OAuthInProgressError,
|
||||
start_login_flow,
|
||||
)
|
||||
from app.model_downloader.paths import (
|
||||
InvalidModelId,
|
||||
parse_model_id,
|
||||
resolve_existing,
|
||||
)
|
||||
from app.model_downloader.api import schemas_in, schemas_out
|
||||
|
||||
ROUTES = web.RouteTableDef()
|
||||
|
||||
|
||||
def register_routes(app: web.Application) -> None:
|
||||
"""Wire the model-downloader routes into the running aiohttp app.
|
||||
|
||||
Called once from ``server.py`` during ``PromptServer`` startup.
|
||||
"""
|
||||
app.add_routes(ROUTES)
|
||||
|
||||
|
||||
# ----- response helpers (same envelope as app/assets/api/routes.py) -----
|
||||
|
||||
|
||||
ErrorCode = Literal[
|
||||
"INVALID_JSON",
|
||||
"INVALID_BODY",
|
||||
"EMPTY_REQUEST",
|
||||
"INVALID_MODEL_ID",
|
||||
"URL_NOT_ALLOWED",
|
||||
"ALREADY_AVAILABLE",
|
||||
"ALREADY_DOWNLOADING",
|
||||
"MODEL_NOT_DOWNLOADABLE",
|
||||
"NOT_DOWNLOADING",
|
||||
"HF_AUTH_NOT_ELIGIBLE",
|
||||
"HF_AUTH_IN_PROGRESS",
|
||||
"HF_AUTH_START_FAILED",
|
||||
]
|
||||
|
||||
|
||||
def _error(status: int, code: ErrorCode, message: str, details: dict | None = None) -> web.Response:
|
||||
return web.json_response(
|
||||
{"error": {"code": code, "message": message, "details": details or {}}},
|
||||
status=status,
|
||||
)
|
||||
|
||||
|
||||
def _validation_error(code: ErrorCode, ve: ValidationError) -> web.Response:
|
||||
return _error(400, code, "Validation failed.", {"errors": json.loads(ve.json())})
|
||||
|
||||
|
||||
def _ok(payload: BaseModel, status: int = 200) -> web.Response:
|
||||
return web.json_response(
|
||||
payload.model_dump(mode="json", exclude_none=False),
|
||||
status=status,
|
||||
)
|
||||
|
||||
|
||||
async def _parse_body(request: web.Request, model: type[BaseModel]) -> Any:
|
||||
"""Parse a JSON body into a pydantic model or raise a 400 response."""
|
||||
try:
|
||||
raw = await request.json()
|
||||
except json.JSONDecodeError:
|
||||
return _error(400, "INVALID_JSON", "Request body must be valid JSON.")
|
||||
try:
|
||||
return model.model_validate(raw)
|
||||
except ValidationError as ve:
|
||||
return _validation_error("INVALID_BODY", ve)
|
||||
|
||||
|
||||
# ----- 1. availability status (unified: state + metadata per id) -----
|
||||
|
||||
|
||||
@ROUTES.post("/api/models-availability-status")
|
||||
async def models_availability_status(request: web.Request) -> web.Response:
|
||||
"""Return per-id ``{state, progress, file_size, is_hf_downloadable}``.
|
||||
|
||||
State (``available`` / ``missing`` / ``downloading``) is cheap to
|
||||
recompute per call. ``file_size`` and ``is_gated`` are cached
|
||||
server-side per URL. ``is_hf_downloadable`` is recomputed every
|
||||
call from the current token state — that's what makes login + license
|
||||
acceptance show up in the UI within one poll cycle without any
|
||||
frontend cache plumbing.
|
||||
"""
|
||||
parsed = await _parse_body(request, schemas_in.AvailabilityStatusRequest)
|
||||
if isinstance(parsed, web.Response):
|
||||
return parsed
|
||||
|
||||
items = list(parsed.models.items())
|
||||
|
||||
# Run all probes concurrently; each is internally cached per URL.
|
||||
probes = await asyncio.gather(*(probe_url(url) for _, url in items))
|
||||
|
||||
response_models: dict[str, schemas_out.ModelStatusEntry] = {}
|
||||
for (model_id, _url), probe in zip(items, probes):
|
||||
try:
|
||||
parse_model_id(model_id)
|
||||
except InvalidModelId:
|
||||
# Ill-formed identifier: report as missing without 400-ing the
|
||||
# whole batch — the workflow author probably typo'd.
|
||||
response_models[model_id] = schemas_out.ModelStatusEntry(
|
||||
state="missing",
|
||||
file_size=probe.file_size,
|
||||
is_hf_downloadable=probe.is_hf_downloadable,
|
||||
)
|
||||
continue
|
||||
|
||||
active = DOWNLOAD_SERVER.get(model_id)
|
||||
if active is not None:
|
||||
response_models[model_id] = schemas_out.ModelStatusEntry(
|
||||
state="downloading",
|
||||
progress=schemas_out.DownloadProgress(
|
||||
bytes_downloaded=active.bytes_downloaded,
|
||||
total_bytes=active.total_bytes,
|
||||
progress=active.progress,
|
||||
),
|
||||
file_size=probe.file_size,
|
||||
is_hf_downloadable=probe.is_hf_downloadable,
|
||||
)
|
||||
continue
|
||||
|
||||
state: schemas_out.ModelState = (
|
||||
"available" if resolve_existing(model_id) is not None else "missing"
|
||||
)
|
||||
response_models[model_id] = schemas_out.ModelStatusEntry(
|
||||
state=state,
|
||||
file_size=probe.file_size,
|
||||
is_hf_downloadable=probe.is_hf_downloadable,
|
||||
)
|
||||
|
||||
return _ok(schemas_out.AvailabilityStatusResponse(
|
||||
models=response_models,
|
||||
hf_auth=schemas_out.HfAuthStatus(
|
||||
token_available=HF_AUTH_STORE.has_token(),
|
||||
eligible=is_hf_auth_eligible(),
|
||||
),
|
||||
))
|
||||
|
||||
|
||||
# ----- 2. start downloads -----
|
||||
|
||||
|
||||
@ROUTES.post("/api/download-models")
|
||||
async def download_models(request: web.Request) -> web.Response:
|
||||
parsed = await _parse_body(request, schemas_in.DownloadModelsRequest)
|
||||
if isinstance(parsed, web.Response):
|
||||
return parsed
|
||||
|
||||
if not parsed.models:
|
||||
return _error(400, "EMPTY_REQUEST", "No models supplied.")
|
||||
|
||||
# ----- precondition pass: validate everything BEFORE registering anything -----
|
||||
# Atomic semantics: if any model fails any precondition (invalid id,
|
||||
# not allow-listed URL, already on disk, already downloading, or gated),
|
||||
# the entire request fails and no state is changed.
|
||||
requested = list(parsed.models.items())
|
||||
|
||||
for model_id, url in requested:
|
||||
try:
|
||||
parse_model_id(model_id)
|
||||
except InvalidModelId as e:
|
||||
return _error(400, "INVALID_MODEL_ID", str(e),
|
||||
{"model_id": model_id})
|
||||
|
||||
if not is_url_allowed(url):
|
||||
return _error(
|
||||
400, "URL_NOT_ALLOWED",
|
||||
"Server-side downloads only accept HuggingFace, Civitai, "
|
||||
"or localhost URLs ending in a known model extension.",
|
||||
{"model_id": model_id, "url": url},
|
||||
)
|
||||
|
||||
if resolve_existing(model_id) is not None:
|
||||
return _error(409, "ALREADY_AVAILABLE",
|
||||
f"Model already exists on disk: {model_id}",
|
||||
{"model_id": model_id})
|
||||
|
||||
if DOWNLOAD_SERVER.is_downloading(model_id):
|
||||
return _error(409, "ALREADY_DOWNLOADING",
|
||||
f"A download for {model_id} is already in progress.",
|
||||
{"model_id": model_id})
|
||||
|
||||
# Reachability check last — it's the only one that talks to the
|
||||
# network. Concurrent probes. For HF URLs ``is_hf_downloadable``
|
||||
# reflects current token access; for non-HF URLs it's None, and we
|
||||
# treat that as "no info, proceed".
|
||||
probes = await asyncio.gather(*(probe_url(url) for _, url in requested))
|
||||
for (model_id, url), probe in zip(requested, probes):
|
||||
if probe.is_hf_downloadable is False:
|
||||
return _error(
|
||||
400, "MODEL_NOT_DOWNLOADABLE",
|
||||
f"Model {model_id} is gated on HuggingFace and the current "
|
||||
f"server token (if any) does not grant access.",
|
||||
{"model_id": model_id, "url": url},
|
||||
)
|
||||
|
||||
# ----- registration pass: try_register is atomic per model_id -----
|
||||
# Defensive: another request might have raced past our pre-check
|
||||
# between the loop above and here. try_register handles that.
|
||||
sessions: list[DownloadSession] = []
|
||||
for model_id, url in requested:
|
||||
session = DOWNLOAD_SERVER.try_register(model_id, url)
|
||||
if session is None:
|
||||
# Race: someone else got in. Roll back what we registered.
|
||||
for s in sessions:
|
||||
DOWNLOAD_SERVER.cancel(s.model_id)
|
||||
return _error(409, "ALREADY_DOWNLOADING",
|
||||
f"A download for {model_id} is already in progress (race).",
|
||||
{"model_id": model_id})
|
||||
sessions.append(session)
|
||||
|
||||
DOWNLOAD_SERVER.sweep_orphan_tmp_files()
|
||||
schedule_batch(sessions)
|
||||
logging.info(
|
||||
"[model_downloader] scheduled %d downloads: %s",
|
||||
len(sessions), [s.model_id for s in sessions],
|
||||
)
|
||||
|
||||
return _ok(schemas_out.DownloadModelsResponse(
|
||||
accepted=True,
|
||||
scheduled=[s.model_id for s in sessions],
|
||||
), status=202)
|
||||
|
||||
|
||||
# ----- 3. cancel a session -----
|
||||
|
||||
|
||||
@ROUTES.post("/api/cancel-model-download-session")
|
||||
async def cancel_model_download_session(request: web.Request) -> web.Response:
|
||||
parsed = await _parse_body(request, schemas_in.CancelDownloadSessionRequest)
|
||||
if isinstance(parsed, web.Response):
|
||||
return parsed
|
||||
|
||||
try:
|
||||
parse_model_id(parsed.model_id)
|
||||
except InvalidModelId as e:
|
||||
return _error(400, "INVALID_MODEL_ID", str(e), {"model_id": parsed.model_id})
|
||||
|
||||
cancelled = DOWNLOAD_SERVER.cancel(parsed.model_id)
|
||||
if not cancelled:
|
||||
return _error(404, "NOT_DOWNLOADING",
|
||||
f"No active download for {parsed.model_id}.",
|
||||
{"model_id": parsed.model_id})
|
||||
|
||||
return _ok(schemas_out.CancelDownloadSessionResponse(cancelled=True))
|
||||
|
||||
|
||||
# ----- 4. HuggingFace OAuth status / login start / logout -----
|
||||
|
||||
|
||||
@ROUTES.get("/api/hf-auth-token-status")
|
||||
async def hf_auth_token_status(request: web.Request) -> web.Response:
|
||||
"""Return whether the server holds a usable HF token + its username.
|
||||
|
||||
Used by the settings UI and (out-of-band) by the frontend on
|
||||
login completion. ``token_available`` is true even if the cached
|
||||
access_token is expired — as long as a refresh_token exists, the
|
||||
user is "logged in" from their perspective.
|
||||
"""
|
||||
token_present = HF_AUTH_STORE.has_token()
|
||||
username: Optional[str] = None
|
||||
if token_present:
|
||||
# Resolve the username via whoami. Done in a worker thread because
|
||||
# huggingface_hub's whoami is synchronous + blocks on a network call.
|
||||
tok = await HF_AUTH_STORE.get_valid_token()
|
||||
if tok is not None:
|
||||
try:
|
||||
username = await asyncio.to_thread(_whoami_username, tok.access_token)
|
||||
except Exception as e:
|
||||
logging.debug("[hf_auth] whoami failed: %s", e)
|
||||
return _ok(schemas_out.HfAuthTokenStatusResponse(
|
||||
token_available=token_present,
|
||||
username=username,
|
||||
))
|
||||
|
||||
|
||||
def _whoami_username(token: str) -> Optional[str]:
|
||||
"""Sync helper: ask HF for the user name attached to a token."""
|
||||
from huggingface_hub import HfApi
|
||||
info = HfApi().whoami(token=token)
|
||||
if isinstance(info, dict):
|
||||
return info.get("name") or info.get("fullname")
|
||||
return None
|
||||
|
||||
|
||||
@ROUTES.post("/api/hf-auth-login-start")
|
||||
async def hf_auth_login_start(request: web.Request) -> web.Response:
|
||||
"""Begin one OAuth attempt: bind the callback port, return the URL.
|
||||
|
||||
Rejected outright if this deployment isn't eligible (we don't
|
||||
surface the option on multi-tenant / public-IP installs).
|
||||
"""
|
||||
if not is_hf_auth_eligible():
|
||||
return _error(
|
||||
403, "HF_AUTH_NOT_ELIGIBLE",
|
||||
"This server is not eligible for interactive HuggingFace login. "
|
||||
"It must be bound to a loopback address and not running in "
|
||||
"--multi-user mode.",
|
||||
)
|
||||
try:
|
||||
url = await start_login_flow()
|
||||
except OAuthInProgressError:
|
||||
return _error(
|
||||
409, "HF_AUTH_IN_PROGRESS",
|
||||
"Another HuggingFace login attempt is in progress. Try again "
|
||||
"after it completes or times out.",
|
||||
)
|
||||
except OAuthCallbackError as e:
|
||||
return _error(
|
||||
503, "HF_AUTH_START_FAILED",
|
||||
f"Could not start the HuggingFace login flow: {e}",
|
||||
)
|
||||
return _ok(schemas_out.HfAuthLoginStartResponse(authorize_url=url))
|
||||
|
||||
|
||||
@ROUTES.post("/api/hf-auth-logout")
|
||||
async def hf_auth_logout(request: web.Request) -> web.Response:
|
||||
"""Drop the in-memory + on-disk HF token."""
|
||||
HF_AUTH_STORE.clear()
|
||||
return _ok(schemas_out.HfAuthLogoutResponse(logged_out=True))
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Request schemas for the model-downloader API.
|
||||
|
||||
Each endpoint accepts a small JSON body. Pydantic enforces the shape at
|
||||
the boundary; route handlers operate only on validated values past that.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class AvailabilityStatusRequest(BaseModel):
|
||||
"""``POST /api/models-availability-status``.
|
||||
|
||||
Sent by the frontend on each poll. Each entry is ``{model_id: url}``;
|
||||
the URL is the one declared in ``properties.models[i].url`` in the
|
||||
workflow JSON and lets the server compute per-id metadata
|
||||
(``file_size`` + ``is_hf_downloadable``) on the same request.
|
||||
"""
|
||||
models: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class DownloadModelsRequest(BaseModel):
|
||||
"""``POST /api/download-models``.
|
||||
|
||||
Same shape as the metadata request — the URL for each model_id.
|
||||
Returns immediately after validation and scheduling.
|
||||
"""
|
||||
models: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class CancelDownloadSessionRequest(BaseModel):
|
||||
"""``POST /api/cancel-model-download-session``."""
|
||||
model_id: str
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AvailabilityStatusRequest",
|
||||
"DownloadModelsRequest",
|
||||
"CancelDownloadSessionRequest",
|
||||
]
|
||||
@@ -0,0 +1,81 @@
|
||||
"""Response schemas for the model-downloader API."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Literal, Optional
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
ModelState = Literal["available", "missing", "downloading"]
|
||||
|
||||
|
||||
class DownloadProgress(BaseModel):
|
||||
"""Embedded in a model entry when its state is ``downloading``."""
|
||||
bytes_downloaded: int
|
||||
total_bytes: Optional[int] = None
|
||||
progress: Optional[float] = None # fraction in [0,1]; null until total known
|
||||
|
||||
|
||||
class ModelStatusEntry(BaseModel):
|
||||
"""Everything the UI needs to render one row, in one shot.
|
||||
|
||||
``state`` reflects what the server has on disk + in-flight; ``file_size``
|
||||
and ``is_hf_downloadable`` come from probes (intrinsic; cached).
|
||||
The HF fields are populated for every poll (cached on the server),
|
||||
so license-acceptance flips show up within one poll interval without
|
||||
any frontend cache invalidation.
|
||||
"""
|
||||
state: ModelState
|
||||
progress: Optional[DownloadProgress] = None
|
||||
file_size: Optional[int] = None
|
||||
# HF-only: True iff the server can fetch this URL with current auth
|
||||
# state. False iff gated and lacking access. None for non-HF URLs.
|
||||
is_hf_downloadable: Optional[bool] = None
|
||||
|
||||
|
||||
class HfAuthStatus(BaseModel):
|
||||
"""Snapshot of HF login state, embedded in availability response."""
|
||||
token_available: bool
|
||||
eligible: bool
|
||||
|
||||
|
||||
class AvailabilityStatusResponse(BaseModel):
|
||||
models: dict[str, ModelStatusEntry]
|
||||
hf_auth: HfAuthStatus
|
||||
|
||||
|
||||
class DownloadModelsResponse(BaseModel):
|
||||
accepted: bool
|
||||
scheduled: list[str]
|
||||
|
||||
|
||||
class CancelDownloadSessionResponse(BaseModel):
|
||||
cancelled: bool
|
||||
|
||||
|
||||
class HfAuthTokenStatusResponse(BaseModel):
|
||||
token_available: bool
|
||||
username: Optional[str] = None
|
||||
|
||||
|
||||
class HfAuthLoginStartResponse(BaseModel):
|
||||
authorize_url: str
|
||||
|
||||
|
||||
class HfAuthLogoutResponse(BaseModel):
|
||||
logged_out: bool
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ModelState",
|
||||
"DownloadProgress",
|
||||
"ModelStatusEntry",
|
||||
"HfAuthStatus",
|
||||
"AvailabilityStatusResponse",
|
||||
"DownloadModelsResponse",
|
||||
"CancelDownloadSessionResponse",
|
||||
"HfAuthTokenStatusResponse",
|
||||
"HfAuthLoginStartResponse",
|
||||
"HfAuthLogoutResponse",
|
||||
]
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Process-wide registry of in-flight model downloads.
|
||||
|
||||
A single instance, ``DOWNLOAD_SERVER``, tracks every currently-running
|
||||
server-side model fetch. Designed to be safe with multiple concurrent
|
||||
clients hitting the API: each model_id has at most one active session,
|
||||
and the API rejects requests that conflict with in-flight downloads.
|
||||
|
||||
Cancellation is cooperative — the download loop checks ``is_active`` on
|
||||
its own session between chunks and raises ``DownloadCancelled`` when the
|
||||
session has been removed from the registry. This avoids the complications
|
||||
of ``Task.cancel()`` from outside the loop while still giving deterministic
|
||||
rollback semantics (the worker is responsible for deleting its own
|
||||
``.tmp`` on the cancel path).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from app.model_downloader.paths import iter_all_tmp_paths
|
||||
|
||||
|
||||
class DownloadCancelled(Exception):
|
||||
"""Raised by the streaming loop when its session has been removed
|
||||
from the registry (cancellation request) and the worker should roll
|
||||
back its ``.tmp`` file."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class DownloadSession:
|
||||
"""One in-flight download.
|
||||
|
||||
``progress`` is a fraction in ``[0.0, 1.0]``; ``None`` until the first
|
||||
byte arrives and we know whether the response carries a
|
||||
``Content-Length``. ``total_bytes`` mirrors that header when present.
|
||||
"""
|
||||
model_id: str
|
||||
url: str
|
||||
progress: Optional[float] = None
|
||||
bytes_downloaded: int = 0
|
||||
total_bytes: Optional[int] = None
|
||||
# Sequence number used solely as identity for the cancellation check —
|
||||
# so that "cancel + restart" doesn't get confused by stale workers.
|
||||
epoch: int = field(default_factory=lambda: 0)
|
||||
|
||||
|
||||
class DownloadServer:
|
||||
"""Singleton registry of active downloads.
|
||||
|
||||
All mutation goes through this object so concurrent route handlers
|
||||
see a consistent view. The ``_lock`` is a plain threading lock
|
||||
because the registry is consulted from both the asyncio event-loop
|
||||
thread (route handlers) and from any worker coroutines spawned to
|
||||
perform downloads.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._sessions: dict[str, DownloadSession] = {}
|
||||
self._epoch_counter = 0
|
||||
self._orphan_sweep_done = False
|
||||
|
||||
# ----- lifecycle -----
|
||||
|
||||
def sweep_orphan_tmp_files(self) -> None:
|
||||
"""Idempotently sweep ``*.tmp`` files left by crashed downloads.
|
||||
|
||||
Deferred off the import path so module load doesn't block on
|
||||
filesystem I/O against potentially-slow mounts. Each route handler
|
||||
that might create a new ``.tmp`` runs this exactly once.
|
||||
"""
|
||||
with self._lock:
|
||||
if self._orphan_sweep_done:
|
||||
return
|
||||
self._orphan_sweep_done = True
|
||||
for path in iter_all_tmp_paths():
|
||||
try:
|
||||
os.remove(path)
|
||||
logging.info("[model_downloader] removed orphan tmp file: %s", path)
|
||||
except OSError as e:
|
||||
logging.warning("[model_downloader] could not remove %s: %s", path, e)
|
||||
|
||||
# ----- queries -----
|
||||
|
||||
def is_downloading(self, model_id: str) -> bool:
|
||||
with self._lock:
|
||||
return model_id in self._sessions
|
||||
|
||||
def get(self, model_id: str) -> Optional[DownloadSession]:
|
||||
with self._lock:
|
||||
return self._sessions.get(model_id)
|
||||
|
||||
def snapshot(self) -> dict[str, DownloadSession]:
|
||||
"""Return a shallow copy of the current sessions map."""
|
||||
with self._lock:
|
||||
return dict(self._sessions)
|
||||
|
||||
# ----- mutations -----
|
||||
|
||||
def try_register(self, model_id: str, url: str) -> Optional[DownloadSession]:
|
||||
"""Atomically register a new session iff none exists for ``model_id``.
|
||||
|
||||
Returns the new session on success, ``None`` if a session is already
|
||||
in flight. Callers must check the return value — the caller is the
|
||||
sole owner of the session it gets back.
|
||||
"""
|
||||
with self._lock:
|
||||
if model_id in self._sessions:
|
||||
return None
|
||||
self._epoch_counter += 1
|
||||
session = DownloadSession(
|
||||
model_id=model_id,
|
||||
url=url,
|
||||
epoch=self._epoch_counter,
|
||||
)
|
||||
self._sessions[model_id] = session
|
||||
return session
|
||||
|
||||
def update_progress(
|
||||
self,
|
||||
session: DownloadSession,
|
||||
bytes_downloaded: int,
|
||||
total_bytes: Optional[int],
|
||||
) -> None:
|
||||
"""Update progress on a session. No-op if the session has been
|
||||
removed (cancelled) — caller should check ``is_active`` separately."""
|
||||
with self._lock:
|
||||
current = self._sessions.get(session.model_id)
|
||||
if current is None or current.epoch != session.epoch:
|
||||
return
|
||||
current.bytes_downloaded = bytes_downloaded
|
||||
current.total_bytes = total_bytes
|
||||
if total_bytes and total_bytes > 0:
|
||||
current.progress = min(1.0, bytes_downloaded / total_bytes)
|
||||
|
||||
def is_active(self, session: DownloadSession) -> bool:
|
||||
"""True iff this exact session is still the registered one for
|
||||
its model_id. False after cancellation, after completion, or if
|
||||
another session has replaced it."""
|
||||
with self._lock:
|
||||
current = self._sessions.get(session.model_id)
|
||||
return current is not None and current.epoch == session.epoch
|
||||
|
||||
def finish(self, session: DownloadSession) -> None:
|
||||
"""Remove a completed (or cancelled) session from the registry.
|
||||
|
||||
Safe to call multiple times. Only removes if the epoch matches —
|
||||
we never accidentally evict a *newer* session for the same model_id.
|
||||
"""
|
||||
with self._lock:
|
||||
current = self._sessions.get(session.model_id)
|
||||
if current is not None and current.epoch == session.epoch:
|
||||
del self._sessions[session.model_id]
|
||||
|
||||
def reset_for_tests(self) -> None:
|
||||
"""Clear all sessions and reset the epoch counter. Test-only."""
|
||||
with self._lock:
|
||||
self._sessions.clear()
|
||||
self._epoch_counter = 0
|
||||
|
||||
def cancel(self, model_id: str) -> bool:
|
||||
"""Remove the session registered for ``model_id``.
|
||||
|
||||
Returns True if there was an active session to cancel. The worker
|
||||
will discover the cancellation on its next ``is_active`` check
|
||||
and roll back its ``.tmp`` file.
|
||||
"""
|
||||
with self._lock:
|
||||
if model_id in self._sessions:
|
||||
del self._sessions[model_id]
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
DOWNLOAD_SERVER = DownloadServer()
|
||||
@@ -0,0 +1,216 @@
|
||||
"""Streaming download worker with progress reporting and cancellation.
|
||||
|
||||
Each download writes to ``<final_path>.tmp`` and atomically renames into
|
||||
place on success. Between chunks the worker checks the registry for
|
||||
cancellation (via ``DownloadServer.is_active``) and rolls back its
|
||||
``.tmp`` on cancel or on any error.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from app.model_downloader.download_server import (
|
||||
DOWNLOAD_SERVER,
|
||||
DownloadCancelled,
|
||||
DownloadSession,
|
||||
)
|
||||
from app.model_downloader.hf_auth.auth_store import HF_AUTH_STORE
|
||||
from app.model_downloader.hf_url import is_hf_url
|
||||
from app.model_downloader.http_client import get_session, parse_content_length
|
||||
from app.model_downloader.paths import resolve_destination
|
||||
|
||||
|
||||
CHUNK_SIZE = 64 * 1024 # 64 KiB — same scale as other ComfyUI download paths.
|
||||
REQUEST_TIMEOUT = aiohttp.ClientTimeout(total=None, sock_connect=30, sock_read=120)
|
||||
|
||||
|
||||
async def stream_to_disk(session: DownloadSession) -> str:
|
||||
"""Run a single download to completion or cancellation.
|
||||
|
||||
Returns the final on-disk path on success. Removes the ``.tmp`` and
|
||||
raises on cancellation or failure. The session is finished
|
||||
(removed from the registry) exactly once, here — callers do not
|
||||
need to call ``DOWNLOAD_SERVER.finish`` themselves.
|
||||
"""
|
||||
final_path, tmp_path = resolve_destination(session.model_id, session.epoch)
|
||||
os.makedirs(os.path.dirname(final_path), exist_ok=True)
|
||||
|
||||
# Wipe any stale .tmp from a previous failed attempt before we start —
|
||||
# otherwise a partial body could masquerade as our completed download
|
||||
# when the rename finally happens.
|
||||
_remove_if_exists(tmp_path)
|
||||
|
||||
bytes_seen = 0
|
||||
try:
|
||||
http = await get_session()
|
||||
headers = _auth_headers_for(session.url)
|
||||
logging.info(
|
||||
"[model_downloader] starting GET %s (auth=%s)",
|
||||
session.url, "yes" if "Authorization" in headers else "no",
|
||||
)
|
||||
async with http.get(
|
||||
session.url,
|
||||
allow_redirects=True,
|
||||
timeout=REQUEST_TIMEOUT,
|
||||
headers=headers,
|
||||
) as resp:
|
||||
if resp.status != 200:
|
||||
# Capture a snippet of the response body so 4xx/5xx aren't
|
||||
# opaque in the logs — HF returns JSON or HTML with a
|
||||
# human-readable reason on failures.
|
||||
body_snippet = await _read_short(resp)
|
||||
logging.warning(
|
||||
"[model_downloader] GET %s failed: status=%d final_url=%s body=%s",
|
||||
session.url, resp.status, str(resp.url), body_snippet,
|
||||
)
|
||||
raise DownloadError(
|
||||
f"unexpected HTTP {resp.status} fetching {session.url}: {body_snippet}",
|
||||
status=resp.status,
|
||||
)
|
||||
|
||||
total = parse_content_length(resp.headers.get("Content-Length"))
|
||||
DOWNLOAD_SERVER.update_progress(session, 0, total)
|
||||
|
||||
with open(tmp_path, "wb") as f:
|
||||
async for chunk in resp.content.iter_chunked(CHUNK_SIZE):
|
||||
# Cancellation check between chunks. Cheap and means
|
||||
# cancellation latency is bounded by one chunk plus
|
||||
# one ``write()`` — typically well under a second
|
||||
# even on slow disks.
|
||||
if not DOWNLOAD_SERVER.is_active(session):
|
||||
raise DownloadCancelled()
|
||||
f.write(chunk)
|
||||
bytes_seen += len(chunk)
|
||||
DOWNLOAD_SERVER.update_progress(session, bytes_seen, total)
|
||||
|
||||
# Final cancellation check before we promote the .tmp to the real
|
||||
# filename — avoids the awkward case where cancel arrives during
|
||||
# the very last chunk and we'd otherwise commit anyway.
|
||||
if not DOWNLOAD_SERVER.is_active(session):
|
||||
raise DownloadCancelled()
|
||||
|
||||
# Size verification before commit. aiohttp already raises
|
||||
# ClientPayloadError on a truncated Content-Length/chunked body,
|
||||
# but this also catches the HTTP/1.0-style case (no Content-Length
|
||||
# + Connection: close) where a short read can masquerade as a
|
||||
# complete download.
|
||||
if total is not None and bytes_seen != total:
|
||||
raise DownloadError(
|
||||
f"size mismatch for {session.model_id}: "
|
||||
f"got {bytes_seen} of {total} bytes from {session.url}"
|
||||
)
|
||||
|
||||
# Atomic rename. os.replace is atomic within the same filesystem,
|
||||
# which is guaranteed here because tmp lives alongside final_path.
|
||||
os.replace(tmp_path, final_path)
|
||||
logging.info(
|
||||
"[model_downloader] downloaded %s (%d bytes) from %s",
|
||||
session.model_id, bytes_seen, session.url,
|
||||
)
|
||||
return final_path
|
||||
|
||||
except DownloadCancelled:
|
||||
logging.info("[model_downloader] cancelled: %s", session.model_id)
|
||||
_remove_if_exists(tmp_path)
|
||||
raise
|
||||
except Exception as e:
|
||||
logging.warning(
|
||||
"[model_downloader] failed: %s from %s: %s: %s",
|
||||
session.model_id, session.url, type(e).__name__, e,
|
||||
exc_info=True,
|
||||
)
|
||||
_remove_if_exists(tmp_path)
|
||||
raise
|
||||
finally:
|
||||
# In all terminal states (success / cancel / error) drop the
|
||||
# session from the registry. Idempotent — only removes if we're
|
||||
# still the live epoch for this model_id.
|
||||
DOWNLOAD_SERVER.finish(session)
|
||||
|
||||
|
||||
class DownloadError(Exception):
|
||||
"""Network / protocol error during a download."""
|
||||
|
||||
def __init__(self, message: str, status: Optional[int] = None) -> None:
|
||||
super().__init__(message)
|
||||
self.status = status
|
||||
|
||||
|
||||
async def _read_short(resp: aiohttp.ClientResponse, limit: int = 512) -> str:
|
||||
"""Read up to ``limit`` bytes of a response body for logging.
|
||||
|
||||
Used to surface the JSON/HTML reason from an HF non-2xx response in
|
||||
server logs instead of just the status code. Best-effort: any
|
||||
error here is swallowed.
|
||||
"""
|
||||
try:
|
||||
raw = await resp.content.read(limit)
|
||||
return raw.decode("utf-8", errors="replace").strip()
|
||||
except Exception:
|
||||
return "<unreadable>"
|
||||
|
||||
|
||||
def _auth_headers_for(url: str) -> dict[str, str]:
|
||||
"""Return any auth headers we should add to the GET for ``url``.
|
||||
|
||||
For HuggingFace URLs we inject the user's OAuth access token as a
|
||||
Bearer header — this is HF's documented way to access gated repos
|
||||
(see ``huggingface_hub.hf_hub_download``'s wire format). For every
|
||||
other host we send no extra headers; allowlisted public files
|
||||
don't need them and we don't want to leak tokens to other hosts.
|
||||
"""
|
||||
if not is_hf_url(url):
|
||||
return {}
|
||||
tok = HF_AUTH_STORE.get_token_sync()
|
||||
if tok is None or not tok.access_token:
|
||||
return {}
|
||||
return {"Authorization": f"Bearer {tok.access_token}"}
|
||||
|
||||
|
||||
def _remove_if_exists(path: str) -> None:
|
||||
try:
|
||||
os.remove(path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except OSError as e:
|
||||
logging.warning("[model_downloader] could not remove %s: %s", path, e)
|
||||
|
||||
|
||||
async def run_batch_sequential(sessions: list[DownloadSession]) -> None:
|
||||
"""Run a list of sessions one after the other.
|
||||
|
||||
Each session is independent: a failure or cancellation of one does
|
||||
not abort the rest. Cancellations are observable via the registry
|
||||
*before* a given download starts, so a session that's been
|
||||
pre-cancelled (cancel before the worker reached it) just gets skipped.
|
||||
"""
|
||||
for session in sessions:
|
||||
# If the session got cancelled before its turn, skip without
|
||||
# touching disk. This is what makes the per-request "sequential
|
||||
# but cancellable" semantic work.
|
||||
if not DOWNLOAD_SERVER.is_active(session):
|
||||
DOWNLOAD_SERVER.finish(session)
|
||||
continue
|
||||
try:
|
||||
await stream_to_disk(session)
|
||||
except DownloadCancelled:
|
||||
# Already logged + tmp removed inside stream_to_disk.
|
||||
continue
|
||||
except Exception:
|
||||
# stream_to_disk already logged. Continue with the rest of the batch.
|
||||
continue
|
||||
|
||||
|
||||
def schedule_batch(sessions: list[DownloadSession]) -> asyncio.Task:
|
||||
"""Kick off ``run_batch_sequential`` on the running event loop.
|
||||
|
||||
Returned task is fire-and-forget; the API handler returns immediately
|
||||
after scheduling and clients observe progress via the polling endpoints.
|
||||
"""
|
||||
return asyncio.create_task(run_batch_sequential(sessions))
|
||||
@@ -0,0 +1,245 @@
|
||||
"""Per-URL probes for the unified availability endpoint.
|
||||
|
||||
Three cached/derived facts per URL:
|
||||
|
||||
- ``is_gated`` intrinsic to the model; cached forever once known.
|
||||
Determined by ``auth_check(repo_id, token=None)``:
|
||||
``GatedRepoError`` → True, success → False.
|
||||
|
||||
- ``is_hf_downloadable`` depends on the *current* token; recomputed every
|
||||
call. For non-gated URLs this is trivially True
|
||||
(no HF call needed). For gated URLs we run
|
||||
``auth_check`` with the stored token each call.
|
||||
|
||||
- ``file_size`` intrinsic to the file. Cached forever once
|
||||
determined (including ``None`` on transient
|
||||
failure — we don't retry). We only attempt the
|
||||
HEAD when we already know the URL is downloadable
|
||||
to us; that way a failed-because-gated probe
|
||||
never lands as a cached ``None``.
|
||||
|
||||
Caches are per-process, in-memory; small, no eviction needed for the
|
||||
workflow-scale (~tens of URLs). Concurrent calls for the same URL
|
||||
deduplicate via per-URL ``asyncio.Lock``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
from huggingface_hub import HfApi
|
||||
from huggingface_hub.errors import (
|
||||
GatedRepoError,
|
||||
HfHubHTTPError,
|
||||
RepositoryNotFoundError,
|
||||
)
|
||||
|
||||
from app.model_downloader.allowlist import is_url_allowed
|
||||
from app.model_downloader.hf_auth.auth_store import HF_AUTH_STORE
|
||||
from app.model_downloader.hf_url import is_hf_url, repo_id_from_url
|
||||
from app.model_downloader.http_client import get_session, parse_content_length
|
||||
|
||||
|
||||
_HEAD_TIMEOUT = aiohttp.ClientTimeout(total=15)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProbeResult:
|
||||
file_size: Optional[int]
|
||||
is_hf_downloadable: Optional[bool]
|
||||
|
||||
|
||||
# --- caches -------------------------------------------------------------- #
|
||||
|
||||
|
||||
# url → bool. Whether this URL's HF repo gates access. Intrinsic to the
|
||||
# model — never changes for a given URL.
|
||||
_is_gated_cache: dict[str, bool] = {}
|
||||
|
||||
# url → Optional[int]. The file's size in bytes, ``None`` if a probe
|
||||
# was attempted and produced no answer. **Only populated when we knew
|
||||
# the URL was downloadable to us at probe time** — so gated-without-
|
||||
# access never lands a ``None`` here that we'd be stuck with after login.
|
||||
_file_size_cache: dict[str, Optional[int]] = {}
|
||||
|
||||
# Per-URL locks for single-flight probes — when multiple polls arrive
|
||||
# in the same tick for the same URL, exactly one of them runs the HF
|
||||
# call and the others wait on the result.
|
||||
_locks: dict[str, asyncio.Lock] = {}
|
||||
|
||||
|
||||
def _lock_for(url: str) -> asyncio.Lock:
|
||||
lock = _locks.get(url)
|
||||
if lock is None:
|
||||
lock = asyncio.Lock()
|
||||
_locks[url] = lock
|
||||
return lock
|
||||
|
||||
|
||||
def clear_caches_for_tests() -> None:
|
||||
"""Test-only: drop everything."""
|
||||
_is_gated_cache.clear()
|
||||
_file_size_cache.clear()
|
||||
_locks.clear()
|
||||
|
||||
|
||||
# --- public entrypoint --------------------------------------------------- #
|
||||
|
||||
|
||||
async def probe_url(url: str) -> ProbeResult:
|
||||
"""Return downloadability + size for one URL, using caches where safe."""
|
||||
if not is_url_allowed(url):
|
||||
return ProbeResult(file_size=None, is_hf_downloadable=None)
|
||||
if not is_hf_url(url):
|
||||
# Non-HF: ``is_hf_downloadable`` is "not applicable" (None).
|
||||
# Size we still cache so we don't HEAD on every poll.
|
||||
size = await _get_or_probe_size(url, token=None)
|
||||
return ProbeResult(file_size=size, is_hf_downloadable=None)
|
||||
|
||||
repo_id = repo_id_from_url(url)
|
||||
if repo_id is None:
|
||||
return ProbeResult(file_size=None, is_hf_downloadable=None)
|
||||
|
||||
# Determine intrinsic gating once.
|
||||
gated = await _resolve_is_gated(url, repo_id)
|
||||
if gated is None:
|
||||
return ProbeResult(file_size=None, is_hf_downloadable=None)
|
||||
|
||||
# Compute current-token downloadability per call.
|
||||
tok = await HF_AUTH_STORE.get_valid_token()
|
||||
token_str: Optional[str] = tok.access_token if tok else None
|
||||
if not gated:
|
||||
is_hf_downloadable: Optional[bool] = True
|
||||
else:
|
||||
is_hf_downloadable = await _auth_check_with_token(repo_id, token_str)
|
||||
|
||||
if is_hf_downloadable is True:
|
||||
size = await _get_or_probe_size(url, token=token_str)
|
||||
else:
|
||||
# Skip the HEAD entirely — would 401 and we'd be stuck with
|
||||
# cached None that survives a later login.
|
||||
size = None
|
||||
|
||||
return ProbeResult(file_size=size, is_hf_downloadable=is_hf_downloadable)
|
||||
|
||||
|
||||
# --- gated/auth probes --------------------------------------------------- #
|
||||
|
||||
|
||||
async def _resolve_is_gated(url: str, repo_id: str) -> Optional[bool]:
|
||||
"""Decide once whether ``repo_id`` is a gated repo."""
|
||||
cached = _is_gated_cache.get(url)
|
||||
if cached is not None:
|
||||
return cached
|
||||
|
||||
async with _lock_for(url):
|
||||
cached = _is_gated_cache.get(url)
|
||||
if cached is not None:
|
||||
return cached
|
||||
# Probe anonymously (token=None) on purpose: an unauthenticated
|
||||
# auth_check is what makes HF raise GatedRepoError for gated repos.
|
||||
# With a token, a gated-but-accepted repo would succeed and look
|
||||
# ungated.
|
||||
try:
|
||||
await asyncio.to_thread(_auth_check_sync, repo_id, None)
|
||||
_is_gated_cache[url] = False
|
||||
return False
|
||||
except GatedRepoError:
|
||||
_is_gated_cache[url] = True
|
||||
return True
|
||||
except RepositoryNotFoundError:
|
||||
# Repo doesn't exist publicly. Treat as gated — we can't
|
||||
# serve it without auth, and an authenticated check might
|
||||
# still succeed if it's a private repo the user can see.
|
||||
_is_gated_cache[url] = True
|
||||
return True
|
||||
except (HfHubHTTPError, Exception) as e:
|
||||
logging.debug(
|
||||
"[hf_auth] is_gated probe failed for %s (will retry): %s",
|
||||
repo_id, e,
|
||||
)
|
||||
return None # don't cache; retry next call
|
||||
|
||||
|
||||
async def _auth_check_with_token(
|
||||
repo_id: str, token: Optional[str]
|
||||
) -> Optional[bool]:
|
||||
"""Auth-check with the supplied token. True/False/None per outcome."""
|
||||
try:
|
||||
await asyncio.to_thread(_auth_check_sync, repo_id, token)
|
||||
return True
|
||||
except GatedRepoError:
|
||||
return False
|
||||
except RepositoryNotFoundError:
|
||||
return False
|
||||
except HfHubHTTPError as e:
|
||||
# 401/403 covers org-SSO-required, revoked tokens, and similar —
|
||||
# all of which mean "can't fetch right now" from the user's POV.
|
||||
status = getattr(getattr(e, "response", None), "status_code", None)
|
||||
if status in (401, 403):
|
||||
return False
|
||||
logging.debug(
|
||||
"[hf_auth] auth_check transient failure for %s: %s", repo_id, e,
|
||||
)
|
||||
return None
|
||||
except Exception as e:
|
||||
logging.warning("[hf_auth] unexpected auth_check error for %s: %s", repo_id, e)
|
||||
return None
|
||||
|
||||
|
||||
def _auth_check_sync(repo_id: str, token: Optional[str]) -> None:
|
||||
"""Thin sync wrapper around ``HfApi.auth_check`` for ``asyncio.to_thread``."""
|
||||
HfApi().auth_check(repo_id, token=token)
|
||||
|
||||
|
||||
# --- size probe ---------------------------------------------------------- #
|
||||
|
||||
|
||||
async def _get_or_probe_size(url: str, token: Optional[str]) -> Optional[int]:
|
||||
"""Return the cached size or HEAD the URL once and cache the result."""
|
||||
if url in _file_size_cache:
|
||||
return _file_size_cache[url]
|
||||
|
||||
async with _lock_for(url):
|
||||
if url in _file_size_cache:
|
||||
return _file_size_cache[url]
|
||||
size = await _probe_size_once(url, token=token)
|
||||
_file_size_cache[url] = size
|
||||
return size
|
||||
|
||||
|
||||
async def _probe_size_once(url: str, token: Optional[str]) -> Optional[int]:
|
||||
"""HEAD the URL and return the file size in bytes, or None on failure.
|
||||
|
||||
HuggingFace serves LFS-tracked files via a 302 to a signed CDN URL.
|
||||
The real file size lives in the non-standard ``X-Linked-Size`` header
|
||||
on that 302 response (``Content-Length`` is the redirect-body length).
|
||||
Disabling redirect-follow lets us read either header on the same
|
||||
response:
|
||||
|
||||
- LFS files: 302 + ``X-Linked-Size``
|
||||
- Small/non-LFS files: 200 + ``Content-Length``
|
||||
"""
|
||||
headers = {"Authorization": f"Bearer {token}"} if token else {}
|
||||
try:
|
||||
session = await get_session()
|
||||
async with session.head(
|
||||
url, allow_redirects=False, timeout=_HEAD_TIMEOUT, headers=headers,
|
||||
) as resp:
|
||||
linked = parse_content_length(resp.headers.get("X-Linked-Size"))
|
||||
if linked is not None:
|
||||
return linked
|
||||
if resp.status == 200:
|
||||
return parse_content_length(resp.headers.get("Content-Length"))
|
||||
return None
|
||||
except (aiohttp.ClientError, TimeoutError, OSError):
|
||||
return None
|
||||
|
||||
|
||||
# Backward-compat shim so consumers that still import the old name keep
|
||||
# building during the refactor; can be removed once routes are updated.
|
||||
MetadataProbeResult = ProbeResult
|
||||
@@ -0,0 +1,121 @@
|
||||
"""In-memory token cache with lazy disk persistence + refresh.
|
||||
|
||||
Public surface is the ``HF_AUTH_STORE`` singleton. Callers ask
|
||||
``get_valid_token()``; the store transparently refreshes from disk
|
||||
on first use, refreshes via the OAuth refresh_token if the cached
|
||||
access_token is expired, and returns ``None`` if neither path works.
|
||||
|
||||
The refresh path imports ``oauth.refresh_access_token`` lazily to
|
||||
avoid an import cycle (oauth needs the store to save tokens it
|
||||
acquires).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from typing import Optional
|
||||
|
||||
from app.model_downloader.hf_auth.token_store import (
|
||||
Token,
|
||||
delete_token,
|
||||
load_token,
|
||||
save_token,
|
||||
)
|
||||
|
||||
|
||||
class HfAuthStore:
|
||||
def __init__(self) -> None:
|
||||
self._lock = threading.Lock()
|
||||
self._token: Optional[Token] = None
|
||||
self._loaded_from_disk = False
|
||||
|
||||
def _ensure_loaded(self) -> None:
|
||||
"""Read the disk token into memory on first access."""
|
||||
if self._loaded_from_disk:
|
||||
return
|
||||
with self._lock:
|
||||
if self._loaded_from_disk:
|
||||
return
|
||||
self._token = load_token()
|
||||
self._loaded_from_disk = True
|
||||
|
||||
def has_token(self) -> bool:
|
||||
"""Cheap check: is there any token in memory?
|
||||
|
||||
Does not attempt refresh; an expired-but-refreshable token still
|
||||
counts as "logged in" from the user's perspective.
|
||||
"""
|
||||
self._ensure_loaded()
|
||||
return self._token is not None
|
||||
|
||||
def _store_token_locked(self, token: Token) -> None:
|
||||
"""Set the in-memory token and persist it to disk.
|
||||
|
||||
Caller must already hold ``self._lock``. Keeping the disk write inside
|
||||
the lock means memory and disk flip together — a concurrent ``clear()``
|
||||
or refresh can't interleave between them.
|
||||
"""
|
||||
self._token = token
|
||||
self._loaded_from_disk = True
|
||||
save_token(token)
|
||||
|
||||
def set_token(self, token: Token) -> None:
|
||||
"""Replace the in-memory token and persist to disk (atomically)."""
|
||||
with self._lock:
|
||||
self._store_token_locked(token)
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Forget the token in memory and on disk (logout)."""
|
||||
with self._lock:
|
||||
self._token = None
|
||||
self._loaded_from_disk = True
|
||||
delete_token()
|
||||
|
||||
def get_token_sync(self) -> Optional[Token]:
|
||||
"""Return the cached token without refreshing.
|
||||
|
||||
Sync callers (e.g. constructing an Authorization header in a
|
||||
non-async path) use this. They accept an expired token over
|
||||
``None``; HF will simply return 401 and the caller can decide
|
||||
what to do.
|
||||
"""
|
||||
self._ensure_loaded()
|
||||
return self._token
|
||||
|
||||
async def get_valid_token(self) -> Optional[Token]:
|
||||
"""Return a fresh token, refreshing via OAuth if necessary.
|
||||
|
||||
Returns ``None`` if there's no cached token at all, or if the
|
||||
cached token is expired and refresh failed. Callers should
|
||||
treat that as "user is not logged in".
|
||||
"""
|
||||
self._ensure_loaded()
|
||||
with self._lock:
|
||||
tok = self._token
|
||||
if tok is None:
|
||||
return None
|
||||
if tok.is_valid():
|
||||
return tok
|
||||
if not tok.refresh_token:
|
||||
return None
|
||||
|
||||
# Lazy import to avoid the oauth ↔ store import cycle.
|
||||
from app.model_downloader.hf_auth.oauth import refresh_access_token
|
||||
|
||||
try:
|
||||
refreshed = await refresh_access_token(tok.refresh_token)
|
||||
except Exception as e:
|
||||
logging.warning("[hf_auth] token refresh failed: %s", e)
|
||||
return None
|
||||
|
||||
with self._lock:
|
||||
# If a logout (clear) or another update replaced the token while we
|
||||
# were awaiting the refresh, don't resurrect the old session.
|
||||
if self._token is not tok:
|
||||
return None
|
||||
self._store_token_locked(refreshed)
|
||||
return refreshed
|
||||
|
||||
|
||||
HF_AUTH_STORE = HfAuthStore()
|
||||
@@ -0,0 +1,55 @@
|
||||
"""Whether this deployment is allowed to do interactive HF OAuth.
|
||||
|
||||
We only let the server hold a HuggingFace token under a strict trust
|
||||
assumption: this is a *single tenant local* install. Concretely:
|
||||
|
||||
- The server is bound to a loopback address. SSH tunneling /
|
||||
reverse-proxies can defeat this, but it's the strongest signal
|
||||
we have without an authentication system.
|
||||
- ``--multi-user`` is off. A shared token used implicitly by multiple
|
||||
declared users would be a footgun — one user's gated downloads
|
||||
would silently authenticate as another.
|
||||
|
||||
Anything else and the frontend hides the HF login UI entirely; gated
|
||||
models continue to show the "acquire it manually" message.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ipaddress
|
||||
import socket
|
||||
|
||||
from comfy.cli_args import args
|
||||
|
||||
|
||||
def _is_loopback(host: str | None) -> bool:
|
||||
"""Duplicates ``server.is_loopback`` (small, no shared module yet).
|
||||
|
||||
Resolves a host or IP literal to whether it lives on the loopback
|
||||
interface (127.0.0.0/8 for IPv4, ::1 for IPv6). Returns False for
|
||||
``0.0.0.0`` / ``::`` because those are bind-all wildcards, not
|
||||
loopback.
|
||||
"""
|
||||
if host is None:
|
||||
return False
|
||||
try:
|
||||
return ipaddress.ip_address(host).is_loopback
|
||||
except ValueError:
|
||||
pass
|
||||
|
||||
loopback = False
|
||||
for family in (socket.AF_INET, socket.AF_INET6):
|
||||
try:
|
||||
r = socket.getaddrinfo(host, None, family, socket.SOCK_STREAM)
|
||||
for _family, _, _, _, sockaddr in r:
|
||||
if not ipaddress.ip_address(sockaddr[0]).is_loopback:
|
||||
return loopback
|
||||
loopback = True
|
||||
except socket.gaierror:
|
||||
pass
|
||||
return loopback
|
||||
|
||||
|
||||
def is_hf_auth_eligible() -> bool:
|
||||
"""True iff this deployment may surface the HF OAuth flow."""
|
||||
return _is_loopback(args.listen) and not args.multi_user
|
||||
@@ -0,0 +1,301 @@
|
||||
"""OAuth 2.0 PKCE flow against HuggingFace's authorization server.
|
||||
|
||||
Wired so that ``POST /api/hf-auth-login-start`` can:
|
||||
1. Generate state + PKCE verifier/challenge in this process.
|
||||
2. Spin up a short-lived loopback HTTP server at port 41954 to
|
||||
receive the redirect callback from HF.
|
||||
3. Return the ``authorize_url`` for the frontend to open in a new tab.
|
||||
|
||||
After the user grants consent on huggingface.co, HF redirects to the
|
||||
local callback URL with ``code`` and ``state``. The callback server
|
||||
validates ``state`` (CSRF), exchanges the code for tokens via PKCE,
|
||||
hands the resulting Token to ``HF_AUTH_STORE.set_token``, and shuts
|
||||
itself down.
|
||||
|
||||
Before this can be exercised end-to-end a maintainer must register a
|
||||
HuggingFace OAuth app and substitute the ``HF_CLIENT_ID`` placeholder
|
||||
below. See the comment above the constant for the exact steps.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import hashlib
|
||||
import logging
|
||||
import secrets
|
||||
import threading
|
||||
import time
|
||||
|
||||
import aiohttp
|
||||
from aiohttp import web
|
||||
|
||||
from app.model_downloader.hf_auth.auth_store import HF_AUTH_STORE
|
||||
from app.model_downloader.hf_auth.token_store import Token
|
||||
from app.model_downloader.http_client import get_session
|
||||
|
||||
|
||||
# --- HF OAuth app registration -------------------------------------------- #
|
||||
# NOTE: The OAuth client_id below is a placeholder. Before this feature can be
|
||||
# exercised end-to-end, a maintainer must register a HuggingFace OAuth app
|
||||
# under a Comfy-Org-controlled HF account and substitute its client_id here.
|
||||
# Detailed walkthrough is in docs/server-side-model-downloads-handover.html
|
||||
# ("HuggingFace OAuth app setup" section). Short version:
|
||||
# 1. huggingface.co → Settings → Connected Apps → "Create app"
|
||||
# 2. Default Scopes: check ``openid`` + ``profile`` (User Info) and
|
||||
# ``gated-repos`` (Repository Access). Leave everything else off.
|
||||
# 3. Redirect URLs: exactly ``http://127.0.0.1:41954/api/auth/huggingface/callback``
|
||||
# — must match ``REDIRECT_URI`` below; change both in lockstep if you
|
||||
# change ``CALLBACK_PORT``.
|
||||
# 4. Save → copy the resulting Client ID into ``HF_CLIENT_ID`` below.
|
||||
# The client_id is not a secret (it travels through the user's browser in
|
||||
# plaintext); HF's "Public app" type means there's no client secret to
|
||||
# manage — PKCE replaces it.
|
||||
HF_CLIENT_ID = "REPLACE_ME_WITH_COMFY_ORG_HF_OAUTH_CLIENT_ID"
|
||||
|
||||
CALLBACK_HOST = "127.0.0.1"
|
||||
CALLBACK_PORT = 41954
|
||||
CALLBACK_PATH = "/api/auth/huggingface/callback"
|
||||
REDIRECT_URI = f"http://{CALLBACK_HOST}:{CALLBACK_PORT}{CALLBACK_PATH}"
|
||||
|
||||
AUTHORIZE_URL = "https://huggingface.co/oauth/authorize"
|
||||
TOKEN_URL = "https://huggingface.co/oauth/token"
|
||||
_TOKEN_REQUEST_TIMEOUT = aiohttp.ClientTimeout(total=30)
|
||||
# Minimal scope set for the feature:
|
||||
# - openid : required by HF when the app uses OIDC at all
|
||||
# - profile : lets ``HfApi.whoami(token=...)`` return a username for the
|
||||
# settings UI; cosmetic but expected
|
||||
# - gated-repos : grants the token enough to call ``auth_check`` and
|
||||
# download files from public gated repos the user has
|
||||
# accepted the license for. The wider ``read-repos`` scope
|
||||
# would also work (it includes ``gated-repos``) but it
|
||||
# additionally grants private-repo read access, which we
|
||||
# don't need and which makes the consent screen scarier
|
||||
# for the user.
|
||||
SCOPE = "openid profile gated-repos"
|
||||
|
||||
# Maximum time the callback server stays up waiting for the user to
|
||||
# complete consent on huggingface.co. Past this, the port closes and
|
||||
# the user has to click "Log in" again.
|
||||
CALLBACK_TIMEOUT_SECS = 300
|
||||
|
||||
|
||||
# Process-wide lock so two simultaneous /api/hf-auth-login-start
|
||||
# requests don't fight over port CALLBACK_PORT.
|
||||
_OAUTH_LOCK = threading.Lock()
|
||||
|
||||
|
||||
class OAuthInProgressError(Exception):
|
||||
"""Another OAuth attempt is already running."""
|
||||
|
||||
|
||||
class OAuthCallbackError(Exception):
|
||||
"""The OAuth callback returned an error (HF denied, port stolen, etc.)."""
|
||||
|
||||
|
||||
# --- PKCE primitives ------------------------------------------------------ #
|
||||
|
||||
|
||||
def _make_pkce() -> tuple[str, str, str]:
|
||||
"""Return ``(verifier, challenge, state)``.
|
||||
|
||||
Verifier never leaves this process. Challenge and state travel
|
||||
through the user's browser. State is checked on the callback to
|
||||
prevent a malicious cross-origin redirect from injecting a token.
|
||||
"""
|
||||
verifier = secrets.token_urlsafe(64)
|
||||
challenge = (
|
||||
base64.urlsafe_b64encode(hashlib.sha256(verifier.encode("ascii")).digest())
|
||||
.rstrip(b"=")
|
||||
.decode("ascii")
|
||||
)
|
||||
state = secrets.token_urlsafe(32)
|
||||
return verifier, challenge, state
|
||||
|
||||
|
||||
def _build_authorize_url(challenge: str, state: str) -> str:
|
||||
from urllib.parse import urlencode
|
||||
|
||||
params = {
|
||||
"client_id": HF_CLIENT_ID,
|
||||
"redirect_uri": REDIRECT_URI,
|
||||
"response_type": "code",
|
||||
"scope": SCOPE,
|
||||
"state": state,
|
||||
"code_challenge": challenge,
|
||||
"code_challenge_method": "S256",
|
||||
}
|
||||
return f"{AUTHORIZE_URL}?{urlencode(params)}"
|
||||
|
||||
|
||||
# --- Token exchange ------------------------------------------------------- #
|
||||
|
||||
|
||||
async def _exchange_code(code: str, verifier: str) -> Token:
|
||||
"""Trade the authorization code for an access+refresh token pair."""
|
||||
data = {
|
||||
"grant_type": "authorization_code",
|
||||
"code": code,
|
||||
"redirect_uri": REDIRECT_URI,
|
||||
"client_id": HF_CLIENT_ID,
|
||||
"code_verifier": verifier,
|
||||
}
|
||||
session = await get_session()
|
||||
async with session.post(TOKEN_URL, data=data, timeout=_TOKEN_REQUEST_TIMEOUT) as resp:
|
||||
resp.raise_for_status()
|
||||
body = await resp.json()
|
||||
return Token(
|
||||
access_token=body["access_token"],
|
||||
refresh_token=body.get("refresh_token"),
|
||||
expires_at=time.time() + float(body.get("expires_in", 3600)),
|
||||
scope=body.get("scope", SCOPE),
|
||||
)
|
||||
|
||||
|
||||
async def refresh_access_token(refresh_token: str) -> Token:
|
||||
"""Trade a refresh_token for a new access (+ possibly refresh) token."""
|
||||
data = {
|
||||
"grant_type": "refresh_token",
|
||||
"refresh_token": refresh_token,
|
||||
"client_id": HF_CLIENT_ID,
|
||||
}
|
||||
session = await get_session()
|
||||
async with session.post(TOKEN_URL, data=data, timeout=_TOKEN_REQUEST_TIMEOUT) as resp:
|
||||
resp.raise_for_status()
|
||||
body = await resp.json()
|
||||
return Token(
|
||||
access_token=body["access_token"],
|
||||
# If HF doesn't rotate refresh tokens, keep using the existing one.
|
||||
refresh_token=body.get("refresh_token", refresh_token),
|
||||
expires_at=time.time() + float(body.get("expires_in", 3600)),
|
||||
scope=body.get("scope", SCOPE),
|
||||
)
|
||||
|
||||
|
||||
# --- Callback server ------------------------------------------------------ #
|
||||
|
||||
|
||||
async def start_login_flow() -> str:
|
||||
"""Begin one OAuth attempt: spawn the callback server, return the URL.
|
||||
|
||||
Returns the URL the frontend should open in a new tab. Raises
|
||||
``OAuthInProgressError`` if another attempt is already running.
|
||||
The callback server runs in the background until the user
|
||||
completes consent or until ``CALLBACK_TIMEOUT_SECS`` elapses;
|
||||
either way the lock + port are released afterward.
|
||||
"""
|
||||
if not _OAUTH_LOCK.acquire(blocking=False):
|
||||
raise OAuthInProgressError()
|
||||
|
||||
try:
|
||||
verifier, challenge, state = _make_pkce()
|
||||
authorize_url = _build_authorize_url(challenge, state)
|
||||
ready: asyncio.Future[None] = asyncio.get_event_loop().create_future()
|
||||
except BaseException:
|
||||
# Failed before handing the lock to the callback-server task: release it
|
||||
# here. (Once the task is spawned, it owns releasing the lock.)
|
||||
_OAUTH_LOCK.release()
|
||||
raise
|
||||
|
||||
asyncio.create_task(_run_callback_server(verifier, state, ready))
|
||||
# Don't return the URL until the callback server is actually bound and
|
||||
# listening — otherwise HF could redirect to a port nothing is serving and
|
||||
# the login would silently dead-end. ``ready`` raises if the bind failed.
|
||||
await ready
|
||||
return authorize_url
|
||||
|
||||
|
||||
async def _run_callback_server(
|
||||
verifier: str, expected_state: str, ready: "asyncio.Future[None]"
|
||||
) -> None:
|
||||
"""Listen for HF's redirect once, capture the token, then shut down.
|
||||
|
||||
Signals ``ready`` once the port is bound (or with an exception if the bind
|
||||
fails), so ``start_login_flow`` only hands back a URL on a live server.
|
||||
"""
|
||||
received: asyncio.Future[Token] = asyncio.get_event_loop().create_future()
|
||||
|
||||
async def handler(request: web.Request) -> web.Response:
|
||||
try:
|
||||
if request.query.get("state") != expected_state:
|
||||
return web.Response(status=400, text="state mismatch")
|
||||
err = request.query.get("error")
|
||||
if err:
|
||||
received.set_exception(OAuthCallbackError(f"HF returned: {err}"))
|
||||
return web.Response(status=400, text=f"OAuth error: {err}")
|
||||
code = request.query.get("code")
|
||||
if not code:
|
||||
return web.Response(status=400, text="missing code")
|
||||
tok = await _exchange_code(code, verifier)
|
||||
if not received.done():
|
||||
received.set_result(tok)
|
||||
return web.Response(
|
||||
content_type="text/html",
|
||||
text=(
|
||||
"<html><body style='font-family:sans-serif;padding:40px'>"
|
||||
"<h2>HuggingFace login successful</h2>"
|
||||
"<p>You can close this tab and return to ComfyUI.</p>"
|
||||
"</body></html>"
|
||||
),
|
||||
)
|
||||
except Exception as exc:
|
||||
if not received.done():
|
||||
received.set_exception(exc)
|
||||
return web.Response(status=500, text=str(exc))
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_get(CALLBACK_PATH, handler)
|
||||
runner = web.AppRunner(app)
|
||||
try:
|
||||
await runner.setup()
|
||||
site = web.TCPSite(runner, CALLBACK_HOST, CALLBACK_PORT, reuse_address=True)
|
||||
await site.start()
|
||||
except Exception as e:
|
||||
# Couldn't bind the callback port (commonly already in use). Tell the
|
||||
# waiting start_login_flow via ``ready`` so it surfaces a clear error
|
||||
# instead of returning a dead URL, and release the lock for next time.
|
||||
logging.warning("[hf_auth] could not start callback server: %s", e)
|
||||
if not ready.done():
|
||||
ready.set_exception(
|
||||
OAuthCallbackError(f"could not bind callback port {CALLBACK_PORT}: {e}")
|
||||
)
|
||||
_OAUTH_LOCK.release()
|
||||
return
|
||||
|
||||
# Bound and listening — now it's safe for start_login_flow to return the URL.
|
||||
if not ready.done():
|
||||
ready.set_result(None)
|
||||
|
||||
try:
|
||||
token = await asyncio.wait_for(received, timeout=CALLBACK_TIMEOUT_SECS)
|
||||
except asyncio.TimeoutError:
|
||||
logging.info("[hf_auth] OAuth login timed out after %ds", CALLBACK_TIMEOUT_SECS)
|
||||
return
|
||||
except OAuthCallbackError as e:
|
||||
logging.warning("[hf_auth] OAuth callback error: %s", e)
|
||||
return
|
||||
except Exception as e:
|
||||
logging.warning("[hf_auth] unexpected OAuth failure: %s", e)
|
||||
return
|
||||
else:
|
||||
HF_AUTH_STORE.set_token(token)
|
||||
logging.info("[hf_auth] OAuth login complete")
|
||||
finally:
|
||||
await runner.cleanup()
|
||||
if _OAUTH_LOCK.locked():
|
||||
_OAUTH_LOCK.release()
|
||||
|
||||
|
||||
def is_login_in_progress() -> bool:
|
||||
"""True iff a callback server is currently bound + waiting."""
|
||||
return _OAUTH_LOCK.locked()
|
||||
|
||||
|
||||
# Re-export for callers that only want the URL builder (e.g. tests).
|
||||
__all__ = [
|
||||
"start_login_flow",
|
||||
"refresh_access_token",
|
||||
"is_login_in_progress",
|
||||
"OAuthInProgressError",
|
||||
"CALLBACK_TIMEOUT_SECS",
|
||||
]
|
||||
@@ -0,0 +1,94 @@
|
||||
"""On-disk persistence for the HuggingFace OAuth token.
|
||||
|
||||
The token shape mirrors what HF returns on the token exchange: an
|
||||
``access_token``, an optional ``refresh_token``, the absolute epoch at
|
||||
which the access token expires, and the granted scope. We persist
|
||||
this so logging in once survives ComfyUI restarts under the internal
|
||||
``__hf_auth`` system-user directory; the file is mode ``0600`` so only
|
||||
the OS user can read it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import stat
|
||||
import time
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Optional
|
||||
|
||||
import folder_paths
|
||||
|
||||
|
||||
# Treat a token as expired this many seconds before its server-reported
|
||||
# ``expires_at`` so we don't try to use a token mid-request only for it
|
||||
# to flip stale between auth_check and the actual GET.
|
||||
EXPIRY_BUFFER_SECS = 60
|
||||
|
||||
TOKEN_FILENAME = "hf_auth_token.json"
|
||||
|
||||
|
||||
@dataclass
|
||||
class Token:
|
||||
"""One OAuth token + the metadata we need to use it."""
|
||||
access_token: str
|
||||
refresh_token: Optional[str]
|
||||
expires_at: float # absolute epoch seconds
|
||||
scope: str = ""
|
||||
|
||||
def is_valid(self) -> bool:
|
||||
"""True iff we can use this token right now."""
|
||||
return (
|
||||
bool(self.access_token)
|
||||
and (self.expires_at - time.time() > EXPIRY_BUFFER_SECS)
|
||||
)
|
||||
|
||||
|
||||
def _token_dir() -> str:
|
||||
return folder_paths.get_system_user_directory("hf_auth")
|
||||
|
||||
|
||||
def _token_path() -> str:
|
||||
return os.path.join(_token_dir(), TOKEN_FILENAME)
|
||||
|
||||
|
||||
def load_token() -> Optional[Token]:
|
||||
"""Read the persisted token, returning ``None`` if absent or corrupt."""
|
||||
path = _token_path()
|
||||
if not os.path.exists(path):
|
||||
return None
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
return Token(**data)
|
||||
except (OSError, json.JSONDecodeError, TypeError) as e:
|
||||
logging.warning("[hf_auth] could not load token at %s: %s", path, e)
|
||||
return None
|
||||
|
||||
|
||||
def save_token(token: Token) -> None:
|
||||
"""Atomically write the token with 0600 permissions."""
|
||||
path = _token_path()
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
tmp = path + ".tmp"
|
||||
fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
json.dump(asdict(token), f)
|
||||
os.replace(tmp, path)
|
||||
try:
|
||||
os.chmod(path, stat.S_IRUSR | stat.S_IWUSR)
|
||||
except OSError as e:
|
||||
# On Windows / weird filesystems chmod may be a no-op; not fatal.
|
||||
logging.debug("[hf_auth] chmod 0600 on %s failed: %s", path, e)
|
||||
|
||||
|
||||
def delete_token() -> None:
|
||||
"""Remove the persisted token; no-op if it doesn't exist."""
|
||||
path = _token_path()
|
||||
try:
|
||||
os.remove(path)
|
||||
except FileNotFoundError:
|
||||
pass
|
||||
except OSError as e:
|
||||
logging.warning("[hf_auth] could not remove token at %s: %s", path, e)
|
||||
@@ -0,0 +1,41 @@
|
||||
"""Parsers for the ``huggingface.co`` URL shape we accept in workflows.
|
||||
|
||||
The download API accepts URLs of the form
|
||||
``https://huggingface.co/<org>/<repo>/resolve/<rev>/<path/to/file>``.
|
||||
We need to recover ``<org>/<repo>`` (the *repo_id*) from such URLs for
|
||||
``huggingface_hub`` API calls (notably ``HfApi.auth_check``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Optional
|
||||
from urllib.parse import urlparse
|
||||
|
||||
_HF_HOST = "huggingface.co"
|
||||
|
||||
|
||||
def is_hf_url(url: str) -> bool:
|
||||
"""Cheap host check — does this URL point at huggingface.co?"""
|
||||
try:
|
||||
return urlparse(url).hostname == _HF_HOST
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def repo_id_from_url(url: str) -> Optional[str]:
|
||||
"""Extract ``<org>/<repo>`` from an HF model file URL.
|
||||
|
||||
Returns ``None`` if the URL isn't on huggingface.co or doesn't look
|
||||
like a model-file URL. The expected shape is
|
||||
``/<org>/<repo>/resolve/<rev>/<path>`` — anything else
|
||||
(datasets, spaces, /tree/, /blob/, …) we treat as out of scope here.
|
||||
"""
|
||||
if not is_hf_url(url):
|
||||
return None
|
||||
parts = urlparse(url).path.lstrip("/").split("/")
|
||||
if len(parts) < 4 or parts[2] != "resolve":
|
||||
return None
|
||||
org, repo = parts[0], parts[1]
|
||||
if not org or not repo:
|
||||
return None
|
||||
return f"{org}/{repo}"
|
||||
@@ -0,0 +1,63 @@
|
||||
"""Lazy module-level aiohttp ClientSession.
|
||||
|
||||
A single shared session means TLS handshakes are reused across HEAD probes
|
||||
and the subsequent GETs to the same host (HuggingFace is the dominant
|
||||
case), which is a noticeable speedup on cold connections.
|
||||
|
||||
We deliberately don't close the session at process exit — aiohttp's
|
||||
warning about unclosed sessions is benign at shutdown, and adding atexit
|
||||
plumbing buys nothing because the OS reclaims the sockets anyway. The
|
||||
session lifetime is the lifetime of the Python process.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import ssl
|
||||
from typing import Optional
|
||||
|
||||
import aiohttp
|
||||
import certifi
|
||||
|
||||
|
||||
# Larger per-host pool than aiohttp's default (=100 total / =0 per host)
|
||||
# so concurrent gated probes + a download to the same host don't queue.
|
||||
_CONNECTOR_LIMIT_PER_HOST = 8
|
||||
|
||||
_session: Optional[aiohttp.ClientSession] = None
|
||||
_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def ssl_context() -> ssl.SSLContext:
|
||||
"""TLS context pinned to certifi's CA bundle.
|
||||
aiohttp's default context uses the OS trust store, which isn't wired up
|
||||
on some Python installs (python.org macOS, slim containers) — there TLS
|
||||
to huggingface.co fails with CERTIFICATE_VERIFY_FAILED.
|
||||
"""
|
||||
return ssl.create_default_context(cafile=certifi.where())
|
||||
|
||||
|
||||
async def get_session() -> aiohttp.ClientSession:
|
||||
"""Return the shared session, creating it on first call."""
|
||||
global _session
|
||||
if _session is not None and not _session.closed:
|
||||
return _session
|
||||
async with _lock:
|
||||
if _session is None or _session.closed:
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit_per_host=_CONNECTOR_LIMIT_PER_HOST,
|
||||
ssl=ssl_context(),
|
||||
)
|
||||
_session = aiohttp.ClientSession(connector=connector)
|
||||
return _session
|
||||
|
||||
|
||||
def parse_content_length(value: Optional[str]) -> Optional[int]:
|
||||
"""Parse a byte-count header value, or None if absent/malformed/negative."""
|
||||
if not value:
|
||||
return None
|
||||
try:
|
||||
n = int(value)
|
||||
except ValueError:
|
||||
return None
|
||||
return n if n >= 0 else None
|
||||
@@ -0,0 +1,111 @@
|
||||
"""Path resolution for model downloads.
|
||||
|
||||
Model identifiers used across the download API are *relative destination
|
||||
paths* of the form ``<directory>/<filename>`` (e.g. ``loras/my_lora.safetensors``).
|
||||
This module turns one of those identifiers into an absolute on-disk path
|
||||
under one of ComfyUI's registered model folders, while rejecting unknown
|
||||
folders, path traversal, and other ill-formed inputs.
|
||||
"""
|
||||
|
||||
import os
|
||||
import re
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import folder_paths
|
||||
|
||||
|
||||
# Constrain components so a model_id can never escape its target directory.
|
||||
# - directory: a single path segment of safe chars
|
||||
# - filename: a single path segment of safe chars, must end with a model extension
|
||||
_SEGMENT_RE = re.compile(r"^[A-Za-z0-9._-]+$")
|
||||
|
||||
# Destination filename must name a model file (same set as the URL allowlist),
|
||||
# so a download can't land as e.g. ``foo.txt`` that ComfyUI won't recognise.
|
||||
_MODEL_EXTENSIONS = (".safetensors", ".sft", ".ckpt", ".pth", ".pt")
|
||||
|
||||
# Distinctive temp suffix so the startup orphan-sweep only removes files THIS
|
||||
# subsystem created — never unrelated ``*.tmp`` files in the model dirs.
|
||||
_TMP_SUFFIX = ".comfy-download.tmp"
|
||||
|
||||
|
||||
class InvalidModelId(ValueError):
|
||||
"""Raised when a model_id is syntactically invalid or refers to an
|
||||
unknown model folder."""
|
||||
|
||||
|
||||
def parse_model_id(model_id: str) -> Tuple[str, str]:
|
||||
"""Split ``<directory>/<filename>`` and validate both components.
|
||||
|
||||
Returns ``(directory, filename)``. Raises ``InvalidModelId`` on
|
||||
malformed input. Does NOT touch the filesystem.
|
||||
"""
|
||||
if not isinstance(model_id, str) or "/" not in model_id:
|
||||
raise InvalidModelId(f"model_id must be '<directory>/<filename>', got {model_id!r}")
|
||||
directory, _, filename = model_id.partition("/")
|
||||
if "/" in filename or not directory or not filename:
|
||||
raise InvalidModelId(f"model_id must be exactly one '/' separator, got {model_id!r}")
|
||||
if not _SEGMENT_RE.match(directory):
|
||||
raise InvalidModelId(f"invalid directory segment {directory!r}")
|
||||
if not _SEGMENT_RE.match(filename):
|
||||
raise InvalidModelId(f"invalid filename segment {filename!r}")
|
||||
if not filename.endswith(_MODEL_EXTENSIONS):
|
||||
raise InvalidModelId(
|
||||
f"filename must end with a model extension {_MODEL_EXTENSIONS}, got {filename!r}"
|
||||
)
|
||||
if directory not in folder_paths.folder_names_and_paths:
|
||||
raise InvalidModelId(f"unknown model folder {directory!r}")
|
||||
return directory, filename
|
||||
|
||||
|
||||
def resolve_existing(model_id: str) -> Optional[str]:
|
||||
"""Return the absolute path of an installed model, or None if missing.
|
||||
|
||||
Honours ``extra_model_paths.yaml`` transparently via
|
||||
``folder_paths.get_full_path``.
|
||||
"""
|
||||
directory, filename = parse_model_id(model_id)
|
||||
return folder_paths.get_full_path(directory, filename)
|
||||
|
||||
|
||||
def resolve_destination(model_id: str, epoch: int = 0) -> Tuple[str, str]:
|
||||
"""Return ``(final_path, tmp_path)`` for a download.
|
||||
|
||||
Downloads land at the first registered path for the model's directory
|
||||
(the "primary" location). The temp sibling is the write target, atomically
|
||||
renamed onto ``final_path`` on success.
|
||||
|
||||
``tmp_path`` embeds the session ``epoch`` so a cancel+retry of the same
|
||||
model never shares a temp path between the old (cancelling) worker and the
|
||||
new attempt — otherwise the old worker's rollback could delete the new
|
||||
worker's in-progress file. The distinctive suffix scopes the orphan sweep.
|
||||
"""
|
||||
directory, filename = parse_model_id(model_id)
|
||||
roots = folder_paths.get_folder_paths(directory)
|
||||
if not roots:
|
||||
raise InvalidModelId(f"no on-disk path registered for folder {directory!r}")
|
||||
root = roots[0]
|
||||
final_path = os.path.join(root, filename)
|
||||
tmp_path = f"{final_path}.{epoch}{_TMP_SUFFIX}"
|
||||
return final_path, tmp_path
|
||||
|
||||
|
||||
def iter_all_tmp_paths():
|
||||
"""Yield this subsystem's temp files under every registered model folder.
|
||||
|
||||
Matches only our distinctive ``_TMP_SUFFIX`` (not every ``*.tmp``) so the
|
||||
startup orphan-sweep can't delete temp files created by other tools.
|
||||
"""
|
||||
seen_roots: set[str] = set()
|
||||
for directory in folder_paths.folder_names_and_paths.keys():
|
||||
for root in folder_paths.get_folder_paths(directory):
|
||||
if root in seen_roots or not os.path.isdir(root):
|
||||
continue
|
||||
seen_roots.add(root)
|
||||
try:
|
||||
for entry in os.scandir(root):
|
||||
if entry.is_file() and entry.name.endswith(_TMP_SUFFIX):
|
||||
yield entry.path
|
||||
except OSError:
|
||||
# Folder might be unreadable / missing on certain mounts —
|
||||
# not fatal, just skip it.
|
||||
continue
|
||||
+26
-2
@@ -50,21 +50,45 @@ class ModelFileManager:
|
||||
@routes.get("/experiment/models/preview/{folder}/{path_index}/{filename:.*}")
|
||||
async def get_model_preview(request):
|
||||
folder_name = request.match_info.get("folder", None)
|
||||
path_index = int(request.match_info.get("path_index", None))
|
||||
filename = request.match_info.get("filename", None)
|
||||
|
||||
if folder_name not in folder_paths.folder_names_and_paths:
|
||||
return web.Response(status=404)
|
||||
|
||||
# The "{filename:.*}" capture also matches the empty string, which
|
||||
# would resolve to the folder itself; reject it explicitly.
|
||||
if not filename:
|
||||
return web.Response(status=400)
|
||||
|
||||
try:
|
||||
path_index = int(request.match_info.get("path_index", None))
|
||||
except (TypeError, ValueError):
|
||||
return web.Response(status=400)
|
||||
|
||||
folders = folder_paths.folder_names_and_paths[folder_name]
|
||||
if path_index < 0 or path_index >= len(folders[0]):
|
||||
return web.Response(status=404)
|
||||
folder = folders[0][path_index]
|
||||
full_filename = os.path.join(folder, filename)
|
||||
full_filename = os.path.normpath(os.path.join(folder, filename))
|
||||
|
||||
# Prevent path traversal: the requested file must stay within the
|
||||
# configured model folder. `filename` is an unrestricted ".*" capture,
|
||||
# so values like "../../../../etc/passwd" would otherwise escape it.
|
||||
if not folder_paths.is_within_directory(folder, full_filename):
|
||||
return web.Response(status=403)
|
||||
|
||||
previews = self.get_model_previews(full_filename)
|
||||
default_preview = previews[0] if len(previews) > 0 else None
|
||||
if default_preview is None or (isinstance(default_preview, str) and not os.path.isfile(default_preview)):
|
||||
return web.Response(status=404)
|
||||
|
||||
# The preview is selected by a glob inside get_model_previews, so a
|
||||
# companion file (e.g. "model.preview.png") could itself be a symlink
|
||||
# resolving outside the model folder. Re-validate the file actually
|
||||
# opened: is_within_directory realpaths it, catching symlink escape.
|
||||
if isinstance(default_preview, str) and not folder_paths.is_within_directory(folder, default_preview):
|
||||
return web.Response(status=403)
|
||||
|
||||
try:
|
||||
with Image.open(default_preview) as img:
|
||||
img_bytes = BytesIO()
|
||||
|
||||
+15
-1
@@ -6,6 +6,7 @@ import glob
|
||||
import shutil
|
||||
import logging
|
||||
import tempfile
|
||||
import mimetypes
|
||||
from aiohttp import web
|
||||
from urllib import parse
|
||||
from comfy.cli_args import args
|
||||
@@ -336,7 +337,20 @@ class UserManager():
|
||||
if not isinstance(path, str):
|
||||
return path
|
||||
|
||||
return web.FileResponse(path)
|
||||
# User data files are arbitrary user-supplied content and are never
|
||||
# meant to render inline. Disable MIME sniffing and force a download
|
||||
# so uploaded markup/scripts can't execute in the app origin (stored
|
||||
# XSS). Content-Disposition: attachment is the load-bearing guard;
|
||||
# the content-type override and nosniff are defence in depth.
|
||||
content_type = mimetypes.guess_type(path)[0] or 'application/octet-stream'
|
||||
if folder_paths.is_dangerous_content_type(content_type):
|
||||
content_type = 'application/octet-stream'
|
||||
|
||||
return web.FileResponse(path, headers={
|
||||
"Content-Type": content_type,
|
||||
"X-Content-Type-Options": "nosniff",
|
||||
"Content-Disposition": "attachment",
|
||||
})
|
||||
|
||||
@routes.post("/userdata/{file}")
|
||||
async def post_userdata(request):
|
||||
|
||||
@@ -240,6 +240,7 @@ database_default_path = os.path.abspath(
|
||||
)
|
||||
parser.add_argument("--database-url", type=str, default=f"sqlite:///{database_default_path}", help="Specify the database URL, e.g. for an in-memory database you can use 'sqlite:///:memory:'.")
|
||||
parser.add_argument("--enable-assets", action="store_true", help="Enable the assets system (API routes, database synchronization, and background scanning).")
|
||||
parser.add_argument("--enable-asset-hashing", action="store_true", help="Compute blake3 content hashes when scanning assets. Hashing enables future asset-portability features (deduplication, cross-machine model resolution) but adds startup cost and per-output cost on large models directories. Off by default; enable to opt in.")
|
||||
parser.add_argument("--feature-flag", type=str, action='append', default=[], metavar="KEY[=VALUE]", help="Set a server feature flag. Use KEY=VALUE to set an explicit value, or bare KEY to set it to true. Can be specified multiple times. Boolean values (true/false) and numbers are auto-converted. Examples: --feature-flag show_signin_button=true or --feature-flag show_signin_button")
|
||||
parser.add_argument("--list-feature-flags", action="store_true", help="Print the registry of known CLI-settable feature flags as JSON and exit.")
|
||||
|
||||
|
||||
+409
-57
@@ -8,6 +8,8 @@ from abc import ABC, abstractmethod
|
||||
import logging
|
||||
import comfy.model_management
|
||||
import comfy.patcher_extension
|
||||
import comfy.utils
|
||||
import comfy.conds
|
||||
if TYPE_CHECKING:
|
||||
from comfy.model_base import BaseModel
|
||||
from comfy.model_patcher import ModelPatcher
|
||||
@@ -51,12 +53,18 @@ class ContextHandlerABC(ABC):
|
||||
|
||||
|
||||
class IndexListContextWindow(ContextWindowABC):
|
||||
def __init__(self, index_list: list[int], dim: int=0, total_frames: int=0):
|
||||
def __init__(self, index_list: list[int], dim: int=0, total_frames: int=0, modality_windows: dict=None, context_overlap: int=0):
|
||||
self.index_list = index_list
|
||||
self.context_length = len(index_list)
|
||||
self.context_overlap = context_overlap
|
||||
self.dim = dim
|
||||
self.total_frames = total_frames
|
||||
self.center_ratio = (min(index_list) + max(index_list)) / (2 * total_frames)
|
||||
self.modality_windows = modality_windows # dict of {mod_idx: IndexListContextWindow}
|
||||
self.guide_frames_indices: list[int] = []
|
||||
self.guide_overlap_info: list[tuple[int, int]] = []
|
||||
self.guide_kf_local_positions: list[int] = []
|
||||
self.guide_downscale_factors: list[int] = []
|
||||
|
||||
def get_tensor(self, full: torch.Tensor, device=None, dim=None, retain_index_list=[]) -> torch.Tensor:
|
||||
if dim is None:
|
||||
@@ -85,6 +93,11 @@ class IndexListContextWindow(ContextWindowABC):
|
||||
region_idx = int(self.center_ratio * num_regions)
|
||||
return min(max(region_idx, 0), num_regions - 1)
|
||||
|
||||
def get_window_for_modality(self, modality_idx: int) -> 'IndexListContextWindow':
|
||||
if modality_idx == 0:
|
||||
return self
|
||||
return self.modality_windows[modality_idx]
|
||||
|
||||
|
||||
class IndexListCallbacks:
|
||||
EVALUATE_CONTEXT_WINDOWS = "evaluate_context_windows"
|
||||
@@ -148,6 +161,172 @@ def slice_cond(cond_value, window: IndexListContextWindow, x_in: torch.Tensor, d
|
||||
return cond_value._copy_with(sliced)
|
||||
|
||||
|
||||
def compute_guide_overlap(guide_entries: list[dict], keyframe_idxs: torch.Tensor, temporal_downscale_ratio: int, window_index_list: list[int]):
|
||||
"""Compute which concatenated guide frames overlap with a context window.
|
||||
|
||||
Each guide's latent-space start is derived from its first token's pixel-t-start
|
||||
in keyframe_idxs (shape (B, [t,h,w], num_tokens, [start, end])), divided by the
|
||||
model's temporal_downscale_ratio.
|
||||
|
||||
Args:
|
||||
guide_entries: list of guide_attention_entry dicts
|
||||
keyframe_idxs: per-token pixel coords cond tensor for the modality
|
||||
temporal_downscale_ratio: model's pixel-to-latent temporal compression ratio
|
||||
window_index_list: the window's frame indices into the video portion
|
||||
|
||||
Returns:
|
||||
suffix_indices: indices into the guide_frames tensor for frame selection
|
||||
overlap_info: list of (entry_idx, overlap_count) for guide_attention_entries adjustment
|
||||
kf_local_positions: window-local frame positions for keyframe_idxs regeneration
|
||||
total_overlap: total number of overlapping guide frames
|
||||
"""
|
||||
window_set = set(window_index_list)
|
||||
window_list = list(window_index_list)
|
||||
suffix_indices = []
|
||||
overlap_info = []
|
||||
kf_local_positions = []
|
||||
suffix_base = 0
|
||||
token_offset = 0
|
||||
|
||||
for entry_idx, entry in enumerate(guide_entries):
|
||||
first_t_pixel = int(keyframe_idxs[0, 0, token_offset, 0].item())
|
||||
latent_start = (first_t_pixel + temporal_downscale_ratio - 1) // temporal_downscale_ratio
|
||||
guide_len = entry["latent_shape"][0]
|
||||
entry_overlap = 0
|
||||
|
||||
for local_offset in range(guide_len):
|
||||
video_pos = latent_start + local_offset
|
||||
if video_pos in window_set:
|
||||
suffix_indices.append(suffix_base + local_offset)
|
||||
kf_local_positions.append(window_list.index(video_pos))
|
||||
entry_overlap += 1
|
||||
|
||||
if entry_overlap > 0:
|
||||
overlap_info.append((entry_idx, entry_overlap))
|
||||
suffix_base += guide_len
|
||||
token_offset += entry["pre_filter_count"]
|
||||
|
||||
return suffix_indices, overlap_info, kf_local_positions, len(suffix_indices)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WindowingState:
|
||||
"""Per-modality context windowing state for each step,
|
||||
built using IndexListContextHandler._build_window_state().
|
||||
For non-multimodal models the lists are length 1
|
||||
"""
|
||||
latents: list[torch.Tensor] # per-modality working latents (guide frames stripped)
|
||||
guide_latents: list[torch.Tensor | None] # per-modality guide frames stripped from latents
|
||||
guide_entries: list[list[dict] | None] # per-modality guide_attention_entry metadata
|
||||
keyframe_idxs: list[torch.Tensor | None] # per-modality keyframe_idxs tensor for guide latent_start derivation
|
||||
latent_shapes: list | None # original packed shapes for unpack/pack (None if not multimodal)
|
||||
dim: int = 0 # primary modality temporal dim for context windowing
|
||||
is_multimodal: bool = False
|
||||
temporal_downscale_ratio: int = 1 # model's pixel-to-latent temporal compression ratio
|
||||
|
||||
def prepare_window(self, window: IndexListContextWindow, model) -> IndexListContextWindow:
|
||||
"""Reformat window for multimodal contexts by deriving per-modality index lists.
|
||||
Non-multimodal contexts return the input window unchanged.
|
||||
"""
|
||||
if not self.is_multimodal:
|
||||
return window
|
||||
|
||||
x = self.latents[0]
|
||||
primary_total = self.latent_shapes[0][self.dim]
|
||||
primary_overlap = window.context_overlap
|
||||
map_shapes = self.latent_shapes
|
||||
if x.size(self.dim) != primary_total:
|
||||
map_shapes = list(self.latent_shapes)
|
||||
video_shape = list(self.latent_shapes[0])
|
||||
video_shape[self.dim] = x.size(self.dim)
|
||||
map_shapes[0] = torch.Size(video_shape)
|
||||
try:
|
||||
per_modality_indices = model.map_context_window_to_modalities(
|
||||
window.index_list, map_shapes, self.dim)
|
||||
except AttributeError:
|
||||
raise NotImplementedError(
|
||||
f"{type(model).__name__} must implement map_context_window_to_modalities for multimodal context windows.")
|
||||
modality_windows = {}
|
||||
for mod_idx in range(1, len(self.latents)):
|
||||
modality_total_frames = self.latents[mod_idx].shape[self.dim]
|
||||
ratio = modality_total_frames / primary_total if primary_total > 0 else 1
|
||||
modality_overlap = max(round(primary_overlap * ratio), 0)
|
||||
modality_windows[mod_idx] = IndexListContextWindow(
|
||||
per_modality_indices[mod_idx], dim=self.dim,
|
||||
total_frames=modality_total_frames,
|
||||
context_overlap=modality_overlap)
|
||||
return IndexListContextWindow(
|
||||
window.index_list, dim=self.dim, total_frames=x.shape[self.dim],
|
||||
modality_windows=modality_windows, context_overlap=primary_overlap)
|
||||
|
||||
def slice_for_window(self, window: IndexListContextWindow, retain_index_list: list[int], device=None) -> tuple[list[torch.Tensor], list[int]]:
|
||||
"""Slice latents for a context window, injecting guide frames where applicable.
|
||||
For multimodal contexts, uses the modality-specific windows derived in prepare_window().
|
||||
"""
|
||||
sliced = []
|
||||
guide_frame_counts = []
|
||||
for idx in range(len(self.latents)):
|
||||
modality_window = window.get_window_for_modality(idx)
|
||||
retain = retain_index_list if idx == 0 else []
|
||||
s = modality_window.get_tensor(self.latents[idx], device, retain_index_list=retain)
|
||||
if self.guide_entries[idx] is not None:
|
||||
s, ng = self._inject_guide_frames(s, modality_window, modality_idx=idx)
|
||||
else:
|
||||
ng = 0
|
||||
sliced.append(s)
|
||||
guide_frame_counts.append(ng)
|
||||
return sliced, guide_frame_counts
|
||||
|
||||
def strip_guide_frames(self, out_per_modality: list[list[torch.Tensor]], guide_frame_counts: list[int], window: IndexListContextWindow):
|
||||
"""Strip injected guide frames from per-cond, per-modality outputs in place."""
|
||||
for idx in range(len(self.latents)):
|
||||
if guide_frame_counts[idx] > 0:
|
||||
window_len = len(window.get_window_for_modality(idx).index_list)
|
||||
for ci in range(len(out_per_modality)):
|
||||
out_per_modality[ci][idx] = out_per_modality[ci][idx].narrow(self.dim, 0, window_len)
|
||||
|
||||
def _inject_guide_frames(self, latent_slice: torch.Tensor, window: IndexListContextWindow, modality_idx: int = 0) -> tuple[torch.Tensor, int]:
|
||||
guide_entries = self.guide_entries[modality_idx]
|
||||
guide_frames = self.guide_latents[modality_idx]
|
||||
keyframe_idxs = self.keyframe_idxs[modality_idx]
|
||||
suffix_idx, overlap_info, kf_local_pos, guide_frame_count = compute_guide_overlap(
|
||||
guide_entries, keyframe_idxs, self.temporal_downscale_ratio, window.index_list)
|
||||
# Shift keyframe positions to account for causal_window_fix anchor occupying sub-pos 0.
|
||||
anchor_idx = getattr(window, 'causal_anchor_index', None)
|
||||
if anchor_idx is not None and anchor_idx >= 0:
|
||||
kf_local_pos = [p + 1 for p in kf_local_pos]
|
||||
window.guide_frames_indices = suffix_idx
|
||||
window.guide_overlap_info = overlap_info
|
||||
window.guide_kf_local_positions = kf_local_pos
|
||||
|
||||
# Derive per-overlap-entry latent_downscale_factor from guide entry latent_shape vs guide frame spatial dims.
|
||||
# guide_frames has full (post-dilation) spatial dims; entry["latent_shape"] has pre-dilation dims.
|
||||
guide_downscale_factors = []
|
||||
if guide_frame_count > 0:
|
||||
full_H = guide_frames.shape[3]
|
||||
for entry_idx, _ in overlap_info:
|
||||
entry_H = guide_entries[entry_idx]["latent_shape"][1]
|
||||
guide_downscale_factors.append(full_H // entry_H)
|
||||
window.guide_downscale_factors = guide_downscale_factors
|
||||
|
||||
if guide_frame_count > 0:
|
||||
idx = tuple([slice(None)] * self.dim + [suffix_idx])
|
||||
return torch.cat([latent_slice, guide_frames[idx]], dim=self.dim), guide_frame_count
|
||||
return latent_slice, 0
|
||||
|
||||
def patch_latent_shapes(self, sub_conds, new_shapes):
|
||||
if not self.is_multimodal:
|
||||
return
|
||||
|
||||
for cond_list in sub_conds:
|
||||
if cond_list is None:
|
||||
continue
|
||||
for cond_dict in cond_list:
|
||||
model_conds = cond_dict.get('model_conds', {})
|
||||
if 'latent_shapes' in model_conds:
|
||||
model_conds['latent_shapes'] = comfy.conds.CONDConstant(new_shapes)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextSchedule:
|
||||
name: str
|
||||
@@ -162,7 +341,7 @@ ContextResults = collections.namedtuple("ContextResults", ['window_idx', 'sub_co
|
||||
class IndexListContextHandler(ContextHandlerABC):
|
||||
def __init__(self, context_schedule: ContextSchedule, fuse_method: ContextFuseMethod, context_length: int=1, context_overlap: int=0, context_stride: int=1,
|
||||
closed_loop: bool=False, dim:int=0, freenoise: bool=False, cond_retain_index_list: list[int]=[], split_conds_to_windows: bool=False,
|
||||
causal_window_fix: bool=True):
|
||||
latent_retain_index_list: list[int]=[], causal_window_fix: bool=True):
|
||||
self.context_schedule = context_schedule
|
||||
self.fuse_method = fuse_method
|
||||
self.context_length = context_length
|
||||
@@ -174,17 +353,118 @@ class IndexListContextHandler(ContextHandlerABC):
|
||||
self.freenoise = freenoise
|
||||
self.cond_retain_index_list = [int(x.strip()) for x in cond_retain_index_list.split(",")] if cond_retain_index_list else []
|
||||
self.split_conds_to_windows = split_conds_to_windows
|
||||
self.latent_retain_index_list = [int(x.strip()) for x in latent_retain_index_list.split(",")] if latent_retain_index_list else []
|
||||
self.causal_window_fix = causal_window_fix
|
||||
|
||||
self.callbacks = {}
|
||||
|
||||
@staticmethod
|
||||
def _get_latent_shapes(conds):
|
||||
for cond_list in conds:
|
||||
if cond_list is None:
|
||||
continue
|
||||
for cond_dict in cond_list:
|
||||
model_conds = cond_dict.get('model_conds', {})
|
||||
if 'latent_shapes' in model_conds:
|
||||
return model_conds['latent_shapes'].cond
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_guide_entries(conds):
|
||||
for cond_list in conds:
|
||||
if cond_list is None:
|
||||
continue
|
||||
for cond_dict in cond_list:
|
||||
model_conds = cond_dict.get('model_conds', {})
|
||||
entries = model_conds.get('guide_attention_entries')
|
||||
if entries is not None and hasattr(entries, 'cond') and entries.cond:
|
||||
return entries.cond
|
||||
return None
|
||||
|
||||
@staticmethod
|
||||
def _get_keyframe_idxs(conds):
|
||||
for cond_list in conds:
|
||||
if cond_list is None:
|
||||
continue
|
||||
for cond_dict in cond_list:
|
||||
model_conds = cond_dict.get('model_conds', {})
|
||||
kf = model_conds.get('keyframe_idxs')
|
||||
if kf is not None and hasattr(kf, 'cond') and kf.cond is not None:
|
||||
return kf.cond
|
||||
return None
|
||||
|
||||
def _apply_freenoise(self, noise: torch.Tensor, conds: list[list[dict]], seed: int) -> torch.Tensor:
|
||||
"""Apply FreeNoise shuffling, scaling context length/overlap per-modality by frame ratio.
|
||||
If guide frames are present on the primary modality, only the video portion is shuffled.
|
||||
"""
|
||||
guide_entries = self._get_guide_entries(conds)
|
||||
guide_count = sum(e["latent_shape"][0] for e in guide_entries) if guide_entries else 0
|
||||
|
||||
latent_shapes = self._get_latent_shapes(conds)
|
||||
if latent_shapes is not None and len(latent_shapes) > 1:
|
||||
modalities = comfy.utils.unpack_latents(noise, latent_shapes)
|
||||
primary_total = latent_shapes[0][self.dim]
|
||||
primary_video_count = modalities[0].size(self.dim) - guide_count
|
||||
apply_freenoise(modalities[0].narrow(self.dim, 0, primary_video_count), self.dim, self.context_length, self.context_overlap, seed)
|
||||
for i in range(1, len(modalities)):
|
||||
mod_total = latent_shapes[i][self.dim]
|
||||
ratio = mod_total / primary_total if primary_total > 0 else 1
|
||||
mod_ctx_len = max(round(self.context_length * ratio), 1)
|
||||
mod_ctx_overlap = max(round(self.context_overlap * ratio), 0)
|
||||
modalities[i] = apply_freenoise(modalities[i], self.dim, mod_ctx_len, mod_ctx_overlap, seed)
|
||||
noise, _ = comfy.utils.pack_latents(modalities)
|
||||
return noise
|
||||
video_count = noise.size(self.dim) - guide_count
|
||||
apply_freenoise(noise.narrow(self.dim, 0, video_count), self.dim, self.context_length, self.context_overlap, seed)
|
||||
return noise
|
||||
|
||||
def _build_window_state(self, x_in: torch.Tensor, conds: list[list[dict]], model: BaseModel) -> WindowingState:
|
||||
"""Build windowing state for the current step, including unpacking latents and extracting guide frame info from conds."""
|
||||
latent_shapes = self._get_latent_shapes(conds)
|
||||
is_multimodal = latent_shapes is not None and len(latent_shapes) > 1
|
||||
unpacked_latents = comfy.utils.unpack_latents(x_in, latent_shapes) if is_multimodal else [x_in]
|
||||
|
||||
unpacked_latents_list = list(unpacked_latents)
|
||||
guide_latents_list = [None] * len(unpacked_latents)
|
||||
guide_entries_list = [None] * len(unpacked_latents)
|
||||
keyframe_idxs_list = [None] * len(unpacked_latents)
|
||||
|
||||
extracted_guide_entries = self._get_guide_entries(conds)
|
||||
extracted_keyframe_idxs = self._get_keyframe_idxs(conds)
|
||||
|
||||
# Strip guide frames (only from first modality for now)
|
||||
if extracted_guide_entries is not None:
|
||||
guide_count = sum(e["latent_shape"][0] for e in extracted_guide_entries)
|
||||
if guide_count > 0:
|
||||
x = unpacked_latents[0]
|
||||
latent_count = x.size(self.dim) - guide_count
|
||||
unpacked_latents_list[0] = x.narrow(self.dim, 0, latent_count)
|
||||
guide_latents_list[0] = x.narrow(self.dim, latent_count, guide_count)
|
||||
guide_entries_list[0] = extracted_guide_entries
|
||||
keyframe_idxs_list[0] = extracted_keyframe_idxs
|
||||
|
||||
|
||||
return WindowingState(
|
||||
latents=unpacked_latents_list,
|
||||
guide_latents=guide_latents_list,
|
||||
guide_entries=guide_entries_list,
|
||||
keyframe_idxs=keyframe_idxs_list,
|
||||
latent_shapes=latent_shapes,
|
||||
dim=self.dim,
|
||||
is_multimodal=is_multimodal,
|
||||
temporal_downscale_ratio=model.latent_format.temporal_downscale_ratio)
|
||||
|
||||
def should_use_context(self, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]) -> bool:
|
||||
# for now, assume first dim is batch - should have stored on BaseModel in actual implementation
|
||||
if x_in.size(self.dim) > self.context_length:
|
||||
logging.info(f"Using context windows {self.context_length} with overlap {self.context_overlap} for {x_in.size(self.dim)} frames.")
|
||||
window_state = self._build_window_state(x_in, conds, model) # build window_state to check frame counts, will be built again in execute
|
||||
total_frame_count = window_state.latents[0].size(self.dim)
|
||||
if total_frame_count > self.context_length:
|
||||
logging.info(f"\nUsing context windows: Context length {self.context_length} with overlap {self.context_overlap} for {total_frame_count} frames.")
|
||||
if self.cond_retain_index_list:
|
||||
logging.info(f"Retaining original cond for indexes: {self.cond_retain_index_list}")
|
||||
if self.latent_retain_index_list:
|
||||
logging.info(f"Retaining original latent for indexes: {self.latent_retain_index_list}")
|
||||
return True
|
||||
logging.info(f"\nNot using context windows since context length ({self.context_length}) exceeds input frames ({total_frame_count}).")
|
||||
return False
|
||||
|
||||
def prepare_control_objects(self, control: ControlBase, device=None) -> ControlBase:
|
||||
@@ -275,7 +555,9 @@ class IndexListContextHandler(ContextHandlerABC):
|
||||
return resized_cond
|
||||
|
||||
def set_step(self, timestep: torch.Tensor, model_options: dict[str]):
|
||||
mask = torch.isclose(model_options["transformer_options"]["sample_sigmas"], timestep[0], rtol=0.0001)
|
||||
sample_sigmas = model_options["transformer_options"]["sample_sigmas"]
|
||||
current_timestep = timestep[0].to(sample_sigmas.dtype)
|
||||
mask = torch.isclose(sample_sigmas, current_timestep, rtol=0.0001)
|
||||
matches = torch.nonzero(mask)
|
||||
if torch.numel(matches) == 0:
|
||||
return # substep from multi-step sampler: keep self._step from the last full step
|
||||
@@ -284,54 +566,98 @@ class IndexListContextHandler(ContextHandlerABC):
|
||||
def get_context_windows(self, model: BaseModel, x_in: torch.Tensor, model_options: dict[str]) -> list[IndexListContextWindow]:
|
||||
full_length = x_in.size(self.dim) # TODO: choose dim based on model
|
||||
context_windows = self.context_schedule.func(full_length, self, model_options)
|
||||
context_windows = [IndexListContextWindow(window, dim=self.dim, total_frames=full_length) for window in context_windows]
|
||||
context_windows = [IndexListContextWindow(window, dim=self.dim, total_frames=full_length, context_overlap=self.context_overlap) for window in context_windows]
|
||||
return context_windows
|
||||
|
||||
def execute(self, calc_cond_batch: Callable, model: BaseModel, conds: list[list[dict]], x_in: torch.Tensor, timestep: torch.Tensor, model_options: dict[str]):
|
||||
self._model = model
|
||||
self.set_step(timestep, model_options)
|
||||
context_windows = self.get_context_windows(model, x_in, model_options)
|
||||
enumerated_context_windows = list(enumerate(context_windows))
|
||||
|
||||
conds_final = [torch.zeros_like(x_in) for _ in conds]
|
||||
window_state = self._build_window_state(x_in, conds, model)
|
||||
num_modalities = len(window_state.latents)
|
||||
|
||||
context_windows = self.get_context_windows(model, window_state.latents[0], model_options)
|
||||
enumerated_context_windows = list(enumerate(context_windows))
|
||||
total_windows = len(enumerated_context_windows)
|
||||
|
||||
# Initialize per-modality accumulators (length 1 for single-modality)
|
||||
accum = [[torch.zeros_like(m) for _ in conds] for m in window_state.latents]
|
||||
if self.fuse_method.name == ContextFuseMethods.RELATIVE:
|
||||
counts_final = [torch.ones(get_shape_for_dim(x_in, self.dim), device=x_in.device) for _ in conds]
|
||||
counts = [[torch.ones(get_shape_for_dim(m, self.dim), device=m.device) for _ in conds] for m in window_state.latents]
|
||||
else:
|
||||
counts_final = [torch.zeros(get_shape_for_dim(x_in, self.dim), device=x_in.device) for _ in conds]
|
||||
biases_final = [([0.0] * x_in.shape[self.dim]) for _ in conds]
|
||||
counts = [[torch.zeros(get_shape_for_dim(m, self.dim), device=m.device) for _ in conds] for m in window_state.latents]
|
||||
biases = [[([0.0] * m.shape[self.dim]) for _ in conds] for m in window_state.latents]
|
||||
|
||||
for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.EXECUTE_START, self.callbacks):
|
||||
callback(self, model, x_in, conds, timestep, model_options)
|
||||
|
||||
# accumulate results from each context window
|
||||
for enum_window in enumerated_context_windows:
|
||||
results = self.evaluate_context_windows(calc_cond_batch, model, x_in, conds, timestep, [enum_window], model_options)
|
||||
results = self.evaluate_context_windows(
|
||||
calc_cond_batch, model, x_in, conds, timestep, [enum_window],
|
||||
model_options, window_state=window_state, total_windows=total_windows)
|
||||
for result in results:
|
||||
self.combine_context_window_results(x_in, result.sub_conds_out, result.sub_conds, result.window, result.window_idx, len(enumerated_context_windows), timestep,
|
||||
conds_final, counts_final, biases_final)
|
||||
# result.sub_conds_out is per-cond, per-modality: list[list[Tensor]]
|
||||
for mod_idx in range(num_modalities):
|
||||
mod_out = [result.sub_conds_out[ci][mod_idx] for ci in range(len(conds))]
|
||||
modality_window = result.window.get_window_for_modality(mod_idx)
|
||||
self.combine_context_window_results(
|
||||
window_state.latents[mod_idx], mod_out, result.sub_conds, modality_window,
|
||||
result.window_idx, total_windows, timestep,
|
||||
accum[mod_idx], counts[mod_idx], biases[mod_idx])
|
||||
|
||||
# fuse accumulated results into final conds
|
||||
try:
|
||||
# finalize conds
|
||||
if self.fuse_method.name == ContextFuseMethods.RELATIVE:
|
||||
# relative is already normalized, so return as is
|
||||
del counts_final
|
||||
return conds_final
|
||||
else:
|
||||
# normalize conds via division by context usage counts
|
||||
for i in range(len(conds_final)):
|
||||
conds_final[i] /= counts_final[i]
|
||||
del counts_final
|
||||
return conds_final
|
||||
result_out = []
|
||||
for ci in range(len(conds)):
|
||||
finalized = []
|
||||
for mod_idx in range(num_modalities):
|
||||
if self.fuse_method.name != ContextFuseMethods.RELATIVE:
|
||||
accum[mod_idx][ci] /= counts[mod_idx][ci]
|
||||
f = accum[mod_idx][ci]
|
||||
|
||||
# if guide frames were injected, append them to the end of the fused latents for the next step
|
||||
if window_state.guide_latents[mod_idx] is not None:
|
||||
f = torch.cat([f, window_state.guide_latents[mod_idx]], dim=self.dim)
|
||||
finalized.append(f)
|
||||
|
||||
# pack modalities together if needed
|
||||
if window_state.is_multimodal and len(finalized) > 1:
|
||||
packed, _ = comfy.utils.pack_latents(finalized)
|
||||
else:
|
||||
packed = finalized[0]
|
||||
|
||||
result_out.append(packed)
|
||||
return result_out
|
||||
finally:
|
||||
for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.EXECUTE_CLEANUP, self.callbacks):
|
||||
callback(self, model, x_in, conds, timestep, model_options)
|
||||
|
||||
def evaluate_context_windows(self, calc_cond_batch: Callable, model: BaseModel, x_in: torch.Tensor, conds, timestep: torch.Tensor, enumerated_context_windows: list[tuple[int, IndexListContextWindow]],
|
||||
model_options, device=None, first_device=None):
|
||||
def evaluate_context_windows(self, calc_cond_batch: Callable, model: BaseModel, x_in: torch.Tensor, conds,
|
||||
timestep: torch.Tensor, enumerated_context_windows: list[tuple[int, IndexListContextWindow]],
|
||||
model_options, window_state: WindowingState, total_windows: int = None,
|
||||
device=None, first_device=None):
|
||||
"""Evaluate context windows and return per-cond, per-modality outputs in ContextResults.sub_conds_out
|
||||
|
||||
For each window:
|
||||
1. Builds windows (for each modality if multimodal)
|
||||
2. Slices window for each modality
|
||||
3. Injects concatenated latent guide frames where present
|
||||
4. Packs together if needed and calls model
|
||||
5. Unpacks and strips any guides from outputs
|
||||
"""
|
||||
x = window_state.latents[0]
|
||||
|
||||
results: list[ContextResults] = []
|
||||
for window_idx, window in enumerated_context_windows:
|
||||
# allow processing to end between context window executions for faster Cancel
|
||||
comfy.model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
# causal_window_fix: prepend a pre-window frame that will be stripped post-forward
|
||||
# prepare the window accounting for multimodal windows
|
||||
window = window_state.prepare_window(window, model)
|
||||
|
||||
# causal_window_fix: prepend a pre-window frame that will be stripped post-forward.
|
||||
# Set anchor before slice_for_window so the latent slice and downstream cond slices both pick it up.
|
||||
anchor_applied = False
|
||||
if self.causal_window_fix:
|
||||
anchor_idx = window.index_list[0] - 1
|
||||
@@ -339,27 +665,46 @@ class IndexListContextHandler(ContextHandlerABC):
|
||||
window.causal_anchor_index = anchor_idx
|
||||
anchor_applied = True
|
||||
|
||||
# slice the window for each modality, injecting guide frames where applicable
|
||||
sliced, guide_frame_counts_per_modality = window_state.slice_for_window(window, self.latent_retain_index_list, device)
|
||||
|
||||
for callback in comfy.patcher_extension.get_all_callbacks(IndexListCallbacks.EVALUATE_CONTEXT_WINDOWS, self.callbacks):
|
||||
callback(self, model, x_in, conds, timestep, model_options, window_idx, window, model_options, device, first_device)
|
||||
|
||||
# update exposed params
|
||||
logging.info(f"Context window {window_idx + 1}/{total_windows or len(enumerated_context_windows)}: frames {window.index_list[0]}-{window.index_list[-1]} of {x.shape[self.dim]}"
|
||||
+ (f" (+{guide_frame_counts_per_modality[0]} guide frames)" if guide_frame_counts_per_modality[0] > 0 else "")
|
||||
)
|
||||
|
||||
# if multimodal, pack modalities together
|
||||
if window_state.is_multimodal and len(sliced) > 1:
|
||||
sub_x, sub_shapes = comfy.utils.pack_latents(sliced)
|
||||
else:
|
||||
sub_x, sub_shapes = sliced[0], [sliced[0].shape]
|
||||
|
||||
# get resized conds for window
|
||||
model_options["transformer_options"]["context_window"] = window
|
||||
# get subsections of x, timestep, conds
|
||||
sub_x = window.get_tensor(x_in, device)
|
||||
sub_timestep = window.get_tensor(timestep, device, dim=0)
|
||||
sub_conds = [self.get_resized_cond(cond, x_in, window, device) for cond in conds]
|
||||
sub_timestep = window.get_tensor(timestep, dim=0)
|
||||
sub_conds = [self.get_resized_cond(cond, x, window) for cond in conds]
|
||||
|
||||
# if multimodal, patch latent_shapes in conds for correct unpacking in model
|
||||
window_state.patch_latent_shapes(sub_conds, sub_shapes)
|
||||
|
||||
# call model on window
|
||||
sub_conds_out = calc_cond_batch(model, sub_conds, sub_x, sub_timestep, model_options)
|
||||
if device is not None:
|
||||
for i in range(len(sub_conds_out)):
|
||||
sub_conds_out[i] = sub_conds_out[i].to(x_in.device)
|
||||
|
||||
# strip causal_window_fix anchor if applied
|
||||
# unpack outputs
|
||||
out_per_modality = [comfy.utils.unpack_latents(sub_conds_out[i], sub_shapes) for i in range(len(sub_conds_out))]
|
||||
|
||||
# strip causal_window_fix anchor from primary modality before guide strip so window_len math stays correct
|
||||
if anchor_applied:
|
||||
for i in range(len(sub_conds_out)):
|
||||
sub_conds_out[i] = sub_conds_out[i].narrow(self.dim, 1, sub_conds_out[i].shape[self.dim] - 1)
|
||||
for ci in range(len(out_per_modality)):
|
||||
t = out_per_modality[ci][0]
|
||||
out_per_modality[ci][0] = t.narrow(self.dim, 1, t.shape[self.dim] - 1)
|
||||
|
||||
results.append(ContextResults(window_idx, sub_conds_out, sub_conds, window))
|
||||
# strip injected guide frames
|
||||
window_state.strip_guide_frames(out_per_modality, guide_frame_counts_per_modality, window)
|
||||
|
||||
results.append(ContextResults(window_idx, out_per_modality, sub_conds, window))
|
||||
return results
|
||||
|
||||
|
||||
@@ -383,7 +728,7 @@ class IndexListContextHandler(ContextHandlerABC):
|
||||
biases_final[i][idx] = bias_total + bias
|
||||
else:
|
||||
# add conds and counts based on weights of fuse method
|
||||
weights = get_context_weights(window.context_length, x_in.shape[self.dim], window.index_list, self, sigma=timestep)
|
||||
weights = get_context_weights(window.context_length, x_in.shape[self.dim], window.index_list, self, sigma=timestep, context_overlap=window.context_overlap)
|
||||
weights_tensor = match_weights_to_dim(weights, x_in, self.dim, device=x_in.device)
|
||||
for i in range(len(sub_conds_out)):
|
||||
window.add_window(conds_final[i], sub_conds_out[i] * weights_tensor)
|
||||
@@ -393,16 +738,22 @@ class IndexListContextHandler(ContextHandlerABC):
|
||||
callback(self, x_in, sub_conds_out, sub_conds, window, window_idx, total_windows, timestep, conds_final, counts_final, biases_final)
|
||||
|
||||
|
||||
def _prepare_sampling_wrapper(executor, model, noise_shape: torch.Tensor, *args, **kwargs):
|
||||
# limit noise_shape length to context_length for more accurate vram use estimation
|
||||
def _prepare_sampling_wrapper(executor, model, noise_shape: torch.Tensor, conds, *args, **kwargs):
|
||||
# Scale noise_shape to a single context window so VRAM estimation budgets per-window.
|
||||
model_options = kwargs.get("model_options", None)
|
||||
if model_options is None:
|
||||
raise Exception("model_options not found in prepare_sampling_wrapper; this should never happen, something went wrong.")
|
||||
handler: IndexListContextHandler = model_options.get("context_handler", None)
|
||||
if handler is not None:
|
||||
noise_shape = list(noise_shape)
|
||||
noise_shape[handler.dim] = min(noise_shape[handler.dim], handler.context_length)
|
||||
return executor(model, noise_shape, *args, **kwargs)
|
||||
is_packed = len(noise_shape) == 3 and noise_shape[1] == 1
|
||||
if is_packed:
|
||||
# TODO: latent_shapes cond isn't attached yet at this point, so we can't compute a
|
||||
# per-window flat latent here. Skipping the clamp over-estimates but prevents immediate OOM.
|
||||
pass
|
||||
elif handler.dim < len(noise_shape) and noise_shape[handler.dim] > handler.context_length:
|
||||
noise_shape[handler.dim] = min(noise_shape[handler.dim], handler.context_length)
|
||||
return executor(model, noise_shape, conds, *args, **kwargs)
|
||||
|
||||
|
||||
def create_prepare_sampling_wrapper(model: ModelPatcher):
|
||||
@@ -422,11 +773,12 @@ def _sampler_sample_wrapper(executor, guider, sigmas, extra_args, callback, nois
|
||||
raise Exception("context_handler not found in sampler_sample_wrapper; this should never happen, something went wrong.")
|
||||
if not handler.freenoise:
|
||||
return executor(guider, sigmas, extra_args, callback, noise, *args, **kwargs)
|
||||
noise = apply_freenoise(noise, handler.dim, handler.context_length, handler.context_overlap, extra_args["seed"])
|
||||
|
||||
conds = [guider.conds.get('positive', guider.conds.get('negative', []))]
|
||||
noise = handler._apply_freenoise(noise, conds, extra_args["seed"])
|
||||
|
||||
return executor(guider, sigmas, extra_args, callback, noise, *args, **kwargs)
|
||||
|
||||
|
||||
def create_sampler_sample_wrapper(model: ModelPatcher):
|
||||
model.add_wrapper_with_key(
|
||||
comfy.patcher_extension.WrappersMP.SAMPLER_SAMPLE,
|
||||
@@ -434,7 +786,6 @@ def create_sampler_sample_wrapper(model: ModelPatcher):
|
||||
_sampler_sample_wrapper
|
||||
)
|
||||
|
||||
|
||||
def match_weights_to_dim(weights: list[float], x_in: torch.Tensor, dim: int, device=None) -> torch.Tensor:
|
||||
total_dims = len(x_in.shape)
|
||||
weights_tensor = torch.Tensor(weights).to(device=device)
|
||||
@@ -580,8 +931,9 @@ def get_matching_context_schedule(context_schedule: str) -> ContextSchedule:
|
||||
return ContextSchedule(context_schedule, func)
|
||||
|
||||
|
||||
def get_context_weights(length: int, full_length: int, idxs: list[int], handler: IndexListContextHandler, sigma: torch.Tensor=None):
|
||||
return handler.fuse_method.func(length, sigma=sigma, handler=handler, full_length=full_length, idxs=idxs)
|
||||
def get_context_weights(length: int, full_length: int, idxs: list[int], handler: IndexListContextHandler, sigma: torch.Tensor=None, context_overlap: int=None):
|
||||
context_overlap = handler.context_overlap if context_overlap is None else context_overlap
|
||||
return handler.fuse_method.func(length, sigma=sigma, handler=handler, full_length=full_length, idxs=idxs, context_overlap=context_overlap)
|
||||
|
||||
|
||||
def create_weights_flat(length: int, **kwargs) -> list[float]:
|
||||
@@ -599,18 +951,18 @@ def create_weights_pyramid(length: int, **kwargs) -> list[float]:
|
||||
weight_sequence = list(range(1, max_weight, 1)) + [max_weight] + list(range(max_weight - 1, 0, -1))
|
||||
return weight_sequence
|
||||
|
||||
def create_weights_overlap_linear(length: int, full_length: int, idxs: list[int], handler: IndexListContextHandler, **kwargs):
|
||||
def create_weights_overlap_linear(length: int, full_length: int, idxs: list[int], context_overlap: int, **kwargs):
|
||||
# based on code in Kijai's WanVideoWrapper: https://github.com/kijai/ComfyUI-WanVideoWrapper/blob/dbb2523b37e4ccdf45127e5ae33e31362f755c8e/nodes.py#L1302
|
||||
# only expected overlap is given different weights
|
||||
weights_torch = torch.ones((length))
|
||||
# blend left-side on all except first window
|
||||
if min(idxs) > 0:
|
||||
ramp_up = torch.linspace(1e-37, 1, handler.context_overlap)
|
||||
weights_torch[:handler.context_overlap] = ramp_up
|
||||
ramp_up = torch.linspace(1e-37, 1, context_overlap)
|
||||
weights_torch[:context_overlap] = ramp_up
|
||||
# blend right-side on all except last window
|
||||
if max(idxs) < full_length-1:
|
||||
ramp_down = torch.linspace(1, 1e-37, handler.context_overlap)
|
||||
weights_torch[-handler.context_overlap:] = ramp_down
|
||||
ramp_down = torch.linspace(1, 1e-37, context_overlap)
|
||||
weights_torch[-context_overlap:] = ramp_down
|
||||
return weights_torch
|
||||
|
||||
class ContextFuseMethods:
|
||||
|
||||
@@ -515,7 +515,7 @@ class Block(nn.Module):
|
||||
h=H,
|
||||
w=W,
|
||||
)
|
||||
x_B_T_H_W_D = x_B_T_H_W_D + gate_self_attn_B_T_1_1_D.to(residual_dtype) * result_B_T_H_W_D.to(residual_dtype)
|
||||
x_B_T_H_W_D = torch.addcmul(x_B_T_H_W_D, gate_self_attn_B_T_1_1_D.to(residual_dtype), result_B_T_H_W_D.to(residual_dtype))
|
||||
|
||||
def _x_fn(
|
||||
_x_B_T_H_W_D: torch.Tensor,
|
||||
@@ -548,7 +548,7 @@ class Block(nn.Module):
|
||||
shift_cross_attn_B_T_1_1_D,
|
||||
transformer_options=transformer_options,
|
||||
)
|
||||
x_B_T_H_W_D = result_B_T_H_W_D.to(residual_dtype) * gate_cross_attn_B_T_1_1_D.to(residual_dtype) + x_B_T_H_W_D
|
||||
x_B_T_H_W_D = torch.addcmul(x_B_T_H_W_D, gate_cross_attn_B_T_1_1_D.to(residual_dtype), result_B_T_H_W_D.to(residual_dtype))
|
||||
|
||||
normalized_x_B_T_H_W_D = _fn(
|
||||
x_B_T_H_W_D,
|
||||
@@ -557,7 +557,7 @@ class Block(nn.Module):
|
||||
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))
|
||||
x_B_T_H_W_D = 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)
|
||||
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
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,290 @@
|
||||
"""Krea 2 (K2) — single-stream MMDiT.
|
||||
|
||||
Text tokens produced by a Qwen3-VL-4B 12-layer ``txtfusion`` adapter and patchified image tokens are
|
||||
concatenated into one sequence and run through ``layers`` shared transformer blocks with
|
||||
AdaLN-single modulation, GQA + per-head QK-norm + sigmoid-gated attention, SwiGLU MLP, and 3-axis RoPE.
|
||||
"""
|
||||
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.patcher_extension
|
||||
import comfy.ldm.common_dit
|
||||
from comfy.ldm.flux.layers import EmbedND, timestep_embedding
|
||||
from comfy.ldm.flux.math import apply_rope
|
||||
from comfy.ldm.modules.attention import optimized_attention_masked
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
"""RMSNorm with the reference ``(1 + scale)`` weight convention (scale stored zero-centered)."""
|
||||
|
||||
def __init__(self, features: int, eps: float = 1e-5, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.scale = nn.Parameter(torch.empty(features, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
dtype = x.dtype
|
||||
weight = comfy.model_management.cast_to(self.scale, dtype=torch.float32, device=x.device) + 1.0
|
||||
return F.rms_norm(x.float(), (x.shape[-1],), weight=weight, eps=self.eps).to(dtype)
|
||||
|
||||
|
||||
class QKNorm(nn.Module):
|
||||
def __init__(self, dim: int, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.qnorm = RMSNorm(dim, device=device, dtype=dtype, operations=operations)
|
||||
self.knorm = RMSNorm(dim, device=device, dtype=dtype, operations=operations)
|
||||
|
||||
def forward(self, q, k):
|
||||
return self.qnorm(q), self.knorm(k)
|
||||
|
||||
|
||||
class SwiGLU(nn.Module):
|
||||
def __init__(self, features: int, multiplier: int, bias: bool = False, multiple: int = 128,
|
||||
device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
mlpdim = int(2 * features / 3) * multiplier
|
||||
mlpdim = multiple * ((mlpdim + multiple - 1) // multiple)
|
||||
self.gate = operations.Linear(features, mlpdim, bias=bias, device=device, dtype=dtype)
|
||||
self.up = operations.Linear(features, mlpdim, bias=bias, device=device, dtype=dtype)
|
||||
self.down = operations.Linear(mlpdim, features, bias=bias, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x):
|
||||
return self.down(F.silu(self.gate(x)).mul_(self.up(x)))
|
||||
|
||||
|
||||
class Attention(nn.Module):
|
||||
def __init__(self, dim: int, heads: int, kvheads: Optional[int] = None, bias: bool = False,
|
||||
device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.heads = heads
|
||||
self.kvheads = kvheads if kvheads is not None else heads
|
||||
self.headdim = dim // self.heads
|
||||
self.wq = operations.Linear(dim, self.headdim * self.heads, bias=bias, device=device, dtype=dtype)
|
||||
self.wk = operations.Linear(dim, self.headdim * self.kvheads, bias=bias, device=device, dtype=dtype)
|
||||
self.wv = operations.Linear(dim, self.headdim * self.kvheads, bias=bias, device=device, dtype=dtype)
|
||||
self.gate = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype)
|
||||
self.qknorm = QKNorm(self.headdim, device=device, dtype=dtype, operations=operations)
|
||||
self.wo = operations.Linear(dim, dim, bias=bias, device=device, dtype=dtype)
|
||||
|
||||
def forward(self, x, freqs=None, mask=None, transformer_options={}):
|
||||
q, k, v, gate = self.wq(x), self.wk(x), self.wv(x), self.gate(x)
|
||||
q = rearrange(q, "B L (H D) -> B H L D", H=self.heads)
|
||||
k = rearrange(k, "B L (H D) -> B H L D", H=self.kvheads)
|
||||
v = rearrange(v, "B L (H D) -> B H L D", H=self.kvheads)
|
||||
q, k = self.qknorm(q, k)
|
||||
if freqs is not None:
|
||||
q, k = apply_rope(q, k, freqs)
|
||||
if self.kvheads != self.heads:
|
||||
rep = self.heads // self.kvheads
|
||||
k = k.repeat_interleave(rep, dim=1)
|
||||
v = v.repeat_interleave(rep, dim=1)
|
||||
out = optimized_attention_masked(q, k, v, self.heads, mask=mask, skip_reshape=True,
|
||||
transformer_options=transformer_options)
|
||||
return self.wo(out * F.sigmoid(gate))
|
||||
|
||||
|
||||
class SimpleModulation(nn.Module):
|
||||
def __init__(self, dim: int, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.lin = nn.Parameter(torch.empty(2, dim, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, vec):
|
||||
out = vec + comfy.model_management.cast_to(self.lin, dtype=vec.dtype, device=vec.device).unsqueeze(0)
|
||||
scale, shift = out.chunk(2, dim=1)
|
||||
return scale, shift
|
||||
|
||||
|
||||
class DoubleSharedModulation(nn.Module):
|
||||
def __init__(self, dim: int, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.lin = nn.Parameter(torch.empty(6 * dim, device=device, dtype=dtype))
|
||||
|
||||
def forward(self, vec):
|
||||
out = vec + comfy.model_management.cast_to(self.lin, dtype=vec.dtype, device=vec.device)
|
||||
return out.chunk(6, dim=-1)
|
||||
|
||||
|
||||
class TextFusionBlock(nn.Module):
|
||||
def __init__(self, features, heads, multiplier, bias=False, kvheads=None, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.prenorm = RMSNorm(features, device=device, dtype=dtype, operations=operations)
|
||||
self.postnorm = RMSNorm(features, device=device, dtype=dtype, operations=operations)
|
||||
self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations)
|
||||
self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations)
|
||||
|
||||
def forward(self, x, mask=None, transformer_options={}):
|
||||
x = x + self.attn(self.prenorm(x), mask=mask, transformer_options=transformer_options)
|
||||
x = x + self.mlp(self.postnorm(x))
|
||||
return x
|
||||
|
||||
|
||||
class TextFusionTransformer(nn.Module):
|
||||
def __init__(self, num_txt_layers, txt_dim, heads, multiplier, bias=False, kvheads=None, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.layerwise_blocks = nn.ModuleList([
|
||||
TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads, device=device, dtype=dtype, operations=operations)
|
||||
for _ in range(2)
|
||||
])
|
||||
self.projector = operations.Linear(num_txt_layers, 1, bias=False, device=device, dtype=dtype)
|
||||
self.refiner_blocks = nn.ModuleList([
|
||||
TextFusionBlock(txt_dim, heads, multiplier, bias, kvheads, device=device, dtype=dtype, operations=operations)
|
||||
for _ in range(2)
|
||||
])
|
||||
|
||||
def forward(self, x, mask=None, transformer_options={}):
|
||||
b, l, n, d = x.shape
|
||||
x = x.reshape(b * l, n, d)
|
||||
for block in self.layerwise_blocks:
|
||||
x = block(x.contiguous(), mask=None, transformer_options=transformer_options)
|
||||
x = rearrange(x, "(b l) n d -> b l d n", b=b, l=l)
|
||||
x = self.projector(x).squeeze(-1)
|
||||
for block in self.refiner_blocks:
|
||||
x = block(x, mask=mask, transformer_options=transformer_options)
|
||||
return x
|
||||
|
||||
|
||||
class SingleStreamBlock(nn.Module):
|
||||
def __init__(self, features, heads, multiplier, bias=False, kvheads=None, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.mod = DoubleSharedModulation(features, device=device, dtype=dtype, operations=operations)
|
||||
self.prenorm = RMSNorm(features, device=device, dtype=dtype, operations=operations)
|
||||
self.postnorm = RMSNorm(features, device=device, dtype=dtype, operations=operations)
|
||||
self.attn = Attention(features, heads, kvheads=kvheads, bias=bias, device=device, dtype=dtype, operations=operations)
|
||||
self.mlp = SwiGLU(features, multiplier, bias, device=device, dtype=dtype, operations=operations)
|
||||
|
||||
def forward(self, x, vec, freqs, mask=None, transformer_options={}):
|
||||
prescale, preshift, pregate, postscale, postshift, postgate = self.mod(vec)
|
||||
x = x + pregate * self.attn((1 + prescale) * self.prenorm(x) + preshift, freqs, mask, transformer_options=transformer_options)
|
||||
x = x + postgate * self.mlp((1 + postscale) * self.postnorm(x) + postshift)
|
||||
return x
|
||||
|
||||
|
||||
class LastLayer(nn.Module):
|
||||
def __init__(self, features, patch, channels, device=None, dtype=None, operations=None):
|
||||
super().__init__()
|
||||
self.norm = RMSNorm(features, device=device, dtype=dtype, operations=operations)
|
||||
self.linear = operations.Linear(features, patch * patch * channels, bias=True, device=device, dtype=dtype)
|
||||
self.modulation = SimpleModulation(features, device=device, dtype=dtype, operations=operations)
|
||||
|
||||
def forward(self, x, tvec):
|
||||
scale, shift = self.modulation(tvec)
|
||||
x = (1 + scale) * self.norm(x) + shift
|
||||
return self.linear(x)
|
||||
|
||||
|
||||
class SingleStreamDiT(nn.Module):
|
||||
def __init__(self, features=6144, tdim=256, txtdim=2560, heads=48, kvheads=12, multiplier=4,
|
||||
layers=28, patch=2, channels=16, bias=False, theta=1e3, txtlayers=12,
|
||||
txtheads=20, txtkvheads=20, image_model=None,
|
||||
device=None, dtype=None, operations=None, **kwargs):
|
||||
super().__init__()
|
||||
self.dtype = dtype
|
||||
self.patch = patch
|
||||
self.channels = channels
|
||||
self.tdim = tdim
|
||||
self.heads = heads
|
||||
self.txtdim = txtdim
|
||||
self.txtlayers = txtlayers
|
||||
|
||||
headdim = features // heads
|
||||
axes = [headdim - 12 * (headdim // 16), 6 * (headdim // 16), 6 * (headdim // 16)]
|
||||
assert sum(axes) == headdim, f"axes {axes} sum != headdim {headdim}"
|
||||
self.pe_embedder = EmbedND(dim=headdim, theta=int(theta), axes_dim=axes)
|
||||
|
||||
self.first = operations.Linear(channels * patch ** 2, features, bias=True, device=device, dtype=dtype)
|
||||
self.blocks = nn.ModuleList([
|
||||
SingleStreamBlock(features, heads, multiplier, bias, kvheads, device=device, dtype=dtype, operations=operations)
|
||||
for _ in range(layers)
|
||||
])
|
||||
self.tmlp = nn.Sequential(
|
||||
operations.Linear(tdim, features, device=device, dtype=dtype),
|
||||
nn.GELU(approximate="tanh"),
|
||||
operations.Linear(features, features, device=device, dtype=dtype),
|
||||
)
|
||||
self.txtfusion = TextFusionTransformer(txtlayers, txtdim, txtheads, multiplier, bias, txtkvheads,
|
||||
device=device, dtype=dtype, operations=operations)
|
||||
self.txtmlp = nn.Sequential(
|
||||
RMSNorm(txtdim, device=device, dtype=dtype, operations=operations),
|
||||
operations.Linear(txtdim, features, device=device, dtype=dtype),
|
||||
nn.GELU(approximate="tanh"),
|
||||
operations.Linear(features, features, device=device, dtype=dtype),
|
||||
)
|
||||
self.last = LastLayer(features, patch, channels, device=device, dtype=dtype, operations=operations)
|
||||
self.tproj = nn.Sequential(
|
||||
nn.GELU(approximate="tanh"),
|
||||
operations.Linear(features, features * 6, device=device, dtype=dtype),
|
||||
)
|
||||
|
||||
def forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
|
||||
return comfy.patcher_extension.WrapperExecutor.new_class_executor(
|
||||
self._forward,
|
||||
self,
|
||||
comfy.patcher_extension.get_all_wrappers(comfy.patcher_extension.WrappersMP.DIFFUSION_MODEL, transformer_options),
|
||||
).execute(x, timesteps, context, attention_mask, transformer_options, **kwargs)
|
||||
|
||||
def _forward(self, x, timesteps, context, attention_mask=None, transformer_options={}, **kwargs):
|
||||
temporal = x.ndim == 5
|
||||
if temporal:
|
||||
b5, c5, t5, h5, w5 = x.shape
|
||||
x = x.reshape(b5 * t5, c5, h5, w5)
|
||||
bs, c, H_orig, W_orig = x.shape
|
||||
patch = self.patch
|
||||
# Pad the latent up to a multiple of patch (as Flux/Lumina/QwenImage do); crop back at the end.
|
||||
x = comfy.ldm.common_dit.pad_to_patch_size(x, (patch, patch))
|
||||
H, W = x.shape[-2], x.shape[-1]
|
||||
h_, w_ = H // patch, W // patch
|
||||
|
||||
# context arrives as (B, seq, txtlayers*txtdim); reshape to (B, txtlayers, seq, txtdim).
|
||||
context = self._unpack_context(context)
|
||||
|
||||
img = rearrange(x, "b c (h ph) (w pw) -> b (h w) (c ph pw)", ph=patch, pw=patch)
|
||||
img = self.first(img)
|
||||
|
||||
t = self.tmlp(timestep_embedding(timesteps, self.tdim).unsqueeze(1).to(img.dtype))
|
||||
tvec = self.tproj(t)
|
||||
|
||||
context = self.txtfusion(context, mask=None, transformer_options=transformer_options)
|
||||
context = self.txtmlp(context)
|
||||
|
||||
txtlen, imglen = context.shape[1], img.shape[1]
|
||||
combined = torch.cat((context, img), dim=1)
|
||||
|
||||
# Position ids: text at 0, image at (0, h_idx, w_idx).
|
||||
device = combined.device
|
||||
txtpos = torch.zeros(bs, txtlen, 3, device=device, dtype=torch.float32)
|
||||
imgids = torch.zeros(h_, w_, 3, device=device, dtype=torch.float32)
|
||||
imgids[..., 1] = torch.arange(h_, device=device, dtype=torch.float32)[:, None]
|
||||
imgids[..., 2] = torch.arange(w_, device=device, dtype=torch.float32)[None, :]
|
||||
imgpos = imgids.reshape(1, h_ * w_, 3).repeat(bs, 1, 1)
|
||||
pos = torch.cat((txtpos, imgpos), dim=1)
|
||||
|
||||
freqs = self.pe_embedder(pos)
|
||||
|
||||
for block in self.blocks:
|
||||
combined = block(combined, tvec, freqs, None, transformer_options=transformer_options)
|
||||
|
||||
final = self.last(combined, t)
|
||||
out = final[:, txtlen:txtlen + imglen, :]
|
||||
out = rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)",
|
||||
h=h_, w=w_, ph=patch, pw=patch, c=self.channels)
|
||||
out = out[:, :, :H_orig, :W_orig] # crop padding back off
|
||||
if temporal:
|
||||
out = out.reshape(b5, t5, self.channels, H_orig, W_orig).movedim(1, 2)
|
||||
return out
|
||||
|
||||
def _unpack_context(self, context):
|
||||
# context: (B, seq, txtlayers*txtdim) -> (B, seq, txtlayers, txtdim).
|
||||
b, seq, fused = context.shape
|
||||
if fused != self.txtlayers * self.txtdim:
|
||||
raise ValueError(
|
||||
f"Krea2 expects conditioning with {self.txtlayers}x{self.txtdim}={self.txtlayers * self.txtdim} "
|
||||
f"features (a {self.txtlayers}-layer Qwen3-VL stack) but got {fused}. "
|
||||
f"Load the text encoder with CLIPLoader type 'krea2'."
|
||||
)
|
||||
return context.reshape(b, seq, self.txtlayers, self.txtdim)
|
||||
@@ -1085,7 +1085,7 @@ class LTXVModel(LTXBaseModel):
|
||||
)
|
||||
|
||||
grid_mask = None
|
||||
if keyframe_idxs is not None:
|
||||
if keyframe_idxs is not None and keyframe_idxs.shape[2] > 0:
|
||||
additional_args.update({ "orig_patchified_shape": list(x.shape)})
|
||||
denoise_mask = self.patchifier.patchify(denoise_mask)[0]
|
||||
grid_mask = ~torch.any(denoise_mask < 0, dim=-1)[0]
|
||||
@@ -1330,7 +1330,7 @@ class LTXVModel(LTXBaseModel):
|
||||
x = x * (1 + scale) + shift
|
||||
x = self.proj_out(x)
|
||||
|
||||
if keyframe_idxs is not None:
|
||||
if keyframe_idxs is not None and keyframe_idxs.shape[2] > 0:
|
||||
grid_mask = kwargs["grid_mask"]
|
||||
orig_patchified_shape = kwargs["orig_patchified_shape"]
|
||||
full_x = torch.zeros(orig_patchified_shape, dtype=x.dtype, device=x.device)
|
||||
|
||||
@@ -326,6 +326,17 @@ def model_lora_keys_unet(model, key_map={}):
|
||||
key_map["transformer.{}".format(key_lora)] = k
|
||||
key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = k #SimpleTuner lycoris format
|
||||
|
||||
if isinstance(model, comfy.model_base.Krea2):
|
||||
diffusers_keys = comfy.utils.krea2_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")
|
||||
for k in diffusers_keys:
|
||||
if k.endswith(".weight"):
|
||||
to = diffusers_keys[k]
|
||||
key_lora = k[:-len(".weight")]
|
||||
key_map["diffusion_model.{}".format(key_lora)] = to
|
||||
key_map["transformer.{}".format(key_lora)] = to
|
||||
key_map["lycoris_{}".format(key_lora.replace(".", "_"))] = to
|
||||
key_map[key_lora] = to
|
||||
|
||||
if isinstance(model, comfy.model_base.Lumina2):
|
||||
diffusers_keys = comfy.utils.z_image_to_diffusers(model.model_config.unet_config, output_prefix="diffusion_model.")
|
||||
for k in diffusers_keys:
|
||||
|
||||
@@ -21,6 +21,7 @@ import comfy.ldm.hunyuan3dv2_1.hunyuandit
|
||||
import torch
|
||||
import logging
|
||||
import comfy.ldm.lightricks.av_model
|
||||
import comfy.ldm.lightricks.symmetric_patchifier
|
||||
import comfy.context_windows
|
||||
from comfy.ldm.modules.diffusionmodules.openaimodel import UNetModel, Timestep
|
||||
from comfy.ldm.cascade.stage_c import StageC
|
||||
@@ -57,6 +58,7 @@ import comfy.ldm.omnigen.omnigen2
|
||||
import comfy.ldm.boogu.model
|
||||
import comfy.ldm.qwen_image.model
|
||||
import comfy.ldm.ideogram4.model
|
||||
import comfy.ldm.krea2.model
|
||||
import comfy.ldm.kandinsky5.model
|
||||
import comfy.ldm.anima.model
|
||||
import comfy.ldm.ace.ace_step15
|
||||
@@ -1204,6 +1206,127 @@ class LTXAV(BaseModel):
|
||||
def scale_latent_inpaint(self, sigma, noise, latent_image, **kwargs):
|
||||
return latent_image
|
||||
|
||||
def map_context_window_to_modalities(self, primary_indices, latent_shapes, dim):
|
||||
result = [primary_indices]
|
||||
if len(latent_shapes) < 2:
|
||||
return result
|
||||
|
||||
video_total = latent_shapes[0][dim]
|
||||
|
||||
for i in range(1, len(latent_shapes)):
|
||||
mod_total = latent_shapes[i][dim]
|
||||
# Map each primary index to its proportional range of modality indices and
|
||||
# concatenate in order. Preserves wrapped/strided geometry so the modality
|
||||
# attends to the same temporal regions as the primary window.
|
||||
mod_indices = []
|
||||
seen = set()
|
||||
for v_idx in primary_indices:
|
||||
a_start = min(int(round(v_idx * mod_total / video_total)), mod_total - 1)
|
||||
a_end = min(int(round((v_idx + 1) * mod_total / video_total)), mod_total)
|
||||
if a_end <= a_start:
|
||||
a_end = a_start + 1
|
||||
for a in range(a_start, a_end):
|
||||
if a not in seen:
|
||||
seen.add(a)
|
||||
mod_indices.append(a)
|
||||
result.append(mod_indices)
|
||||
|
||||
return result
|
||||
|
||||
@staticmethod
|
||||
def _get_guide_entries(conds):
|
||||
for cond_list in conds:
|
||||
if cond_list is None:
|
||||
continue
|
||||
for cond_dict in cond_list:
|
||||
model_conds = cond_dict.get('model_conds', {})
|
||||
entries = model_conds.get('guide_attention_entries')
|
||||
if entries is not None and hasattr(entries, 'cond') and entries.cond:
|
||||
return entries.cond
|
||||
return None
|
||||
|
||||
def resize_cond_for_context_window(self, cond_key, cond_value, window, x_in, device, retain_index_list=[]):
|
||||
# Audio denoise mask — slice using audio modality window
|
||||
if cond_key == "audio_denoise_mask" and hasattr(window, 'modality_windows') and window.modality_windows:
|
||||
audio_window = window.modality_windows.get(1)
|
||||
if audio_window is not None and hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor):
|
||||
sliced = audio_window.get_tensor(cond_value.cond, device, dim=2)
|
||||
return cond_value._copy_with(sliced)
|
||||
|
||||
# Video denoise mask — split into video + guide portions, slice each
|
||||
if cond_key == "denoise_mask" and hasattr(cond_value, "cond") and isinstance(cond_value.cond, torch.Tensor):
|
||||
cond_tensor = cond_value.cond
|
||||
guide_count = cond_tensor.size(window.dim) - x_in.size(window.dim)
|
||||
if guide_count > 0:
|
||||
T_video = x_in.size(window.dim)
|
||||
video_mask = cond_tensor.narrow(window.dim, 0, T_video)
|
||||
guide_mask = cond_tensor.narrow(window.dim, T_video, guide_count)
|
||||
sliced_video = window.get_tensor(video_mask, device, retain_index_list=retain_index_list)
|
||||
suffix_indices = window.guide_frames_indices
|
||||
if suffix_indices:
|
||||
idx = tuple([slice(None)] * window.dim + [suffix_indices])
|
||||
sliced_guide = guide_mask[idx].to(device)
|
||||
return cond_value._copy_with(torch.cat([sliced_video, sliced_guide], dim=window.dim))
|
||||
else:
|
||||
return cond_value._copy_with(sliced_video)
|
||||
|
||||
# Keyframe indices — regenerate pixel coords for window, select guide positions
|
||||
if cond_key == "keyframe_idxs":
|
||||
kf_local_pos = window.guide_kf_local_positions
|
||||
if not kf_local_pos:
|
||||
return cond_value._copy_with(cond_value.cond[:, :, :0, :]) # empty
|
||||
H, W = x_in.shape[3], x_in.shape[4]
|
||||
window_len = len(window.index_list)
|
||||
# account for causal_window_fix anchor in coord space size
|
||||
anchor_idx = getattr(window, 'causal_anchor_index', None)
|
||||
if anchor_idx is not None and anchor_idx >= 0:
|
||||
window_len += 1
|
||||
patchifier = self.diffusion_model.patchifier
|
||||
latent_coords = patchifier.get_latent_coords(window_len, H, W, 1, cond_value.cond.device)
|
||||
scale_factors = self.diffusion_model.vae_scale_factors
|
||||
pixel_coords = comfy.ldm.lightricks.symmetric_patchifier.latent_to_pixel_coords(
|
||||
latent_coords,
|
||||
scale_factors,
|
||||
causal_fix=self.diffusion_model.causal_temporal_positioning)
|
||||
tokens = []
|
||||
for pos in kf_local_pos:
|
||||
tokens.extend(range(pos * H * W, (pos + 1) * H * W))
|
||||
pixel_coords = pixel_coords[:, :, tokens, :]
|
||||
|
||||
# Adjust spatial end positions for dilated (downscaled) guides.
|
||||
# Each guide entry may have a different downscale factor; expand the
|
||||
# per-entry factor to cover all tokens belonging to that entry.
|
||||
downscale_factors = window.guide_downscale_factors
|
||||
overlap_info = window.guide_overlap_info
|
||||
if downscale_factors:
|
||||
per_token_factor = []
|
||||
for (entry_idx, overlap_count), dsf in zip(overlap_info, downscale_factors):
|
||||
per_token_factor.extend([dsf] * (overlap_count * H * W))
|
||||
factor_tensor = torch.tensor(per_token_factor, device=pixel_coords.device, dtype=pixel_coords.dtype)
|
||||
spatial_end_offset = (factor_tensor.unsqueeze(0).unsqueeze(0).unsqueeze(-1) - 1) * torch.tensor(
|
||||
scale_factors[1:], device=pixel_coords.device, dtype=pixel_coords.dtype,
|
||||
).view(1, -1, 1, 1)
|
||||
pixel_coords[:, 1:, :, 1:] += spatial_end_offset
|
||||
|
||||
B = cond_value.cond.shape[0]
|
||||
if B > 1:
|
||||
pixel_coords = pixel_coords.expand(B, -1, -1, -1)
|
||||
return cond_value._copy_with(pixel_coords)
|
||||
|
||||
# Guide attention entries — adjust per-guide counts based on window overlap
|
||||
if cond_key == "guide_attention_entries":
|
||||
overlap_info = window.guide_overlap_info
|
||||
H, W = x_in.shape[3], x_in.shape[4]
|
||||
new_entries = []
|
||||
for entry_idx, overlap_count in overlap_info:
|
||||
e = cond_value.cond[entry_idx]
|
||||
new_entries.append({**e,
|
||||
"pre_filter_count": overlap_count * H * W,
|
||||
"latent_shape": [overlap_count, H, W]})
|
||||
return cond_value._copy_with(new_entries)
|
||||
|
||||
return None
|
||||
|
||||
class HunyuanVideo(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan_video.model.HunyuanVideo)
|
||||
@@ -2156,6 +2279,17 @@ class Ideogram4(BaseModel):
|
||||
out['c_crossattn'] = comfy.conds.CONDRegular(cross_attn)
|
||||
return out
|
||||
|
||||
class Krea2(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLUX, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.krea2.model.SingleStreamDiT)
|
||||
|
||||
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)
|
||||
return out
|
||||
|
||||
class HunyuanImage21(BaseModel):
|
||||
def __init__(self, model_config, model_type=ModelType.FLOW, device=None):
|
||||
super().__init__(model_config, model_type, device=device, unet_model=comfy.ldm.hunyuan_video.model.HunyuanVideo)
|
||||
|
||||
@@ -834,6 +834,21 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
|
||||
dit_config["num_layers"] = count_blocks(state_dict_keys, '{}layers.'.format(key_prefix) + '{}.')
|
||||
return dit_config
|
||||
|
||||
if '{}txtfusion.projector.weight'.format(key_prefix) in state_dict_keys: # Krea 2 (K2)
|
||||
dit_config = {}
|
||||
dit_config["image_model"] = "krea2"
|
||||
head_dim = 128
|
||||
first_w = state_dict['{}first.weight'.format(key_prefix)] # (features, channels*patch^2)
|
||||
dit_config["features"] = first_w.shape[0]
|
||||
dit_config["channels"] = first_w.shape[1] // (2 * 2) # patch=2
|
||||
dit_config["patch"] = 2
|
||||
dit_config["layers"] = count_blocks(state_dict_keys, '{}blocks.'.format(key_prefix) + '{}.')
|
||||
dit_config["heads"] = state_dict['{}blocks.0.attn.wq.weight'.format(key_prefix)].shape[0] // head_dim
|
||||
dit_config["kvheads"] = state_dict['{}blocks.0.attn.wk.weight'.format(key_prefix)].shape[0] // head_dim
|
||||
dit_config["txtlayers"] = state_dict['{}txtfusion.projector.weight'.format(key_prefix)].shape[1]
|
||||
dit_config["txtdim"] = state_dict['{}txtfusion.layerwise_blocks.0.prenorm.scale'.format(key_prefix)].shape[0]
|
||||
return dit_config
|
||||
|
||||
if '{}visual_transformer_blocks.0.cross_attention.key_norm.weight'.format(key_prefix) in state_dict_keys: # Kandinsky 5
|
||||
dit_config = {}
|
||||
model_dim = state_dict['{}visual_embeddings.in_layer.bias'.format(key_prefix)].shape[0]
|
||||
|
||||
+56
-8
@@ -256,7 +256,7 @@ def resolve_cast_module_with_vbar(s, dtype, device, bias_dtype, compute_dtype, w
|
||||
if (want_requant and len(fns) == 0 or update_weight):
|
||||
seed = comfy.utils.string_to_seed(s.seed_key)
|
||||
if isinstance(orig, QuantizedTensor):
|
||||
y = QuantizedTensor.from_float(x, s.layout_type, scale="recalculate", stochastic_rounding=seed)
|
||||
y = orig.requantize_from_float(x, scale="recalculate", stochastic_rounding=seed)
|
||||
else:
|
||||
y = comfy.float.stochastic_rounding(x, orig.dtype, seed=seed)
|
||||
if want_requant and len(fns) == 0:
|
||||
@@ -1089,6 +1089,19 @@ def _load_quantized_module(module, super_load, state_dict, prefix, local_metadat
|
||||
if ts is None or bs is None:
|
||||
raise ValueError(f"Missing NVFP4 scales for layer {layer_name}")
|
||||
scales = {"scale": ts, "block_scale": bs}
|
||||
elif module.quant_format == "int8_tensorwise":
|
||||
scale = pop_scale("weight_scale")
|
||||
if scale is None:
|
||||
raise ValueError(f"Missing INT8 weight scale for layer {layer_name}")
|
||||
scales = {"scale": scale}
|
||||
params_conf = layer_conf.get("params", {})
|
||||
if not isinstance(params_conf, dict):
|
||||
params_conf = {}
|
||||
if layer_conf.get("convrot", params_conf.get("convrot", False)):
|
||||
scales["convrot"] = True
|
||||
scales["convrot_groupsize"] = int(
|
||||
layer_conf.get("convrot_groupsize", params_conf.get("convrot_groupsize", 256))
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported quantization format: {module.quant_format}")
|
||||
|
||||
@@ -1131,6 +1144,10 @@ def _quantized_weight_state_dict(module, sd, prefix, extra_quant_conf=None, extr
|
||||
quant_conf = {"format": module.quant_format}
|
||||
if getattr(module, '_full_precision_mm_config', False):
|
||||
quant_conf["full_precision_matrix_mult"] = True
|
||||
params = getattr(module.weight, "_params", None)
|
||||
if module.quant_format == "int8_tensorwise" and getattr(params, "convrot", False):
|
||||
quant_conf["convrot"] = True
|
||||
quant_conf["convrot_groupsize"] = getattr(params, "convrot_groupsize", 256)
|
||||
if extra_quant_conf:
|
||||
quant_conf.update(extra_quant_conf)
|
||||
sd[f"{prefix}comfy_quant"] = torch.tensor(list(json.dumps(quant_conf).encode("utf-8")), dtype=torch.uint8)
|
||||
@@ -1183,8 +1200,33 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
def _forward(self, input, weight, bias):
|
||||
return torch.nn.functional.linear(input, weight, bias)
|
||||
|
||||
def forward_comfy_cast_weights(self, input, compute_dtype=None, want_requant=False):
|
||||
weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True, compute_dtype=compute_dtype, want_requant=want_requant)
|
||||
def forward_comfy_cast_weights(
|
||||
self,
|
||||
input,
|
||||
compute_dtype=None,
|
||||
want_requant=False,
|
||||
weight_only_quant=False,
|
||||
):
|
||||
if weight_only_quant:
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
self,
|
||||
input=None,
|
||||
dtype=self.weight.dtype,
|
||||
device=input.device,
|
||||
bias_dtype=input.dtype,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=True,
|
||||
)
|
||||
weight = weight.to(dtype=input.dtype)
|
||||
else:
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
self,
|
||||
input,
|
||||
offloadable=True,
|
||||
compute_dtype=compute_dtype,
|
||||
want_requant=want_requant,
|
||||
)
|
||||
x = self._forward(input, weight, bias)
|
||||
uncast_bias_weight(self, weight, bias, offload_stream)
|
||||
return x
|
||||
@@ -1203,9 +1245,10 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
not getattr(self, 'comfy_force_cast_weights', False) and
|
||||
len(self.weight_function) == 0 and len(self.bias_function) == 0
|
||||
)
|
||||
quantize_input = QUANT_ALGOS.get(getattr(self, 'quant_format', None), {}).get("quantize_input", True)
|
||||
|
||||
# Training path: quantized forward with compute_dtype backward via autograd function
|
||||
if (input.requires_grad and _use_quantized):
|
||||
if (input.requires_grad and _use_quantized and quantize_input):
|
||||
|
||||
weight, bias, offload_stream = cast_bias_weight(
|
||||
self,
|
||||
@@ -1227,7 +1270,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
return output
|
||||
|
||||
# Inference path (unchanged)
|
||||
if _use_quantized:
|
||||
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
|
||||
@@ -1241,7 +1284,13 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
scale = comfy.model_management.cast_to_device(scale, input.device, None)
|
||||
input = QuantizedTensor.from_float(input_reshaped, self.layout_type, scale=scale)
|
||||
|
||||
output = self.forward_comfy_cast_weights(input, compute_dtype, want_requant=isinstance(input, QuantizedTensor))
|
||||
weight_only_quant = _use_quantized and not quantize_input and isinstance(self.weight, QuantizedTensor)
|
||||
output = self.forward_comfy_cast_weights(
|
||||
input,
|
||||
compute_dtype,
|
||||
want_requant=isinstance(input, QuantizedTensor),
|
||||
weight_only_quant=weight_only_quant,
|
||||
)
|
||||
|
||||
# Reshape output back to 3D if input was 3D
|
||||
if reshaped_3d:
|
||||
@@ -1257,8 +1306,7 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
|
||||
|
||||
def set_weight(self, weight, inplace_update=False, seed=None, return_weight=False, **kwargs):
|
||||
if getattr(self, 'layout_type', None) is not None:
|
||||
# dtype is now implicit in the layout class
|
||||
weight = QuantizedTensor.from_float(weight, self.layout_type, scale="recalculate", stochastic_rounding=seed, inplace_ops=True).to(self.weight.dtype)
|
||||
weight = self.weight.requantize_from_float(weight, scale="recalculate", stochastic_rounding=seed, inplace_ops=True).to(self.weight.dtype)
|
||||
else:
|
||||
weight = weight.to(self.weight.dtype)
|
||||
if return_weight:
|
||||
|
||||
@@ -10,6 +10,7 @@ try:
|
||||
QuantizedLayout,
|
||||
TensorCoreFP8Layout as _CKFp8Layout,
|
||||
TensorCoreNVFP4Layout as _CKNvfp4Layout,
|
||||
TensorWiseINT8Layout as _CKTensorWiseINT8Layout,
|
||||
register_layout_op,
|
||||
register_layout_class,
|
||||
get_layout_class,
|
||||
@@ -47,6 +48,9 @@ except ImportError as e:
|
||||
class _CKNvfp4Layout:
|
||||
pass
|
||||
|
||||
class _CKTensorWiseINT8Layout:
|
||||
pass
|
||||
|
||||
def register_layout_class(name, cls):
|
||||
pass
|
||||
|
||||
@@ -174,6 +178,7 @@ class TensorCoreFP8E5M2Layout(_TensorCoreFP8LayoutBase):
|
||||
|
||||
# Backward compatibility alias - default to E4M3
|
||||
TensorCoreFP8Layout = TensorCoreFP8E4M3Layout
|
||||
TensorWiseINT8Layout = _CKTensorWiseINT8Layout
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
@@ -184,6 +189,7 @@ register_layout_class("TensorCoreFP8Layout", TensorCoreFP8Layout)
|
||||
register_layout_class("TensorCoreFP8E4M3Layout", TensorCoreFP8E4M3Layout)
|
||||
register_layout_class("TensorCoreFP8E5M2Layout", TensorCoreFP8E5M2Layout)
|
||||
register_layout_class("TensorCoreNVFP4Layout", TensorCoreNVFP4Layout)
|
||||
register_layout_class("TensorWiseINT8Layout", _CKTensorWiseINT8Layout)
|
||||
if _CK_MXFP8_AVAILABLE:
|
||||
register_layout_class("TensorCoreMXFP8Layout", TensorCoreMXFP8Layout)
|
||||
|
||||
@@ -214,6 +220,13 @@ if _CK_MXFP8_AVAILABLE:
|
||||
"group_size": 32,
|
||||
}
|
||||
|
||||
QUANT_ALGOS["int8_tensorwise"] = {
|
||||
"storage_t": torch.int8,
|
||||
"parameters": {"weight_scale"},
|
||||
"comfy_tensor_layout": "TensorWiseINT8Layout",
|
||||
"quantize_input": False,
|
||||
}
|
||||
|
||||
|
||||
# ==============================================================================
|
||||
# Re-exports for backward compatibility
|
||||
@@ -226,6 +239,7 @@ __all__ = [
|
||||
"TensorCoreFP8E4M3Layout",
|
||||
"TensorCoreFP8E5M2Layout",
|
||||
"TensorCoreNVFP4Layout",
|
||||
"TensorWiseINT8Layout",
|
||||
"QUANT_ALGOS",
|
||||
"register_layout_op",
|
||||
]
|
||||
|
||||
@@ -58,6 +58,7 @@ import comfy.text_encoders.omnigen2
|
||||
import comfy.text_encoders.qwen_image
|
||||
import comfy.text_encoders.hunyuan_image
|
||||
import comfy.text_encoders.z_image
|
||||
import comfy.text_encoders.krea2
|
||||
import comfy.text_encoders.ideogram4
|
||||
import comfy.text_encoders.ovis
|
||||
import comfy.text_encoders.kandinsky5
|
||||
@@ -1303,6 +1304,7 @@ class CLIPType(Enum):
|
||||
PIXELDIT = 29
|
||||
IDEOGRAM4 = 30
|
||||
BOOGU = 31
|
||||
KREA2 = 32
|
||||
|
||||
|
||||
|
||||
@@ -1628,6 +1630,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.boogu.te(**llama_detect(clip_data))
|
||||
clip_target.tokenizer = comfy.text_encoders.boogu.BooguTokenizer
|
||||
elif clip_type == CLIPType.KREA2 and te_model == TEModel.QWEN3VL_4B: # Krea2: full Qwen3-VL-4B (12-layer tap for conditioning + multimodal generate).
|
||||
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 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)
|
||||
|
||||
+11
-5
@@ -543,18 +543,24 @@ class SDTokenizer:
|
||||
def _try_get_embedding(self, embedding_name:str):
|
||||
'''
|
||||
Takes a potential embedding name and tries to retrieve it.
|
||||
Returns a Tuple consisting of the embedding and any leftover string, embedding can be None.
|
||||
Returns a Tuple consisting of the embedding, the cleaned embedding name, and any leftover string, embedding can be None.
|
||||
'''
|
||||
split_embed = embedding_name.split()
|
||||
embedding_name = split_embed[0]
|
||||
leftover = ' '.join(split_embed[1:])
|
||||
|
||||
match = re.search(r'[<\[]', embedding_name)
|
||||
if match is not None:
|
||||
leftover = embedding_name[match.start():] + (" " + leftover if leftover else "")
|
||||
embedding_name = embedding_name[:match.start()]
|
||||
|
||||
embed = load_embed(embedding_name, self.embedding_directory, self.embedding_size, self.embedding_key)
|
||||
if embed is None:
|
||||
stripped = embedding_name.strip(',')
|
||||
if len(stripped) < len(embedding_name):
|
||||
embed = load_embed(stripped, self.embedding_directory, self.embedding_size, self.embedding_key)
|
||||
return (embed, "{} {}".format(embedding_name[len(stripped):], leftover))
|
||||
return (embed, leftover)
|
||||
return (embed, embedding_name, "{} {}".format(embedding_name[len(stripped):], leftover))
|
||||
return (embed, embedding_name, leftover)
|
||||
|
||||
def pad_tokens(self, tokens, amount):
|
||||
if self.pad_left:
|
||||
@@ -585,7 +591,7 @@ class SDTokenizer:
|
||||
tokens = []
|
||||
for weighted_segment, weight in parsed_weights:
|
||||
to_tokenize = unescape_important(weighted_segment)
|
||||
split = re.split(' {0}|\n{0}'.format(self.embedding_identifier), to_tokenize)
|
||||
split = re.split(r'(?<=\s){}'.format(re.escape(self.embedding_identifier)), to_tokenize)
|
||||
to_tokenize = [split[0]]
|
||||
for i in range(1, len(split)):
|
||||
to_tokenize.append("{}{}".format(self.embedding_identifier, split[i]))
|
||||
@@ -595,7 +601,7 @@ class SDTokenizer:
|
||||
# if we find an embedding, deal with the embedding
|
||||
if word.startswith(self.embedding_identifier) and self.embedding_directory is not None:
|
||||
embedding_name = word[len(self.embedding_identifier):].strip('\n')
|
||||
embed, leftover = self._try_get_embedding(embedding_name)
|
||||
embed, embedding_name, leftover = self._try_get_embedding(embedding_name)
|
||||
if embed is None:
|
||||
logging.warning(f"warning, embedding:{embedding_name} does not exist, ignoring")
|
||||
else:
|
||||
|
||||
@@ -26,6 +26,7 @@ import comfy.text_encoders.kandinsky5
|
||||
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.anima
|
||||
import comfy.text_encoders.ace15
|
||||
import comfy.text_encoders.longcat_image
|
||||
@@ -1818,6 +1819,35 @@ class Ideogram4(supported_models_base.BASE):
|
||||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_8b.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.ideogram4.Ideogram4Tokenizer, comfy.text_encoders.ideogram4.te(**hunyuan_detect))
|
||||
|
||||
|
||||
class Krea2(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "krea2",
|
||||
}
|
||||
|
||||
sampling_settings = {
|
||||
"multiplier": 1.0,
|
||||
"shift": 1.15,
|
||||
}
|
||||
|
||||
memory_usage_factor = 2.2
|
||||
|
||||
latent_format = latent_formats.Wan21
|
||||
|
||||
supported_inference_dtypes = [torch.bfloat16, torch.float16, torch.float32]
|
||||
|
||||
vae_key_prefix = ["vae."]
|
||||
text_encoder_key_prefix = ["text_encoders."]
|
||||
|
||||
def get_model(self, state_dict, prefix="", device=None):
|
||||
out = model_base.Krea2(self, device=device)
|
||||
return out
|
||||
|
||||
def clip_target(self, state_dict={}):
|
||||
pref = self.text_encoder_key_prefix[0]
|
||||
hunyuan_detect = comfy.text_encoders.hunyuan_video.llama_detect(state_dict, "{}qwen3vl_4b.transformer.".format(pref))
|
||||
return supported_models_base.ClipTarget(comfy.text_encoders.krea2.Krea2Tokenizer, comfy.text_encoders.krea2.te(**hunyuan_detect))
|
||||
|
||||
class QwenImage(supported_models_base.BASE):
|
||||
unet_config = {
|
||||
"image_model": "qwen_image",
|
||||
@@ -2325,6 +2355,7 @@ models = [
|
||||
Boogu,
|
||||
QwenImage,
|
||||
Ideogram4,
|
||||
Krea2,
|
||||
Flux2,
|
||||
Lens,
|
||||
Kandinsky5Image,
|
||||
|
||||
@@ -0,0 +1,84 @@
|
||||
"""Krea 2 (K2) text encoder: Qwen3-VL-4B, 12-layer tap.
|
||||
|
||||
K2 conditions on a stack of hidden states from 12 layers of Qwen3-VL-4B
|
||||
(reference taps ``hidden_states[2,5,8,...,35]``), kept as a ``(B, 12, seq, 2560)`` tensor and
|
||||
consumed by the DiT's internal ``txtfusion`` adapter. Comfy carries conditioning as a 3D tensor,
|
||||
so the 12-layer stack is flattened to ``(B, seq, 12*2560)`` here and unpacked inside the model.
|
||||
"""
|
||||
|
||||
import numbers
|
||||
|
||||
import torch
|
||||
|
||||
import comfy.text_encoders.qwen3vl
|
||||
from comfy import sd1_clip
|
||||
|
||||
# tap k == hidden_states[k] (no offset).
|
||||
KREA2_TAP_LAYERS = [2, 5, 8, 11, 14, 17, 20, 23, 26, 29, 32, 35]
|
||||
|
||||
# Identical system template to Qwen-Image; Krea2 strips the system+user-opening prefix.
|
||||
KREA2_TEMPLATE = "<|im_start|>system\nDescribe the image by detailing the color, shape, size, texture, quantity, text, spatial relationships of the objects and background:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n"
|
||||
|
||||
|
||||
class Krea2Tokenizer(comfy.text_encoders.qwen3vl.Qwen3VLTokenizer):
|
||||
def __init__(self, embedding_directory=None, tokenizer_data={}):
|
||||
super().__init__(embedding_directory=embedding_directory, tokenizer_data=tokenizer_data, model_type="qwen3vl_4b")
|
||||
self.llama_template = KREA2_TEMPLATE # conditioning template; image text-gen uses qwen3vl's default image template.
|
||||
|
||||
def tokenize_with_weights(self, text, return_word_ids=False, llama_template=None, images=[], prevent_empty_text=False, thinking=True, **kwargs):
|
||||
# Krea2 conditions on the no-think template; thinking=True drops the empty <think> block qwen3vl adds.
|
||||
return super().tokenize_with_weights(text, return_word_ids=return_word_ids, llama_template=llama_template, images=images, prevent_empty_text=prevent_empty_text, thinking=thinking, **kwargs)
|
||||
|
||||
|
||||
class Krea2Qwen3VLClipModel(comfy.text_encoders.qwen3vl.Qwen3VLClipModel):
|
||||
def __init__(self, device="cpu", dtype=None, attention_mask=True, model_options={}):
|
||||
super().__init__(device=device, layer=KREA2_TAP_LAYERS, layer_idx=None, dtype=dtype,
|
||||
attention_mask=attention_mask, model_options=model_options, model_type="qwen3vl_4b")
|
||||
|
||||
|
||||
class Krea2TEModel(sd1_clip.SD1ClipModel):
|
||||
def __init__(self, device="cpu", dtype=None, model_options={}):
|
||||
super().__init__(device=device, dtype=dtype, name="qwen3vl_4b", clip_model=Krea2Qwen3VLClipModel, model_options=model_options)
|
||||
|
||||
def encode_token_weights(self, token_weight_pairs, template_end=-1):
|
||||
out, pooled, extra = super().encode_token_weights(token_weight_pairs) # out: (B, 12, seq, 2560)
|
||||
tok_pairs = token_weight_pairs["qwen3vl_4b"][0]
|
||||
|
||||
# Strip the system + user-opening prefix
|
||||
count_im_start = 0
|
||||
if template_end == -1:
|
||||
for i, v in enumerate(tok_pairs):
|
||||
elem = v[0]
|
||||
if not torch.is_tensor(elem) and isinstance(elem, numbers.Integral):
|
||||
if elem == 151644 and count_im_start < 2:
|
||||
template_end = i
|
||||
count_im_start += 1
|
||||
if out.shape[2] > (template_end + 3):
|
||||
if tok_pairs[template_end + 1][0] == 872: # "user"
|
||||
if tok_pairs[template_end + 2][0] == 198: # "\n"
|
||||
template_end += 3
|
||||
|
||||
out = out[:, :, template_end:]
|
||||
|
||||
b, n, seq, h = out.shape
|
||||
# Flatten the 12-layer axis into the feature dim: (B, seq, 12*2560). Unpacked in the model.
|
||||
out = out.permute(0, 2, 1, 3).reshape(b, seq, n * h)
|
||||
|
||||
if "attention_mask" in extra:
|
||||
extra["attention_mask"] = extra["attention_mask"][:, template_end:]
|
||||
if extra["attention_mask"].sum() == torch.numel(extra["attention_mask"]):
|
||||
extra.pop("attention_mask")
|
||||
|
||||
return out, pooled, extra
|
||||
|
||||
|
||||
def te(dtype_llama=None, llama_quantization_metadata=None):
|
||||
class Krea2TEModel_(Krea2TEModel):
|
||||
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 Krea2TEModel_
|
||||
@@ -167,7 +167,7 @@ class Qwen3VLTokenizer(sd1_clip.SD1Tokenizer):
|
||||
embed_count = 0
|
||||
for r in tokens[key_name]:
|
||||
for i in range(len(r)):
|
||||
if r[i][0] == 151655: # <|image_pad|>
|
||||
if isinstance(r[i][0], (int, float)) and r[i][0] == 151655: # <|image_pad|>
|
||||
if len(images) > embed_count:
|
||||
r[i] = ({"type": "image", "data": images[embed_count], "original_type": "image"},) + r[i][1:]
|
||||
embed_count += 1
|
||||
|
||||
@@ -818,6 +818,44 @@ def z_image_to_diffusers(mmdit_config, output_prefix=""):
|
||||
|
||||
return key_map
|
||||
|
||||
def krea2_to_diffusers(mmdit_config, output_prefix=""):
|
||||
n_layers = mmdit_config.get("layers", 0)
|
||||
n_txt_layerwise = 2 # TextFusionTransformer hardcodes 2 layerwise + 2 refiner blocks
|
||||
n_txt_refiner = 2
|
||||
key_map = {}
|
||||
|
||||
def add_block(prefix_to, prefix_from):
|
||||
block_map = {
|
||||
"attn.to_q": "attn.wq", "attn.to_k": "attn.wk", "attn.to_v": "attn.wv",
|
||||
"attn.to_gate": "attn.gate", "attn.to_out.0": "attn.wo",
|
||||
"attn.to_out": "attn.wo", # some tools drop the ".0" on to_out
|
||||
"ff.gate": "mlp.gate", "ff.up": "mlp.up", "ff.down": "mlp.down",
|
||||
}
|
||||
for d, c in block_map.items():
|
||||
key_map["{}.{}.weight".format(prefix_to, d)] = "{}{}.{}.weight".format(output_prefix, prefix_from, c)
|
||||
|
||||
for i in range(n_layers):
|
||||
add_block("transformer_blocks.{}".format(i), "blocks.{}".format(i))
|
||||
for i in range(n_txt_layerwise):
|
||||
add_block("text_fusion.layerwise_blocks.{}".format(i), "txtfusion.layerwise_blocks.{}".format(i))
|
||||
for i in range(n_txt_refiner):
|
||||
add_block("text_fusion.refiner_blocks.{}".format(i), "txtfusion.refiner_blocks.{}".format(i))
|
||||
|
||||
MAP_BASIC = [
|
||||
("img_in", "first"),
|
||||
("time_embed.linear_1", "tmlp.0"),
|
||||
("time_embed.linear_2", "tmlp.2"),
|
||||
("time_mod_proj", "tproj.1"),
|
||||
("txt_in.linear_1", "txtmlp.1"),
|
||||
("txt_in.linear_2", "txtmlp.3"),
|
||||
("text_fusion.projector", "txtfusion.projector"),
|
||||
("final_layer.linear", "last.linear"),
|
||||
]
|
||||
for d, c in MAP_BASIC:
|
||||
key_map["{}.weight".format(d)] = "{}{}.weight".format(output_prefix, c)
|
||||
|
||||
return key_map
|
||||
|
||||
def repeat_to_batch_size(tensor, batch_size, dim=0):
|
||||
if tensor.shape[dim] > batch_size:
|
||||
return tensor.narrow(dim, 0, batch_size)
|
||||
|
||||
@@ -98,12 +98,24 @@ def _parse_cli_feature_flags() -> dict[str, Any]:
|
||||
|
||||
|
||||
# Default server capabilities
|
||||
def _hf_auth_eligible_at_startup() -> bool:
|
||||
"""Snapshot eligibility once at feature-flag init time.
|
||||
|
||||
Imports lazily because the flags module loads very early in the
|
||||
server boot sequence — earlier than the model_downloader package.
|
||||
"""
|
||||
from app.model_downloader.hf_auth.eligibility import is_hf_auth_eligible
|
||||
return is_hf_auth_eligible()
|
||||
|
||||
|
||||
_CORE_FEATURE_FLAGS: dict[str, Any] = {
|
||||
"supports_preview_metadata": True,
|
||||
"max_upload_size": args.max_upload_size * 1024 * 1024, # Convert MB to bytes
|
||||
"extension": {"manager": {"supports_v4": True}},
|
||||
"node_replacements": True,
|
||||
"assets": args.enable_assets,
|
||||
"server_side_model_downloads": True,
|
||||
"hf_auth_eligible": _hf_auth_eligible_at_startup(),
|
||||
}
|
||||
|
||||
# CLI-provided flags cannot overwrite core flags
|
||||
|
||||
@@ -891,6 +891,14 @@ class Tracks(ComfyTypeIO):
|
||||
track_visibility: torch.Tensor
|
||||
Type = TrackDict
|
||||
|
||||
@comfytype(io_type="DICT")
|
||||
class Dict(ComfyTypeIO):
|
||||
Type = dict
|
||||
|
||||
@comfytype(io_type="ARRAY")
|
||||
class Array(ComfyTypeIO):
|
||||
Type = list
|
||||
|
||||
@comfytype(io_type="COMFY_MULTITYPED_V3")
|
||||
class MultiType:
|
||||
Type = Any
|
||||
@@ -1279,6 +1287,19 @@ class Color(ComfyTypeIO):
|
||||
def as_dict(self):
|
||||
return super().as_dict()
|
||||
|
||||
|
||||
@comfytype(io_type="COLORS")
|
||||
class Colors(ComfyTypeIO):
|
||||
Type = list[Color.Type]
|
||||
|
||||
class Input(WidgetInput):
|
||||
def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None,
|
||||
socketless: bool=True, default: list[str]=None, advanced: bool=None):
|
||||
super().__init__(id, display_name, optional, tooltip, None, default, socketless, None, None, None, None, advanced)
|
||||
if default is None:
|
||||
self.default = []
|
||||
|
||||
|
||||
@comfytype(io_type="BOUNDING_BOX")
|
||||
class BoundingBox(ComfyTypeIO):
|
||||
class BoundingBoxDict(TypedDict):
|
||||
@@ -1326,6 +1347,20 @@ class Curve(ComfyTypeIO):
|
||||
return d
|
||||
|
||||
|
||||
@comfytype(io_type="BOUNDING_BOXES")
|
||||
class BoundingBoxes(ComfyTypeIO):
|
||||
class BoundingBoxWithMetadata(BoundingBox.BoundingBoxDict):
|
||||
metadata: dict
|
||||
Type = list[BoundingBoxWithMetadata]
|
||||
|
||||
class Input(WidgetInput):
|
||||
def __init__(self, id: str, display_name: str=None, optional=False, tooltip: str=None,
|
||||
socketless: bool=True, default: list[dict]=None, advanced: bool=None):
|
||||
super().__init__(id, display_name, optional, tooltip, None, default, socketless, None, None, None, None, advanced)
|
||||
if default is None:
|
||||
self.default = []
|
||||
|
||||
|
||||
@comfytype(io_type="HISTOGRAM")
|
||||
class Histogram(ComfyTypeIO):
|
||||
"""A histogram represented as a list of bin counts."""
|
||||
@@ -2376,6 +2411,8 @@ __all__ = [
|
||||
"AnyType",
|
||||
"MultiType",
|
||||
"Tracks",
|
||||
"Dict",
|
||||
"Array",
|
||||
"Color",
|
||||
# Dynamic Types
|
||||
"MatchType",
|
||||
@@ -2394,6 +2431,8 @@ __all__ = [
|
||||
"PriceBadgeDepends",
|
||||
"PriceBadge",
|
||||
"BoundingBox",
|
||||
"BoundingBoxes",
|
||||
"Colors",
|
||||
"Curve",
|
||||
"Histogram",
|
||||
"Range",
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import Literal
|
||||
from typing import Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
@@ -163,15 +163,31 @@ class SeedanceVirtualLibraryCreateAssetRequest(BaseModel):
|
||||
asset_type: str | None = Field(None, description="BytePlus asset type. Defaults to Image server-side when omitted.")
|
||||
|
||||
|
||||
# Dollars per 1K tokens, keyed by (model_id, has_video_input).
|
||||
# Dollars per 1K tokens, keyed by (model_id, has_video_input, resolution).
|
||||
SEEDANCE2_PRICE_PER_1K_TOKENS = {
|
||||
("dreamina-seedance-2-0-260128", False): 0.007,
|
||||
("dreamina-seedance-2-0-260128", True): 0.0043,
|
||||
("dreamina-seedance-2-0-fast-260128", False): 0.0056,
|
||||
("dreamina-seedance-2-0-fast-260128", True): 0.0033,
|
||||
("dreamina-seedance-2-0-260128", False, "480p"): 0.007,
|
||||
("dreamina-seedance-2-0-260128", True, "480p"): 0.0043,
|
||||
("dreamina-seedance-2-0-260128", False, "720p"): 0.007,
|
||||
("dreamina-seedance-2-0-260128", True, "720p"): 0.0043,
|
||||
("dreamina-seedance-2-0-260128", False, "1080p"): 0.0077,
|
||||
("dreamina-seedance-2-0-260128", True, "1080p"): 0.0047,
|
||||
("dreamina-seedance-2-0-260128", False, "4k"): 0.004,
|
||||
("dreamina-seedance-2-0-260128", True, "4k"): 0.0024,
|
||||
("dreamina-seedance-2-0-fast-260128", False, "480p"): 0.0056,
|
||||
("dreamina-seedance-2-0-fast-260128", True, "480p"): 0.0033,
|
||||
("dreamina-seedance-2-0-fast-260128", False, "720p"): 0.0056,
|
||||
("dreamina-seedance-2-0-fast-260128", True, "720p"): 0.0033,
|
||||
("dreamina-seedance-2-0-mini", False, "480p"): 0.0035,
|
||||
("dreamina-seedance-2-0-mini", True, "480p"): 0.0021,
|
||||
("dreamina-seedance-2-0-mini", False, "720p"): 0.0035,
|
||||
("dreamina-seedance-2-0-mini", True, "720p"): 0.0021,
|
||||
}
|
||||
|
||||
|
||||
def seedance2_price_per_1k_tokens(model_id: str, has_video_input: bool, resolution: str) -> float | None:
|
||||
return SEEDANCE2_PRICE_PER_1K_TOKENS.get((model_id, has_video_input, resolution))
|
||||
|
||||
|
||||
RECOMMENDED_PRESETS = [
|
||||
("1024x1024 (1:1)", 1024, 1024),
|
||||
("864x1152 (3:4)", 864, 1152),
|
||||
@@ -266,6 +282,10 @@ SEEDANCE2_REF_VIDEO_PIXEL_LIMITS = {
|
||||
"480p": {"min": 409_600, "max": 927_408},
|
||||
"720p": {"min": 409_600, "max": 927_408},
|
||||
},
|
||||
"dreamina-seedance-2-0-mini": {
|
||||
"480p": {"min": 409_600, "max": 927_408},
|
||||
"720p": {"min": 409_600, "max": 927_408},
|
||||
},
|
||||
}
|
||||
|
||||
# The time in this dictionary are given for 10 seconds duration.
|
||||
@@ -296,3 +316,36 @@ VIDEO_TASKS_EXECUTION_TIME = {
|
||||
"1080p": 150,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class SeedAudioConfig(BaseModel):
|
||||
format: str = Field(default="mp3")
|
||||
sample_rate: int = Field(default=24000)
|
||||
speech_rate: int = Field(default=0)
|
||||
loudness_rate: int = Field(default=0)
|
||||
pitch_rate: int = Field(default=0)
|
||||
|
||||
|
||||
class SeedAudioReference(BaseModel):
|
||||
speaker: str | None = Field(default=None)
|
||||
audio_data: str | None = Field(default=None)
|
||||
audio_url: str | None = Field(default=None)
|
||||
image_data: str | None = Field(default=None)
|
||||
image_url: str | None = Field(default=None)
|
||||
|
||||
|
||||
class SeedAudioRequest(BaseModel):
|
||||
model: str = Field(default="seed-audio-1.0")
|
||||
text_prompt: str = Field(...)
|
||||
references: list[SeedAudioReference] | None = Field(default=None)
|
||||
audio_config: SeedAudioConfig = Field(default_factory=SeedAudioConfig)
|
||||
watermark: dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class SeedAudioResponse(BaseModel):
|
||||
audio: str | None = Field(default=None)
|
||||
url: str | None = Field(default=None)
|
||||
duration: float | None = Field(default=None)
|
||||
original_duration: float | None = Field(default=None)
|
||||
code: int | None = Field(default=None)
|
||||
message: str | None = Field(default=None)
|
||||
|
||||
@@ -121,6 +121,7 @@ class GeminiGenerationConfig(BaseModel):
|
||||
topK: int | None = Field(None, ge=1)
|
||||
topP: float | None = Field(None, ge=0.0, le=1.0)
|
||||
thinkingConfig: GeminiThinkingConfig | None = Field(None)
|
||||
responseModalities: list[str] | None = Field(None)
|
||||
|
||||
|
||||
class GeminiImageOutputOptions(BaseModel):
|
||||
|
||||
@@ -33,53 +33,6 @@ class IdeogramColorPalette(
|
||||
)
|
||||
|
||||
|
||||
class ImageRequest(BaseModel):
|
||||
aspect_ratio: Optional[str] = Field(
|
||||
None,
|
||||
description="Optional. The aspect ratio (e.g., 'ASPECT_16_9', 'ASPECT_1_1'). Cannot be used with resolution. Defaults to 'ASPECT_1_1' if unspecified.",
|
||||
)
|
||||
color_palette: Optional[Dict[str, Any]] = Field(
|
||||
None, description='Optional. Color palette object. Only for V_2, V_2_TURBO.'
|
||||
)
|
||||
magic_prompt_option: Optional[str] = Field(
|
||||
None, description="Optional. MagicPrompt usage ('AUTO', 'ON', 'OFF')."
|
||||
)
|
||||
model: str = Field(..., description="The model used (e.g., 'V_2', 'V_2A_TURBO')")
|
||||
negative_prompt: Optional[str] = Field(
|
||||
None,
|
||||
description='Optional. Description of what to exclude. Only for V_1, V_1_TURBO, V_2, V_2_TURBO.',
|
||||
)
|
||||
num_images: Optional[int] = Field(
|
||||
1,
|
||||
description='Optional. Number of images to generate (1-8). Defaults to 1.',
|
||||
ge=1,
|
||||
le=8,
|
||||
)
|
||||
prompt: str = Field(
|
||||
..., description='Required. The prompt to use to generate the image.'
|
||||
)
|
||||
resolution: Optional[str] = Field(
|
||||
None,
|
||||
description="Optional. Resolution (e.g., 'RESOLUTION_1024_1024'). Only for model V_2. Cannot be used with aspect_ratio.",
|
||||
)
|
||||
seed: Optional[int] = Field(
|
||||
None,
|
||||
description='Optional. A number between 0 and 2147483647.',
|
||||
ge=0,
|
||||
le=2147483647,
|
||||
)
|
||||
style_type: Optional[str] = Field(
|
||||
None,
|
||||
description="Optional. Style type ('AUTO', 'GENERAL', 'REALISTIC', 'DESIGN', 'RENDER_3D', 'ANIME'). Only for models V_2 and above.",
|
||||
)
|
||||
|
||||
|
||||
class IdeogramGenerateRequest(BaseModel):
|
||||
image_request: ImageRequest = Field(
|
||||
..., description='The image generation request parameters.'
|
||||
)
|
||||
|
||||
|
||||
class Datum(BaseModel):
|
||||
is_image_safe: Optional[bool] = Field(
|
||||
None, description='Indicates whether the image is considered safe.'
|
||||
@@ -113,20 +66,6 @@ class StyleCode(RootModel[str]):
|
||||
root: str = Field(..., pattern='^[0-9A-Fa-f]{8}$')
|
||||
|
||||
|
||||
class Datum1(BaseModel):
|
||||
is_image_safe: Optional[bool] = None
|
||||
prompt: Optional[str] = None
|
||||
resolution: Optional[str] = None
|
||||
seed: Optional[int] = None
|
||||
style_type: Optional[str] = None
|
||||
url: Optional[str] = None
|
||||
|
||||
|
||||
class IdeogramV3IdeogramResponse(BaseModel):
|
||||
created: Optional[datetime] = None
|
||||
data: Optional[List[Datum1]] = None
|
||||
|
||||
|
||||
class RenderingSpeed1(str, Enum):
|
||||
TURBO = 'TURBO'
|
||||
DEFAULT = 'DEFAULT'
|
||||
|
||||
@@ -10,6 +10,7 @@ from pydantic import BaseModel, Field, confloat
|
||||
class LumaIO:
|
||||
LUMA_REF = "LUMA_REF"
|
||||
LUMA_CONCEPTS = "LUMA_CONCEPTS"
|
||||
LUMA_RAY32_KEYFRAME = "LUMA_RAY32_KEYFRAME"
|
||||
|
||||
|
||||
class LumaReference:
|
||||
@@ -20,13 +21,14 @@ class LumaReference:
|
||||
def create_api_model(self, download_url: str):
|
||||
return LumaImageRef(url=download_url, weight=self.weight)
|
||||
|
||||
|
||||
class LumaReferenceChain:
|
||||
def __init__(self, first_ref: LumaReference=None):
|
||||
def __init__(self, first_ref: LumaReference = None):
|
||||
self.refs: list[LumaReference] = []
|
||||
if first_ref:
|
||||
self.refs.append(first_ref)
|
||||
|
||||
def add(self, luma_ref: LumaReference=None):
|
||||
def add(self, luma_ref: LumaReference = None):
|
||||
self.refs.append(luma_ref)
|
||||
|
||||
def create_api_model(self, download_urls: list[str], max_refs=4):
|
||||
@@ -124,7 +126,7 @@ def get_luma_concepts(include_none=False):
|
||||
"pull_out",
|
||||
"aerial",
|
||||
"crane_up",
|
||||
"eye_level"
|
||||
"eye_level",
|
||||
]
|
||||
|
||||
|
||||
@@ -162,8 +164,8 @@ class LumaVideoModelOutputDuration(str, Enum):
|
||||
|
||||
|
||||
class LumaGenerationType(str, Enum):
|
||||
video = 'video'
|
||||
image = 'image'
|
||||
video = "video"
|
||||
image = "image"
|
||||
|
||||
|
||||
class LumaState(str, Enum):
|
||||
@@ -174,86 +176,109 @@ class LumaState(str, Enum):
|
||||
|
||||
|
||||
class LumaAssets(BaseModel):
|
||||
video: Optional[str] = Field(None, description='The URL of the video')
|
||||
image: Optional[str] = Field(None, description='The URL of the image')
|
||||
progress_video: Optional[str] = Field(None, description='The URL of the progress video')
|
||||
video: Optional[str] = Field(None, description="The URL of the video")
|
||||
image: Optional[str] = Field(None, description="The URL of the image")
|
||||
progress_video: Optional[str] = Field(None, description="The URL of the progress video")
|
||||
|
||||
|
||||
class LumaImageRef(BaseModel):
|
||||
"""Used for image gen"""
|
||||
url: str = Field(..., description='The URL of the image reference')
|
||||
weight: confloat(ge=0.0, le=1.0) = Field(..., description='The weight of the image reference')
|
||||
|
||||
url: str = Field(..., description="The URL of the image reference")
|
||||
weight: confloat(ge=0.0, le=1.0) = Field(..., description="The weight of the image reference")
|
||||
|
||||
|
||||
class LumaImageReference(BaseModel):
|
||||
"""Used for video gen"""
|
||||
type: Optional[str] = Field('image', description='Input type, defaults to image')
|
||||
url: str = Field(..., description='The URL of the image')
|
||||
|
||||
type: Optional[str] = Field("image", description="Input type, defaults to image")
|
||||
url: str = Field(..., description="The URL of the image")
|
||||
|
||||
|
||||
class LumaModifyImageRef(BaseModel):
|
||||
url: str = Field(..., description='The URL of the image reference')
|
||||
weight: confloat(ge=0.0, le=1.0) = Field(..., description='The weight of the image reference')
|
||||
url: str = Field(..., description="The URL of the image reference")
|
||||
weight: confloat(ge=0.0, le=1.0) = Field(..., description="The weight of the image reference")
|
||||
|
||||
|
||||
class LumaCharacterRef(BaseModel):
|
||||
identity0: LumaImageIdentity = Field(..., description='The image identity object')
|
||||
identity0: LumaImageIdentity = Field(..., description="The image identity object")
|
||||
|
||||
|
||||
class LumaImageIdentity(BaseModel):
|
||||
images: list[str] = Field(..., description='The URLs of the image identity')
|
||||
images: list[str] = Field(..., description="The URLs of the image identity")
|
||||
|
||||
|
||||
class LumaGenerationReference(BaseModel):
|
||||
type: str = Field('generation', description='Input type, defaults to generation')
|
||||
id: str = Field(..., description='The ID of the generation')
|
||||
type: str = Field("generation", description="Input type, defaults to generation")
|
||||
id: str = Field(..., description="The ID of the generation")
|
||||
|
||||
|
||||
class LumaKeyframes(BaseModel):
|
||||
frame0: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description='')
|
||||
frame1: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description='')
|
||||
frame0: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description="")
|
||||
frame1: Optional[Union[LumaImageReference, LumaGenerationReference]] = Field(None, description="")
|
||||
|
||||
|
||||
class LumaConceptObject(BaseModel):
|
||||
key: str = Field(..., description='Camera Concept name')
|
||||
key: str = Field(..., description="Camera Concept name")
|
||||
|
||||
|
||||
class LumaImageGenerationRequest(BaseModel):
|
||||
prompt: str = Field(..., description='The prompt of the generation')
|
||||
model: LumaImageModel = Field(LumaImageModel.photon_1, description='The image model used for the generation')
|
||||
aspect_ratio: Optional[LumaAspectRatio] = Field(LumaAspectRatio.ratio_16_9, description='The aspect ratio of the generation')
|
||||
image_ref: Optional[list[LumaImageRef]] = Field(None, description='List of image reference objects')
|
||||
style_ref: Optional[list[LumaImageRef]] = Field(None, description='List of style reference objects')
|
||||
character_ref: Optional[LumaCharacterRef] = Field(None, description='The image identity object')
|
||||
modify_image_ref: Optional[LumaModifyImageRef] = Field(None, description='The modify image reference object')
|
||||
prompt: str = Field(..., description="The prompt of the generation")
|
||||
model: LumaImageModel = Field(LumaImageModel.photon_1, description="The image model used for the generation")
|
||||
aspect_ratio: Optional[LumaAspectRatio] = Field(LumaAspectRatio.ratio_16_9)
|
||||
image_ref: Optional[list[LumaImageRef]] = Field(None, description="List of image reference objects")
|
||||
style_ref: Optional[list[LumaImageRef]] = Field(None, description="List of style reference objects")
|
||||
character_ref: Optional[LumaCharacterRef] = Field(None, description="The image identity object")
|
||||
modify_image_ref: Optional[LumaModifyImageRef] = Field(None, description="The modify image reference object")
|
||||
|
||||
|
||||
class LumaGenerationRequest(BaseModel):
|
||||
prompt: str = Field(..., description='The prompt of the generation')
|
||||
model: LumaVideoModel = Field(LumaVideoModel.ray_2, description='The video model used for the generation')
|
||||
duration: Optional[LumaVideoModelOutputDuration] = Field(None, description='The duration of the generation')
|
||||
aspect_ratio: Optional[LumaAspectRatio] = Field(None, description='The aspect ratio of the generation')
|
||||
resolution: Optional[LumaVideoOutputResolution] = Field(None, description='The resolution of the generation')
|
||||
loop: Optional[bool] = Field(None, description='Whether to loop the video')
|
||||
keyframes: Optional[LumaKeyframes] = Field(None, description='The keyframes of the generation')
|
||||
concepts: Optional[list[LumaConceptObject]] = Field(None, description='Camera Concepts to apply to generation')
|
||||
prompt: str = Field(..., description="The prompt of the generation")
|
||||
model: LumaVideoModel = Field(LumaVideoModel.ray_2, description="The video model used for the generation")
|
||||
duration: Optional[LumaVideoModelOutputDuration] = Field(None, description="The duration of the generation")
|
||||
aspect_ratio: Optional[LumaAspectRatio] = Field(None, description="The aspect ratio of the generation")
|
||||
resolution: Optional[LumaVideoOutputResolution] = Field(None, description="The resolution of the generation")
|
||||
loop: Optional[bool] = Field(None, description="Whether to loop the video")
|
||||
keyframes: Optional[LumaKeyframes] = Field(None, description="The keyframes of the generation")
|
||||
concepts: Optional[list[LumaConceptObject]] = Field(None, description="Camera Concepts to apply to generation")
|
||||
|
||||
|
||||
class LumaGeneration(BaseModel):
|
||||
id: str = Field(..., description='The ID of the generation')
|
||||
generation_type: LumaGenerationType = Field(..., description='Generation type, image or video')
|
||||
state: LumaState = Field(..., description='The state of the generation')
|
||||
failure_reason: Optional[str] = Field(None, description='The reason for the state of the generation')
|
||||
created_at: str = Field(..., description='The date and time when the generation was created')
|
||||
assets: Optional[LumaAssets] = Field(None, description='The assets of the generation')
|
||||
model: str = Field(..., description='The model used for the generation')
|
||||
request: Union[LumaGenerationRequest, LumaImageGenerationRequest] = Field(..., description="The request used for the generation")
|
||||
id: str = Field(..., description="The ID of the generation")
|
||||
generation_type: LumaGenerationType = Field(..., description="Generation type, image or video")
|
||||
state: LumaState = Field(..., description="The state of the generation")
|
||||
failure_reason: Optional[str] = Field(None, description="The reason for the state of the generation")
|
||||
created_at: str = Field(..., description="The date and time when the generation was created")
|
||||
assets: Optional[LumaAssets] = Field(None, description="The assets of the generation")
|
||||
model: str = Field(..., description="The model used for the generation")
|
||||
request: Union[LumaGenerationRequest, LumaImageGenerationRequest] = Field(...)
|
||||
|
||||
|
||||
class Luma2ImageRef(BaseModel):
|
||||
url: str | None = None
|
||||
data: str | None = None
|
||||
media_type: str | None = None
|
||||
generation_id: str | None = Field(None, description="reference a prior generation (extend / source reuse)")
|
||||
|
||||
|
||||
class Luma2VideoEdit(BaseModel):
|
||||
"""Edit controls for Ray 3.2 ``video_edit`` generations."""
|
||||
|
||||
auto_controls: bool | None = Field(None, description="derive a conditioning schedule from the source (recommended)")
|
||||
strength: str | None = Field(None, description="'adhere_1' .. 'reimagine_3'; constrained by IO.Combo")
|
||||
|
||||
|
||||
class Luma2VideoOptions(BaseModel):
|
||||
"""Ray 3.2 ``video`` output settings (text / image / keyframe / edit / extend)."""
|
||||
|
||||
resolution: str | None = Field(None, description="360p | 540p | 720p | 1080p")
|
||||
duration: str | None = Field(None, description="5s | 10s")
|
||||
loop: bool | None = Field(None)
|
||||
start_frame: Luma2ImageRef | None = Field(None)
|
||||
end_frame: Luma2ImageRef | None = Field(None)
|
||||
keyframes: list[Luma2ImageRef] | None = Field(None)
|
||||
keyframe_indexes: list[int] | None = Field(None)
|
||||
edit: Luma2VideoEdit | None = Field(None)
|
||||
|
||||
|
||||
class Luma2GenerationRequest(BaseModel):
|
||||
@@ -266,6 +291,7 @@ class Luma2GenerationRequest(BaseModel):
|
||||
web_search: bool | None = None
|
||||
image_ref: list[Luma2ImageRef] | None = None
|
||||
source: Luma2ImageRef | None = None
|
||||
video: Luma2VideoOptions | None = Field(None)
|
||||
|
||||
|
||||
class Luma2Generation(BaseModel):
|
||||
@@ -277,3 +303,31 @@ class Luma2Generation(BaseModel):
|
||||
output: list[LumaImageReference] | None = None
|
||||
failure_reason: str | None = None
|
||||
failure_code: str | None = None
|
||||
|
||||
|
||||
# --- Ray 3.2 multi-keyframe chain ---
|
||||
|
||||
LUMA_KEYFRAME_MODE_FRACTION = "fraction" # value in [0.0, 1.0] of the output video duration
|
||||
LUMA_KEYFRAME_MODE_SECONDS = "seconds" # absolute time, in seconds, from the start of the output
|
||||
|
||||
|
||||
class LumaRay32KeyframeItem:
|
||||
"""One guide image anchored at a position on the Ray 3.2 output timeline."""
|
||||
|
||||
def __init__(self, image: torch.Tensor, mode: str, value: float):
|
||||
self.image = image
|
||||
self.mode = mode # LUMA_KEYFRAME_MODE_FRACTION | LUMA_KEYFRAME_MODE_SECONDS
|
||||
self.value = value
|
||||
|
||||
|
||||
class LumaRay32KeyframeChain:
|
||||
def __init__(self):
|
||||
self.items: list[LumaRay32KeyframeItem] = []
|
||||
|
||||
def add(self, item: LumaRay32KeyframeItem) -> None:
|
||||
self.items.append(item)
|
||||
|
||||
def clone(self) -> "LumaRay32KeyframeChain":
|
||||
c = LumaRay32KeyframeChain()
|
||||
c.items = list(self.items)
|
||||
return c
|
||||
|
||||
@@ -1,147 +0,0 @@
|
||||
from enum import Enum
|
||||
from typing import Optional
|
||||
|
||||
from pydantic import BaseModel, Field, confloat
|
||||
|
||||
|
||||
class StabilityFormat(str, Enum):
|
||||
png = 'png'
|
||||
jpeg = 'jpeg'
|
||||
webp = 'webp'
|
||||
|
||||
|
||||
class StabilityAspectRatio(str, Enum):
|
||||
ratio_1_1 = "1:1"
|
||||
ratio_16_9 = "16:9"
|
||||
ratio_9_16 = "9:16"
|
||||
ratio_3_2 = "3:2"
|
||||
ratio_2_3 = "2:3"
|
||||
ratio_5_4 = "5:4"
|
||||
ratio_4_5 = "4:5"
|
||||
ratio_21_9 = "21:9"
|
||||
ratio_9_21 = "9:21"
|
||||
|
||||
|
||||
def get_stability_style_presets(include_none=True):
|
||||
presets = []
|
||||
if include_none:
|
||||
presets.append("None")
|
||||
return presets + [x.value for x in StabilityStylePreset]
|
||||
|
||||
|
||||
class StabilityStylePreset(str, Enum):
|
||||
_3d_model = "3d-model"
|
||||
analog_film = "analog-film"
|
||||
anime = "anime"
|
||||
cinematic = "cinematic"
|
||||
comic_book = "comic-book"
|
||||
digital_art = "digital-art"
|
||||
enhance = "enhance"
|
||||
fantasy_art = "fantasy-art"
|
||||
isometric = "isometric"
|
||||
line_art = "line-art"
|
||||
low_poly = "low-poly"
|
||||
modeling_compound = "modeling-compound"
|
||||
neon_punk = "neon-punk"
|
||||
origami = "origami"
|
||||
photographic = "photographic"
|
||||
pixel_art = "pixel-art"
|
||||
tile_texture = "tile-texture"
|
||||
|
||||
|
||||
class Stability_SD3_5_Model(str, Enum):
|
||||
sd3_5_large = "sd3.5-large"
|
||||
# sd3_5_large_turbo = "sd3.5-large-turbo"
|
||||
sd3_5_medium = "sd3.5-medium"
|
||||
|
||||
|
||||
class Stability_SD3_5_GenerationMode(str, Enum):
|
||||
text_to_image = "text-to-image"
|
||||
image_to_image = "image-to-image"
|
||||
|
||||
|
||||
class StabilityStable3_5Request(BaseModel):
|
||||
model: str = Field(...)
|
||||
mode: str = Field(...)
|
||||
prompt: str = Field(...)
|
||||
negative_prompt: Optional[str] = Field(None)
|
||||
aspect_ratio: Optional[str] = Field(None)
|
||||
seed: Optional[int] = Field(None)
|
||||
output_format: Optional[str] = Field(StabilityFormat.png.value)
|
||||
image: Optional[str] = Field(None)
|
||||
style_preset: Optional[str] = Field(None)
|
||||
cfg_scale: float = Field(...)
|
||||
strength: Optional[confloat(ge=0.0, le=1.0)] = Field(None)
|
||||
|
||||
|
||||
class StabilityUpscaleConservativeRequest(BaseModel):
|
||||
prompt: str = Field(...)
|
||||
negative_prompt: Optional[str] = Field(None)
|
||||
seed: Optional[int] = Field(None)
|
||||
output_format: Optional[str] = Field(StabilityFormat.png.value)
|
||||
image: Optional[str] = Field(None)
|
||||
creativity: Optional[confloat(ge=0.2, le=0.5)] = Field(None)
|
||||
|
||||
|
||||
class StabilityUpscaleCreativeRequest(BaseModel):
|
||||
prompt: str = Field(...)
|
||||
negative_prompt: Optional[str] = Field(None)
|
||||
seed: Optional[int] = Field(None)
|
||||
output_format: Optional[str] = Field(StabilityFormat.png.value)
|
||||
image: Optional[str] = Field(None)
|
||||
creativity: Optional[confloat(ge=0.1, le=0.5)] = Field(None)
|
||||
style_preset: Optional[str] = Field(None)
|
||||
|
||||
|
||||
class StabilityStableUltraRequest(BaseModel):
|
||||
prompt: str = Field(...)
|
||||
negative_prompt: Optional[str] = Field(None)
|
||||
aspect_ratio: Optional[str] = Field(None)
|
||||
seed: Optional[int] = Field(None)
|
||||
output_format: Optional[str] = Field(StabilityFormat.png.value)
|
||||
image: Optional[str] = Field(None)
|
||||
style_preset: Optional[str] = Field(None)
|
||||
strength: Optional[confloat(ge=0.0, le=1.0)] = Field(None)
|
||||
|
||||
|
||||
class StabilityStableUltraResponse(BaseModel):
|
||||
image: Optional[str] = Field(None)
|
||||
finish_reason: Optional[str] = Field(None)
|
||||
seed: Optional[int] = Field(None)
|
||||
|
||||
|
||||
class StabilityResultsGetResponse(BaseModel):
|
||||
image: Optional[str] = Field(None)
|
||||
finish_reason: Optional[str] = Field(None)
|
||||
seed: Optional[int] = Field(None)
|
||||
id: Optional[str] = Field(None)
|
||||
name: Optional[str] = Field(None)
|
||||
errors: Optional[list[str]] = Field(None)
|
||||
status: Optional[str] = Field(None)
|
||||
result: Optional[str] = Field(None)
|
||||
|
||||
|
||||
class StabilityAsyncResponse(BaseModel):
|
||||
id: Optional[str] = Field(None)
|
||||
|
||||
|
||||
class StabilityTextToAudioRequest(BaseModel):
|
||||
model: str = Field(...)
|
||||
prompt: str = Field(...)
|
||||
duration: int = Field(190, ge=1, le=190)
|
||||
seed: int = Field(0, ge=0, le=4294967294)
|
||||
steps: int = Field(8, ge=4, le=8)
|
||||
output_format: str = Field("wav")
|
||||
|
||||
|
||||
class StabilityAudioToAudioRequest(StabilityTextToAudioRequest):
|
||||
strength: float = Field(0.01, ge=0.01, le=1.0)
|
||||
|
||||
|
||||
class StabilityAudioInpaintRequest(StabilityTextToAudioRequest):
|
||||
mask_start: int = Field(30, ge=0, le=190)
|
||||
mask_end: int = Field(190, ge=0, le=190)
|
||||
|
||||
|
||||
class StabilityAudioResponse(BaseModel):
|
||||
audio: Optional[str] = Field(None)
|
||||
@@ -1,3 +1,4 @@
|
||||
import base64
|
||||
import hashlib
|
||||
import logging
|
||||
import math
|
||||
@@ -15,12 +16,15 @@ from comfy_api_nodes.apis.bytedance import (
|
||||
RECOMMENDED_PRESETS_SEEDREAM_4_0,
|
||||
RECOMMENDED_PRESETS_SEEDREAM_4_5,
|
||||
RECOMMENDED_PRESETS_SEEDREAM_5_LITE,
|
||||
SEEDANCE2_PRICE_PER_1K_TOKENS,
|
||||
SEEDANCE2_REF_VIDEO_PIXEL_LIMITS,
|
||||
VIDEO_TASKS_EXECUTION_TIME,
|
||||
GetAssetResponse,
|
||||
Image2VideoTaskCreationRequest,
|
||||
ImageTaskCreationResponse,
|
||||
SeedAudioConfig,
|
||||
SeedAudioReference,
|
||||
SeedAudioRequest,
|
||||
SeedAudioResponse,
|
||||
Seedance2TaskCreationRequest,
|
||||
SeedanceCreateAssetRequest,
|
||||
SeedanceCreateAssetResponse,
|
||||
@@ -40,9 +44,12 @@ from comfy_api_nodes.apis.bytedance import (
|
||||
TaskVideoContentUrl,
|
||||
Text2ImageTaskCreationRequest,
|
||||
Text2VideoTaskCreationRequest,
|
||||
seedance2_price_per_1k_tokens,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
ApiEndpoint,
|
||||
audio_bytes_to_audio_input,
|
||||
audio_input_to_mp3,
|
||||
download_url_to_image_tensor,
|
||||
download_url_to_video_output,
|
||||
downscale_image_tensor_by_max_side,
|
||||
@@ -51,11 +58,14 @@ from comfy_api_nodes.util import (
|
||||
image_tensor_pair_to_batch,
|
||||
poll_op,
|
||||
sync_op,
|
||||
tensor_to_base64_string,
|
||||
upload_audio_to_comfyapi,
|
||||
upload_image_to_comfyapi,
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
upscale_image_tensor_to_min_pixels,
|
||||
upscale_video_to_min_pixels,
|
||||
validate_audio_duration,
|
||||
validate_image_aspect_ratio,
|
||||
validate_image_dimensions,
|
||||
validate_string,
|
||||
@@ -89,6 +99,7 @@ BYTEPLUS_SEEDANCE2_TASK_STATUS_ENDPOINT = "/proxy/byteplus-seedance2/api/v3/cont
|
||||
SEEDANCE_MODELS = {
|
||||
"Seedance 2.0": "dreamina-seedance-2-0-260128",
|
||||
"Seedance 2.0 Fast": "dreamina-seedance-2-0-fast-260128",
|
||||
"Seedance 2.0 Mini": "dreamina-seedance-2-0-mini",
|
||||
}
|
||||
|
||||
DEPRECATED_MODELS = {"seedance-1-0-lite-t2v-250428", "seedance-1-0-lite-i2v-250428"}
|
||||
@@ -141,7 +152,7 @@ SEEDANCE2_RATIO_WH = {
|
||||
"9:16": (9, 16),
|
||||
"21:9": (21, 9),
|
||||
}
|
||||
SEEDANCE2_RES_SHORT_SIDE = {"480p": 480, "720p": 720, "1080p": 1080}
|
||||
SEEDANCE2_RES_SHORT_SIDE = {"480p": 480, "720p": 720, "1080p": 1080, "4k": 2160}
|
||||
|
||||
|
||||
def _seedance2_target_dims(resolution: str, ratio: str, image: torch.Tensor) -> tuple[int, int]:
|
||||
@@ -377,9 +388,9 @@ async def _seedance_virtual_library_upload_video_asset(
|
||||
return f"asset://{create_resp.asset_id}"
|
||||
|
||||
|
||||
def _seedance2_price_extractor(model_id: str, has_video_input: bool):
|
||||
def _seedance2_price_extractor(model_id: str, has_video_input: bool, resolution: str):
|
||||
"""Returns a price_extractor closure for Seedance 2.0 poll_op."""
|
||||
rate = SEEDANCE2_PRICE_PER_1K_TOKENS.get((model_id, has_video_input))
|
||||
rate = seedance2_price_per_1k_tokens(model_id, has_video_input, resolution)
|
||||
if rate is None:
|
||||
return None
|
||||
|
||||
@@ -1621,10 +1632,12 @@ class ByteDance2TextToVideoNode(IO.ComfyNode):
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("Seedance 2.0", _seedance2_text_inputs(["480p", "720p", "1080p"])),
|
||||
IO.DynamicCombo.Option("Seedance 2.0", _seedance2_text_inputs(["480p", "720p", "1080p", "4k"])),
|
||||
IO.DynamicCombo.Option("Seedance 2.0 Fast", _seedance2_text_inputs(["480p", "720p"])),
|
||||
IO.DynamicCombo.Option("Seedance 2.0 Mini", _seedance2_text_inputs(["480p", "720p"])),
|
||||
],
|
||||
tooltip="Seedance 2.0 for maximum quality; Seedance 2.0 Fast for speed optimization.",
|
||||
tooltip="Seedance 2.0 for maximum quality; Fast for speed optimization; "
|
||||
"Mini for the fastest, lowest-cost generation.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
@@ -1660,11 +1673,16 @@ class ByteDance2TextToVideoNode(IO.ComfyNode):
|
||||
$rate480 := 10044;
|
||||
$rate720 := 21600;
|
||||
$rate1080 := 48800;
|
||||
$rate4k := 195200;
|
||||
$m := widgets.model;
|
||||
$pricePer1K := $contains($m, "fast") ? 0.008008 : 0.01001;
|
||||
$res := $lookup(widgets, "model.resolution");
|
||||
$dur := $lookup(widgets, "model.duration");
|
||||
$rate := $res = "1080p" ? $rate1080 :
|
||||
$pricePer1K := $res = "4k" ? 0.00572 :
|
||||
$res = "1080p" ? 0.011011 :
|
||||
$contains($m, "mini") ? 0.005005 :
|
||||
$contains($m, "fast") ? 0.008008 : 0.01001;
|
||||
$rate := $res = "4k" ? $rate4k :
|
||||
$res = "1080p" ? $rate1080 :
|
||||
$res = "720p" ? $rate720 :
|
||||
$rate480;
|
||||
$cost := $dur * $rate * $pricePer1K / 1000;
|
||||
@@ -1703,7 +1721,7 @@ class ByteDance2TextToVideoNode(IO.ComfyNode):
|
||||
ApiEndpoint(path=f"{BYTEPLUS_SEEDANCE2_TASK_STATUS_ENDPOINT}/{initial_response.id}"),
|
||||
response_model=TaskStatusResponse,
|
||||
status_extractor=lambda r: r.status,
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=False),
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=False, resolution=model["resolution"]),
|
||||
poll_interval=9,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
@@ -1724,14 +1742,19 @@ class ByteDance2FirstLastFrameNode(IO.ComfyNode):
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Seedance 2.0",
|
||||
_seedance2_text_inputs(["480p", "720p", "1080p"], default_ratio="adaptive"),
|
||||
_seedance2_text_inputs(["480p", "720p", "1080p", "4k"], default_ratio="adaptive"),
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Seedance 2.0 Fast",
|
||||
_seedance2_text_inputs(["480p", "720p"], default_ratio="adaptive"),
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Seedance 2.0 Mini",
|
||||
_seedance2_text_inputs(["480p", "720p"], default_ratio="adaptive"),
|
||||
),
|
||||
],
|
||||
tooltip="Seedance 2.0 for maximum quality; Seedance 2.0 Fast for speed optimization.",
|
||||
tooltip="Seedance 2.0 for maximum quality; Fast for speed optimization; "
|
||||
"Mini for the fastest, lowest-cost generation.",
|
||||
),
|
||||
IO.Image.Input(
|
||||
"first_frame",
|
||||
@@ -1791,11 +1814,16 @@ class ByteDance2FirstLastFrameNode(IO.ComfyNode):
|
||||
$rate480 := 10044;
|
||||
$rate720 := 21600;
|
||||
$rate1080 := 48800;
|
||||
$rate4k := 195200;
|
||||
$m := widgets.model;
|
||||
$pricePer1K := $contains($m, "fast") ? 0.008008 : 0.01001;
|
||||
$res := $lookup(widgets, "model.resolution");
|
||||
$dur := $lookup(widgets, "model.duration");
|
||||
$rate := $res = "1080p" ? $rate1080 :
|
||||
$pricePer1K := $res = "4k" ? 0.00572 :
|
||||
$res = "1080p" ? 0.011011 :
|
||||
$contains($m, "mini") ? 0.005005 :
|
||||
$contains($m, "fast") ? 0.008008 : 0.01001;
|
||||
$rate := $res = "4k" ? $rate4k :
|
||||
$res = "1080p" ? $rate1080 :
|
||||
$res = "720p" ? $rate720 :
|
||||
$rate480;
|
||||
$cost := $dur * $rate * $pricePer1K / 1000;
|
||||
@@ -1913,7 +1941,7 @@ class ByteDance2FirstLastFrameNode(IO.ComfyNode):
|
||||
ApiEndpoint(path=f"{BYTEPLUS_SEEDANCE2_TASK_STATUS_ENDPOINT}/{initial_response.id}"),
|
||||
response_model=TaskStatusResponse,
|
||||
status_extractor=lambda r: r.status,
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=False),
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=False, resolution=model["resolution"]),
|
||||
poll_interval=9,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
@@ -2010,14 +2038,19 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Seedance 2.0",
|
||||
_seedance2_reference_inputs(["480p", "720p", "1080p"], default_ratio="adaptive"),
|
||||
_seedance2_reference_inputs(["480p", "720p", "1080p", "4k"], default_ratio="adaptive"),
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Seedance 2.0 Fast",
|
||||
_seedance2_reference_inputs(["480p", "720p"], default_ratio="adaptive"),
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Seedance 2.0 Mini",
|
||||
_seedance2_reference_inputs(["480p", "720p"], default_ratio="adaptive"),
|
||||
),
|
||||
],
|
||||
tooltip="Seedance 2.0 for maximum quality; Seedance 2.0 Fast for speed optimization.",
|
||||
tooltip="Seedance 2.0 for maximum quality; Fast for speed optimization; "
|
||||
"Mini for the fastest, lowest-cost generation.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
@@ -2056,13 +2089,21 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
|
||||
$rate480 := 10044;
|
||||
$rate720 := 21600;
|
||||
$rate1080 := 48800;
|
||||
$rate4k := 195200;
|
||||
$m := widgets.model;
|
||||
$hasVideo := $lookup(inputGroups, "model.reference_videos") > 0;
|
||||
$noVideoPricePer1K := $contains($m, "fast") ? 0.008008 : 0.01001;
|
||||
$videoPricePer1K := $contains($m, "fast") ? 0.004719 : 0.006149;
|
||||
$res := $lookup(widgets, "model.resolution");
|
||||
$dur := $lookup(widgets, "model.duration");
|
||||
$rate := $res = "1080p" ? $rate1080 :
|
||||
$noVideoPricePer1K := $res = "4k" ? 0.00572 :
|
||||
$res = "1080p" ? 0.011011 :
|
||||
$contains($m, "mini") ? 0.005005 :
|
||||
$contains($m, "fast") ? 0.008008 : 0.01001;
|
||||
$videoPricePer1K := $res = "4k" ? 0.003432 :
|
||||
$res = "1080p" ? 0.006721 :
|
||||
$contains($m, "mini") ? 0.003003 :
|
||||
$contains($m, "fast") ? 0.004719 : 0.006149;
|
||||
$rate := $res = "4k" ? $rate4k :
|
||||
$res = "1080p" ? $rate1080 :
|
||||
$res = "720p" ? $rate720 :
|
||||
$rate480;
|
||||
$noVideoCost := $dur * $rate * $noVideoPricePer1K / 1000;
|
||||
@@ -2258,7 +2299,9 @@ class ByteDance2ReferenceNode(IO.ComfyNode):
|
||||
ApiEndpoint(path=f"{BYTEPLUS_SEEDANCE2_TASK_STATUS_ENDPOINT}/{initial_response.id}"),
|
||||
response_model=TaskStatusResponse,
|
||||
status_extractor=lambda r: r.status,
|
||||
price_extractor=_seedance2_price_extractor(model_id, has_video_input=has_video_input),
|
||||
price_extractor=_seedance2_price_extractor(
|
||||
model_id, has_video_input=has_video_input, resolution=model["resolution"]
|
||||
),
|
||||
poll_interval=9,
|
||||
)
|
||||
return IO.NodeOutput(await download_url_to_video_output(response.content.video_url))
|
||||
@@ -2441,6 +2484,311 @@ class ByteDanceCreateVideoAsset(IO.ComfyNode):
|
||||
return IO.NodeOutput(asset_id, resolved_group)
|
||||
|
||||
|
||||
MODE_TEXT = "text only"
|
||||
MODE_AUDIO = "audio reference"
|
||||
MODE_IMAGE = "image reference"
|
||||
MODE_SPEAKER = "preset voice"
|
||||
|
||||
# (speaker_id, display_label) for built-in TTS 2.0 voices; resolvable ids are account-scoped.
|
||||
SEED_AUDIO_PRESET_VOICES: list[tuple[str, str]] = [
|
||||
("zh_female_vv_uranus_bigtts", "Vivi (Female, multilingual)"),
|
||||
("zh_female_xiaohe_uranus_bigtts", "Mindy (Female, multilingual)"),
|
||||
("en_female_stokie_uranus_bigtts", "Stokie (Female, English)"),
|
||||
("en_female_dacey_uranus_bigtts", "Dacey (Female, English)"),
|
||||
("en_male_tim_uranus_bigtts", "Tim (Male, English)"),
|
||||
("zh_male_m191_uranus_bigtts", "Kian (Male, multilingual)"),
|
||||
("zh_male_taocheng_uranus_bigtts", "Cedric (Male, multilingual)"),
|
||||
("zh_male_sophie_uranus_bigtts", "Sophie (Female, multilingual)"),
|
||||
("zh_female_yingyujiaoxue_uranus_bigtts", "Jean (Female, multilingual)"),
|
||||
("zh_male_dayi_uranus_bigtts", "Magnus (Male, multilingual)"),
|
||||
("zh_female_mizai_uranus_bigtts", "Mabel (Female, multilingual)"),
|
||||
("zh_female_jitangnv_uranus_bigtts", "Nadia (Female, multilingual)"),
|
||||
("zh_female_meilinvyou_uranus_bigtts", "Opal (Female, multilingual)"),
|
||||
("zh_female_liuchangnv_uranus_bigtts", "Pearl (Female, multilingual)"),
|
||||
("zh_male_ruyayichen_uranus_bigtts", "Quentin (Male, multilingual)"),
|
||||
("zh_female_vivo_uranus_bigtts", "Vienna (Female, multilingual)"),
|
||||
("zh_female_xiaoai_uranus_bigtts", "Alina (Female, multilingual)"),
|
||||
("zh_female_cancan_uranus_bigtts", "Corinne (Female, multilingual)"),
|
||||
("zh_female_tianmeixiaoyuan_uranus_bigtts", "Esther (Female, multilingual)"),
|
||||
("zh_female_tianmeitaozi_uranus_bigtts", "Freya (Female, multilingual)"),
|
||||
("zh_female_shuangkuaisisi_uranus_bigtts", "Gigi (Female, multilingual)"),
|
||||
("zh_female_peiqi_uranus_bigtts", "Holly (Female, multilingual)"),
|
||||
("zh_female_xiaoxue_uranus_bigtts", "Lyla (Female, multilingual)"),
|
||||
("zh_female_yuanqi_uranus_bigtts", "Daisy (Female, multilingual)"),
|
||||
("zh_female_kefunvsheng_uranus_bigtts", "Tracy (Female, multilingual)"),
|
||||
("zh_male_shaonianzixin_uranus_bigtts", "Jess (Male, multilingual)"),
|
||||
("zh_female_linjianvhai_uranus_bigtts", "Pinky (Female, multilingual)"),
|
||||
("zh_female_kiwi_uranus_bigtts", "Sweety (Female, multilingual)"),
|
||||
("zh_female_sajiaoxuemei_uranus_bigtts", "Sandy (Female, multilingual)"),
|
||||
("de_male_seven_uranus_bigtts", "Sven (Male, German)"),
|
||||
("jp_female_minimi_uranus_bigtts", "Minimi (Female, Japanese)"),
|
||||
("fr_male_usseau_uranus_bigtts", "Usseau (Male, French)"),
|
||||
("es_male_felipe_uranus_bigtts", "Felipe (Male, Spanish)"),
|
||||
("id_male_han_uranus_bigtts", "Han (Male, Indonesian)"),
|
||||
("pt_male_martins_uranus_bigtts", "Martins (Male, Portuguese)"),
|
||||
("it_male_enzo_uranus_bigtts", "Enzo (Male, Italian)"),
|
||||
("kr_male_shane_uranus_bigtts", "Shane (Male, Korean)"),
|
||||
("zh_male_liufei_uranus_bigtts", "Felix (Male, Chinese)"),
|
||||
("zh_female_qingxinnvsheng_uranus_bigtts", "Celeste (Female, Chinese)"),
|
||||
("zh_male_sunwukong_uranus_bigtts", "Monkey King (Male, Chinese)"),
|
||||
]
|
||||
SEED_AUDIO_VOICE_OPTIONS = [label for _, label in SEED_AUDIO_PRESET_VOICES]
|
||||
SEED_AUDIO_VOICE_MAP = {label: speaker_id for speaker_id, label in SEED_AUDIO_PRESET_VOICES}
|
||||
|
||||
_AUDIO_TAG_RE = re.compile(r"@Audio(\d+)", re.IGNORECASE)
|
||||
|
||||
|
||||
def max_audio_tag(prompt: str) -> int:
|
||||
"""Highest N referenced as @AudioN in the prompt (0 if none)."""
|
||||
nums = [int(m) for m in _AUDIO_TAG_RE.findall(prompt or "")]
|
||||
return max(nums) if nums else 0
|
||||
|
||||
|
||||
def connected_audio_indices(reference_mode: dict) -> list[int]:
|
||||
"""Indices (1-based) of connected reference_audio sockets, in order."""
|
||||
return [
|
||||
i
|
||||
for i in range(1, 3 + 1)
|
||||
if reference_mode.get(f"reference_audio_{i}") is not None
|
||||
]
|
||||
|
||||
|
||||
def validate_seed_audio_inputs(
|
||||
text_prompt: str,
|
||||
mode: str,
|
||||
audio_indices: list[int],
|
||||
has_image: bool,
|
||||
preset_voice: str | None = None,
|
||||
) -> None:
|
||||
validate_string(text_prompt, field_name="text_prompt", min_length=1, max_length=3000)
|
||||
max_tag = max_audio_tag(text_prompt)
|
||||
|
||||
if mode == MODE_TEXT:
|
||||
if max_tag:
|
||||
raise ValueError(
|
||||
f"The prompt references @Audio{max_tag}, but reference mode is '{MODE_TEXT}'. "
|
||||
f"Switch to '{MODE_AUDIO}' and connect the reference clip(s)."
|
||||
)
|
||||
elif mode == MODE_AUDIO:
|
||||
if not audio_indices:
|
||||
raise ValueError(
|
||||
f"Reference mode '{MODE_AUDIO}' requires at least one reference_audio input "
|
||||
f"(or switch to '{MODE_TEXT}')."
|
||||
)
|
||||
if audio_indices != list(range(1, len(audio_indices) + 1)):
|
||||
raise ValueError(
|
||||
"Connect reference_audio inputs in order without gaps: reference_audio_1, then _2, then _3."
|
||||
)
|
||||
if max_tag > len(audio_indices):
|
||||
raise ValueError(
|
||||
f"The prompt references @Audio{max_tag}, but only {len(audio_indices)} "
|
||||
f"reference audio(s) are connected."
|
||||
)
|
||||
elif mode == MODE_IMAGE:
|
||||
if not has_image:
|
||||
raise ValueError(f"Reference mode '{MODE_IMAGE}' requires a reference_image input.")
|
||||
if max_tag:
|
||||
raise ValueError(
|
||||
f"@AudioN tags are not used in '{MODE_IMAGE}' mode; the prompt should contain "
|
||||
f"only the text to synthesize."
|
||||
)
|
||||
elif mode == MODE_SPEAKER:
|
||||
if not preset_voice or preset_voice not in SEED_AUDIO_VOICE_MAP:
|
||||
raise ValueError(f"Reference mode '{MODE_SPEAKER}' requires selecting a preset voice.")
|
||||
if max_tag > 1:
|
||||
raise ValueError(
|
||||
f"'{MODE_SPEAKER}' mode uses a single voice, so @Audio{max_tag} is out of range. "
|
||||
f"Remove the @AudioN tags — the whole prompt is read in the selected voice."
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown reference mode: {mode!r}")
|
||||
|
||||
|
||||
class ByteDanceSeedAudioNode(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="ByteDanceSeedAudio",
|
||||
display_name="ByteDance Seed Audio 1.0",
|
||||
category="partner/audio/ByteDance",
|
||||
description=(
|
||||
"Generate speech, music, sound effects and multi-speaker dialogue from a single prompt "
|
||||
"with ByteDance Seed Audio 1.0. Describe the voice(s), emotion, ambience, background music "
|
||||
"and sound effects in the prompt, and include the lines to speak. Optionally pick a built-in "
|
||||
"preset voice, clone voices from up to 3 reference clips (tagged @Audio1-3 in the prompt), "
|
||||
"or derive a voice from a character image. Up to 2 minutes of audio per run."
|
||||
),
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"text_prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip=(
|
||||
"Describe the voice(s), emotion, pacing, ambience, background music and sound "
|
||||
"effects, and include the lines to speak (name characters inline for dialogue). "
|
||||
"In 'audio reference' mode, refer to connected clips by order as @Audio1, @Audio2, "
|
||||
"@Audio3. Maximum 3000 characters."
|
||||
),
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"reference_mode",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(MODE_TEXT, []),
|
||||
IO.DynamicCombo.Option(
|
||||
MODE_AUDIO,
|
||||
[
|
||||
IO.Audio.Input(
|
||||
"reference_audio_1",
|
||||
optional=True,
|
||||
tooltip="Reference clip for voice cloning, tagged @Audio1 in the prompt. "
|
||||
"Up to 30s.",
|
||||
),
|
||||
IO.Audio.Input(
|
||||
"reference_audio_2",
|
||||
optional=True,
|
||||
tooltip="Reference clip tagged @Audio2 in the prompt. Up to 30s.",
|
||||
),
|
||||
IO.Audio.Input(
|
||||
"reference_audio_3",
|
||||
optional=True,
|
||||
tooltip="Reference clip tagged @Audio3 in the prompt. Up to 30s.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
MODE_IMAGE,
|
||||
[
|
||||
IO.Image.Input(
|
||||
"reference_image",
|
||||
optional=True,
|
||||
tooltip="A single character image; the model derives a voice from it. "
|
||||
"Cannot be combined with reference audio.",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
MODE_SPEAKER,
|
||||
[
|
||||
IO.Combo.Input(
|
||||
"preset_voice",
|
||||
options=SEED_AUDIO_VOICE_OPTIONS,
|
||||
default=SEED_AUDIO_VOICE_OPTIONS[0],
|
||||
tooltip="A built-in TTS 2.0 voice that reads the prompt. No reference "
|
||||
"clip needed, and @AudioN tags are not used in this mode.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip=(
|
||||
"How to condition the voice: 'text only' (describe everything in the prompt), "
|
||||
"'audio reference' (clone up to 3 voices, tagged @Audio1-3), 'image reference' "
|
||||
"(derive a voice from one character image), or 'preset voice' (pick a built-in "
|
||||
"named voice that reads the prompt)."
|
||||
),
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"sample_rate",
|
||||
options=["8000", "16000", "24000", "32000", "44100", "48000"],
|
||||
default="24000",
|
||||
tooltip="Output sample rate in Hz.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"speech_rate",
|
||||
default=0,
|
||||
min=-50,
|
||||
max=100,
|
||||
tooltip="Speaking speed. 0 = normal, 100 = 2.0x, -50 = 0.5x.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"loudness_rate",
|
||||
default=0,
|
||||
min=-50,
|
||||
max=100,
|
||||
tooltip="Loudness. 0 = normal, 100 = 2.0x, -50 = 0.5x.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"pitch_rate",
|
||||
default=0,
|
||||
min=-12,
|
||||
max=12,
|
||||
tooltip="Pitch shift in semitones (-12 to 12).",
|
||||
),
|
||||
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.",
|
||||
),
|
||||
],
|
||||
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.2145, "format":{"suffix":"/minute","approximate":true}}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
text_prompt: str,
|
||||
reference_mode: dict,
|
||||
sample_rate: str,
|
||||
speech_rate: int,
|
||||
loudness_rate: int,
|
||||
pitch_rate: int,
|
||||
seed: int,
|
||||
) -> IO.NodeOutput:
|
||||
mode = reference_mode["reference_mode"]
|
||||
audio_indices = connected_audio_indices(reference_mode)
|
||||
image = reference_mode.get("reference_image")
|
||||
preset_voice = reference_mode.get("preset_voice")
|
||||
validate_seed_audio_inputs(text_prompt, mode, audio_indices, image is not None, preset_voice)
|
||||
|
||||
references: list[SeedAudioReference] | None = None
|
||||
if mode == MODE_AUDIO:
|
||||
references = []
|
||||
for i in audio_indices:
|
||||
clip = reference_mode[f"reference_audio_{i}"]
|
||||
validate_audio_duration(clip, max_duration=30.0)
|
||||
mp3_bytes = audio_input_to_mp3(clip).getvalue()
|
||||
references.append(SeedAudioReference(audio_data=base64.b64encode(mp3_bytes).decode("utf-8")))
|
||||
elif mode == MODE_IMAGE:
|
||||
image = upscale_image_tensor_to_min_pixels(image, 160_000)
|
||||
references = [SeedAudioReference(image_data=tensor_to_base64_string(image, mime_type="image/png"))]
|
||||
elif mode == MODE_SPEAKER:
|
||||
references = [SeedAudioReference(speaker=SEED_AUDIO_VOICE_MAP[preset_voice])]
|
||||
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/byteplus/api/v3/tts/create", method="POST"),
|
||||
response_model=SeedAudioResponse,
|
||||
data=SeedAudioRequest(
|
||||
text_prompt=text_prompt,
|
||||
references=references,
|
||||
audio_config=SeedAudioConfig(
|
||||
sample_rate=int(sample_rate),
|
||||
speech_rate=speech_rate,
|
||||
loudness_rate=loudness_rate,
|
||||
pitch_rate=pitch_rate,
|
||||
),
|
||||
),
|
||||
)
|
||||
if not response.audio:
|
||||
raise Exception(
|
||||
f"Seed Audio returned no audio (code={response.code}): {response.message}"
|
||||
)
|
||||
return IO.NodeOutput(audio_bytes_to_audio_input(base64.b64decode(response.audio)))
|
||||
|
||||
|
||||
class ByteDanceExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
@@ -2457,6 +2805,7 @@ class ByteDanceExtension(ComfyExtension):
|
||||
ByteDance2ReferenceNode,
|
||||
ByteDanceCreateImageAsset,
|
||||
ByteDanceCreateVideoAsset,
|
||||
ByteDanceSeedAudioNode,
|
||||
]
|
||||
|
||||
|
||||
|
||||
+201
-34
@@ -5,7 +5,6 @@ See: https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/infer
|
||||
|
||||
import base64
|
||||
import os
|
||||
from enum import Enum
|
||||
from fnmatch import fnmatch
|
||||
from io import BytesIO
|
||||
from typing import Any, Literal
|
||||
@@ -14,7 +13,7 @@ import torch
|
||||
from typing_extensions import override
|
||||
|
||||
import folder_paths
|
||||
from comfy_api.latest import IO, ComfyExtension, Input, Types
|
||||
from comfy_api.latest import IO, ComfyExtension, Input, InputImpl, Types
|
||||
from comfy_api_nodes.apis.gemini import (
|
||||
GeminiContent,
|
||||
GeminiFileData,
|
||||
@@ -38,6 +37,7 @@ from comfy_api_nodes.util import (
|
||||
audio_to_base64_string,
|
||||
bytesio_to_image_tensor,
|
||||
download_url_to_image_tensor,
|
||||
download_url_to_video_output,
|
||||
get_number_of_images,
|
||||
sync_op,
|
||||
tensor_to_base64_string,
|
||||
@@ -46,6 +46,7 @@ from comfy_api_nodes.util import (
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_string,
|
||||
validate_video_duration,
|
||||
video_to_base64_string,
|
||||
)
|
||||
|
||||
@@ -78,15 +79,6 @@ GEMINI_IMAGE_2_PRICE_BADGE = IO.PriceBadge(
|
||||
)
|
||||
|
||||
|
||||
class GeminiImageModel(str, Enum):
|
||||
"""
|
||||
Gemini Image Model Names allowed by comfy-api
|
||||
"""
|
||||
|
||||
gemini_2_5_flash_image_preview = "gemini-2.5-flash-image-preview"
|
||||
gemini_2_5_flash_image = "gemini-2.5-flash-image"
|
||||
|
||||
|
||||
async def create_image_parts(
|
||||
cls: type[IO.ComfyNode],
|
||||
images: Input.Image | list[Input.Image],
|
||||
@@ -239,25 +231,38 @@ 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
|
||||
if response.modelVersion in ("gemini-2.5-pro-preview-05-06", "gemini-2.5-pro"):
|
||||
output_video_tokens_price = 0.0
|
||||
if response.modelVersion == "gemini-2.5-pro":
|
||||
input_tokens_price = 1.25
|
||||
output_text_tokens_price = 10.0
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion in (
|
||||
"gemini-2.5-flash-preview-04-17",
|
||||
"gemini-2.5-flash",
|
||||
):
|
||||
elif response.modelVersion == "gemini-2.5-flash":
|
||||
input_tokens_price = 0.30
|
||||
output_text_tokens_price = 2.50
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion in (
|
||||
"gemini-2.5-flash-image-preview",
|
||||
"gemini-2.5-flash-image",
|
||||
):
|
||||
elif response.modelVersion == "gemini-2.5-flash-image":
|
||||
input_tokens_price = 0.30
|
||||
output_text_tokens_price = 2.50
|
||||
output_image_tokens_price = 30.0
|
||||
@@ -265,18 +270,27 @@ def calculate_tokens_price(response: GeminiGenerateContentResponse) -> float | N
|
||||
input_tokens_price = 2
|
||||
output_text_tokens_price = 12.0
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion == "gemini-3.1-flash-lite-preview":
|
||||
elif response.modelVersion in ("gemini-3.1-flash-lite-preview", "gemini-3.1-flash-lite"):
|
||||
input_tokens_price = 0.25
|
||||
output_text_tokens_price = 1.50
|
||||
output_image_tokens_price = 0.0
|
||||
elif response.modelVersion == "gemini-3-pro-image-preview":
|
||||
elif response.modelVersion in ("gemini-3-pro-image-preview", "gemini-3-pro-image"):
|
||||
input_tokens_price = 2
|
||||
output_text_tokens_price = 12.0
|
||||
output_image_tokens_price = 120.0
|
||||
elif response.modelVersion == "gemini-3.1-flash-image-preview":
|
||||
elif response.modelVersion in ("gemini-3.1-flash-image-preview", "gemini-3.1-flash-image"):
|
||||
input_tokens_price = 0.5
|
||||
output_text_tokens_price = 3.0
|
||||
output_image_tokens_price = 60.0
|
||||
elif response.modelVersion == "gemini-3.1-flash-lite-image":
|
||||
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
|
||||
@@ -284,6 +298,8 @@ 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:
|
||||
@@ -455,8 +471,6 @@ class GeminiNode(IO.ComfyNode):
|
||||
IO.Combo.Input(
|
||||
"model",
|
||||
options=[
|
||||
"gemini-2.5-pro-preview-05-06",
|
||||
"gemini-2.5-flash-preview-04-17",
|
||||
"gemini-2.5-pro",
|
||||
"gemini-2.5-flash",
|
||||
"gemini-3-pro-preview",
|
||||
@@ -904,8 +918,7 @@ class GeminiImage(IO.ComfyNode):
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"model",
|
||||
options=GeminiImageModel,
|
||||
default=GeminiImageModel.gemini_2_5_flash_image,
|
||||
options=["gemini-2.5-flash-image"],
|
||||
tooltip="The Gemini model to use for generating responses.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
@@ -1321,7 +1334,7 @@ class GeminiNanoBanana2(IO.ComfyNode):
|
||||
)
|
||||
|
||||
|
||||
def _nano_banana_2_v2_model_inputs():
|
||||
def _nano_banana_2_v2_model_inputs(resolutions: list[str]):
|
||||
return [
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
@@ -1348,8 +1361,8 @@ def _nano_banana_2_v2_model_inputs():
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["1K", "2K", "4K"],
|
||||
tooltip="Target output resolution. For 2K/4K the native Gemini upscaler is used.",
|
||||
options=resolutions,
|
||||
tooltip="Target output resolution.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"thinking_level",
|
||||
@@ -1395,7 +1408,11 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Nano Banana 2 (Gemini 3.1 Flash Image)",
|
||||
_nano_banana_2_v2_model_inputs(),
|
||||
_nano_banana_2_v2_model_inputs(resolutions=["1K", "2K", "4K"]),
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Nano Banana 2 Lite",
|
||||
_nano_banana_2_v2_model_inputs(resolutions=["1K"]),
|
||||
),
|
||||
],
|
||||
),
|
||||
@@ -1464,9 +1481,13 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["model", "model.resolution"]),
|
||||
expr="""
|
||||
(
|
||||
$r := $lookup(widgets, "model.resolution");
|
||||
$prices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
|
||||
{"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
|
||||
$contains(widgets.model, "lite")
|
||||
? {"type":"usd","usd": 0.034, "format":{"suffix":"/Image","approximate":true}}
|
||||
: (
|
||||
$r := $lookup(widgets, "model.resolution");
|
||||
$prices := {"1k": 0.0696, "2k": 0.1014, "4k": 0.154};
|
||||
{"type":"usd","usd": $lookup($prices, $r), "format":{"suffix":"/Image","approximate":true}}
|
||||
)
|
||||
)
|
||||
""",
|
||||
),
|
||||
@@ -1487,6 +1508,8 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
||||
model_choice = model["model"]
|
||||
if model_choice == "Nano Banana 2 (Gemini 3.1 Flash Image)":
|
||||
model_id = "gemini-3.1-flash-image-preview"
|
||||
elif model_choice == "Nano Banana 2 Lite":
|
||||
model_id = "gemini-3.1-flash-lite-image"
|
||||
else:
|
||||
model_id = model_choice
|
||||
|
||||
@@ -1536,6 +1559,149 @@ class GeminiNanoBanana2V2(IO.ComfyNode):
|
||||
)
|
||||
|
||||
|
||||
OMNI_MAX_IMAGES = 14
|
||||
OMNI_MAX_VIDEOS = 3
|
||||
|
||||
OMNI_MODELS: dict[str, str] = {
|
||||
"Omni Flash": "gemini-omni-flash-preview",
|
||||
}
|
||||
|
||||
|
||||
def _omni_flash_inputs() -> list[Input]:
|
||||
"""Per-model inputs for the Omni video DynamicCombo (prompt + reference media + sampling)."""
|
||||
return [
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Describe the video to generate. Specify the length and aspect ratio directly in the "
|
||||
'prompt, e.g. "a 6-second clip in 16:9". Length may be 3-10 seconds; the aspect ratio must be '
|
||||
"16:9 (landscape) or 9:16 (portrait). The output is 720p, 24 FPS, with audio.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"images",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Image.Input("image"),
|
||||
names=[f"image_{i}" for i in range(1, OMNI_MAX_IMAGES + 1)],
|
||||
min=0,
|
||||
),
|
||||
tooltip=f"Optional reference image(s) to guide or animate the video. Up to {OMNI_MAX_IMAGES} images.",
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"videos",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Video.Input("video"),
|
||||
names=[f"video_{i}" for i in range(1, OMNI_MAX_VIDEOS + 1)],
|
||||
min=0,
|
||||
),
|
||||
tooltip=f"Optional reference video(s) to guide or edit. Up to {OMNI_MAX_VIDEOS} videos, "
|
||||
f"each up to 10 seconds long.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"temperature",
|
||||
default=1.0,
|
||||
min=0.0,
|
||||
max=2.0,
|
||||
step=0.01,
|
||||
tooltip="Controls randomness. Lower is more focused/deterministic, higher is more varied.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"top_p",
|
||||
default=0.95,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip="Nucleus sampling: sample from the smallest token set whose cumulative probability reaches top_p.",
|
||||
advanced=True,
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
class GeminiVideoOmni(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="GeminiVideoOmni",
|
||||
display_name="Google Gemini Omni (Video)",
|
||||
category="partner/video/Gemini",
|
||||
essentials_category="Video Generation",
|
||||
description="Generate a video with audio from a text prompt using Google's Gemini Omni Flash model. "
|
||||
"Optionally provide reference images and/or videos to guide or edit the result. Describe the desired "
|
||||
"length (3-10s) and aspect ratio (16:9 or 9:16) directly in the prompt.",
|
||||
inputs=[
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("Omni Flash", _omni_flash_inputs()),
|
||||
],
|
||||
tooltip="The Gemini video model used to generate the video.",
|
||||
),
|
||||
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.",
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr='{"type":"usd","usd":0.146,"format":{"suffix":"/second","approximate":true}}'
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, model: dict, seed: int) -> IO.NodeOutput:
|
||||
prompt = model.get("prompt") or ""
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1)
|
||||
model_id = OMNI_MODELS[model["model"]]
|
||||
|
||||
images = [t for t in (model.get("images") or {}).values() if t is not None]
|
||||
videos = [v for v in (model.get("videos") or {}).values() if v is not None]
|
||||
if sum(get_number_of_images(t) for t in images) > OMNI_MAX_IMAGES:
|
||||
raise ValueError(f"The current maximum number of supported images is {OMNI_MAX_IMAGES}.")
|
||||
if len(videos) > OMNI_MAX_VIDEOS:
|
||||
raise ValueError(f"The current maximum number of supported videos is {OMNI_MAX_VIDEOS}.")
|
||||
for video in videos:
|
||||
validate_video_duration(video, max_duration=10)
|
||||
|
||||
parts: list[GeminiPart] = []
|
||||
if images or videos:
|
||||
parts.extend(await build_gemini_media_parts(cls, images, [], videos))
|
||||
parts.append(GeminiPart(text=prompt))
|
||||
response = 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"],
|
||||
temperature=model.get("temperature", 1.0),
|
||||
topP=model.get("top_p", 0.95),
|
||||
),
|
||||
),
|
||||
response_model=GeminiGenerateContentResponse,
|
||||
price_extractor=calculate_tokens_price,
|
||||
)
|
||||
return IO.NodeOutput(
|
||||
await get_video_from_response(response, cls=cls),
|
||||
get_text_from_response(response),
|
||||
)
|
||||
|
||||
|
||||
class GeminiExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
@@ -1546,6 +1712,7 @@ class GeminiExtension(ComfyExtension):
|
||||
GeminiImage2,
|
||||
GeminiNanoBanana2,
|
||||
GeminiNanoBanana2V2,
|
||||
GeminiVideoOmni,
|
||||
GeminiInputFiles,
|
||||
]
|
||||
|
||||
|
||||
@@ -30,7 +30,7 @@ from comfy_api_nodes.util import (
|
||||
|
||||
|
||||
_GROK_VIDEO_MODEL_API_IDS = {
|
||||
"grok-imagine-video-1.5": "grok-imagine-video-1.5-preview",
|
||||
"grok-imagine-video-1.5": "grok-imagine-video-1.5",
|
||||
}
|
||||
|
||||
|
||||
@@ -521,8 +521,8 @@ class GrokVideoNode(IO.ComfyNode):
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["480p", "720p"],
|
||||
tooltip="The resolution of the output video.",
|
||||
options=["480p", "720p", "1080p"],
|
||||
tooltip="The resolution of the output video. 1080p is only available for grok-imagine-video-1.5.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
@@ -570,11 +570,12 @@ class GrokVideoNode(IO.ComfyNode):
|
||||
(
|
||||
$is15 := $contains(widgets.model, "1.5");
|
||||
$rate := $is15
|
||||
? (widgets.resolution = "720p" ? 0.2002 : 0.1144)
|
||||
? (widgets.resolution = "1080p" ? 0.25 : (widgets.resolution = "720p" ? 0.14 : 0.08))
|
||||
: (widgets.resolution = "720p" ? 0.07 : 0.05);
|
||||
$imgCost := $is15 ? 0.0143 : 0.002;
|
||||
$imgCost := $is15 ? 0.01 : 0.002;
|
||||
$base := $rate * widgets.duration;
|
||||
{"type":"usd","usd": inputs.image.connected ? $base + $imgCost : $base}
|
||||
$total := inputs.image.connected ? $base + $imgCost : $base;
|
||||
{"type":"usd","usd": $is15 ? $total * 1.43 : $total}
|
||||
)
|
||||
""",
|
||||
),
|
||||
@@ -593,6 +594,8 @@ class GrokVideoNode(IO.ComfyNode):
|
||||
) -> IO.NodeOutput:
|
||||
if image is None and model == "grok-imagine-video-1.5":
|
||||
raise ValueError(f"The '{model}' model requires an input image; connect one to the 'image' input.")
|
||||
if resolution == "1080p" and model != "grok-imagine-video-1.5":
|
||||
raise ValueError(f"1080p resolution is only available for grok-imagine-video-1.5, not '{model}'.")
|
||||
image_url = None
|
||||
if image is not None:
|
||||
if get_number_of_images(image) != 1:
|
||||
|
||||
@@ -5,9 +5,7 @@ from PIL import Image
|
||||
import numpy as np
|
||||
import torch
|
||||
from comfy_api_nodes.apis.ideogram import (
|
||||
IdeogramGenerateRequest,
|
||||
IdeogramGenerateResponse,
|
||||
ImageRequest,
|
||||
IdeogramV3Request,
|
||||
IdeogramV3EditRequest,
|
||||
IdeogramV4Request,
|
||||
@@ -21,101 +19,6 @@ from comfy_api_nodes.util import (
|
||||
validate_string,
|
||||
)
|
||||
|
||||
V1_V1_RES_MAP = {
|
||||
"Auto":"AUTO",
|
||||
"512 x 1536":"RESOLUTION_512_1536",
|
||||
"576 x 1408":"RESOLUTION_576_1408",
|
||||
"576 x 1472":"RESOLUTION_576_1472",
|
||||
"576 x 1536":"RESOLUTION_576_1536",
|
||||
"640 x 1024":"RESOLUTION_640_1024",
|
||||
"640 x 1344":"RESOLUTION_640_1344",
|
||||
"640 x 1408":"RESOLUTION_640_1408",
|
||||
"640 x 1472":"RESOLUTION_640_1472",
|
||||
"640 x 1536":"RESOLUTION_640_1536",
|
||||
"704 x 1152":"RESOLUTION_704_1152",
|
||||
"704 x 1216":"RESOLUTION_704_1216",
|
||||
"704 x 1280":"RESOLUTION_704_1280",
|
||||
"704 x 1344":"RESOLUTION_704_1344",
|
||||
"704 x 1408":"RESOLUTION_704_1408",
|
||||
"704 x 1472":"RESOLUTION_704_1472",
|
||||
"720 x 1280":"RESOLUTION_720_1280",
|
||||
"736 x 1312":"RESOLUTION_736_1312",
|
||||
"768 x 1024":"RESOLUTION_768_1024",
|
||||
"768 x 1088":"RESOLUTION_768_1088",
|
||||
"768 x 1152":"RESOLUTION_768_1152",
|
||||
"768 x 1216":"RESOLUTION_768_1216",
|
||||
"768 x 1232":"RESOLUTION_768_1232",
|
||||
"768 x 1280":"RESOLUTION_768_1280",
|
||||
"768 x 1344":"RESOLUTION_768_1344",
|
||||
"832 x 960":"RESOLUTION_832_960",
|
||||
"832 x 1024":"RESOLUTION_832_1024",
|
||||
"832 x 1088":"RESOLUTION_832_1088",
|
||||
"832 x 1152":"RESOLUTION_832_1152",
|
||||
"832 x 1216":"RESOLUTION_832_1216",
|
||||
"832 x 1248":"RESOLUTION_832_1248",
|
||||
"864 x 1152":"RESOLUTION_864_1152",
|
||||
"896 x 960":"RESOLUTION_896_960",
|
||||
"896 x 1024":"RESOLUTION_896_1024",
|
||||
"896 x 1088":"RESOLUTION_896_1088",
|
||||
"896 x 1120":"RESOLUTION_896_1120",
|
||||
"896 x 1152":"RESOLUTION_896_1152",
|
||||
"960 x 832":"RESOLUTION_960_832",
|
||||
"960 x 896":"RESOLUTION_960_896",
|
||||
"960 x 1024":"RESOLUTION_960_1024",
|
||||
"960 x 1088":"RESOLUTION_960_1088",
|
||||
"1024 x 640":"RESOLUTION_1024_640",
|
||||
"1024 x 768":"RESOLUTION_1024_768",
|
||||
"1024 x 832":"RESOLUTION_1024_832",
|
||||
"1024 x 896":"RESOLUTION_1024_896",
|
||||
"1024 x 960":"RESOLUTION_1024_960",
|
||||
"1024 x 1024":"RESOLUTION_1024_1024",
|
||||
"1088 x 768":"RESOLUTION_1088_768",
|
||||
"1088 x 832":"RESOLUTION_1088_832",
|
||||
"1088 x 896":"RESOLUTION_1088_896",
|
||||
"1088 x 960":"RESOLUTION_1088_960",
|
||||
"1120 x 896":"RESOLUTION_1120_896",
|
||||
"1152 x 704":"RESOLUTION_1152_704",
|
||||
"1152 x 768":"RESOLUTION_1152_768",
|
||||
"1152 x 832":"RESOLUTION_1152_832",
|
||||
"1152 x 864":"RESOLUTION_1152_864",
|
||||
"1152 x 896":"RESOLUTION_1152_896",
|
||||
"1216 x 704":"RESOLUTION_1216_704",
|
||||
"1216 x 768":"RESOLUTION_1216_768",
|
||||
"1216 x 832":"RESOLUTION_1216_832",
|
||||
"1232 x 768":"RESOLUTION_1232_768",
|
||||
"1248 x 832":"RESOLUTION_1248_832",
|
||||
"1280 x 704":"RESOLUTION_1280_704",
|
||||
"1280 x 720":"RESOLUTION_1280_720",
|
||||
"1280 x 768":"RESOLUTION_1280_768",
|
||||
"1280 x 800":"RESOLUTION_1280_800",
|
||||
"1312 x 736":"RESOLUTION_1312_736",
|
||||
"1344 x 640":"RESOLUTION_1344_640",
|
||||
"1344 x 704":"RESOLUTION_1344_704",
|
||||
"1344 x 768":"RESOLUTION_1344_768",
|
||||
"1408 x 576":"RESOLUTION_1408_576",
|
||||
"1408 x 640":"RESOLUTION_1408_640",
|
||||
"1408 x 704":"RESOLUTION_1408_704",
|
||||
"1472 x 576":"RESOLUTION_1472_576",
|
||||
"1472 x 640":"RESOLUTION_1472_640",
|
||||
"1472 x 704":"RESOLUTION_1472_704",
|
||||
"1536 x 512":"RESOLUTION_1536_512",
|
||||
"1536 x 576":"RESOLUTION_1536_576",
|
||||
"1536 x 640":"RESOLUTION_1536_640",
|
||||
}
|
||||
|
||||
V1_V2_RATIO_MAP = {
|
||||
"1:1":"ASPECT_1_1",
|
||||
"4:3":"ASPECT_4_3",
|
||||
"3:4":"ASPECT_3_4",
|
||||
"16:9":"ASPECT_16_9",
|
||||
"9:16":"ASPECT_9_16",
|
||||
"2:1":"ASPECT_2_1",
|
||||
"1:2":"ASPECT_1_2",
|
||||
"3:2":"ASPECT_3_2",
|
||||
"2:3":"ASPECT_2_3",
|
||||
"4:5":"ASPECT_4_5",
|
||||
"5:4":"ASPECT_5_4",
|
||||
}
|
||||
|
||||
V3_RATIO_MAP = {
|
||||
"1:3":"1x3",
|
||||
@@ -229,298 +132,6 @@ async def download_and_process_images(image_urls):
|
||||
return stacked_tensors
|
||||
|
||||
|
||||
class IdeogramV1(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="IdeogramV1",
|
||||
display_name="Ideogram V1",
|
||||
category="partner/image/Ideogram",
|
||||
description="Generates images using the Ideogram V1 model.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Prompt for the image generation",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"turbo",
|
||||
default=False,
|
||||
tooltip="Whether to use turbo mode (faster generation, potentially lower quality)",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=list(V1_V2_RATIO_MAP.keys()),
|
||||
default="1:1",
|
||||
tooltip="The aspect ratio for image generation.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"magic_prompt_option",
|
||||
options=["AUTO", "ON", "OFF"],
|
||||
default="AUTO",
|
||||
tooltip="Determine if MagicPrompt should be used in generation",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
step=1,
|
||||
control_after_generate=True,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
optional=True,
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Description of what to exclude from the image",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"num_images",
|
||||
default=1,
|
||||
min=1,
|
||||
max=8,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
optional=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["num_images", "turbo"]),
|
||||
expr="""
|
||||
(
|
||||
$n := widgets.num_images;
|
||||
$base := (widgets.turbo = true) ? 0.0286 : 0.0858;
|
||||
{"type":"usd","usd": $round($base * $n, 2)}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt,
|
||||
turbo=False,
|
||||
aspect_ratio="1:1",
|
||||
magic_prompt_option="AUTO",
|
||||
seed=0,
|
||||
negative_prompt="",
|
||||
num_images=1,
|
||||
):
|
||||
# Determine the model based on turbo setting
|
||||
aspect_ratio = V1_V2_RATIO_MAP.get(aspect_ratio, None)
|
||||
model = "V_1_TURBO" if turbo else "V_1"
|
||||
|
||||
response = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/ideogram/generate", method="POST"),
|
||||
response_model=IdeogramGenerateResponse,
|
||||
data=IdeogramGenerateRequest(
|
||||
image_request=ImageRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
num_images=num_images,
|
||||
seed=seed,
|
||||
aspect_ratio=aspect_ratio if aspect_ratio != "ASPECT_1_1" else None,
|
||||
magic_prompt_option=(magic_prompt_option if magic_prompt_option != "AUTO" else None),
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
)
|
||||
),
|
||||
max_retries=1,
|
||||
)
|
||||
|
||||
if not response.data or len(response.data) == 0:
|
||||
raise Exception("No images were generated in the response")
|
||||
|
||||
image_urls = [image_data.url for image_data in response.data if image_data.url]
|
||||
if not image_urls:
|
||||
raise Exception("No image URLs were generated in the response")
|
||||
return IO.NodeOutput(await download_and_process_images(image_urls))
|
||||
|
||||
|
||||
class IdeogramV2(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="IdeogramV2",
|
||||
display_name="Ideogram V2",
|
||||
category="partner/image/Ideogram",
|
||||
description="Generates images using the Ideogram V2 model.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Prompt for the image generation",
|
||||
),
|
||||
IO.Boolean.Input(
|
||||
"turbo",
|
||||
default=False,
|
||||
tooltip="Whether to use turbo mode (faster generation, potentially lower quality)",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=list(V1_V2_RATIO_MAP.keys()),
|
||||
default="1:1",
|
||||
tooltip="The aspect ratio for image generation. Ignored if resolution is not set to AUTO.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=list(V1_V1_RES_MAP.keys()),
|
||||
default="Auto",
|
||||
tooltip="The resolution for image generation. "
|
||||
"If not set to AUTO, this overrides the aspect_ratio setting.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"magic_prompt_option",
|
||||
options=["AUTO", "ON", "OFF"],
|
||||
default="AUTO",
|
||||
tooltip="Determine if MagicPrompt should be used in generation",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=2147483647,
|
||||
step=1,
|
||||
control_after_generate=True,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
optional=True,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"style_type",
|
||||
options=["AUTO", "GENERAL", "REALISTIC", "DESIGN", "RENDER_3D", "ANIME"],
|
||||
default="NONE",
|
||||
tooltip="Style type for generation (V2 only)",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Description of what to exclude from the image",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"num_images",
|
||||
default=1,
|
||||
min=1,
|
||||
max=8,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
optional=True,
|
||||
),
|
||||
#"color_palette": (
|
||||
# IO.STRING,
|
||||
# {
|
||||
# "multiline": False,
|
||||
# "default": "",
|
||||
# "tooltip": "Color palette preset name or hex colors with weights",
|
||||
# },
|
||||
#),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["num_images", "turbo"]),
|
||||
expr="""
|
||||
(
|
||||
$n := widgets.num_images;
|
||||
$base := (widgets.turbo = true) ? 0.0715 : 0.1144;
|
||||
{"type":"usd","usd": $round($base * $n, 2)}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt,
|
||||
turbo=False,
|
||||
aspect_ratio="1:1",
|
||||
resolution="Auto",
|
||||
magic_prompt_option="AUTO",
|
||||
seed=0,
|
||||
style_type="NONE",
|
||||
negative_prompt="",
|
||||
num_images=1,
|
||||
color_palette="",
|
||||
):
|
||||
aspect_ratio = V1_V2_RATIO_MAP.get(aspect_ratio, None)
|
||||
resolution = V1_V1_RES_MAP.get(resolution, None)
|
||||
# Determine the model based on turbo setting
|
||||
model = "V_2_TURBO" if turbo else "V_2"
|
||||
|
||||
# Handle resolution vs aspect_ratio logic
|
||||
# If resolution is not AUTO, it overrides aspect_ratio
|
||||
final_resolution = None
|
||||
final_aspect_ratio = None
|
||||
|
||||
if resolution != "AUTO":
|
||||
final_resolution = resolution
|
||||
else:
|
||||
final_aspect_ratio = aspect_ratio if aspect_ratio != "ASPECT_1_1" else None
|
||||
|
||||
response = await sync_op(
|
||||
cls,
|
||||
endpoint=ApiEndpoint(path="/proxy/ideogram/generate", method="POST"),
|
||||
response_model=IdeogramGenerateResponse,
|
||||
data=IdeogramGenerateRequest(
|
||||
image_request=ImageRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
num_images=num_images,
|
||||
seed=seed,
|
||||
aspect_ratio=final_aspect_ratio,
|
||||
resolution=final_resolution,
|
||||
magic_prompt_option=(magic_prompt_option if magic_prompt_option != "AUTO" else None),
|
||||
style_type=style_type if style_type != "NONE" else None,
|
||||
negative_prompt=negative_prompt if negative_prompt else None,
|
||||
color_palette=color_palette if color_palette else None,
|
||||
)
|
||||
),
|
||||
max_retries=1,
|
||||
)
|
||||
if not response.data or len(response.data) == 0:
|
||||
raise Exception("No images were generated in the response")
|
||||
|
||||
image_urls = [image_data.url for image_data in response.data if image_data.url]
|
||||
if not image_urls:
|
||||
raise Exception("No image URLs were generated in the response")
|
||||
return IO.NodeOutput(await download_and_process_images(image_urls))
|
||||
|
||||
|
||||
class IdeogramV3(IO.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
@@ -917,8 +528,6 @@ class IdeogramExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
IdeogramV1,
|
||||
IdeogramV2,
|
||||
IdeogramV3,
|
||||
IdeogramV4,
|
||||
]
|
||||
|
||||
+554
-10
@@ -3,9 +3,13 @@ from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import IO, ComfyExtension, Input
|
||||
from comfy_api_nodes.apis.luma import (
|
||||
LUMA_KEYFRAME_MODE_FRACTION,
|
||||
LUMA_KEYFRAME_MODE_SECONDS,
|
||||
Luma2Generation,
|
||||
Luma2GenerationRequest,
|
||||
Luma2ImageRef,
|
||||
Luma2VideoEdit,
|
||||
Luma2VideoOptions,
|
||||
LumaAspectRatio,
|
||||
LumaCharacterRef,
|
||||
LumaConceptChain,
|
||||
@@ -18,6 +22,8 @@ from comfy_api_nodes.apis.luma import (
|
||||
LumaIO,
|
||||
LumaKeyframes,
|
||||
LumaModifyImageRef,
|
||||
LumaRay32KeyframeChain,
|
||||
LumaRay32KeyframeItem,
|
||||
LumaReference,
|
||||
LumaReferenceChain,
|
||||
LumaVideoModel,
|
||||
@@ -33,6 +39,7 @@ from comfy_api_nodes.util import (
|
||||
sync_op,
|
||||
upload_image_to_comfyapi,
|
||||
upload_images_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_string,
|
||||
)
|
||||
|
||||
@@ -692,7 +699,10 @@ async def _luma2_upload_image_refs(
|
||||
async def _luma2_submit_and_poll(
|
||||
cls: type[IO.ComfyNode],
|
||||
request: Luma2GenerationRequest,
|
||||
) -> Input.Image:
|
||||
*,
|
||||
estimated_duration: int | None = None,
|
||||
) -> Luma2Generation:
|
||||
"""Submit a Luma Agents generation and poll until done; returns the completed generation."""
|
||||
initial = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/luma_2/generations", method="POST"),
|
||||
@@ -700,21 +710,21 @@ async def _luma2_submit_and_poll(
|
||||
data=request,
|
||||
)
|
||||
if not initial.id:
|
||||
raise RuntimeError("Luma 2 API did not return a generation id.")
|
||||
raise RuntimeError("Luma API did not return a generation id.")
|
||||
final = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/luma_2/generations/{initial.id}", method="GET"),
|
||||
response_model=Luma2Generation,
|
||||
status_extractor=lambda r: r.state,
|
||||
progress_extractor=lambda r: None,
|
||||
estimated_duration=estimated_duration,
|
||||
)
|
||||
if not final.output:
|
||||
if not final.output or not final.output[0].url:
|
||||
msg = final.failure_reason or "no output returned"
|
||||
raise RuntimeError(f"Luma 2 generation failed: {msg}")
|
||||
url = final.output[0].url
|
||||
if not url:
|
||||
raise RuntimeError("Luma 2 generation completed without an output URL.")
|
||||
return await download_url_to_image_tensor(url)
|
||||
if final.failure_code:
|
||||
msg = f"{msg} [{final.failure_code}]"
|
||||
raise RuntimeError(f"Luma generation failed: {msg}")
|
||||
return final
|
||||
|
||||
|
||||
class LumaImageNode(IO.ComfyNode):
|
||||
@@ -843,7 +853,8 @@ class LumaImageNode(IO.ComfyNode):
|
||||
web_search=model["web_search"],
|
||||
image_ref=await _luma2_upload_image_refs(cls, model.get("image_ref"), max_count=9),
|
||||
)
|
||||
return IO.NodeOutput(await _luma2_submit_and_poll(cls, request))
|
||||
final = await _luma2_submit_and_poll(cls, request)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final.output[0].url))
|
||||
|
||||
|
||||
class LumaImageEditNode(IO.ComfyNode):
|
||||
@@ -929,7 +940,533 @@ class LumaImageEditNode(IO.ComfyNode):
|
||||
web_search=model["web_search"],
|
||||
image_ref=await _luma2_upload_image_refs(cls, model.get("image_ref"), max_count=8),
|
||||
)
|
||||
return IO.NodeOutput(await _luma2_submit_and_poll(cls, request))
|
||||
final = await _luma2_submit_and_poll(cls, request)
|
||||
return IO.NodeOutput(await download_url_to_image_tensor(final.output[0].url))
|
||||
|
||||
|
||||
_BADGE_RAY32_VIDEO = IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["resolution", "duration"]),
|
||||
expr="""
|
||||
(
|
||||
$p := {
|
||||
"360p": {"5s": 0.06, "10s": 0.18},
|
||||
"540p": {"5s": 0.15, "10s": 0.45},
|
||||
"720p": {"5s": 0.3, "10s": 0.9},
|
||||
"1080p": {"5s": 1.2, "10s": 3.6}
|
||||
};
|
||||
{"type": "usd", "usd": $lookup($lookup($p, widgets.resolution), widgets.duration)}
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
_BADGE_RAY32_VIDEO_5S = IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["resolution"]),
|
||||
expr="""
|
||||
(
|
||||
$p := {"360p": 0.06, "540p": 0.15, "720p": 0.3, "1080p": 1.2};
|
||||
{"type": "usd", "usd": $lookup($p, widgets.resolution)}
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
_BADGE_RAY32_EDIT = IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["resolution"]),
|
||||
expr="""
|
||||
(
|
||||
$p := {
|
||||
"360p": {"min": 0.54, "max": 1.08},
|
||||
"540p": {"min": 0.72, "max": 1.44},
|
||||
"720p": {"min": 1.08, "max": 2.16},
|
||||
"1080p": {"min": 2.16, "max": 4.32}
|
||||
};
|
||||
$r := $lookup($p, widgets.resolution);
|
||||
{"type": "range_usd", "min_usd": $r.min, "max_usd": $r.max, "format": {"note": "(by source length)"}}
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
_BADGE_RAY32_REFRAME = IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["resolution"]),
|
||||
expr="""
|
||||
(
|
||||
$p := {"360p": 0.03, "540p": 0.06, "720p": 0.12, "1080p": 0.36};
|
||||
{"type": "usd", "usd": $lookup($p, widgets.resolution), "format": {"suffix": "/second"}}
|
||||
)
|
||||
""",
|
||||
)
|
||||
|
||||
|
||||
def _ray32_seed_input() -> IO.Input:
|
||||
return IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=0xFFFFFFFFFFFFFFFF,
|
||||
control_after_generate=True,
|
||||
tooltip="Seed to determine if node should re-run; results are nondeterministic regardless of seed.",
|
||||
)
|
||||
|
||||
|
||||
async def _ray32_generate(cls: type[IO.ComfyNode], request: Luma2GenerationRequest) -> IO.NodeOutput:
|
||||
"""Run a ray-3.2 generation and return (video, generation_id)."""
|
||||
final = await _luma2_submit_and_poll(cls, request, estimated_duration=120)
|
||||
video = await download_url_to_video_output(final.output[0].url)
|
||||
return IO.NodeOutput(video, final.id or "")
|
||||
|
||||
|
||||
class LumaRay32TextToVideoNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32TextToVideoNode",
|
||||
display_name="Luma Ray 3.2 Text to Video",
|
||||
category="partner/video/Luma",
|
||||
description="Generate a video from a text prompt using Luma's Ray 3.2 model.",
|
||||
inputs=[
|
||||
IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the video generation."),
|
||||
IO.Combo.Input("aspect_ratio", options=["16:9", "9:16", "1:1", "4:3", "3:4", "21:9"]),
|
||||
IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
|
||||
IO.Combo.Input("duration", options=["5s", "10s"]),
|
||||
IO.Boolean.Input(
|
||||
"loop",
|
||||
default=False,
|
||||
tooltip="Make the video loop seamlessly. Only available with 5s duration.",
|
||||
),
|
||||
_ray32_seed_input(),
|
||||
],
|
||||
outputs=[IO.Video.Output(), IO.String.Output(display_name="generation_id")],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=_BADGE_RAY32_VIDEO,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls, prompt: str, aspect_ratio: str, resolution: str, duration: str, loop: bool, seed: int
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
|
||||
if loop and duration == "10s":
|
||||
raise ValueError("Looping is only available with 5s duration on Ray 3.2.")
|
||||
request = Luma2GenerationRequest(
|
||||
prompt=prompt,
|
||||
model="ray-3.2",
|
||||
type="video",
|
||||
aspect_ratio=aspect_ratio,
|
||||
video=Luma2VideoOptions(resolution=resolution, duration=duration, loop=loop or None),
|
||||
)
|
||||
return await _ray32_generate(cls, request)
|
||||
|
||||
|
||||
class LumaRay32ImageToVideoNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32ImageToVideoNode",
|
||||
display_name="Luma Ray 3.2 Image to Video",
|
||||
category="partner/video/Luma",
|
||||
description="Generate a video from a start and/or end frame using Luma's Ray 3.2 model. "
|
||||
"Image-anchored generations are always 5 seconds.",
|
||||
inputs=[
|
||||
IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the video generation."),
|
||||
IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
|
||||
IO.Boolean.Input(
|
||||
"loop",
|
||||
default=False,
|
||||
tooltip="Make the video loop seamlessly. Not available when an end_frame is set.",
|
||||
),
|
||||
_ray32_seed_input(),
|
||||
IO.Image.Input("start_frame", optional=True, tooltip="First frame of the generated video."),
|
||||
IO.Image.Input("end_frame", optional=True, tooltip="Last frame of the generated video."),
|
||||
],
|
||||
outputs=[IO.Video.Output(), IO.String.Output(display_name="generation_id")],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=_BADGE_RAY32_VIDEO_5S,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
resolution: str,
|
||||
loop: bool,
|
||||
seed: int,
|
||||
start_frame: torch.Tensor | None = None,
|
||||
end_frame: torch.Tensor | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
|
||||
if start_frame is None and end_frame is None:
|
||||
raise ValueError("Provide at least one of start_frame / end_frame.")
|
||||
if loop and end_frame is not None:
|
||||
raise ValueError("Looping is not available when an end_frame is set.")
|
||||
video = Luma2VideoOptions(resolution=resolution, duration="5s", loop=loop or None)
|
||||
if start_frame is not None:
|
||||
url = await upload_image_to_comfyapi(cls, start_frame, mime_type="image/png")
|
||||
video.start_frame = Luma2ImageRef(url=url)
|
||||
if end_frame is not None:
|
||||
url = await upload_image_to_comfyapi(cls, end_frame, mime_type="image/png")
|
||||
video.end_frame = Luma2ImageRef(url=url)
|
||||
request = Luma2GenerationRequest(prompt=prompt, model="ray-3.2", type="video", video=video)
|
||||
return await _ray32_generate(cls, request)
|
||||
|
||||
|
||||
class LumaRay32KeyframeNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32KeyframeNode",
|
||||
display_name="Luma Ray 3.2 Keyframe",
|
||||
category="partner/video/Luma",
|
||||
description="Anchor a guide image to a position on the Ray 3.2 output video timeline. Connect this to "
|
||||
"the 'keyframes' input of the Luma Ray 3.2 Keyframes to Video node; chain several together via the "
|
||||
"optional 'keyframes' input below.",
|
||||
inputs=[
|
||||
IO.Image.Input("image", tooltip="Guide image to place at the chosen moment of the output video."),
|
||||
IO.DynamicCombo.Input(
|
||||
"position",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Fraction of duration (0.0-1.0)",
|
||||
[
|
||||
IO.Float.Input(
|
||||
"fraction",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Where in the output video this image applies " "(0.0 = start, 1.0 = end).",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"Absolute time (seconds)",
|
||||
[
|
||||
IO.Float.Input(
|
||||
"seconds",
|
||||
default=0.0,
|
||||
min=0.0,
|
||||
max=10.0,
|
||||
step=0.1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
tooltip="Time in seconds from the start of the output video where this "
|
||||
"image applies.",
|
||||
),
|
||||
],
|
||||
),
|
||||
],
|
||||
tooltip="How to place this image on the output video's timeline.",
|
||||
),
|
||||
IO.Custom(LumaIO.LUMA_RAY32_KEYFRAME).Input(
|
||||
"keyframes",
|
||||
optional=True,
|
||||
tooltip="Optional earlier keyframes to chain with this one.",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Custom(LumaIO.LUMA_RAY32_KEYFRAME).Output(display_name="keyframes")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(
|
||||
cls,
|
||||
image: torch.Tensor,
|
||||
position: dict,
|
||||
keyframes: LumaRay32KeyframeChain | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
chain = keyframes.clone() if keyframes is not None else LumaRay32KeyframeChain()
|
||||
if position["position"] == "Absolute time (seconds)":
|
||||
mode, value = LUMA_KEYFRAME_MODE_SECONDS, float(position["seconds"])
|
||||
else:
|
||||
mode, value = LUMA_KEYFRAME_MODE_FRACTION, float(position["fraction"])
|
||||
chain.add(LumaRay32KeyframeItem(image=image, mode=mode, value=value))
|
||||
return IO.NodeOutput(chain)
|
||||
|
||||
|
||||
class LumaRay32KeyframesToVideoNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32KeyframesToVideoNode",
|
||||
display_name="Luma Ray 3.2 Keyframes to Video",
|
||||
category="partner/video/Luma",
|
||||
description="Generate a video that interpolates through a sequence of guide images, each anchored to a "
|
||||
"position on the timeline, using Luma Ray 3.2. Build the sequence with Luma Ray 3.2 Keyframe nodes "
|
||||
"(at least 2).",
|
||||
inputs=[
|
||||
IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the video generation."),
|
||||
IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
|
||||
IO.Combo.Input("duration", options=["5s", "10s"]),
|
||||
_ray32_seed_input(),
|
||||
IO.Custom(LumaIO.LUMA_RAY32_KEYFRAME).Input(
|
||||
"keyframes",
|
||||
tooltip="Keyframe sequence from Luma Ray 3.2 Keyframe nodes (at least 2).",
|
||||
),
|
||||
],
|
||||
outputs=[IO.Video.Output(), IO.String.Output(display_name="generation_id")],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=_BADGE_RAY32_VIDEO,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
resolution: str,
|
||||
duration: str,
|
||||
seed: int,
|
||||
keyframes: LumaRay32KeyframeChain | None = None,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
|
||||
items = keyframes.items if keyframes is not None else []
|
||||
if len(items) < 2:
|
||||
raise ValueError(
|
||||
"Connect at least 2 Luma Ray 3.2 Keyframe nodes "
|
||||
"(use Luma Ray 3.2 Image to Video for a single start/end frame)."
|
||||
)
|
||||
if len(items) > 64:
|
||||
raise ValueError(f"Ray 3.2 supports at most 64 keyframes; got {len(items)}.")
|
||||
maxframe = 120 if duration == "5s" else 240
|
||||
duration_seconds = maxframe / 24 # 5.0 or 10.0
|
||||
# Resolve each keyframe to an output-frame index, then order by position
|
||||
# (so the user can chain keyframes in any order — the position is what places them)
|
||||
placed: list[tuple[int, torch.Tensor]] = []
|
||||
for item in items:
|
||||
if item.mode == LUMA_KEYFRAME_MODE_SECONDS:
|
||||
if item.value > duration_seconds:
|
||||
raise ValueError(
|
||||
f"Keyframe position {item.value:g}s is past the end of the {duration} video; "
|
||||
f"use 0-{duration_seconds:g}s (or switch the keyframe to fraction mode)."
|
||||
)
|
||||
idx = round(item.value * 24)
|
||||
else:
|
||||
idx = round(item.value * maxframe)
|
||||
placed.append((max(0, min(maxframe, idx)), item.image))
|
||||
placed.sort(key=lambda p: p[0])
|
||||
indexes = [idx for idx, _ in placed]
|
||||
for a, b in zip(indexes, indexes[1:]):
|
||||
if a == b:
|
||||
raise ValueError(
|
||||
f"Two keyframes resolve to the same output frame ({a}) for a {duration} video "
|
||||
f"(valid range 0-{maxframe}); give each keyframe a distinct position."
|
||||
)
|
||||
refs: list[Luma2ImageRef] = []
|
||||
for _, image in placed:
|
||||
url = await upload_image_to_comfyapi(cls, image, mime_type="image/png")
|
||||
refs.append(Luma2ImageRef(url=url))
|
||||
request = Luma2GenerationRequest(
|
||||
prompt=prompt,
|
||||
model="ray-3.2",
|
||||
type="video",
|
||||
video=Luma2VideoOptions(resolution=resolution, duration=duration, keyframes=refs, keyframe_indexes=indexes),
|
||||
)
|
||||
return await _ray32_generate(cls, request)
|
||||
|
||||
|
||||
class LumaRay32VideoEditNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32VideoEditNode",
|
||||
display_name="Luma Ray 3.2 Video Edit",
|
||||
category="partner/video/Luma",
|
||||
description="Re-render an existing video under a new prompt using Luma Ray 3.2 (restyle, relight, add "
|
||||
"or remove elements) while keeping the original motion. Source video up to 18 seconds; the edited "
|
||||
"video keeps the source's length.",
|
||||
inputs=[
|
||||
IO.Video.Input("video", tooltip="Source video to edit. Up to 18 seconds."),
|
||||
IO.String.Input("prompt", multiline=True, default="", tooltip="Describes the desired edit."),
|
||||
IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
|
||||
IO.Combo.Input(
|
||||
"strength",
|
||||
options=[
|
||||
"auto",
|
||||
"adhere_1",
|
||||
"adhere_2",
|
||||
"adhere_3",
|
||||
"flex_1",
|
||||
"flex_2",
|
||||
"flex_3",
|
||||
"reimagine_1",
|
||||
"reimagine_2",
|
||||
"reimagine_3",
|
||||
],
|
||||
default="auto",
|
||||
tooltip="How strongly to preserve vs. reimagine the source. 'auto' lets Ray 3.2 choose; "
|
||||
"adhere_* preserves the most, flex_* is balanced, reimagine_* changes the most.",
|
||||
),
|
||||
_ray32_seed_input(),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="generation_id"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=_BADGE_RAY32_EDIT,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls, video: Input.Video, prompt: str, resolution: str, strength: str, seed: int
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=True, min_length=1, max_length=6000)
|
||||
try:
|
||||
duration = "5s" if video.get_duration() <= 5.0 else "10s"
|
||||
except Exception:
|
||||
duration = "10s"
|
||||
source_url = await upload_video_to_comfyapi(cls, video, max_duration=18)
|
||||
edit = Luma2VideoEdit(auto_controls=True) if strength == "auto" else Luma2VideoEdit(strength=strength)
|
||||
request = Luma2GenerationRequest(
|
||||
prompt=prompt,
|
||||
model="ray-3.2",
|
||||
type="video_edit",
|
||||
source=Luma2ImageRef(url=source_url, media_type="video/mp4"),
|
||||
video=Luma2VideoOptions(resolution=resolution, duration=duration, edit=edit),
|
||||
)
|
||||
return await _ray32_generate(cls, request)
|
||||
|
||||
|
||||
class LumaRay32VideoReframeNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32VideoReframeNode",
|
||||
display_name="Luma Ray 3.2 Video Reframe",
|
||||
category="partner/video/Luma",
|
||||
description="Change the aspect ratio of an existing video, using Luma Ray 3.2 to fill the newly "
|
||||
"exposed canvas areas. Source video up to 30 seconds. Billed per second of output.",
|
||||
inputs=[
|
||||
IO.Video.Input("video", tooltip="Source video to reframe. Up to 30 seconds."),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Describes how the newly exposed canvas areas should be filled.",
|
||||
),
|
||||
IO.Combo.Input("aspect_ratio", options=["16:9", "9:16", "1:1", "4:3", "3:4", "21:9"]),
|
||||
IO.Combo.Input("resolution", options=["360p", "540p", "720p", "1080p"], default="720p"),
|
||||
_ray32_seed_input(),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="generation_id"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=_BADGE_RAY32_REFRAME,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls, video: Input.Video, prompt: str, aspect_ratio: str, resolution: str, seed: int
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=False, min_length=1, max_length=6000)
|
||||
if resolution == "1080p" and aspect_ratio in {"9:16", "3:4"}:
|
||||
raise ValueError("1080p is not available for vertical aspect ratios (9:16, 3:4) when reframing.")
|
||||
source_url = await upload_video_to_comfyapi(cls, video, max_duration=30)
|
||||
request = Luma2GenerationRequest(
|
||||
prompt=prompt,
|
||||
model="ray-3.2",
|
||||
type="video_reframe",
|
||||
aspect_ratio=aspect_ratio,
|
||||
source=Luma2ImageRef(url=source_url, media_type="video/mp4"),
|
||||
video=Luma2VideoOptions(resolution=resolution),
|
||||
)
|
||||
return await _ray32_generate(cls, request)
|
||||
|
||||
|
||||
class LumaRay32ExtendVideoNode(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> IO.Schema:
|
||||
return IO.Schema(
|
||||
node_id="LumaRay32ExtendVideoNode",
|
||||
display_name="Luma Ray 3.2 Extend Video",
|
||||
category="partner/video/Luma",
|
||||
description="Extend a previous Ray 3.2 generation forward (continue after it) or backward (lead-in "
|
||||
"before it). Connect the generation_id output of a prior Luma Ray 3.2 node."
|
||||
" Extensions are always 5 seconds.",
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"source_generation_id",
|
||||
default="",
|
||||
tooltip="generation_id of the prior Ray 3.2 video to extend."
|
||||
" Connect the generation_id output of another Luma Ray 3.2 node.",
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"direction",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"Forward (continue after)",
|
||||
[
|
||||
IO.Boolean.Input(
|
||||
"loop",
|
||||
default=False,
|
||||
tooltip="Loop the extended video seamlessly (forward extend only).",
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option("Backward (lead-in before)", []),
|
||||
],
|
||||
tooltip="Forward continues after the prior clip; backward is prepended before it.",
|
||||
),
|
||||
IO.String.Input("prompt", multiline=True, default="", tooltip="Text prompt for the new content."),
|
||||
IO.Combo.Input("resolution", options=["540p", "720p", "1080p"], default="720p"),
|
||||
_ray32_seed_input(),
|
||||
],
|
||||
outputs=[
|
||||
IO.Video.Output(),
|
||||
IO.String.Output(display_name="generation_id"),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=_BADGE_RAY32_VIDEO_5S,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls, source_generation_id: str, direction: dict, prompt: str, resolution: str, seed: int
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=False, min_length=1, max_length=6000)
|
||||
gen_id = (source_generation_id or "").strip()
|
||||
if not gen_id:
|
||||
raise ValueError(
|
||||
"source_generation_id is required (connect the generation_id output of a prior Luma Ray 3.2 node)."
|
||||
)
|
||||
video = Luma2VideoOptions(resolution=resolution, duration="5s")
|
||||
ref = Luma2ImageRef(generation_id=gen_id)
|
||||
if direction["direction"] == "Forward (continue after)":
|
||||
video.start_frame = ref
|
||||
if direction.get("loop"):
|
||||
video.loop = True
|
||||
else:
|
||||
video.end_frame = ref
|
||||
request = Luma2GenerationRequest(prompt=prompt, model="ray-3.2", type="video", video=video)
|
||||
return await _ray32_generate(cls, request)
|
||||
|
||||
|
||||
class LumaExtension(ComfyExtension):
|
||||
@@ -944,6 +1481,13 @@ class LumaExtension(ComfyExtension):
|
||||
LumaConceptsNode,
|
||||
LumaImageNode,
|
||||
LumaImageEditNode,
|
||||
LumaRay32TextToVideoNode,
|
||||
LumaRay32ImageToVideoNode,
|
||||
LumaRay32KeyframeNode,
|
||||
LumaRay32KeyframesToVideoNode,
|
||||
LumaRay32VideoEditNode,
|
||||
LumaRay32VideoReframeNode,
|
||||
LumaRay32ExtendVideoNode,
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -1,932 +0,0 @@
|
||||
from inspect import cleandoc
|
||||
from typing import Optional
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, Input, IO
|
||||
from comfy_api_nodes.apis.stability import (
|
||||
StabilityUpscaleConservativeRequest,
|
||||
StabilityUpscaleCreativeRequest,
|
||||
StabilityAsyncResponse,
|
||||
StabilityResultsGetResponse,
|
||||
StabilityStable3_5Request,
|
||||
StabilityStableUltraRequest,
|
||||
StabilityStableUltraResponse,
|
||||
StabilityAspectRatio,
|
||||
Stability_SD3_5_Model,
|
||||
Stability_SD3_5_GenerationMode,
|
||||
get_stability_style_presets,
|
||||
StabilityTextToAudioRequest,
|
||||
StabilityAudioToAudioRequest,
|
||||
StabilityAudioInpaintRequest,
|
||||
StabilityAudioResponse,
|
||||
)
|
||||
from comfy_api_nodes.util import (
|
||||
validate_audio_duration,
|
||||
validate_string,
|
||||
audio_input_to_mp3,
|
||||
bytesio_to_image_tensor,
|
||||
tensor_to_bytesio,
|
||||
audio_bytes_to_audio_input,
|
||||
sync_op,
|
||||
poll_op,
|
||||
ApiEndpoint,
|
||||
)
|
||||
|
||||
import torch
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from enum import Enum
|
||||
|
||||
|
||||
class StabilityPollStatus(str, Enum):
|
||||
finished = "finished"
|
||||
in_progress = "in_progress"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
def get_async_dummy_status(x: StabilityResultsGetResponse):
|
||||
if x.name is not None or x.errors is not None:
|
||||
return StabilityPollStatus.failed
|
||||
elif x.finish_reason is not None:
|
||||
return StabilityPollStatus.finished
|
||||
return StabilityPollStatus.in_progress
|
||||
|
||||
|
||||
class StabilityStableImageUltraNode(IO.ComfyNode):
|
||||
"""
|
||||
Generates images synchronously based on prompt and resolution.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityStableImageUltraNode",
|
||||
display_name="Stability AI Stable Image Ultra",
|
||||
category="partner/image/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines" +
|
||||
"elements, colors, and subjects will lead to better results. " +
|
||||
"To control the weight of a given word use the format `(word:weight)`," +
|
||||
"where `word` is the word you'd like to control the weight of and `weight`" +
|
||||
"is a value between 0 and 1. For example: `The sky was a crisp (blue:0.3) and (green:0.8)`" +
|
||||
"would convey a sky that was blue and green, but more green than blue.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=StabilityAspectRatio,
|
||||
default=StabilityAspectRatio.ratio_1_1,
|
||||
tooltip="Aspect ratio of generated image.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"style_preset",
|
||||
options=get_stability_style_presets(),
|
||||
tooltip="Optional desired style of generated image.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for creating the noise.",
|
||||
),
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
optional=True,
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
default="",
|
||||
tooltip="A blurb of text describing what you do not wish to see in the output image. This is an advanced feature.",
|
||||
force_input=True,
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"image_denoise",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip="Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.",
|
||||
optional=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.08}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
prompt: str,
|
||||
aspect_ratio: str,
|
||||
style_preset: str,
|
||||
seed: int,
|
||||
image: Optional[torch.Tensor] = None,
|
||||
negative_prompt: str = "",
|
||||
image_denoise: Optional[float] = 0.5,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=False)
|
||||
# prepare image binary if image present
|
||||
image_binary = None
|
||||
if image is not None:
|
||||
image_binary = tensor_to_bytesio(image, total_pixels=1504*1504).read()
|
||||
else:
|
||||
image_denoise = None
|
||||
|
||||
if not negative_prompt:
|
||||
negative_prompt = None
|
||||
if style_preset == "None":
|
||||
style_preset = None
|
||||
|
||||
files = {
|
||||
"image": image_binary
|
||||
}
|
||||
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/stable-image/generate/ultra", method="POST"),
|
||||
response_model=StabilityStableUltraResponse,
|
||||
data=StabilityStableUltraRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seed=seed,
|
||||
strength=image_denoise,
|
||||
style_preset=style_preset,
|
||||
),
|
||||
files=files,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
|
||||
if response_api.finish_reason != "SUCCESS":
|
||||
raise Exception(f"Stable Image Ultra generation failed: {response_api.finish_reason}.")
|
||||
|
||||
image_data = base64.b64decode(response_api.image)
|
||||
returned_image = bytesio_to_image_tensor(BytesIO(image_data))
|
||||
|
||||
return IO.NodeOutput(returned_image)
|
||||
|
||||
|
||||
class StabilityStableImageSD_3_5Node(IO.ComfyNode):
|
||||
"""
|
||||
Generates images synchronously based on prompt and resolution.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityStableImageSD_3_5Node",
|
||||
display_name="Stability AI Stable Diffusion 3.5 Image",
|
||||
category="partner/image/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"model",
|
||||
options=Stability_SD3_5_Model,
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"aspect_ratio",
|
||||
options=StabilityAspectRatio,
|
||||
default=StabilityAspectRatio.ratio_1_1,
|
||||
tooltip="Aspect ratio of generated image.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"style_preset",
|
||||
options=get_stability_style_presets(),
|
||||
tooltip="Optional desired style of generated image.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"cfg_scale",
|
||||
default=4.0,
|
||||
min=1.0,
|
||||
max=10.0,
|
||||
step=0.1,
|
||||
tooltip="How strictly the diffusion process adheres to the prompt text (higher values keep your image closer to your prompt)",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for creating the noise.",
|
||||
),
|
||||
IO.Image.Input(
|
||||
"image",
|
||||
optional=True,
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
default="",
|
||||
tooltip="Keywords of what you do not wish to see in the output image. This is an advanced feature.",
|
||||
force_input=True,
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"image_denoise",
|
||||
default=0.5,
|
||||
min=0.0,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
tooltip="Denoise of input image; 0.0 yields image identical to input, 1.0 is as if no image was provided at all.",
|
||||
optional=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
depends_on=IO.PriceBadgeDepends(widgets=["model"]),
|
||||
expr="""
|
||||
(
|
||||
$contains(widgets.model,"large")
|
||||
? {"type":"usd","usd":0.065}
|
||||
: {"type":"usd","usd":0.035}
|
||||
)
|
||||
""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
model: str,
|
||||
prompt: str,
|
||||
aspect_ratio: str,
|
||||
style_preset: str,
|
||||
seed: int,
|
||||
cfg_scale: float,
|
||||
image: Optional[torch.Tensor] = None,
|
||||
negative_prompt: str = "",
|
||||
image_denoise: Optional[float] = 0.5,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=False)
|
||||
# prepare image binary if image present
|
||||
image_binary = None
|
||||
mode = Stability_SD3_5_GenerationMode.text_to_image
|
||||
if image is not None:
|
||||
image_binary = tensor_to_bytesio(image, total_pixels=1504*1504).read()
|
||||
mode = Stability_SD3_5_GenerationMode.image_to_image
|
||||
aspect_ratio = None
|
||||
else:
|
||||
image_denoise = None
|
||||
|
||||
if not negative_prompt:
|
||||
negative_prompt = None
|
||||
if style_preset == "None":
|
||||
style_preset = None
|
||||
|
||||
files = {
|
||||
"image": image_binary
|
||||
}
|
||||
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/stable-image/generate/sd3", method="POST"),
|
||||
response_model=StabilityStableUltraResponse,
|
||||
data=StabilityStable3_5Request(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seed=seed,
|
||||
strength=image_denoise,
|
||||
style_preset=style_preset,
|
||||
cfg_scale=cfg_scale,
|
||||
model=model,
|
||||
mode=mode,
|
||||
),
|
||||
files=files,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
|
||||
if response_api.finish_reason != "SUCCESS":
|
||||
raise Exception(f"Stable Diffusion 3.5 Image generation failed: {response_api.finish_reason}.")
|
||||
|
||||
image_data = base64.b64decode(response_api.image)
|
||||
returned_image = bytesio_to_image_tensor(BytesIO(image_data))
|
||||
|
||||
return IO.NodeOutput(returned_image)
|
||||
|
||||
|
||||
class StabilityUpscaleConservativeNode(IO.ComfyNode):
|
||||
"""
|
||||
Upscale image with minimal alterations to 4K resolution.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityUpscaleConservativeNode",
|
||||
display_name="Stability AI Upscale Conservative",
|
||||
category="partner/image/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"creativity",
|
||||
default=0.35,
|
||||
min=0.2,
|
||||
max=0.5,
|
||||
step=0.01,
|
||||
tooltip="Controls the likelihood of creating additional details not heavily conditioned by the init image.",
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for creating the noise.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
default="",
|
||||
tooltip="Keywords of what you do not wish to see in the output image. This is an advanced feature.",
|
||||
force_input=True,
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.4}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: torch.Tensor,
|
||||
prompt: str,
|
||||
creativity: float,
|
||||
seed: int,
|
||||
negative_prompt: str = "",
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=False)
|
||||
image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read()
|
||||
|
||||
if not negative_prompt:
|
||||
negative_prompt = None
|
||||
|
||||
files = {
|
||||
"image": image_binary
|
||||
}
|
||||
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/stable-image/upscale/conservative", method="POST"),
|
||||
response_model=StabilityStableUltraResponse,
|
||||
data=StabilityUpscaleConservativeRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
creativity=round(creativity,2),
|
||||
seed=seed,
|
||||
),
|
||||
files=files,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
|
||||
if response_api.finish_reason != "SUCCESS":
|
||||
raise Exception(f"Stability Upscale Conservative generation failed: {response_api.finish_reason}.")
|
||||
|
||||
image_data = base64.b64decode(response_api.image)
|
||||
returned_image = bytesio_to_image_tensor(BytesIO(image_data))
|
||||
|
||||
return IO.NodeOutput(returned_image)
|
||||
|
||||
|
||||
class StabilityUpscaleCreativeNode(IO.ComfyNode):
|
||||
"""
|
||||
Upscale image with minimal alterations to 4K resolution.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityUpscaleCreativeNode",
|
||||
display_name="Stability AI Upscale Creative",
|
||||
category="partner/image/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="What you wish to see in the output image. A strong, descriptive prompt that clearly defines elements, colors, and subjects will lead to better results.",
|
||||
),
|
||||
IO.Float.Input(
|
||||
"creativity",
|
||||
default=0.3,
|
||||
min=0.1,
|
||||
max=0.5,
|
||||
step=0.01,
|
||||
tooltip="Controls the likelihood of creating additional details not heavily conditioned by the init image.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"style_preset",
|
||||
options=get_stability_style_presets(),
|
||||
tooltip="Optional desired style of generated image.",
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for creating the noise.",
|
||||
),
|
||||
IO.String.Input(
|
||||
"negative_prompt",
|
||||
default="",
|
||||
tooltip="Keywords of what you do not wish to see in the output image. This is an advanced feature.",
|
||||
force_input=True,
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.6}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
image: torch.Tensor,
|
||||
prompt: str,
|
||||
creativity: float,
|
||||
style_preset: str,
|
||||
seed: int,
|
||||
negative_prompt: str = "",
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, strip_whitespace=False)
|
||||
image_binary = tensor_to_bytesio(image, total_pixels=1024*1024).read()
|
||||
|
||||
if not negative_prompt:
|
||||
negative_prompt = None
|
||||
if style_preset == "None":
|
||||
style_preset = None
|
||||
|
||||
files = {
|
||||
"image": image_binary
|
||||
}
|
||||
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/stable-image/upscale/creative", method="POST"),
|
||||
response_model=StabilityAsyncResponse,
|
||||
data=StabilityUpscaleCreativeRequest(
|
||||
prompt=prompt,
|
||||
negative_prompt=negative_prompt,
|
||||
creativity=round(creativity,2),
|
||||
style_preset=style_preset,
|
||||
seed=seed,
|
||||
),
|
||||
files=files,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
|
||||
response_poll = await poll_op(
|
||||
cls,
|
||||
ApiEndpoint(path=f"/proxy/stability/v2beta/results/{response_api.id}"),
|
||||
response_model=StabilityResultsGetResponse,
|
||||
poll_interval=3,
|
||||
status_extractor=lambda x: get_async_dummy_status(x),
|
||||
)
|
||||
|
||||
if response_poll.finish_reason != "SUCCESS":
|
||||
raise Exception(f"Stability Upscale Creative generation failed: {response_poll.finish_reason}.")
|
||||
|
||||
image_data = base64.b64decode(response_poll.result)
|
||||
returned_image = bytesio_to_image_tensor(BytesIO(image_data))
|
||||
|
||||
return IO.NodeOutput(returned_image)
|
||||
|
||||
|
||||
class StabilityUpscaleFastNode(IO.ComfyNode):
|
||||
"""
|
||||
Quickly upscales an image via Stability API call to 4x its original size; intended for upscaling low-quality/compressed images.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityUpscaleFastNode",
|
||||
display_name="Stability AI Upscale Fast",
|
||||
category="partner/image/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.Image.Input("image"),
|
||||
],
|
||||
outputs=[
|
||||
IO.Image.Output(),
|
||||
],
|
||||
hidden=[
|
||||
IO.Hidden.auth_token_comfy_org,
|
||||
IO.Hidden.api_key_comfy_org,
|
||||
IO.Hidden.unique_id,
|
||||
],
|
||||
is_api_node=True,
|
||||
price_badge=IO.PriceBadge(
|
||||
expr="""{"type":"usd","usd":0.02}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, image: torch.Tensor) -> IO.NodeOutput:
|
||||
image_binary = tensor_to_bytesio(image, total_pixels=4096*4096).read()
|
||||
|
||||
files = {
|
||||
"image": image_binary
|
||||
}
|
||||
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/stable-image/upscale/fast", method="POST"),
|
||||
response_model=StabilityStableUltraResponse,
|
||||
files=files,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
|
||||
if response_api.finish_reason != "SUCCESS":
|
||||
raise Exception(f"Stability Upscale Fast failed: {response_api.finish_reason}.")
|
||||
|
||||
image_data = base64.b64decode(response_api.image)
|
||||
returned_image = bytesio_to_image_tensor(BytesIO(image_data))
|
||||
|
||||
return IO.NodeOutput(returned_image)
|
||||
|
||||
|
||||
class StabilityTextToAudio(IO.ComfyNode):
|
||||
"""Generates high-quality music and sound effects from text descriptions."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityTextToAudio",
|
||||
display_name="Stability AI Text To Audio",
|
||||
category="partner/audio/Stability AI",
|
||||
essentials_category="Audio",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.Combo.Input(
|
||||
"model",
|
||||
options=["stable-audio-2.5"],
|
||||
),
|
||||
IO.String.Input("prompt", multiline=True, default=""),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=190,
|
||||
min=1,
|
||||
max=190,
|
||||
step=1,
|
||||
tooltip="Controls the duration in seconds of the generated audio.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for generation.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"steps",
|
||||
default=8,
|
||||
min=4,
|
||||
max=8,
|
||||
step=1,
|
||||
tooltip="Controls the number of sampling steps.",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
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.2}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(cls, model: str, prompt: str, duration: int, seed: int, steps: int) -> IO.NodeOutput:
|
||||
validate_string(prompt, max_length=10000)
|
||||
payload = StabilityTextToAudioRequest(prompt=prompt, model=model, duration=duration, seed=seed, steps=steps)
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/audio/stable-audio-2/text-to-audio", method="POST"),
|
||||
response_model=StabilityAudioResponse,
|
||||
data=payload,
|
||||
content_type="multipart/form-data",
|
||||
)
|
||||
if not response_api.audio:
|
||||
raise ValueError("No audio file was received in response.")
|
||||
return IO.NodeOutput(audio_bytes_to_audio_input(base64.b64decode(response_api.audio)))
|
||||
|
||||
|
||||
class StabilityAudioToAudio(IO.ComfyNode):
|
||||
"""Transforms existing audio samples into new high-quality compositions using text instructions."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityAudioToAudio",
|
||||
display_name="Stability AI Audio To Audio",
|
||||
category="partner/audio/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.Combo.Input(
|
||||
"model",
|
||||
options=["stable-audio-2.5"],
|
||||
),
|
||||
IO.String.Input("prompt", multiline=True, default=""),
|
||||
IO.Audio.Input("audio", tooltip="Audio must be between 6 and 190 seconds long."),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=190,
|
||||
min=1,
|
||||
max=190,
|
||||
step=1,
|
||||
tooltip="Controls the duration in seconds of the generated audio.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for generation.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"steps",
|
||||
default=8,
|
||||
min=4,
|
||||
max=8,
|
||||
step=1,
|
||||
tooltip="Controls the number of sampling steps.",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Float.Input(
|
||||
"strength",
|
||||
default=1,
|
||||
min=0.01,
|
||||
max=1.0,
|
||||
step=0.01,
|
||||
display_mode=IO.NumberDisplay.slider,
|
||||
tooltip="Parameter controls how much influence the audio parameter has on the generated audio.",
|
||||
optional=True,
|
||||
),
|
||||
],
|
||||
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.2}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls, model: str, prompt: str, audio: Input.Audio, duration: int, seed: int, steps: int, strength: float
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, max_length=10000)
|
||||
validate_audio_duration(audio, 6, 190)
|
||||
payload = StabilityAudioToAudioRequest(
|
||||
prompt=prompt, model=model, duration=duration, seed=seed, steps=steps, strength=strength
|
||||
)
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
ApiEndpoint(path="/proxy/stability/v2beta/audio/stable-audio-2/audio-to-audio", method="POST"),
|
||||
response_model=StabilityAudioResponse,
|
||||
data=payload,
|
||||
content_type="multipart/form-data",
|
||||
files={"audio": audio_input_to_mp3(audio)},
|
||||
)
|
||||
if not response_api.audio:
|
||||
raise ValueError("No audio file was received in response.")
|
||||
return IO.NodeOutput(audio_bytes_to_audio_input(base64.b64decode(response_api.audio)))
|
||||
|
||||
|
||||
class StabilityAudioInpaint(IO.ComfyNode):
|
||||
"""Transforms part of existing audio sample using text instructions."""
|
||||
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return IO.Schema(
|
||||
node_id="StabilityAudioInpaint",
|
||||
display_name="Stability AI Audio Inpaint",
|
||||
category="partner/audio/Stability AI",
|
||||
description=cleandoc(cls.__doc__ or ""),
|
||||
inputs=[
|
||||
IO.Combo.Input(
|
||||
"model",
|
||||
options=["stable-audio-2.5"],
|
||||
),
|
||||
IO.String.Input("prompt", multiline=True, default=""),
|
||||
IO.Audio.Input("audio", tooltip="Audio must be between 6 and 190 seconds long."),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=190,
|
||||
min=1,
|
||||
max=190,
|
||||
step=1,
|
||||
tooltip="Controls the duration in seconds of the generated audio.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"seed",
|
||||
default=0,
|
||||
min=0,
|
||||
max=4294967294,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
control_after_generate=True,
|
||||
tooltip="The random seed used for generation.",
|
||||
optional=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"steps",
|
||||
default=8,
|
||||
min=4,
|
||||
max=8,
|
||||
step=1,
|
||||
tooltip="Controls the number of sampling steps.",
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"mask_start",
|
||||
default=30,
|
||||
min=0,
|
||||
max=190,
|
||||
step=1,
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
IO.Int.Input(
|
||||
"mask_end",
|
||||
default=190,
|
||||
min=0,
|
||||
max=190,
|
||||
step=1,
|
||||
optional=True,
|
||||
advanced=True,
|
||||
),
|
||||
],
|
||||
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.2}""",
|
||||
),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
async def execute(
|
||||
cls,
|
||||
model: str,
|
||||
prompt: str,
|
||||
audio: Input.Audio,
|
||||
duration: int,
|
||||
seed: int,
|
||||
steps: int,
|
||||
mask_start: int,
|
||||
mask_end: int,
|
||||
) -> IO.NodeOutput:
|
||||
validate_string(prompt, max_length=10000)
|
||||
if mask_end <= mask_start:
|
||||
raise ValueError(f"Value of mask_end({mask_end}) should be greater then mask_start({mask_start})")
|
||||
validate_audio_duration(audio, 6, 190)
|
||||
|
||||
payload = StabilityAudioInpaintRequest(
|
||||
prompt=prompt,
|
||||
model=model,
|
||||
duration=duration,
|
||||
seed=seed,
|
||||
steps=steps,
|
||||
mask_start=mask_start,
|
||||
mask_end=mask_end,
|
||||
)
|
||||
response_api = await sync_op(
|
||||
cls,
|
||||
endpoint=ApiEndpoint(path="/proxy/stability/v2beta/audio/stable-audio-2/inpaint", method="POST"),
|
||||
response_model=StabilityAudioResponse,
|
||||
data=payload,
|
||||
content_type="multipart/form-data",
|
||||
files={"audio": audio_input_to_mp3(audio)},
|
||||
)
|
||||
if not response_api.audio:
|
||||
raise ValueError("No audio file was received in response.")
|
||||
return IO.NodeOutput(audio_bytes_to_audio_input(base64.b64decode(response_api.audio)))
|
||||
|
||||
|
||||
class StabilityExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
StabilityStableImageUltraNode,
|
||||
StabilityStableImageSD_3_5Node,
|
||||
StabilityUpscaleConservativeNode,
|
||||
StabilityUpscaleCreativeNode,
|
||||
StabilityUpscaleFastNode,
|
||||
StabilityTextToAudio,
|
||||
StabilityAudioToAudio,
|
||||
StabilityAudioInpaint,
|
||||
]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> StabilityExtension:
|
||||
return StabilityExtension()
|
||||
@@ -48,10 +48,13 @@ from comfy_api_nodes.util import (
|
||||
upload_image_to_comfyapi,
|
||||
upload_video_to_comfyapi,
|
||||
validate_audio_duration,
|
||||
validate_image_aspect_ratio,
|
||||
validate_image_dimensions,
|
||||
validate_string,
|
||||
validate_video_duration,
|
||||
)
|
||||
|
||||
|
||||
RES_IN_PARENS = re.compile(r"\((\d+)\s*[x×]\s*(\d+)\)")
|
||||
|
||||
|
||||
@@ -1657,6 +1660,44 @@ class HappyHorseTextToVideoApi(IO.ComfyNode):
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"happyhorse-1.1-t2v",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Prompt describing the elements and visual features. "
|
||||
"Supports English and Chinese.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720P", "1080P"],
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"ratio",
|
||||
options=[
|
||||
"16:9",
|
||||
"9:16",
|
||||
"1:1",
|
||||
"4:3",
|
||||
"3:4",
|
||||
"21:9",
|
||||
"9:21",
|
||||
"5:4",
|
||||
"4:5",
|
||||
],
|
||||
),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=5,
|
||||
min=3,
|
||||
max=15,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"happyhorse-1.0-t2v",
|
||||
[
|
||||
@@ -1719,7 +1760,9 @@ class HappyHorseTextToVideoApi(IO.ComfyNode):
|
||||
(
|
||||
$res := $lookup(widgets, "model.resolution");
|
||||
$dur := $lookup(widgets, "model.duration");
|
||||
$ppsTable := { "720p": 0.14, "1080p": 0.24 };
|
||||
$ppsTable := $contains(widgets.model, "1.1")
|
||||
? { "720p": 0.2002, "1080p": 0.2574 }
|
||||
: { "720p": 0.14, "1080p": 0.24 };
|
||||
$pps := $lookup($ppsTable, $res);
|
||||
{ "type": "usd", "usd": $pps * $dur }
|
||||
)
|
||||
@@ -1781,6 +1824,30 @@ class HappyHorseImageToVideoApi(IO.ComfyNode):
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"happyhorse-1.1-i2v",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Prompt describing the elements and visual features. "
|
||||
"Supports English and Chinese.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720P", "1080P"],
|
||||
),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=5,
|
||||
min=3,
|
||||
max=15,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"happyhorse-1.0-i2v",
|
||||
[
|
||||
@@ -1843,7 +1910,9 @@ class HappyHorseImageToVideoApi(IO.ComfyNode):
|
||||
(
|
||||
$res := $lookup(widgets, "model.resolution");
|
||||
$dur := $lookup(widgets, "model.duration");
|
||||
$ppsTable := { "720p": 0.14, "1080p": 0.24 };
|
||||
$ppsTable := $contains(widgets.model, "1.1")
|
||||
? { "720p": 0.2002, "1080p": 0.2574 }
|
||||
: { "720p": 0.14, "1080p": 0.24 };
|
||||
$pps := $lookup($ppsTable, $res);
|
||||
{ "type": "usd", "usd": $pps * $dur }
|
||||
)
|
||||
@@ -1859,6 +1928,8 @@ class HappyHorseImageToVideoApi(IO.ComfyNode):
|
||||
seed: int,
|
||||
watermark: bool,
|
||||
):
|
||||
validate_image_dimensions(first_frame, min_width=300, min_height=300)
|
||||
validate_image_aspect_ratio(first_frame, (1, 2.5), (2.5, 1), strict=False)
|
||||
media = [
|
||||
Wan27MediaItem(
|
||||
type="first_frame",
|
||||
@@ -2053,6 +2124,62 @@ class HappyHorseReferenceVideoApi(IO.ComfyNode):
|
||||
IO.DynamicCombo.Input(
|
||||
"model",
|
||||
options=[
|
||||
IO.DynamicCombo.Option(
|
||||
"happyhorse-1.1-r2v",
|
||||
[
|
||||
IO.String.Input(
|
||||
"prompt",
|
||||
multiline=True,
|
||||
default="",
|
||||
tooltip="Prompt describing the video. Use identifiers such as 'character1' and "
|
||||
"'character2' to refer to the reference characters.",
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"resolution",
|
||||
options=["720P", "1080P"],
|
||||
),
|
||||
IO.Combo.Input(
|
||||
"ratio",
|
||||
options=[
|
||||
"16:9",
|
||||
"9:16",
|
||||
"1:1",
|
||||
"4:3",
|
||||
"3:4",
|
||||
"21:9",
|
||||
"9:21",
|
||||
"5:4",
|
||||
"4:5",
|
||||
],
|
||||
),
|
||||
IO.Int.Input(
|
||||
"duration",
|
||||
default=5,
|
||||
min=3,
|
||||
max=15,
|
||||
step=1,
|
||||
display_mode=IO.NumberDisplay.number,
|
||||
),
|
||||
IO.Autogrow.Input(
|
||||
"reference_images",
|
||||
template=IO.Autogrow.TemplateNames(
|
||||
IO.Image.Input("reference_image"),
|
||||
names=[
|
||||
"image1",
|
||||
"image2",
|
||||
"image3",
|
||||
"image4",
|
||||
"image5",
|
||||
"image6",
|
||||
"image7",
|
||||
"image8",
|
||||
"image9",
|
||||
],
|
||||
min=1,
|
||||
),
|
||||
),
|
||||
],
|
||||
),
|
||||
IO.DynamicCombo.Option(
|
||||
"happyhorse-1.0-r2v",
|
||||
[
|
||||
@@ -2133,7 +2260,9 @@ class HappyHorseReferenceVideoApi(IO.ComfyNode):
|
||||
(
|
||||
$res := $lookup(widgets, "model.resolution");
|
||||
$dur := $lookup(widgets, "model.duration");
|
||||
$ppsTable := { "720p": 0.14, "1080p": 0.24 };
|
||||
$ppsTable := $contains(widgets.model, "1.1")
|
||||
? { "720p": 0.2002, "1080p": 0.2574 }
|
||||
: { "720p": 0.14, "1080p": 0.24 };
|
||||
$pps := $lookup($ppsTable, $res);
|
||||
{ "type": "usd", "usd": $pps * $dur }
|
||||
)
|
||||
@@ -2149,8 +2278,11 @@ class HappyHorseReferenceVideoApi(IO.ComfyNode):
|
||||
watermark: bool,
|
||||
):
|
||||
validate_string(model["prompt"], strip_whitespace=False, min_length=1)
|
||||
media = []
|
||||
reference_images = model.get("reference_images", {})
|
||||
for key in reference_images:
|
||||
validate_image_dimensions(reference_images[key], min_width=400, min_height=400)
|
||||
validate_image_aspect_ratio(reference_images[key], (1, 2.5), (2.5, 1), strict=False)
|
||||
media = []
|
||||
for key in reference_images:
|
||||
media.append(
|
||||
Wan27MediaItem(
|
||||
@@ -2159,7 +2291,7 @@ class HappyHorseReferenceVideoApi(IO.ComfyNode):
|
||||
)
|
||||
)
|
||||
if not media:
|
||||
raise ValueError("At least one reference reference image must be provided.")
|
||||
raise ValueError("At least one reference image must be provided.")
|
||||
|
||||
initial_response = await sync_op(
|
||||
cls,
|
||||
|
||||
@@ -26,6 +26,7 @@ from .conversions import (
|
||||
text_filepath_to_base64_string,
|
||||
text_filepath_to_data_uri,
|
||||
trim_video,
|
||||
upscale_image_tensor_to_min_pixels,
|
||||
upscale_video_to_min_pixels,
|
||||
video_to_base64_string,
|
||||
)
|
||||
@@ -99,6 +100,7 @@ __all__ = [
|
||||
"text_filepath_to_base64_string",
|
||||
"text_filepath_to_data_uri",
|
||||
"trim_video",
|
||||
"upscale_image_tensor_to_min_pixels",
|
||||
"upscale_video_to_min_pixels",
|
||||
"video_to_base64_string",
|
||||
# Validation utilities
|
||||
|
||||
@@ -4,6 +4,8 @@ import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Callable
|
||||
from datetime import datetime, timezone
|
||||
from email.utils import parsedate_to_datetime
|
||||
from io import BytesIO
|
||||
|
||||
from yarl import URL
|
||||
@@ -91,6 +93,32 @@ async def sleep_with_interrupt(
|
||||
await asyncio.sleep(min(1.0, end - now))
|
||||
|
||||
|
||||
def _retry_after_wait(value: str | None, fallback: float, max_wait: float) -> float:
|
||||
"""Delay before the next retry, honoring a server ``Retry-After`` header."""
|
||||
|
||||
seconds: float | None = None
|
||||
if value is not None:
|
||||
value = value.strip()
|
||||
if value.isascii() and value.isdigit():
|
||||
# delay-seconds form. The ASCII-digit guard keeps exotic Unicode "digit" characters away from float()
|
||||
# an all-digit string always converts (huge values become inf, never raising).
|
||||
seconds = float(value)
|
||||
elif value:
|
||||
# HTTP-date form. parsedate_to_datetime raises OverflowError (not a ValueError) on absurd years/offsets
|
||||
try:
|
||||
parsed = parsedate_to_datetime(value)
|
||||
except (TypeError, ValueError, OverflowError):
|
||||
parsed = None
|
||||
if parsed is not None:
|
||||
if parsed.tzinfo is None: # naive datetime: HTTP-date is UTC
|
||||
parsed = parsed.replace(tzinfo=timezone.utc)
|
||||
delta = (parsed - datetime.now(timezone.utc)).total_seconds()
|
||||
seconds = delta if delta > 0 else 0.0
|
||||
if seconds is None:
|
||||
return fallback
|
||||
return min(seconds, max_wait)
|
||||
|
||||
|
||||
def mimetype_to_extension(mime_type: str) -> str:
|
||||
"""Converts a MIME type to a file extension."""
|
||||
return mime_type.split("/")[-1].lower()
|
||||
|
||||
@@ -21,6 +21,7 @@ from server import PromptServer
|
||||
|
||||
from . import request_logger
|
||||
from ._helpers import (
|
||||
_retry_after_wait,
|
||||
default_base_url,
|
||||
get_comfy_api_headers,
|
||||
get_node_id,
|
||||
@@ -82,6 +83,7 @@ class _PollUIState:
|
||||
|
||||
|
||||
_RETRY_STATUS = {408, 500, 502, 503, 504} # status 429 is handled separately
|
||||
_MAX_RETRY_AFTER_WAIT = 150.0 # Cap a server Retry-After at this many seconds so a large hint can't block execution
|
||||
COMPLETED_STATUSES = ["succeeded", "succeed", "success", "completed", "finished", "done", "complete"]
|
||||
FAILED_STATUSES = ["cancelled", "canceled", "canceling", "fail", "failed", "error"]
|
||||
QUEUED_STATUSES = ["created", "queued", "queueing", "submitted", "initializing", "wait", "in_queue"]
|
||||
@@ -747,6 +749,7 @@ async def _request_base(cfg: _RequestConfig, expect_binary: bool):
|
||||
should_retry = True
|
||||
|
||||
if should_retry:
|
||||
wait_time = _retry_after_wait(resp.headers.get("Retry-After"), wait_time, _MAX_RETRY_AFTER_WAIT)
|
||||
logging.warning(
|
||||
"HTTP %s %s -> %s. Waiting %.2fs (%s).",
|
||||
method,
|
||||
|
||||
@@ -448,6 +448,15 @@ def _compute_upscale_dims(src_w: int, src_h: int, total_pixels: int) -> tuple[in
|
||||
return new_w, new_h
|
||||
|
||||
|
||||
def upscale_image_tensor_to_min_pixels(image: torch.Tensor, total_pixels: int) -> torch.Tensor:
|
||||
samples = image.movedim(-1, 1)
|
||||
dims = _compute_upscale_dims(samples.shape[3], samples.shape[2], int(total_pixels))
|
||||
if dims is None:
|
||||
return image
|
||||
new_w, new_h = dims
|
||||
return common_upscale(samples, new_w, new_h, "lanczos", "disabled").movedim(1, -1)
|
||||
|
||||
|
||||
def upscale_video_to_min_pixels(video: Input.Video, min_pixels: int) -> Input.Video:
|
||||
"""Upscale a video to meet at least ``min_pixels`` (w * h), preserving aspect ratio.
|
||||
|
||||
|
||||
+80
-1
@@ -4,11 +4,22 @@ Provides normalization and helper functions for job status tracking.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from typing import Optional
|
||||
from typing import Callable, Optional
|
||||
|
||||
from comfy_api.internal import prune_dict
|
||||
|
||||
|
||||
# Result of classifying a job for cancellation.
|
||||
# 'running' -> job is currently executing (interrupt it)
|
||||
# 'pending' -> job is queued but not started (dequeue it)
|
||||
# 'terminal' -> job already finished (present in history); cancel is a no-op
|
||||
# 'unknown' -> job id is not present anywhere
|
||||
CANCEL_RUNNING = 'running'
|
||||
CANCEL_PENDING = 'pending'
|
||||
CANCEL_TERMINAL = 'terminal'
|
||||
CANCEL_UNKNOWN = 'unknown'
|
||||
|
||||
|
||||
class JobStatus:
|
||||
"""Job status constants."""
|
||||
PENDING = 'pending'
|
||||
@@ -407,3 +418,71 @@ def get_all_jobs(
|
||||
jobs = jobs[:limit]
|
||||
|
||||
return (jobs, total_count)
|
||||
|
||||
|
||||
def classify_job_for_cancel(prompt_id: str, running: list, queued: list, history: dict) -> str:
|
||||
"""Classify a job id for cancellation.
|
||||
|
||||
Returns one of CANCEL_RUNNING, CANCEL_PENDING, CANCEL_TERMINAL, CANCEL_UNKNOWN.
|
||||
|
||||
Queue items are tuples whose second element (index 1) is the prompt_id.
|
||||
History is a dict keyed by prompt_id, so a job present there has already
|
||||
finished and cancelling it is a no-op.
|
||||
"""
|
||||
for item in running:
|
||||
if item[1] == prompt_id:
|
||||
return CANCEL_RUNNING
|
||||
for item in queued:
|
||||
if item[1] == prompt_id:
|
||||
return CANCEL_PENDING
|
||||
if prompt_id in history:
|
||||
return CANCEL_TERMINAL
|
||||
return CANCEL_UNKNOWN
|
||||
|
||||
|
||||
def cancel_job(
|
||||
prompt_id: str,
|
||||
running: list,
|
||||
queued: list,
|
||||
history: dict,
|
||||
interrupt: Callable[[str], bool],
|
||||
dequeue: Callable[[str], bool],
|
||||
) -> str:
|
||||
"""Cancel a single job by id, regardless of state.
|
||||
|
||||
Maps the cancel onto the runtime's existing mechanics:
|
||||
- a running job is interrupted via ``interrupt``
|
||||
- a pending job is removed from the queue via ``dequeue``
|
||||
- a job that already finished (terminal) is a no-op
|
||||
- an unknown id is a no-op (callers that need fail-fast behaviour should
|
||||
validate ids up front with ``classify_job_for_cancel``)
|
||||
|
||||
Both ``interrupt`` and ``dequeue`` take the prompt id and return whether
|
||||
they acted on a job that was *actually* in that state, so the value returned
|
||||
here reflects what truly happened rather than the (possibly stale)
|
||||
classification. This matters around the narrow TOCTOU windows where a job
|
||||
changes state between the caller's snapshot and the action:
|
||||
|
||||
- a job classified RUNNING may have finished before ``interrupt`` fires:
|
||||
``interrupt`` returns False and this returns CANCEL_UNKNOWN (no-op).
|
||||
- a job classified PENDING may have started executing before ``dequeue``
|
||||
fires: ``dequeue`` returns False, ``interrupt`` then catches the now-
|
||||
running job and this returns CANCEL_RUNNING. If it had simply finished
|
||||
instead, both return False and this returns CANCEL_UNKNOWN.
|
||||
|
||||
``interrupt`` must be atomic — interrupt the job only if it is still the one
|
||||
running — so a cancel can never land on an unrelated prompt that started in
|
||||
the meantime (see ``execution.PromptQueue.interrupt_if_running``).
|
||||
"""
|
||||
classification = classify_job_for_cancel(prompt_id, running, queued, history)
|
||||
if classification == CANCEL_RUNNING:
|
||||
return CANCEL_RUNNING if interrupt(prompt_id) else CANCEL_UNKNOWN
|
||||
if classification == CANCEL_PENDING:
|
||||
if dequeue(prompt_id):
|
||||
return CANCEL_PENDING
|
||||
# Left the pending queue between classification and dequeue: if it
|
||||
# started executing, interrupt the now-running job; otherwise it has
|
||||
# already finished and the cancel is a genuine no-op.
|
||||
return CANCEL_RUNNING if interrupt(prompt_id) else CANCEL_UNKNOWN
|
||||
# CANCEL_TERMINAL and CANCEL_UNKNOWN are intentional no-ops.
|
||||
return classification
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
def hex_to_rgb(value: str) -> tuple[int, int, int]:
|
||||
h = value.lstrip("#")
|
||||
if len(h) != 6:
|
||||
return (255, 255, 255)
|
||||
try:
|
||||
return (int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16))
|
||||
except ValueError:
|
||||
return (255, 255, 255)
|
||||
|
||||
|
||||
def readable_color(rgb: tuple[int, int, int]) -> tuple[int, int, int]:
|
||||
r, g, b = rgb
|
||||
lum = 0.299 * r + 0.587 * g + 0.114 * b
|
||||
if lum >= 130:
|
||||
return (r, g, b)
|
||||
t = (130 - lum) / (255 - lum)
|
||||
return (round(r + (255 - r) * t), round(g + (255 - g) * t), round(b + (255 - b) * t))
|
||||
|
||||
|
||||
def normalize_palette(colors) -> list[str]:
|
||||
if isinstance(colors, dict):
|
||||
colors = colors.values()
|
||||
return [c.upper() for c in colors if isinstance(c, str) and c]
|
||||
+17
-18
@@ -158,7 +158,7 @@ class SaveAudio(IO.ComfyNode):
|
||||
return IO.Schema(
|
||||
node_id="SaveAudio",
|
||||
search_aliases=["export flac"],
|
||||
display_name="Save Audio (FLAC) (Deprecated)",
|
||||
display_name="Save Audio (FLAC) (DEPRECATED)",
|
||||
category="audio",
|
||||
essentials_category="Audio",
|
||||
inputs=[
|
||||
@@ -166,8 +166,9 @@ class SaveAudio(IO.ComfyNode):
|
||||
IO.String.Input("filename_prefix", default="audio/ComfyUI"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
is_deprecated=True,
|
||||
is_output_node=True,
|
||||
outputs=[IO.Audio.Output("audio")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -175,11 +176,10 @@ class SaveAudio(IO.ComfyNode):
|
||||
if audio is None:
|
||||
raise ValueError("SaveAudio: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(
|
||||
audio,
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=format)
|
||||
)
|
||||
|
||||
save_flac = execute # TODO: remove
|
||||
|
||||
|
||||
class SaveAudioMP3(IO.ComfyNode):
|
||||
@classmethod
|
||||
@@ -187,7 +187,7 @@ class SaveAudioMP3(IO.ComfyNode):
|
||||
return IO.Schema(
|
||||
node_id="SaveAudioMP3",
|
||||
search_aliases=["export mp3"],
|
||||
display_name="Save Audio (MP3) (Deprecated)",
|
||||
display_name="Save Audio (MP3) (DEPRECATED)",
|
||||
category="audio",
|
||||
essentials_category="Audio",
|
||||
inputs=[
|
||||
@@ -196,8 +196,9 @@ class SaveAudioMP3(IO.ComfyNode):
|
||||
IO.Combo.Input("quality", options=["V0", "128k", "320k"], default="V0"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
is_deprecated=True,
|
||||
is_output_node=True,
|
||||
outputs=[IO.Audio.Output("audio")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -205,13 +206,12 @@ class SaveAudioMP3(IO.ComfyNode):
|
||||
if audio is None:
|
||||
raise ValueError("SaveAudioMP3: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(
|
||||
audio,
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(
|
||||
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
|
||||
)
|
||||
)
|
||||
|
||||
save_mp3 = execute # TODO: remove
|
||||
|
||||
|
||||
class SaveAudioOpus(IO.ComfyNode):
|
||||
@classmethod
|
||||
@@ -219,7 +219,7 @@ class SaveAudioOpus(IO.ComfyNode):
|
||||
return IO.Schema(
|
||||
node_id="SaveAudioOpus",
|
||||
search_aliases=["export opus"],
|
||||
display_name="Save Audio (Opus) (Deprecated)",
|
||||
display_name="Save Audio (Opus) (DEPRECATED)",
|
||||
category="audio",
|
||||
inputs=[
|
||||
IO.Audio.Input("audio"),
|
||||
@@ -227,8 +227,9 @@ class SaveAudioOpus(IO.ComfyNode):
|
||||
IO.Combo.Input("quality", options=["64k", "96k", "128k", "192k", "320k"], default="128k"),
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
is_deprecated=True,
|
||||
is_output_node=True,
|
||||
outputs=[IO.Audio.Output("audio")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -236,13 +237,12 @@ class SaveAudioOpus(IO.ComfyNode):
|
||||
if audio is None:
|
||||
raise ValueError("SaveAudioOpus: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(
|
||||
audio,
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(
|
||||
audio, filename_prefix=filename_prefix, cls=cls, format=format, quality=quality
|
||||
)
|
||||
)
|
||||
|
||||
save_opus = execute # TODO: remove
|
||||
|
||||
|
||||
class SaveAudioAdvanced(IO.ComfyNode):
|
||||
@classmethod
|
||||
@@ -258,10 +258,7 @@ class SaveAudioAdvanced(IO.ComfyNode):
|
||||
IO.String.Input(
|
||||
"filename_prefix",
|
||||
default="audio/ComfyUI",
|
||||
tooltip=(
|
||||
"The prefix for the file to save. May include formatting tokens "
|
||||
"such as %date:yyyy-MM-dd%."
|
||||
),
|
||||
tooltip=("The prefix for the file to save. May include formatting tokens such as %date:yyyy-MM-dd%."),
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"format",
|
||||
@@ -279,6 +276,7 @@ class SaveAudioAdvanced(IO.ComfyNode):
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.Audio.Output("audio")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -289,7 +287,7 @@ class SaveAudioAdvanced(IO.ComfyNode):
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=file_format, quality=quality)
|
||||
else:
|
||||
ui=UI.AudioSaveHelper.get_save_audio_ui(audio, filename_prefix=filename_prefix, cls=cls, format=file_format)
|
||||
return IO.NodeOutput(ui=ui)
|
||||
return IO.NodeOutput(audio, ui=ui)
|
||||
|
||||
|
||||
class PreviewAudio(IO.ComfyNode):
|
||||
@@ -305,13 +303,14 @@ class PreviewAudio(IO.ComfyNode):
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.Audio.Output("audio")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, audio) -> IO.NodeOutput:
|
||||
if audio is None:
|
||||
raise ValueError("PreviewAudio: input audio is None (source video may have no audio track).")
|
||||
return IO.NodeOutput(ui=UI.PreviewAudio(audio, cls=cls))
|
||||
return IO.NodeOutput(audio, ui=UI.PreviewAudio(audio, cls=cls))
|
||||
|
||||
save_flac = execute # TODO: remove
|
||||
|
||||
|
||||
@@ -0,0 +1,253 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageDraw, ImageEnhance, ImageFont
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
from comfy_extras.color_util import hex_to_rgb, normalize_palette, readable_color
|
||||
|
||||
_PREVIEW_LONG_EDGE = 1024
|
||||
_PREVIEW_DIM = 0.25
|
||||
|
||||
|
||||
def pixels_to_fractions(box: dict, width: int, height: int) -> dict:
|
||||
w = width or 1
|
||||
h = height or 1
|
||||
return {
|
||||
"x": box.get("x", 0) / w,
|
||||
"y": box.get("y", 0) / h,
|
||||
"w": box.get("width", 0) / w,
|
||||
"h": box.get("height", 0) / h,
|
||||
}
|
||||
|
||||
|
||||
def fractions_to_pixels(box: dict, width: int, height: int) -> dict:
|
||||
x, y = box.get("x", 0.0), box.get("y", 0.0)
|
||||
w, h = box.get("w", 0.0), box.get("h", 0.0)
|
||||
if w < 0:
|
||||
x, w = x + w, -w
|
||||
if h < 0:
|
||||
y, h = y + h, -h
|
||||
return {
|
||||
"x": round(x * width),
|
||||
"y": round(y * height),
|
||||
"width": round(w * width),
|
||||
"height": round(h * height),
|
||||
}
|
||||
|
||||
|
||||
def fractions_to_bbox_frame(boxes: list, width: int, height: int) -> list:
|
||||
pixels = [
|
||||
fractions_to_pixels(box, width, height)
|
||||
for box in boxes
|
||||
if isinstance(box, dict)
|
||||
]
|
||||
return [pixels] if pixels else []
|
||||
|
||||
|
||||
def _font(size: int):
|
||||
try:
|
||||
return ImageFont.load_default(size)
|
||||
except Exception:
|
||||
return ImageFont.load_default()
|
||||
|
||||
|
||||
def _wrap(draw, text: str, font, max_w: float) -> list[str]:
|
||||
lines = []
|
||||
for para in text.split("\n"):
|
||||
line = ""
|
||||
for word in para.split():
|
||||
test = word if not line else line + " " + word
|
||||
if line and draw.textlength(test, font=font) > max_w:
|
||||
lines.append(line)
|
||||
line = word
|
||||
else:
|
||||
line = test
|
||||
lines.append(line)
|
||||
return lines
|
||||
|
||||
|
||||
def _bg_from_image(image) -> Image.Image | None:
|
||||
if image is None:
|
||||
return None
|
||||
try:
|
||||
arr = (image[0].detach().cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
|
||||
return Image.fromarray(arr)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def render_preview(regions, width, height, bg=None):
|
||||
if bg is not None:
|
||||
iw, ih = bg.size
|
||||
long_edge = max(iw, ih) or 1
|
||||
scale = min(1.0, _PREVIEW_LONG_EDGE / long_edge)
|
||||
rw, rh = max(1, round(iw * scale)), max(1, round(ih * scale))
|
||||
base = bg.convert("RGB").resize((rw, rh), Image.LANCZOS)
|
||||
base = ImageEnhance.Brightness(base).enhance(_PREVIEW_DIM)
|
||||
img = base.convert("RGBA")
|
||||
else:
|
||||
long_edge = max(width, height) or 1
|
||||
scale = min(1.0, _PREVIEW_LONG_EDGE / long_edge)
|
||||
rw, rh = max(1, round(width * scale)), max(1, round(height * scale))
|
||||
grey = round(_PREVIEW_DIM * 128)
|
||||
img = Image.new("RGBA", (rw, rh), (grey, grey, grey, 255))
|
||||
|
||||
overlay = Image.new("RGBA", (rw, rh), (0, 0, 0, 0))
|
||||
draw = ImageDraw.Draw(overlay)
|
||||
fs = max(10, round(rh / 64))
|
||||
font = _font(fs)
|
||||
tag_font = _font(max(9, fs - 2))
|
||||
line_h = fs + 2
|
||||
|
||||
for i, region in enumerate(regions):
|
||||
if not isinstance(region, dict):
|
||||
continue
|
||||
palette = [c for c in (region.get("palette") or []) if c]
|
||||
r, g, b = hex_to_rgb(palette[0]) if palette else (140, 140, 140)
|
||||
x1 = max(0, min(rw, round(region.get("x", 0) * rw)))
|
||||
y1 = max(0, min(rh, round(region.get("y", 0) * rh)))
|
||||
x2 = max(0, min(rw, round((region.get("x", 0) + region.get("w", 0)) * rw)))
|
||||
y2 = max(0, min(rh, round((region.get("y", 0) + region.get("h", 0)) * rh)))
|
||||
if x2 < x1:
|
||||
x1, x2 = x2, x1
|
||||
if y2 < y1:
|
||||
y1, y2 = y2, y1
|
||||
|
||||
draw.rectangle([x1, y1, x2, y2], outline=(r, g, b, 255), width=2)
|
||||
|
||||
swatches = palette[:5]
|
||||
if swatches and (x2 - x1) > 2:
|
||||
sh = max(5, fs // 2)
|
||||
seg = (x2 - x1) / len(swatches)
|
||||
for p, hexc in enumerate(swatches):
|
||||
sx = x1 + round(p * seg)
|
||||
draw.rectangle([sx, y1, x1 + round((p + 1) * seg), y1 + sh], fill=hex_to_rgb(hexc))
|
||||
|
||||
etype = "text" if region.get("type") == "text" else "obj"
|
||||
tag = str(i + 1).zfill(2)
|
||||
tw = draw.textlength(tag, font=tag_font)
|
||||
draw.rectangle([x1, y1, x1 + tw + 6, y1 + fs + 2], fill=(r, g, b, 255))
|
||||
tag_fill = (0, 0, 0, 255) if (0.299 * r + 0.587 * g + 0.114 * b) > 140 else (255, 255, 255, 255)
|
||||
draw.text((x1 + 3, y1 + 1), tag, fill=tag_fill, font=tag_font)
|
||||
|
||||
body = region.get("desc", "") or ""
|
||||
if etype == "text" and region.get("text"):
|
||||
body = '"%s"%s' % (region["text"], " — " + body if body else "")
|
||||
if body and (x2 - x1) > 8:
|
||||
ty = y1 + fs + 5
|
||||
for line in _wrap(draw, body, font, x2 - x1 - 8):
|
||||
if ty > y2:
|
||||
break
|
||||
draw.text((x1 + 4, ty), line, fill=readable_color((r, g, b)) + (255,), font=font)
|
||||
ty += line_h
|
||||
|
||||
composed = Image.alpha_composite(img, overlay).convert("RGB")
|
||||
arr = np.asarray(composed, dtype=np.float32) / 255.0
|
||||
return torch.from_numpy(arr).unsqueeze(0)
|
||||
|
||||
|
||||
def boxes_to_regions(boxes, width: int, height: int) -> list:
|
||||
regions: list = []
|
||||
if not isinstance(boxes, list):
|
||||
return regions
|
||||
for box in boxes:
|
||||
if not isinstance(box, dict):
|
||||
continue
|
||||
meta = box.get("metadata")
|
||||
meta = meta if isinstance(meta, dict) else {}
|
||||
regions.append({
|
||||
**pixels_to_fractions(box, width, height),
|
||||
"type": meta.get("type", "obj"),
|
||||
"text": meta.get("text", ""),
|
||||
"desc": meta.get("desc", ""),
|
||||
"palette": meta.get("palette", []),
|
||||
})
|
||||
return regions
|
||||
|
||||
|
||||
def _norm_bbox(region: dict) -> list[int]:
|
||||
def grid(value: float) -> int:
|
||||
return max(0, min(1000, round(value * 1000)))
|
||||
|
||||
x, y = region.get("x", 0.0), region.get("y", 0.0)
|
||||
w, h = region.get("w", 0.0), region.get("h", 0.0)
|
||||
ymin, xmin, ymax, xmax = grid(y), grid(x), grid(y + h), grid(x + w)
|
||||
if ymin > ymax:
|
||||
ymin, ymax = ymax, ymin
|
||||
if xmin > xmax:
|
||||
xmin, xmax = xmax, xmin
|
||||
return [ymin, xmin, ymax, xmax]
|
||||
|
||||
|
||||
def build_elements(regions: list) -> list:
|
||||
elements = []
|
||||
for region in regions:
|
||||
if not isinstance(region, dict):
|
||||
continue
|
||||
etype = "text" if region.get("type") == "text" else "obj"
|
||||
element = {"type": etype}
|
||||
element["bbox"] = _norm_bbox(region)
|
||||
if etype == "text":
|
||||
element["text"] = region.get("text", "")
|
||||
element["desc"] = region.get("desc", "")
|
||||
palette = normalize_palette(region.get("palette", []))
|
||||
if palette:
|
||||
element["color_palette"] = palette[:5]
|
||||
elements.append(element)
|
||||
return elements
|
||||
|
||||
|
||||
class CreateBoundingBoxes(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
editor_state = io.BoundingBoxes.Input(
|
||||
"editor_state",
|
||||
socketless=False,
|
||||
tooltip="Draw bounding boxes and set each box type, text, description, color palette. Start with background element first and foreground last.",
|
||||
)
|
||||
return io.Schema(
|
||||
node_id="CreateBoundingBoxes",
|
||||
display_name="Create Bounding Boxes",
|
||||
category="utilities",
|
||||
description="Draw bounding boxes in a canvas. Outputs Ideogram prompt elements, pixel-space bounding boxes, and a preview image.",
|
||||
inputs=[
|
||||
io.Image.Input(
|
||||
"background",
|
||||
optional=True,
|
||||
tooltip="Optional image used as background in the canvas and preview.",
|
||||
),
|
||||
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,
|
||||
],
|
||||
outputs=[
|
||||
io.Image.Output(display_name="preview"),
|
||||
io.BoundingBox.Output(display_name="bboxes"),
|
||||
io.Array.Output(display_name="elements"),
|
||||
],
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, width, height, editor_state=None, background=None) -> io.NodeOutput:
|
||||
regions = boxes_to_regions(editor_state, width, height)
|
||||
preview = render_preview(regions, width, height, _bg_from_image(background))
|
||||
return io.NodeOutput(
|
||||
preview,
|
||||
fractions_to_bbox_frame(regions, width, height),
|
||||
build_elements(regions),
|
||||
ui={"dims": [width, height]},
|
||||
)
|
||||
|
||||
|
||||
class BoundingBoxesExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [CreateBoundingBoxes]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> BoundingBoxesExtension:
|
||||
return BoundingBoxesExtension()
|
||||
@@ -1,5 +1,6 @@
|
||||
from typing_extensions import override
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
from comfy_extras.color_util import hex_to_rgb
|
||||
|
||||
|
||||
class ColorToRGBInt(io.ComfyNode):
|
||||
@@ -15,21 +16,30 @@ class ColorToRGBInt(io.ComfyNode):
|
||||
],
|
||||
outputs=[
|
||||
io.Int.Output(display_name="rgb_int"),
|
||||
io.Color.Output(display_name="hex")
|
||||
io.Color.Output(display_name="hex"),
|
||||
io.Float.Output(display_name="alpha"),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, color: str) -> io.NodeOutput:
|
||||
# expect format #RRGGBB
|
||||
if len(color) != 7 or color[0] != "#":
|
||||
raise ValueError("Color must be in format #RRGGBB")
|
||||
r = int(color[1:3], 16)
|
||||
g = int(color[3:5], 16)
|
||||
b = int(color[5:7], 16)
|
||||
# expect format #RRGGBB or #RRGGBBAA
|
||||
if len(color) not in (7, 9) or color[0] != "#":
|
||||
raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA")
|
||||
try:
|
||||
int(color[1:], 16)
|
||||
except ValueError:
|
||||
raise ValueError("Color must be in format #RRGGBB or #RRGGBBAA") from None
|
||||
|
||||
alpha = 1.0
|
||||
if len(color) == 9:
|
||||
alpha = int(color[7:9], 16) / 255.0
|
||||
color = color[:7]
|
||||
|
||||
r, g, b = hex_to_rgb(color)
|
||||
|
||||
rgb_int = r * 256 * 256 + g * 256 + b
|
||||
return io.NodeOutput(rgb_int, color)
|
||||
return io.NodeOutput(rgb_int, color, alpha)
|
||||
|
||||
|
||||
class ColorExtension(ComfyExtension):
|
||||
|
||||
@@ -8,7 +8,8 @@ class CLIPTextEncodeControlnet(io.ComfyNode):
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="CLIPTextEncodeControlnet",
|
||||
category="experimental/conditioning",
|
||||
display_name="CLIP Text Encode (Controlnet)",
|
||||
category="model/conditioning",
|
||||
inputs=[
|
||||
io.Clip.Input("clip"),
|
||||
io.Conditioning.Input("conditioning"),
|
||||
@@ -35,11 +36,12 @@ class T5TokenizerOptions(io.ComfyNode):
|
||||
def define_schema(cls) -> io.Schema:
|
||||
return io.Schema(
|
||||
node_id="T5TokenizerOptions",
|
||||
category="experimental/conditioning",
|
||||
display_name="T5 Tokenizer Options",
|
||||
category="model/conditioning",
|
||||
inputs=[
|
||||
io.Clip.Input("clip"),
|
||||
io.Int.Input("min_padding", default=0, min=0, max=10000, step=1, advanced=True),
|
||||
io.Int.Input("min_length", default=0, min=0, max=10000, step=1, advanced=True),
|
||||
io.Int.Input("min_padding", default=0, min=0, max=10000, step=1),
|
||||
io.Int.Input("min_length", default=0, min=0, max=10000, step=1),
|
||||
],
|
||||
outputs=[io.Clip.Output()],
|
||||
is_experimental=True,
|
||||
|
||||
@@ -13,21 +13,22 @@ class ContextWindowsManualNode(io.ComfyNode):
|
||||
description="Manually set context windows.",
|
||||
inputs=[
|
||||
io.Model.Input("model", tooltip="The model to apply context windows to during sampling."),
|
||||
io.Int.Input("context_length", min=1, default=16, tooltip="The length of the context window.", advanced=True),
|
||||
io.Int.Input("context_overlap", min=0, default=4, tooltip="The overlap of the context window.", advanced=True),
|
||||
io.Int.Input("context_length", min=1, default=16, tooltip="The length of the context window."),
|
||||
io.Int.Input("context_overlap", min=0, default=4, tooltip="The overlap of the context window."),
|
||||
io.Combo.Input("context_schedule", options=[
|
||||
comfy.context_windows.ContextSchedules.STATIC_STANDARD,
|
||||
comfy.context_windows.ContextSchedules.UNIFORM_STANDARD,
|
||||
comfy.context_windows.ContextSchedules.UNIFORM_LOOPED,
|
||||
comfy.context_windows.ContextSchedules.BATCHED,
|
||||
], tooltip="The stride of the context window."),
|
||||
io.Int.Input("context_stride", min=1, default=1, tooltip="The stride of the context window; only applicable to uniform schedules.", advanced=True),
|
||||
], default=comfy.context_windows.ContextSchedules.STATIC_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
|
||||
io.Int.Input("context_stride", min=1, default=1, tooltip="The stride of the context window; only applicable to uniform schedules."),
|
||||
io.Boolean.Input("closed_loop", default=False, tooltip="Whether to close the context window loop; only applicable to looped schedules."),
|
||||
io.Combo.Input("fuse_method", options=comfy.context_windows.ContextFuseMethods.LIST_STATIC, default=comfy.context_windows.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
|
||||
io.Int.Input("dim", min=0, max=5, default=0, tooltip="The dimension to apply the context windows to."),
|
||||
io.Boolean.Input("freenoise", default=False, tooltip="Whether to apply FreeNoise noise shuffling, improves window blending."),
|
||||
io.String.Input("cond_retain_index_list", default="", tooltip="List of latent indices to retain in the conditioning tensors for each window, for example setting this to '0' will use the initial start image for each window."),
|
||||
io.String.Input("cond_retain_index_list", default="", tooltip="List of latent indices to retain in the conditioning tensors for each window. For concat-style I2V models (e.g. Wan I2V, HunyuanVideo I2V, Cosmos I2V, SVD) the encoded start image lives in the c_concat conditioning channels; setting this to '0' will retain that start image content at sub-pos 0 of every window."),
|
||||
io.Boolean.Input("split_conds_to_windows", default=False, tooltip="Whether to split multiple conditionings (created by ConditionCombine) to each window based on region index."),
|
||||
io.String.Input("latent_retain_index_list", default="", tooltip="List of latent indices to retain in the noise latent itself for each window. Use for workflows where reference content (e.g. a start image) lives directly in the noise latent rather than in separate conditioning channels (e.g. inplace-style I2V like LTXV, AnimateDiff). Independent of cond_retain_index_list."),
|
||||
io.Boolean.Input("causal_window_fix", default=True, tooltip="Whether to add a causal fix frame to non-0-indexed context windows."),
|
||||
],
|
||||
outputs=[
|
||||
@@ -38,7 +39,7 @@ class ContextWindowsManualNode(io.ComfyNode):
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model: io.Model.Type, context_length: int, context_overlap: int, context_schedule: str, context_stride: int, closed_loop: bool, fuse_method: str, dim: int, freenoise: bool,
|
||||
cond_retain_index_list: list[int]=[], split_conds_to_windows: bool=False, causal_window_fix: bool=True) -> io.Model:
|
||||
cond_retain_index_list: list[int]=[], split_conds_to_windows: bool=False, latent_retain_index_list: list[int]=[], causal_window_fix: bool=True) -> io.Model:
|
||||
model = model.clone()
|
||||
model.model_options["context_handler"] = comfy.context_windows.IndexListContextHandler(
|
||||
context_schedule=comfy.context_windows.get_matching_context_schedule(context_schedule),
|
||||
@@ -51,6 +52,7 @@ class ContextWindowsManualNode(io.ComfyNode):
|
||||
freenoise=freenoise,
|
||||
cond_retain_index_list=cond_retain_index_list,
|
||||
split_conds_to_windows=split_conds_to_windows,
|
||||
latent_retain_index_list=latent_retain_index_list,
|
||||
causal_window_fix=causal_window_fix,
|
||||
)
|
||||
# make memory usage calculation only take into account the context window latents
|
||||
@@ -65,33 +67,71 @@ class WanContextWindowsManualNode(ContextWindowsManualNode):
|
||||
schema = super().define_schema()
|
||||
schema.node_id = "WanContextWindowsManual"
|
||||
schema.display_name = "WAN Context Windows (Manual)"
|
||||
schema.description = "Manually set context windows for WAN-like models (dim=2)."
|
||||
schema.display_name = "Wan Context Windows"
|
||||
schema.description = "Set context windows for Wan-like models."
|
||||
schema.category="model/patch/wan"
|
||||
schema.inputs = [
|
||||
io.Model.Input("model", tooltip="The model to apply context windows to during sampling."),
|
||||
io.Int.Input("context_length", min=1, max=nodes.MAX_RESOLUTION, step=4, default=81, tooltip="The length of the context window.", advanced=True),
|
||||
io.Int.Input("context_overlap", min=0, default=30, tooltip="The overlap of the context window.", advanced=True),
|
||||
io.Int.Input("context_length", min=1, max=nodes.MAX_RESOLUTION, step=4, default=81, tooltip="The length of the context window in real frames. Must be 4*n + 1."),
|
||||
io.Int.Input("context_overlap", min=0, default=30, tooltip="The overlap of the context window in real frames."),
|
||||
io.Combo.Input("context_schedule", options=[
|
||||
comfy.context_windows.ContextSchedules.STATIC_STANDARD,
|
||||
comfy.context_windows.ContextSchedules.UNIFORM_STANDARD,
|
||||
comfy.context_windows.ContextSchedules.UNIFORM_LOOPED,
|
||||
comfy.context_windows.ContextSchedules.BATCHED,
|
||||
], tooltip="The stride of the context window."),
|
||||
], default=comfy.context_windows.ContextSchedules.UNIFORM_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
|
||||
io.Int.Input("context_stride", min=1, default=1, tooltip="The stride of the context window; only applicable to uniform schedules.", advanced=True),
|
||||
io.Boolean.Input("closed_loop", default=False, tooltip="Whether to close the context window loop; only applicable to looped schedules."),
|
||||
io.Boolean.Input("closed_loop", default=False, tooltip="Whether to close the context window loop; only applicable to looped schedules.", advanced=True),
|
||||
io.Combo.Input("fuse_method", options=comfy.context_windows.ContextFuseMethods.LIST_STATIC, default=comfy.context_windows.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
|
||||
io.Boolean.Input("freenoise", default=False, tooltip="Whether to apply FreeNoise noise shuffling, improves window blending."),
|
||||
#io.String.Input("cond_retain_index_list", default="", tooltip="List of latent indices to retain in the conditioning tensors for each window, for example setting this to '0' will use the initial start image for each window."),
|
||||
#io.Boolean.Input("split_conds_to_windows", default=False, tooltip="Whether to split multiple conditionings (created by ConditionCombine) to each window based on region index."),
|
||||
io.Boolean.Input("freenoise", default=True, tooltip="Whether to apply FreeNoise noise shuffling, improves window blending.", advanced=True),
|
||||
io.Boolean.Input("retain_first_frame", default=False, tooltip="Retain the first I2V frame in every context window (may help retain initial reference)."),
|
||||
io.Boolean.Input("split_conds_to_windows", default=False, tooltip="Whether to split multiple conditionings (created by ConditionCombine) to each window based on region index.", advanced=True),
|
||||
]
|
||||
return schema
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model: io.Model.Type, context_length: int, context_overlap: int, context_schedule: str, context_stride: int, closed_loop: bool, fuse_method: str, freenoise: bool,
|
||||
cond_retain_index_list: list[int]=[], split_conds_to_windows: bool=False) -> io.Model:
|
||||
context_length = max(((context_length - 1) // 4) + 1, 1) # at least length 1
|
||||
context_overlap = max(((context_overlap - 1) // 4) + 1, 0) # at least overlap 0
|
||||
return super().execute(model, context_length, context_overlap, context_schedule, context_stride, closed_loop, fuse_method, dim=2, freenoise=freenoise, cond_retain_index_list=cond_retain_index_list, split_conds_to_windows=split_conds_to_windows)
|
||||
retain_first_frame: bool=False, split_conds_to_windows: bool=False) -> io.Model:
|
||||
context_length = max(((context_length - 1) // 4) + 1, 1) # at least length 1
|
||||
context_overlap = max(context_overlap // 4, 0) # at least overlap 0
|
||||
retain_index_list = "0" if retain_first_frame else ""
|
||||
return super().execute(model, context_length, context_overlap, context_schedule, context_stride, closed_loop, fuse_method, dim=2, freenoise=freenoise, cond_retain_index_list=retain_index_list, split_conds_to_windows=split_conds_to_windows)
|
||||
|
||||
|
||||
class LTXVContextWindowsNode(ContextWindowsManualNode):
|
||||
@classmethod
|
||||
def define_schema(cls) -> io.Schema:
|
||||
schema = super().define_schema()
|
||||
schema.node_id = "LTXVContextWindows"
|
||||
schema.display_name = "LTXV Context Windows"
|
||||
schema.description = "Set context windows for LTXV-like models."
|
||||
schema.inputs = [
|
||||
io.Model.Input("model", tooltip="The model to apply context windows to during sampling."),
|
||||
io.Int.Input("context_length", min=1, max=nodes.MAX_RESOLUTION, step=8, default=145, tooltip="The length of the context window in real frames. Must be 8*n + 1."),
|
||||
io.Int.Input("context_overlap", min=0, step=8, default=40, tooltip="The overlap of the context window in real frames."),
|
||||
io.Combo.Input("context_schedule", options=[
|
||||
comfy.context_windows.ContextSchedules.STATIC_STANDARD,
|
||||
comfy.context_windows.ContextSchedules.UNIFORM_STANDARD,
|
||||
comfy.context_windows.ContextSchedules.UNIFORM_LOOPED,
|
||||
comfy.context_windows.ContextSchedules.BATCHED,
|
||||
], default=comfy.context_windows.ContextSchedules.UNIFORM_STANDARD, tooltip="Step-dependent scheduling algorithm for context windows."),
|
||||
io.Int.Input("context_stride", min=1, default=1, tooltip="The stride of the context window; only applicable to uniform schedules.", advanced=True),
|
||||
io.Boolean.Input("closed_loop", default=False, tooltip="Whether to close the context window loop; only applicable to looped schedules.", advanced=True),
|
||||
io.Combo.Input("fuse_method", options=comfy.context_windows.ContextFuseMethods.LIST_STATIC, default=comfy.context_windows.ContextFuseMethods.PYRAMID, tooltip="The method to use to fuse the context windows."),
|
||||
io.Boolean.Input("freenoise", default=True, tooltip="Whether to apply FreeNoise noise shuffling, improves window blending.", advanced=True),
|
||||
io.Boolean.Input("retain_first_frame", default=False, tooltip="Retain the first latent frame in every context window (may help retain initial reference)."),
|
||||
io.Boolean.Input("split_conds_to_windows", default=False, tooltip="Whether to split multiple conditionings (created by ConditionCombine) to each window based on region index.", advanced=True),
|
||||
]
|
||||
return schema
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model: io.Model.Type, context_length: int, context_overlap: int, context_schedule: str, fuse_method: str, freenoise: bool,
|
||||
retain_first_frame: bool=False, split_conds_to_windows: bool=False, context_stride: int=1, closed_loop: bool=False) -> io.Model:
|
||||
context_length = max(((context_length - 1) // 8) + 1, 1) # at least length 1
|
||||
context_overlap = max(context_overlap // 8, 0) # at least overlap 0
|
||||
retain_index_list = "0" if retain_first_frame else ""
|
||||
return super().execute(model, context_length, context_overlap, context_schedule, context_stride, closed_loop, fuse_method, dim=2, freenoise=freenoise,
|
||||
cond_retain_index_list=retain_index_list, latent_retain_index_list=retain_index_list, split_conds_to_windows=split_conds_to_windows)
|
||||
|
||||
|
||||
class ContextWindowsExtension(ComfyExtension):
|
||||
@@ -99,6 +139,7 @@ class ContextWindowsExtension(ComfyExtension):
|
||||
return [
|
||||
ContextWindowsManualNode,
|
||||
WanContextWindowsManualNode,
|
||||
LTXVContextWindowsNode,
|
||||
]
|
||||
|
||||
def comfy_entrypoint():
|
||||
|
||||
@@ -1070,7 +1070,7 @@ class AddNoise(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="AddNoise",
|
||||
category="experimental/custom_sampling/noise",
|
||||
category="model/sampling/noise",
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
io.Model.Input("model"),
|
||||
@@ -1120,7 +1120,7 @@ class ManualSigmas(io.ComfyNode):
|
||||
return io.Schema(
|
||||
node_id="ManualSigmas",
|
||||
search_aliases=["custom noise schedule", "define sigmas"],
|
||||
category="experimental/custom_sampling",
|
||||
category="model/sampling/sigmas",
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
io.String.Input("sigmas", default="1, 0.5", multiline=False)
|
||||
|
||||
@@ -1583,7 +1583,7 @@ class LoadTrainingDataset(io.ComfyNode):
|
||||
shard_path = os.path.join(dataset_dir, shard_file)
|
||||
|
||||
with open(shard_path, "rb") as f:
|
||||
shard_data = torch.load(f)
|
||||
shard_data = torch.load(f, weights_only=True)
|
||||
|
||||
all_latents.extend(shard_data["latents"])
|
||||
all_conditioning.extend(shard_data["conditioning"])
|
||||
|
||||
@@ -77,7 +77,7 @@ class FrameInterpolate(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="FrameInterpolate",
|
||||
display_name="Frame Interpolate",
|
||||
display_name="Run Frame Interpolation Model",
|
||||
category="video",
|
||||
search_aliases=["rife", "film", "frame interpolation", "slow motion", "interpolate frames", "vfi"],
|
||||
inputs=[
|
||||
|
||||
+237
-349
@@ -1,85 +1,68 @@
|
||||
import os
|
||||
import sys
|
||||
import re
|
||||
import ctypes
|
||||
import logging
|
||||
import ctypes.util
|
||||
import importlib.util
|
||||
from typing import TypedDict
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import nodes
|
||||
import comfy_angle
|
||||
from comfy_api.latest import ComfyExtension, io, ui
|
||||
from typing_extensions import override
|
||||
from utils.install_util import get_missing_requirements_message
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def _check_opengl_availability():
|
||||
"""Early check for OpenGL availability. Raises RuntimeError if unlikely to work."""
|
||||
logger.debug("_check_opengl_availability: starting")
|
||||
missing = []
|
||||
def _preload_angle():
|
||||
egl_path = comfy_angle.get_egl_path()
|
||||
gles_path = comfy_angle.get_glesv2_path()
|
||||
|
||||
# Check Python packages (using find_spec to avoid importing)
|
||||
logger.debug("_check_opengl_availability: checking for glfw package")
|
||||
if importlib.util.find_spec("glfw") is None:
|
||||
missing.append("glfw")
|
||||
if sys.platform == "win32":
|
||||
angle_dir = comfy_angle.get_lib_dir()
|
||||
os.add_dll_directory(angle_dir)
|
||||
os.environ["PATH"] = angle_dir + os.pathsep + os.environ.get("PATH", "")
|
||||
|
||||
logger.debug("_check_opengl_availability: checking for OpenGL package")
|
||||
if importlib.util.find_spec("OpenGL") is None:
|
||||
missing.append("PyOpenGL")
|
||||
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
f"OpenGL dependencies not available.\n{get_missing_requirements_message()}\n"
|
||||
)
|
||||
|
||||
# On Linux without display, check if headless backends are available
|
||||
logger.debug(f"_check_opengl_availability: platform={sys.platform}")
|
||||
if sys.platform.startswith("linux"):
|
||||
has_display = os.environ.get("DISPLAY") or os.environ.get("WAYLAND_DISPLAY")
|
||||
logger.debug(f"_check_opengl_availability: has_display={bool(has_display)}")
|
||||
if not has_display:
|
||||
# Check for EGL or OSMesa libraries
|
||||
logger.debug("_check_opengl_availability: checking for EGL library")
|
||||
has_egl = ctypes.util.find_library("EGL")
|
||||
logger.debug("_check_opengl_availability: checking for OSMesa library")
|
||||
has_osmesa = ctypes.util.find_library("OSMesa")
|
||||
|
||||
# Error disabled for CI as it fails this check
|
||||
# if not has_egl and not has_osmesa:
|
||||
# raise RuntimeError(
|
||||
# "GLSL Shader node: No display and no headless backend (EGL/OSMesa) found.\n"
|
||||
# "See error below for installation instructions."
|
||||
# )
|
||||
logger.debug(f"Headless mode: EGL={'yes' if has_egl else 'no'}, OSMesa={'yes' if has_osmesa else 'no'}")
|
||||
|
||||
logger.debug("_check_opengl_availability: completed")
|
||||
mode = 0 if sys.platform == "win32" else ctypes.RTLD_GLOBAL
|
||||
ctypes.CDLL(str(egl_path), mode=mode)
|
||||
ctypes.CDLL(str(gles_path), mode=mode)
|
||||
|
||||
|
||||
# Run early check at import time
|
||||
logger.debug("nodes_glsl: running _check_opengl_availability at import time")
|
||||
_check_opengl_availability()
|
||||
|
||||
# OpenGL modules - initialized lazily when context is created
|
||||
gl = None
|
||||
glfw = None
|
||||
EGL = None
|
||||
# Pre-load ANGLE *before* any PyOpenGL import so that the EGL platform
|
||||
# plugin picks up ANGLE's libEGL / libGLESv2 instead of system libs.
|
||||
_preload_angle()
|
||||
os.environ.setdefault("PYOPENGL_PLATFORM", "egl")
|
||||
|
||||
|
||||
def _import_opengl():
|
||||
"""Import OpenGL module. Called after context is created."""
|
||||
global gl
|
||||
if gl is None:
|
||||
logger.debug("_import_opengl: importing OpenGL.GL")
|
||||
import OpenGL.GL as _gl
|
||||
gl = _gl
|
||||
logger.debug("_import_opengl: import completed")
|
||||
return gl
|
||||
import OpenGL
|
||||
OpenGL.USE_ACCELERATE = False
|
||||
|
||||
|
||||
def _patch_find_library():
|
||||
"""PyOpenGL's EGL platform looks for 'EGL' and 'GLESv2' by short name
|
||||
via ctypes.util.find_library, but ANGLE ships as 'libEGL' and
|
||||
'libGLESv2'. Patch find_library to return the full ANGLE paths so
|
||||
PyOpenGL loads the same libraries we pre-loaded."""
|
||||
if sys.platform == "linux":
|
||||
return
|
||||
import ctypes.util
|
||||
_orig = ctypes.util.find_library
|
||||
def _patched(name):
|
||||
if name == 'EGL':
|
||||
return comfy_angle.get_egl_path()
|
||||
if name == 'GLESv2':
|
||||
return comfy_angle.get_glesv2_path()
|
||||
return _orig(name)
|
||||
ctypes.util.find_library = _patched
|
||||
|
||||
|
||||
_patch_find_library()
|
||||
|
||||
from OpenGL import EGL
|
||||
from OpenGL import GLES3 as gl
|
||||
|
||||
class SizeModeInput(TypedDict):
|
||||
size_mode: str
|
||||
width: int
|
||||
@@ -102,7 +85,7 @@ MAX_OUTPUTS = 4 # fragColor0-3 (MRT)
|
||||
# (-1,-1)---(3,-1)
|
||||
#
|
||||
# v_texCoord is computed from clip space: * 0.5 + 0.5 maps (-1,1) -> (0,1)
|
||||
VERTEX_SHADER = """#version 330 core
|
||||
VERTEX_SHADER = """#version 300 es
|
||||
out vec2 v_texCoord;
|
||||
void main() {
|
||||
vec2 verts[3] = vec2[](vec2(-1, -1), vec2(3, -1), vec2(-1, 3));
|
||||
@@ -126,14 +109,99 @@ void main() {
|
||||
"""
|
||||
|
||||
|
||||
def _convert_es_to_desktop(source: str) -> str:
|
||||
"""Convert GLSL ES (WebGL) shader source to desktop GLSL 330 core."""
|
||||
# Remove any existing #version directive
|
||||
source = re.sub(r"#version\s+\d+(\s+es)?\s*\n?", "", source, flags=re.IGNORECASE)
|
||||
# Remove precision qualifiers (not needed in desktop GLSL)
|
||||
source = re.sub(r"precision\s+(lowp|mediump|highp)\s+\w+\s*;\s*\n?", "", source)
|
||||
# Prepend desktop GLSL version
|
||||
return "#version 330 core\n" + source
|
||||
|
||||
def _egl_attribs(*values):
|
||||
"""Build an EGL_NONE-terminated EGLint attribute array."""
|
||||
vals = list(values) + [EGL.EGL_NONE]
|
||||
return (ctypes.c_int32 * len(vals))(*vals)
|
||||
|
||||
|
||||
# EGL platform extension constants
|
||||
EGL_PLATFORM_ANGLE_ANGLE = 0x3202
|
||||
EGL_PLATFORM_ANGLE_TYPE_ANGLE = 0x3203
|
||||
EGL_PLATFORM_ANGLE_TYPE_VULKAN_ANGLE = 0x3450
|
||||
EGL_MESA_PLATFORM_SURFACELESS = 0x31DD
|
||||
|
||||
|
||||
_eglGetPlatformDisplayEXT = None
|
||||
|
||||
def _get_egl_platform_display_ext(platform, native_display, attribs):
|
||||
"""Call eglGetPlatformDisplayEXT via ctypes (extension, not in PyOpenGL)."""
|
||||
global _eglGetPlatformDisplayEXT
|
||||
if _eglGetPlatformDisplayEXT is None:
|
||||
from OpenGL import platform as _plat
|
||||
egl_lib = _plat.PLATFORM.EGL
|
||||
_get_proc = egl_lib.eglGetProcAddress
|
||||
_get_proc.restype = ctypes.c_void_p
|
||||
_get_proc.argtypes = [ctypes.c_char_p]
|
||||
ptr = _get_proc(b"eglGetPlatformDisplayEXT")
|
||||
if not ptr:
|
||||
return None
|
||||
func_type = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.c_uint32, ctypes.c_void_p, ctypes.c_void_p)
|
||||
_eglGetPlatformDisplayEXT = func_type(ptr)
|
||||
|
||||
raw = _eglGetPlatformDisplayEXT(platform, native_display, attribs)
|
||||
if not raw:
|
||||
return None
|
||||
return ctypes.cast(raw, EGL.EGLDisplay)
|
||||
|
||||
|
||||
def _get_egl_display():
|
||||
"""Get an EGL display, trying the default first then ANGLE's Vulkan
|
||||
platform for headless environments without a display server."""
|
||||
failures = []
|
||||
|
||||
# Try the default display first (works when X11/Wayland is available)
|
||||
display = EGL.eglGetDisplay(EGL.EGL_DEFAULT_DISPLAY)
|
||||
if display:
|
||||
major, minor = ctypes.c_int32(0), ctypes.c_int32(0)
|
||||
try:
|
||||
if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)):
|
||||
return display, major.value, minor.value
|
||||
except Exception as e:
|
||||
failures.append(f"default: {e}")
|
||||
|
||||
logger.info("Default EGL display unavailable, trying headless fallbacks")
|
||||
|
||||
# Headless fallback strategies, tried in order:
|
||||
headless_strategies = [
|
||||
("surfaceless", EGL_MESA_PLATFORM_SURFACELESS, None, None),
|
||||
("ANGLE Vulkan", EGL_PLATFORM_ANGLE_ANGLE, None,
|
||||
_egl_attribs(EGL_PLATFORM_ANGLE_TYPE_ANGLE, EGL_PLATFORM_ANGLE_TYPE_VULKAN_ANGLE)),
|
||||
]
|
||||
|
||||
for name, platform, native_display, attribs in headless_strategies:
|
||||
display = _get_egl_platform_display_ext(platform, native_display, attribs)
|
||||
if not display:
|
||||
failures.append(f"{name}: eglGetPlatformDisplayEXT returned no display")
|
||||
continue
|
||||
major, minor = ctypes.c_int32(0), ctypes.c_int32(0)
|
||||
try:
|
||||
if EGL.eglInitialize(display, ctypes.byref(major), ctypes.byref(minor)):
|
||||
logger.info(f"Using EGL {name} platform (headless)")
|
||||
return display, major.value, minor.value
|
||||
failures.append(f"{name}: eglInitialize returned false")
|
||||
except Exception as e:
|
||||
failures.append(f"{name}: {e}")
|
||||
continue
|
||||
|
||||
details = "\n".join(f" - {f}" for f in failures)
|
||||
raise RuntimeError(
|
||||
"Failed to initialize EGL display.\n"
|
||||
"No display server and no headless EGL platform available.\n"
|
||||
f"Tried:\n{details}\n"
|
||||
"Ensure GPU drivers are installed or set DISPLAY for a virtual framebuffer."
|
||||
)
|
||||
|
||||
|
||||
def _gl_str(name):
|
||||
"""Get an OpenGL string parameter."""
|
||||
v = gl.glGetString(name)
|
||||
if not v:
|
||||
return "Unknown"
|
||||
if isinstance(v, bytes):
|
||||
return v.decode(errors="replace")
|
||||
return ctypes.string_at(v).decode(errors="replace")
|
||||
|
||||
|
||||
def _detect_output_count(source: str) -> int:
|
||||
@@ -159,163 +227,8 @@ def _detect_pass_count(source: str) -> int:
|
||||
return 1
|
||||
|
||||
|
||||
def _init_glfw():
|
||||
"""Initialize GLFW. Returns (window, glfw_module). Raises RuntimeError on failure."""
|
||||
logger.debug("_init_glfw: starting")
|
||||
# On macOS, glfw.init() must be called from main thread or it hangs forever
|
||||
if sys.platform == "darwin":
|
||||
logger.debug("_init_glfw: skipping on macOS")
|
||||
raise RuntimeError("GLFW backend not supported on macOS")
|
||||
|
||||
logger.debug("_init_glfw: importing glfw module")
|
||||
import glfw as _glfw
|
||||
|
||||
logger.debug("_init_glfw: calling glfw.init()")
|
||||
if not _glfw.init():
|
||||
raise RuntimeError("glfw.init() failed")
|
||||
|
||||
try:
|
||||
logger.debug("_init_glfw: setting window hints")
|
||||
_glfw.window_hint(_glfw.VISIBLE, _glfw.FALSE)
|
||||
_glfw.window_hint(_glfw.CONTEXT_VERSION_MAJOR, 3)
|
||||
_glfw.window_hint(_glfw.CONTEXT_VERSION_MINOR, 3)
|
||||
_glfw.window_hint(_glfw.OPENGL_PROFILE, _glfw.OPENGL_CORE_PROFILE)
|
||||
|
||||
logger.debug("_init_glfw: calling create_window()")
|
||||
window = _glfw.create_window(64, 64, "ComfyUI GLSL", None, None)
|
||||
if not window:
|
||||
raise RuntimeError("glfw.create_window() failed")
|
||||
|
||||
logger.debug("_init_glfw: calling make_context_current()")
|
||||
_glfw.make_context_current(window)
|
||||
logger.debug("_init_glfw: completed successfully")
|
||||
return window, _glfw
|
||||
except Exception:
|
||||
logger.debug("_init_glfw: failed, terminating glfw")
|
||||
_glfw.terminate()
|
||||
raise
|
||||
|
||||
|
||||
def _init_egl():
|
||||
"""Initialize EGL for headless rendering. Returns (display, context, surface, EGL_module). Raises RuntimeError on failure."""
|
||||
logger.debug("_init_egl: starting")
|
||||
from OpenGL import EGL as _EGL
|
||||
from OpenGL.EGL import (
|
||||
eglGetDisplay, eglInitialize, eglChooseConfig, eglCreateContext,
|
||||
eglMakeCurrent, eglCreatePbufferSurface, eglBindAPI,
|
||||
eglTerminate, eglDestroyContext, eglDestroySurface,
|
||||
EGL_DEFAULT_DISPLAY, EGL_NO_CONTEXT, EGL_NONE,
|
||||
EGL_SURFACE_TYPE, EGL_PBUFFER_BIT, EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT,
|
||||
EGL_RED_SIZE, EGL_GREEN_SIZE, EGL_BLUE_SIZE, EGL_ALPHA_SIZE, EGL_DEPTH_SIZE,
|
||||
EGL_WIDTH, EGL_HEIGHT, EGL_OPENGL_API,
|
||||
)
|
||||
logger.debug("_init_egl: imports completed")
|
||||
|
||||
display = None
|
||||
context = None
|
||||
surface = None
|
||||
|
||||
try:
|
||||
logger.debug("_init_egl: calling eglGetDisplay()")
|
||||
display = eglGetDisplay(EGL_DEFAULT_DISPLAY)
|
||||
if display == _EGL.EGL_NO_DISPLAY:
|
||||
raise RuntimeError("eglGetDisplay() failed")
|
||||
|
||||
logger.debug("_init_egl: calling eglInitialize()")
|
||||
major, minor = _EGL.EGLint(), _EGL.EGLint()
|
||||
if not eglInitialize(display, major, minor):
|
||||
display = None # Not initialized, don't terminate
|
||||
raise RuntimeError("eglInitialize() failed")
|
||||
logger.debug(f"_init_egl: EGL version {major.value}.{minor.value}")
|
||||
|
||||
config_attribs = [
|
||||
EGL_SURFACE_TYPE, EGL_PBUFFER_BIT,
|
||||
EGL_RENDERABLE_TYPE, EGL_OPENGL_BIT,
|
||||
EGL_RED_SIZE, 8, EGL_GREEN_SIZE, 8, EGL_BLUE_SIZE, 8, EGL_ALPHA_SIZE, 8,
|
||||
EGL_DEPTH_SIZE, 0, EGL_NONE
|
||||
]
|
||||
configs = (_EGL.EGLConfig * 1)()
|
||||
num_configs = _EGL.EGLint()
|
||||
if not eglChooseConfig(display, config_attribs, configs, 1, num_configs) or num_configs.value == 0:
|
||||
raise RuntimeError("eglChooseConfig() failed")
|
||||
config = configs[0]
|
||||
logger.debug(f"_init_egl: config chosen, num_configs={num_configs.value}")
|
||||
|
||||
if not eglBindAPI(EGL_OPENGL_API):
|
||||
raise RuntimeError("eglBindAPI() failed")
|
||||
|
||||
logger.debug("_init_egl: calling eglCreateContext()")
|
||||
context_attribs = [
|
||||
_EGL.EGL_CONTEXT_MAJOR_VERSION, 3,
|
||||
_EGL.EGL_CONTEXT_MINOR_VERSION, 3,
|
||||
_EGL.EGL_CONTEXT_OPENGL_PROFILE_MASK, _EGL.EGL_CONTEXT_OPENGL_CORE_PROFILE_BIT,
|
||||
EGL_NONE
|
||||
]
|
||||
context = eglCreateContext(display, config, EGL_NO_CONTEXT, context_attribs)
|
||||
if context == EGL_NO_CONTEXT:
|
||||
raise RuntimeError("eglCreateContext() failed")
|
||||
|
||||
logger.debug("_init_egl: calling eglCreatePbufferSurface()")
|
||||
pbuffer_attribs = [EGL_WIDTH, 64, EGL_HEIGHT, 64, EGL_NONE]
|
||||
surface = eglCreatePbufferSurface(display, config, pbuffer_attribs)
|
||||
if surface == _EGL.EGL_NO_SURFACE:
|
||||
raise RuntimeError("eglCreatePbufferSurface() failed")
|
||||
|
||||
logger.debug("_init_egl: calling eglMakeCurrent()")
|
||||
if not eglMakeCurrent(display, surface, surface, context):
|
||||
raise RuntimeError("eglMakeCurrent() failed")
|
||||
|
||||
logger.debug("_init_egl: completed successfully")
|
||||
return display, context, surface, _EGL
|
||||
|
||||
except Exception:
|
||||
logger.debug("_init_egl: failed, cleaning up")
|
||||
# Clean up any resources on failure
|
||||
if surface is not None:
|
||||
eglDestroySurface(display, surface)
|
||||
if context is not None:
|
||||
eglDestroyContext(display, context)
|
||||
if display is not None:
|
||||
eglTerminate(display)
|
||||
raise
|
||||
|
||||
|
||||
def _init_osmesa():
|
||||
"""Initialize OSMesa for software rendering. Returns (context, buffer). Raises RuntimeError on failure."""
|
||||
import ctypes
|
||||
|
||||
logger.debug("_init_osmesa: starting")
|
||||
os.environ["PYOPENGL_PLATFORM"] = "osmesa"
|
||||
|
||||
logger.debug("_init_osmesa: importing OpenGL.osmesa")
|
||||
from OpenGL import GL as _gl
|
||||
from OpenGL.osmesa import (
|
||||
OSMesaCreateContextExt, OSMesaMakeCurrent, OSMesaDestroyContext,
|
||||
OSMESA_RGBA,
|
||||
)
|
||||
logger.debug("_init_osmesa: imports completed")
|
||||
|
||||
ctx = OSMesaCreateContextExt(OSMESA_RGBA, 24, 0, 0, None)
|
||||
if not ctx:
|
||||
raise RuntimeError("OSMesaCreateContextExt() failed")
|
||||
|
||||
width, height = 64, 64
|
||||
buffer = (ctypes.c_ubyte * (width * height * 4))()
|
||||
|
||||
logger.debug("_init_osmesa: calling OSMesaMakeCurrent()")
|
||||
if not OSMesaMakeCurrent(ctx, buffer, _gl.GL_UNSIGNED_BYTE, width, height):
|
||||
OSMesaDestroyContext(ctx)
|
||||
raise RuntimeError("OSMesaMakeCurrent() failed")
|
||||
|
||||
logger.debug("_init_osmesa: completed successfully")
|
||||
return ctx, buffer
|
||||
|
||||
|
||||
class GLContext:
|
||||
"""Manages OpenGL context and resources for shader execution.
|
||||
|
||||
Tries backends in order: GLFW (desktop) → EGL (headless GPU) → OSMesa (software).
|
||||
"""
|
||||
"""Manages an OpenGL ES 3.0 context via EGL/ANGLE (singleton)."""
|
||||
|
||||
_instance = None
|
||||
_initialized = False
|
||||
@@ -327,131 +240,105 @@ class GLContext:
|
||||
|
||||
def __init__(self):
|
||||
if GLContext._initialized:
|
||||
logger.debug("GLContext.__init__: already initialized, skipping")
|
||||
return
|
||||
|
||||
logger.debug("GLContext.__init__: starting initialization")
|
||||
|
||||
global glfw, EGL
|
||||
|
||||
import time
|
||||
start = time.perf_counter()
|
||||
|
||||
self._backend = None
|
||||
self._window = None
|
||||
self._egl_display = None
|
||||
self._egl_context = None
|
||||
self._egl_surface = None
|
||||
self._osmesa_ctx = None
|
||||
self._osmesa_buffer = None
|
||||
self._display = None
|
||||
self._surface = None
|
||||
self._context = None
|
||||
self._vao = None
|
||||
|
||||
# Try backends in order: GLFW → EGL → OSMesa
|
||||
errors = []
|
||||
|
||||
logger.debug("GLContext.__init__: trying GLFW backend")
|
||||
try:
|
||||
self._window, glfw = _init_glfw()
|
||||
self._backend = "glfw"
|
||||
logger.debug("GLContext.__init__: GLFW backend succeeded")
|
||||
except Exception as e:
|
||||
logger.debug(f"GLContext.__init__: GLFW backend failed: {e}")
|
||||
errors.append(("GLFW", e))
|
||||
self._display, self._egl_major, self._egl_minor = _get_egl_display()
|
||||
|
||||
if self._backend is None:
|
||||
logger.debug("GLContext.__init__: trying EGL backend")
|
||||
try:
|
||||
self._egl_display, self._egl_context, self._egl_surface, EGL = _init_egl()
|
||||
self._backend = "egl"
|
||||
logger.debug("GLContext.__init__: EGL backend succeeded")
|
||||
except Exception as e:
|
||||
logger.debug(f"GLContext.__init__: EGL backend failed: {e}")
|
||||
errors.append(("EGL", e))
|
||||
if not EGL.eglBindAPI(EGL.EGL_OPENGL_ES_API):
|
||||
raise RuntimeError("eglBindAPI(EGL_OPENGL_ES_API) failed")
|
||||
|
||||
if self._backend is None:
|
||||
logger.debug("GLContext.__init__: trying OSMesa backend")
|
||||
try:
|
||||
self._osmesa_ctx, self._osmesa_buffer = _init_osmesa()
|
||||
self._backend = "osmesa"
|
||||
logger.debug("GLContext.__init__: OSMesa backend succeeded")
|
||||
except Exception as e:
|
||||
logger.debug(f"GLContext.__init__: OSMesa backend failed: {e}")
|
||||
errors.append(("OSMesa", e))
|
||||
config = EGL.EGLConfig()
|
||||
n_configs = ctypes.c_int32(0)
|
||||
if not EGL.eglChooseConfig(
|
||||
self._display,
|
||||
_egl_attribs(
|
||||
EGL.EGL_RENDERABLE_TYPE, EGL.EGL_OPENGL_ES3_BIT,
|
||||
EGL.EGL_SURFACE_TYPE, EGL.EGL_PBUFFER_BIT,
|
||||
EGL.EGL_RED_SIZE, 8, EGL.EGL_GREEN_SIZE, 8,
|
||||
EGL.EGL_BLUE_SIZE, 8, EGL.EGL_ALPHA_SIZE, 8,
|
||||
),
|
||||
ctypes.byref(config), 1, ctypes.byref(n_configs),
|
||||
) or n_configs.value == 0:
|
||||
raise RuntimeError("eglChooseConfig() failed")
|
||||
|
||||
if self._backend is None:
|
||||
if sys.platform == "win32":
|
||||
platform_help = (
|
||||
"Windows: Ensure GPU drivers are installed and display is available.\n"
|
||||
" CPU-only/headless mode is not supported on Windows."
|
||||
)
|
||||
elif sys.platform == "darwin":
|
||||
platform_help = (
|
||||
"macOS: GLFW is not supported.\n"
|
||||
" Install OSMesa via Homebrew: brew install mesa\n"
|
||||
" Then: pip install PyOpenGL PyOpenGL-accelerate"
|
||||
)
|
||||
else:
|
||||
platform_help = (
|
||||
"Linux: Install one of these backends:\n"
|
||||
" Desktop: sudo apt install libgl1-mesa-glx libglfw3\n"
|
||||
" Headless with GPU: sudo apt install libegl1-mesa libgl1-mesa-dri\n"
|
||||
" Headless (CPU): sudo apt install libosmesa6"
|
||||
)
|
||||
|
||||
error_details = "\n".join(f" {name}: {err}" for name, err in errors)
|
||||
raise RuntimeError(
|
||||
f"Failed to create OpenGL context.\n\n"
|
||||
f"Backend errors:\n{error_details}\n\n"
|
||||
f"{platform_help}"
|
||||
self._surface = EGL.eglCreatePbufferSurface(
|
||||
self._display, config,
|
||||
_egl_attribs(EGL.EGL_WIDTH, 64, EGL.EGL_HEIGHT, 64),
|
||||
)
|
||||
if not self._surface:
|
||||
raise RuntimeError("eglCreatePbufferSurface() failed")
|
||||
|
||||
# Now import OpenGL.GL (after context is current)
|
||||
logger.debug("GLContext.__init__: importing OpenGL.GL")
|
||||
_import_opengl()
|
||||
self._context = EGL.eglCreateContext(
|
||||
self._display, config, EGL.EGL_NO_CONTEXT,
|
||||
_egl_attribs(EGL.EGL_CONTEXT_CLIENT_VERSION, 3),
|
||||
)
|
||||
if not self._context:
|
||||
raise RuntimeError("eglCreateContext() failed")
|
||||
|
||||
# Create VAO (required for core profile, but OSMesa may use compat profile)
|
||||
logger.debug("GLContext.__init__: creating VAO")
|
||||
try:
|
||||
vao = gl.glGenVertexArrays(1)
|
||||
gl.glBindVertexArray(vao)
|
||||
self._vao = vao # Only store after successful bind
|
||||
logger.debug("GLContext.__init__: VAO created successfully")
|
||||
except Exception as e:
|
||||
logger.debug(f"GLContext.__init__: VAO creation failed (may be expected for OSMesa): {e}")
|
||||
# OSMesa with older Mesa may not support VAOs
|
||||
# Clean up if we created but couldn't bind
|
||||
if vao:
|
||||
try:
|
||||
gl.glDeleteVertexArrays(1, [vao])
|
||||
except Exception:
|
||||
pass
|
||||
if not EGL.eglMakeCurrent(self._display, self._surface, self._surface, self._context):
|
||||
raise RuntimeError("eglMakeCurrent() failed")
|
||||
|
||||
self._vao = gl.glGenVertexArrays(1)
|
||||
gl.glBindVertexArray(self._vao)
|
||||
|
||||
except Exception:
|
||||
self._cleanup()
|
||||
raise
|
||||
|
||||
elapsed = (time.perf_counter() - start) * 1000
|
||||
|
||||
# Log device info
|
||||
renderer = gl.glGetString(gl.GL_RENDERER)
|
||||
vendor = gl.glGetString(gl.GL_VENDOR)
|
||||
version = gl.glGetString(gl.GL_VERSION)
|
||||
renderer = renderer.decode() if renderer else "Unknown"
|
||||
vendor = vendor.decode() if vendor else "Unknown"
|
||||
version = version.decode() if version else "Unknown"
|
||||
renderer = _gl_str(gl.GL_RENDERER)
|
||||
vendor = _gl_str(gl.GL_VENDOR)
|
||||
version = _gl_str(gl.GL_VERSION)
|
||||
|
||||
GLContext._initialized = True
|
||||
logger.info(f"GLSL context initialized in {elapsed:.1f}ms ({self._backend}) - {renderer} ({vendor}), GL {version}")
|
||||
logger.info(f"GLSL context initialized in {elapsed:.1f}ms - EGL {self._egl_major}.{self._egl_minor}, {renderer} ({vendor}), GL {version}")
|
||||
|
||||
def make_current(self):
|
||||
if self._backend == "glfw":
|
||||
glfw.make_context_current(self._window)
|
||||
elif self._backend == "egl":
|
||||
from OpenGL.EGL import eglMakeCurrent
|
||||
eglMakeCurrent(self._egl_display, self._egl_surface, self._egl_surface, self._egl_context)
|
||||
elif self._backend == "osmesa":
|
||||
from OpenGL.osmesa import OSMesaMakeCurrent
|
||||
OSMesaMakeCurrent(self._osmesa_ctx, self._osmesa_buffer, gl.GL_UNSIGNED_BYTE, 64, 64)
|
||||
|
||||
if not EGL.eglMakeCurrent(self._display, self._surface, self._surface, self._context):
|
||||
err = EGL.eglGetError()
|
||||
raise RuntimeError(f"eglMakeCurrent() failed (EGL error: 0x{err:04X})")
|
||||
if self._vao is not None:
|
||||
gl.glBindVertexArray(self._vao)
|
||||
|
||||
def _cleanup(self):
|
||||
if not self._display:
|
||||
return
|
||||
try:
|
||||
if self._vao is not None:
|
||||
gl.glDeleteVertexArrays(1, [self._vao])
|
||||
self._vao = None
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
EGL.eglMakeCurrent(self._display, EGL.EGL_NO_SURFACE, EGL.EGL_NO_SURFACE, EGL.EGL_NO_CONTEXT)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if self._context:
|
||||
EGL.eglDestroyContext(self._display, self._context)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
if self._surface:
|
||||
EGL.eglDestroySurface(self._display, self._surface)
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
EGL.eglTerminate(self._display)
|
||||
except Exception:
|
||||
pass
|
||||
self._display = None
|
||||
|
||||
|
||||
def _compile_shader(source: str, shader_type: int) -> int:
|
||||
"""Compile a shader and return its ID."""
|
||||
@@ -459,8 +346,10 @@ def _compile_shader(source: str, shader_type: int) -> int:
|
||||
gl.glShaderSource(shader, source)
|
||||
gl.glCompileShader(shader)
|
||||
|
||||
if gl.glGetShaderiv(shader, gl.GL_COMPILE_STATUS) != gl.GL_TRUE:
|
||||
error = gl.glGetShaderInfoLog(shader).decode()
|
||||
if not gl.glGetShaderiv(shader, gl.GL_COMPILE_STATUS):
|
||||
error = gl.glGetShaderInfoLog(shader)
|
||||
if isinstance(error, bytes):
|
||||
error = error.decode(errors="replace")
|
||||
gl.glDeleteShader(shader)
|
||||
raise RuntimeError(f"Shader compilation failed:\n{error}")
|
||||
|
||||
@@ -484,8 +373,10 @@ def _create_program(vertex_source: str, fragment_source: str) -> int:
|
||||
gl.glDeleteShader(vertex_shader)
|
||||
gl.glDeleteShader(fragment_shader)
|
||||
|
||||
if gl.glGetProgramiv(program, gl.GL_LINK_STATUS) != gl.GL_TRUE:
|
||||
error = gl.glGetProgramInfoLog(program).decode()
|
||||
if not gl.glGetProgramiv(program, gl.GL_LINK_STATUS):
|
||||
error = gl.glGetProgramInfoLog(program)
|
||||
if isinstance(error, bytes):
|
||||
error = error.decode(errors="replace")
|
||||
gl.glDeleteProgram(program)
|
||||
raise RuntimeError(f"Program linking failed:\n{error}")
|
||||
|
||||
@@ -530,9 +421,6 @@ def _render_shader_batch(
|
||||
ctx = GLContext()
|
||||
ctx.make_current()
|
||||
|
||||
# Convert from GLSL ES to desktop GLSL 330
|
||||
fragment_source = _convert_es_to_desktop(fragment_code)
|
||||
|
||||
# Detect how many outputs the shader actually uses
|
||||
num_outputs = _detect_output_count(fragment_code)
|
||||
|
||||
@@ -558,9 +446,9 @@ def _render_shader_batch(
|
||||
try:
|
||||
# Compile shaders (once for all batches)
|
||||
try:
|
||||
program = _create_program(VERTEX_SHADER, fragment_source)
|
||||
program = _create_program(VERTEX_SHADER, fragment_code)
|
||||
except RuntimeError:
|
||||
logger.error(f"Fragment shader:\n{fragment_source}")
|
||||
logger.error(f"Fragment shader:\n{fragment_code}")
|
||||
raise
|
||||
|
||||
gl.glUseProgram(program)
|
||||
@@ -723,13 +611,13 @@ def _render_shader_batch(
|
||||
gl.glDrawArrays(gl.GL_TRIANGLES, 0, 3)
|
||||
|
||||
# Read back outputs for this batch
|
||||
# (glGetTexImage is synchronous, implicitly waits for rendering)
|
||||
gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, fbo)
|
||||
batch_outputs = []
|
||||
for tex in output_textures:
|
||||
gl.glBindTexture(gl.GL_TEXTURE_2D, tex)
|
||||
data = gl.glGetTexImage(gl.GL_TEXTURE_2D, 0, gl.GL_RGBA, gl.GL_FLOAT)
|
||||
img = np.frombuffer(data, dtype=np.float32).reshape(height, width, 4)
|
||||
batch_outputs.append(img[::-1, :, :].copy())
|
||||
for i in range(num_outputs):
|
||||
gl.glReadBuffer(gl.GL_COLOR_ATTACHMENT0 + i)
|
||||
buf = np.empty((height, width, 4), dtype=np.float32)
|
||||
gl.glReadPixels(0, 0, width, height, gl.GL_RGBA, gl.GL_FLOAT, buf)
|
||||
batch_outputs.append(buf[::-1, :, :].copy())
|
||||
|
||||
# Pad with black images for unused outputs
|
||||
black_img = np.zeros((height, width, 4), dtype=np.float32)
|
||||
@@ -750,18 +638,18 @@ def _render_shader_batch(
|
||||
gl.glBindFramebuffer(gl.GL_FRAMEBUFFER, 0)
|
||||
gl.glUseProgram(0)
|
||||
|
||||
for tex in input_textures:
|
||||
gl.glDeleteTextures(int(tex))
|
||||
for tex in curve_textures:
|
||||
gl.glDeleteTextures(int(tex))
|
||||
for tex in output_textures:
|
||||
gl.glDeleteTextures(int(tex))
|
||||
for tex in ping_pong_textures:
|
||||
gl.glDeleteTextures(int(tex))
|
||||
if input_textures:
|
||||
gl.glDeleteTextures(len(input_textures), input_textures)
|
||||
if curve_textures:
|
||||
gl.glDeleteTextures(len(curve_textures), curve_textures)
|
||||
if output_textures:
|
||||
gl.glDeleteTextures(len(output_textures), output_textures)
|
||||
if ping_pong_textures:
|
||||
gl.glDeleteTextures(len(ping_pong_textures), ping_pong_textures)
|
||||
if fbo is not None:
|
||||
gl.glDeleteFramebuffers(1, [fbo])
|
||||
for pp_fbo in ping_pong_fbos:
|
||||
gl.glDeleteFramebuffers(1, [pp_fbo])
|
||||
if ping_pong_fbos:
|
||||
gl.glDeleteFramebuffers(len(ping_pong_fbos), ping_pong_fbos)
|
||||
if program is not None:
|
||||
gl.glDeleteProgram(program)
|
||||
|
||||
|
||||
@@ -214,11 +214,13 @@ class SaveAnimatedWEBP(IO.ComfyNode):
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.Image.Output(display_name="images")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, fps, filename_prefix, lossless, quality, method, num_frames=0) -> IO.NodeOutput:
|
||||
return IO.NodeOutput(
|
||||
images,
|
||||
ui=UI.ImageSaveHelper.get_save_animated_webp_ui(
|
||||
images=images,
|
||||
filename_prefix=filename_prefix,
|
||||
@@ -230,8 +232,6 @@ class SaveAnimatedWEBP(IO.ComfyNode):
|
||||
)
|
||||
)
|
||||
|
||||
save_images = execute # TODO: remove
|
||||
|
||||
|
||||
class SaveAnimatedPNG(IO.ComfyNode):
|
||||
|
||||
@@ -249,11 +249,13 @@ class SaveAnimatedPNG(IO.ComfyNode):
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.Image.Output(display_name="images")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, images, fps, compress_level, filename_prefix="ComfyUI") -> IO.NodeOutput:
|
||||
return IO.NodeOutput(
|
||||
images,
|
||||
ui=UI.ImageSaveHelper.get_save_animated_png_ui(
|
||||
images=images,
|
||||
filename_prefix=filename_prefix,
|
||||
@@ -263,8 +265,6 @@ class SaveAnimatedPNG(IO.ComfyNode):
|
||||
)
|
||||
)
|
||||
|
||||
save_images = execute # TODO: remove
|
||||
|
||||
|
||||
class ImageStitch(IO.ComfyNode):
|
||||
"""Upstreamed from https://github.com/kijai/ComfyUI-KJNodes"""
|
||||
@@ -513,6 +513,7 @@ class SaveSVGNode(IO.ComfyNode):
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.SVG.Output("svg")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -562,9 +563,7 @@ class SaveSVGNode(IO.ComfyNode):
|
||||
|
||||
results.append(UI.SavedResult(filename=file, subfolder=subfolder, type=IO.FolderType.output))
|
||||
counter += 1
|
||||
return IO.NodeOutput(ui={"images": results})
|
||||
|
||||
save_svg = execute # TODO: remove
|
||||
return IO.NodeOutput(svg, ui={"images": results})
|
||||
|
||||
|
||||
class GetImageSize(IO.ComfyNode):
|
||||
@@ -1157,40 +1156,27 @@ class SaveImageAdvanced(IO.ComfyNode):
|
||||
IO.String.Input(
|
||||
"filename_prefix",
|
||||
default="ComfyUI",
|
||||
tooltip=(
|
||||
"The prefix for the file to save. May include formatting tokens "
|
||||
"such as %date:yyyy-MM-dd% or %Empty Latent Image.width%."
|
||||
),
|
||||
tooltip=("The prefix for the file to save. May include formatting tokens such as %date:yyyy-MM-dd% or %Empty Latent Image.width%."),
|
||||
),
|
||||
IO.DynamicCombo.Input(
|
||||
"format",
|
||||
options=[
|
||||
IO.DynamicCombo.Option("png", [
|
||||
IO.Combo.Input("bit_depth", options=["8-bit", "16-bit"],
|
||||
default="8-bit", advanced=True),
|
||||
IO.Combo.Input("input_color_space", options=["sRGB"],
|
||||
default="sRGB", advanced=True),
|
||||
IO.Combo.Input("bit_depth", options=["8-bit", "16-bit"], default="8-bit", advanced=True),
|
||||
IO.Combo.Input("input_color_space", options=["sRGB"], default="sRGB", advanced=True),
|
||||
]),
|
||||
IO.DynamicCombo.Option("exr", [
|
||||
IO.Combo.Input("bit_depth", options=["32-bit float"],
|
||||
default="32-bit float", advanced=True),
|
||||
IO.Combo.Input("bit_depth", options=["32-bit float"], default="32-bit float", advanced=True),
|
||||
IO.Combo.Input(
|
||||
"input_color_space",
|
||||
options=["sRGB", "HDR", "linear"],
|
||||
default="sRGB",
|
||||
advanced=True,
|
||||
tooltip=(
|
||||
"Colorspace of the input tensor. The EXR is "
|
||||
"always written as scene-linear in the matching "
|
||||
"gamut.\n"
|
||||
" 'sRGB' — input is sRGB-encoded Rec.709; "
|
||||
"the inverse sRGB EOTF is applied.\n"
|
||||
" 'HDR' — input is HLG-encoded Rec.2020 "
|
||||
"(BT.2100); the inverse HLG OETF is applied "
|
||||
"to get scene-linear light.\n"
|
||||
" 'linear' — input is already scene-linear "
|
||||
"(Rec.709 primaries); written through unchanged. "
|
||||
"Use this for renderer/compositor output."
|
||||
"Colorspace of the input tensor. The EXR is always written as scene-linear in the matching gamut.\n"
|
||||
"sRGB — input is sRGB-encoded Rec.709; the inverse sRGB EOTF is applied.\n"
|
||||
"HDR — input is HLG-encoded Rec.2020 (BT.2100); the inverse HLG OETF is applied to get scene-linear light.\n"
|
||||
"linear — input is already scene-linear (Rec.709 primaries); written through unchanged. Use this for renderer/compositor output."
|
||||
),
|
||||
),
|
||||
]),
|
||||
@@ -1200,6 +1186,7 @@ class SaveImageAdvanced(IO.ComfyNode):
|
||||
],
|
||||
hidden=[IO.Hidden.prompt, IO.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[IO.Image.Output(display_name="images")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -1237,7 +1224,7 @@ class SaveImageAdvanced(IO.ComfyNode):
|
||||
results.append({"filename": file, "subfolder": subfolder, "type": "output"})
|
||||
counter += 1
|
||||
|
||||
return IO.NodeOutput(ui={"images": results})
|
||||
return IO.NodeOutput(images, ui={"images": results})
|
||||
|
||||
|
||||
class ImagesExtension(ComfyExtension):
|
||||
|
||||
@@ -0,0 +1,77 @@
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
from comfy_extras.color_util import normalize_palette
|
||||
|
||||
|
||||
class BuildJsonPromptIdeogram(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
color_palette = io.Colors.Input(
|
||||
"color_palette",
|
||||
socketless=False,
|
||||
tooltip="Hex color codes that steer the image's dominant colors. Up to 16 entries.",
|
||||
)
|
||||
return io.Schema(
|
||||
node_id="BuildJsonPromptIdeogram",
|
||||
display_name="Build JSON Prompt (Ideogram)",
|
||||
category="text",
|
||||
description="Build a JSON prompt for the Ideogram 4 model.",
|
||||
inputs=[
|
||||
io.Array.Input("element", tooltip="Prompt elements from the node Create Bounding Boxes."),
|
||||
io.String.Input("high_level_description", multiline=True, default="",
|
||||
tooltip="Optional description of the image in one or two sentences. Strongly recommended."),
|
||||
io.String.Input("background", multiline=True, default="",
|
||||
tooltip="Mandatory description of the image background or environment."),
|
||||
io.DynamicCombo.Input("style", options=[
|
||||
io.DynamicCombo.Option("none", []),
|
||||
io.DynamicCombo.Option("photo", [io.String.Input("photo", default="", tooltip="Camera or lens details for photographic outputs (e.g. 35mm, f/1.4, bokeh).")]),
|
||||
io.DynamicCombo.Option("art_style", [io.String.Input("art_style", default="", tooltip="Art style description (e.g. flat vector illustration, bold outlines).")]),
|
||||
]),
|
||||
io.String.Input("aesthetics", default="", tooltip="Mandatory aesthetic keywords (e.g. moody, cinematic, desaturated)."),
|
||||
io.String.Input("lighting", default="", tooltip="Mandatory lighting description (e.g. golden hour, rim light, dramatic shadows)."),
|
||||
io.String.Input("medium", default="", tooltip="Mandatory medium type (e.g. photograph, illustration, 3d_render, painting, graphic_design). When style = photo, set to photograph."),
|
||||
color_palette,
|
||||
],
|
||||
outputs=[io.Dict.Output(display_name="prompt")],
|
||||
is_experimental=True,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, element, style, high_level_description="", background="",
|
||||
aesthetics="", lighting="", medium="", color_palette=None) -> io.NodeOutput:
|
||||
elements = element if isinstance(element, list) else []
|
||||
kind = style.get("style", "none") if isinstance(style, dict) else "none"
|
||||
photo = style.get("photo", "") if isinstance(style, dict) else ""
|
||||
art_style = style.get("art_style", "") if isinstance(style, dict) else ""
|
||||
palette = normalize_palette(color_palette or [])
|
||||
|
||||
caption: dict = {}
|
||||
if high_level_description.strip():
|
||||
caption["high_level_description"] = high_level_description
|
||||
if kind != "none":
|
||||
style_desc: dict = {"aesthetics": aesthetics, "lighting": lighting}
|
||||
if kind == "photo":
|
||||
style_desc["photo"] = photo
|
||||
style_desc["medium"] = medium
|
||||
else:
|
||||
style_desc["medium"] = medium
|
||||
style_desc["art_style"] = art_style
|
||||
if palette:
|
||||
style_desc["color_palette"] = palette
|
||||
caption["style_description"] = style_desc
|
||||
caption["compositional_deconstruction"] = {
|
||||
"background": background,
|
||||
"elements": elements,
|
||||
}
|
||||
return io.NodeOutput(caption)
|
||||
|
||||
|
||||
class JsonPromptExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [BuildJsonPromptIdeogram]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> JsonPromptExtension:
|
||||
return JsonPromptExtension()
|
||||
@@ -317,11 +317,74 @@ class PreviewPointCloud(IO.ComfyNode):
|
||||
)
|
||||
|
||||
|
||||
MESH_EXTENSIONS = {'.gltf', '.glb', '.obj', '.fbx', '.stl'}
|
||||
|
||||
|
||||
class Load3DAdvanced(IO.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
input_dir = os.path.join(folder_paths.get_input_directory(), "3d")
|
||||
os.makedirs(input_dir, exist_ok=True)
|
||||
|
||||
input_path = Path(input_dir)
|
||||
base_path = Path(folder_paths.get_input_directory())
|
||||
|
||||
files = [
|
||||
normalize_path(str(file_path.relative_to(base_path)))
|
||||
for file_path in input_path.rglob("*")
|
||||
if file_path.suffix.lower() in MESH_EXTENSIONS
|
||||
]
|
||||
return IO.Schema(
|
||||
node_id="Load3DAdvanced",
|
||||
display_name="Load 3D (Advanced)",
|
||||
category="3d",
|
||||
search_aliases=[
|
||||
"load mesh",
|
||||
"load gltf",
|
||||
"load glb",
|
||||
"load obj",
|
||||
"load fbx",
|
||||
"load stl",
|
||||
],
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
IO.Combo.Input("model_file", options=["none"] + sorted(files), upload=IO.UploadType.model),
|
||||
IO.Load3D.Input("viewport_state"),
|
||||
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 validate_inputs(cls, model_file, **kwargs) -> bool | str:
|
||||
if not model_file or model_file == "none":
|
||||
return True
|
||||
if not folder_paths.exists_annotated_filepath(model_file):
|
||||
return f"Invalid 3D model file: {model_file}"
|
||||
return True
|
||||
|
||||
@classmethod
|
||||
def execute(cls, model_file, viewport_state, width: int, height: int, **kwargs) -> IO.NodeOutput:
|
||||
file_3d = None
|
||||
if model_file and model_file != "none":
|
||||
file_3d = Types.File3D(folder_paths.get_annotated_filepath(model_file))
|
||||
model_3d_info = viewport_state.get('model_3d_info', [])
|
||||
return IO.NodeOutput(file_3d, model_3d_info, viewport_state['camera_info'], width, height)
|
||||
|
||||
|
||||
class Load3DExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[IO.ComfyNode]]:
|
||||
return [
|
||||
Load3D,
|
||||
Load3DAdvanced,
|
||||
Preview3D,
|
||||
Preview3DAdvanced,
|
||||
PreviewGaussianSplat,
|
||||
|
||||
@@ -89,7 +89,8 @@ class SwitchNode(io.ComfyNode):
|
||||
template = io.MatchType.Template("switch")
|
||||
return io.Schema(
|
||||
node_id="ComfySwitchNode",
|
||||
display_name="Switch",
|
||||
search_aliases=["if", "then", "switch", "conditional", "branch"],
|
||||
display_name="If/Else Switch",
|
||||
category="utilities/logic",
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
|
||||
@@ -337,6 +337,36 @@ class ModelMergeQwenImage(comfy_extras.nodes_model_merging.ModelMergeBlocks):
|
||||
|
||||
return {"required": arg_dict}
|
||||
|
||||
class ModelMergeKrea2(comfy_extras.nodes_model_merging.ModelMergeBlocks):
|
||||
CATEGORY = "model/merging/model specific"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
arg_dict = { "model1": ("MODEL",),
|
||||
"model2": ("MODEL",)}
|
||||
|
||||
argument = ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01})
|
||||
|
||||
arg_dict["first."] = argument
|
||||
arg_dict["tmlp."] = argument
|
||||
arg_dict["txtmlp."] = argument
|
||||
arg_dict["tproj."] = argument
|
||||
|
||||
for i in range(2):
|
||||
arg_dict["txtfusion.layerwise_blocks.{}.".format(i)] = argument
|
||||
|
||||
arg_dict["txtfusion.projector."] = argument
|
||||
|
||||
for i in range(2):
|
||||
arg_dict["txtfusion.refiner_blocks.{}.".format(i)] = argument
|
||||
|
||||
for i in range(28):
|
||||
arg_dict["blocks.{}.".format(i)] = argument
|
||||
|
||||
arg_dict["last."] = argument
|
||||
|
||||
return {"required": arg_dict}
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"ModelMergeSD1": ModelMergeSD1,
|
||||
"ModelMergeSD2": ModelMergeSD1, #SD1 and SD2 have the same blocks
|
||||
@@ -353,4 +383,5 @@ NODE_CLASS_MAPPINGS = {
|
||||
"ModelMergeCosmosPredict2_2B": ModelMergeCosmosPredict2_2B,
|
||||
"ModelMergeCosmosPredict2_14B": ModelMergeCosmosPredict2_14B,
|
||||
"ModelMergeQwenImage": ModelMergeQwenImage,
|
||||
"ModelMergeKrea2": ModelMergeKrea2,
|
||||
}
|
||||
|
||||
@@ -123,7 +123,8 @@ class PhotoMakerLoader(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="PhotoMakerLoader",
|
||||
category="experimental/photomaker",
|
||||
display_name="Load PhotoMaker Model",
|
||||
category="model/loaders",
|
||||
inputs=[
|
||||
io.Combo.Input("photomaker_model_name", options=folder_paths.get_filename_list("photomaker")),
|
||||
],
|
||||
@@ -149,7 +150,8 @@ class PhotoMakerEncode(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="PhotoMakerEncode",
|
||||
category="experimental/photomaker",
|
||||
display_name="PhotoMaker Encode",
|
||||
category="model/conditioning/photomaker",
|
||||
inputs=[
|
||||
io.Photomaker.Input("photomaker"),
|
||||
io.Image.Input("image"),
|
||||
|
||||
@@ -10,12 +10,11 @@ class String(io.ComfyNode):
|
||||
return io.Schema(
|
||||
node_id="PrimitiveString",
|
||||
search_aliases=["text", "string", "text box", "prompt"],
|
||||
display_name="Text String",
|
||||
display_name="Text String (DEPRECATED)",
|
||||
category="utilities/primitive",
|
||||
inputs=[
|
||||
io.String.Input("value"),
|
||||
],
|
||||
inputs=[io.String.Input("value")],
|
||||
outputs=[io.String.Output()],
|
||||
is_deprecated=True
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -29,12 +28,10 @@ class StringMultiline(io.ComfyNode):
|
||||
return io.Schema(
|
||||
node_id="PrimitiveStringMultiline",
|
||||
search_aliases=["text", "string", "text multiline", "string multiline", "text box", "prompt"],
|
||||
display_name="Text String (Multiline)",
|
||||
display_name="Input Text",
|
||||
category="utilities/primitive",
|
||||
essentials_category="Basics",
|
||||
inputs=[
|
||||
io.String.Input("value", multiline=True),
|
||||
],
|
||||
inputs=[io.String.Input("value", multiline=True)],
|
||||
outputs=[io.String.Output()],
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,33 @@
|
||||
import sys
|
||||
from typing_extensions import override
|
||||
|
||||
from comfy_api.latest import ComfyExtension, io
|
||||
|
||||
|
||||
class SeedNode(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="SeedNode",
|
||||
display_name="Seed",
|
||||
search_aliases=["seed", "random"],
|
||||
category="utilities",
|
||||
inputs=[
|
||||
io.Int.Input("seed", min=0, max=sys.maxsize, control_after_generate=io.ControlAfterGenerate.fixed),
|
||||
],
|
||||
outputs=[io.Int.Output(display_name="seed")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, seed: int) -> io.NodeOutput:
|
||||
return io.NodeOutput(seed)
|
||||
|
||||
|
||||
class SeedExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
return [SeedNode]
|
||||
|
||||
|
||||
async def comfy_entrypoint() -> SeedExtension:
|
||||
return SeedExtension()
|
||||
@@ -119,7 +119,7 @@ class StableCascade_SuperResolutionControlnet(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="StableCascade_SuperResolutionControlnet",
|
||||
category="experimental/stable_cascade",
|
||||
category="experimental/stable cascade",
|
||||
is_experimental=True,
|
||||
inputs=[
|
||||
io.Image.Input("image"),
|
||||
|
||||
@@ -440,6 +440,57 @@ class JsonExtractString(io.ComfyNode):
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return io.NodeOutput("")
|
||||
|
||||
|
||||
def _dump_json(value, indent):
|
||||
return json.dumps(value, ensure_ascii=False, indent=indent or None)
|
||||
|
||||
|
||||
class ConvertDictionaryToString(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="ConvertDictionaryToString",
|
||||
display_name="Convert Dictionary to String",
|
||||
category="text",
|
||||
search_aliases=["json", "dict to json", "stringify", "serialize", "dict to string"],
|
||||
inputs=[
|
||||
io.Dict.Input("dictionary"),
|
||||
io.Int.Input("indent", default=2, min=0, max=8,
|
||||
tooltip="Spaces per indent level. 0 produces compact single-line string."),
|
||||
],
|
||||
outputs=[
|
||||
io.String.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, dictionary, indent=2):
|
||||
return io.NodeOutput(_dump_json(dictionary, indent))
|
||||
|
||||
|
||||
class ConvertArrayToString(io.ComfyNode):
|
||||
@classmethod
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="ConvertArrayToString",
|
||||
display_name="Convert Array to String",
|
||||
category="text",
|
||||
search_aliases=["json", "list to json", "stringify", "serialize", "list to string", "array to json"],
|
||||
inputs=[
|
||||
io.Array.Input("array"),
|
||||
io.Int.Input("indent", default=2, min=0, max=8,
|
||||
tooltip="Spaces per indent level. 0 produces compact single-line string."),
|
||||
],
|
||||
outputs=[
|
||||
io.String.Output(),
|
||||
],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def execute(cls, array, indent=2):
|
||||
return io.NodeOutput(_dump_json(array, indent))
|
||||
|
||||
|
||||
class StringExtension(ComfyExtension):
|
||||
@override
|
||||
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
||||
@@ -457,6 +508,8 @@ class StringExtension(ComfyExtension):
|
||||
RegexExtract,
|
||||
RegexReplace,
|
||||
JsonExtractString,
|
||||
ConvertDictionaryToString,
|
||||
ConvertArrayToString,
|
||||
]
|
||||
|
||||
async def comfy_entrypoint() -> StringExtension:
|
||||
|
||||
@@ -143,7 +143,7 @@ class VAEDecodeTripoSplat(IO.ComfyNode):
|
||||
return IO.Schema(
|
||||
node_id="VAEDecodeTripoSplat",
|
||||
display_name="TripoSplat Decode",
|
||||
category="3d/latent",
|
||||
category="model/latent/triposplat",
|
||||
description="Decode the sampled TripoSplat latent into a 3D gaussian splat. "
|
||||
"Modify the number of gaussians to vary the density.",
|
||||
inputs=[
|
||||
@@ -188,7 +188,7 @@ class TripoSplatSamplingPreview(IO.ComfyNode):
|
||||
return IO.Schema(
|
||||
node_id="TripoSplatSamplingPreview",
|
||||
display_name="TripoSplat Sampling Preview",
|
||||
category="3d/latent",
|
||||
category="model/latent/triposplat",
|
||||
description="Patch the TripoSplat model for the standard Ksampler node to show a live decoded "
|
||||
"gaussian splat preview at each step.",
|
||||
inputs=[
|
||||
|
||||
@@ -27,6 +27,7 @@ class SaveWEBM(io.ComfyNode):
|
||||
],
|
||||
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[io.Image.Output(display_name="images")]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -69,7 +70,7 @@ class SaveWEBM(io.ComfyNode):
|
||||
container.mux(stream.encode())
|
||||
container.close()
|
||||
|
||||
return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
|
||||
return io.NodeOutput(images, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
|
||||
|
||||
class SaveVideo(io.ComfyNode):
|
||||
@classmethod
|
||||
@@ -89,6 +90,7 @@ class SaveVideo(io.ComfyNode):
|
||||
],
|
||||
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
|
||||
is_output_node=True,
|
||||
outputs=[io.Video.Output("video")],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -117,7 +119,7 @@ class SaveVideo(io.ComfyNode):
|
||||
metadata=saved_metadata
|
||||
)
|
||||
|
||||
return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
|
||||
return io.NodeOutput(video, ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
|
||||
|
||||
|
||||
class CreateVideo(io.ComfyNode):
|
||||
@@ -233,13 +235,8 @@ class VideoSlice(io.ComfyNode):
|
||||
def define_schema(cls):
|
||||
return io.Schema(
|
||||
node_id="Video Slice",
|
||||
display_name="Video Slice",
|
||||
search_aliases=[
|
||||
"trim video duration",
|
||||
"skip first frames",
|
||||
"frame load cap",
|
||||
"start time",
|
||||
],
|
||||
display_name="Trim Video",
|
||||
search_aliases=["trim video duration", "skip first frames", "frame load cap", "start time"],
|
||||
category="video",
|
||||
essentials_category="Video Tools",
|
||||
inputs=[
|
||||
|
||||
+1
-1
@@ -1,3 +1,3 @@
|
||||
# This file is automatically generated by the build process when version is
|
||||
# updated in pyproject.toml.
|
||||
__version__ = "0.25.0"
|
||||
__version__ = "0.27.0"
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1308,6 +1308,25 @@ class PromptQueue:
|
||||
queued = copy.copy(self.queue)
|
||||
return (running, queued)
|
||||
|
||||
def interrupt_if_running(self, prompt_id):
|
||||
"""Interrupt the running prompt with this id, atomically.
|
||||
|
||||
Checks the live running set and signals the interrupt under the queue
|
||||
mutex, so the worker cannot move the job to done (and start the next
|
||||
prompt) in between. Returns True if a matching job was running and an
|
||||
interrupt was signalled, False otherwise. The atomicity is what keeps a
|
||||
cancel from landing on an unrelated prompt that started after a separate
|
||||
is-running check: the global interrupt flag is reset at the start of
|
||||
every prompt (execute_async), so a job that finishes before consuming
|
||||
the flag cannot leak the interrupt onto its successor.
|
||||
"""
|
||||
with self.mutex:
|
||||
for item in self.currently_running.values():
|
||||
if item[1] == prompt_id:
|
||||
nodes.interrupt_processing()
|
||||
return True
|
||||
return False
|
||||
|
||||
def get_tasks_remaining(self):
|
||||
with self.mutex:
|
||||
return len(self.queue) + len(self.currently_running)
|
||||
|
||||
@@ -8,21 +8,37 @@
|
||||
# # You can use is_default to mark that these folders should be listed first, and used as the default dirs for eg downloads
|
||||
# #is_default: true
|
||||
# checkpoints: models/checkpoints/
|
||||
# configs: models/configs/
|
||||
# loras: models/loras/
|
||||
# vae: models/vae/
|
||||
# text_encoders: |
|
||||
# models/text_encoders/
|
||||
# models/clip/ # legacy location still supported
|
||||
# clip_vision: models/clip_vision/
|
||||
# configs: models/configs/
|
||||
# controlnet: models/controlnet/
|
||||
# models/clip/
|
||||
# diffusion_models: |
|
||||
# models/diffusion_models
|
||||
# models/unet
|
||||
# models/unet/
|
||||
# models/diffusion_models/
|
||||
# clip_vision: models/clip_vision/
|
||||
# style_models: models/style_models/
|
||||
# embeddings: models/embeddings/
|
||||
# loras: models/loras/
|
||||
# diffusers: models/diffusers/
|
||||
# vae_approx: models/vae_approx/
|
||||
# controlnet: |
|
||||
# models/controlnet/
|
||||
# models/t2i_adapter/
|
||||
# gligen: models/gligen/
|
||||
# upscale_models: models/upscale_models/
|
||||
# vae: models/vae/
|
||||
# audio_encoders: models/audio_encoders/
|
||||
# latent_upscale_models: models/latent_upscale_models/
|
||||
# custom_nodes: custom_nodes/
|
||||
# hypernetworks: models/hypernetworks/
|
||||
# photomaker: models/photomaker/
|
||||
# classifiers: models/classifiers/
|
||||
# model_patches: models/model_patches/
|
||||
# audio_encoders: models/audio_encoders/
|
||||
# background_removal: models/background_removal/
|
||||
# frame_interpolation: models/frame_interpolation/
|
||||
# geometry_estimation: models/geometry_estimation/
|
||||
# optical_flow: models/optical_flow/
|
||||
# detection: models/detection/
|
||||
|
||||
|
||||
#config for a1111 ui
|
||||
@@ -45,8 +61,7 @@
|
||||
# controlnet: models/ControlNet
|
||||
|
||||
|
||||
# For a full list of supported keys (style_models, vae_approx, hypernetworks, photomaker,
|
||||
# model_patches, audio_encoders, classifiers, etc.) see folder_paths.py.
|
||||
# For the canonical list of supported keys and extensions, see folder_paths.py.
|
||||
|
||||
#other_ui:
|
||||
# base_path: path/to/ui
|
||||
|
||||
+63
-2
@@ -264,6 +264,59 @@ def annotated_filepath(name: str) -> tuple[str, str | None]:
|
||||
return name, base_dir
|
||||
|
||||
|
||||
# Content types a browser may execute or render inline. File endpoints that
|
||||
# serve user-controlled content must force these to download (and ideally set
|
||||
# Content-Disposition: attachment) to avoid stored XSS. Centralised here so the
|
||||
# /view and /userdata handlers can't drift apart. mimetypes.guess_type may
|
||||
# return either the text/* or application/* spelling depending on platform, so
|
||||
# both are listed.
|
||||
DANGEROUS_CONTENT_TYPES = {
|
||||
'text/html', 'text/html-sandboxed', 'application/xhtml+xml',
|
||||
'text/javascript', 'application/javascript', 'application/x-javascript',
|
||||
'application/ecmascript', 'text/css',
|
||||
'image/svg+xml', 'application/xml', 'text/xml',
|
||||
# message/rfc822 (.mht/.mhtml) can carry script in some browsers.
|
||||
'message/rfc822',
|
||||
}
|
||||
|
||||
|
||||
def is_dangerous_content_type(content_type: str | None) -> bool:
|
||||
"""Return True if a browser may execute or render `content_type` inline.
|
||||
|
||||
Normalises before matching so the check can't be slipped past with a
|
||||
charset/boundary parameter (``text/html; charset=utf-8``) or casing
|
||||
(``TEXT/HTML``). Any XML dialect (``*+xml`` or ``*/xml``) is treated as
|
||||
dangerous because XML can carry inline script via stylesheet/entity tricks,
|
||||
which also covers the ``application/{xslt,rss,atom,rdf}+xml`` family without
|
||||
enumerating each one. Endpoints serving user-controlled content should route
|
||||
a dangerous type to ``application/octet-stream`` + ``Content-Disposition:
|
||||
attachment`` + ``X-Content-Type-Options: nosniff``.
|
||||
"""
|
||||
if not content_type:
|
||||
return False
|
||||
normalized = content_type.split(';', 1)[0].strip().lower()
|
||||
if normalized in DANGEROUS_CONTENT_TYPES:
|
||||
return True
|
||||
return normalized.endswith('+xml') or normalized.endswith('/xml')
|
||||
|
||||
|
||||
def is_within_directory(directory: str, target: str) -> bool:
|
||||
"""Return True if `target` resolves to a path inside `directory`.
|
||||
|
||||
Uses realpath on both operands so that a symlink placed inside `directory`
|
||||
that points elsewhere cannot escape the containment check at open time.
|
||||
"""
|
||||
try:
|
||||
directory = os.path.realpath(directory)
|
||||
target = os.path.realpath(target)
|
||||
return os.path.commonpath((directory, target)) == directory
|
||||
except ValueError:
|
||||
# ValueError is raised by realpath() on a path with an embedded null
|
||||
# byte, and by commonpath() on Windows when the paths are on different
|
||||
# drives. In either case the target is not safely within the directory.
|
||||
return False
|
||||
|
||||
|
||||
def get_annotated_filepath(name: str, default_dir: str | None=None) -> str:
|
||||
name, base_dir = annotated_filepath(name)
|
||||
|
||||
@@ -273,7 +326,12 @@ def get_annotated_filepath(name: str, default_dir: str | None=None) -> str:
|
||||
else:
|
||||
base_dir = get_input_directory() # fallback path
|
||||
|
||||
return os.path.join(base_dir, name)
|
||||
filepath = os.path.abspath(os.path.join(base_dir, name))
|
||||
# Prevent path traversal: the resolved path must stay within base_dir.
|
||||
# repr() the name in the message so a crafted value can't inject log lines.
|
||||
if not is_within_directory(base_dir, filepath):
|
||||
raise ValueError("Invalid file path: {!r}".format(name))
|
||||
return filepath
|
||||
|
||||
|
||||
def exists_annotated_filepath(name) -> bool:
|
||||
@@ -282,7 +340,10 @@ def exists_annotated_filepath(name) -> bool:
|
||||
if base_dir is None:
|
||||
base_dir = get_input_directory() # fallback path
|
||||
|
||||
filepath = os.path.join(base_dir, name)
|
||||
filepath = os.path.abspath(os.path.join(base_dir, name))
|
||||
# Treat traversal attempts as non-existent rather than probing the filesystem.
|
||||
if not is_within_directory(base_dir, filepath):
|
||||
return False
|
||||
return os.path.exists(filepath)
|
||||
|
||||
|
||||
|
||||
@@ -403,7 +403,7 @@ def prompt_worker(q, server_instance):
|
||||
hook_breaker_ac10a0.restore_functions()
|
||||
|
||||
if not asset_seeder.is_disabled():
|
||||
asset_seeder.enqueue_enrich(roots=("output",), compute_hashes=True)
|
||||
asset_seeder.enqueue_enrich(roots=("output",), compute_hashes=args.enable_asset_hashing)
|
||||
asset_seeder.resume()
|
||||
|
||||
|
||||
@@ -458,7 +458,7 @@ def setup_database():
|
||||
if dependencies_available():
|
||||
init_db()
|
||||
if args.enable_assets:
|
||||
if asset_seeder.start(roots=("models", "input", "output"), prune_first=True, compute_hashes=True):
|
||||
if asset_seeder.start(roots=("models", "input", "output"), prune_first=True, compute_hashes=args.enable_asset_hashing):
|
||||
logging.info("Background asset scan initiated for models, input, output")
|
||||
except Exception as e:
|
||||
if "database is locked" in str(e):
|
||||
@@ -557,8 +557,13 @@ if __name__ == "__main__":
|
||||
logging.warning("WARNING: You are using a python version older than 3.10, please upgrade to a newer one. 3.12 and above is recommended.")
|
||||
|
||||
if args.disable_dynamic_vram:
|
||||
logging.warning("Dynamic vram disabled with argument. If you have any issues with dynamic vram enabled please give us a detailed reports as this argument will be removed soon.")
|
||||
|
||||
logging.warning(
|
||||
"Dynamic vram disabled with argument. If you have any issues with "
|
||||
"dynamic vram enabled please give us a detailed reports as this "
|
||||
"argument will be removed soon. If you use gguf we recommend keeping "
|
||||
"dynamic vram enabled and using native ComfyUI model formats instead. "
|
||||
"ComfyUI native formats like fp8 will be faster even if they are larger than your memory."
|
||||
)
|
||||
event_loop, _, start_all_func = start_comfyui()
|
||||
try:
|
||||
x = start_all_func()
|
||||
|
||||
@@ -20,8 +20,6 @@ from PIL.PngImagePlugin import PngInfo
|
||||
import numpy as np
|
||||
import safetensors.torch
|
||||
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
|
||||
|
||||
import comfy.diffusers_load
|
||||
import comfy.samplers
|
||||
import comfy.sample
|
||||
@@ -161,6 +159,29 @@ class ConditioningConcat:
|
||||
|
||||
return (out, )
|
||||
|
||||
class ConditioningMultiply:
|
||||
SEARCH_ALIASES = ["scale conditioning", "scale prompt", "multiply conditioning", "multiply prompt"]
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"conditioning": ("CONDITIONING", ),
|
||||
"multiplier": ("FLOAT", {"default": 1.0, "min": -100.0, "max": 100.0, "step": 0.01})
|
||||
}}
|
||||
RETURN_TYPES = ("CONDITIONING",)
|
||||
FUNCTION = "multiply"
|
||||
CATEGORY = "model/conditioning/transform"
|
||||
|
||||
def multiply(self, conditioning, multiplier):
|
||||
c = []
|
||||
for t in conditioning:
|
||||
values = {}
|
||||
pooled_output = t[1].get("pooled_output", None)
|
||||
if pooled_output is not None:
|
||||
values["pooled_output"] = pooled_output * multiplier
|
||||
scaled = node_helpers.conditioning_set_values([[t[0] * multiplier, t[1]]], values)[0]
|
||||
c.append(scaled)
|
||||
return (c,)
|
||||
|
||||
class ConditioningSetArea:
|
||||
SEARCH_ALIASES = ["regional prompt", "area prompt", "spatial conditioning", "localized prompt"]
|
||||
|
||||
@@ -328,7 +349,7 @@ class VAEDecodeTiled:
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
FUNCTION = "decode"
|
||||
|
||||
CATEGORY = "experimental"
|
||||
CATEGORY = "model/latent"
|
||||
|
||||
def decode(self, vae, samples, tile_size, overlap=64, temporal_size=64, temporal_overlap=8):
|
||||
if tile_size < overlap * 4:
|
||||
@@ -375,7 +396,7 @@ class VAEEncodeTiled:
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
FUNCTION = "encode"
|
||||
|
||||
CATEGORY = "experimental"
|
||||
CATEGORY = "model/latent"
|
||||
|
||||
def encode(self, vae, pixels, tile_size, overlap, temporal_size=64, temporal_overlap=8):
|
||||
t = vae.encode_tiled(pixels, tile_x=tile_size, tile_y=tile_size, overlap=overlap, tile_t=temporal_size, overlap_t=temporal_overlap)
|
||||
@@ -482,16 +503,18 @@ class SaveLatent:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {"required": { "samples": ("LATENT", ),
|
||||
"filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
RETURN_TYPES = ()
|
||||
return { "required": {
|
||||
"samples": ("LATENT",),
|
||||
"filename_prefix": ("STRING", {"default": "latents/ComfyUI"})},
|
||||
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
||||
}
|
||||
RETURN_TYPES = ("LATENT",)
|
||||
RETURN_NAMES = ("samples",)
|
||||
FUNCTION = "save"
|
||||
|
||||
OUTPUT_NODE = True
|
||||
|
||||
CATEGORY = "experimental"
|
||||
CATEGORY = "model/latent"
|
||||
|
||||
def save(self, samples, filename_prefix="ComfyUI", prompt=None, extra_pnginfo=None):
|
||||
full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir)
|
||||
@@ -524,7 +547,7 @@ class SaveLatent:
|
||||
output["latent_format_version_0"] = torch.tensor([])
|
||||
|
||||
comfy.utils.save_torch_file(output, file, metadata=metadata)
|
||||
return { "ui": { "latents": results } }
|
||||
return { "ui": { "latents": results }, "result": (samples,) }
|
||||
|
||||
|
||||
class LoadLatent:
|
||||
@@ -536,7 +559,7 @@ class LoadLatent:
|
||||
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.endswith(".latent")]
|
||||
return {"required": {"latent": [sorted(files), ]}, }
|
||||
|
||||
CATEGORY = "experimental"
|
||||
CATEGORY = "model/latent"
|
||||
|
||||
RETURN_TYPES = ("LATENT", )
|
||||
FUNCTION = "load"
|
||||
@@ -969,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"], ),
|
||||
"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"], ),
|
||||
},
|
||||
"optional": {
|
||||
"device": (["default", "cpu"], {"advanced": True}),
|
||||
@@ -1629,14 +1652,18 @@ class SaveImage:
|
||||
return {
|
||||
"required": {
|
||||
"images": ("IMAGE", {"tooltip": "The images to save."}),
|
||||
"filename_prefix": ("STRING", {"default": "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."})
|
||||
"filename_prefix": ("STRING", {
|
||||
"default": "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."
|
||||
})
|
||||
},
|
||||
"hidden": {
|
||||
"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ()
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("images",)
|
||||
FUNCTION = "save_images"
|
||||
|
||||
OUTPUT_NODE = True
|
||||
@@ -1672,7 +1699,7 @@ class SaveImage:
|
||||
})
|
||||
counter += 1
|
||||
|
||||
return { "ui": { "images": results } }
|
||||
return { "ui": { "images": results }, "result" : (images,) }
|
||||
|
||||
class PreviewImage(SaveImage):
|
||||
def __init__(self):
|
||||
@@ -2046,6 +2073,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"ConditioningAverage": ConditioningAverage,
|
||||
"ConditioningCombine": ConditioningCombine,
|
||||
"ConditioningConcat": ConditioningConcat,
|
||||
"ConditioningMultiply": ConditioningMultiply,
|
||||
"ConditioningSetArea": ConditioningSetArea,
|
||||
"ConditioningSetAreaPercentage": ConditioningSetAreaPercentage,
|
||||
"ConditioningSetAreaStrength": ConditioningSetAreaStrength,
|
||||
@@ -2117,6 +2145,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ConditioningAverage ": "Conditioning (Average)",
|
||||
"ConditioningAverage": "Conditioning (Average)",
|
||||
"ConditioningConcat": "Conditioning (Concat)",
|
||||
"ConditioningMultiply": "Conditioning (Multiply)",
|
||||
"ConditioningSetArea": "Conditioning (Set Area)",
|
||||
"ConditioningSetAreaPercentage": "Conditioning (Set Area with Percentage)",
|
||||
"ConditioningSetAreaStrength": "Conditioning (Set Area Strength)",
|
||||
@@ -2126,6 +2155,8 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"GLIGENTextBoxApply": "Apply GLIGEN Text Box",
|
||||
"ConditioningZeroOut": "Conditioning Zero Out",
|
||||
# Latent
|
||||
"LoadLatent": "Load Latent",
|
||||
"SaveLatent": "Save Latent",
|
||||
"VAEEncodeForInpaint": "VAE Encode (for Inpainting)",
|
||||
"SetLatentNoiseMask": "Set Latent Noise Mask",
|
||||
"VAEDecode": "VAE Decode",
|
||||
@@ -2160,7 +2191,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"ImageSharpen": "Sharpen Image",
|
||||
"ImageScaleToTotalPixels": "Scale Image to Total Pixels",
|
||||
"GetImageSize": "Get Image Size",
|
||||
# experimental
|
||||
"VAEDecodeTiled": "VAE Decode (Tiled)",
|
||||
"VAEEncodeTiled": "VAE Encode (Tiled)",
|
||||
}
|
||||
@@ -2299,6 +2329,9 @@ async def init_external_custom_nodes():
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# TODO: remove at some point when custom nodes don't break.
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy"))
|
||||
|
||||
base_node_names = set(NODE_CLASS_MAPPINGS.keys())
|
||||
node_paths = folder_paths.get_folder_paths("custom_nodes")
|
||||
node_import_times = []
|
||||
@@ -2367,6 +2400,8 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_images.py",
|
||||
"nodes_video_model.py",
|
||||
"nodes_ideogram4.py",
|
||||
"nodes_bounding_boxes.py",
|
||||
"nodes_json_prompt.py",
|
||||
"nodes_train.py",
|
||||
"nodes_dataset.py",
|
||||
"nodes_sag.py",
|
||||
@@ -2466,6 +2501,7 @@ async def init_builtin_extra_nodes():
|
||||
"nodes_gaussian_splat.py",
|
||||
"nodes_triposplat.py",
|
||||
"nodes_depth_anything_3.py",
|
||||
"nodes_seed.py",
|
||||
]
|
||||
|
||||
import_failed = []
|
||||
|
||||
+561
-8
@@ -55,6 +55,12 @@ components:
|
||||
description: URL for asset preview/thumbnail
|
||||
format: uri
|
||||
type: string
|
||||
short_url:
|
||||
description: Durable, owner-gated short link to this asset's content (relative `/api/s/{id}` path). Stable across the underlying signed URL's expiry — resolving it re-mints a fresh signed URL on every request — so it is safe to persist or share into chat, unlike `preview_url`. Only the minting user can resolve it. Omitted when the short-link surface is disabled or the asset has no resolvable content hash.
|
||||
nullable: true
|
||||
type: string
|
||||
x-runtime:
|
||||
- cloud
|
||||
size:
|
||||
description: Size of the asset in bytes
|
||||
format: int64
|
||||
@@ -182,6 +188,49 @@ components:
|
||||
- id
|
||||
- updated_at
|
||||
type: object
|
||||
AvailabilityStatusRequest:
|
||||
description: |
|
||||
Models to query — each entry is `model_id → URL`. The URL lets
|
||||
the server compute file_size + is_hf_downloadable on the same
|
||||
request, eliminating the need for a separate metadata endpoint.
|
||||
properties:
|
||||
models:
|
||||
additionalProperties:
|
||||
type: string
|
||||
description: model_id → URL declared in the workflow.
|
||||
type: object
|
||||
required:
|
||||
- models
|
||||
type: object
|
||||
AvailabilityStatusResponse:
|
||||
description: Per-model state + metadata + HF auth snapshot.
|
||||
properties:
|
||||
hf_auth:
|
||||
$ref: '#/components/schemas/HfAuthStatus'
|
||||
models:
|
||||
additionalProperties:
|
||||
$ref: '#/components/schemas/ModelStatusEntry'
|
||||
type: object
|
||||
required:
|
||||
- models
|
||||
- hf_auth
|
||||
type: object
|
||||
CancelDownloadSessionRequest:
|
||||
description: Request to cancel an in-flight download for a given model_id.
|
||||
properties:
|
||||
model_id:
|
||||
type: string
|
||||
required:
|
||||
- model_id
|
||||
type: object
|
||||
CancelDownloadSessionResponse:
|
||||
description: Result of a cancellation request.
|
||||
properties:
|
||||
cancelled:
|
||||
type: boolean
|
||||
required:
|
||||
- cancelled
|
||||
type: object
|
||||
CreateWorkflowRequest:
|
||||
description: Request body for creating a new saved workflow.
|
||||
properties:
|
||||
@@ -224,6 +273,51 @@ components:
|
||||
- base_version
|
||||
- workflow_json
|
||||
type: object
|
||||
DownloadModelsRequest:
|
||||
description: Map of model_id → URL of files to fetch into the model folders.
|
||||
properties:
|
||||
models:
|
||||
additionalProperties:
|
||||
type: string
|
||||
description: model_id → URL of models to download.
|
||||
type: object
|
||||
required:
|
||||
- models
|
||||
type: object
|
||||
DownloadModelsResponse:
|
||||
description: Acknowledgement that downloads have been scheduled.
|
||||
properties:
|
||||
accepted:
|
||||
description: Always true; the request was scheduled.
|
||||
type: boolean
|
||||
scheduled:
|
||||
description: The list of model_ids whose downloads are now in-flight.
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
required:
|
||||
- accepted
|
||||
- scheduled
|
||||
type: object
|
||||
DownloadProgress:
|
||||
description: In-flight download progress; embedded in ModelStatusEntry.
|
||||
properties:
|
||||
bytes_downloaded:
|
||||
format: int64
|
||||
type: integer
|
||||
progress:
|
||||
description: Fraction in [0,1]; null until total_bytes is known.
|
||||
format: float
|
||||
nullable: true
|
||||
type: number
|
||||
total_bytes:
|
||||
description: Content-Length when supplied by the source.
|
||||
format: int64
|
||||
nullable: true
|
||||
type: integer
|
||||
required:
|
||||
- bytes_downloaded
|
||||
type: object
|
||||
ErrorResponse:
|
||||
description: Standard error response with a machine-readable code and human-readable message.
|
||||
properties:
|
||||
@@ -388,6 +482,46 @@ components:
|
||||
- name
|
||||
- info
|
||||
type: object
|
||||
HfAuthLoginStartResponse:
|
||||
description: URL the frontend should open in a new tab to complete login.
|
||||
properties:
|
||||
authorize_url:
|
||||
type: string
|
||||
required:
|
||||
- authorize_url
|
||||
type: object
|
||||
HfAuthLogoutResponse:
|
||||
description: Result of the logout call (always logged_out = true).
|
||||
properties:
|
||||
logged_out:
|
||||
type: boolean
|
||||
required:
|
||||
- logged_out
|
||||
type: object
|
||||
HfAuthStatus:
|
||||
description: Inline snapshot of the server's HuggingFace OAuth state.
|
||||
properties:
|
||||
eligible:
|
||||
description: True iff this deployment can surface interactive HF login.
|
||||
type: boolean
|
||||
token_available:
|
||||
description: True iff a token (possibly expired but refreshable) is stored.
|
||||
type: boolean
|
||||
required:
|
||||
- token_available
|
||||
- eligible
|
||||
type: object
|
||||
HfAuthTokenStatusResponse:
|
||||
description: Whether the server holds an HF OAuth token + resolved username.
|
||||
properties:
|
||||
token_available:
|
||||
type: boolean
|
||||
username:
|
||||
nullable: true
|
||||
type: string
|
||||
required:
|
||||
- token_available
|
||||
type: object
|
||||
HistoryDetailEntry:
|
||||
description: History entry with full prompt data
|
||||
properties:
|
||||
@@ -673,6 +807,35 @@ components:
|
||||
- created_at
|
||||
- updated_at
|
||||
type: object
|
||||
JobsCancelRequest:
|
||||
additionalProperties: false
|
||||
description: Request to cancel multiple jobs by ID.
|
||||
properties:
|
||||
job_ids:
|
||||
description: Job identifiers (UUIDs) to cancel.
|
||||
items:
|
||||
format: uuid
|
||||
type: string
|
||||
maxItems: 100
|
||||
minItems: 1
|
||||
type: array
|
||||
required:
|
||||
- job_ids
|
||||
type: object
|
||||
JobsCancelResponse:
|
||||
description: Response for POST /api/jobs/cancel.
|
||||
properties:
|
||||
cancelled:
|
||||
description: |
|
||||
Job IDs for which a cancel event was successfully dispatched by this
|
||||
call. Jobs already in a terminal or cancelling state are idempotently
|
||||
skipped and will not appear here.
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
required:
|
||||
- cancelled
|
||||
type: object
|
||||
JobsListResponse:
|
||||
description: Paginated list of jobs for the authenticated user.
|
||||
properties:
|
||||
@@ -763,6 +926,41 @@ components:
|
||||
- name
|
||||
- folders
|
||||
type: object
|
||||
ModelStatusEntry:
|
||||
description: Everything the UI needs to render one row of the model.
|
||||
properties:
|
||||
file_size:
|
||||
description: Bytes, when known. Cached server-side per URL.
|
||||
format: int64
|
||||
nullable: true
|
||||
type: integer
|
||||
is_hf_downloadable:
|
||||
description: |
|
||||
HuggingFace-only signal. True if the server can fetch this
|
||||
URL with its current auth state (public, or gated-with-access).
|
||||
False if gated and lacking access. Null for non-HF URLs and
|
||||
for HF URLs whose probe failed entirely.
|
||||
nullable: true
|
||||
type: boolean
|
||||
progress:
|
||||
type: object
|
||||
allOf:
|
||||
- $ref: '#/components/schemas/DownloadProgress'
|
||||
description: Present when `state == downloading`.
|
||||
nullable: true
|
||||
state:
|
||||
description: |
|
||||
`available` — file is on disk.
|
||||
`missing` — not on disk and no download in flight.
|
||||
`downloading` — server is currently fetching the file.
|
||||
enum:
|
||||
- available
|
||||
- missing
|
||||
- downloading
|
||||
type: string
|
||||
required:
|
||||
- state
|
||||
type: object
|
||||
NodeInfo:
|
||||
description: Metadata describing a single ComfyUI node type and its inputs/outputs.
|
||||
properties:
|
||||
@@ -1006,7 +1204,7 @@ components:
|
||||
description: If true, clear all pending jobs from the queue
|
||||
type: boolean
|
||||
delete:
|
||||
description: Array of PENDING job IDs to cancel
|
||||
description: Array of job IDs to cancel; pending and running jobs transition to cancelled
|
||||
items:
|
||||
type: string
|
||||
type: array
|
||||
@@ -1657,6 +1855,12 @@ paths:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unsupported media type
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Validation error (e.g., disallowed model_type tag)
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
@@ -1822,6 +2026,83 @@ paths:
|
||||
summary: Update asset metadata
|
||||
tags:
|
||||
- file
|
||||
/api/assets/{id}/content:
|
||||
get:
|
||||
description: |
|
||||
Returns the binary content of an asset by ID.
|
||||
|
||||
The contract is the same across runtimes — "GET this path and you
|
||||
receive the asset's bytes" — but the mechanism differs:
|
||||
- **Local ComfyUI** streams the bytes directly (`200`,
|
||||
`application/octet-stream`).
|
||||
- **Cloud** does not proxy large files; it responds `302` with a
|
||||
`Location` redirect to a short-lived signed storage URL. Clients that
|
||||
follow redirects (browsers, `fetch`/XHR, `<img>`/`<video>`) receive
|
||||
the bytes transparently.
|
||||
|
||||
Prefer this over the filename-addressed `/api/view` when you have an
|
||||
asset ID.
|
||||
operationId: getAssetContent
|
||||
parameters:
|
||||
- description: Asset ID
|
||||
in: path
|
||||
name: id
|
||||
required: true
|
||||
schema:
|
||||
type: string
|
||||
- description: |
|
||||
Content-Disposition for the response: `attachment` (download) or
|
||||
`inline` (render in browser). Defaults to `attachment`.
|
||||
in: query
|
||||
name: disposition
|
||||
schema:
|
||||
default: attachment
|
||||
enum:
|
||||
- inline
|
||||
- attachment
|
||||
type: string
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/octet-stream:
|
||||
schema:
|
||||
format: binary
|
||||
type: string
|
||||
description: Asset content stream (local runtime streams the bytes directly)
|
||||
"302":
|
||||
description: Redirect to a signed storage URL (cloud runtime)
|
||||
headers:
|
||||
Cache-Control:
|
||||
description: Private caching directive scoped to the signed URL lifetime
|
||||
schema:
|
||||
type: string
|
||||
Location:
|
||||
description: Short-lived signed URL to the asset content in storage
|
||||
schema:
|
||||
type: string
|
||||
Vary:
|
||||
description: Partitions any cached redirect by auth credentials so a private redirect is not reused across users
|
||||
schema:
|
||||
type: string
|
||||
"404":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Asset not found
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error
|
||||
security:
|
||||
- ApiKeyAuth: []
|
||||
- BearerAuth: []
|
||||
- CookieAuth: []
|
||||
summary: Get asset content
|
||||
tags:
|
||||
- file
|
||||
/api/assets/{id}/tags:
|
||||
delete:
|
||||
description: Removes one or more tags from an existing asset
|
||||
@@ -2025,6 +2306,12 @@ paths:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Source asset with given hash not found
|
||||
"422":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Validation error (e.g., disallowed model_type tag)
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
@@ -2226,6 +2513,72 @@ paths:
|
||||
summary: Get tag histogram for filtered assets
|
||||
tags:
|
||||
- file
|
||||
/api/cancel-model-download-session:
|
||||
post:
|
||||
description: |
|
||||
Cancels the download session for the given model_id. The worker
|
||||
observes the cancellation between chunks, removes its partial `.tmp`
|
||||
file, and exits without writing the destination path.
|
||||
operationId: postCancelModelDownloadSession
|
||||
requestBody:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/CancelDownloadSessionRequest'
|
||||
required: true
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/CancelDownloadSessionResponse'
|
||||
description: Session cancelled.
|
||||
"404":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: No active download for that model_id.
|
||||
summary: Cancel an in-flight server-side model download
|
||||
tags:
|
||||
- model
|
||||
/api/download-models:
|
||||
post:
|
||||
description: |
|
||||
Schedules downloads for every model_id in the request map. Returns
|
||||
immediately after validation; progress is observed via
|
||||
`/api/models-availability-status`. Fails atomically if any model
|
||||
is already on disk, already downloading, gated, or has a URL that
|
||||
is not on the server's allowlist (HuggingFace, Civitai, localhost).
|
||||
operationId: postDownloadModels
|
||||
requestBody:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/DownloadModelsRequest'
|
||||
required: true
|
||||
responses:
|
||||
"202":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/DownloadModelsResponse'
|
||||
description: Downloads accepted and scheduled.
|
||||
"400":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: One of the requested models is invalid, gated, or has a non-allowed URL.
|
||||
"409":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: One of the requested models is already on disk or downloading.
|
||||
summary: Start a server-side download of one or more models
|
||||
tags:
|
||||
- model
|
||||
/api/embeddings:
|
||||
get:
|
||||
description: Returns the list of text-encoder embeddings available on disk.
|
||||
@@ -2245,6 +2598,10 @@ paths:
|
||||
description: |
|
||||
Returns a list of model folders available in the system.
|
||||
This is an experimental endpoint that replaces the legacy /models endpoint.
|
||||
Each folder's name is the identifier to pass to /api/experiment/models/{folder}.
|
||||
Once the model_type migration is active the names are model_type folder_names
|
||||
(e.g. `ultralytics_bbox`); a folder with no folder_name mapping is returned by
|
||||
its directory path.
|
||||
operationId: getModelFolders
|
||||
responses:
|
||||
"200":
|
||||
@@ -2527,6 +2884,74 @@ paths:
|
||||
summary: Get a specific subgraph blueprint
|
||||
tags:
|
||||
- workflow
|
||||
/api/hf-auth-login-start:
|
||||
post:
|
||||
description: |
|
||||
Spawns a short-lived loopback callback server (port 41954) and
|
||||
returns the URL the frontend should open in a new tab. After the
|
||||
user grants consent, HF redirects back to the callback URL with
|
||||
an authorization code; the server exchanges that for a token and
|
||||
persists it. Subsequent `/api/hf-auth-token-status` calls will
|
||||
return `token_available: true`. Rejected with 403 if the
|
||||
deployment is not eligible (not loopback or in --multi-user mode);
|
||||
409 if another login attempt is already in progress.
|
||||
operationId: postHfAuthLoginStart
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HfAuthLoginStartResponse'
|
||||
description: Login flow started; `authorize_url` is ready.
|
||||
"403":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Deployment is not eligible for interactive HF login.
|
||||
"409":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Another login attempt is already in progress.
|
||||
summary: Begin a HuggingFace OAuth login flow
|
||||
tags:
|
||||
- model
|
||||
/api/hf-auth-logout:
|
||||
post:
|
||||
description: Clears the in-memory cache and removes the on-disk token file.
|
||||
operationId: postHfAuthLogout
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HfAuthLogoutResponse'
|
||||
description: Logged out (idempotent — succeeds even if no token was held).
|
||||
summary: Drop the stored HuggingFace OAuth token
|
||||
tags:
|
||||
- model
|
||||
/api/hf-auth-token-status:
|
||||
get:
|
||||
description: |
|
||||
Returns `token_available: true` when the server has a token
|
||||
in memory (or on disk) for HuggingFace, irrespective of whether
|
||||
the access_token is currently fresh — an expired one with a
|
||||
refresh_token still counts as "logged in" because we'll refresh
|
||||
transparently on next use. If a username is resolvable via
|
||||
`HfApi.whoami` we return that too, for the settings UI.
|
||||
operationId: getHfAuthTokenStatus
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/HfAuthTokenStatusResponse'
|
||||
description: Token status.
|
||||
summary: Whether the server holds a usable HuggingFace OAuth token
|
||||
tags:
|
||||
- model
|
||||
/api/history:
|
||||
post:
|
||||
deprecated: true
|
||||
@@ -2675,14 +3100,20 @@ paths:
|
||||
summary: Get internationalisation translation strings
|
||||
/api/interrupt:
|
||||
post:
|
||||
deprecated: true
|
||||
description: |
|
||||
Cancel all currently RUNNING jobs for the authenticated user.
|
||||
This will interrupt any job that is currently in 'in_progress' status.
|
||||
Note: This endpoint only affects running jobs. To cancel pending jobs, use /api/queue.
|
||||
Deprecated. Prefer the jobs-namespace cancel endpoints:
|
||||
POST /api/jobs/{job_id}/cancel for a single job, or
|
||||
POST /api/jobs/cancel to cancel jobs by ID.
|
||||
|
||||
Cancels the first active job for the authenticated user (the currently
|
||||
running job if there is one, otherwise the next pending job). Takes no
|
||||
body and cannot target a specific job — use the jobs-namespace endpoints
|
||||
for that.
|
||||
operationId: interruptJob
|
||||
responses:
|
||||
"200":
|
||||
description: Success - Job interrupted or no running job found
|
||||
description: Success - first active job cancelled, or no active job found
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
@@ -2695,7 +3126,7 @@ paths:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error
|
||||
summary: Interrupt currently running jobs
|
||||
summary: Interrupt the first active job
|
||||
tags:
|
||||
- queue
|
||||
/api/job/{job_id}/status:
|
||||
@@ -2869,6 +3300,17 @@ paths:
|
||||
schema:
|
||||
format: uuid
|
||||
type: string
|
||||
- description: |
|
||||
When present, each output item in the response receives a `short_url` field containing a short link for that asset. Omit this parameter (the default) to receive a response identical to the no-param baseline. The value selects the link's lifetime and auth model: use `ephemeral_tool_chain` for short-lived (≤5 minute) machine-to-machine handoffs — these are public bearer links where the link ID itself is the credential, so anyone holding the link can resolve it (intended for pasting into an agent/MCP tool chain); use `default` for durable (30 day) human-revisitable links, which are owner-gated and resolvable only by the authenticated owner. Links are always minted under the authenticated request owner's identity; the auth model is selected by the server and is never settable by the caller.
|
||||
in: query
|
||||
name: short_link
|
||||
schema:
|
||||
enum:
|
||||
- ephemeral_tool_chain
|
||||
- default
|
||||
type: string
|
||||
x-runtime:
|
||||
- cloud
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
@@ -2954,6 +3396,106 @@ paths:
|
||||
summary: Cancel a job
|
||||
tags:
|
||||
- workflow
|
||||
/api/jobs/cancel:
|
||||
post:
|
||||
description: |
|
||||
Cancel one or more jobs for the authenticated user in a single request.
|
||||
|
||||
State-agnostic: cancels both pending and running jobs (both transition to
|
||||
the cancelled state via the same mechanism as the single-job endpoint).
|
||||
|
||||
Idempotent per job: a job already in a terminal or cancelling state is a
|
||||
no-op and simply will not appear in the returned `cancelled` list.
|
||||
|
||||
Fail-fast on unknown IDs: if any provided job ID does not exist for this
|
||||
user, the request returns 404 and no jobs are cancelled. This surfaces
|
||||
bad IDs to the caller rather than silently dropping them.
|
||||
|
||||
This is the canonical batch-cancel endpoint. The delete operation on
|
||||
POST /api/queue is deprecated in favour of this.
|
||||
operationId: cancelJobs
|
||||
requestBody:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/JobsCancelRequest'
|
||||
required: true
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/JobsCancelResponse'
|
||||
description: Success - cancel requests dispatched (or jobs were already terminal)
|
||||
"400":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Bad Request - job_ids is missing, empty, exceeds the maximum count, or contains an invalid UUID
|
||||
"401":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Unauthorized - Authentication required
|
||||
"404":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: One or more job IDs not found for this user (no jobs cancelled)
|
||||
"500":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Internal server error - cancellation failed
|
||||
summary: Cancel multiple jobs
|
||||
tags:
|
||||
- workflow
|
||||
/api/models-availability-status:
|
||||
post:
|
||||
description: |
|
||||
Given a map of `{model_id: url}` (model_id is
|
||||
`<directory>/<filename>`), returns per-id state plus the
|
||||
metadata the UI needs to render the row:
|
||||
|
||||
- `state` — one of `available` / `missing` / `downloading`
|
||||
- `progress` — embedded when `state == downloading`
|
||||
- `file_size` — bytes (when known)
|
||||
- `is_hf_downloadable` — for HF URLs only: true if the
|
||||
server can currently fetch the file with its stored auth
|
||||
state, false if gated and lacking access, null otherwise
|
||||
|
||||
Designed for 1 Hz polling. `file_size` and the intrinsic
|
||||
"is this model gated" check are cached server-side per URL;
|
||||
`is_hf_downloadable` is recomputed per call so license
|
||||
acceptance and login/logout transitions show up within one
|
||||
poll interval without any client-side cache plumbing.
|
||||
operationId: postModelsAvailabilityStatus
|
||||
requestBody:
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/AvailabilityStatusRequest'
|
||||
required: true
|
||||
responses:
|
||||
"200":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/AvailabilityStatusResponse'
|
||||
description: Per-model status and metadata.
|
||||
"400":
|
||||
content:
|
||||
application/json:
|
||||
schema:
|
||||
$ref: '#/components/schemas/ErrorResponse'
|
||||
description: Malformed request body.
|
||||
summary: Unified per-model status + metadata for the polling UI
|
||||
tags:
|
||||
- model
|
||||
/api/node_replacements:
|
||||
get:
|
||||
description: |
|
||||
@@ -3104,9 +3646,18 @@ paths:
|
||||
tags:
|
||||
- queue
|
||||
post:
|
||||
deprecated: true
|
||||
description: |
|
||||
Cancel specific PENDING jobs by ID or clear all pending jobs in the queue.
|
||||
Note: This endpoint only affects pending jobs. To cancel running jobs, use /api/interrupt.
|
||||
Deprecated. Prefer the jobs-namespace cancel endpoints:
|
||||
POST /api/jobs/cancel for cancelling jobs by ID, and
|
||||
POST /api/jobs/{job_id}/cancel for a single job.
|
||||
|
||||
Cancel specific jobs by ID (the `delete` field) or clear all pending
|
||||
jobs in the queue (the `clear` field). Despite the `delete` naming, this
|
||||
does not delete anything — listed jobs transition to the cancelled state,
|
||||
and `delete` cancels both pending and running jobs (not pending-only as
|
||||
previously documented). Job-by-ID cancellation is superseded by
|
||||
POST /api/jobs/cancel; `clear` has no jobs-namespace replacement yet.
|
||||
operationId: manageQueue
|
||||
requestBody:
|
||||
content:
|
||||
@@ -4891,3 +5442,5 @@ tags:
|
||||
name: queue
|
||||
- description: Job lifecycle queries
|
||||
name: job
|
||||
- description: Server-side model availability and downloads
|
||||
name: model
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "ComfyUI"
|
||||
version = "0.25.0"
|
||||
version = "0.27.0"
|
||||
readme = "README.md"
|
||||
license = { file = "LICENSE" }
|
||||
requires-python = ">=3.10"
|
||||
|
||||
+7
-6
@@ -1,6 +1,6 @@
|
||||
comfyui-frontend-package==1.45.15
|
||||
comfyui-workflow-templates==0.10.0
|
||||
comfyui-embedded-docs==0.5.4
|
||||
comfyui-frontend-package==1.45.20
|
||||
comfyui-workflow-templates==0.11.2
|
||||
comfyui-embedded-docs==0.5.6
|
||||
torch
|
||||
torchsde
|
||||
torchvision
|
||||
@@ -9,6 +9,7 @@ numpy>=1.25.0
|
||||
einops
|
||||
transformers>=4.50.3
|
||||
tokenizers>=0.13.3
|
||||
huggingface_hub
|
||||
sentencepiece
|
||||
safetensors>=0.4.2
|
||||
aiohttp>=3.11.8
|
||||
@@ -22,7 +23,7 @@ alembic
|
||||
SQLAlchemy>=2.0.0
|
||||
filelock
|
||||
av>=16.0.0
|
||||
comfy-kitchen==0.2.10
|
||||
comfy-kitchen==0.2.16
|
||||
comfy-aimdo==0.4.10
|
||||
requests
|
||||
simpleeval>=1.0.0
|
||||
@@ -33,5 +34,5 @@ kornia>=0.7.1
|
||||
spandrel
|
||||
pydantic~=2.0
|
||||
pydantic-settings~=2.0
|
||||
PyOpenGL
|
||||
glfw
|
||||
PyOpenGL>=3.1.8
|
||||
comfy-angle
|
||||
|
||||
@@ -8,7 +8,15 @@ import time
|
||||
import nodes
|
||||
import folder_paths
|
||||
import execution
|
||||
from comfy_execution.jobs import JobStatus, get_job, get_all_jobs, validate_job_id
|
||||
from comfy_execution.jobs import (
|
||||
JobStatus,
|
||||
get_job,
|
||||
get_all_jobs,
|
||||
validate_job_id,
|
||||
cancel_job,
|
||||
CANCEL_PENDING,
|
||||
CANCEL_RUNNING,
|
||||
)
|
||||
import uuid
|
||||
import urllib
|
||||
import json
|
||||
@@ -39,6 +47,7 @@ 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.asset_management import resolve_hash_to_path
|
||||
from app.model_downloader.api.routes import register_routes as register_model_downloader_routes
|
||||
|
||||
from app.user_manager import UserManager
|
||||
from app.model_manager import ModelFileManager
|
||||
@@ -119,6 +128,7 @@ def create_cors_middleware(allowed_origin: str):
|
||||
|
||||
return cors_middleware
|
||||
|
||||
|
||||
def is_loopback(host):
|
||||
if host is None:
|
||||
return False
|
||||
@@ -248,6 +258,7 @@ class PromptServer():
|
||||
else:
|
||||
register_assets_routes(self.app)
|
||||
asset_seeder.disable()
|
||||
register_model_downloader_routes(self.app)
|
||||
routes = web.RouteTableDef()
|
||||
self.routes = routes
|
||||
self.last_node_id = None
|
||||
@@ -608,15 +619,30 @@ class PromptServer():
|
||||
or 'application/octet-stream'
|
||||
)
|
||||
|
||||
# For security, force certain mimetypes to download instead of display
|
||||
if content_type in {'text/html', 'text/html-sandboxed', 'application/xhtml+xml', 'text/javascript', 'text/css'}:
|
||||
content_type = 'application/octet-stream' # Forces download
|
||||
# For security, force renderable/active types (HTML, JS,
|
||||
# CSS, SVG, XML — anything that can carry inline <script>
|
||||
# and execute in the page origin) to download instead of
|
||||
# displaying inline, preventing stored XSS. The
|
||||
# attachment disposition is the load-bearing guard: a
|
||||
# bare filename= hint does not force a download per
|
||||
# RFC 6266, so we only attach it on the dangerous branch
|
||||
# to avoid breaking inline display of legitimate images.
|
||||
# Escape backslash/quote per RFC 6266 quoted-string so a
|
||||
# filename containing a double quote (which passes the
|
||||
# ".."/leading-slash filter above) can't break out of the
|
||||
# header's quoted-string and malform the disposition.
|
||||
safe_filename = filename.replace("\\", "\\\\").replace('"', '\\"')
|
||||
disposition = f"filename=\"{safe_filename}\""
|
||||
if folder_paths.is_dangerous_content_type(content_type):
|
||||
content_type = 'application/octet-stream'
|
||||
disposition = f"attachment; filename=\"{safe_filename}\""
|
||||
|
||||
return web.FileResponse(
|
||||
file,
|
||||
headers={
|
||||
"Content-Disposition": f"filename=\"{filename}\"",
|
||||
"Content-Type": content_type
|
||||
"Content-Disposition": disposition,
|
||||
"Content-Type": content_type,
|
||||
"X-Content-Type-Options": "nosniff"
|
||||
}
|
||||
)
|
||||
|
||||
@@ -899,6 +925,107 @@ class PromptServer():
|
||||
|
||||
return web.json_response(job)
|
||||
|
||||
def _cancel_job_by_id(job_id):
|
||||
"""Cancel a single job by id using the queue's existing mechanics.
|
||||
|
||||
Running jobs are interrupted (same mechanism as /interrupt); pending
|
||||
jobs are dequeued (same mechanism as /queue {"delete": [...]}).
|
||||
Already-finished or unknown ids are no-ops. State-agnostic.
|
||||
|
||||
Returns True when a cancel was actually dispatched (running or
|
||||
pending job), False when the call was a no-op (terminal/unknown id).
|
||||
"""
|
||||
running, queued = self.prompt_queue.get_current_queue()
|
||||
history = self.prompt_queue.get_history()
|
||||
|
||||
def interrupt(prompt_id):
|
||||
logging.info(f"Cancelling running prompt {prompt_id}")
|
||||
# Atomic: only interrupts if the job is still the one running,
|
||||
# so a cancel can't land on a prompt that started in the gap
|
||||
# since the snapshot above. Returns whether it actually fired.
|
||||
return self.prompt_queue.interrupt_if_running(prompt_id)
|
||||
|
||||
def dequeue(prompt_id):
|
||||
logging.info(f"Cancelling pending prompt {prompt_id}")
|
||||
return self.prompt_queue.delete_queue_item(lambda a: a[1] == prompt_id)
|
||||
|
||||
classification = cancel_job(job_id, running, queued, history, interrupt, dequeue)
|
||||
return classification in (CANCEL_RUNNING, CANCEL_PENDING)
|
||||
|
||||
@routes.post("/api/jobs/{job_id}/cancel")
|
||||
async def cancel_job_by_id(request):
|
||||
"""Cancel a single job by id, regardless of state.
|
||||
|
||||
Idempotent: cancelling a job that has already finished, or an id
|
||||
that is not known, returns 200 with {"cancelled": false} rather
|
||||
than an error.
|
||||
"""
|
||||
job_id = request.match_info.get("job_id", None)
|
||||
if not job_id:
|
||||
return web.json_response(
|
||||
{"error": "job_id is required"},
|
||||
status=400
|
||||
)
|
||||
|
||||
cancelled = _cancel_job_by_id(job_id)
|
||||
return web.json_response({"cancelled": cancelled})
|
||||
|
||||
@routes.post("/api/jobs/cancel")
|
||||
async def cancel_jobs_batch(request):
|
||||
"""Cancel a batch of jobs by id.
|
||||
|
||||
Body: {"job_ids": ["<uuid>", ...]}
|
||||
|
||||
Best-effort and idempotent: every well-formed id is cancelled if it
|
||||
is running or pending; ids that are already finished or unknown are
|
||||
no-ops, not errors. A batch of all no-ops still returns 200 with
|
||||
{"cancelled": false}. This matches the single-cancel endpoint and
|
||||
means "cancel all" still cancels the in-progress jobs even if some
|
||||
finished between the client's snapshot and the request. Malformed
|
||||
ids are still rejected up front with 400 (see below).
|
||||
"""
|
||||
try:
|
||||
json_data = await request.json()
|
||||
except json.JSONDecodeError:
|
||||
return web.json_response(
|
||||
{"error": "Request body must be valid JSON"},
|
||||
status=400
|
||||
)
|
||||
|
||||
job_ids = json_data.get("job_ids") if isinstance(json_data, dict) else None
|
||||
if not isinstance(job_ids, list):
|
||||
return web.json_response(
|
||||
{"error": "job_ids must be a list"},
|
||||
status=400
|
||||
)
|
||||
|
||||
# Validate that every element is a well-formed job id before doing
|
||||
# anything else. An unhashable element (e.g. a nested dict or list)
|
||||
# would cause a TypeError when used as a history dict key; a
|
||||
# non-string or non-UUID value is never a valid id. Reject early
|
||||
# with 400 rather than letting the classify loop raise 500.
|
||||
invalid_ids = []
|
||||
for jid in job_ids:
|
||||
try:
|
||||
validate_job_id(jid)
|
||||
except (ValueError, AttributeError):
|
||||
invalid_ids.append(jid if isinstance(jid, str) else repr(jid))
|
||||
if invalid_ids:
|
||||
return web.json_response(
|
||||
{"error": "job_ids contains invalid id(s)", "invalid_ids": invalid_ids},
|
||||
status=400,
|
||||
)
|
||||
|
||||
# Best-effort: cancel each id that is still running/pending; an id
|
||||
# that has finished or never existed is a no-op rather than a reason
|
||||
# to fail the whole batch.
|
||||
cancelled = False
|
||||
for jid in job_ids:
|
||||
if _cancel_job_by_id(jid):
|
||||
cancelled = True
|
||||
|
||||
return web.json_response({"cancelled": cancelled})
|
||||
|
||||
@routes.get("/history")
|
||||
async def get_history(request):
|
||||
max_items = request.rel_url.query.get("max_items", None)
|
||||
|
||||
@@ -0,0 +1,708 @@
|
||||
"""Unit tests for the HuggingFace auth subsystem.
|
||||
|
||||
Covers:
|
||||
- token store: save/load roundtrip, chmod 0600, atomic write, delete
|
||||
- eligibility under various CLI-arg combinations
|
||||
- URL parsing (huggingface.co host detection + repo_id extraction)
|
||||
- HF-aware gated_detection.probe_url (mocked auth_check)
|
||||
- HF auth routes (token status, login start with eligibility gate, logout)
|
||||
- PKCE primitives + authorize URL shape
|
||||
|
||||
The OAuth callback server itself isn't exercised end-to-end here — that
|
||||
requires a real HF server. We test the components (state checking,
|
||||
URL building, code-exchange request shape) instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import stat
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from app.model_downloader.api.routes import register_routes
|
||||
from app.model_downloader.hf_auth import oauth
|
||||
from app.model_downloader.hf_auth.auth_store import HF_AUTH_STORE, HfAuthStore
|
||||
from app.model_downloader.hf_auth.token_store import (
|
||||
EXPIRY_BUFFER_SECS,
|
||||
Token,
|
||||
delete_token,
|
||||
load_token,
|
||||
save_token,
|
||||
)
|
||||
from app.model_downloader.hf_url import is_hf_url, repo_id_from_url
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_user_dir(tmp_path):
|
||||
"""Redirect ``folder_paths.get_user_directory`` so the token file
|
||||
lands in an isolated tmp_path instead of the real user dir."""
|
||||
user_dir = tmp_path / "user"
|
||||
user_dir.mkdir()
|
||||
with patch("folder_paths.get_user_directory", return_value=str(user_dir)):
|
||||
yield user_dir
|
||||
|
||||
|
||||
def _token_file_path(user_dir) -> str:
|
||||
return os.path.join(user_dir, "__hf_auth", "hf_auth_token.json")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_auth_store():
|
||||
"""Wipe singleton state between tests: auth + probe caches."""
|
||||
from app.model_downloader import gated_detection
|
||||
|
||||
HF_AUTH_STORE._token = None
|
||||
HF_AUTH_STORE._loaded_from_disk = False
|
||||
gated_detection.clear_caches_for_tests()
|
||||
yield HF_AUTH_STORE
|
||||
HF_AUTH_STORE._token = None
|
||||
HF_AUTH_STORE._loaded_from_disk = False
|
||||
gated_detection.clear_caches_for_tests()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app(patched_user_dir, fresh_auth_store):
|
||||
app = web.Application()
|
||||
register_routes(app)
|
||||
return app
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# URL parsing
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_is_hf_url_recognises_huggingface_co():
|
||||
assert is_hf_url("https://huggingface.co/x/y/resolve/main/z.safetensors")
|
||||
assert is_hf_url("https://huggingface.co/abc")
|
||||
assert not is_hf_url("https://hf-mirror.com/x/y/resolve/main/z.safetensors")
|
||||
assert not is_hf_url("https://civitai.com/x.safetensors")
|
||||
|
||||
|
||||
def test_repo_id_from_url_extracts_org_and_repo():
|
||||
url = "https://huggingface.co/Lightricks/LTX-2.3-22b-IC-LoRA-HDR/resolve/main/x.safetensors"
|
||||
assert repo_id_from_url(url) == "Lightricks/LTX-2.3-22b-IC-LoRA-HDR"
|
||||
|
||||
|
||||
def test_repo_id_from_url_handles_nested_path():
|
||||
url = "https://huggingface.co/Comfy-Org/ltx-2.3/resolve/main/split_files/loras/x.safetensors"
|
||||
assert repo_id_from_url(url) == "Comfy-Org/ltx-2.3"
|
||||
|
||||
|
||||
def test_repo_id_from_url_returns_none_for_non_hf():
|
||||
assert repo_id_from_url("https://civitai.com/x.safetensors") is None
|
||||
|
||||
|
||||
def test_repo_id_from_url_returns_none_for_non_resolve_paths():
|
||||
assert repo_id_from_url("https://huggingface.co/org/repo/blob/main/x.safetensors") is None
|
||||
assert repo_id_from_url("https://huggingface.co/org") is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Token store
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_token_store_roundtrip(patched_user_dir):
|
||||
tok = Token(
|
||||
access_token="hf_abc",
|
||||
refresh_token="rf_def",
|
||||
expires_at=9999999999.0,
|
||||
scope="openid profile",
|
||||
)
|
||||
save_token(tok)
|
||||
loaded = load_token()
|
||||
assert loaded == tok
|
||||
|
||||
|
||||
def test_token_store_writes_0600(patched_user_dir):
|
||||
tok = Token(access_token="x", refresh_token=None, expires_at=0.0)
|
||||
save_token(tok)
|
||||
path = _token_file_path(patched_user_dir)
|
||||
mode = stat.S_IMODE(os.stat(path).st_mode)
|
||||
# On Windows we silently no-op chmod; allow either the intended
|
||||
# mode or whatever umask the OS gave us.
|
||||
if os.name == "posix":
|
||||
assert mode == 0o600
|
||||
|
||||
|
||||
def test_token_store_delete_removes_file(patched_user_dir):
|
||||
tok = Token(access_token="x", refresh_token=None, expires_at=0.0)
|
||||
save_token(tok)
|
||||
delete_token()
|
||||
path = _token_file_path(patched_user_dir)
|
||||
assert not os.path.exists(path)
|
||||
# Idempotent: second delete is fine.
|
||||
delete_token()
|
||||
|
||||
|
||||
def test_token_store_load_returns_none_for_missing_file(patched_user_dir):
|
||||
assert load_token() is None
|
||||
|
||||
|
||||
def test_token_store_load_returns_none_for_corrupt_file(patched_user_dir):
|
||||
path = _token_file_path(patched_user_dir)
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "w") as f:
|
||||
f.write("not json {")
|
||||
assert load_token() is None
|
||||
|
||||
|
||||
def test_token_is_valid_uses_buffer(patched_user_dir):
|
||||
import time
|
||||
|
||||
fresh = Token(access_token="x", refresh_token=None, expires_at=time.time() + 3600)
|
||||
nearly_expired = Token(
|
||||
access_token="x",
|
||||
refresh_token=None,
|
||||
expires_at=time.time() + EXPIRY_BUFFER_SECS - 1,
|
||||
)
|
||||
assert fresh.is_valid()
|
||||
assert not nearly_expired.is_valid()
|
||||
|
||||
|
||||
def test_token_is_valid_rejects_empty_access_token():
|
||||
import time
|
||||
|
||||
tok = Token(access_token="", refresh_token=None, expires_at=time.time() + 3600)
|
||||
assert not tok.is_valid()
|
||||
|
||||
|
||||
def test_token_is_valid_rejects_at_exact_buffer_boundary():
|
||||
import time
|
||||
|
||||
tok = Token(
|
||||
access_token="x",
|
||||
refresh_token=None,
|
||||
expires_at=time.time() + EXPIRY_BUFFER_SECS,
|
||||
)
|
||||
assert not tok.is_valid()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Auth store
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_auth_store_loads_lazily(patched_user_dir):
|
||||
tok = Token(access_token="x", refresh_token=None, expires_at=9999999999.0)
|
||||
save_token(tok)
|
||||
store = HfAuthStore()
|
||||
assert store.has_token()
|
||||
assert store.get_token_sync() == tok
|
||||
|
||||
|
||||
def test_auth_store_set_persists(patched_user_dir):
|
||||
store = HfAuthStore()
|
||||
tok = Token(access_token="x", refresh_token=None, expires_at=9999999999.0)
|
||||
store.set_token(tok)
|
||||
# Token is on disk now — a fresh store sees it.
|
||||
assert HfAuthStore().get_token_sync() == tok
|
||||
|
||||
|
||||
def test_auth_store_clear_removes_in_memory_and_on_disk(patched_user_dir):
|
||||
store = HfAuthStore()
|
||||
tok = Token(access_token="x", refresh_token=None, expires_at=9999999999.0)
|
||||
store.set_token(tok)
|
||||
store.clear()
|
||||
assert not store.has_token()
|
||||
assert HfAuthStore().get_token_sync() is None
|
||||
|
||||
|
||||
def test_auth_store_has_token_true_when_expired_but_refreshable(patched_user_dir):
|
||||
import time
|
||||
|
||||
store = HfAuthStore()
|
||||
expired = Token(
|
||||
access_token="old",
|
||||
refresh_token="rf",
|
||||
expires_at=time.time() - 100,
|
||||
)
|
||||
store.set_token(expired)
|
||||
assert store.has_token()
|
||||
assert not expired.is_valid()
|
||||
|
||||
|
||||
def test_auth_store_get_token_sync_returns_expired_without_refresh(patched_user_dir):
|
||||
import time
|
||||
|
||||
store = HfAuthStore()
|
||||
expired = Token(
|
||||
access_token="old",
|
||||
refresh_token=None,
|
||||
expires_at=time.time() - 100,
|
||||
)
|
||||
store.set_token(expired)
|
||||
assert store.get_token_sync() == expired
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_store_get_valid_returns_none_when_expired_without_refresh(
|
||||
patched_user_dir,
|
||||
):
|
||||
import time
|
||||
|
||||
store = HfAuthStore()
|
||||
expired = Token(
|
||||
access_token="old",
|
||||
refresh_token=None,
|
||||
expires_at=time.time() - 100,
|
||||
)
|
||||
store.set_token(expired)
|
||||
with patch(
|
||||
"app.model_downloader.hf_auth.oauth.refresh_access_token",
|
||||
new=AsyncMock(),
|
||||
) as refresh_mock:
|
||||
result = await store.get_valid_token()
|
||||
assert result is None
|
||||
refresh_mock.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_store_get_valid_returns_fresh_token(patched_user_dir):
|
||||
store = HfAuthStore()
|
||||
import time
|
||||
|
||||
tok = Token(access_token="x", refresh_token=None, expires_at=time.time() + 3600)
|
||||
store.set_token(tok)
|
||||
fetched = await store.get_valid_token()
|
||||
assert fetched == tok
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_store_get_valid_refresh_on_expired(patched_user_dir):
|
||||
store = HfAuthStore()
|
||||
import time
|
||||
|
||||
expired = Token(
|
||||
access_token="old",
|
||||
refresh_token="rf",
|
||||
expires_at=time.time() - 100,
|
||||
)
|
||||
store.set_token(expired)
|
||||
refreshed = Token(
|
||||
access_token="new",
|
||||
refresh_token="rf",
|
||||
expires_at=time.time() + 3600,
|
||||
)
|
||||
with patch(
|
||||
"app.model_downloader.hf_auth.oauth.refresh_access_token",
|
||||
new=AsyncMock(return_value=refreshed),
|
||||
):
|
||||
result = await store.get_valid_token()
|
||||
assert result == refreshed
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_store_get_valid_token_does_not_resurrect_after_logout(
|
||||
patched_user_dir,
|
||||
):
|
||||
"""A logout landing *during* an in-flight refresh must not be undone by
|
||||
the refresh writing the token back (the resurrection race)."""
|
||||
store = HfAuthStore()
|
||||
import time
|
||||
|
||||
expired = Token(
|
||||
access_token="old", refresh_token="rf", expires_at=time.time() - 100
|
||||
)
|
||||
store.set_token(expired)
|
||||
refreshed = Token(
|
||||
access_token="new", refresh_token="rf", expires_at=time.time() + 3600
|
||||
)
|
||||
|
||||
async def fake_refresh(_refresh_token):
|
||||
# Simulate the user clicking "Log out" while the refresh is in flight.
|
||||
store.clear()
|
||||
return refreshed
|
||||
|
||||
with patch(
|
||||
"app.model_downloader.hf_auth.oauth.refresh_access_token",
|
||||
new=fake_refresh,
|
||||
):
|
||||
result = await store.get_valid_token()
|
||||
|
||||
# The refresh result is discarded — logout wins, in memory and on disk.
|
||||
assert result is None
|
||||
assert not store.has_token()
|
||||
assert load_token() is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_auth_store_get_valid_returns_none_on_refresh_failure(patched_user_dir):
|
||||
store = HfAuthStore()
|
||||
import time
|
||||
|
||||
expired = Token(
|
||||
access_token="old",
|
||||
refresh_token="rf",
|
||||
expires_at=time.time() - 100,
|
||||
)
|
||||
store.set_token(expired)
|
||||
with patch(
|
||||
"app.model_downloader.hf_auth.oauth.refresh_access_token",
|
||||
new=AsyncMock(side_effect=RuntimeError("HF down")),
|
||||
):
|
||||
result = await store.get_valid_token()
|
||||
assert result is None
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Eligibility
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"listen,multi_user,expected",
|
||||
[
|
||||
("127.0.0.1", False, True),
|
||||
("127.0.0.1", True, False), # multi-user disables it
|
||||
("0.0.0.0", False, False), # bind-all is not loopback
|
||||
("0.0.0.0", True, False),
|
||||
("192.168.1.5", False, False), # LAN address
|
||||
("::1", False, True), # IPv6 loopback
|
||||
],
|
||||
)
|
||||
def test_eligibility(listen, multi_user, expected, monkeypatch):
|
||||
from app.model_downloader.hf_auth import eligibility
|
||||
from comfy.cli_args import args
|
||||
|
||||
monkeypatch.setattr(args, "listen", listen)
|
||||
monkeypatch.setattr(args, "multi_user", multi_user)
|
||||
assert eligibility.is_hf_auth_eligible() is expected
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# gated_detection HF probe
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_url_hf_public(fresh_auth_store):
|
||||
"""auth_check succeeds with no token → is_hf_downloadable = True."""
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
|
||||
url = "https://huggingface.co/public/repo/resolve/main/x.safetensors"
|
||||
with patch("app.model_downloader.gated_detection._auth_check_sync"), patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=1024),
|
||||
):
|
||||
result = await probe_url(url)
|
||||
assert result.is_hf_downloadable is True
|
||||
assert result.file_size == 1024
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_url_hf_gated_no_access(fresh_auth_store):
|
||||
"""auth_check raises GatedRepoError → is_hf_downloadable = False."""
|
||||
from huggingface_hub.errors import GatedRepoError
|
||||
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
|
||||
url = "https://huggingface.co/gated/repo/resolve/main/x.safetensors"
|
||||
fake_response = MagicMock(status_code=403)
|
||||
with patch(
|
||||
"app.model_downloader.gated_detection._auth_check_sync",
|
||||
side_effect=GatedRepoError("gated", response=fake_response),
|
||||
), patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=None),
|
||||
):
|
||||
result = await probe_url(url)
|
||||
assert result.is_hf_downloadable is False
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_url_non_hf_skips_auth_check():
|
||||
"""Non-HF URLs never call auth_check; is_hf_downloadable stays None."""
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
|
||||
url = "https://civitai.com/api/download/models/1.safetensors"
|
||||
with patch(
|
||||
"app.model_downloader.gated_detection._auth_check_sync",
|
||||
) as mocked, patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=2048),
|
||||
):
|
||||
result = await probe_url(url)
|
||||
assert result.is_hf_downloadable is None
|
||||
assert result.file_size == 2048
|
||||
mocked.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_gated_cached_across_calls(fresh_auth_store):
|
||||
"""Intrinsic ``is_gated`` should be determined exactly once per URL.
|
||||
|
||||
Subsequent ``probe_url`` calls for the same URL must not re-issue
|
||||
the null-token auth_check — that's the whole point of the cache."""
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
|
||||
url = "https://huggingface.co/public/repo/resolve/main/x.safetensors"
|
||||
with patch(
|
||||
"app.model_downloader.gated_detection._auth_check_sync"
|
||||
) as mocked, patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=1024),
|
||||
):
|
||||
await probe_url(url)
|
||||
await probe_url(url)
|
||||
await probe_url(url)
|
||||
# Three probe_url calls × public-only-needs-1-auth_check = 1 call total.
|
||||
assert mocked.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_size_cached_across_calls(fresh_auth_store):
|
||||
"""Once a successful HEAD lands, subsequent calls don't re-HEAD."""
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
|
||||
url = "https://huggingface.co/public/repo/resolve/main/x.safetensors"
|
||||
with patch(
|
||||
"app.model_downloader.gated_detection._auth_check_sync"
|
||||
), patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=2048),
|
||||
) as size_probe:
|
||||
r1 = await probe_url(url)
|
||||
r2 = await probe_url(url)
|
||||
assert r1.file_size == 2048
|
||||
assert r2.file_size == 2048
|
||||
assert size_probe.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_size_not_probed_for_gated_no_access(fresh_auth_store):
|
||||
"""When ``is_hf_downloadable`` is False we must NOT HEAD the URL —
|
||||
otherwise a 401-due-to-gating would land as a cached ``None`` that
|
||||
survives a later successful login."""
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
from huggingface_hub.errors import GatedRepoError
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
url = "https://huggingface.co/gated/repo/resolve/main/x.safetensors"
|
||||
fake_resp = MagicMock(status_code=403)
|
||||
with patch(
|
||||
"app.model_downloader.gated_detection._auth_check_sync",
|
||||
side_effect=GatedRepoError("gated", response=fake_resp),
|
||||
), patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=None),
|
||||
) as size_probe:
|
||||
result = await probe_url(url)
|
||||
assert result.is_hf_downloadable is False
|
||||
assert result.file_size is None
|
||||
assert size_probe.call_count == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_probe_url_passes_token_when_available(fresh_auth_store, patched_user_dir):
|
||||
"""For a gated URL, auth_check runs twice: once with token=None to
|
||||
determine the intrinsic ``is_gated`` flag (cached forever), and once
|
||||
with the stored access_token to determine ``is_hf_downloadable`` for
|
||||
the current user."""
|
||||
from app.model_downloader import gated_detection
|
||||
from app.model_downloader.gated_detection import probe_url
|
||||
from huggingface_hub.errors import GatedRepoError
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
gated_detection.clear_caches_for_tests()
|
||||
fresh_auth_store.set_token(Token(
|
||||
access_token="hf_test_token",
|
||||
refresh_token=None,
|
||||
expires_at=9999999999.0,
|
||||
))
|
||||
url = "https://huggingface.co/private/repo/resolve/main/x.safetensors"
|
||||
|
||||
fake_resp = MagicMock(status_code=403)
|
||||
|
||||
def fake_auth_check(repo_id, token):
|
||||
# Null-token call → repo is gated. Subsequent call with the real
|
||||
# token succeeds (user has access).
|
||||
if token is None:
|
||||
raise GatedRepoError("gated", response=fake_resp)
|
||||
|
||||
with patch(
|
||||
"app.model_downloader.gated_detection._auth_check_sync",
|
||||
side_effect=fake_auth_check,
|
||||
) as mocked, patch(
|
||||
"app.model_downloader.gated_detection._probe_size_once",
|
||||
new=AsyncMock(return_value=None),
|
||||
):
|
||||
result = await probe_url(url)
|
||||
|
||||
# is_hf_downloadable should be True (token-authed call succeeded).
|
||||
assert result.is_hf_downloadable is True
|
||||
# Two calls: (repo_id, None) then (repo_id, <token>).
|
||||
assert mocked.call_count == 2
|
||||
assert mocked.call_args_list[0].args == ("private/repo", None)
|
||||
assert mocked.call_args_list[1].args == ("private/repo", "hf_test_token")
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# OAuth primitives
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_make_pkce_returns_distinct_high_entropy_values():
|
||||
verifier1, challenge1, state1 = oauth._make_pkce()
|
||||
verifier2, challenge2, state2 = oauth._make_pkce()
|
||||
assert verifier1 != verifier2
|
||||
assert challenge1 != challenge2
|
||||
assert state1 != state2
|
||||
# Verifier should be at least 43 chars per PKCE spec.
|
||||
assert len(verifier1) >= 43
|
||||
|
||||
|
||||
def test_build_authorize_url_includes_pkce_and_state():
|
||||
url = oauth._build_authorize_url("challenge123", "state456")
|
||||
assert url.startswith(oauth.AUTHORIZE_URL)
|
||||
assert "client_id=" + oauth.HF_CLIENT_ID in url
|
||||
assert "code_challenge=challenge123" in url
|
||||
assert "code_challenge_method=S256" in url
|
||||
assert "state=state456" in url
|
||||
assert "response_type=code" in url
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Routes
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_token_status_empty(aiohttp_client, app):
|
||||
"""No token set → token_available=false, username=null."""
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.get("/api/hf-auth-token-status")
|
||||
assert resp.status == 200
|
||||
data = await resp.json()
|
||||
assert data == {"token_available": False, "username": None}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_token_status_with_token(
|
||||
aiohttp_client, app, fresh_auth_store, patched_user_dir
|
||||
):
|
||||
"""Token present, whoami works → username is returned."""
|
||||
fresh_auth_store.set_token(Token(
|
||||
access_token="x", refresh_token=None, expires_at=9999999999.0,
|
||||
))
|
||||
with patch(
|
||||
"app.model_downloader.api.routes._whoami_username",
|
||||
return_value="alice",
|
||||
):
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.get("/api/hf-auth-token-status")
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"token_available": True, "username": "alice"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_login_start_403_when_ineligible(aiohttp_client, app, monkeypatch):
|
||||
"""Not loopback / multi-user → 403."""
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.is_hf_auth_eligible",
|
||||
lambda: False,
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post("/api/hf-auth-login-start")
|
||||
assert resp.status == 403
|
||||
assert (await resp.json())["error"]["code"] == "HF_AUTH_NOT_ELIGIBLE"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_login_start_returns_authorize_url(aiohttp_client, app, monkeypatch):
|
||||
"""Eligible + first attempt → 200 with authorize_url."""
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.is_hf_auth_eligible",
|
||||
lambda: True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.start_login_flow",
|
||||
AsyncMock(return_value="https://huggingface.co/oauth/authorize?fake=1"),
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post("/api/hf-auth-login-start")
|
||||
assert resp.status == 200
|
||||
assert (await resp.json())["authorize_url"].startswith(
|
||||
"https://huggingface.co/oauth/authorize"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_login_start_409_when_in_progress(aiohttp_client, app, monkeypatch):
|
||||
"""Lock already held → 409."""
|
||||
from app.model_downloader.hf_auth.oauth import OAuthInProgressError
|
||||
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.is_hf_auth_eligible",
|
||||
lambda: True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.start_login_flow",
|
||||
AsyncMock(side_effect=OAuthInProgressError()),
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post("/api/hf-auth-login-start")
|
||||
assert resp.status == 409
|
||||
assert (await resp.json())["error"]["code"] == "HF_AUTH_IN_PROGRESS"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_login_start_503_when_callback_bind_fails(
|
||||
aiohttp_client, app, monkeypatch
|
||||
):
|
||||
"""Callback server failed to bind (e.g. port busy) → 503, not a dead URL."""
|
||||
from app.model_downloader.hf_auth.oauth import OAuthCallbackError
|
||||
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.is_hf_auth_eligible",
|
||||
lambda: True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.start_login_flow",
|
||||
AsyncMock(side_effect=OAuthCallbackError("could not bind callback port")),
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post("/api/hf-auth-login-start")
|
||||
assert resp.status == 503
|
||||
assert (await resp.json())["error"]["code"] == "HF_AUTH_START_FAILED"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_hf_auth_logout_clears_store(
|
||||
aiohttp_client, app, fresh_auth_store, patched_user_dir
|
||||
):
|
||||
fresh_auth_store.set_token(Token(
|
||||
access_token="x", refresh_token=None, expires_at=9999999999.0,
|
||||
))
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post("/api/hf-auth-logout")
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"logged_out": True}
|
||||
assert not fresh_auth_store.has_token()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_availability_includes_hf_auth_snapshot(aiohttp_client, app, monkeypatch):
|
||||
"""The availability response embeds {token_available, eligible}."""
|
||||
monkeypatch.setattr(
|
||||
"app.model_downloader.api.routes.is_hf_auth_eligible",
|
||||
lambda: True,
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/models-availability-status",
|
||||
json={"models": {}},
|
||||
)
|
||||
assert resp.status == 200
|
||||
data = await resp.json()
|
||||
assert "hf_auth" in data
|
||||
assert data["hf_auth"] == {"token_available": False, "eligible": True}
|
||||
@@ -0,0 +1,514 @@
|
||||
"""Unit tests for the server-side model download subsystem.
|
||||
|
||||
Covers the pieces that don't require talking to a real network:
|
||||
|
||||
- path parsing & allowlist (pure functions)
|
||||
- DownloadServer registry lifecycle (in-memory state)
|
||||
- API routes via aiohttp_client + folder_paths/probe_url patches
|
||||
|
||||
Streaming downloads themselves are exercised indirectly — the route-level
|
||||
tests stub out the network probe so we can verify the gating logic in
|
||||
``download_models`` without making real HTTP calls.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import patch, AsyncMock
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from app.model_downloader.allowlist import is_url_allowed
|
||||
from app.model_downloader.api.routes import register_routes
|
||||
from app.model_downloader.download_server import DownloadServer
|
||||
from app.model_downloader.gated_detection import MetadataProbeResult
|
||||
from app.model_downloader.paths import (
|
||||
InvalidModelId,
|
||||
parse_model_id,
|
||||
resolve_destination,
|
||||
resolve_existing,
|
||||
)
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Fixtures
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def model_root(tmp_path):
|
||||
"""A fake ``models/`` root with two registered folder types."""
|
||||
loras_dir = tmp_path / "loras"
|
||||
checkpoints_dir = tmp_path / "checkpoints"
|
||||
loras_dir.mkdir()
|
||||
checkpoints_dir.mkdir()
|
||||
return tmp_path, loras_dir, checkpoints_dir
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def patched_folder_paths(model_root):
|
||||
"""Point folder_paths at our fake roots for the duration of one test."""
|
||||
_root, loras_dir, checkpoints_dir = model_root
|
||||
mapping = {
|
||||
"loras": ([str(loras_dir)], {".safetensors"}),
|
||||
"checkpoints": ([str(checkpoints_dir)], {".safetensors"}),
|
||||
}
|
||||
with patch(
|
||||
"folder_paths.folder_names_and_paths", mapping
|
||||
), patch(
|
||||
"folder_paths.get_folder_paths",
|
||||
side_effect=lambda name: mapping.get(name, ([], set()))[0],
|
||||
):
|
||||
yield mapping
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fresh_download_server():
|
||||
"""Reset the module-level singleton between tests so registry state
|
||||
doesn't leak across tests sharing the singleton."""
|
||||
from app.model_downloader.download_server import DOWNLOAD_SERVER
|
||||
|
||||
DOWNLOAD_SERVER.reset_for_tests()
|
||||
yield DOWNLOAD_SERVER
|
||||
DOWNLOAD_SERVER.reset_for_tests()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app(patched_folder_paths, fresh_download_server):
|
||||
app = web.Application()
|
||||
register_routes(app)
|
||||
return app
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Pure helpers: allowlist + path parsing
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_allowlist_accepts_hf_safetensors():
|
||||
assert is_url_allowed("https://huggingface.co/x/y/resolve/main/z.safetensors")
|
||||
|
||||
|
||||
def test_allowlist_accepts_civitai_pth():
|
||||
assert is_url_allowed("https://civitai.com/api/download/models/123.pth")
|
||||
|
||||
|
||||
def test_allowlist_rejects_unknown_host():
|
||||
assert not is_url_allowed("https://example.com/x.safetensors")
|
||||
|
||||
|
||||
def test_allowlist_rejects_api_path_on_hf():
|
||||
# On an allowlisted host but not pointing at a model file.
|
||||
assert not is_url_allowed("https://huggingface.co/api/models")
|
||||
|
||||
|
||||
def test_allowlist_rejects_non_https_except_localhost():
|
||||
assert not is_url_allowed("http://huggingface.co/x/y.safetensors")
|
||||
assert is_url_allowed("http://localhost:8000/x.safetensors")
|
||||
|
||||
|
||||
def test_parse_model_id_valid(patched_folder_paths):
|
||||
assert parse_model_id("loras/foo.safetensors") == ("loras", "foo.safetensors")
|
||||
|
||||
|
||||
def test_parse_model_id_rejects_traversal(patched_folder_paths):
|
||||
with pytest.raises(InvalidModelId):
|
||||
parse_model_id("../etc/passwd")
|
||||
|
||||
|
||||
def test_parse_model_id_rejects_unknown_folder(patched_folder_paths):
|
||||
with pytest.raises(InvalidModelId):
|
||||
parse_model_id("nope/x.safetensors")
|
||||
|
||||
|
||||
def test_parse_model_id_rejects_double_slash(patched_folder_paths):
|
||||
with pytest.raises(InvalidModelId):
|
||||
parse_model_id("loras/sub/x.safetensors")
|
||||
|
||||
|
||||
def test_resolve_existing_returns_path_when_present(model_root, patched_folder_paths):
|
||||
_root, loras_dir, _ = model_root
|
||||
target = loras_dir / "foo.safetensors"
|
||||
target.write_bytes(b"x")
|
||||
assert resolve_existing("loras/foo.safetensors") == str(target)
|
||||
|
||||
|
||||
def test_resolve_existing_returns_none_when_absent(patched_folder_paths):
|
||||
assert resolve_existing("loras/missing.safetensors") is None
|
||||
|
||||
|
||||
def test_resolve_destination_returns_tmp_pair(model_root, patched_folder_paths):
|
||||
_root, loras_dir, _ = model_root
|
||||
final, tmp = resolve_destination("loras/foo.safetensors", epoch=7)
|
||||
assert final == str(loras_dir / "foo.safetensors")
|
||||
# Temp path embeds the session epoch (so cancel+retry can't collide on it)
|
||||
# and uses the subsystem-specific suffix the startup sweep matches.
|
||||
assert tmp == f"{final}.7.comfy-download.tmp"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# DownloadServer registry: lifecycle, races, cancellation epoch semantics
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
def test_register_is_exclusive():
|
||||
server = DownloadServer()
|
||||
s1 = server.try_register("loras/x.safetensors", "https://huggingface.co/a")
|
||||
s2 = server.try_register("loras/x.safetensors", "https://huggingface.co/b")
|
||||
assert s1 is not None
|
||||
assert s2 is None
|
||||
assert server.is_downloading("loras/x.safetensors")
|
||||
|
||||
|
||||
def test_cancel_removes_session():
|
||||
server = DownloadServer()
|
||||
server.try_register("loras/x.safetensors", "https://huggingface.co/a")
|
||||
assert server.cancel("loras/x.safetensors") is True
|
||||
assert not server.is_downloading("loras/x.safetensors")
|
||||
|
||||
|
||||
def test_cancel_returns_false_when_absent():
|
||||
server = DownloadServer()
|
||||
assert server.cancel("loras/never.safetensors") is False
|
||||
|
||||
|
||||
def test_finish_only_clears_matching_epoch():
|
||||
"""If a session is cancelled and a new one for the same id is
|
||||
registered, ``finish`` from the original worker must not evict the
|
||||
newer session."""
|
||||
server = DownloadServer()
|
||||
s_old = server.try_register("loras/x.safetensors", "u1")
|
||||
server.cancel("loras/x.safetensors")
|
||||
s_new = server.try_register("loras/x.safetensors", "u2")
|
||||
assert s_new is not None and s_new.epoch != s_old.epoch
|
||||
# Old worker's late finish() is a no-op:
|
||||
server.finish(s_old)
|
||||
assert server.is_downloading("loras/x.safetensors")
|
||||
server.finish(s_new)
|
||||
assert not server.is_downloading("loras/x.safetensors")
|
||||
|
||||
|
||||
def test_is_active_follows_cancellation():
|
||||
server = DownloadServer()
|
||||
s = server.try_register("loras/x.safetensors", "u")
|
||||
assert server.is_active(s)
|
||||
server.cancel("loras/x.safetensors")
|
||||
assert not server.is_active(s)
|
||||
|
||||
|
||||
def test_update_progress_tracks_fraction():
|
||||
server = DownloadServer()
|
||||
s = server.try_register("loras/x.safetensors", "u")
|
||||
server.update_progress(s, 50, 100)
|
||||
snap = server.snapshot()["loras/x.safetensors"]
|
||||
assert snap.bytes_downloaded == 50
|
||||
assert snap.total_bytes == 100
|
||||
assert snap.progress == 0.5
|
||||
|
||||
|
||||
def test_update_progress_with_unknown_total_keeps_progress_none():
|
||||
server = DownloadServer()
|
||||
s = server.try_register("loras/x.safetensors", "u")
|
||||
server.update_progress(s, 50, None)
|
||||
assert server.snapshot()["loras/x.safetensors"].progress is None
|
||||
|
||||
|
||||
def test_cleanup_orphan_tmp_files(model_root):
|
||||
"""Orphan temp left by a crashed download must be swept on first use,
|
||||
while unrelated *.tmp files in the model dir are left untouched."""
|
||||
_root, loras_dir, _ = model_root
|
||||
orphan = loras_dir / "stale.safetensors.3.comfy-download.tmp"
|
||||
orphan.write_bytes(b"partial")
|
||||
unrelated = loras_dir / "someothertool.tmp"
|
||||
unrelated.write_bytes(b"not ours")
|
||||
mapping = {"loras": ([str(loras_dir)], {".safetensors"})}
|
||||
with patch("folder_paths.folder_names_and_paths", mapping), patch(
|
||||
"folder_paths.get_folder_paths",
|
||||
side_effect=lambda name: mapping.get(name, ([], set()))[0],
|
||||
):
|
||||
server = DownloadServer()
|
||||
assert orphan.exists(), "sweep must not run at construction time"
|
||||
server.sweep_orphan_tmp_files()
|
||||
assert not orphan.exists()
|
||||
assert unrelated.exists(), "unrelated .tmp must not be swept"
|
||||
# Idempotent — a second call is a cheap no-op.
|
||||
server.sweep_orphan_tmp_files()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Route: POST /api/models-availability-status
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_availability_partitions_correctly(
|
||||
aiohttp_client, app, model_root, fresh_download_server
|
||||
):
|
||||
_root, loras_dir, _ = model_root
|
||||
(loras_dir / "present.safetensors").write_bytes(b"x")
|
||||
fresh_download_server.try_register(
|
||||
"loras/inflight.safetensors", "http://localhost:8000/x.safetensors"
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
|
||||
# Stub probes — we're testing state assignment, not network calls.
|
||||
with patch(
|
||||
"app.model_downloader.api.routes.probe_url",
|
||||
new=AsyncMock(return_value=MetadataProbeResult(
|
||||
file_size=None, is_hf_downloadable=None,
|
||||
)),
|
||||
):
|
||||
body = {
|
||||
"models": {
|
||||
"loras/present.safetensors": "http://localhost:8000/p.safetensors",
|
||||
"loras/missing.safetensors": "http://localhost:8000/m.safetensors",
|
||||
"loras/inflight.safetensors": "http://localhost:8000/x.safetensors",
|
||||
}
|
||||
}
|
||||
resp = await client.post("/api/models-availability-status", json=body)
|
||||
assert resp.status == 200
|
||||
data = await resp.json()
|
||||
models = data["models"]
|
||||
assert models["loras/present.safetensors"]["state"] == "available"
|
||||
assert models["loras/missing.safetensors"]["state"] == "missing"
|
||||
assert models["loras/inflight.safetensors"]["state"] == "downloading"
|
||||
assert "hf_auth" in data
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_availability_invalid_id_classified_as_missing(aiohttp_client, app):
|
||||
client = await aiohttp_client(app)
|
||||
with patch(
|
||||
"app.model_downloader.api.routes.probe_url",
|
||||
new=AsyncMock(return_value=MetadataProbeResult(
|
||||
file_size=None, is_hf_downloadable=None,
|
||||
)),
|
||||
):
|
||||
resp = await client.post(
|
||||
"/api/models-availability-status",
|
||||
json={"models": {"../etc/passwd": "http://localhost:8000/x.safetensors"}},
|
||||
)
|
||||
assert resp.status == 200
|
||||
data = await resp.json()
|
||||
assert data["models"]["../etc/passwd"]["state"] == "missing"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Route: POST /api/download-models — precondition gating
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_rejects_url_not_in_allowlist(aiohttp_client, app):
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={"models": {"loras/x.safetensors": "https://evil.com/x.safetensors"}},
|
||||
)
|
||||
assert resp.status == 400
|
||||
err = (await resp.json())["error"]
|
||||
assert err["code"] == "URL_NOT_ALLOWED"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_rejects_already_available(
|
||||
aiohttp_client, app, model_root
|
||||
):
|
||||
_root, loras_dir, _ = model_root
|
||||
(loras_dir / "x.safetensors").write_bytes(b"x")
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={"models": {
|
||||
"loras/x.safetensors": "https://huggingface.co/a/b/resolve/main/x.safetensors"
|
||||
}},
|
||||
)
|
||||
assert resp.status == 409
|
||||
assert (await resp.json())["error"]["code"] == "ALREADY_AVAILABLE"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_rejects_already_downloading(
|
||||
aiohttp_client, app, fresh_download_server
|
||||
):
|
||||
fresh_download_server.try_register(
|
||||
"loras/x.safetensors", "https://huggingface.co/u.safetensors"
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={"models": {
|
||||
"loras/x.safetensors": "https://huggingface.co/a/b/resolve/main/x.safetensors"
|
||||
}},
|
||||
)
|
||||
assert resp.status == 409
|
||||
assert (await resp.json())["error"]["code"] == "ALREADY_DOWNLOADING"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_rejects_gated_model(aiohttp_client, app):
|
||||
client = await aiohttp_client(app)
|
||||
with patch(
|
||||
"app.model_downloader.api.routes.probe_url",
|
||||
new=AsyncMock(return_value=MetadataProbeResult(file_size=None, is_hf_downloadable=False)),
|
||||
):
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={"models": {
|
||||
"loras/x.safetensors": "https://huggingface.co/g/r/resolve/main/x.safetensors"
|
||||
}},
|
||||
)
|
||||
assert resp.status == 400
|
||||
assert (await resp.json())["error"]["code"] == "MODEL_NOT_DOWNLOADABLE"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_rejects_invalid_model_id(aiohttp_client, app):
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={"models": {"../etc/passwd": "https://huggingface.co/x.safetensors"}},
|
||||
)
|
||||
assert resp.status == 400
|
||||
assert (await resp.json())["error"]["code"] == "INVALID_MODEL_ID"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_atomic_failure_does_not_register_partial(
|
||||
aiohttp_client, app, model_root, fresh_download_server
|
||||
):
|
||||
"""If one model in a batch fails, none get registered."""
|
||||
_root, loras_dir, _ = model_root
|
||||
(loras_dir / "already.safetensors").write_bytes(b"x")
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={
|
||||
"models": {
|
||||
"loras/already.safetensors":
|
||||
"https://huggingface.co/a/b/resolve/main/already.safetensors",
|
||||
"loras/new.safetensors":
|
||||
"https://huggingface.co/a/b/resolve/main/new.safetensors",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert resp.status == 409
|
||||
# The "new" model should not have been registered as part of the
|
||||
# failed batch.
|
||||
assert not fresh_download_server.is_downloading("loras/new.safetensors")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_download_schedules_when_all_preconditions_pass(
|
||||
aiohttp_client, app, fresh_download_server
|
||||
):
|
||||
"""Verify the precondition pass, registration pass, and async
|
||||
scheduling all wire up correctly. We patch the streamer to avoid
|
||||
real HTTP while still letting the route execute end-to-end."""
|
||||
started = asyncio.Event()
|
||||
finish_signal = asyncio.Event()
|
||||
|
||||
async def fake_stream(session):
|
||||
started.set()
|
||||
await finish_signal.wait()
|
||||
from app.model_downloader.download_server import DOWNLOAD_SERVER
|
||||
DOWNLOAD_SERVER.finish(session)
|
||||
return "/dev/null"
|
||||
|
||||
with patch(
|
||||
"app.model_downloader.api.routes.probe_url",
|
||||
new=AsyncMock(return_value=MetadataProbeResult(file_size=42, is_hf_downloadable=True)),
|
||||
), patch(
|
||||
"app.model_downloader.downloader.stream_to_disk", new=fake_stream
|
||||
):
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/download-models",
|
||||
json={"models": {
|
||||
"loras/new.safetensors":
|
||||
"https://huggingface.co/a/b/resolve/main/new.safetensors"
|
||||
}},
|
||||
)
|
||||
assert resp.status == 202
|
||||
body = await resp.json()
|
||||
assert body["accepted"] is True
|
||||
assert body["scheduled"] == ["loras/new.safetensors"]
|
||||
# Wait for the worker to actually start.
|
||||
await asyncio.wait_for(started.wait(), timeout=2.0)
|
||||
assert fresh_download_server.is_downloading("loras/new.safetensors")
|
||||
finish_signal.set()
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Route: POST /api/cancel-model-download-session
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_removes_active_session(
|
||||
aiohttp_client, app, fresh_download_server
|
||||
):
|
||||
fresh_download_server.try_register(
|
||||
"loras/x.safetensors", "https://huggingface.co/u.safetensors"
|
||||
)
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/cancel-model-download-session",
|
||||
json={"model_id": "loras/x.safetensors"},
|
||||
)
|
||||
assert resp.status == 200
|
||||
assert (await resp.json())["cancelled"] is True
|
||||
assert not fresh_download_server.is_downloading("loras/x.safetensors")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_returns_404_when_none(aiohttp_client, app):
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/cancel-model-download-session",
|
||||
json={"model_id": "loras/nothing.safetensors"},
|
||||
)
|
||||
assert resp.status == 404
|
||||
assert (await resp.json())["error"]["code"] == "NOT_DOWNLOADING"
|
||||
|
||||
|
||||
# --------------------------------------------------------------------------- #
|
||||
# Unified availability response embeds metadata per id
|
||||
# --------------------------------------------------------------------------- #
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_availability_embeds_metadata(aiohttp_client, app):
|
||||
"""``file_size`` + ``is_hf_downloadable`` come back on the same
|
||||
request as the state — no separate metadata endpoint."""
|
||||
results = {
|
||||
"https://huggingface.co/a/b/resolve/main/free.safetensors":
|
||||
MetadataProbeResult(file_size=1024, is_hf_downloadable=True),
|
||||
"https://huggingface.co/g/r/resolve/main/gated.safetensors":
|
||||
MetadataProbeResult(file_size=None, is_hf_downloadable=False),
|
||||
}
|
||||
|
||||
async def fake_probe(url):
|
||||
return results[url]
|
||||
|
||||
with patch(
|
||||
"app.model_downloader.api.routes.probe_url", new=fake_probe
|
||||
):
|
||||
client = await aiohttp_client(app)
|
||||
resp = await client.post(
|
||||
"/api/models-availability-status",
|
||||
json={
|
||||
"models": {
|
||||
"loras/free.safetensors":
|
||||
"https://huggingface.co/a/b/resolve/main/free.safetensors",
|
||||
"loras/gated.safetensors":
|
||||
"https://huggingface.co/g/r/resolve/main/gated.safetensors",
|
||||
}
|
||||
},
|
||||
)
|
||||
assert resp.status == 200
|
||||
models = (await resp.json())["models"]
|
||||
assert models["loras/free.safetensors"]["file_size"] == 1024
|
||||
assert models["loras/free.safetensors"]["is_hf_downloadable"] is True
|
||||
assert models["loras/gated.safetensors"]["file_size"] is None
|
||||
assert models["loras/gated.safetensors"]["is_hf_downloadable"] is False
|
||||
@@ -1,3 +1,5 @@
|
||||
import contextlib
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
@@ -9,6 +11,40 @@ import requests
|
||||
from helpers import get_asset_filename, trigger_sync_seed_assets
|
||||
|
||||
|
||||
def test_download_svg_forced_to_attachment(http: requests.Session, api_base: str):
|
||||
"""GHSA-779p-m5rp-r4h4 CISA-5 (sibling route): an uploaded SVG must never be
|
||||
served inline from GET /api/assets/{id}/content, or an inline <script> runs
|
||||
in the app origin (stored XSS). Even with disposition=inline requested, a
|
||||
dangerous content type must be forced to application/octet-stream +
|
||||
Content-Disposition: attachment + nosniff. Regression guard for the stale
|
||||
inline blocklist that previously omitted image/svg+xml and ignored the
|
||||
centralized folder_paths.is_dangerous_content_type check.
|
||||
"""
|
||||
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"]),
|
||||
"name": "evil.svg",
|
||||
}
|
||||
up = http.post(api_base + "/api/assets", files=files, data=form_data, timeout=120)
|
||||
body = up.json()
|
||||
assert up.status_code in (200, 201), body
|
||||
aid = body["id"]
|
||||
try:
|
||||
r = http.get(f"{api_base}/api/assets/{aid}/content?disposition=inline", timeout=120)
|
||||
r.content
|
||||
assert r.status_code == 200
|
||||
ct = r.headers.get("Content-Type", "").lower()
|
||||
cd = r.headers.get("Content-Disposition", "").lower()
|
||||
assert "svg" not in ct, f"SVG served with a renderable content type: {ct!r}"
|
||||
assert ct.startswith("application/octet-stream"), f"expected octet-stream, got {ct!r}"
|
||||
assert "attachment" in cd, f"inline disposition not overridden to attachment: {cd!r}"
|
||||
assert r.headers.get("X-Content-Type-Options", "").lower() == "nosniff"
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
http.delete(f"{api_base}/api/assets/{aid}", timeout=30)
|
||||
|
||||
|
||||
def test_download_attachment_and_inline(http: requests.Session, api_base: str, seeded_asset: dict):
|
||||
aid = seeded_asset["id"]
|
||||
|
||||
|
||||
@@ -228,6 +228,62 @@ class TestMixedPrecisionOps(unittest.TestCase):
|
||||
with self.assertRaises(KeyError):
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
|
||||
def test_int8_convrot_metadata_loads_into_params(self):
|
||||
"""ConvRot metadata must reach TensorWiseINT8Layout params."""
|
||||
torch.manual_seed(123)
|
||||
layer_quant_config = {
|
||||
"layer": {
|
||||
"format": "int8_tensorwise",
|
||||
"convrot": True,
|
||||
"convrot_groupsize": 256,
|
||||
}
|
||||
}
|
||||
weight = torch.randn(16, 256, dtype=torch.bfloat16)
|
||||
bias = torch.randn(16, dtype=torch.bfloat16)
|
||||
q_weight = QuantizedTensor.from_float(
|
||||
weight,
|
||||
"TensorWiseINT8Layout",
|
||||
per_channel=True,
|
||||
convrot=True,
|
||||
convrot_groupsize=256,
|
||||
)
|
||||
state_dict = {
|
||||
"layer.weight": q_weight._qdata,
|
||||
"layer.bias": bias,
|
||||
"layer.weight_scale": q_weight._params.scale,
|
||||
}
|
||||
|
||||
state_dict, _ = comfy.utils.convert_old_quants(
|
||||
state_dict,
|
||||
metadata={"_quantization_metadata": json.dumps({"layers": layer_quant_config})},
|
||||
)
|
||||
model = torch.nn.Module()
|
||||
model.layer = ops.mixed_precision_ops({}).Linear(256, 16, device="cpu", dtype=torch.bfloat16)
|
||||
model.load_state_dict(state_dict, strict=False)
|
||||
|
||||
self.assertIsInstance(model.layer.weight, QuantizedTensor)
|
||||
self.assertEqual(model.layer.weight._layout_cls, "TensorWiseINT8Layout")
|
||||
self.assertTrue(model.layer.weight._params.convrot)
|
||||
self.assertEqual(model.layer.weight._params.convrot_groupsize, 256)
|
||||
|
||||
input_tensor = torch.randn(4, 256, dtype=torch.bfloat16)
|
||||
loaded_out = model.layer(input_tensor)
|
||||
ref_out = torch.nn.functional.linear(input_tensor, q_weight, bias)
|
||||
self.assertTrue(torch.equal(loaded_out, ref_out))
|
||||
|
||||
fp16_input = input_tensor.to(torch.float16)
|
||||
loaded_fp16_out = model.layer(fp16_input)
|
||||
ref_fp16_out = torch.nn.functional.linear(
|
||||
fp16_input,
|
||||
q_weight.to(dtype=torch.float16),
|
||||
bias.to(dtype=torch.float16),
|
||||
)
|
||||
self.assertTrue(torch.equal(loaded_fp16_out, ref_fp16_out))
|
||||
|
||||
saved = model.state_dict()
|
||||
saved_conf = json.loads(saved["layer.comfy_quant"].numpy().tobytes())
|
||||
self.assertTrue(saved_conf["convrot"])
|
||||
self.assertEqual(saved_conf["convrot_groupsize"], 256)
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
@@ -53,8 +53,11 @@ def test_annotated_filepath():
|
||||
|
||||
def test_get_annotated_filepath():
|
||||
default_dir = "/default/dir"
|
||||
assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.join(default_dir, "test.txt")
|
||||
assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.join(folder_paths.get_output_directory(), "test.txt")
|
||||
# get_annotated_filepath now normalizes with os.path.abspath (part of the
|
||||
# GHSA-779p traversal hardening), so compare against the normalized form —
|
||||
# on Windows abspath also prepends the current drive letter.
|
||||
assert folder_paths.get_annotated_filepath("test.txt", default_dir) == os.path.abspath(os.path.join(default_dir, "test.txt"))
|
||||
assert folder_paths.get_annotated_filepath("test.txt [output]") == os.path.abspath(os.path.join(folder_paths.get_output_directory(), "test.txt"))
|
||||
|
||||
def test_add_model_folder_path_append(clear_folder_paths):
|
||||
folder_paths.add_model_folder_path("test_folder", "/default/path", is_default=True)
|
||||
|
||||
@@ -0,0 +1,453 @@
|
||||
"""Tests for the jobs-namespace cancel endpoints.
|
||||
|
||||
Covers both layers:
|
||||
|
||||
* the pure cancel helpers in ``comfy_execution.jobs``
|
||||
(``classify_job_for_cancel`` / ``cancel_job``), which hold the business
|
||||
logic of mapping a cancel onto interrupt-vs-dequeue, and
|
||||
|
||||
* the HTTP contract of ``POST /api/jobs/{job_id}/cancel`` and
|
||||
``POST /api/jobs/cancel`` (status codes, single-cancel idempotency, and
|
||||
best-effort batch cancellation that treats unknown/finished ids as no-ops
|
||||
while still rejecting malformed ids with 400).
|
||||
|
||||
The HTTP layer is exercised against a small aiohttp app whose handlers are a
|
||||
faithful copy of the wiring in ``server.py`` driven by a fake queue that
|
||||
mirrors ``execution.PromptQueue`` (``get_current_queue`` / ``get_history`` /
|
||||
``delete_queue_item``). This keeps the test free of the heavy ComfyUI runtime
|
||||
(torch, nodes, ...) while still testing the real cancel logic.
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
import pytest
|
||||
from aiohttp import web
|
||||
|
||||
from comfy_execution.jobs import (
|
||||
CANCEL_PENDING,
|
||||
CANCEL_RUNNING,
|
||||
CANCEL_TERMINAL,
|
||||
CANCEL_UNKNOWN,
|
||||
cancel_job,
|
||||
classify_job_for_cancel,
|
||||
validate_job_id,
|
||||
)
|
||||
|
||||
# Classifications for which a cancel was actually dispatched (vs a no-op).
|
||||
_CANCELLED = (CANCEL_RUNNING, CANCEL_PENDING)
|
||||
|
||||
# Canonical UUID ids for HTTP-layer tests (the batch endpoint validates UUID format).
|
||||
_UUID_A = "aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa"
|
||||
_UUID_B = "bbbbbbbb-bbbb-4bbb-bbbb-bbbbbbbbbbbb"
|
||||
_UUID_C = "cccccccc-cccc-4ccc-cccc-cccccccccccc"
|
||||
_UUID_D = "dddddddd-dddd-4ddd-dddd-dddddddddddd"
|
||||
_UUID_MISSING = "ffffffff-ffff-4fff-ffff-ffffffffffff"
|
||||
|
||||
|
||||
def make_queue_item(prompt_id, number=0):
|
||||
"""Build a queue tuple shaped like the real ones: index 1 is the id."""
|
||||
return (number, prompt_id, {}, {}, [])
|
||||
|
||||
|
||||
class FakePromptQueue:
|
||||
"""Minimal stand-in for execution.PromptQueue for the cancel paths.
|
||||
|
||||
Tracks interrupts and dequeues so tests can assert side effects.
|
||||
"""
|
||||
|
||||
def __init__(self, running=None, pending=None, history=None):
|
||||
self._running = list(running or [])
|
||||
self._pending = list(pending or [])
|
||||
self._history = dict(history or {})
|
||||
self.interrupt_count = 0
|
||||
|
||||
def get_current_queue(self):
|
||||
return (list(self._running), list(self._pending))
|
||||
|
||||
def get_history(self, prompt_id=None):
|
||||
if prompt_id is None:
|
||||
return dict(self._history)
|
||||
if prompt_id in self._history:
|
||||
return {prompt_id: self._history[prompt_id]}
|
||||
return {}
|
||||
|
||||
def delete_queue_item(self, function):
|
||||
for i, item in enumerate(self._pending):
|
||||
if function(item):
|
||||
self._pending.pop(i)
|
||||
return True
|
||||
return False
|
||||
|
||||
def interrupt_if_running(self, prompt_id):
|
||||
# Mirrors execution.PromptQueue.interrupt_if_running: only signals an
|
||||
# interrupt when the id is actually in the running set.
|
||||
if any(item[1] == prompt_id for item in self._running):
|
||||
self.interrupt_count += 1
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def build_app(queue):
|
||||
"""Build an aiohttp app exposing the cancel routes against ``queue``.
|
||||
|
||||
Handler bodies mirror server.py exactly.
|
||||
"""
|
||||
|
||||
def _cancel_job_by_id(job_id):
|
||||
running, pending = queue.get_current_queue()
|
||||
history = queue.get_history()
|
||||
|
||||
def interrupt(prompt_id):
|
||||
return queue.interrupt_if_running(prompt_id)
|
||||
|
||||
def dequeue(prompt_id):
|
||||
return queue.delete_queue_item(lambda a: a[1] == prompt_id)
|
||||
|
||||
classification = cancel_job(
|
||||
job_id, running, pending, history, interrupt, dequeue
|
||||
)
|
||||
return classification in _CANCELLED
|
||||
|
||||
async def cancel_job_by_id(request):
|
||||
job_id = request.match_info.get("job_id", None)
|
||||
if not job_id:
|
||||
return web.json_response({"error": "job_id is required"}, status=400)
|
||||
cancelled = _cancel_job_by_id(job_id)
|
||||
return web.json_response({"cancelled": cancelled})
|
||||
|
||||
async def cancel_jobs_batch(request):
|
||||
try:
|
||||
json_data = await request.json()
|
||||
except json.JSONDecodeError:
|
||||
return web.json_response(
|
||||
{"error": "Request body must be valid JSON"}, status=400
|
||||
)
|
||||
|
||||
job_ids = json_data.get("job_ids") if isinstance(json_data, dict) else None
|
||||
if not isinstance(job_ids, list):
|
||||
return web.json_response({"error": "job_ids must be a list"}, status=400)
|
||||
|
||||
invalid_ids = []
|
||||
for jid in job_ids:
|
||||
try:
|
||||
validate_job_id(jid)
|
||||
except (ValueError, AttributeError):
|
||||
invalid_ids.append(jid if isinstance(jid, str) else repr(jid))
|
||||
if invalid_ids:
|
||||
return web.json_response(
|
||||
{"error": "job_ids contains invalid id(s)", "invalid_ids": invalid_ids},
|
||||
status=400,
|
||||
)
|
||||
|
||||
cancelled = False
|
||||
for jid in job_ids:
|
||||
if _cancel_job_by_id(jid):
|
||||
cancelled = True
|
||||
return web.json_response({"cancelled": cancelled})
|
||||
|
||||
app = web.Application()
|
||||
app.router.add_post("/api/jobs/{job_id}/cancel", cancel_job_by_id)
|
||||
app.router.add_post("/api/jobs/cancel", cancel_jobs_batch)
|
||||
return app
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure helper tests: classification + cancel side effects
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestClassifyJobForCancel:
|
||||
def test_running(self):
|
||||
running = [make_queue_item("a")]
|
||||
assert classify_job_for_cancel("a", running, [], {}) == CANCEL_RUNNING
|
||||
|
||||
def test_pending(self):
|
||||
pending = [make_queue_item("b")]
|
||||
assert classify_job_for_cancel("b", [], pending, {}) == CANCEL_PENDING
|
||||
|
||||
def test_terminal(self):
|
||||
history = {"c": {"prompt": make_queue_item("c"), "outputs": {}, "status": {}}}
|
||||
assert classify_job_for_cancel("c", [], [], history) == CANCEL_TERMINAL
|
||||
|
||||
def test_unknown(self):
|
||||
assert classify_job_for_cancel("z", [], [], {}) == CANCEL_UNKNOWN
|
||||
|
||||
|
||||
class TestCancelJobHelper:
|
||||
"""``interrupt`` and ``dequeue`` both take the id and return whether they
|
||||
actually acted, so cancel_job's return reflects the real outcome."""
|
||||
|
||||
def test_running_is_interrupted_not_dequeued(self):
|
||||
interrupts = []
|
||||
dequeues = []
|
||||
result = cancel_job(
|
||||
"a", [make_queue_item("a")], [], {},
|
||||
interrupt=lambda pid: interrupts.append(pid) or True,
|
||||
dequeue=lambda pid: dequeues.append(pid) or True,
|
||||
)
|
||||
assert result == CANCEL_RUNNING
|
||||
assert interrupts == ["a"]
|
||||
assert dequeues == []
|
||||
|
||||
def test_pending_is_dequeued_not_interrupted(self):
|
||||
interrupts = []
|
||||
dequeues = []
|
||||
result = cancel_job(
|
||||
"b", [], [make_queue_item("b")], {},
|
||||
interrupt=lambda pid: interrupts.append(pid) or True,
|
||||
dequeue=lambda pid: dequeues.append(pid) or True,
|
||||
)
|
||||
assert result == CANCEL_PENDING
|
||||
assert dequeues == ["b"]
|
||||
assert interrupts == []
|
||||
|
||||
def test_terminal_is_noop(self):
|
||||
history = {"c": {"prompt": make_queue_item("c"), "outputs": {}, "status": {}}}
|
||||
interrupts = []
|
||||
dequeues = []
|
||||
result = cancel_job(
|
||||
"c", [], [], history,
|
||||
interrupt=lambda pid: interrupts.append(pid) or True,
|
||||
dequeue=lambda pid: dequeues.append(pid) or True,
|
||||
)
|
||||
assert result == CANCEL_TERMINAL
|
||||
assert interrupts == []
|
||||
assert dequeues == []
|
||||
|
||||
def test_unknown_is_noop(self):
|
||||
interrupts = []
|
||||
dequeues = []
|
||||
result = cancel_job(
|
||||
"z", [], [], {},
|
||||
interrupt=lambda pid: interrupts.append(pid) or True,
|
||||
dequeue=lambda pid: dequeues.append(pid) or True,
|
||||
)
|
||||
assert result == CANCEL_UNKNOWN
|
||||
assert interrupts == []
|
||||
assert dequeues == []
|
||||
|
||||
def test_running_but_finished_before_interrupt_returns_unknown(self):
|
||||
"""Classified RUNNING from a stale snapshot, but the job finished before
|
||||
the atomic interrupt fired (interrupt returns False). cancel_job reports
|
||||
UNKNOWN rather than claiming a cancel that did not happen — and the
|
||||
atomic interrupt guarantees no unrelated job was hit."""
|
||||
interrupts = []
|
||||
result = cancel_job(
|
||||
"a", [make_queue_item("a")], [], {},
|
||||
interrupt=lambda pid: interrupts.append(pid) or False,
|
||||
dequeue=lambda pid: True,
|
||||
)
|
||||
assert result == CANCEL_UNKNOWN
|
||||
assert interrupts == ["a"] # interrupt was attempted atomically
|
||||
|
||||
def test_pending_started_running_is_interrupted(self):
|
||||
"""Pending->running race: the job leaves the queue (dequeue False)
|
||||
because it started executing. The atomic interrupt catches the now-
|
||||
running job, so cancel_job interrupts it and reports CANCEL_RUNNING."""
|
||||
interrupts = []
|
||||
dequeues = []
|
||||
result = cancel_job(
|
||||
"b", [], [make_queue_item("b")], {},
|
||||
interrupt=lambda pid: interrupts.append(pid) or True,
|
||||
dequeue=lambda pid: (dequeues.append(pid), False)[1],
|
||||
)
|
||||
assert result == CANCEL_RUNNING
|
||||
assert dequeues == ["b"] # dequeue attempted first
|
||||
assert interrupts == ["b"] # then the now-running job was interrupted
|
||||
|
||||
def test_pending_dequeue_miss_not_running_returns_unknown(self):
|
||||
"""Dequeue miss where the job is not running anymore (it finished): the
|
||||
atomic interrupt finds nothing to interrupt and returns False, so
|
||||
cancel_job is a no-op reporting UNKNOWN — never reporting a cancel that
|
||||
did not happen, and never interrupting a bystander."""
|
||||
interrupts = []
|
||||
dequeues = []
|
||||
result = cancel_job(
|
||||
"b", [], [make_queue_item("b")], {},
|
||||
interrupt=lambda pid: interrupts.append(pid) or False,
|
||||
dequeue=lambda pid: (dequeues.append(pid), False)[1],
|
||||
)
|
||||
assert result == CANCEL_UNKNOWN
|
||||
assert dequeues == ["b"]
|
||||
assert interrupts == ["b"] # interrupt attempted, found nothing running
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP contract tests: POST /api/jobs/{job_id}/cancel
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestSingleCancelEndpoint:
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_running_job_interrupts(self, aiohttp_client):
|
||||
queue = FakePromptQueue(running=[make_queue_item("a")])
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/a/cancel")
|
||||
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": True}
|
||||
assert queue.interrupt_count == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_pending_job_dequeues(self, aiohttp_client):
|
||||
queue = FakePromptQueue(pending=[make_queue_item("b")])
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/b/cancel")
|
||||
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": True}
|
||||
# Pending job removed from the queue; nothing interrupted.
|
||||
assert queue.get_current_queue()[1] == []
|
||||
assert queue.interrupt_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_terminal_job_is_idempotent_noop(self, aiohttp_client):
|
||||
history = {"c": {"prompt": make_queue_item("c"), "outputs": {}, "status": {}}}
|
||||
queue = FakePromptQueue(history=history)
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/c/cancel")
|
||||
|
||||
# Already-finished job: 200 no-op (cancelled=false), not an error.
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": False}
|
||||
assert queue.interrupt_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_unknown_id_is_200_noop(self, aiohttp_client):
|
||||
queue = FakePromptQueue()
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/does-not-exist/cancel")
|
||||
|
||||
# Single-cancel of an unknown id is treated as an idempotent no-op.
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": False}
|
||||
assert queue.interrupt_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cancel_pending_that_started_running_interrupts(self, aiohttp_client):
|
||||
"""Pending->running race end to end: the job is pending at snapshot time
|
||||
but starts executing by the time we dequeue (delete misses). The live
|
||||
re-check sees it running and interrupts it, so the cancel is not dropped
|
||||
and the caller still gets cancelled=True."""
|
||||
|
||||
class RacingQueue(FakePromptQueue):
|
||||
def delete_queue_item(self, function):
|
||||
# The worker picked the job up just before we removed it: it
|
||||
# leaves the pending queue (delete misses) and is now running.
|
||||
self._running = list(self._pending)
|
||||
self._pending = []
|
||||
return False
|
||||
|
||||
queue = RacingQueue(pending=[make_queue_item("b")])
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/b/cancel")
|
||||
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": True}
|
||||
assert queue.interrupt_count == 1
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# HTTP contract tests: POST /api/jobs/cancel (batch)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestBatchCancelEndpoint:
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_happy_path(self, aiohttp_client):
|
||||
queue = FakePromptQueue(
|
||||
running=[make_queue_item(_UUID_A)],
|
||||
pending=[make_queue_item(_UUID_B, number=1)],
|
||||
)
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/cancel", json={"job_ids": [_UUID_A, _UUID_B]})
|
||||
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": True}
|
||||
assert queue.interrupt_count == 1 # running job interrupted
|
||||
assert queue.get_current_queue()[1] == [] # pending job dequeued
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_best_effort_skips_unknown_id(self, aiohttp_client):
|
||||
"""An unknown id in the batch is a no-op, not a reason to abort: the
|
||||
running and pending jobs are still cancelled (200, cancelled=true). This
|
||||
is the "cancel all as a job finishes" case from review."""
|
||||
queue = FakePromptQueue(
|
||||
running=[make_queue_item(_UUID_A)],
|
||||
pending=[make_queue_item(_UUID_B, number=1)],
|
||||
)
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post(
|
||||
"/api/jobs/cancel", json={"job_ids": [_UUID_A, _UUID_MISSING, _UUID_B]}
|
||||
)
|
||||
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": True}
|
||||
assert queue.interrupt_count == 1 # running job interrupted
|
||||
assert queue.get_current_queue()[1] == [] # pending job dequeued
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_all_terminal_is_idempotent_noop(self, aiohttp_client):
|
||||
history = {
|
||||
_UUID_C: {"prompt": make_queue_item(_UUID_C), "outputs": {}, "status": {}},
|
||||
_UUID_D: {"prompt": make_queue_item(_UUID_D), "outputs": {}, "status": {}},
|
||||
}
|
||||
queue = FakePromptQueue(history=history)
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/cancel", json={"job_ids": [_UUID_C, _UUID_D]})
|
||||
|
||||
# All known but terminal: 200 with cancelled=false, nothing dispatched.
|
||||
assert resp.status == 200
|
||||
assert (await resp.json()) == {"cancelled": False}
|
||||
assert queue.interrupt_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_missing_job_ids_is_400(self, aiohttp_client):
|
||||
queue = FakePromptQueue()
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/cancel", json={})
|
||||
|
||||
assert resp.status == 400
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_unhashable_element_is_400_not_500(self, aiohttp_client):
|
||||
"""An unhashable element such as a dict or list must yield 400, not 500.
|
||||
|
||||
Previously, passing e.g. {"job_ids": [{}]} would reach the classify
|
||||
loop where ``prompt_id in history`` raises TypeError on an unhashable
|
||||
type, resulting in an unhandled 500. The input-validation guard must
|
||||
catch this before any queue or history access.
|
||||
"""
|
||||
queue = FakePromptQueue()
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post("/api/jobs/cancel", json={"job_ids": [{}]})
|
||||
|
||||
assert resp.status == 400
|
||||
body = await resp.json()
|
||||
assert "invalid_ids" in body
|
||||
# No queue side effects.
|
||||
assert queue.interrupt_count == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_batch_non_uuid_string_element_is_400(self, aiohttp_client):
|
||||
"""A string that is not a valid UUID must be rejected with 400."""
|
||||
queue = FakePromptQueue()
|
||||
client = await aiohttp_client(build_app(queue))
|
||||
|
||||
resp = await client.post(
|
||||
"/api/jobs/cancel", json={"job_ids": ["not-a-uuid"]}
|
||||
)
|
||||
|
||||
assert resp.status == 400
|
||||
body = await resp.json()
|
||||
assert "invalid_ids" in body
|
||||
@@ -0,0 +1,192 @@
|
||||
"""CI unit tests for FIX #2 of GHSA-779p-m5rp-r4h4.
|
||||
|
||||
Path traversal / hardening in app/model_manager.py get_model_preview
|
||||
(route /experiment/models/preview/{folder}/{path_index}/{filename:.*}).
|
||||
|
||||
Reference: https://github.com/Comfy-Org/ComfyUI/security/advisories/GHSA-779p-m5rp-r4h4
|
||||
"""
|
||||
import pytest
|
||||
import yarl
|
||||
from io import BytesIO
|
||||
from PIL import Image
|
||||
from aiohttp import web
|
||||
from unittest.mock import patch
|
||||
from app.model_manager import ModelFileManager
|
||||
|
||||
pytestmark = (
|
||||
pytest.mark.asyncio
|
||||
) # This applies the asyncio mark to all test functions in the module
|
||||
|
||||
@pytest.fixture
|
||||
def model_manager():
|
||||
return ModelFileManager()
|
||||
|
||||
@pytest.fixture
|
||||
def app(model_manager):
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
model_manager.add_routes(routes)
|
||||
app.add_routes(routes)
|
||||
return app
|
||||
|
||||
|
||||
async def test_legit_preview_returns_200(aiohttp_client, app, tmp_path):
|
||||
"""Sanity: a real preview PNG inside the model folder is served as webp 200."""
|
||||
img = Image.new('RGB', (16, 16), color=(255, 0, 128))
|
||||
img.save(tmp_path / "test_model.png", format='PNG')
|
||||
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(tmp_path)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get('/experiment/models/preview/test_folder/0/test_model.png')
|
||||
|
||||
assert response.status == 200
|
||||
assert response.content_type == 'image/webp'
|
||||
|
||||
img_bytes = BytesIO(await response.read())
|
||||
served = Image.open(img_bytes)
|
||||
assert served.format
|
||||
assert served.format.lower() == 'webp'
|
||||
served.close()
|
||||
|
||||
|
||||
async def test_non_integer_path_index_returns_400(aiohttp_client, app, tmp_path):
|
||||
"""A non-integer path_index segment must be rejected with 400."""
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(tmp_path)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get('/experiment/models/preview/test_folder/abc/test_model.png')
|
||||
|
||||
assert response.status == 400
|
||||
|
||||
|
||||
async def test_out_of_range_path_index_returns_404(aiohttp_client, app, tmp_path):
|
||||
"""A path_index beyond the configured folder list must return 404."""
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(tmp_path)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get('/experiment/models/preview/test_folder/99/test_model.png')
|
||||
|
||||
assert response.status == 404
|
||||
|
||||
|
||||
async def test_empty_filename_returns_400(aiohttp_client, app, tmp_path):
|
||||
"""The "{filename:.*}" capture also matches the empty string (trailing
|
||||
slash). It would resolve to the folder itself and must be rejected with 400."""
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(tmp_path)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get('/experiment/models/preview/test_folder/0/')
|
||||
|
||||
assert response.status == 400
|
||||
|
||||
|
||||
async def test_path_traversal_in_filename_returns_403(aiohttp_client, app, tmp_path):
|
||||
"""Path traversal in {filename} must be rejected with 403 and must NOT read
|
||||
a file outside the configured model directory.
|
||||
|
||||
GOTCHA: aiohttp/yarl collapses literal ``../`` dot-segments out of the URL
|
||||
path before it reaches the handler, which would make this test vacuously
|
||||
pass (the request would hit a different/non-existent route). We percent-encode
|
||||
the dots and slashes (``%2e%2e%2f``) and send the URL with
|
||||
``yarl.URL(..., encoded=True)`` so the bytes survive client-side normalization
|
||||
untouched; aiohttp's router then percent-decodes them into ``match_info``,
|
||||
delivering the literal ``../`` traversal to the handler's ``{filename:.*}``
|
||||
capture.
|
||||
|
||||
Without the fix the handler computes
|
||||
``os.path.normpath(os.path.join(folder, "../../../../etc/hosts"))``, which
|
||||
escapes ``tmp_path`` and would be passed straight to get_model_previews ->
|
||||
Image.open, serving bytes from outside the model dir (200/served bytes). The
|
||||
is_within_directory() containment check is the load-bearing fix that turns
|
||||
that escape into a 403.
|
||||
"""
|
||||
# Sanity-anchor: a legit preview exists inside tmp_path, so a 200 path is
|
||||
# genuinely reachable — proving the 403 below is the containment check
|
||||
# firing, not an unrelated 404.
|
||||
img = Image.new('RGB', (16, 16), color=(255, 0, 128))
|
||||
img.save(tmp_path / "test_model.png", format='PNG')
|
||||
|
||||
# Percent-encoded "../../../../etc/hosts" so yarl does not collapse the
|
||||
# dot-segments before the request leaves the client.
|
||||
encoded_traversal = '%2e%2e%2f' * 4 + 'etc%2fhosts'
|
||||
raw_path = '/experiment/models/preview/test_folder/0/' + encoded_traversal
|
||||
url = yarl.URL(raw_path, encoded=True)
|
||||
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(tmp_path)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get(url)
|
||||
|
||||
# Confirm the traversal actually reached the handler intact: a 200 here
|
||||
# would mean either normalization stripped the ``../`` (vacuous pass) or
|
||||
# the containment check failed open and served outside-dir bytes.
|
||||
assert response.status == 403, (
|
||||
f"expected 403 from is_within_directory() containment check, "
|
||||
f"got {response.status}; traversal may have been normalized away "
|
||||
f"or the fix failed open"
|
||||
)
|
||||
body = await response.read()
|
||||
assert body == b"", "403 response must not carry any file bytes"
|
||||
|
||||
|
||||
async def test_symlink_companion_preview_returns_403(aiohttp_client, app, tmp_path):
|
||||
"""A companion preview file is selected by a glob inside get_model_previews
|
||||
and then opened. If that companion is a symlink whose path is in-dir but
|
||||
whose target escapes the model folder, it must be rejected with 403 — not
|
||||
served. The requested path itself stays in-dir (so the first containment
|
||||
check passes); the load-bearing fix is the SECOND is_within_directory check
|
||||
on the file actually opened.
|
||||
"""
|
||||
model_dir = tmp_path / "models"
|
||||
model_dir.mkdir()
|
||||
secret_dir = tmp_path / "secret"
|
||||
secret_dir.mkdir()
|
||||
# A real image OUTSIDE the model dir — valid, so without the fix Image.open
|
||||
# would succeed and its bytes would be served (200).
|
||||
secret = secret_dir / "secret.png"
|
||||
Image.new('RGB', (8, 8), color=(0, 0, 0)).save(secret, format='PNG')
|
||||
# Companion preview, in-dir by name but a symlink escaping the model dir.
|
||||
# (No real model file is needed — get_model_previews globs companions by
|
||||
# basename, and omitting a .safetensors avoids the metadata-header read.)
|
||||
companion = model_dir / "model.preview.png"
|
||||
try:
|
||||
companion.symlink_to(secret)
|
||||
except (OSError, NotImplementedError):
|
||||
pytest.skip("symlinks not supported on this platform/filesystem")
|
||||
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(model_dir)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get('/experiment/models/preview/test_folder/0/model.safetensors')
|
||||
|
||||
assert response.status == 403, (
|
||||
f"expected 403 — the globbed companion preview is a symlink resolving "
|
||||
f"outside the model dir and must not be served; got {response.status}"
|
||||
)
|
||||
assert await response.read() == b""
|
||||
|
||||
|
||||
async def test_null_byte_in_filename_no_500(aiohttp_client, app, tmp_path):
|
||||
"""A NUL byte in the filename must yield a clean client rejection, not a 500
|
||||
from an uncaught ValueError in is_within_directory's realpath() call."""
|
||||
raw_path = '/experiment/models/preview/test_folder/0/' + 'a%00b'
|
||||
url = yarl.URL(raw_path, encoded=True)
|
||||
|
||||
with patch('folder_paths.folder_names_and_paths', {
|
||||
'test_folder': ([str(tmp_path)], None)
|
||||
}):
|
||||
client = await aiohttp_client(app)
|
||||
response = await client.get(url)
|
||||
|
||||
assert response.status != 500, (
|
||||
f"NUL byte produced a 500 (uncaught ValueError); expected a clean "
|
||||
f"4xx rejection, got {response.status}"
|
||||
)
|
||||
assert 400 <= response.status < 500
|
||||
@@ -0,0 +1,165 @@
|
||||
"""Security tests for GHSA-779p-m5rp-r4h4 — FIX #3.
|
||||
|
||||
Path traversal in folder_paths.get_annotated_filepath / exists_annotated_filepath,
|
||||
plus the shared is_within_directory() containment helper.
|
||||
|
||||
These are pure-function tests (no running server). The input/output/temp
|
||||
directories are pointed at tmp_path via the folder_paths setters, so a crafted
|
||||
name containing `../`, an absolute path, or a symlink that escapes the base
|
||||
directory must be rejected.
|
||||
|
||||
Reference: https://github.com/Comfy-Org/ComfyUI/security/advisories/GHSA-779p-m5rp-r4h4
|
||||
"""
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
import folder_paths
|
||||
from comfy.options import enable_args_parsing
|
||||
enable_args_parsing()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sandbox(tmp_path):
|
||||
"""Point folder_paths' input/output/temp dirs at a real temp sandbox.
|
||||
|
||||
Yields the realpath'd base, input, output and temp directories. The original
|
||||
directory values are restored afterward so tests stay isolated.
|
||||
"""
|
||||
base = os.path.realpath(str(tmp_path))
|
||||
input_dir = os.path.join(base, "input")
|
||||
output_dir = os.path.join(base, "output")
|
||||
temp_dir = os.path.join(base, "temp")
|
||||
for d in (input_dir, output_dir, temp_dir):
|
||||
os.makedirs(d, exist_ok=True)
|
||||
|
||||
orig_input = folder_paths.get_input_directory()
|
||||
orig_output = folder_paths.get_output_directory()
|
||||
orig_temp = folder_paths.get_temp_directory()
|
||||
|
||||
folder_paths.set_input_directory(input_dir)
|
||||
folder_paths.set_output_directory(output_dir)
|
||||
folder_paths.set_temp_directory(temp_dir)
|
||||
|
||||
yield {
|
||||
"base": base,
|
||||
"input": input_dir,
|
||||
"output": output_dir,
|
||||
"temp": temp_dir,
|
||||
}
|
||||
|
||||
folder_paths.set_input_directory(orig_input)
|
||||
folder_paths.set_output_directory(orig_output)
|
||||
folder_paths.set_temp_directory(orig_temp)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_within_directory() — the shared containment helper
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_is_within_directory_legit_child(sandbox):
|
||||
base = sandbox["input"]
|
||||
child = os.path.join(base, "sub", "image.png")
|
||||
assert folder_paths.is_within_directory(base, child) is True
|
||||
|
||||
|
||||
def test_is_within_directory_dotdot_escape(sandbox):
|
||||
base = sandbox["input"]
|
||||
escape = os.path.join(base, "..", "..", "etc", "passwd")
|
||||
assert folder_paths.is_within_directory(base, escape) is False
|
||||
|
||||
|
||||
def test_is_within_directory_symlink_escape(sandbox):
|
||||
"""A symlink created INSIDE base that points OUTSIDE base must not pass.
|
||||
|
||||
This is the key new hardening: is_within_directory realpath()s both operands,
|
||||
so a symlink planted in the base directory can't be used to read files
|
||||
elsewhere. We create a real on-disk symlink and a real secret target to
|
||||
verify the check actually resolves the link.
|
||||
"""
|
||||
base = sandbox["input"]
|
||||
|
||||
# A directory living outside the base, holding a secret file.
|
||||
outside = os.path.join(sandbox["base"], "outside_secret_dir")
|
||||
os.makedirs(outside, exist_ok=True)
|
||||
secret = os.path.join(outside, "secret.txt")
|
||||
with open(secret, "w") as f:
|
||||
f.write("top secret")
|
||||
|
||||
# Plant a symlink inside base that points at the outside directory.
|
||||
# symlink creation can require elevated privileges / Developer Mode on
|
||||
# Windows, so skip cleanly where it isn't available (same guard as the
|
||||
# sibling test in test_ghsa_779p_02_preview_traversal.py).
|
||||
link = os.path.join(base, "escape_link")
|
||||
try:
|
||||
os.symlink(outside, link)
|
||||
except (OSError, NotImplementedError):
|
||||
pytest.skip("symlinks not supported on this platform/filesystem")
|
||||
|
||||
# Accessing the secret "through" the in-base symlink must be rejected.
|
||||
target_via_link = os.path.join(link, "secret.txt")
|
||||
assert folder_paths.is_within_directory(base, target_via_link) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# get_annotated_filepath()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_get_annotated_filepath_legit_name(sandbox):
|
||||
result = folder_paths.get_annotated_filepath("image.png")
|
||||
assert result == os.path.join(sandbox["input"], "image.png")
|
||||
assert folder_paths.is_within_directory(sandbox["input"], result)
|
||||
|
||||
|
||||
def test_get_annotated_filepath_input_annotation(sandbox):
|
||||
result = folder_paths.get_annotated_filepath("image.png [input]")
|
||||
assert result == os.path.join(sandbox["input"], "image.png")
|
||||
|
||||
|
||||
def test_get_annotated_filepath_output_annotation(sandbox):
|
||||
result = folder_paths.get_annotated_filepath("image.png [output]")
|
||||
assert result == os.path.join(sandbox["output"], "image.png")
|
||||
|
||||
|
||||
def test_get_annotated_filepath_temp_annotation(sandbox):
|
||||
result = folder_paths.get_annotated_filepath("image.png [temp]")
|
||||
assert result == os.path.join(sandbox["temp"], "image.png")
|
||||
|
||||
|
||||
def test_get_annotated_filepath_dotdot_raises(sandbox):
|
||||
with pytest.raises(ValueError):
|
||||
folder_paths.get_annotated_filepath("../etc/passwd")
|
||||
|
||||
|
||||
def test_get_annotated_filepath_dotdot_with_annotation_raises(sandbox):
|
||||
with pytest.raises(ValueError):
|
||||
folder_paths.get_annotated_filepath("../../etc/passwd [output]")
|
||||
|
||||
|
||||
def test_get_annotated_filepath_absolute_escape_raises(sandbox):
|
||||
with pytest.raises(ValueError):
|
||||
folder_paths.get_annotated_filepath("/etc/passwd")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# exists_annotated_filepath()
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def test_exists_annotated_filepath_existing_legit_file(sandbox):
|
||||
real = os.path.join(sandbox["input"], "real.png")
|
||||
with open(real, "w") as f:
|
||||
f.write("data")
|
||||
assert folder_paths.exists_annotated_filepath("real.png") is True
|
||||
|
||||
|
||||
def test_exists_annotated_filepath_traversal_returns_false(sandbox):
|
||||
"""A traversal name must return False without raising and without probing
|
||||
outside the base directory (must never reach os.path.exists for the escape).
|
||||
"""
|
||||
# /etc/passwd exists on POSIX; the function must still report False because
|
||||
# the resolved path escapes the input directory.
|
||||
assert folder_paths.exists_annotated_filepath("../../../../../../etc/passwd") is False
|
||||
|
||||
|
||||
def test_exists_annotated_filepath_absolute_returns_false(sandbox):
|
||||
assert folder_paths.exists_annotated_filepath("/etc/passwd") is False
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user