model_patcher: synchronize and release offloaded Hook weights

[liminfei-amd]

Summary

Keep plain-weight Hook delegates on synchronous transfers after Hook writeback
changes prepared model storage, while preserving each module's declared model
dtype. Release the retained ordinary Hook model when DynamicVRAM actually
unloads its owning model.

This was discovered while following up #14382 and is independent of #14413's
quantized state-key fix.

Problem

BF16 model LoRAs work through LoraLoaderModelOnly, but the equivalent mapped
WeightHook can complete with an all-black/non-finite image after the model has
been prepared for DynamicVRAM offload.

The Hook path uses a plain delegate. After mapped weight writeback changes
prepared storage, that delegate still uses model-level pinned asynchronous
transfers. A mapped zero-valued Hook also fails, while an unmatched zero-patch
Hook is normal.

After the transfer fix, the reporter confirmed full BF16 Hook sampling works,
but the retained delegate survived model unload: VAE load reported
0 models unloaded, left about 20.96 GiB allocated on a 24 GiB card, and both
regular and tiled decode OOMed. Restarting and decoding the saved latent worked.

Change

• Keep a plain-weight Hook delegate on synchronous weight transfers after its
  first mapped write.
• Restore declared model dtype casting where storage dtype differs.
• Leave quantized setter-backed weights on their separate replacement path.
• Keep the ordinary delegate for normal Hook reuse, but finalize it after
  free_memory() removes the real LoadedModel.
• Finalize only models that were actually removed; keep_loaded models are
  unchanged.

Tests

The PR branch passes 13/13 focused tests covering transfer selection, model
dtype restoration, delegate retention/recreation, post-removal finalization,
and multi-model keep_loaded. Its full CPU suite is baseline-equivalent:
1232 passed, 10 skipped, 2 failed, with both failures in container MIME
registration tests.

The combined test branch (#14413, #15532, #15533, and #15456) passes 32/32
focused/VAE tests and a baseline-equivalent full suite
(1241 passed, 10 skipped, 2 failed).

On one physical gfx1100 GPU with PyTorch 2.9.1 / ROCm 7.2.1 and
comfy-kitchen==0.2.30:

• before the unload change, unload_all_models() left the delegate 1 -> 1
  and following ordinary prompts took about 29 seconds;
• after the change, the delegate reaches zero and following prompts return to
  15-16 seconds;
• Hook/no-LoRA/ordinary outputs for seeds 12345-12348 remain pixel-identical,
  every final dynamic pin count is zero, and logs contain no Hook, LoRA, OOM,
  HIP, traceback, or index error;
• the mapped BF16 subset remains pixel-identical to the prior fix and finite,
  nonblack.

The earlier PyTorch 2.12 / ROCm 7.13 validation also confirms that no-LoRA
and ordinary BF16 controls remain finite, a mapped zero-valued Hook matches
the no-LoRA control pixel-for-pixel, and two repeated mapped-Hook prompts
complete without the previous black-image false success.

A later Hook recreates its delegate in unit coverage. A full
Hook -> unload -> Hook GPU run reaches the second delegate creation without a
code error but exceeds the available 30 GiB host-memory envelope.

The #14382 reporter
then completed two consecutive full BF16 sampling and VAE runs on the original
RX 7900 XTX workflow using combined branch dad437146b. This closes the
reporter boundary for VAE decode and a second generation after the real unload.
The branch also contains #15456, so this is end-to-end integration evidence,
not an isolated claim for #15456's generic VAE retry.

This draft branch is stacked on #14413 so the combined follow-up remains
testable before that PR merges. The combined test branch additionally includes
#15532 and existing VAE retry PR #15456; #15456's comfy/sd.py changes are not
part of this PR.

Addresses #15531.

---

This is an experimental, AI-generated code change, reviewed by AMD engineers before submission.

Made with Cursor