Make a context manager for cast_bias_weight and use it.

[blepping]

This pull adds a context manager for cast_bias_weight to ensure that uncast_bias_weight is always called when necessary, even if an exception occurs. It possibly also fixes a couple potential issues:

1. fp8_linear accesses the weight after uncasting it.
2. Probably a bigger issue, the Linear module in mixed precision ops converts the dtype of the weight, but it calls uncast_bias_weight on the converted tensor rather than the original. I'm not sure what happens when you try to uncast something that wasn't cast but I'm guessing it's a potential problem.

The changes are mostly straightforward but there are a lot of code paths and I haven't tested all of them. I have checked and triple checked the patch, ran it by a few LLMs and as far as I know there aren't any current problems.

I tried to maintain the existing code style as much as possible.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro Plus

Run ID: 5a858091-aec8-4810-b5ec-4e261ba2ee67

📥 Commits

Reviewing files that changed from the base of the PR and between 77917ed and 528c44d.

📒 Files selected for processing (4)

• comfy/background_removal/birefnet.py
• comfy/controlnet.py
• comfy/ops.py
• comfy/text_encoders/llama.py

---

📝 Walkthrough

Walkthrough

Changes

A new CastBiasWeightContext context manager is added in comfy/ops.py, wrapping the existing cast_bias_weight/uncast_bias_weight pair. Numerous call sites across comfy/ops.py (Linear, Conv1d/2d/3d, GroupNorm, BatchNorm2d, LayerNorm, RMSNorm, ConvTranspose1d/2d, Embedding, FP8, cuBLAS, mixed-precision quantized ops, MoE expert weight offloading) are refactored to use the new context manager instead of manual cast/uncast calls. Similar refactors are applied in comfy/controlnet.py, comfy/background_removal/birefnet.py, and comfy/text_encoders/llama.py. No public API signatures change besides the new class.

Sequence Diagram(s)

Not applicable — the change is a mechanical refactor of internal weight-casting call patterns without introducing new external interaction flows.

Related PRs: None identified from the provided context.

Suggested labels: refactor, comfy/ops

Suggested reviewers: Reviewers familiar with comfy/ops.py weight-casting/offload internals

Poem:
A rabbit hopped through cast and weight,
Found manual uncasts, thought "not great!"
Wrapped them snug in context's care,
One tidy with, cleanup laid bare.
Hop, hop, refactor — clean and light! 🐇✨

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title accurately summarizes the main change: introducing and using a cast_bias_weight context manager.
Description check	✅ Passed	The description matches the PR’s purpose and mentions the same context-manager and weight-handling fixes.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[comfyanonymous]

@rattus128 can you take a look?

[blepping]

This could also be done with a context manager decorator, though it would require adding an import for contextlib. Something like:

@contextlib.contextmanager
def cast_bias_weight_context(*args, **kwargs):
    slf = args[0] if len(args) else None
    state = (None, None) if slf is None else cast_bias_weight(*args, **kwargs)
    del args, kwargs
    if len(state) < 3 or state[2] is None:
        state = [state]
        yield state.pop()
        return
    try:
        yield state[:2]
    finally:
        uncast_bias_weight(slf, *state)

Not sure if that's better. Using the list and popping is so the context manager can be sure it's not holding any references in the case it doesn't need to uncast.