fix(model cache): keep patch models alive while they are patched in

[lstein]

Stacked on #9263 — the branch sits on top of lstein/feat/multi-gpu, so the diff below includes that PR's commits. Base is main and it stays a draft until #9263 merges, at which point it is rebased and marked ready — same pattern as #9387–#9389. Review only the last commit; everything else belongs to #9263.

The problem

Patch models are used without ever locking their cache record. A patch iterator calls context.models.load(), reads .model, yields the ModelPatchRaw and lets the LoadedModel wrapper go — and LayerPatcher.apply_smart_model_patches materializes the iterator before applying anything, so every wrapper is already gone before the first patch lands, while the patches stay applied for the whole denoise.

The LoadedModel wrapper is the cache's only signal that such a record is in use: the lock count is zero, so ModelCache.release_first_use_grace (added in 7ffee4d on #9263) treats the record as abandoned as soon as the wrapper is collected, and a peer's budget reconcile is then free to evict it. Because ModelPatchRaw is not an nn.Module, _delete_cache_entry takes the remove_non_shared path and drops the bytes from the shared RamBudget while the RAM is still resident — the cap under-counts and over-admits, which is the exact failure RamBudget exists to prevent — plus a disk reload on the LoRA's next use.

This is a regression introduced by that commit: with a single LoRA there is no later put() on that cache, so the post-admission grace previously survived the whole generation.

The fix

Make the wrapper outlive the patch. LayerPatcher gains a PatchSpec type alias:

PatchSpec = Union[Tuple[ModelPatchRaw, float], Tuple[ModelPatchRaw, float, object]]

Producers pass the LoadedModel as an optional third element instead of del-ing it, and apply_smart_model_patches — which already materializes the iterable and holds it for the entire patched region — pins it for exactly as long as the patches are applied. When the patched region exits, patches is released, the wrappers die, and the grace is released normally.

Two-element specs remain valid, so callers whose patches are not cache-backed (the patcher tests, and anything owning the model's lifetime itself) are unaffected.

Why it touches so many files

There is no central place to do this: only the producer holds the LoadedModel. The change is 16 patch iterators across 15 invocation files — every model family's denoise and text-encoder path (FLUX, FLUX2, Z-Image, Qwen, Anima, WAN, SD3, SDXL/compel, tiled MD) — plus layer_patcher.py. Each edit is identical: append lora_info to the yielded tuple, drop the del, retype the iterator to Iterator[PatchSpec].

The WAN expert-swapper path works unchanged: LoRAIteratorFactory is retyped and produces a fresh iterator per swap, so each re-entry pins its own wrappers.

Tests

tests/backend/patches/test_layer_patcher_shared_weights.py:

• test_patch_spec_cache_handle_is_held_for_the_patched_region — the spec's handle must stay alive for the whole with block and be released after it. Verified to fail (the patch's cache handle was released while its patch was still applied) when the pinning is reverted by dropping the third element before materializing.
• test_two_element_patch_specs_are_still_accepted — back-compat for non-cache-backed callers.

Full suite green; ruff check and format clean.

🤖 Generated with Claude Code

[lstein]

Superseded by JPPhoto's 1142430 on the #9263 branch, which lands a strictly stronger version of this fix: instead of keeping the LoadedModel wrapper alive (which only preserved the post-admission grace), the patch specs now carry a model_in_ram() context manager that takes a real cache lock for the patched region — protecting warm records from peer eviction too, across the same producers plus the model families added since (anima, krea2, flux2). Reviewed and hardened in d63c1f3 / 272a2ef.