Merge 528c44de2d into 7cf4e78335

.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.

CLAUDE.md

@@ -1 +0,0 @@
-AGENTS.md

comfy/background_removal/birefnet.py

@@ -433,19 +433,16 @@ class DeformableConv2d(nn.Module):
     def forward(self, x):
         offset = self.offset_conv(x)
         modulator = 2. * torch.sigmoid(self.modulator_conv(x))
-        weight, bias, offload_info = comfy.ops.cast_bias_weight(self.regular_conv, x, offloadable=True)
-
-        x = deform_conv2d(
-            input=x,
-            offset=offset,
-            weight=weight,
-            bias=None,
-            padding=self.padding,
-            mask=modulator,
-            stride=self.stride,
-        )
-        comfy.ops.uncast_bias_weight(self.regular_conv, weight, bias, offload_info)
-        return x
+        with comfy.ops.CastBiasWeightContext(self.regular_conv, x, offloadable=True) as (weight, _bias):
+            return deform_conv2d(
+                input=x,
+                offset=offset,
+                weight=weight,
+                bias=None,
+                padding=self.padding,
+                mask=modulator,
+                stride=self.stride,
+            )
 
 class BasicDecBlk(nn.Module):
     def __init__(self, in_channels=64, out_channels=64, inter_channels=64, device=None, dtype=None, operations=None):

comfy/controlnet.py

@@ -381,13 +381,10 @@ class ControlLoraOps:
             self.bias = None
 
         def forward(self, input):
