Commit Graph
33 Commits
Author SHA1 Message Date
Adrien Toupet 2911b78288 feat: Separate Flash Attention 2/3 and SageAttention 2/3 backends
- Rename attention modes: flash_attn→flash_attn_2/3, sa2/sa3→sageattn_2/3
- Add separate detection and wrappers for FA2, FA3, SA2, SA3 in compatibility.py
- FA3: Filter unsupported params (dropout_p, window_size), return tuple[0]
- SA2/SA3: Add half-precision dtype handling (convert fp32/fp8→bf16)
- SA3: Add varlen-to-batched conversion with SA2 fallback for non-uniform seqs
- Add fallback chains: FA3→FA2→SDPA, SA3→SA2→SDPA
- Update debug.py to show granular availability: FlashAttn / SageAttn
- Update all references: README, CLI, ComfyUI nodes, docstrings
2025-12-10 15:26:16 -05:00
Adrien Toupet c010deeea1 Remove ineffective allow_vram_overflow setting
- PyTorch's set_per_process_memory_fraction cannot prevent WDDM paging on Windows
- Keep overflow detection and warning when VRAM exceeds physical limit
- Simplify peak memory formatting
- Remove setting from CLI, ComfyUI node, and memory_manager
2025-12-10 00:56:27 -05:00
Adrien Toupet 7cbf025561 Fix VRAM peak tracking: separate allocated vs reserved, Windows-only overflow
- Track both peak_allocated (tensor usage) and peak_reserved (cache pool) per phase
- peak_allocated resets properly between phases via reset_peak_memory_stats()
- Overflow detection/warnings now Windows-only (WDDM paging behavior)
- Remove get_memory_architecture() - replaced with simple is_mps + platform checks
- Phase summary shows: VRAM XGB allocated, YGB reserved | RAM ZGB
- Simplify MPS path (unified memory has no overflow concept)
2025-12-09 23:51:51 -05:00
Adrien Toupet 5c60716c47 Refactor: centralize backend detection, fix architecture-aware VRAM overflow reporting 2025-12-09 21:06:12 -05:00
Adrien Toupet 77a00f651a Fix: OOM regression from 2.5.14 strict VRAM limit (#367)
Add allow_vram_overflow option (default: False) to make strict VRAM limit configurable.

The 2.5.14 change 'Enforce physical VRAM limit' prevented PyTorch from
overflowing to system RAM, causing OOM on workflows that previously
worked.

- Add allow_vram_overflow parameter to DiT Model Loader node
- Add --allow_vram_overflow CLI flag
- Show warning when enabled, track mid-session changes
- Suppress swap detection warning when user explicitly allows overflow

Note: Enabling overflow is a last resort - performance degrades severely
when physical VRAM is exceeded. Optimizing settings (BlockSwap, VAE tiling,
batch size, resolution, model size...) is always recommended.
2025-12-09 17:12:10 -05:00
Adrien Toupet 30bc924043 Update header logo design (thanks @naxci1, closes #378) 2025-12-09 14:07:46 -05:00
Adrien Toupet 18b44d66e1 feat: add environment info display in debug mode to help with issue reporting 2025-12-05 11:11:18 -05:00
Adrien Toupet ff937756a3 Add VRAM swap detection - show GPU+swap breakdown in peak stats, warn when swap detected 2025-11-30 23:47:20 -05:00
Adrien Toupet 849f6b5997 Add peak RAM tracking alongside VRAM in debug summary 2025-11-28 15:32:07 -05:00
Adrien Toupet be2efd474f fix: use canonical torch.backends.mps.is_available() for reliable MPS detection
Replaces torch.mps.is_available() with torch.backends.mps.is_available()
across all files. The torch.backends API is the official PyTorch method
for MPS detection (since PyTorch 1.12) and works reliably on:
- macOS with Apple Silicon (returns True when MPS available)
- macOS without MPS support (returns False)
- Windows/Linux (hasattr guard prevents AttributeError)
2025-11-28 11:49:43 -05:00
Adrien Toupet 40d37b7248 Release v2.5.9: Bug fixes and enhancements
- Fix: OpenCV memory layout error in tile debug visualization (#283)
- Fix: macOS MPS allocator fallback for safetensors loading (#290)
- Fix: Windows log buffering with flush=True (#278)
- Fix: ComfyUI registry icon URL (raw.githubusercontent)
- Feature: Version display in node name and CLI/ComfyUI header
- Feature: GitHub Sponsors support link
- License: Migrate from MIT to Apache 2.0 to match Bytedance Seed official repo
2025-11-12 14:27:40 -05:00
Adrien Toupet 9715d3e37a Fix: torch.mps AttributeError on Windows
Add defensive checks for torch.mps.is_available() to handle PyTorch versions where the method doesn't exist on non-Mac platforms. Resolves AttributeError: module 'torch.mps' has no attribute 'is_available'
2025-11-07 23:37:35 -05:00
Adrien Toupet afd5950d37 feat: torch.compile optimization and memory tracking improvements
- Eliminate graph breaks in na.py and attention.py for full torch.compile support
  * Replace cumsum-based tensor slicing with _tensor_split to avoid .item() calls
  * Use torch.tensor_split with .long().cpu() for PyTorch API requirements
  * Replace reshape-based averaging with split-stack-mean pattern
- Fix phase peak VRAM tracking to capture peaks during OOM retry cycles
- Standardize phase4 naming to 'postprocessing' for consistency
- Update RoPE docstrings for accuracy
- Update example workflows with icon
2025-11-06 17:48:39 -05:00
Adrien Toupet f8627298f7 Add peak VRAM tracking and summary display by phase 2025-11-05 21:21:23 -05:00
Adrien Toupet ad020d3803 feat(cli): improve UX with auto-format detection, FPS tracking, and consistent messaging with ComfyUI implementation
- Auto-detect output format per file type (mp4 for videos, png for images)
- Add visual separators between processed files for better readability
- Simplify FPS calculation to use wall-clock time for real-world throughput
- Consolidate banner/footer into shared Debug methods
- Update offload device args to support multi-GPU (cpu/cuda:N)
- Standardize terminology: 'upscaling' instead of 'video upscaling'
- Remove code duplication between CLI and ComfyUI implementations
- Consistent quote style (double quotes) throughout CLI
2025-11-04 11:54:57 -05:00
Adrien Toupet 9893a1d7b0 refactor: Add indent_level parameter to debug logging system
- Add indent_level param to debug.log() and manage_tensor()
- Replace all hardcoded spaces with indent_level (0/1/2)
- Fix encode_all_batches memory flow: encode before storage
- Minor whitespace cleanup
2025-10-19 10:13:17 -04:00
Adrien Toupet 97c7b9cd12 refactor: Overhaul device and memory management architecture
Major improvements to device handling, model offloading, and code clarity:

Device Management:
- Replace string-based devices with torch.device objects throughout
- Add explicit offload device parameters: dit_offload_device, vae_offload_device, tensor_offload_device
- Remove preserve_vram in favor of explicit offload control
- Improve get_device() to return torch.device objects consistently
- Enhance get_device_list() with smart MPS-only system handling

Context & Pipeline:
- Merge setup_device_environment and prepare_generation_context into single setup_generation_context
- Simplify LOCAL_RANK handling (set to '0' for single-GPU mode)
- Store device configuration on runner for submodule access

Tensor & Model Management:
- Add manage_tensor_device() for consistent tensor movement with logging
- Update manage_model_device() to use torch.device objects
- Add validation for BlockSwap and caching configurations
- Rename cache_in_ram to cache_model (more accurate naming)

CLI & Interface:
- Remove --preserve_vram flag (breaking change)
- Add --vae_offload_device and --tensor_offload_device flags
- Update ComfyUI node parameters with validation and better tooltips
- Improve device selection UI with offload_device options
2025-10-17 17:24:53 -04:00
Adrien Toupet e48c2da776 feat: global model cache, alpha upscaling rewrite, and quality improvements
Major refactor introducing:

- GlobalModelCache for cross-node model sharing with dynamic config updates
  * Models cached by node ID, reused across different upscaler instances
  * Config changes (torch.compile, BlockSwap, tiling) handled dynamically
  * Runner templates cached when both models present

- Complete alpha channel processing rewrite
  * New edge-guided alpha upscaling in src/core/alpha_upscaling.py
  * Remove broken VAE RGBA adapter approach
  * Alpha properly extracted, upscaled, and merged with RGB

- BlockSwap enhancements
  * Support I/O-only swapping without transformer block offloading
  * Separate memory reporting for I/O components vs transformer blocks
  * Clearer logging of GPU vs CPU placement
  * New is_blockswap_enabled() utility

- Memory management improvements
  * Add release_tensor_collection() for batch tensor cleanup
  * Proper tensor memory release throughout pipeline
  * Better cleanup in postprocess phase

- Code quality and refactoring
  * Extensive docstrings with type hints across all modules
  * Generic _update_model_config() reduces code duplication
  * Better function signatures and parameter documentation
  * Improved separation of concerns in model_manager.py
  * Clearer debug logging categories and messages
2025-10-15 00:51:56 -04:00
Adrien Toupet d6e04aa108 refactor: modularize node architecture and enhance BlockSwap flexibility
Core Changes:
- Split monolithic comfyui_node.py into modular per-node files:
  * video_upscaler.py - main upscaler node
  * dit_model_loader.py - DiT model configuration
  * vae_model_loader.py - VAE model configuration
  * torch_compile_settings.py - torch.compile settings
  * __init__.py - centralized node registry

BlockSwap Enhancements:
- Fixed swap_io_components to work independently of blocks_to_swap
- Both features can now be enabled separately or together
- Updated validation logic throughout (model_manager, blockswap, generation)
- Improved config description to reflect combined capabilities

Logging Improvements:
- Standardized indentation and improved debug messages clarity
- Added footer with support links (YouTube, GitHub) for community engagement
2025-10-11 17:37:58 -04:00
Adrien Toupet 9ee244d71b feat: implement torch.compile optimization with BlockSwap compatibility
Core Features:
- Add torch.compile support for DiT (20-40% speedup) and VAE (15-25% speedup)
- Ensure BlockSwap compatibility by applying BlockSwap before torch.compile
- Add SeedVR2TorchCompileSettings node for ComfyUI configuration
- Add CLI arguments: --compile_dit, --compile_vae, --compile_backend, --compile_mode, --compile_fullgraph, --compile_dynamic, --compile_dynamo_cache_size_limit, --compile_dynamo_recompile_limit

Optimizations:
- Optimize na.py and other backend files for torch.compile compatibility (replace .tolist() with tensor operations)
- Remove .item() calls from attention modules (max_seqlen_q/k now accept tensors)
- Add @torch._dynamo.disable decorators to timing/debug methods to prevent compilation warnings
- Replace torch.repeat with torch.repeat_interleave in modulation.py for better performance

Bug Fixes:
- Fix timing stack cleanup in memory_manager.py (end timer before early returns)
- Fix BlockSwap timing reporting (_get_swap_start_time and _log_swap_timing now properly excluded from compilation)

Code Quality:
- Add comprehensive type hints and docstrings across all modified modules
- Standardize BlockSwap timer names (blockswap_block_*, blockswap_io_*)
2025-10-10 12:34:49 -04:00
Adrien Toupet 1518ecc1a5 refactor: Fix ComfyUI node conflicts via relative imports and clearer model structure
- Renamed model directories for clarity: dit -> dit_7b, dit_v2 -> dit_3b
- Converted all absolute imports to relative imports throughout codebase
- Removed sys.path.append() manipulations that caused namespace conflicts
- Updated YAML configs to reference renamed model directories
- Simplified model variant detection logic using new directory names
- Standardized function calls with named arguments for better clarity
- Fixed generation context initialization and interrupt handling

This resolves import conflicts with other ComfyUI nodes (e.g., Basic data handling)
that use sys.path manipulation, making the module properly isolated and compatible.

Fixes #29, #114, #136
2025-09-16 16:45:31 -04:00
Adrien Toupet 477f57fd5a Refactor: Three-phase batch processing pipeline for improved performance
Major architectural change to minimize model swapping overhead by processing
all batches in three distinct phases instead of sequential per-batch processing:
- Phase 1: Encode all batches with VAE
- Phase 2: Upscale all latents with DiT
- Phase 3: Decode all latents with VAE

Core changes:
- Split monolithic generation_loop into modular functions:
  - prepare_generation_context(): Shared state management
  - setup_device_environment(): Device configuration
  - prepare_runner(): Model loading with cache support
  - encode_all_batches(): Batch VAE encoding
  - upscale_all_batches(): Batch DiT upscaling
  - decode_all_batches(): Batch VAE decoding
- Removed generation_step function (logic integrated into upscale phase)
- Added lazy precision initialization to avoid redundant setup

Performance improvements:
- Pre-allocated lists for memory efficiency
- Better cleanup of intermediate storage between phases
- Added unique timer names to clear_memory() to avoid naming conflicts
- Improved model state management with change detection and caching

UI/UX enhancements:
- Switched to ComfyUI's native ProgressBar with weighted phase progress
- Changed from per-batch FPS to overall average FPS (always visible)
- Improved log clarity with clear phase separators
- Added ASCII art logo to clearly identify SeedVR2 process start
- Better progress tracking with weighted percentages across three phases

Code cleanup:
- Removed deprecated timer_context from Debug class
- Removed unused time imports across multiple files
- Fixed LOCAL_RANK environment variable to handle string conversion properly
- Improved error handling with try/except/finally blocks in all phases
2025-09-16 14:38:57 -04:00
Adrien Toupet 66bcf436ab perf: optimize debug logging to reduce performance overhead
- 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
2025-08-28 23:50:49 -04:00
Adrien Toupet 89bde29cbe perf: optimize model initialization with meta device for 90% speedup
- 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
2025-08-26 13:23:30 -04:00
Adrien Toupet 5595d58597 refactor: streamline generation pipeline and improve dtype handling
- 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"
2025-08-25 17:11:29 -04:00
Adrien Toupet 47552c0ed5 refactor: unify model configuration and improve logging precision
- 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
2025-08-25 13:45:13 -04:00
Adrien Toupet 8333fb856e refactor(WIP): complete memory management overhaul with proper error handling
Memory Management:
- Add type hints to all memory functions for better IDE support and maintainability
- Replace silent exception handling with debug logging across all operations
- Removed unnecessary CPU transfers for GPU memory release
- Introduce unified manage_model_device() for consistent device management
- Remove torch._C._clear_cache() private API usage (incompatible across PyTorch versions)

Performance & Debugging:
- Add debug timers to critical operations (clear_memory, clear_runtime_caches, etc.)
- Consolidate all cleanup code into core functions: manage_*, release_*, clear_*, complete_cleanup
- Improved timer log messages for clarity

Code Quality:
- Remove unused imports
- Remove redundant RoPE cache clearing (clear_runtime_caches handles it)
- Simplify configure_runner() by eliminating duplicate code paths
- Update infer.py to use generic device management functions
- Add release_text_embeddings() helper to deduplicate embedding cleanup
2025-08-23 01:13:45 -04:00
Adrien Toupet 2053a80f38 refactor(WIP): Improve debug logging consistency and reduce redundancy
- 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
2025-08-22 15:38:00 -04:00
Adrien Toupet 9b7c113681 refactor(WIP): Memory Management Overhaul
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.
2025-08-22 09:20:41 -04:00
lihaoyun6 176562f829 Restored compatibility with FP8 safetensors for MPS backend; Changed the way to detect MPS device 2025-08-13 00:53:05 +08:00
Benjamin Herb fe67b2a042 feat: Add BlockSwap support for standalone CLI and update logging
- Add BlockSwap CLI arguments and activation
- Fix imports to work without ComfyUI dependencies
- Update to new debugging framework
2025-08-09 02:27:42 +02:00
Adrien Toupet 1da2cdaf12 Fix VRAM memory leak when using cache_model, added log_memory_state detailed_tensors debug output for troubleshooting 2025-08-08 00:51:02 +02:00
Adrien Toupet d075dde799 feat: Extract model caching to main node + add hierarchical debug system
- Move cache_model and enable_debug options from BlockSwap to main SeedVR2 node

- cache_model: Keep models in RAM between runs (skip reload for faster iterations)

- enable_debug: Show detailed memory/timing tracking with hierarchical display

- Replace hardcoded debug statements with unified Debug class

- Add visual icons for log categorization

- Implement parent-child timing relationships for operation breakdown

- Add VRAM/RAM usage tracking with before/after comparisons

Model caching now available for all workflows, not just BlockSwap users.

Debug system provides granular performance insights when enabled.
2025-08-08 00:49:46 +02:00