memory use optimisations; disable grad calc for text encoder

[bghira]

tl;dr fixes memory use for text encoders not evicting fully.

This pull request introduces improvements to memory management and pipeline handling for text embedding caches and model components, with a focus on reducing VRAM usage and improving resource cleanup. The changes add mechanisms to track and clear cached pipelines and text encoders, enhance reporting of CUDA memory usage, and ensure that all text embedding caches are properly cleared, including those not directly tracked by the main data backend registry. Additionally, new tests are provided to verify the correct unloading of text encoders and cache components.

Memory and resource management improvements:

• Added a WeakSet registry to the TextEmbeddingCache class to track all active cache instances and a new active_caches class method for retrieving them. This enables comprehensive cleanup of all caches, including those not referenced by the main backend registry. [1] [2]
• Implemented the _report_cuda_usage method in Trainer to log which model and cache components still hold CUDA tensors, aiding in pinpointing VRAM leaks. Also added _clear_pipeline_caches to drop cached pipelines before training.
• Enhanced the init_unload_text_encoder method in Trainer to clear text encoders and pipelines from all TextEmbeddingCache instances, not just those in the main registry, and to report memory usage before and after unloading.

Pipeline initialization and placeholder logic:

• Updated pipeline initialization in ModelFoundation to use placeholder modules when the base model is not loaded, preventing errors from unavailable components. Also improved scheduler handling by reconstructing or deep-copying the scheduler if needed. [1] [2] [3] [4] [5]
• Modified the encode_wan_prompt method in TextEmbeddingCache to obtain pipelines via the model using the new registry and placeholder logic, improving robustness.

Testing and reliability:

• Added a new test (tests/test_text_encoder_unload.py) to verify that untracked text embed caches are properly cleared when unloading text encoders, ensuring reliability of the new cleanup logic.

These changes collectively improve the robustness of memory management, resource cleanup, and pipeline initialization in the training workflow, helping to prevent VRAM leaks and ensuring that all caches are properly handled.