-            weight, bias, offload_stream = comfy.ops.cast_bias_weight(self, input, offloadable=True)
-            if self.up is not None:
-                x = torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
-            else:
-                x = torch.nn.functional.linear(input, weight, bias)
-            comfy.ops.uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with comfy.ops.CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                if self.up is None:
+                    return torch.nn.functional.linear(input, weight, bias)
+                return torch.nn.functional.linear(input, weight + (torch.mm(self.up.flatten(start_dim=1), self.down.flatten(start_dim=1))).reshape(self.weight.shape).type(input.dtype), bias)
 
     class Conv2d(torch.nn.Module, comfy.ops.CastWeightBiasOp):
         def __init__(

comfy/ops.py

@@ -402,6 +402,26 @@ def uncast_bias_weight(s, weight, bias, offload_stream):
             device = bias_a.device
     os.wait_stream(comfy.model_management.current_stream(device))
 
+class CastBiasWeightContext:
+    # When initialized with no arguments or the first is None, the context
+    # will return the tuple (None, None).
+    def __init__(self, *args, **kwargs):
+        self.slf = args[0] if len(args) else None
+        self.state = (None, None) if self.slf is None else cast_bias_weight(*args, **kwargs)
+
+    def __enter__(self):
+        result = self.state
+        if len(result) < 3 or result[2] is None:
+            # Not offloaded, immediately drop references.
+            self.state = self.slf = None
+        return result[:2]
+
+    def __exit__(self, *_args) -> None:
+        if not self.slf:
+            return
+        slf, state = self.slf, self.state
+        self.state = self.slf = None
+        uncast_bias_weight(slf, *state)
 
 class CastWeightBiasOp:
     comfy_cast_weights = False
@@ -490,10 +510,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.linear(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.linear(input, weight, bias)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -507,10 +525,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = self._conv_forward(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return self._conv_forward(input, weight, bias)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -524,10 +540,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = self._conv_forward(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return self._conv_forward(input, weight, bias)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -552,10 +566,8 @@ class disable_weight_init:
                 return super()._conv_forward(input, weight, bias, *args, **kwargs)
 
         def forward_comfy_cast_weights(self, input, autopad=None):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = self._conv_forward(input, weight, bias, autopad=autopad)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return self._conv_forward(input, weight, bias, autopad=autopad)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -569,10 +581,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.group_norm(input, self.num_groups, weight, bias, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -586,12 +596,10 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
-            running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
-            x = torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                running_mean = self.running_mean.to(device=input.device, dtype=weight.dtype) if self.running_mean is not None else None
+                running_var = self.running_var.to(device=input.device, dtype=weight.dtype) if self.running_var is not None else None
+                return torch.nn.functional.batch_norm(input, running_mean, running_var, weight, bias, self.training, self.momentum, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -605,15 +613,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            if self.weight is not None:
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            else:
-                weight = None
-                bias = None
-                offload_stream = None
-            x = torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.layer_norm(input, self.normalized_shape, weight, bias, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -628,15 +629,8 @@ class disable_weight_init:
             return None
 
         def forward_comfy_cast_weights(self, input):
-            if self.weight is not None:
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            else:
-                weight = None
-                bias = None
-                offload_stream = None
-            x = torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self if self.weight is not None else None, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.rms_norm(input, self.normalized_shape, weight, self.eps)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -655,12 +649,10 @@ class disable_weight_init:
                 input, output_size, self.stride, self.padding, self.kernel_size,
                 num_spatial_dims, self.dilation)
 
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.conv_transpose2d(
-                input, weight, bias, self.stride, self.padding,
-                output_padding, self.groups, self.dilation)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.conv_transpose2d(
+                       input, weight, bias, self.stride, self.padding,
+                       output_padding, self.groups, self.dilation)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -679,12 +671,10 @@ class disable_weight_init:
                 input, output_size, self.stride, self.padding, self.kernel_size,
                 num_spatial_dims, self.dilation)
 
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.conv_transpose1d(
-                input, weight, bias, self.stride, self.padding,
-                output_padding, self.groups, self.dilation)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.conv_transpose1d(
+                       input, weight, bias, self.stride, self.padding,
+                       output_padding, self.groups, self.dilation)
 
         def forward(self, *args, **kwargs):
             run_every_op()
@@ -749,10 +739,8 @@ class disable_weight_init:
             output_dtype = out_dtype
             if self.weight.dtype == torch.float16 or self.weight.dtype == torch.bfloat16:
                 out_dtype = None
-            weight, bias, offload_stream = cast_bias_weight(self, device=input.device, dtype=out_dtype, offloadable=True)
-            x = torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, device=input.device, dtype=out_dtype, offloadable=True) as (weight, bias):
+                return torch.nn.functional.embedding(input, weight, self.padding_idx, self.max_norm, self.norm_type, self.scale_grad_by_freq, self.sparse).to(dtype=output_dtype)
 
 
         def forward(self, *args, **kwargs):
@@ -828,7 +816,6 @@ def fp8_linear(self, input):
     if input.ndim != 2:
         return None
     lora_compute_dtype=comfy.model_management.lora_compute_dtype(input.device)
-    w, bias, offload_stream = cast_bias_weight(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True)
     scale_weight = torch.ones((), device=input.device, dtype=torch.float32)
 
     scale_input = torch.ones((), device=input.device, dtype=torch.float32)
@@ -837,15 +824,16 @@ def fp8_linear(self, input):
     layout_params_input = TensorCoreFP8Layout.Params(scale=scale_input, orig_dtype=input_dtype, orig_shape=tuple(input_fp8.shape))
     quantized_input = QuantizedTensor(input_fp8, "TensorCoreFP8Layout", layout_params_input)
 
-    # Wrap weight in QuantizedTensor - this enables unified dispatch
-    # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
-    layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=tuple(w.shape))
-    quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
-    o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
+    with CastBiasWeightContext(self, input, dtype=dtype, bias_dtype=input_dtype, offloadable=True, compute_dtype=lora_compute_dtype, want_requant=True) as (w, bias):
+        # Wrap weight in QuantizedTensor - this enables unified dispatch
+        # Call F.linear - __torch_dispatch__ routes to fp8_linear handler in quant_ops.py!
+        w_shape = tuple(w.shape)
+        layout_params_weight = TensorCoreFP8Layout.Params(scale=scale_weight, orig_dtype=input_dtype, orig_shape=w_shape)
+        quantized_weight = QuantizedTensor(w, "TensorCoreFP8Layout", layout_params_weight)
+        o = torch.nn.functional.linear(quantized_input, quantized_weight, bias)
 
-    uncast_bias_weight(self, w, bias, offload_stream)
     if tensor_3d:
-        o = o.reshape((input_shape[0], input_shape[1], w.shape[0]))
+        o = o.reshape((input_shape[0], input_shape[1], w_shape[0]))
 
     return o
 
@@ -865,10 +853,8 @@ class fp8_ops(manual_cast):
                 except Exception as e:
                     logging.info("Exception during fp8 op: {}".format(e))
 
-            weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-            x = torch.nn.functional.linear(input, weight, bias)
-            uncast_bias_weight(self, weight, bias, offload_stream)
-            return x
+            with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                return torch.nn.functional.linear(input, weight, bias)
 
 CUBLAS_IS_AVAILABLE = False
 try:
@@ -884,10 +870,8 @@ if CUBLAS_IS_AVAILABLE:
                 return None
 
             def forward_comfy_cast_weights(self, input):
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                x = cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
-                uncast_bias_weight(self, weight, bias, offload_stream)
-                return x
+                with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
+                    return cublas_half_matmul(input, weight, bias, self._epilogue_str, self.has_bias)
 
             def forward(self, *args, **kwargs):
                 run_every_op()
@@ -1207,29 +1191,28 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 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(
+                if not weight_only_quant:
+                    with CastBiasWeightContext(
                         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
+                    ) as (weight, bias):
+                        return self._forward(input, weight, bias)
+
+                with CastBiasWeightContext(
+                    self,
+                    input=None,
+                    dtype=self.weight.dtype,
+                    device=input.device,
+                    bias_dtype=input.dtype,
+                    offloadable=True,
+                    compute_dtype=compute_dtype,
+                    want_requant=True,
+                ) as (weight, bias):
+                    weight = weight.to(dtype=input.dtype)
+                    return self._forward(input, weight, bias)
 
             def forward(self, input, *args, **kwargs):
                 run_every_op()
@@ -1249,25 +1232,20 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
 
                 # Training path: quantized forward with compute_dtype backward via autograd function
                 if (input.requires_grad and _use_quantized and quantize_input):
-
-                    weight, bias, offload_stream = cast_bias_weight(
+                    with CastBiasWeightContext(
                         self,
                         input,
                         offloadable=True,
                         compute_dtype=compute_dtype,
                         want_requant=True
-                    )
+                    ) as (weight, bias):
+                        scale = getattr(self, 'input_scale', None)
+                        if scale is not None:
+                            scale = comfy.model_management.cast_to_device(scale, input.device, None)
 
-                    scale = getattr(self, 'input_scale', None)
-                    if scale is not None:
-                        scale = comfy.model_management.cast_to_device(scale, input.device, None)
-
-                    output = QuantLinearFunc.apply(
-                        input, weight, bias, self.layout_type, scale, compute_dtype
-                    )
-
-                    uncast_bias_weight(self, weight, bias, offload_stream)
-                    return output
+                        return QuantLinearFunc.apply(
+                            input, weight, bias, self.layout_type, scale, compute_dtype
+                        )
 
                 # Inference path (unchanged)
                 if _use_quantized and quantize_input:
@@ -1378,13 +1356,11 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 """Cast the whole bank once; expert_linear inside reuses the cast.
                 Not re-entrant — do not nest calls on the same instance.
                 """
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                self._resident_bank = (weight, bias)
-                try:
-                    yield self
-                finally:
-                    self._resident_bank = None
-                    uncast_bias_weight(self, weight, bias, offload_stream)
+                with CastBiasWeightContext(self, input, offloadable=True) as self._resident_bank:
+                    try:
+                        yield self
+                    finally:
+                        self._resident_bank = None
 
             def expert_linear(self, input: torch.Tensor, i: int) -> torch.Tensor:
                 """Linear against expert i's weight (with optional bias)."""
@@ -1392,11 +1368,8 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
                 if resident is not None:
                     weight, bias = resident
                     return self._expert_linear_impl(input, weight, bias, i)
-                weight, bias, offload_stream = cast_bias_weight(self, input, offloadable=True)
-                try:
+                with CastBiasWeightContext(self, input, offloadable=True) as (weight, bias):
                     return self._expert_linear_impl(input, weight, bias, i)
-                finally:
-                    uncast_bias_weight(self, weight, bias, offload_stream)
 
             def _expert_linear_impl(self, input, weight, bias, i):
                 if isinstance(weight, QuantizedTensor):
@@ -1487,17 +1460,16 @@ def mixed_precision_ops(quant_config={}, compute_dtype=torch.bfloat16, full_prec
 
                 # Optimized path: lookup in fp8, dequantize only the selected rows.
                 if isinstance(weight, QuantizedTensor) and len(self.weight_function) == 0:
-                    qdata, _, offload_stream = cast_bias_weight(self, device=input.device, dtype=weight.dtype, offloadable=True)
-                    if isinstance(qdata, QuantizedTensor):
-                        scale = qdata._params.scale
-                        qdata = qdata._qdata
-                    else:
-                        scale = None
+                    with CastBiasWeightContext(self, device=input.device, dtype=weight.dtype, offloadable=True) as (qdata, _bias):
+                        if isinstance(qdata, QuantizedTensor):
+                            scale = qdata._params.scale
+                            qdata = qdata._qdata
+                        else:
+                            scale = None
 
-                    x = torch.nn.functional.embedding(
-                        input, qdata, self.padding_idx, self.max_norm,
-                        self.norm_type, self.scale_grad_by_freq, self.sparse)
-                    uncast_bias_weight(self, qdata, None, offload_stream)
+                        x = torch.nn.functional.embedding(
+                            input, qdata, self.padding_idx, self.max_norm,
+                            self.norm_type, self.scale_grad_by_freq, self.sparse)
                     target_dtype = out_dtype if out_dtype is not None else weight._params.orig_dtype
                     x = x.to(dtype=target_dtype)
                     if scale is not None and scale != 1.0:

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"):

comfy/text_encoders/llama.py

@@ -859,16 +859,10 @@ class BaseGenerate:
         else:
             module = self.model.embed_tokens
 
-        offload_stream = None
-        if module.comfy_cast_weights:
-            weight, _, offload_stream = comfy.ops.cast_bias_weight(module, input, offloadable=True)
-        else:
-            weight = self.model.embed_tokens.weight.to(x)
-
-        x = torch.nn.functional.linear(input, weight, None)
-
-        comfy.ops.uncast_bias_weight(module, weight, None, offload_stream)
-        return x
+        if not module.comfy_cast_weights:
+            return torch.nn.functional.linear(input, self.model.embed_tokens.weight.to(x), None)
+        with comfy.ops.CastBiasWeightContext(module, input, offloadable=True) as (weight, _bias):
+            return torch.nn.functional.linear(input, weight, None)
 
     def init_kv_cache(self, batch, max_cache_len, device, execution_dtype):
         model_config = self.model.config

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):