diff --git a/.coderabbit.yaml b/.coderabbit.yaml index 0d1e49270..08629ed8e 100644 --- a/.coderabbit.yaml +++ b/.coderabbit.yaml @@ -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" diff --git a/.github/workflows/cla.yml b/.github/workflows/cla.yml new file mode 100644 index 000000000..b75397e50 --- /dev/null +++ b/.github/workflows/cla.yml @@ -0,0 +1,91 @@ +name: CLA Assistant + +on: + issue_comment: + types: [created] + pull_request_target: + types: [opened, synchronize, closed] + +permissions: + actions: write + contents: read # 'read' is enough because signatures live in a REMOTE repo + pull-requests: write + statuses: write + +jobs: + cla-assistant: + runs-on: ubuntu-latest + steps: + # The CLA action normally requires every commit author in a PR to sign. + # We only want the PR author to sign, so we allowlist all other committers + # by computing them from the PR's commits and excluding the PR author. + - name: Build author-only allowlist + id: allowlist + if: > + github.event_name == 'pull_request_target' || + (github.event_name == 'issue_comment' && github.event.issue.pull_request && ( + github.event.comment.body == 'recheck' || + github.event.comment.body == 'I have read and agree to the Contributor License Agreement' + )) + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + PR_NUMBER: ${{ github.event.pull_request.number || github.event.issue.number }} + PR_AUTHOR: ${{ github.event.pull_request.user.login || github.event.issue.user.login }} + BASE_ALLOWLIST: action@github.com,actions-user,ampagent,claude,comfy-pr-bot,GitHub Action,github-actions,github-actions[bot],Glary Bot,Glary-Bot,*[bot] + run: | + others=$(gh api "repos/${{ github.repository }}/pulls/${PR_NUMBER}/commits" --paginate \ + --jq '.[] | (.author.login // empty), (.committer.login // empty)' \ + | sort -u | grep -vix "${PR_AUTHOR}" | paste -sd, -) + if [ -n "$others" ]; then + echo "allowlist=${BASE_ALLOWLIST},${others}" >> "$GITHUB_OUTPUT" + else + echo "allowlist=${BASE_ALLOWLIST}" >> "$GITHUB_OUTPUT" + fi + + - name: CLA Assistant + # Run on PR events, on "recheck" comment, or when someone posts the exact signing phrase. + # IMPORTANT: this phrase must match `custom-pr-sign-comment` below. + if: > + github.event_name == 'pull_request_target' || + (github.event_name == 'issue_comment' && github.event.issue.pull_request && ( + github.event.comment.body == 'recheck' || + github.event.comment.body == 'I have read and agree to the Contributor License Agreement' + )) + uses: contributor-assistant/github-action@ca4a40a7d1004f18d9960b404b97e5f30a505a08 # v2.6.1 + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + # PAT required to write to the centralized signatures repo. + PERSONAL_ACCESS_TOKEN: ${{ secrets.PERSONAL_ACCESS_TOKEN }} + with: + # Where the CLA document lives (shown to contributors) + path-to-document: https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md + + # Centralized signature storage + remote-organization-name: comfy-org + remote-repository-name: comfy-cla + path-to-signatures: signatures/cla.json + branch: main + + # Only the PR author must sign: bots plus every non-author committer + # are allowlisted via the "Build author-only allowlist" step above. + # *[bot] is a catch-all for any GitHub App bot account. + allowlist: ${{ steps.allowlist.outputs.allowlist }} + + # Custom PR comment messages + custom-notsigned-prcomment: | + 🎉 Thank you for your contribution, we really appreciate it! 🎉 + + Like many open source projects, we require contributors to sign our [Contributor License Agreement (CLA)](https://github.com/Comfy-Org/comfy-cla/blob/main/comfyui_icla.md). A CLA makes the ownership of contributions explicit, so contributors and the project share a clear understanding of how the code can be used. By signing, you: + + - Confirm that you own your contribution. + - Keep the right to reuse your own code. + - Grant us a copyright license to include and share it within our projects. + + CLAs are standard practice across major open source projects including those under the Apache Software Foundation and the Linux Foundation. Ours is based on the Apache Software Foundation's CLA. Most importantly, it would enable us to relicense the project under a more permissive license in the future, giving the project and its community greater flexibility. + + ✍ **To sign, please post a new comment on this PR with exactly the following text:** ✍ + + custom-pr-sign-comment: I have read and agree to the Contributor License Agreement + + custom-allsigned-prcomment: | + ✅ All contributors have signed the CLA. Thank you! This PR is ready to be merged. diff --git a/AGENTS.md b/AGENTS.md index a8bacbd5e..05efd834b 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -127,6 +127,8 @@ - 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. +- If a library version is pinned in `requirements.txt`, do not add code to + ComfyUI to handle older versions of that library. - 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 diff --git a/CLAUDE.md b/CLAUDE.md deleted file mode 120000 index 47dc3e3d8..000000000 --- a/CLAUDE.md +++ /dev/null @@ -1 +0,0 @@ -AGENTS.md \ No newline at end of file diff --git a/README.md b/README.md index bcec86377..14c8d2cb2 100644 --- a/README.md +++ b/README.md @@ -229,7 +229,7 @@ Python 3.14 works but some custom nodes may have issues. The free threaded varia Python 3.13 is very well supported. If you have trouble with some custom node dependencies on 3.13 you can try 3.12 -torch 2.4 and above is supported but some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. +torch 2.5 is minimally supported but using a newer version is extremely recommended. Some features and optimizations might only work on newer versions. We generally recommend using the latest major version of pytorch with the latest cuda version unless it is less than 2 weeks old. If your pytorch is more than 6 months old, please update it. ### Instructions: diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 4bef096fb..0d7df5e13 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -225,6 +225,7 @@ parser.add_argument( ) parser.add_argument("--user-directory", type=is_valid_directory, default=None, help="Set the ComfyUI user directory with an absolute path. Overrides --base-directory.") +parser.add_argument("--models-directory", type=is_valid_directory, default=None, help="Set the ComfyUI models directory. Overrides the models folder in --base-directory.") parser.add_argument("--enable-compress-response-body", action="store_true", help="Enable compressing response body.") diff --git a/comfy/ldm/ace/ace_step15.py b/comfy/ldm/ace/ace_step15.py index 2ca2d26c4..02182c49f 100644 --- a/comfy/ldm/ace/ace_step15.py +++ b/comfy/ldm/ace/ace_step15.py @@ -217,10 +217,7 @@ class AceStepAttention(nn.Module): cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) - n_rep = self.num_heads // self.num_kv_heads - if n_rep > 1: - key_states = key_states.repeat_interleave(n_rep, dim=1) - value_states = value_states.repeat_interleave(n_rep, dim=1) + gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {} attn_bias = None if self.sliding_window is not None and not self.is_cross_attention: @@ -244,7 +241,7 @@ class AceStepAttention(nn.Module): else: attn_bias = window_bias - attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False) + attn_output = optimized_attention(query_states, key_states, value_states, self.num_heads, attn_bias, skip_reshape=True, low_precision_attention=False, **gqa_kwargs) attn_output = self.o_proj(attn_output) return attn_output diff --git a/comfy/ldm/audio/dit.py b/comfy/ldm/audio/dit.py index c28be5b49..b0759a240 100644 --- a/comfy/ldm/audio/dit.py +++ b/comfy/ldm/audio/dit.py @@ -425,19 +425,16 @@ class Attention(nn.Module): if n == 1 and causal: causal = False - if h != kv_h: - # Repeat interleave kv_heads to match q_heads - heads_per_kv_head = h // kv_h - k, v = map(lambda t: t.repeat_interleave(heads_per_kv_head, dim = 1), (k, v)) + gqa_kwargs = {"enable_gqa": True} if h != kv_h else {} if self.differential: q, q_diff = q.unbind(dim=1) k, k_diff = k.unbind(dim=1) - out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options) - out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options) + out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs) + out_diff = optimized_attention(q_diff, k_diff, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs) out = out - out_diff else: - out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options) + out = optimized_attention(q, k, v, h, skip_reshape=True, low_precision_attention=False, transformer_options=transformer_options, **gqa_kwargs) out = self.to_out(out) diff --git a/comfy/ldm/boogu/model.py b/comfy/ldm/boogu/model.py index 966f3c583..ca88bdeb1 100644 --- a/comfy/ldm/boogu/model.py +++ b/comfy/ldm/boogu/model.py @@ -74,11 +74,8 @@ class BooguDoubleStreamProcessor(nn.Module): key = key.transpose(1, 2) value = value.transpose(1, 2) - if attn.kv_heads < attn.heads: - key = key.repeat_interleave(attn.heads // attn.kv_heads, dim=1) - value = value.repeat_interleave(attn.heads // attn.kv_heads, dim=1) - - hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options) + gqa_kwargs = {"enable_gqa": True} if attn.kv_heads < attn.heads else {} + hidden_states = optimized_attention_masked(query, key, value, attn.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs) # Split back to instruction/image, apply per-stream output projections, recombine. instruct_hidden_states = self.instruct_out(hidden_states[:, :L_instruct]) diff --git a/comfy/ldm/modules/attention.py b/comfy/ldm/modules/attention.py index 55360535a..2411aff5c 100644 --- a/comfy/ldm/modules/attention.py +++ b/comfy/ldm/modules/attention.py @@ -1,5 +1,6 @@ import math import sys +import inspect import torch import torch.nn.functional as F @@ -14,16 +15,16 @@ from .sub_quadratic_attention import efficient_dot_product_attention from comfy import model_management -TORCH_HAS_GQA = model_management.torch_version_numeric >= (2, 5) - if model_management.xformers_enabled(): import xformers import xformers.ops SAGE_ATTENTION_IS_AVAILABLE = False +SAGE_ATTENTION_SUPPORTS_MASK = False try: from sageattention import sageattn SAGE_ATTENTION_IS_AVAILABLE = True + SAGE_ATTENTION_SUPPORTS_MASK = "attn_mask" in inspect.signature(sageattn).parameters except ImportError as e: if model_management.sage_attention_enabled(): if e.name == "sageattention": @@ -89,6 +90,44 @@ def default(val, d): return val return d +def _gqa_repeat_factor(query_heads, key_heads, value_heads): + if key_heads != value_heads: + raise ValueError(f"Key/value head count mismatch for GQA: {key_heads} != {value_heads}") + if query_heads == key_heads: + return 1 + if query_heads % key_heads != 0: + raise ValueError(f"Query heads must be divisible by key/value heads for GQA: {query_heads} vs {key_heads}") + return query_heads // key_heads + +def _repeat_kv_for_gqa(k, v, query_heads, head_dim): + n_rep = _gqa_repeat_factor(query_heads, k.shape[head_dim], v.shape[head_dim]) + if n_rep > 1: + k = k.repeat_interleave(n_rep, dim=head_dim) + v = v.repeat_interleave(n_rep, dim=head_dim) + return k, v + +def _heads_from_dim(tensor, dim_head, name): + inner_dim = tensor.shape[-1] + if inner_dim % dim_head != 0: + raise ValueError(f"{name} inner dimension {inner_dim} is not divisible by head dimension {dim_head}") + return inner_dim // dim_head + +def _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, enable_gqa=False, expand_kv=True): + q = q.unsqueeze(3).reshape(b, -1, heads, dim_head) + if enable_gqa: + key_heads = _heads_from_dim(k, dim_head, "Key") + value_heads = _heads_from_dim(v, dim_head, "Value") + else: + key_heads = heads + value_heads = heads + k = k.unsqueeze(3).reshape(b, -1, key_heads, dim_head) + v = v.unsqueeze(3).reshape(b, -1, value_heads, dim_head) + if enable_gqa: + _gqa_repeat_factor(heads, key_heads, value_heads) + if expand_kv: + k, v = _repeat_kv_for_gqa(k, v, heads, -2) + return q, k, v + # feedforward class GEGLU(nn.Module): @@ -152,28 +191,19 @@ def attention_basic(q, k, v, heads, mask=None, attn_precision=None, skip_reshape b, _, dim_head = q.shape dim_head //= heads - if kwargs.get("enable_gqa", False) and q.shape[-3] != k.shape[-3]: - n_rep = q.shape[-3] // k.shape[-3] - k = k.repeat_interleave(n_rep, dim=-3) - v = v.repeat_interleave(n_rep, dim=-3) - scale = kwargs.get("scale", dim_head ** -0.5) h = heads if skip_reshape: - q, k, v = map( + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3) + q, k, v = map( lambda t: t.reshape(b * heads, -1, dim_head), (q, k, v), ) else: - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, -1, heads, dim_head) - .permute(0, 2, 1, 3) - .reshape(b * heads, -1, dim_head) - .contiguous(), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) + q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v)) # force cast to fp32 to avoid overflowing if attn_precision == torch.float32: @@ -231,13 +261,16 @@ def attention_sub_quad(query, key, value, heads, mask=None, attn_precision=None, query = query * (kwargs["scale"] * dim_head ** 0.5) if skip_reshape: + if kwargs.get("enable_gqa", False): + key, value = _repeat_kv_for_gqa(key, value, query.shape[-3], -3) query = query.reshape(b * heads, -1, dim_head) value = value.reshape(b * heads, -1, dim_head) key = key.reshape(b * heads, -1, dim_head).movedim(1, 2) else: - query = query.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) - value = value.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) - key = key.unsqueeze(3).reshape(b, -1, heads, dim_head).permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1) + query, key, value = _reshape_qkv_to_heads(query, key, value, b, heads, dim_head, kwargs.get("enable_gqa", False)) + query = query.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) + value = value.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head) + key = key.permute(0, 2, 3, 1).reshape(b * heads, dim_head, -1) dtype = query.dtype @@ -304,19 +337,15 @@ def attention_split(q, k, v, heads, mask=None, attn_precision=None, skip_reshape scale = kwargs.get("scale", dim_head ** -0.5) if skip_reshape: - q, k, v = map( + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3) + q, k, v = map( lambda t: t.reshape(b * heads, -1, dim_head), (q, k, v), ) else: - q, k, v = map( - lambda t: t.unsqueeze(3) - .reshape(b, -1, heads, dim_head) - .permute(0, 2, 1, 3) - .reshape(b * heads, -1, dim_head) - .contiguous(), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) + q, k, v = map(lambda t: t.permute(0, 2, 1, 3).reshape(b * heads, -1, dim_head).contiguous(), (q, k, v)) r1 = torch.zeros(q.shape[0], q.shape[1], v.shape[2], device=q.device, dtype=q.dtype) @@ -438,7 +467,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh disabled_xformers = True if disabled_xformers: - return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, **kwargs) + return attention_pytorch(q, k, v, heads, mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs) if skip_reshape: # b h k d -> b k h d @@ -446,13 +475,12 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh lambda t: t.permute(0, 2, 1, 3), (q, k, v), ) + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-2], -2) # actually do the reshaping else: dim_head //= heads - q, k, v = map( - lambda t: t.reshape(b, -1, heads, dim_head), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) if mask is not None: # add a singleton batch dimension @@ -474,7 +502,7 @@ def attention_xformers(q, k, v, heads, mask=None, attn_precision=None, skip_resh mask = mask_out[..., :mask.shape[-1]] mask = mask.expand(b, heads, -1, -1) - out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask) + out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=mask, scale=kwargs.get("scale", None)) if skip_output_reshape: out = out.permute(0, 2, 1, 3) @@ -498,10 +526,8 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha else: b, _, dim_head = q.shape dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False) + q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) if mask is not None: # add a batch dimension if there isn't already one @@ -511,9 +537,7 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha if mask.ndim == 3: mask = mask.unsqueeze(1) - # Pass through extra SDPA kwargs (scale, enable_gqa) if provided - # enable_gqa requires PyTorch 2.5+; older versions use manual KV expansion above - sdpa_keys = ("scale", "enable_gqa") if TORCH_HAS_GQA else ("scale",) + sdpa_keys = ("scale", "enable_gqa") sdpa_extra = {k: v for k, v in kwargs.items() if k in sdpa_keys} if SDP_BATCH_LIMIT >= b: @@ -541,20 +565,19 @@ def attention_pytorch(q, k, v, heads, mask=None, attn_precision=None, skip_resha @wrap_attn def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape=False, skip_output_reshape=False, **kwargs): - if kwargs.get("low_precision_attention", True) is False: + if kwargs.get("low_precision_attention", True) is False or (mask is not None and not SAGE_ATTENTION_SUPPORTS_MASK): return attention_pytorch(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs) exception_fallback = False if skip_reshape: b, _, _, dim_head = q.shape tensor_layout = "HND" + if kwargs.get("enable_gqa", False): + k, v = _repeat_kv_for_gqa(k, v, q.shape[-3], -3) else: b, _, dim_head = q.shape dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False)) tensor_layout = "NHD" if mask is not None: @@ -565,8 +588,12 @@ def attention_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape= if mask.ndim == 3: mask = mask.unsqueeze(1) + sage_kwargs = {"is_causal": False, "tensor_layout": tensor_layout, "sm_scale": kwargs.get("scale", None), "smooth_k": False} + if mask is not None: + sage_kwargs["attn_mask"] = mask + try: - out = sageattn(q, k, v, attn_mask=mask, is_causal=False, tensor_layout=tensor_layout) + out = sageattn(q, k, v, **sage_kwargs) except Exception as e: logging.error("Error running sage attention: {}, using pytorch attention instead.".format(e)) exception_fallback = True @@ -616,7 +643,6 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape skip_output_reshape=skip_output_reshape, **kwargs ) - q_s, k_s, v_s = q, k, v N = q.shape[2] dim_head = D else: @@ -642,11 +668,15 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape **kwargs ) - if not skip_reshape: - q_s, k_s, v_s = map( - lambda t: t.view(B, -1, heads, dim_head).permute(0, 2, 1, 3).contiguous(), - (q, k, v), - ) + if skip_reshape: + q_s = q + if kwargs.get("enable_gqa", False): + k_s, v_s = _repeat_kv_for_gqa(k, v, H, -3) + else: + k_s, v_s = k, v + else: + q_s, k_s, v_s = _reshape_qkv_to_heads(q, k, v, B, heads, dim_head, kwargs.get("enable_gqa", False)) + q_s, k_s, v_s = map(lambda t: t.permute(0, 2, 1, 3).contiguous(), (q_s, k_s, v_s)) B, H, L, D = q_s.shape try: @@ -662,7 +692,7 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape q, k, v, heads, mask=mask, attn_precision=attn_precision, - skip_reshape=False, + skip_reshape=skip_reshape, skip_output_reshape=skip_output_reshape, **kwargs ) @@ -681,19 +711,20 @@ def attention3_sage(q, k, v, heads, mask=None, attn_precision=None, skip_reshape try: @torch.library.custom_op("flash_attention::flash_attn", mutates_args=()) def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor: - return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal) + dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor: + softmax_scale_arg = None if softmax_scale == -1.0 else softmax_scale + return flash_attn_func(q, k, v, dropout_p=dropout_p, causal=causal, softmax_scale=softmax_scale_arg) @flash_attn_wrapper.register_fake - def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False): + def flash_attn_fake(q, k, v, dropout_p=0.0, causal=False, softmax_scale=-1.0): # Output shape is the same as q return q.new_empty(q.shape) except AttributeError as error: FLASH_ATTN_ERROR = error def flash_attn_wrapper(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, - dropout_p: float = 0.0, causal: bool = False) -> torch.Tensor: + dropout_p: float = 0.0, causal: bool = False, softmax_scale: float = -1.0) -> torch.Tensor: assert False, f"Could not define flash_attn_wrapper: {FLASH_ATTN_ERROR}" @wrap_attn @@ -703,10 +734,8 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape else: b, _, dim_head = q.shape dim_head //= heads - q, k, v = map( - lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), - (q, k, v), - ) + q, k, v = _reshape_qkv_to_heads(q, k, v, b, heads, dim_head, kwargs.get("enable_gqa", False), expand_kv=False) + q, k, v = map(lambda t: t.transpose(1, 2), (q, k, v)) if mask is not None: # add a batch dimension if there isn't already one @@ -725,10 +754,16 @@ def attention_flash(q, k, v, heads, mask=None, attn_precision=None, skip_reshape v.transpose(1, 2), dropout_p=0.0, causal=False, + softmax_scale=kwargs.get("scale", -1.0), ).transpose(1, 2) except Exception as e: logging.warning(f"Flash Attention failed, using default SDPA: {e}") - out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False) + sdpa_extra = {} + if kwargs.get("enable_gqa", False): + sdpa_extra["enable_gqa"] = True + if "scale" in kwargs: + sdpa_extra["scale"] = kwargs["scale"] + out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False, **sdpa_extra) if not skip_output_reshape: out = ( out.transpose(1, 2).reshape(b, -1, heads * dim_head) @@ -1209,5 +1244,3 @@ class SpatialVideoTransformer(SpatialTransformer): x = self.proj_out(x) out = x + x_in return out - - diff --git a/comfy/ldm/omnigen/omnigen2.py b/comfy/ldm/omnigen/omnigen2.py index b8da4cf39..d18a9f6d0 100644 --- a/comfy/ldm/omnigen/omnigen2.py +++ b/comfy/ldm/omnigen/omnigen2.py @@ -141,11 +141,8 @@ class Attention(nn.Module): key = key.transpose(1, 2) value = value.transpose(1, 2) - if self.kv_heads < self.heads: - key = key.repeat_interleave(self.heads // self.kv_heads, dim=1) - value = value.repeat_interleave(self.heads // self.kv_heads, dim=1) - - hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options) + gqa_kwargs = {"enable_gqa": True} if self.kv_heads < self.heads else {} + hidden_states = optimized_attention_masked(query, key, value, self.heads, attention_mask, skip_reshape=True, transformer_options=transformer_options, **gqa_kwargs) hidden_states = self.to_out[0](hidden_states) return hidden_states diff --git a/comfy/ops.py b/comfy/ops.py index 69d32e254..35a1ee31e 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -174,6 +174,8 @@ def cast_modules_with_vbar(comfy_modules, dtype, device, bias_dtype, non_blockin elif xfer_dest2 is not None: xfer_source.prepare(xfer_dest2, stream, copy=True, commit=False) return + else: + return comfy.model_management.cast_to_gathered(xfer_source, xfer_dest, non_blocking=non_blocking, stream=stream, r2=xfer_dest2) def handle_pin(m, pin, source, dest, subset="weights", size=None): diff --git a/comfy/sd.py b/comfy/sd.py index 610c4e2b8..faf3104f3 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -468,6 +468,9 @@ class CLIP: def decode(self, token_ids, skip_special_tokens=True): return self.tokenizer.decode(token_ids, skip_special_tokens=skip_special_tokens) + def is_dynamic(self): + return self.patcher.is_dynamic() + class VAE: def __init__(self, sd=None, device=None, config=None, dtype=None, metadata=None): if 'decoder.up_blocks.0.resnets.0.norm1.weight' in sd.keys(): #diffusers format @@ -1251,6 +1254,8 @@ class VAE: except: return None + def is_dynamic(self): + return self.patcher.is_dynamic() class StyleModel: def __init__(self, model, device="cpu"): diff --git a/comfy/text_encoders/gpt_oss.py b/comfy/text_encoders/gpt_oss.py index d596ef9a0..066796b6a 100644 --- a/comfy/text_encoders/gpt_oss.py +++ b/comfy/text_encoders/gpt_oss.py @@ -12,7 +12,7 @@ import torch.nn.functional as F import comfy.ops from comfy import sd1_clip -from comfy.ldm.modules.attention import TORCH_HAS_GQA, optimized_attention_for_device +from comfy.ldm.modules.attention import optimized_attention_for_device from comfy.text_encoders.llama import RMSNorm, apply_rope @@ -110,10 +110,6 @@ def _attention_with_sinks(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, sin putting the sink logit in the mask at that column. """ - if num_kv_groups > 1 and not TORCH_HAS_GQA: - k = k.repeat_interleave(num_kv_groups, dim=1) - v = v.repeat_interleave(num_kv_groups, dim=1) - B, _, S_q, D = q.shape H_kv = k.shape[1] S_kv = k.shape[-2] diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py index e9f38a9a2..3f98fb0a5 100644 --- a/comfy/text_encoders/llama.py +++ b/comfy/text_encoders/llama.py @@ -550,10 +550,8 @@ class Attention(nn.Module): xv = xv[:, :, -sliding_window:] attention_mask = attention_mask[..., -sliding_window:] if attention_mask is not None else None - xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - - output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True) + gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {} + output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs) return self.o_proj(output), present_key_value class MLP(nn.Module): @@ -937,22 +935,41 @@ class BaseGenerate: return torch.argmax(logits, dim=-1, keepdim=True) # Sampling mode - if repetition_penalty != 1.0: - for i in range(logits.shape[0]): - for token_id in set(token_history): - logits[i, token_id] *= repetition_penalty if logits[i, token_id] < 0 else 1/repetition_penalty - - if presence_penalty is not None and presence_penalty != 0.0: - for i in range(logits.shape[0]): - for token_id in set(token_history): - logits[i, token_id] -= presence_penalty + if len(token_history) > 0 and (repetition_penalty != 1.0 or (presence_penalty is not None and presence_penalty != 0.0)): + token_ids = torch.tensor(list(set(token_history)), device=logits.device) + token_logits = logits[:, token_ids] + if repetition_penalty != 1.0: + token_logits = torch.where(token_logits < 0, token_logits * repetition_penalty, token_logits / repetition_penalty) + if presence_penalty is not None and presence_penalty != 0.0: + token_logits = token_logits - presence_penalty + logits[:, token_ids] = token_logits if temperature != 1.0: logits = logits / temperature if top_k > 0: - indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None] - logits[indices_to_remove] = torch.finfo(logits.dtype).min + top_k = min(top_k, logits.shape[-1]) + logits, top_indices = torch.topk(logits, top_k) + + if min_p > 0.0: + probs_before_filter = torch.nn.functional.softmax(logits, dim=-1) + top_probs, _ = probs_before_filter.max(dim=-1, keepdim=True) + min_threshold = min_p * top_probs + indices_to_remove = probs_before_filter < min_threshold + logits[indices_to_remove] = torch.finfo(logits.dtype).min + + if top_p < 1.0: + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = torch.cumsum(torch.nn.functional.softmax(sorted_logits, dim=-1), dim=-1) + sorted_indices_to_remove = cumulative_probs > top_p + sorted_indices_to_remove[..., 0] = False + indices_to_remove = torch.zeros_like(logits, dtype=torch.bool) + indices_to_remove.scatter_(1, sorted_indices, sorted_indices_to_remove) + logits[indices_to_remove] = torch.finfo(logits.dtype).min + + probs = torch.nn.functional.softmax(logits, dim=-1) + next_token = torch.multinomial(probs, num_samples=1, generator=generator) + return top_indices.gather(1, next_token) if min_p > 0.0: probs_before_filter = torch.nn.functional.softmax(logits, dim=-1) diff --git a/comfy/text_encoders/qwen35.py b/comfy/text_encoders/qwen35.py index 71a17990f..304a4357f 100644 --- a/comfy/text_encoders/qwen35.py +++ b/comfy/text_encoders/qwen35.py @@ -366,12 +366,8 @@ class GatedAttention(nn.Module): xv = torch.cat((past_value[:, :, :index], xv), dim=2) present_key_value = (xk, xv, index + num_tokens) - # Expand KV heads for GQA - if self.num_heads != self.num_kv_heads: - xk = xk.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - xv = xv.repeat_interleave(self.num_heads // self.num_kv_heads, dim=1) - - output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True) + gqa_kwargs = {"enable_gqa": True} if self.num_heads != self.num_kv_heads else {} + output = optimized_attention(xq, xk, xv, self.num_heads, mask=attention_mask, skip_reshape=True, **gqa_kwargs) output = output * gate.sigmoid() return self.o_proj(output), present_key_value diff --git a/comfy_api/latest/_input_impl/video_types.py b/comfy_api/latest/_input_impl/video_types.py index 6c69256ab..bc95a5b99 100644 --- a/comfy_api/latest/_input_impl/video_types.py +++ b/comfy_api/latest/_input_impl/video_types.py @@ -281,11 +281,18 @@ class VideoFromFile(VideoInput): video_done = False audio_done = True - if len(container.streams.audio): - audio_stream = container.streams.audio[-1] + # Use the last decodable audio stream. Streams FFmpeg has no decoder for have no codec context, + # and decoding their packets crashes the process. (e.g. APAC spatial-audio track in iPhone) + audio_stream = next( + (s for s in reversed(container.streams.audio) if s.codec_context is not None), + None, + ) + if audio_stream is not None: streams += [audio_stream] resampler = av.audio.resampler.AudioResampler(format='fltp') audio_done = False + elif len(container.streams.audio): + logging.warning("No decodable audio stream found in video; ignoring audio.") for packet in container.demux(*streams): if video_done and audio_done: @@ -457,10 +464,13 @@ class VideoFromFile(VideoInput): else: output_container.metadata[key] = json.dumps(value) - # Add streams to the new container + # Add streams to the new container. Streams with no codec context cannot be used as an output template. stream_map = {} for stream in streams: if isinstance(stream, (av.VideoStream, av.AudioStream, SubtitleStream)): + if stream.codec_context is None: + logging.warning("Skipping %s stream %d with unsupported codec", stream.type, stream.index) + continue out_stream = output_container.add_stream_from_template(template=stream, opaque=True) stream_map[stream] = out_stream diff --git a/comfy_api_nodes/util/request_logger.py b/comfy_api_nodes/util/request_logger.py index fe0543d9b..70ecaf41a 100644 --- a/comfy_api_nodes/util/request_logger.py +++ b/comfy_api_nodes/util/request_logger.py @@ -9,6 +9,7 @@ from typing import Any import folder_paths logger = logging.getLogger(__name__) +_SENSITIVE_HEADERS = {"authorization", "x-api-key"} def get_log_directory(): @@ -73,6 +74,10 @@ def _format_data_for_logging(data: Any) -> str: return str(data) +def _redact_headers(headers: dict) -> dict: + return {k: ("***" if k.lower() in _SENSITIVE_HEADERS else v) for k, v in headers.items()} + + def log_request_response( operation_id: str, request_method: str, @@ -101,7 +106,7 @@ def log_request_response( log_content.append(f"Method: {request_method}") log_content.append(f"URL: {request_url}") if request_headers: - log_content.append(f"Headers:\n{_format_data_for_logging(request_headers)}") + log_content.append(f"Headers:\n{_format_data_for_logging(_redact_headers(request_headers))}") if request_params: log_content.append(f"Params:\n{_format_data_for_logging(request_params)}") if request_data is not None: diff --git a/comfy_api_nodes/util/upload_helpers.py b/comfy_api_nodes/util/upload_helpers.py index 6d1d107a1..f7029ee78 100644 --- a/comfy_api_nodes/util/upload_helpers.py +++ b/comfy_api_nodes/util/upload_helpers.py @@ -158,7 +158,14 @@ async def upload_video_to_comfyapi( # Convert VideoInput to BytesIO using specified container/codec video_bytes_io = BytesIO() - video.save_to(video_bytes_io, format=container, codec=codec) + try: + video.save_to(video_bytes_io, format=container, codec=codec) + except Exception as e: + raise ValueError( + f"Could not convert the input video to {container.value.upper()} for upload; " + f"the file may be corrupted or use an unsupported codec. " + f"Try re-exporting it as MP4 (H.264). Original error: {e}" + ) from e video_bytes_io.seek(0) return await upload_file_to_comfyapi(cls, video_bytes_io, filename, upload_mime_type, wait_label) diff --git a/comfy_execution/caching.py b/comfy_execution/caching.py index ba1e8bc84..ad75a0e50 100644 --- a/comfy_execution/caching.py +++ b/comfy_execution/caching.py @@ -503,6 +503,21 @@ RAM_CACHE_DEFAULT_RAM_USAGE = 0.05 RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER = 1.3 + +def all_outputs_dynamic(outputs): + if outputs is None: + return False + + for output in outputs: + if isinstance(output, (list, tuple)): + if not all_outputs_dynamic(output): + return False + elif not hasattr(output, "is_dynamic") or not output.is_dynamic(): + return False + + return True + + class RAMPressureCache(LRUCache): def __init__(self, key_class, enable_providers=False): @@ -533,7 +548,11 @@ class RAMPressureCache(LRUCache): for key, cache_entry in self.cache.items(): if not free_active and self.used_generation[key] == self.generation: continue - oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key]) + + if all_outputs_dynamic(cache_entry.outputs) and self.used_generation[key] == self.generation: + continue + + oom_score = RAM_CACHE_OLD_WORKFLOW_OOM_MULTIPLIER ** (self.generation - self.used_generation[key]) ram_usage = RAM_CACHE_DEFAULT_RAM_USAGE def scan_list_for_ram_usage(outputs): diff --git a/comfy_extras/nodes_color.py b/comfy_extras/nodes_color.py index f58e51bff..6d10b26f4 100644 --- a/comfy_extras/nodes_color.py +++ b/comfy_extras/nodes_color.py @@ -16,23 +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") + # 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") from None + 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): diff --git a/comfy_extras/nodes_text_overlay.py b/comfy_extras/nodes_text_overlay.py new file mode 100644 index 000000000..4c5cdae60 --- /dev/null +++ b/comfy_extras/nodes_text_overlay.py @@ -0,0 +1,150 @@ +import numpy as np +import torch +from PIL import Image as PILImage, ImageColor, ImageDraw, ImageFont +from typing_extensions import override + +from comfy_api.latest import ComfyExtension, IO + + +class TextOverlay(IO.ComfyNode): + @classmethod + def define_schema(cls): + return IO.Schema( + node_id="TextOverlay", + display_name="Draw Text Overlay", + category="text", + description="Draw text overlay on an image or batch of images.", + search_aliases=["text", "label", "caption", "subtitle", "watermark", "title", "addlabel", "overlay"], + inputs=[ + IO.Image.Input("images"), + IO.String.Input("text", multiline=True, default=""), + IO.Float.Input("font_size", default=5.0, min=0.5, max=50.0, step=0.5, tooltip="Font size as a percentage of the image height."), + IO.Color.Input("color", default="#ffffff", tooltip="Color of the text."), + IO.Combo.Input("position", options=["top", "bottom"], default="top"), + IO.Combo.Input("align", options=["left", "center", "right"], default="left"), + IO.Boolean.Input("outline", default=True, tooltip="Draw a black outline around the text."), + ], + outputs=[IO.Image.Output(display_name="images")], + ) + + @classmethod + def execute(cls, images, text, font_size, color, position, align, outline) -> IO.NodeOutput: + if text.strip() == "": + return IO.NodeOutput(images) + + text = text.replace("\\n", "\n").replace("\\t", "\t") + + text_rgba = cls.parse_color_to_rgba(color) + outline_rgba = (0, 0, 0, 255) if outline else (0, 0, 0, 0) + + # Render the overlay once and composite it across all frames in the batch + height = images.shape[1] + width = images.shape[2] + overlay_rgb, overlay_alpha = cls.render_overlay_text(width, height, text, position, align, font_size, text_rgba, outline_rgba) + overlay_rgb = overlay_rgb.to(device=images.device, dtype=images.dtype) + overlay_alpha = overlay_alpha.to(device=images.device, dtype=images.dtype) + + result = images * (1.0 - overlay_alpha) + overlay_rgb * overlay_alpha + return IO.NodeOutput(result) + + @staticmethod + def parse_color_to_rgba(color_string): + parsed = ImageColor.getrgb(color_string) + + if len(parsed) == 3: + return (*parsed, 255) + + return parsed + + @classmethod + def render_overlay_text(cls, width, height, text, position, align, font_size, text_rgba, outline_rgba): + line_spacing = 1.2 + margin_percent = 1.0 + min_font_percent = 2.0 + min_font_pixels = 10 + outline_thickness_factor = 0.04 + + # Draw onto a transparent layer so the result can be alpha-composited over any frame. + layer = PILImage.new("RGBA", (width, height), (0, 0, 0, 0)) + draw = ImageDraw.Draw(layer) + + margin = int(round(margin_percent / 100.0 * min(width, height))) + max_width = max(1, width - 2 * margin) + max_height = max(1, height - 2 * margin) + + # Font scales with resolution, then shrinks to fit the height. + size = max(1, int(round(font_size / 100.0 * height))) + floor = min(size, max(min_font_pixels, int(round(min_font_percent / 100.0 * height)))) + + while True: + font = ImageFont.load_default(size=size) + stroke = max(1, int(round(size * outline_thickness_factor))) if outline_rgba[3] > 0 else 0 + block = "\n".join(cls.wrap_text(text, font, max_width)) + # convert line spacing to pixel spacing + single = draw.textbbox((0, 0), "Ay", font=font, stroke_width=stroke) + double = draw.multiline_textbbox((0, 0), "Ay\nAy", font=font, spacing=0, stroke_width=stroke) + natural_advance = (double[3] - double[1]) - (single[3] - single[1]) + pixel_spacing = int(round(size * line_spacing - natural_advance)) + box = draw.multiline_textbbox((0, 0), block, font=font, spacing=pixel_spacing, stroke_width=stroke) + block_height = box[3] - box[1] + + if block_height <= max_height or size <= floor: + break + + size = max(floor, int(size * 0.9)) + + anchor_h, x = {"left": ("l", margin), "center": ("m", width / 2), "right": ("r", width - margin)}[align] + + # Offset y so the rendered text sits flush against the margin + if position == "bottom": + y = height - margin - box[3] + else: + y = margin - box[1] + + draw.multiline_text((x, y), block, font=font, fill=text_rgba, anchor=anchor_h + "a", + align=align, spacing=pixel_spacing, stroke_width=stroke, stroke_fill=outline_rgba) + + overlay = np.array(layer).astype(np.float32) / 255.0 + overlay_rgb = torch.from_numpy(overlay[:, :, :3]) + overlay_alpha = torch.from_numpy(overlay[:, :, 3:4]) + return overlay_rgb, overlay_alpha + + @staticmethod + def wrap_text(text, font, max_width): + lines = [] + for raw_line in text.split("\n"): + words = raw_line.split() + if not words: + lines.append("") + continue + current = "" + # Break the line into words and split words that are too long + for word in words: + while font.getlength(word) > max_width and len(word) > 1: + cut = 1 + while cut < len(word) and font.getlength(word[:cut + 1]) <= max_width: + cut += 1 + if current: + lines.append(current) + current = "" + lines.append(word[:cut]) + word = word[cut:] + candidate = word if not current else current + " " + word + if not current or font.getlength(candidate) <= max_width: + current = candidate + else: + lines.append(current) + current = word + if current: + lines.append(current) + return lines + + +class TextOverlayExtension(ComfyExtension): + @override + async def get_node_list(self) -> list[type[IO.ComfyNode]]: + return [TextOverlay] + + +async def comfy_entrypoint() -> TextOverlayExtension: + return TextOverlayExtension() diff --git a/folder_paths.py b/folder_paths.py index ee048b0f2..937428c18 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -17,7 +17,11 @@ if args.base_directory: else: base_path = os.path.dirname(os.path.realpath(__file__)) -models_dir = os.path.join(base_path, "models") +if args.models_directory: + models_dir = os.path.abspath(args.models_directory) +else: + models_dir = os.path.join(base_path, "models") + folder_names_and_paths["checkpoints"] = ([os.path.join(models_dir, "checkpoints")], supported_pt_extensions) folder_names_and_paths["configs"] = ([os.path.join(models_dir, "configs")], [".yaml"]) diff --git a/main.py b/main.py index 20ec83c9e..580074b19 100644 --- a/main.py +++ b/main.py @@ -131,6 +131,10 @@ def apply_custom_paths(): if args.base_directory: logging.info(f"Setting base directory to: {folder_paths.base_path}") + # --models-directory + if args.models_directory: + logging.info(f"Setting models directory to: {folder_paths.models_dir}") + # --output-directory, --input-directory, --user-directory if args.output_directory: output_dir = os.path.abspath(args.output_directory) diff --git a/nodes.py b/nodes.py index 9043a8d0a..e126576fe 100644 --- a/nodes.py +++ b/nodes.py @@ -2478,6 +2478,7 @@ async def init_builtin_extra_nodes(): "nodes_glsl.py", "nodes_lora_debug.py", "nodes_textgen.py", + "nodes_text_overlay.py", "nodes_color.py", "nodes_toolkit.py", "nodes_replacements.py", diff --git a/requirements.txt b/requirements.txt index 34af2ce39..978411b3e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ comfyui-frontend-package==1.45.20 comfyui-workflow-templates==0.11.2 -comfyui-embedded-docs==0.5.6 +comfyui-embedded-docs==0.5.7 torch torchsde torchvision diff --git a/tests-unit/comfy_test/folder_path_test.py b/tests-unit/comfy_test/folder_path_test.py index 3b398e60b..a0ef17a4c 100644 --- a/tests-unit/comfy_test/folder_path_test.py +++ b/tests-unit/comfy_test/folder_path_test.py @@ -163,3 +163,20 @@ def test_base_path_change_clears_old(set_base_dir): for name in ["controlnet", "diffusion_models", "text_encoders"]: assert len(folder_paths.get_folder_paths(name)) == 2 + + +def test_models_directory_cli_and_getters(temp_dir): + try: + with patch.object(sys, 'argv', ["main.py", "--models-directory", temp_dir]): + reload(comfy.cli_args) + reload(folder_paths) + + assert folder_paths.models_dir == os.path.abspath(temp_dir) + + with pytest.raises(Exception): + comfy.cli_args.is_valid_directory(os.path.join(temp_dir, "non_existent_folder_path")) + finally: + with patch.object(sys, 'argv', ["main.py"]): + reload(comfy.cli_args) + reload(folder_paths) +