- Set VAE to BFloat16 on MPS to prevent dtype mismatch errors
- Add override_dtype parameter to model loading pipeline
- Convert VAE weights to target dtype during load for efficiency
### 🚀 Performance Optimizations
#### Model Initialization & Loading
- **Faster model initialization** using meta device to avoid unnecessary memory allocation during DiT and VAE creation
- **Memory-mapped VAE loading** with `assign=True` to reference tensors directly instead of copying, matching DiT's loading strategy
- **Optimized state dict management** with explicit deletion after loading to free memory immediately
- **Removed slow preinitialize_rope_cache()** bottleneck during model preparation
#### Generation Pipeline
- **Streamlined batch processing** by merging redundant log steps and unifying "Model Configuration" and "Input Preparation" into single "Generation Setup" step
- **Optimized memory usage** by passing compute_dtype through entire diffusion pipeline, ensuring timesteps use bfloat16 instead of float32
- **Code deduplication** by reusing existing functions: `prepare_video_transforms()`, `load_text_embeddings()`, `calculate_optimal_batch_params()`
- **Improved batch recommendations** with `calculate_optimal_batch_params()` now prioritizing temporal stability by suggesting largest valid 4n+1 batch size
#### VAE Operations
- **Optimized tile blending** with pre-computed cosine ramps instead of recalculating per-tile
- **In-place operations** using `mul_` and `div_` to reduce memory allocations
- **Reduced redundant transfers** by skipping device moves when tensors are already on correct device
- **Improved slicing operations** with consistent memory cleanup for both encode and decode paths
### 💾 Memory Management
#### Critical Fixes
- **Fixed major VRAM leak** in VAE decoder by clearing InflatedCausalConv3d memory buffers when moving VAE to CPU
- **Fixed RoPE OOM issues** by computing only required dimensions instead of maximum (1024x128x128), enabling higher resolution upscaling
- **Added OOM retry mechanism** for VAE operations (Upsample3D, ResnetBlock3D, InflatedCausalConv3d, GroupNorm) with single retry on failure
- **Added cuBLAS workspace clearing** for better GPU memory management
#### BlockSwap Integration
- **Full preserve_vram compatibility** with bypass mechanism to temporarily disable BlockSwap protection during offload
- **Transparent handling** in `manage_model_device()` for BlockSwap models
- **Proper state restoration** of blocks to correct devices (partial GPU config) on reload
- **Simplified preserve_vram logic** by removing BlockSwap checks from infer.py
### 🏗️ Architecture Refactoring
#### Code Organization
- **Unified model configuration** by merging `configure_dit_model_inference` and `configure_vae_model_inference` into single `configure_model_inference` function
- **Focused helper functions** extracted from `configure_model_inference()`
- **Removed redundant code**: gradient checkpointing (training-only), model_weight variables, platform-specific VAE dtype logic
- **Simplified FP8 logic** by removing unnecessary conversion code from `load_quantized_state_dict`
#### Dtype Handling
- **Unified dtype flow**: VAE now uses configured dtype consistently, removing redundant vae_dtype variable
- **Fixed FP16 black frames** by using BFloat16 compute/autocast for FP16 models
- **Natural dtype propagation** from generation_loop → configure_diffusion → sampling, eliminating manual conversions
- **Removed nested autocast contexts** relying on outer context from generation_step
### 🐛 Bug Fixes & Improvements
#### Debug & Logging
- **Optimized debug performance** by making tensor counting opt-in (`show_tensors=False` by default)
- **Conditional logging** with `reset_vram_peak()` only logging when debug is enabled
- **Early return pattern** in log() method to avoid string operations when disabled
- **Standardized formatting**: consistent .2f float precision, uppercase device names, unified category names
- **Improved clarity**: better frame padding logs, clearer FP8 RoPE conversion messages
- Change log_memory_state() to default show_tensors=False (opt-in tensor counting), skipping expensive tensor counting during intermediate operations
- Enable detailed tensor analysis only at final cleanup
- Add early return pattern in log() method to avoid string operations when disabled
- Make reset_vram_peak() logging conditional on debug.enabled state
- Remove unused vram_info call in model_manager.py
- Add InflatedCausalConv3d memory cleanup to slicing_decode to match slicing_encode
- Skip redundant device transfers when tensors are already on correct device
- Pre-compute cosine ramps for tile blending instead of recalculating per-tile
- Use in-place operations (mul_, div_) to reduce memory allocations
- Reduce debug logging frequency and clarify tile progress messages
- Filter modules with memory before clearing to avoid unnecessary iterations
Minor performance improvements with no memory increase.
- Fix NaMMRotaryEmbedding3d to compute only required dimensions instead of maximum (1024x128x128), enabling much higher resolution upscaling with 3B model without running OOM
- Remove slow preinitialize_rope_cache() which was causing bottleneck during model preparation - no longer needed with NaMMRotaryEmbedding3d fixed
- Clean up unnecessary memory clearing calls and improve debug logging clarity
Dtype Consistency:
- Pass compute_dtype through entire diffusion pipeline to ensure timesteps use bfloat16 instead of float32, reducing memory usage and eliminating dtype conversion overhead
- Propagate dtype from generation_loop → configure_diffusion → sampling timesteps creation
Code Deduplication:
- Removed duplicated code and made use of existing reusable functions: prepare_video_transforms(), load_text_embeddings(), calculate_optimal_batch_params()
- Improve calculate_optimal_batch_params() to prioritize temporal stability by recommending largest valid 4n+1 batch size
Memory Optimizations:
- Remove unnecessary GPU↔CPU offloading for text embeddings and timesteps during preserve_vram mode (added overhead without meaningful memory benefits)
- Simplify text embeddings cleanup to work directly with dictionary structure
- Add cuBLAS workspace clearing for better GPU memory management
Code Cleanup:
- Remove redundant self.runner = None in _internal_execute (handled by cleanup())
- Enhance user messaging for batch padding waste with clearer explanations
- Add bypass mechanism to temporarily disable BlockSwap protection during offload
- Refactor manage_model_device() to handle BlockSwap models transparently
- Restore blocks to correct devices (partial GPU config) on reload
- Simplify preserve_vram logic in infer.py by removing BlockSwap checks
This allows BlockSwap to work correctly with preserve_vram, offloading the entire model to CPU after inference and restoring only necessary blocks to GPU before next inference.
- Use meta device initialization to avoid unnecessary memory allocation during model creation
- Reduces DiT and VAE initialization time by ~90% when loading to CPU
- Explicitly delete state dicts after loading to free memory immediately
- Refactor configure_model_inference() into focused helper functions for DRY code
- Improve debug logging with timestamps and cleaner section separators for readability
- Implement retry_on_oom() helper with single retry for all VAE operations (Upsample3D, ResnetBlock3D, InflatedCausalConv3d, GroupNorm)
- Fix regression: use BFloat16 compute/autocast for FP16 models to prevent black frames
- Improve frame padding logs to clarify model constraint (frames % 4 == 1)
- Remove hardcoded target_dtype and manual tensor conversions in inference()
- Remove redundant inner autocast context - rely on outer context from generation_step
- Align compute_dtype and autocast_dtype for FP16 models (both now float16)
- Remove unused block_swap_config parameter
- Fix device logging to properly display device names
The dtype is now determined once in generation_loop based on model weights and flows naturally through the pipeline via the outer autocast context, eliminating unnecessary conversions and improving performance.
- Merge generation setup log steps: combine "Model Configuration" and "Input Preparation" into unified "Generation Setup" step
- Simplify dtype handling: remove redundant vae_dtype variable, VAE now uses configured dtype consistently
- Remove platform-specific VAE dtype logic (MPS special case)
- Unify VAE encode/decode: remove autocast wrapper and target_dtype parameter, both now use configured dtype
- Add consistent docstrings for vae_encode and vae_decode methods
- Improve precision logging: show both DiT and VAE dtypes, rename model_dtype to dit_dtype
- Optimize text embeddings movement: only move when preserve_vram is active (BlockSwap handles model layers separately)
- Load text embeddings directly to CPU when preserve_vram is enabled
- Update logging consistency: uppercase device names, unified category names, clearer messages
- Clarify FP8 RoPE conversion log message: specify "from FP8 to BFloat16"
- Merge configure_dit_model_inference and configure_vae_model_inference into single configure_model_inference function
- Simplify load_quantized_state_dict by removing unnecessary FP8 conversion logic
- Remove redundant gradient checkpointing call (training-only feature)
- Remove redundant model_weight variables
- Standardize float formatting to .2f across all debug logs for consistency
- Eliminate redundant device movement operations already handled during model creation
- Match DiT's load_state_dict(assign=True) for VAE to reference memory-mapped tensors directly instead of copying, reducing peak RAM
- Move post-inference memory logging after offloading for accurate metrics
Updated manage_model_device() to clear InflatedCausalConv3d memory buffers
when moving VAE to CPU. These buffers were not automatically released with
.to('cpu'), causing intermediate tensors to remain in VRAM after batch processing.
Preserve BlockSwap's GPU/CPU distribution when cache_model=True to avoid unnecessary reconfiguration on subsequent runs. Also fixes regression error where tensors were not on the correct device after caching.
Preserve BlockSwap's GPU/CPU distribution when cache_model=True to avoid unnecessary reconfiguration on subsequent runs. Also fixes regression error where tensors were not on the correct device after caching.
Regression from previous commit - cleanup_blockswap was removing forward wrappers/RoPE patches
but leaving blocks in position, then re-applying BlockSwap duplicated memory allocations
causing OOM on second generation.
Now when caching with same config:
- Blocks remain in their GPU/CPU positions for next inference
- Forward wrappers and RoPE patches stay intact
- Only toggles _blockswap_active flag for safety
- Reuses exact same VRAM allocation as previous run
Full cleanup only happens when cache_model=False or config changes.
Fixes OOM errors and eliminates redundant GPU transfers between cached generations.
#### Memory Management & Leak Fixes
- **Fixed BlockSwap pinned memory leak**: Removed `non_blocking` parameter that was preventing pinned memory release
- **Eliminated logging-induced memory leaks**: Modified `log_memory_state()` to avoid holding tensor references and added checkpoint history limits
- **Complete model cleanup**: Ensured proper deletion and tensor storage release when `cache_model=False`
- **Fixed lingering GPU tensors**: Resolved issues with scalar tensors and text embeddings remaining on GPU between batches with `preserve_vram`
#### Performance Optimizations
- **Direct-to-CPU model loading**: DiT/VAE weights load directly to CPU when `preserve_vram` or BlockSwap is active, avoiding VRAM spikes
- **Reactive OOM handling**: Replaced proactive memory clearing during VAE decode with reactive OOM recovery, fixing performance regression (removed the need for the keep_vae_loaded_during_decode flag)
- **Optimized cleanup functions**: Consolidated all cleanup into core functions (`manage_*`, `release_*`, `clear_*`)
#### Enhanced Debugging & Monitoring
- **Unified debug system**: Improved logging categories, icons, and hierarchical timers
- **Comprehensive memory tracking**: Per-interval VRAM peak tracking, accurate process-specific RAM monitoring
- **Debug propagation**: Debug instance now properly propagates to all submodules (sampler, VAE, etc.)
- **Improved log clarity**: Standardized message formatting, removed redundant logs, added critical operation snapshots
#### Code Architecture Improvements
- **Type hints **: Added type hints to all memory management functions
- **Centralized device management**: New `manage_model_device()` function for consistent device movement logic
- **Error handling**: Added try/except blocks to all critical memory and device operations
- **Removed private APIs**: Eliminated `torch._C._clear_cache()` usage that was triggering warning
#### Bug Fixes
- Resolved duplicate dtype detection in generation loop
- Corrected CPU offloading logic for text embeddings and sampling timesteps
- Fixed timer namespace collisions in batch processing
- Removed redundant RoPE cache clearing operations
#### Various
- Remove "Force move weights to device" in FP8CompatibleDiT forward() to avoid clash with blockswap - This was used for preserve_vram but will refactor preserve_vram in a separate commit
- Remove unused imports and legacy code/comments
- Remove "Force move weights to device" in FP8CompatibleDiT forward() to avoid clash with blockswap - This was used for preserve_vram but will refactor preserve_vram in a separate commit
- Remove duplicate dtype detection in generation_step (now passed from generation_loop)
- Add device checks before CPU moves to avoid redundant operations
- Improve debug.log and debug.log_memory_state() formatting, content, and categories for better visibility
- Remove unused imports and excessive clear_memory() calls
- Clean up non-essential logs from always display
This is a work-in-progress commit that consolidates a series of changes to fix memory leaks, optimize VRAM/RAM usage, improve performance, and enhance code maintainability.
* **BlockSwap Pinned Memory:** Disabled `use_non_blocking=True` for CPU-to-GPU transfers to resolve a memory leak where pinned memory was not being released.
* **Logging-Induced Leaks:** Modified `log_memory_state()` to avoid holding references to tensors during analysis and added a history limit to the checkpoint system to prevent unbounded memory growth.
* **Incomplete Model Cleanup:** Ensured models are completely deleted and their tensor storage is released when `cache_model=False`.
* **Lingering Tensors:** Fixed an issue where a scalar tensor from sampling timesteps and text embeddings remained on the GPU between batches when `preserve_vram` is active.
* **Centralized Cleanup Functions:** Introduced `clear_memory()` to replace `clear_vram_cache()` and all manual `torch.cuda.empty_cache()` calls, providing consistent VRAM/RAM cleanup logic. The function features a `full` parameter to distinguish between a fast, GPU-only cache clear (~1-5ms) for frequent operations and a full cleanup with garbage collection (~10-50ms) for critical stages.
* **Direct-to-CPU Model Loading:** Modified DiT/VAE weight loading to load directly onto the CPU when `preserve_vram` or `BlockSwap` is active, avoiding unnecessary VRAM spikes during model preparation.
* **VAE Device Management:** Created the `manage_vae_device()` helper function to centralize the logic for moving the VAE between the CPU and GPU, reducing code duplication. This also fixed a bug that incorrectly kept the VAE on the GPU when `preserve_vram` was active.
* **CPU Offloading:** Implemented logic to move text embeddings and sampling timesteps to the CPU after each batch when `preserve_vram` is active, reducing idle VRAM usage.
* **VAE Decode Performance:** Replaced proactive, frequent memory clearing during VAE decode with a reactive Out-of-Memory (OOM) handling system. This fixed a significant performance regression and eliminated the need for the `keep_vae_loaded_during_decode` flag.
* **Reduced Overhead:** Removed redundant `gc.collect()` calls from multiple locations to decrease unnecessary processing overhead.
* **Interval-Based VRAM Tracking:** Modified the logging system to reset peak VRAM statistics after each `log_memory_state()` call, enabling accurate tracking of peak memory usage for specific processing intervals (e.g., encode, inference, decode).
* **Accurate RAM Monitoring:** Added the `get_ram_usage()` function for correct process-specific RAM tracking.
* **Efficient Log Refactoring:** Refactored `log_memory_state()` into modular helper methods, optimizing tensor analysis into a single-pass `gc` iteration to improve both performance and maintainability.
* **Log Clarity:** Refined memory state and debug logging to remove redundant snapshots and add new ones for critical operations like model loading, weight loading, VAE encoding, and decoding. Standardized log message conventions.
* **Per-Batch Timers:** Implemented timer namespacing to ensure that performance timers for each batch are logged correctly without overwriting one another.
* **Error Handling:** Added `try/except` blocks to key memory and device management functions to handle edge cases and improve robustness.
* **Code Cleanup:** Removed deprecated code and outdated comments throughout the related modules.
* **Documentation:** Updated comments and function docstrings to reflect the new memory management architecture.
- Replace ambiguous 'block' terminology with 'batch merge' in generation.py to distinguish from BlockSwap transformer blocks
- Use consistent 1-based indexing for batch numbers in merge operations
- Improve merge operation messages for better user clarity
- Fix debug log categories: batch merging -> 'video', model prep -> 'model'
Remove tile size/overlap max limits, set min=64/step=32 for finer control, update tooltips for clarity, and fix runtime parameter updates when cache_model=True
Keep weights and arguments on the same device when 'preserve_vram' is ON
Restored compatibility with FP8 safetensors for MPS backend
Added multi-gpu support for ComfyUI nodes
Added FPS calculation
Set use_non_blocking to False by default to prevent memory leaks (temp fix)
Moved some parms from main node to extra args node
Adrien Toupet