Merge pull request #441 from AInVFX/main

v2.5.24: Restores the MPS memory leak workaround that was accidentally removed during code cleanup in v2.5.23

.github/FUNDING.yml

@@ -0,0 +1 @@
+github: [adrientoupet, numz]

CONTRIBUTING.md

@@ -16,15 +16,9 @@ We welcome code contributions from the community. Here's how to propose changes:
 1. Fork this repository to your own GitHub account.
 2. Create a new branch on your fork for your changes.
 3. Make your changes in this branch.
-4. When you are ready, submit a pull request to either:
+4. When you are ready, submit a pull request to the **`main`** branch.
-   - **`main` branch** - For stable features and bug fixes
-   - **`nightly` branch** - For experimental features or if you want faster integration (note: this branch changes more frequently and may be less stable)
 
-### Branch Guidelines
+We use the GitHub Flow workflow.
-- **Main Branch**: Use this for production-ready changes, critical bug fixes, and well-tested features
-- **Nightly Branch**: Use this for experimental features, cutting-edge improvements, or if you need your changes integrated quickly. Be aware that this branch is actively under development and may contain unstable features.
-
-We use the GitHub Flow workflow. Choose the appropriate target branch based on the nature of your contribution.
 
 Before submitting a pull request, please make sure your code adheres to the project's coding conventions and it has passed all tests. If you are adding features, please also add appropriate tests.
 

LICENSE

@@ -1,21 +1,201 @@
-MIT License
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
 
-Copyright (c) 2025 NumZ, AInVFX and Contributors
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
 
-Permission is hereby granted, free of charge, to any person obtaining a copy
+   1. Definitions.
-of this software and associated documentation files (the "Software"), to deal
-in the Software without restriction, including without limitation the rights
-to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-copies of the Software, and to permit persons to whom the Software is
-furnished to do so, subject to the following conditions:
 
-The above copyright notice and this permission notice shall be included in all
+      "License" shall mean the terms and conditions for use, reproduction,
-copies or substantial portions of the Software.
+      and distribution as defined by Sections 1 through 9 of this document.
 
-THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+      "Licensor" shall mean the copyright owner or entity authorized by
-IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+      the copyright owner that is granting the License.
-FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+
-AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+      "Legal Entity" shall mean the union of the acting entity and all
-LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+      other entities that control, are controlled by, or are under common
-OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+      control with that entity. For the purposes of this definition,
-SOFTWARE.
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "{}"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2025 seed
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -4,7 +4,7 @@
 
 Official release of [SeedVR2](https://github.com/ByteDance-Seed/SeedVR) for ComfyUI that enables high-quality video and image upscaling.
 
-Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#️-run-as-standalone-cli) section.
+Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#-run-as-standalone-cli) section.
 
 [![SeedVR2 v2.5 Deep Dive Tutorial](https://img.youtube.com/vi/MBtWYXq_r60/maxresdefault.jpg)](https://youtu.be/MBtWYXq_r60)
 
@@ -14,8 +14,8 @@ Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#
 
 ## 📋 Quick Access
 
-- [🆙 Future Releases](#-future-releases)
+- [🆙 Future Work](#-future-work)
-- [🚀 Updates](#-updates)
+- [🚀 Release Notes](#-release-notes)
 - [🎯 Features](#-features)
 - [🔧 Requirements](#-requirements)
 - [📦 Installation](#-installation)
@@ -26,7 +26,7 @@ Can run as **Multi-GPU standalone CLI** too, see [🖥️ Run as Standalone](#
 - [🙏 Credits](#-credits)
 - [📜 License](#-license)
 
-## 🆙 Future Releases
+## 🆙 Future Work
 
 We're actively working on improvements and new features. To stay informed:
 
@@ -34,7 +34,156 @@ We're actively working on improvements and new features. To stay informed:
 - **💬 Join the Community**: Learn from others, share your workflows, and get help in the [Discussions](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions)
 - **🔮 Next Model Survey**: We're looking for community input on the next open-source super-powerful generic restoration model. Share your suggestions in [Issue #164](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues/164)
 
-## 🚀 Updates
+## 🚀 Release Notes
+
+**2025.12.24 - Version 2.5.24**
+
+- **🍎 Fix: MPS memory leak regression** - Restored MPS cache clearing after VAE encode/decode operations that was accidentally removed during code cleanup in v2.5.23
+
+**2025.12.24 - Version 2.5.23**
+
+- **🔒 Security: Prevent code execution in model loading** - Added protection against malicious .pth files by restricting deserialization to tensors only
+- **🎥 Fix: FFmpeg video writer reliability** - Resolved ffmpeg process hanging issues by redirecting stderr and adding buffer flush, with improved error messages for debugging *(thanks [@thehhmdb](https://github.com/thehhmdb))*
+- **⚡ Fix: GGUF VAE model support** - Enabled automatic weight dequantization for convolution operations, making GGUF-quantized VAE models fully functional *(thanks [@naxci1](https://github.com/naxci1))*
+- **🛡️ Fix: VAE slicing edge cases** - Protected against division by zero crashes when using small split sizes with high temporal downsampling *(thanks [@naxci1](https://github.com/naxci1))*
+- **🎨 Fix: LAB color transfer precision** - Resolved dtype mismatch errors during video upscaling by ensuring consistent float types before matrix operations
+- **🔧 Fix: PyTorch 2.9+ compatibility** - Extended Conv3d memory workaround to all PyTorch 2.9+ versions, fixing 3x VRAM usage on newer PyTorch releases
+- **📦 Fix: Bitsandbytes compatibility** - Added ValueError exception handling for Intel Gaudi version detection failures on non-Gaudi systems
+- **🍎 MPS: Memory optimization** - Reduced memory usage during encode/decode operations on Apple Silicon *(thanks [@s-cerevisiae](https://github.com/s-cerevisiae))*
+
+
+**2025.12.13 - Version 2.5.22**
+
+- **🎬 CLI: FFmpeg video backend with 10-bit support** - New `--video_backend ffmpeg` and `--10bit` flags enable x265 encoding with 10-bit color depth, reducing banding artifacts in gradients compared to 8-bit OpenCV output *(based on PR by [@thehhmdb](https://github.com/thehhmdb) - thank you!)*
+- **🍎 Fix: MPS bicubic upscaling compatibility** - Added CPU fallback for bicubic+antialias interpolation on PyTorch versions before 2.8.0, resolving RGBA alpha upscaling errors on Apple Silicon
+- **⚡ Fix: Cross-platform histogram matching** - Replaced scatter_ operation with argsort+index_select for improved reliability across CUDA, ROCm, and MPS backends
+- **🧹 MPS: Remove sync overhead** - Reverted unnecessary `torch.mps.synchronize()` calls introduced in v2.5.21 for consistent behavior with CUDA pipeline
+
+**2025.12.12 - Version 2.5.21**
+
+- **🛠️ Fix: GGUF dequantization error on MPS** - Resolved shape mismatch error introduced in 2.5.20 by skipping GGUF quantized buffers in precision conversion - these must remain in packed format for on-the-fly dequantization during inference
+- **🍎 MPS: Eliminate CPU sync overhead** - Skip unnecessary CPU tensor offload on Apple Silicon unified memory architecture, preventing sync stalls that caused slowdowns. Input images and output video now stay on MPS device throughout the pipeline
+- **⚡ MPS: Preload text embeddings** - Load text embeddings before Phase 1 encoding to avoid sync stall at Phase 2 start, improving timing accuracy and throughput
+- **🧹 MPS: Optimized model cleanup** - Skip redundant CPU movement before model deletion on unified memory
+
+**2025.12.12 - Version 2.5.20**
+
+- **⚡ Expanded attention backends** - Full support for Flash Attention 2 (Ampere+), Flash Attention 3 (Hopper+), SageAttention 2, and SageAttention 3 (Blackwell/RTX 50xx), with automatic fallback chains to PyTorch SDPA when unavailable *(based on PR by [@naxci1](https://github.com/naxci1) - thank you!)*
+- **🍎 macOS/Apple Silicon compatibility** - Replaced MPS autocast with explicit dtype conversion throughout VAE and DiT pipelines, resolving hangs and crashes on M-series Macs. BlockSwap now auto-disables with warning (unified memory makes it meaningless)
+- **🛡️ Flash Attention graceful fallback** - Added compatibility shims for corrupted or partially installed flash_attn/xformers DLLs, preventing startup crashes
+- **🛡️ AMD ROCm: bitsandbytes conflict fix** - Prevent kernel registration errors when diffusers attempts to re-import broken bitsandbytes installations
+- **📦 ComfyUI Manager: macOS classifier fix** - Removed NVIDIA CUDA classifier causing false "GPU not supported" warnings on macOS
+- **📚 Documentation updates** - Updated README with attention backend details, BlockSwap macOS notes, and clarified model caching descriptions
+
+**2025.12.10 - Version 2.5.19**
+
+- **🎨 New header logo design** - Refreshed ASCII art banner *(thanks [@naxci1](https://github.com/naxci1))*
+- **🧹 Remove dead flash attention wrapper** - Removed legacy code from FP8CompatibleDiT; FlashAttentionVarlen already handles backend switching via its `attention_mode` attribute
+- **🛡️ Fix graceful fallback from flash-attn** - Add compatibility shims for corrupted flash_attn/xformers DLLs, preventing startup crashes when CUDA extensions are broken
+- **📊 Improved VRAM tracking** - Separate allocated vs reserved memory tracking, Windows-only overflow detection (WDDM paging behavior)
+- **♻️ Centralize backend detection** - Unified `is_mps_available()`, `is_cuda_available()`, `get_gpu_backend()` helpers across codebase
+- **🔄 Revert 2.5.14 VRAM limit enforcement** - Removed `set_per_process_memory_fraction` call; Overflow detection and warnings remain.
+
+**2025.12.09 - Version 2.5.18**
+
+- **🚀 CLI: Streaming mode for long videos** - New `--chunk_size` flag processes videos in memory-bounded chunks, enabling arbitrarily long videos without RAM limits. Works with model caching (`--cache_dit`/`--cache_vae`) for chunk-to-chunk reuse *(inspired by [disk02](https://github.com/disk02) PR contribution)*
+- **⚡ CLI: Multi-GPU streaming** - Each GPU now streams its segment internally with independent model caching, improving memory efficiency and enabling `--temporal_overlap` blending at GPU boundaries
+- **🔧 CLI: Fix large video MemoryError** - Shared memory transfer replaces numpy pickling, preventing crashes on high-resolution/long video outputs *(inspired by  [FurkanGozukara](https://github.com/FurkanGozukara) PR contribution)*
+
+**2025.12.05 - Version 2.5.17**
+
+- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Runtime bf16 CUBLAS probe replaces compute capability heuristics, correctly detecting unsupported GPUs without affecting RTX 20XX
+
+**2025.12.05 - Version 2.5.16**
+
+- **🔧 Fix: Older GPU compatibility (GTX 970, etc.)** - Automatic fallback for GPUs without bfloat16 support
+- **🐛 Fix: Quality regression** - Reverted bfloat16 detection that was causing artifact issues
+- **📋 Debug: Environment info display** - Shows system info in debug mode to help with issue reporting
+- **📚 Docs: Simplified contribution workflow** - Streamlined to main branch only
+
+**2025.12.03 - Version 2.5.15**
+
+- **🍎 Fix: MPS compatibility** - Disable antialias for MPS tensors and fix bfloat16 arange issues
+- **⚡ Fix: Autocast device type** - Use proper device type attribute to prevent autocast errors
+- **📊 Memory: Accurate VRAM tracking** - Use max_memory_reserved for more precise peak reporting
+- **🔧 Fix: Triton compatibility** - Add shim for bitsandbytes 0.45+ / triton 3.0+ (fixes PyTorch 2.7 installation errors)
+
+**2025.12.01 - Version 2.5.14**
+
+- **🍎 Fix: MPS device comparison** - Normalize device strings to prevent unnecessary tensor movements
+- **📊 Memory: VRAM swap detection** - Peak stats now show GPU+swap breakdown when overflow occurs, with warning when swap detected
+- **🛡️ Memory: Enforce physical VRAM limit** - PyTorch now OOMs instead of silently swapping to shared memory (prevents extreme slowdowns on Windows)
+
+**2025.11.30 - Version 2.5.13**
+
+- **🔧 Fix: PyTorch 2.7+ triton import error** - Resolved installation crash caused by triton.ops import chain on newer triton versions
+- **💾 Fix: OOM on float32 conversion for long videos** - Graceful fallback to native dtype when insufficient memory for float32 conversion
+- **🍎 Fix: CLI watermark error on macOS** - Resolved MPS-related watermark processing crash on Apple Silicon
+
+**2025.11.28 - Version 2.5.12**
+
+- **🐛 Fix: Color artifacts regression** - Reverted in-place tensor operations in video transform pipeline that caused color artifacts on some images
+
+**2025.11.28 - Version 2.5.11**
+
+- **⚡ Feature: CUDNN attention backend** - Added support for PyTorch 2.3+ CUDNN_ATTENTION backend with automatic fallback for older versions (thanks @eadwu)
+- **💾 Fix: Memory spike for long videos** - VAE decode now streams directly to pre-allocated tensor, eliminating OOM errors during long video processing
+- **🎨 Fix: LAB color correction artifacts** - Resolved tile boundary artifacts using wavelet reconstruction preprocessing
+- **🎨 Fix: Color reference misalignment** - Fixed color correction frame alignment with temporal overlap
+- **🍎 Fix: MPS detection reliability** - Switched to canonical `torch.backends.mps.is_available()` API for consistent Apple Silicon detection
+- **🖥️ Fix: Mac subprocess error** - CLI now uses direct processing on Mac to avoid MPS allocator failures in child processes
+- **🖥️ Fix: Multi-GPU device assignment** - CUDA_VISIBLE_DEVICES now set before spawn for proper worker inheritance
+- **📊 Fix: BlockSwap logging** - Now shows effective/total blocks (e.g., 32/32) instead of raw requested value
+- **🔧 Feature: Auto bfloat16 detection** - Automatically detects bfloat16 support to prevent CUBLAS errors on older GPUs
+- **📊 Feature: Peak RAM tracking** - Added RAM usage alongside VRAM in debug summary
+- **⚡ Performance: In-place tensor ops** - Reduced memory allocation overhead with in-place operations throughout pipeline
+- **📖 Docs: Multi-GPU clarification** - Clarified frame-level parallelism behavior expectations for multi-GPU setups
+
+**2025.11.13 - Version 2.5.10**
+
+- **🎯 Fix: Deterministic generation** - Identical images with the same seed now produce identical results across different sessions and batch positions
+- **🔧 Fix: Model caching with BlockSwap** - Resolved issue where cached DiT models wouldn't properly reload when VAE caching state changed
+- **💾 Fix: Runner caching optimization** - Runner templates now correctly cache whenever both DiT and VAE are cached, regardless of caching order
+- **📁 Fix: Case-insensitive model paths** - Extra model paths in YAML config now work regardless of case (seedvr2, SEEDVR2, SeedVR2, etc.)
+- **🐛 Fix: High resolution tile debug crash** - Fixed "NoneType has no attribute log" error when using maximum resolution with VAE tiling
+- **📊 Fix: Temporal overlap logging** - Corrected frame count reporting when temporal overlap is automatically adjusted
+- **🔍 Feature: Enhanced model path debugging** - Added detailed logging to help troubleshoot model loading issues (visible in debug mode)
+
+**2025.11.12 - Version 2.5.9**
+
+- **🐛 Fix: Tile debug visualization crash** - Fixed OpenCV error when using VAE tile debug mode on certain systems.
+- **🍎 Fix: macOS MPS loading error** - Added automatic CPU fallback for MPS allocator issues on certain PyTorch/macOS versions.
+- **🖥️ Fix: Windows log buffering** - Added flush to print statements for real-time log visibility in ComfyUI on Windows
+- **📦 Fix: ComfyUI Registry logo** - Updated icon URL to display properly in ComfyUI node registry
+- **ℹ️ Feature: Version display** - Added version number to node name and CLI/ComfyUI header for better tracking
+- **💝 Feature: GitHub Sponsors** - Added sponsor button to support project development. Thank you everyone for your support!
+- **📜 License: Apache 2.0** - Reverted License from MIT to Apache 2.0 to match ByteDance Seed project
+
+**2025.11.10 - Version 2.5.8**
+
+- **🐛 Fix (CLI): Windows batch processing duplicate files** - Fixed CLI batch mode processing each file twice on Windows due to case-insensitive filesystem. Improved directory scanning performance by 2-3x
+- **📁 Fix(CLI): Output folder location** - Output files now created in sensible locations: batch mode creates `{folder_name}_upscaled/` sibling folder with original filenames preserved; single file mode adds `_upscaled` suffix in same directory. All logs now show absolute paths for clarity
+- **🎨 Fix(CLI): RGBA alpha channel support** - PNG images with transparency are now properly detected and preserved through the upscaling pipeline, matching ComfyUI behavior
+
+**2025.11.10 - Version 2.5.7**
+
+- **🔧 Fix: Conv3d workaround compatibility** - Enhanced platform detection and added graceful fallback to prevent errors on PyTorch dev builds and AMD ROCm systems
+
+**2025.11.09 - Version 2.5.6**
+
+- 🎨 **Fix: Restored natural look for 7b model** - Corrected torch.compile optimization that was causing overly plastic/ high-specular appearance in upscaled videos with 7b model.
+
+- 💾 **Memory: Fixed RAM leak for long videos** - On-demand reconstruction with lightweight batch indices instead of storing full transformed videos, fixed release_tensor_memory to handle CPU/CUDA/MPS consistently, and refactored batch processing helpers
+
+**2025.11.08 - Version 2.5.4**
+
+- 🎨 **Fix: AdaIN color correction** - Replace `.view()` with `.reshape()` to handle non-contiguous tensors after spatial padding, resolving "view size is not compatible with input tensor's size and stride" error
+- 🔴 **Fix: AMD ROCm compatibility** - Add cuDNN availability check in Conv3d workaround to prevent "ATen not compiled with cuDNN support" error on ROCm systems (AMD GPUs on Windows/Linux)
+
+**2025.11.08 - Version 2.5.3**
+
+- 🍎 **Fix: Apple Silicon MPS device handling** - Corrected MPS device enumeration to use `"mps"` instead of `"mps:0"`, resolving invalid device errors on M-series Macs
+- 🪟 **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
 
 **2025.11.07 - Version 2.5.0** 🎉
 
@@ -122,7 +271,7 @@ We're actively working on improvements and new features. To stay informed:
 
 **2025.07.03**
 
-- 🛠️ Can run as **standalone mode** with **Multi GPU** see [🖥️ Run as Standalone](#️-run-as-standalone-cli)
+- 🛠️ Can run as **standalone mode** with **Multi GPU** see [🖥️ Run as Standalone](#run-as-standalone-cli)
 
 **2025.06.30**
 
@@ -169,8 +318,8 @@ We're actively working on improvements and new features. To stay informed:
 ### Performance Features
 - **torch.compile Integration**: Optional 20-40% DiT speedup and 15-25% VAE speedup with PyTorch 2.0+ compilation
 - **Multi-GPU CLI**: Distribute workload across multiple GPUs with automatic temporal overlap blending
-- **Model Caching**: Keep models loaded in memory for faster batch processing
+- **Model Caching**: Keep models loaded between generations for single-GPU directory processing or multi-GPU streaming
-- **Flexible Attention Backends**: Choose between PyTorch SDPA (stable, always available) or Flash Attention 2 (faster on supported hardware)
+- **Flexible Attention Backends**: Choose between PyTorch SDPA (stable, always available), Flash Attention 2/3, or SageAttention 2/3 for faster computation on supported hardware
 
 ### Quality Control
 - **Advanced Color Correction**: Five methods including LAB (recommended for highest fidelity), wavelet, wavelet adaptive, HSV, and AdaIN
@@ -199,7 +348,7 @@ With the current optimizations (tiling, BlockSwap, GGUF quantization), SeedVR2 c
 - **Python**: 3.12+ (Python 3.12 and 3.13 tested and recommended)
 - **PyTorch**: 2.0+ for torch.compile support (optional but recommended)
 - **Triton**: Required for torch.compile with inductor backend (optional)
-- **Flash Attention 2**: Provides faster attention computation on supported hardware (optional, falls back to PyTorch SDPA)
+- **Flash Attention / SageAttention**: Flash Attention 2 (Ampere+), Flash Attention 3 (Hopper+), SageAttention 2 or SageAttention 3 (Blackwell) provide faster attention computation on supported hardware (optional, falls back to PyTorch SDPA)
 
 ## 📦 Installation
 
@@ -314,14 +463,21 @@ Configure the DiT (Diffusion Transformer) model for video upscaling.
   - Requires offload_device to be set and different from device
 
 - **attention_mode**: Attention computation backend
-  - `sdpa`: PyTorch scaled_dot_product_attention (default, stable, always available)
+  - `sdpa`: PyTorch scaled_dot_product_attention (default, always available)
-  - `flash_attn`: Flash Attention 2 (faster on supported hardware, requires flash-attn package)
+  - `flash_attn_2`: Flash Attention 2 (Ampere+, requires flash-attn package)
+  - `flash_attn_3`: Flash Attention 3 (Hopper+, requires flash-attn with FA3 support)
+  - `sageattn_2`: SageAttention 2 (requires sageattention package)
+  - `sageattn_3`: SageAttention 3 (Blackwell/RTX 50xx, requires sageattn3 package)
 
 - **torch_compile_args**: Connect to SeedVR2 Torch Compile Settings node for 20-40% speedup
 
 **BlockSwap Explained:**
 
-BlockSwap enables running large models on GPUs with limited VRAM by dynamically swapping transformer blocks between GPU and CPU memory during inference. Here's how it works:
+BlockSwap enables running large models on GPUs with limited VRAM by dynamically swapping transformer blocks between GPU and CPU memory during inference.
+
+> **Note:** BlockSwap is not available on macOS. Apple Silicon Macs use unified memory architecture where GPU and CPU share the same memory pool, making BlockSwap meaningless. The option will be automatically disabled with a warning if requested on macOS.
+
+Here's how it works:
 
 - **What it does**: Keeps only the currently-needed transformer blocks on the GPU, while storing the rest on CPU or another device
 - **When to use it**: When you get OOM (Out of Memory) errors during the upscaling phase
@@ -676,9 +832,20 @@ The CLI provides comprehensive options for single-GPU, multi-GPU, and batch proc
 # Basic image upscaling
 python inference_cli.py image.jpg
 
-# Basic video video upscaling with temporal consistency
+# Basic video upscaling with temporal consistency
 python inference_cli.py video.mp4 --resolution 720 --batch_size 33
 
+# Streaming mode for long videos (memory-efficient) with 10-bit video output (requires FFMPEG)
+# Processes video in chunks of 330 frames to avoid loading entire video into RAM
+# Use --temporal_overlap to ensure smooth transitions between chunks
+python inference_cli.py long_video.mp4 \
+    --resolution 1080 \
+    --batch_size 33 \
+    --chunk_size 330 \
+    --temporal_overlap 3 \
+    --video_backend ffmpeg \
+    --10bit
+
 # Multi-GPU processing with temporal overlap
 python inference_cli.py video.mp4 \
     --cuda_device 0,1 \
@@ -724,6 +891,8 @@ python inference_cli.py media_folder/ \
 - `<input>`: Input file (.mp4, .avi, .png, .jpg, etc.) or directory
 - `--output`: Output path (default: auto-generated in 'output/' directory)
 - `--output_format`: Output format: 'mp4' (video) or 'png' (image sequence). Default: auto-detect from input type
+- `--video_backend`: Video encoder backend: 'opencv' (default) or 'ffmpeg' (requires ffmpeg in PATH)
+- `--10bit`: Save 10-bit video with x265 codec and yuv420p10le pixel format (reduces banding in gradients). Without this flag, ffmpeg uses x264 (yuv420p) for maximum compatibility. Requires --video_backend ffmpeg
 - `--model_dir`: Model directory (default: ./models/SEEDVR2)
 
 **Model Selection:**
@@ -735,7 +904,8 @@ python inference_cli.py media_folder/ \
 - `--batch_size`: Frames per batch (must follow 4n+1: 1, 5, 9, 13, 17, 21...). Ideally matches shot length for best temporal consistency (default: 5)
 - `--seed`: Random seed for reproducibility (default: 42)
 - `--skip_first_frames`: Skip N initial frames (default: 0)
-- `--load_cap`: Load maximum N frames from video. 0 = load all (default: 0)
+- `--load_cap`: Maximum total frames to load from video. 0 = load all (default: 0)
+- `--chunk_size`: Frames per chunk for streaming mode. When > 0, processes video in memory-bounded chunks of N frames, writing each chunk before loading the next. Essential for long videos that would otherwise exceed RAM. Use with `--temporal_overlap` for seamless chunk transitions. 0 = load all frames at once (default: 0)
 - `--prepend_frames`: Prepend N reversed frames to reduce start artifacts (auto-removed) (default: 0)
 - `--temporal_overlap`: Frames to overlap between batches/GPUs for smooth blending (default: 0)
 
@@ -747,9 +917,8 @@ python inference_cli.py media_folder/ \
 **Memory Management:**
 - `--dit_offload_device`: Device to offload DiT model: 'none' (keep on GPU), 'cpu', or 'cuda:X' (default: none)
 - `--vae_offload_device`: Device to offload VAE model: 'none', 'cpu', or 'cuda:X' (default: none)
-- `--blocks_to_swap`: Number of transformer blocks to swap (0=disabled, 3B: 0-32, 7B: 0-36). Requires dit_offload_device (default: 0)
+- `--blocks_to_swap`: Number of transformer blocks to swap (0=disabled, 3B: 0-32, 7B: 0-36). Requires dit_offload_device (default: 0). Not available on macOS.
-- `--swap_io_components`: Offload I/O components for additional VRAM savings. Requires dit_offload_device
+- `--swap_io_components`: Offload I/O components for additional VRAM savings. Requires dit_offload_device. Not available on macOS.
-- `--use_non_blocking`: Use non-blocking memory transfers for BlockSwap (recommended)
 
 **VAE Tiling:**
 - `--vae_encode_tiled`: Enable VAE encode tiling to reduce VRAM during encoding
@@ -761,7 +930,8 @@ python inference_cli.py media_folder/ \
 - `--tile_debug`: Visualize tiles: 'false' (default), 'encode', or 'decode'
 
 **Performance Optimization:**
-- `--attention_mode`: Attention backend: 'sdpa' (default, stable) or 'flash_attn' (faster, requires package)
+- `--allow_vram_overflow`: Allow VRAM overflow to system RAM. Prevents OOM but may cause severe slowdown
+- `--attention_mode`: Attention backend: 'sdpa' (default), 'flash_attn_2' (Ampere+), 'flash_attn_3' (Hopper+), 'sageattn_2', or 'sageattn_3' (Blackwell)
 - `--compile_dit`: Enable torch.compile for DiT model (20-40% speedup, requires PyTorch 2.0+ and Triton)
 - `--compile_vae`: Enable torch.compile for VAE model (15-25% speedup, requires PyTorch 2.0+ and Triton)
 - `--compile_backend`: Compilation backend: 'inductor' (full optimization) or 'cudagraphs' (lightweight) (default: inductor)
@@ -772,8 +942,8 @@ python inference_cli.py media_folder/ \
 - `--compile_dynamo_recompile_limit`: Max recompilation attempts before fallback (default: 128)
 
 **Model Caching (batch processing):**
-- `--cache_dit`: Cache DiT model between files (single GPU only, speeds up directory processing)
+- `--cache_dit`: Keep DiT model in memory between generations. Works with single-GPU directory processing or multi-GPU streaming (`--chunk_size`). Requires `--dit_offload_device`
-- `--cache_vae`: Cache VAE model between files (single GPU only, speeds up directory processing)
+- `--cache_vae`: Keep VAE model in memory between generations. Works with single-GPU directory processing or multi-GPU streaming (`--chunk_size`). Requires `--vae_offload_device`
 
 **Multi-GPU:**
 - `--cuda_device`: CUDA device id(s). Single id (e.g., '0') or comma-separated list '0,1' for multi-GPU
@@ -783,26 +953,40 @@ python inference_cli.py media_folder/ \
 
 ### Multi-GPU Processing Explained
 
-The CLI's multi-GPU mode automatically distributes the workload across multiple GPUs with intelligent temporal overlap handling:
+The CLI's multi-GPU mode uses **frame-level parallelism**: the video is split into chunks and each GPU processes its chunk independently through all 4 phases (encode → upscale → decode → postprocess). This is ideal for long videos where you want to reduce total processing time by dividing the workload.
 
 **How it works:**
-1. Video is split into chunks, one per GPU
+1. Video frames are split evenly across GPUs (e.g., 100 frames on 2 GPUs → 50 frames each)
-2. Each GPU processes its chunk independently
+2. Each GPU loads its own copy of the models and processes its chunk independently
-3. Chunks overlap by `--temporal_overlap` frames
+3. When `--temporal_overlap` is set, chunks include overlapping frames for seamless blending
-4. Results are blended together seamlessly using the overlap region
+4. Results are concatenated (and blended at overlap regions) into the final video
 
-**Example for 2 GPUs with temporal_overlap=4:**
+**Example for 100 frames on 2 GPUs with temporal_overlap=4:**
 ```
-GPU 0: Frames 0-50 (includes 4 overlap frames at end)
+GPU 0: Frames 0-53 (50 base + 4 overlap at end, processed as independent video)
-GPU 1: Frames 46-100 (includes 4 overlap frames at beginning)
+GPU 1: Frames 50-99 (50 frames, 4 overlap at start, processed as independent video)
-Result: Frames 0-100 with smooth transition at frame 48
+Result: Frames 0-99 with smooth blending at the transition point
 ```
 
+**Important considerations:**
+- Each GPU processes its chunk as a separate video with its own batch splitting
+- `batch_size` controls batching *within* each GPU's chunk, not across GPUs
+- For short videos (< 100 frames), single GPU is often more efficient due to model loading overhead
+- Multi-GPU doubles VRAM usage (each GPU loads full models) but roughly halves processing time
+
+**When to use multi-GPU:**
+- Long videos (100+ frames) where splitting provides significant time savings
+- When you have multiple GPUs with sufficient VRAM each
+
+**When to use single GPU:**
+- Short videos where model loading overhead outweighs parallel gains
+- When you want all frames processed together for maximum temporal coherence
+
 **Best practices:**
-- Set `--temporal_overlap` to 2-8 frames for smooth blending
+- Set `--temporal_overlap` to 2-4 frames for smooth blending between GPU chunks
 - Higher overlap = smoother transitions but more redundant processing
 - Use `--prepend_frames` to reduce artifacts at video start
-- batch_size should divide evenly into chunk sizes for best results
+- For optimal quality on short videos, use single GPU with `batch_size` matching your shot length
 
 ## ⚠️ Limitations
 
@@ -850,7 +1034,7 @@ For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
 2. Create your feature branch (`git checkout -b feature/AmazingFeature`)
 3. Commit your changes (`git commit -m 'Add some AmazingFeature'`)
 4. Push to the branch (`git push origin feature/AmazingFeature`)
-5. Open a Pull Request to **main** branch for stable features or **nightly** branch for experimental features
+5. Open a Pull Request to the **main** branch
 
 **Get Help:**
 - YouTube: [AInVFX Channel](https://www.youtube.com/@AInVFX)
@@ -862,8 +1046,8 @@ For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md).
 
 This ComfyUI implementation is a collaborative project by **[NumZ](https://github.com/numz)** and **[AInVFX](https://www.youtube.com/@AInVFX)** (Adrien Toupet), based on the original [SeedVR2](https://github.com/ByteDance-Seed/SeedVR) by ByteDance Seed Team.
 
-Special thanks to our community contributors including [benjaminherb](https://github.com/benjaminherb), [cmeka](https://github.com/cmeka), [FurkanGozukara](https://github.com/FurkanGozukara), [JohnAlcatraz](https://github.com/JohnAlcatraz), [lihaoyun6](https://github.com/lihaoyun6), [Luchuanzhao](https://github.com/Luchuanzhao), [Luke2642](https://github.com/Luke2642), [naxci1](https://github.com/naxci1), [q5sys](https://github.com/q5sys), and many others for their improvements, bug fixes, and testing.
+Special thanks to our community contributors including [naxci1](https://github.com/naxci1), [thehhmdb](https://github.com/thehhmdb), [s-cerevisiae](https://github.com/s-cerevisiae), [benjaminherb](https://github.com/benjaminherb), [cmeka](https://github.com/cmeka), [FurkanGozukara](https://github.com/FurkanGozukara), [JohnAlcatraz](https://github.com/JohnAlcatraz), [lihaoyun6](https://github.com/lihaoyun6), [Luchuanzhao](https://github.com/Luchuanzhao), [Luke2642](https://github.com/Luke2642), [proxyid](https://github.com/proxyid), [q5sys](https://github.com/q5sys), and many others for their improvements, bug fixes, and testing.
 
 ## 📜 License
 
-The code in this repository is released under the MIT license as found in the [LICENSE](LICENSE) file.
+The code in this repository is released under the Apache 2.0 license as found in the [LICENSE](LICENSE) file.

__init__.py

@@ -3,6 +3,7 @@ ComfyUI-SeedVR2_VideoUpscaler
 Official SeedVR2 integration for ComfyUI
 """
 
+from .src.optimization.compatibility import ensure_triton_compat  # noqa: F401
 from .src.interfaces import comfy_entrypoint, SeedVR2Extension
 
 __all__ = ["comfy_entrypoint", "SeedVR2Extension"]

pyproject.toml

@@ -1,15 +1,14 @@
 [project]
 name = "seedvr2_videoupscaler"
 description = "SeedVR2 official ComfyUI integration: ByteDance-Seed's one-step diffusion-based video/image upscaling with memory-efficient inference"
-version = "2.5.0"
+version = "2.5.24"
 authors = [
     {name = "numz"},
     {name = "adrientoupet"}
 ]
 license = {file = "LICENSE"} 
 classifiers = [
-     "Operating System :: OS Independent",
+     "Operating System :: OS Independent"
-     "Environment :: GPU :: NVIDIA CUDA"
 ]
 dependencies = [
     "torch",
@@ -37,5 +36,5 @@ Forum = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions"
 [tool.comfy]
 PublisherId = "ainvfx"
 DisplayName = "ComfyUI-SeedVR2_VideoUpscaler"
-Icon = ""
+Icon = "https://raw.githubusercontent.com/numz/ComfyUI-SeedVR2_VideoUpscaler/refs/heads/main/docs/seedvr_logo.png"
 includes = []

src/common/diffusion/timesteps/sampling/trailing.py

@@ -36,7 +36,7 @@ class UniformTrailingSamplingTimesteps(SamplingTimesteps):
         dtype: torch.dtype = torch.float32,
     ):
         # Create trailing timesteps with specified dtype
-        timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device=device, dtype=dtype)
+        timesteps = torch.arange(1.0, 0.0, -1.0 / steps, device='cpu').to(device=device, dtype=dtype)
 
         # Shift timesteps.
         timesteps = shift * timesteps / (1 + (shift - 1) * timesteps)

src/common/distributed/basic.py

@@ -21,6 +21,7 @@ from datetime import timedelta
 import torch
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel
+from ...optimization.memory_manager import is_mps_available
 
 def get_global_rank() -> int:
     """
@@ -47,7 +48,7 @@ def get_device() -> torch.device:
     """
     Get current rank device.
     """
-    if torch.mps.is_available():
+    if is_mps_available():
         return torch.device("mps")
     return torch.device("cuda", get_local_rank())
 

src/core/alpha_upscaling.py

@@ -337,13 +337,23 @@ def edge_guided_alpha_upscale(
     rgb_edges = detect_edges_batch(images=rgb_normalized, method='sobel', debug=debug)
     
     # Step 1: Initial bicubic upscale provides smooth base before edge refinement
-    alpha_upscaled = F.interpolate(
+    # MPS on PyTorch < 2.8 doesn't support bicubic+antialias - use CPU fallback
-        input_alpha,
+    try:
-        size=(H_out, W_out),
+        alpha_upscaled = F.interpolate(
-        mode='bicubic',
+            input_alpha,
-        align_corners=False,
+            size=(H_out, W_out),
-        antialias=True
+            mode='bicubic',
-    ).clamp(0, 1)
+            align_corners=False,
+            antialias=True
+        ).clamp(0, 1)
+    except NotImplementedError:
+        alpha_upscaled = F.interpolate(
+            input_alpha.cpu(),
+            size=(H_out, W_out),
+            mode='bicubic',
+            align_corners=False,
+            antialias=True
+        ).to(device).clamp(0, 1)
     
     if is_binary_mask:
         if debug:

src/core/generation_utils.py

@@ -36,6 +36,7 @@ from .model_configuration import configure_runner
 from .infer import VideoDiffusionInfer
 from ..data.image.transforms.divisible_crop import DivisiblePad
 from ..data.image.transforms.na_resize import NaResize
+from ..optimization.compatibility import COMPUTE_DTYPE, BFLOAT16_SUPPORTED
 from ..optimization.memory_manager import manage_tensor
 from ..utils.constants import get_script_directory
 
@@ -349,7 +350,12 @@ def setup_generation_context(
     vae_device = _normalize_device(vae_device)
     dit_offload_device = _normalize_device(dit_offload_device) if dit_offload_device is not None else None
     vae_offload_device = _normalize_device(vae_offload_device) if vae_offload_device is not None else None
-    tensor_offload_device = _normalize_device(tensor_offload_device) if tensor_offload_device is not None else None
+    # MPS unified memory: CPU offload causes sync overhead with no memory benefit
+    is_mps = dit_device.type == 'mps' or vae_device.type == 'mps'
+    if is_mps and tensor_offload_device is not None and str(tensor_offload_device) == 'cpu':
+        tensor_offload_device = None
+    else:
+        tensor_offload_device = _normalize_device(tensor_offload_device) if tensor_offload_device is not None else None
     
     # Set LOCAL_RANK to 0 for single-GPU inference mode
     # CLI multi-GPU uses CUDA_VISIBLE_DEVICES to restrict visibility per worker
@@ -371,17 +377,15 @@ def setup_generation_context(
         'dit_offload_device': dit_offload_device,
         'vae_offload_device': vae_offload_device,
         'tensor_offload_device': tensor_offload_device,
-        'compute_dtype': torch.bfloat16, # Hardcoded - gives the best compromise between memory & quality without artifacts
+        'compute_dtype': COMPUTE_DTYPE,
         'interrupt_fn': interrupt_fn,
         'video_transform': None,
         'text_embeds': None,
-        'all_transformed_videos': [],
         'all_latents': [],
         'all_upscaled_latents': [],
         'batch_samples': [],
         'final_video': None,
         'comfyui_available': comfyui_available,
-        'interrupt_fn': interrupt_fn,
     }
     
     if debug:
@@ -403,7 +407,12 @@ def setup_generation_context(
             f"LOCAL_RANK={os.environ['LOCAL_RANK']}",
             category="setup"
         )
-        reason = "quality" if ctx['compute_dtype'] == torch.float32 else "compatibility"
+        if ctx['compute_dtype'] == torch.float32:
+            reason = "quality"
+        elif not BFLOAT16_SUPPORTED:
+            reason = "compatibility (GPU lacks bfloat16 CUBLAS - 7B models unsupported, 3B may have artifacts)"
+        else:
+            reason = "performance"
         debug.log(f"Unified compute dtype: {ctx['compute_dtype']} across entire pipeline for maximum {reason}", category="precision")
     
     return ctx
@@ -453,7 +462,7 @@ def prepare_runner(
         decode_tile_size: Tile size for decoding (height, width)
         decode_tile_overlap: Tile overlap for decoding (height, width)
         tile_debug: Tile visualization mode (false/encode/decode)
-        attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
+        attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
         torch_compile_args_dit: Optional torch.compile configuration for DiT model
         torch_compile_args_vae: Optional torch.compile configuration for VAE model
         
@@ -525,8 +534,8 @@ def load_text_embeddings(script_directory: str, device: torch.device,
         - Memory-efficient embedding preparation
         - Consistent movement logging
     """
-    text_pos_embeds = torch.load(os.path.join(script_directory, 'pos_emb.pt'))
+    text_pos_embeds = torch.load(os.path.join(script_directory, 'pos_emb.pt'), weights_only=True)
-    text_neg_embeds = torch.load(os.path.join(script_directory, 'neg_emb.pt'))
+    text_neg_embeds = torch.load(os.path.join(script_directory, 'neg_emb.pt'), weights_only=True)
     
     text_pos_embeds = manage_tensor(
         tensor=text_pos_embeds,
@@ -724,7 +733,7 @@ def _draw_tile_boundaries(image: torch.Tensor, debug: 'Debug', tile_boundaries:
     annotated_frames = []
     for frame_idx in range(T):
         # Convert frame to numpy (handle RGB and RGBA)
-        img = (image[frame_idx].float().cpu().numpy() * 255).astype(np.uint8)  # [H, W, C]
+        img = np.ascontiguousarray((image[frame_idx].float().cpu().numpy() * 255).astype(np.uint8))  # [H, W, C]
         
         # Draw boundary lines inside each tile
         for idx, tile_info in enumerate(tile_boundaries):
@@ -815,4 +824,4 @@ def ensure_precision_initialized(
             debug.log(f"Model precision: {', '.join(parts)}", category="precision")
             
     except Exception as e:
-        debug.log(f"Could not log model dtypes: {e}", level="WARNING", category="precision", force=True)
+        debug.log(f"Could not log model dtypes: {e}", level="WARNING", category="precision", force=True)

src/core/infer.py

@@ -26,41 +26,13 @@ from ..common.diffusion import (
 from ..common.distributed import (
     get_device,
 )
+from ..optimization.performance import (
+    optimized_channels_to_last,
+    optimized_channels_to_second
+)
 from ..models.dit_3b import na
 
 
-def optimized_channels_to_last(tensor):
-    """🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
-    Moves channels from position 1 to last position using PyTorch native operations.
-    """
-    if tensor.ndim == 3:  # [batch, channels, spatial]
-        return tensor.permute(0, 2, 1)
-    elif tensor.ndim == 4:  # [batch, channels, height, width]
-        return tensor.permute(0, 2, 3, 1)
-    elif tensor.ndim == 5:  # [batch, channels, depth, height, width]
-        return tensor.permute(0, 2, 3, 4, 1)
-    else:
-        # Fallback for other dimensions - move channel (dim=1) to last
-        dims = list(range(tensor.ndim))
-        dims = [dims[0]] + dims[2:] + [dims[1]]  # [0, 2, 3, ..., 1]
-        return tensor.permute(*dims)
-
-def optimized_channels_to_second(tensor):
-    """🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
-    Moves channels from last position to position 1 using PyTorch native operations.
-    """
-    if tensor.ndim == 3:  # [batch, spatial, channels]
-        return tensor.permute(0, 2, 1)
-    elif tensor.ndim == 4:  # [batch, height, width, channels]
-        return tensor.permute(0, 3, 1, 2)
-    elif tensor.ndim == 5:  # [batch, depth, height, width, channels]
-        return tensor.permute(0, 4, 1, 2, 3)
-    else:
-        # Fallback for other dimensions - move last dim to position 1
-        dims = list(range(tensor.ndim))
-        dims = [dims[0], dims[-1]] + dims[1:-1]  # [0, -1, 1, 2, ..., -2]
-        return tensor.permute(*dims)
-
 class VideoDiffusionInfer():
     def __init__(self, config: DictConfig, debug: 'Debug',
                  encode_tiled: bool = False, encode_tile_size: Tuple[int, int] = (512, 512), 
@@ -182,15 +154,26 @@ class VideoDiffusionInfer():
                     vae_dtype = dtype  # Fallback
 
                 # Use autocast if VAE dtype differs from input dtype
+                # Skip autocast on MPS (only supports bf16, unified memory = no benefit)
+                # Instead, explicitly convert input to model dtype
                 if vae_dtype != sample.dtype:
-                    with torch.autocast(str(device), sample.dtype, enabled=True):
+                    if device.type == 'mps':
+                        # MPS: explicit dtype conversion instead of autocast
+                        sample = sample.to(vae_dtype)
                         if use_sample:
                             latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size, 
                                                     tile_overlap=self.encode_tile_overlap).latent
                         else:
-                            # Deterministic vae encode, only used for i2v inference (optionally)
                             latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
                                                 tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
+                    else:
+                        with torch.autocast(device.type, sample.dtype, enabled=True):
+                            if use_sample:
+                                latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size, 
+                                                        tile_overlap=self.encode_tile_overlap).latent
+                            else:
+                                latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size,
+                                                    tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
                 else:
                     if use_sample:
                         latent = self.vae.encode(sample, tiled=self.encode_tiled, tile_size=self.encode_tile_size, 
@@ -201,7 +184,7 @@ class VideoDiffusionInfer():
                                             tile_overlap=self.encode_tile_overlap).posterior.mode().squeeze(2)
 
                 latent = latent.unsqueeze(2) if latent.ndim == 4 else latent
-                latent = rearrange(latent, "b c ... -> b ... c")
+                latent = optimized_channels_to_last(latent)
                 latent = (latent - shift) * scale
                 latents.append(latent)
 
@@ -248,7 +231,7 @@ class VideoDiffusionInfer():
 
             for i, latent in enumerate(latents):
                 latent = latent / scale + shift
-                latent = rearrange(latent, "b ... c -> b c ...")
+                latent = optimized_channels_to_second(latent)
                 latent = latent.squeeze(2)
 
                 # Detect VAE model dtype
@@ -258,13 +241,23 @@ class VideoDiffusionInfer():
                     vae_dtype = dtype  # Fallback
 
                 # Use autocast if VAE dtype differs from latent dtype
+                # Skip autocast on MPS (only supports bf16, unified memory = no benefit)
                 if vae_dtype != latent.dtype:
-                    with torch.autocast(str(device), latent.dtype, enabled=True):
+                    if device.type == 'mps':
+                        # MPS: explicit dtype conversion instead of autocast
+                        latent = latent.to(vae_dtype)
                         sample = self.vae.decode(
                             latent,
                             tiled=self.decode_tiled, tile_size=self.decode_tile_size,
                             tile_overlap=self.decode_tile_overlap
                         ).sample
+                    else:
+                        with torch.autocast(device.type, latent.dtype, enabled=True):
+                            sample = self.vae.decode(
+                                latent,
+                                tiled=self.decode_tiled, tile_size=self.decode_tile_size,
+                                tile_overlap=self.decode_tile_overlap
+                            ).sample
                 else:
                     sample = self.vae.decode(
                         latent,

src/core/model_configuration.py

@@ -70,11 +70,11 @@ from .model_cache import get_global_cache
 from ..common.config import load_config
 from ..models.video_vae_v3.modules.causal_inflation_lib import InflatedCausalConv3d
 from ..optimization.compatibility import (
-    FP8CompatibleDiT,
+    CompatibleDiT,
     TRITON_AVAILABLE,
-    validate_flash_attention_availability
+    validate_attention_mode
 )
-from ..optimization.blockswap import is_blockswap_enabled, apply_block_swap_to_dit, cleanup_blockswap
+from ..optimization.blockswap import is_blockswap_enabled, validate_blockswap_config, apply_block_swap_to_dit, cleanup_blockswap
 from ..optimization.memory_manager import cleanup_dit, cleanup_vae
 from ..utils.constants import find_model_file
 
@@ -171,7 +171,7 @@ def _describe_attention_mode(attention_mode: Optional[str]) -> str:
     Generate human-readable description of attention mode configuration.
     
     Args:
-        attention_mode: Attention mode string ('sdpa' or 'flash_attn')
+        attention_mode: Attention mode string ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
         
     Returns:
         Human-readable description string
@@ -181,7 +181,10 @@ def _describe_attention_mode(attention_mode: Optional[str]) -> str:
     
     mode_descriptions = {
         'sdpa': 'PyTorch SDPA',
-        'flash_attn': 'Flash Attention 2'
+        'flash_attn_2': 'Flash Attention 2',
+        'flash_attn_3': 'Flash Attention 3',
+        'sageattn_2': 'SageAttention 2',
+        'sageattn_3': 'SageAttention 3 (Blackwell)'
     }
     
     return mode_descriptions.get(attention_mode, attention_mode)
@@ -436,7 +439,7 @@ def _update_dit_config(
             - dynamic: bool - Enable dynamic shapes
             - dynamo_cache_size_limit: int - Cache size limit
             - dynamo_recompile_limit: int - Recompilation limit
-        attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
+        attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
         debug: Debug instance for logging
         
     Returns:
@@ -771,7 +774,7 @@ def configure_runner(
         decode_tile_size: Tile size for decoding (height, width)
         decode_tile_overlap: Tile overlap for decoding (height, width)
         tile_debug: Tile visualization mode (false/encode/decode)
-        attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
+        attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
         torch_compile_args_dit: Optional torch.compile configuration for DiT model
         torch_compile_args_vae: Optional torch.compile configuration for VAE model
         
@@ -792,6 +795,14 @@ def configure_runner(
     if debug is None:
         raise ValueError("Debug instance must be provided to configure_runner")
     
+    # Validate BlockSwap configuration early (before any model loading)
+    block_swap_config = validate_blockswap_config(
+        block_swap_config=block_swap_config,
+        dit_device=ctx['dit_device'],
+        dit_offload_device=ctx.get('dit_offload_device'),
+        debug=debug
+    )
+    
     # Phase 1: Initialize cache and get cached models
     cache_context = _initialize_cache_context(
         dit_cache, vae_cache, dit_id, vae_id, 
@@ -857,7 +868,7 @@ def _configure_runner_settings(
         decode_tile_size: Tile dimensions (height, width) for decoding in pixels
         decode_tile_overlap: Overlap dimensions (height, width) between decoding tiles
         tile_debug: Tile visualization mode (false/encode/decode)
-        attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
+        attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
         torch_compile_args_dit: torch.compile configuration for DiT model or None
         torch_compile_args_vae: torch.compile configuration for VAE model or None
         block_swap_config: BlockSwap configuration for DiT model or None
@@ -1165,23 +1176,23 @@ def apply_model_specific_config(model: torch.nn.Module, runner: VideoDiffusionIn
     """
     if is_dit:
         # DiT-specific
-        # Apply FP8 compatibility wrapper with compute_dtype
+        # Apply compatibility wrapper with compute_dtype
-        if not isinstance(model, FP8CompatibleDiT):
+        if not isinstance(model, CompatibleDiT):
-            debug.log("Applying FP8/RoPE compatibility wrapper to DiT model", category="setup")
+            debug.log("Applying DiT compatibility wrapper", category="setup")
-            debug.start_timer("FP8CompatibleDiT")
+            debug.start_timer("CompatibleDiT")
             # Get compute_dtype from runner if available, fallback to bfloat16
             compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16)
-            model = FP8CompatibleDiT(model, debug, compute_dtype=compute_dtype, skip_conversion=False)
+            model = CompatibleDiT(model, debug, compute_dtype=compute_dtype, skip_conversion=False)
-            debug.end_timer("FP8CompatibleDiT", "FP8/RoPE compatibility wrapper application")
+            debug.end_timer("CompatibleDiT", "Compatibility wrapper application")
         else:
-            debug.log("Reusing existing FP8/RoPE compatibility wrapper", category="reuse")
+            debug.log("Reusing existing DiT compatibility wrapper", category="reuse")
         
         # Apply attention mode and compute_dtype to all FlashAttentionVarlen modules
         if hasattr(runner, '_dit_attention_mode'):
             requested_attention_mode = runner._dit_attention_mode or 'sdpa'
             
             # Validate and get final attention_mode (with warning if fallback needed)
-            attention_mode = validate_flash_attention_availability(requested_attention_mode, debug)
+            attention_mode = validate_attention_mode(requested_attention_mode, debug)
             
             # Get compute_dtype from runner
             compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16)            
@@ -1466,4 +1477,4 @@ def _propagate_debug_to_modules(module: torch.nn.Module, debug: 'Debug') -> None
     for name, submodule in module.named_modules():
         if submodule.__class__.__name__ in target_modules:
             if not hasattr(submodule, 'debug'):  # Only set if not already present
-                submodule.debug = debug
+                submodule.debug = debug

src/core/model_loader.py

@@ -115,7 +115,28 @@ def load_quantized_state_dict(checkpoint_path: str, device: torch.device = torch
                 debug.log("This is a one-time installation that will enable loading of .safetensors files", 
                          level="INFO", category="info", force=True)
             raise ImportError(error_msg)
-        state = load_safetensors_file(checkpoint_path, device=device_str)
+        
+        # Try direct device loading first (optimal path)
+        try:
+            state = load_safetensors_file(checkpoint_path, device=device_str)
+        except RuntimeError as e:
+            # MPS allocator fallback: some PyTorch/macOS versions have issues with
+            # direct MPS loading (allocation failures, watermark errors, etc.)
+            error_msg = str(e).lower()
+            is_mps_alloc_error = device.type == "mps" and any(
+                keyword in error_msg for keyword in ["watermark", "allocat", "memory"]
+            )
+            
+            if is_mps_alloc_error:
+                # Transparent fallback - only log if debug enabled
+                if debug:
+                    debug.log("Using CPU intermediate loading for MPS compatibility", 
+                            category="info", indent_level=1)
+                state = load_safetensors_file(checkpoint_path, device="cpu")
+                # Tensors will be moved to MPS during model.load_state_dict()
+            else:
+                # Re-raise if it's a different error (file corruption, etc.)
+                raise
     elif checkpoint_path.endswith('.gguf'):
         validate_gguf_availability(f"load {os.path.basename(checkpoint_path)}", debug)
         state = _load_gguf_state(
@@ -125,7 +146,7 @@ def load_quantized_state_dict(checkpoint_path: str, device: torch.device = torch
                     handle_prefix="model.diffusion_model."
                 )
     elif checkpoint_path.endswith('.pth'):
-        state = torch.load(checkpoint_path, map_location=device_str, mmap=True)
+        state = torch.load(checkpoint_path, map_location=device_str, mmap=True, weights_only=True)
     else:
         raise ValueError(f"Unsupported checkpoint format. Expected .safetensors or .pth, got: {checkpoint_path}")
     
@@ -372,6 +393,20 @@ class GGUFTensor(torch.Tensor):
                 if debug:
                     debug.log(f"Error in {func.__name__} dequantization: {e}", level="WARNING", category="dit", force=True)
                 raise
+
+        # Handle conv2d/conv3d operations (critical for GGUF VAE models)
+        # Conv3d layers (InflatedCausalConv3d) are not replaced by layer replacement
+        if func in {torch.nn.functional.conv2d, torch.nn.functional.conv3d}:
+            if len(args) >= 2 and isinstance(args[1], cls):  # weight is second arg
+                try:
+                    weight_tensor = args[1]
+                    dequantized_weight = weight_tensor.dequantize(device=args[0].device, dtype=args[0].dtype)
+                    new_args = (args[0], dequantized_weight) + args[2:]
+                    return func(*new_args, **kwargs)
+                except Exception as e:
+                    if debug:
+                        debug.log(f"Error in conv dequantization: {e}", level="WARNING", category="dit", force=True)
+                    raise
         
         # For ALL other operations, delegate to parent WITHOUT dequantization
         # This includes .cpu(), .to(), .device, .dtype, .shape, etc.

src/data/image/transforms/area_resize.py

@@ -19,6 +19,7 @@ import torch
 from PIL import Image
 from torchvision.transforms import functional as TVF
 from torchvision.transforms.functional import InterpolationMode
+from ....optimization.memory_manager import is_mps_available
 
 
 class AreaResize:
@@ -31,7 +32,7 @@ class AreaResize:
         self.max_area = max_area
         self.downsample_only = downsample_only
         self.interpolation = interpolation
-        if torch.mps.is_available():
+        if is_mps_available():
             self.interpolation = InterpolationMode.BILINEAR
 
     def __call__(self, image: Union[torch.Tensor, Image.Image]):
@@ -50,10 +51,12 @@ class AreaResize:
 
         resized_height, resized_width = round(height * scale), round(width * scale)
 
+        antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
         return TVF.resize(
             image,
             size=(resized_height, resized_width),
             interpolation=self.interpolation,
+            antialias=antialias,
         )
 
 

src/data/image/transforms/na_resize.py

@@ -18,6 +18,7 @@ from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resiz
 
 from .area_resize import AreaResize
 from .side_resize import SideResize
+from ....optimization.memory_manager import is_mps_available
 
 def NaResize(
     resolution: int,
@@ -26,7 +27,7 @@ def NaResize(
     max_resolution: int = 0,
     interpolation: InterpolationMode = InterpolationMode.BICUBIC,
 ):
-    Interpolation = InterpolationMode.BILINEAR if torch.mps.is_available() else interpolation
+    Interpolation = InterpolationMode.BILINEAR if is_mps_available() else interpolation
     if mode == "area":
         return AreaResize(
             max_area=resolution**2,

src/data/image/transforms/side_resize.py

@@ -17,6 +17,7 @@ import torch
 from PIL import Image
 from torchvision.transforms import InterpolationMode
 from torchvision.transforms import functional as TVF
+from ....optimization.memory_manager import is_mps_available
 
 class SideResize:
     def __init__(
@@ -30,7 +31,7 @@ class SideResize:
         self.max_size = max_size
         self.downsample_only = downsample_only
         self.interpolation = interpolation
-        if torch.mps.is_available():
+        if is_mps_available():
             self.interpolation = InterpolationMode.BILINEAR
 
     def __call__(self, image: Union[torch.Tensor, Image.Image]):
@@ -56,8 +57,9 @@ class SideResize:
         else:
             size = self.size
 
-        # Resize to shortest edge
+        # Resize to shortest edge (disable antialias only for MPS tensors - not supported)
-        resized = TVF.resize(image, size, self.interpolation)
+        antialias = not (isinstance(image, torch.Tensor) and image.device.type == 'mps')
+        resized = TVF.resize(image, size, self.interpolation, antialias=antialias)
         
         # Apply max_size constraint if specified
         if self.max_size > 0:
@@ -69,6 +71,6 @@ class SideResize:
             if max(h, w) > self.max_size:
                 scale = self.max_size / max(h, w)
                 new_h, new_w = round(h * scale), round(w * scale)
-                resized = TVF.resize(resized, (new_h, new_w), self.interpolation)
+                resized = TVF.resize(resized, (new_h, new_w), self.interpolation, antialias=antialias)
         
         return resized

src/interfaces/dit_model_loader.py

@@ -66,7 +66,8 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
                         "• 3B model: 0-32 blocks\n"
                         "• 7B model: 0-36 blocks\n"
                         "\n"
-                        "Requires offload_device to be set and different from device."
+                        "Requires offload_device to be set and different from device.\n"
+                        "Not available on macOS (unified memory architecture)."
                     )
                 ),
                 io.Boolean.Input("swap_io_components",
@@ -74,7 +75,8 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
                     optional=True,
                     tooltip=(
                         "Offload input/output embeddings and normalization layers to reduce VRAM.\n"
-                        "Requires offload_device to be set and different from device."
+                        "Requires offload_device to be set and different from device.\n"
+                        "Not available on macOS (unified memory architecture)."
                     )
                 ),
                 io.Combo.Input("offload_device",
@@ -100,16 +102,19 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
                     )
                 ),
                 io.Combo.Input("attention_mode",
-                    options=["sdpa", "flash_attn"],
+                    options=["sdpa", "flash_attn_2", "flash_attn_3", "sageattn_2", "sageattn_3"],
                     default="sdpa",
                     optional=True,
                     tooltip=(
                         "Attention computation backend:\n"
                         "• sdpa: PyTorch scaled_dot_product_attention (default, stable, always available)\n"
-                        "• flash_attn: Flash Attention 2 (faster on supported hardware, requires flash-attn package)\n"
+                        "• flash_attn_2: Flash Attention 2 (Ampere+, requires flash-attn package)\n"
+                        "• flash_attn_3: Flash Attention 3 (Hopper+, requires flash-attn with FA3 support)\n"
+                        "• sageattn_2: SageAttention 2 (requires sageattention package)\n"
+                        "• sageattn_3: SageAttention 3 (Blackwell/RTX 50xx only, requires sageattn3 package)\n"
                         "\n"
                         "SDPA is recommended - stable and works everywhere.\n"
-                        "Flash Attention provides speedup through optimized CUDA kernels on compatible GPUs."
+                        "Flash Attention and SageAttention provide speedup through optimized CUDA kernels on compatible GPUs."
                     )
                 ),
                 io.Custom("TORCH_COMPILE_ARGS").Input("torch_compile_args",
@@ -142,23 +147,15 @@ class SeedVR2LoadDiTModel(io.ComfyNode):
             cache_model: Whether to keep model loaded between runs
             blocks_to_swap: Number of transformer blocks to swap (requires offload_device != device)
             swap_io_components: Whether to offload I/O components (requires offload_device != device)
-            attention_mode: Attention computation backend ('sdpa' or 'flash_attn')
+            attention_mode: Attention computation backend ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3')
             torch_compile_args: Optional torch.compile configuration from settings node
             
         Returns:
             NodeOutput containing configuration dictionary for SeedVR2 main node
             
         Raises:
-            ValueError: If BlockSwap is enabled but offload_device is invalid
+            ValueError: If cache_model is enabled but offload_device is not set
         """
-        # Validate BlockSwap configuration
-        if (blocks_to_swap > 0 or swap_io_components) and (offload_device == "none" or offload_device == device):
-            raise ValueError(
-                "BlockSwap requires offload_device to be set and different from device. "
-                f"Current: device='{device}', offload_device='{offload_device}'. "
-                "Please set offload_device to a different device (e.g., 'cpu' or another GPU)."
-            )
-        
         # Validate cache_model configuration
         if cache_model and offload_device == "none":
             raise ValueError(

src/interfaces/video_upscaler.py

@@ -6,7 +6,7 @@ Main ComfyUI node for high-quality video upscaling using diffusion models
 import torch
 from comfy_api.latest import io
 from typing import Tuple, Dict, Any, Optional
-from ..utils.constants import get_base_cache_dir
+from ..utils.constants import get_base_cache_dir, __version__
 from ..utils.downloads import download_weight
 from ..utils.debug import Debug
 from ..core.generation_phases import (
@@ -19,7 +19,9 @@ from ..core.generation_utils import (
     setup_generation_context, 
     prepare_runner,
     compute_generation_info,
-    log_generation_start
+    log_generation_start,
+    load_text_embeddings,
+    script_directory
 )
 from ..optimization.memory_manager import (
     cleanup_text_embeddings,
@@ -50,7 +52,7 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
     def define_schema(cls) -> io.Schema:        
         return io.Schema(
             node_id="SeedVR2VideoUpscaler",
-            display_name="SeedVR2 Video Upscaler",
+            display_name=f"SeedVR2 Video Upscaler (v{__version__})",
             category="SEEDVR2",
             description=(
                 "SeedVR2 main upscaling node: processes video frames using DiT and VAE models with diffusion-based enhancement. "
@@ -349,13 +351,12 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
 
         block_swap_config = None
         if blocks_to_swap > 0 or swap_io_components:
-            # Convert offload device string to torch.device for BlockSwap
+            block_swap_config = {
+                "blocks_to_swap": blocks_to_swap,
+                "swap_io_components": swap_io_components,
+            }
             if dit_offload_str != "none":
-                block_swap_config = {
+                block_swap_config["offload_device"] = torch.device(dit_offload_str)
-                    "blocks_to_swap": blocks_to_swap,
-                    "swap_io_components": swap_io_components,
-                    "offload_device": torch.device(dit_offload_str)
-                }
 
         # Device configuration for offloading - convert "none" to None, else torch.device
         vae_offload_str = vae.get("offload_device", "none")
@@ -438,6 +439,10 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
             # Store cache context in ctx for use in generation phases
             ctx['cache_context'] = cache_context
 
+            # Preload text embeddings before Phase 1 to avoid sync stall in Phase 2
+            ctx['text_embeds'] = load_text_embeddings(script_directory, ctx['dit_device'], ctx['compute_dtype'], debug)
+            debug.log("Loaded text embeddings for DiT", category="dit")
+
             debug.log_memory_state("After model preparation", show_tensors=False, detailed_tensors=False)
             debug.end_timer("model_preparation", "Model preparation", force=True, show_breakdown=True)
             
@@ -509,15 +514,21 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
             )
 
             sample = ctx['final_video']
-            
+            debug.log("", category="none", force=True)
+
             # Ensure CPU tensor in float32 for maximum ComfyUI compatibility
             if torch.is_tensor(sample):
                 if sample.is_cuda or sample.is_mps:
                     sample = sample.cpu()
                 if sample.dtype != torch.float32:
-                    sample = sample.to(torch.float32)
+                    src_dtype = sample.dtype
+                    try:
+                        sample = sample.to(torch.float32)
+                        debug.log(f"Converted output from {src_dtype} to float32", category="precision")
+                    except Exception as e:
+                        debug.log(f"Could not convert to float32: {e}. Output is {src_dtype}, compatibility with other nodes not guaranteed", 
+                                  level="WARNING", category="precision", force=True)
 
-            debug.log("", category="none", force=True)
             debug.log("Upscaling completed successfully!", category="success", force=True)
             debug.end_timer("generation", "Video generation")
 
@@ -529,7 +540,7 @@ class SeedVR2VideoUpscaler(io.ComfyNode):
             debug.log_memory_state("After all phases complete", show_tensors=False, detailed_tensors=False)
             
             # Final peak vram summary
-            debug.log_peak_vram_summary()
+            debug.log_peak_memory_summary()
 
             # Final timing summary
             debug.log("", category="none")

src/models/dit_3b/attention.py

@@ -15,8 +15,11 @@
 import torch
 import torch.nn.functional as F
 
-# Import flash_attn with automatic fallback from compatibility layer
+# Import flash/sage attn with automatic fallback from compatibility layer
-from ...optimization.compatibility import flash_attn_varlen_func, FLASH_ATTN_AVAILABLE
+from ...optimization.compatibility import (
+    call_flash_attn_2_varlen, call_flash_attn_3_varlen,
+    call_sage_attn_2_varlen, call_sage_attn_3_varlen
+)
 
 from torch import nn
 
@@ -61,37 +64,6 @@ def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=N
     return torch.cat(output_splits, dim=0)
 
 
-@torch._dynamo.disable
-def _call_flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
-    """
-    Wrapper for flash_attn_varlen_func that handles tensor-to-scalar conversion.
-    
-    This function is excluded from torch.compile because:
-    1. flash_attn is a C++ extension that can't be compiled anyway
-    2. It requires Python int scalars for max_seqlen parameters
-    3. Disabling compilation here keeps the rest of the model compilable
-    """
-    if not FLASH_ATTN_AVAILABLE:
-        raise ImportError("flash_attn is not available")
-    
-    # Convert tensor max_seqlen to Python int if needed
-    if torch.is_tensor(max_seqlen_q):
-        max_seqlen_q = int(max_seqlen_q.item())
-    if torch.is_tensor(max_seqlen_k):
-        max_seqlen_k = int(max_seqlen_k.item())
-    
-    return flash_attn_varlen_func(
-        q=q,
-        k=k,
-        v=v,
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-        **kwargs
-    )
-
-
 class TorchAttention(nn.Module):
     def tflops(self, args, kwargs, output) -> float:
         assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
@@ -107,12 +79,16 @@ class TorchAttention(nn.Module):
 
 class FlashAttentionVarlen(nn.Module):
     """
-    Variable-length attention with configurable backend (Flash Attention or PyTorch SDPA).
+    Variable-length attention with configurable backend.
     
-    Backend selection is validated during model configuration.
+    Supported backends:
-    Compilation behavior:
+    - sdpa: PyTorch SDPA (fully compilable, always available)
-    - SDPA: Fully compilable, optimal performance
+    - flash_attn_2: Flash Attention 2 (Ampere+)
-    - Flash Attention: Uses @torch._dynamo.disable wrapper (C++ extension)
+    - flash_attn_3: Flash Attention 3 (Hopper+)
+    - sageattn_2: SageAttention 2
+    - sageattn_3: SageAttention 3 (Blackwell/RTX 50xx)
+    
+    All non-SDPA backends use @torch._dynamo.disable wrapper (C++ extensions).
     """
 
     def __init__(self, attention_mode: str = 'sdpa', compute_dtype: torch.dtype = None):
@@ -120,7 +96,7 @@ class FlashAttentionVarlen(nn.Module):
         Initialize with specified attention backend.
         
         Args:
-            attention_mode: 'flash_attn' or 'sdpa' (validated externally by validate_flash_attention_availability)
+            attention_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
             compute_dtype: Compute dtype for attention (set by pipeline, defaults to None for auto-detection)
         """
         super().__init__()
@@ -144,11 +120,26 @@ class FlashAttentionVarlen(nn.Module):
             k = k.to(self.compute_dtype)
             v = v.to(self.compute_dtype)
         
-        if self.attention_mode == 'flash_attn':
+        if self.attention_mode == 'flash_attn_3':
-            return _call_flash_attn_varlen_func(
+            return call_flash_attn_3_varlen(
                 q, k, v, cu_seqlens_q, cu_seqlens_k, 
                 max_seqlen_q, max_seqlen_k, **kwargs
             )
+        elif self.attention_mode == 'flash_attn_2':
+            return call_flash_attn_2_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k, 
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
+        elif self.attention_mode == 'sageattn_3':
+            return call_sage_attn_3_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
+        elif self.attention_mode == 'sageattn_2':
+            return call_sage_attn_2_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
         else:
             # PyTorch SDPA
             return pytorch_varlen_attention(

src/models/dit_3b/na.py

@@ -400,22 +400,19 @@ def repeat_concat_idx(
         The text features appear multiple times (once per window) and need
         to be averaged to produce a single set of text features.
         
-        COMPILE OPTIMIZATION: Uses split-stack-mean instead of reshape to avoid .item()
+        COMPILE OPTIMIZATION: Uses unflatten with tensor dims (compile-friendly)
         """
         vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len])
         
-        # Coalesce repeated text using pure tensor operations (no .item())
+        # Coalesce repeated text using unflatten and mean
         txt_splits = _tensor_split(txt_out, repeat_txt_len, dim=0)
         txt_out_coalesced = []
         
-        for txt, base_len in zip(txt_splits, txt_len):
+        for txt in txt_splits:
             # txt has shape (base_len * num_repeats, *other_dims)
-            # Split into num_repeats equal pieces using tensor operations
+            # unflatten to (base_len, num_repeats, *other_dims) then average dim 1
-            split_lens = base_len.unsqueeze(0).repeat_interleave(num_repeats_tensor.squeeze())
+            txt = txt.unflatten(0, (-1, num_repeats_tensor.squeeze())).mean(1)
-            pieces = _tensor_split(txt, split_lens, dim=0)
+            txt_out_coalesced.append(txt)
-            # Stack and average across repetitions
-            txt_avg = torch.stack(pieces, dim=0).mean(0)
-            txt_out_coalesced.append(txt_avg)
         
         return vid_out, torch.cat(txt_out_coalesced)
 

src/models/dit_7b/attention.py

@@ -15,8 +15,11 @@
 import torch
 import torch.nn.functional as F
 
-# Import flash_attn with automatic fallback from compatibility layer
+# Import flash/sage attn with automatic fallback from compatibility layer
-from ...optimization.compatibility import flash_attn_varlen_func, FLASH_ATTN_AVAILABLE
+from ...optimization.compatibility import (
+    call_flash_attn_2_varlen, call_flash_attn_3_varlen,
+    call_sage_attn_2_varlen, call_sage_attn_3_varlen
+)
 
 from torch import nn
 
@@ -61,37 +64,6 @@ def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=N
     return torch.cat(output_splits, dim=0)
 
 
-@torch._dynamo.disable
-def _call_flash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
-    """
-    Wrapper for flash_attn_varlen_func that handles tensor-to-scalar conversion.
-    
-    This function is excluded from torch.compile because:
-    1. flash_attn is a C++ extension that can't be compiled anyway
-    2. It requires Python int scalars for max_seqlen parameters
-    3. Disabling compilation here keeps the rest of the model compilable
-    """
-    if not FLASH_ATTN_AVAILABLE:
-        raise ImportError("flash_attn is not available")
-    
-    # Convert tensor max_seqlen to Python int if needed
-    if torch.is_tensor(max_seqlen_q):
-        max_seqlen_q = int(max_seqlen_q.item())
-    if torch.is_tensor(max_seqlen_k):
-        max_seqlen_k = int(max_seqlen_k.item())
-    
-    return flash_attn_varlen_func(
-        q=q,
-        k=k,
-        v=v,
-        cu_seqlens_q=cu_seqlens_q,
-        cu_seqlens_k=cu_seqlens_k,
-        max_seqlen_q=max_seqlen_q,
-        max_seqlen_k=max_seqlen_k,
-        **kwargs
-    )
-
-
 class TorchAttention(nn.Module):
     def tflops(self, args, kwargs, output) -> float:
         assert len(args) == 0 or len(args) > 2, "query, key should both provided by args / kwargs"
@@ -107,12 +79,16 @@ class TorchAttention(nn.Module):
 
 class FlashAttentionVarlen(nn.Module):
     """
-    Variable-length attention with configurable backend (Flash Attention or PyTorch SDPA).
+    Variable-length attention with configurable backend.
     
-    Backend selection is validated during model configuration.
+    Supported backends:
-    Compilation behavior:
+    - sdpa: PyTorch SDPA (fully compilable, always available)
-    - SDPA: Fully compilable, optimal performance
+    - flash_attn_2: Flash Attention 2 (Ampere+)
-    - Flash Attention: Uses @torch._dynamo.disable wrapper (C++ extension)
+    - flash_attn_3: Flash Attention 3 (Hopper+)
+    - sageattn_2: SageAttention 2
+    - sageattn_3: SageAttention 3 (Blackwell/RTX 50xx)
+    
+    All non-SDPA backends use @torch._dynamo.disable wrapper (C++ extensions).
     """
 
     def __init__(self, attention_mode: str = 'sdpa', compute_dtype: torch.dtype = None):
@@ -120,7 +96,7 @@ class FlashAttentionVarlen(nn.Module):
         Initialize with specified attention backend.
         
         Args:
-            attention_mode: 'flash_attn' or 'sdpa' (validated externally by validate_flash_attention_availability)
+            attention_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
             compute_dtype: Compute dtype for attention (set by pipeline, defaults to None for auto-detection)
         """
         super().__init__()
@@ -144,11 +120,26 @@ class FlashAttentionVarlen(nn.Module):
             k = k.to(self.compute_dtype)
             v = v.to(self.compute_dtype)
         
-        if self.attention_mode == 'flash_attn':
+        if self.attention_mode == 'flash_attn_3':
-            return _call_flash_attn_varlen_func(
+            return call_flash_attn_3_varlen(
                 q, k, v, cu_seqlens_q, cu_seqlens_k, 
                 max_seqlen_q, max_seqlen_k, **kwargs
             )
+        elif self.attention_mode == 'flash_attn_2':
+            return call_flash_attn_2_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k, 
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
+        elif self.attention_mode == 'sageattn_3':
+            return call_sage_attn_3_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
+        elif self.attention_mode == 'sageattn_2':
+            return call_sage_attn_2_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
         else:
             # PyTorch SDPA
             return pytorch_varlen_attention(

src/models/dit_7b/na.py

@@ -400,22 +400,19 @@ def repeat_concat_idx(
         The text features appear multiple times (once per window) and need
         to be averaged to produce a single set of text features.
         
-        COMPILE OPTIMIZATION: Uses split-stack-mean instead of reshape to avoid .item()
+        COMPILE OPTIMIZATION: Uses unflatten with tensor dims (compile-friendly)
         """
         vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len])
         
-        # Coalesce repeated text using pure tensor operations (no .item())
+        # Coalesce repeated text using unflatten and mean
         txt_splits = _tensor_split(txt_out, repeat_txt_len, dim=0)
         txt_out_coalesced = []
         
-        for txt, base_len in zip(txt_splits, txt_len):
+        for txt in txt_splits:
             # txt has shape (base_len * num_repeats, *other_dims)
-            # Split into num_repeats equal pieces using tensor operations
+            # unflatten to (base_len, num_repeats, *other_dims) then average dim 1
-            split_lens = base_len.unsqueeze(0).repeat_interleave(num_repeats_tensor.squeeze())
+            txt = txt.unflatten(0, (-1, num_repeats_tensor.squeeze())).mean(1)
-            pieces = _tensor_split(txt, split_lens, dim=0)
+            txt_out_coalesced.append(txt)
-            # Stack and average across repetitions
-            txt_avg = torch.stack(pieces, dim=0).mean(0)
-            txt_out_coalesced.append(txt_avg)
         
         return vid_out, torch.cat(txt_out_coalesced)
 

src/models/video_vae_v3/modules/attn_video_vae.py

@@ -1093,7 +1093,7 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
     ):
         extra_cond_dim = kwargs.pop("extra_cond_dim") if "extra_cond_dim" in kwargs else None
         self.slicing_sample_min_size = slicing_sample_min_size
-        self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
+        self.slicing_latent_min_size = max(1, slicing_sample_min_size // (2**temporal_scale_num))
 
         super().__init__(
             in_channels=in_channels,
@@ -1224,6 +1224,10 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         
         output = causal_conv_gather_outputs(output)
         
+        # MPS memory leak workaround (pytorch/pytorch#155060)
+        if self.device.type == 'mps':
+            torch.mps.empty_cache()
+        
         # Only transfer back if needed
         return output if output.device == x.device else output.to(x.device)
 
@@ -1240,6 +1244,10 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         output = self.decoder(_z, memory_state=memory_state)
         output = causal_conv_gather_outputs(output)
         
+        # MPS memory leak workaround (pytorch/pytorch#155060)
+        if self.device.type == 'mps':
+            torch.mps.empty_cache()
+        
         # Only transfer back if needed
         return output if output.device == z.device else output.to(z.device)
 
@@ -1308,7 +1316,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         if H <= tile_h and W <= tile_w:
             return self.slicing_encode(x)
         else:
-            self.debug.log(f"Using VAE tiled encoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
+            if self.debug:
+                self.debug.log(f"Using VAE tiled encoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
 
         # Spatial scale factor (output/latent)
         scale_factor = self.spatial_downsample_factor
@@ -1481,7 +1490,8 @@ class VideoAutoencoderKL(diffusers.AutoencoderKL):
         if H <= latent_tile_h and W <= latent_tile_w:
             return self.slicing_decode(z)
         else:
-            self.debug.log(f"Using VAE tiled decoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
+            if self.debug:
+                self.debug.log(f"Using VAE tiled decoding (Tile: {tile_size}, Overlap: {tile_overlap})", category="vae", force=True, indent_level=1)
         
         latent_overlap_h = max(0, min((overlap_h // scale_factor), latent_tile_h - 1))
         latent_overlap_w = max(0, min((overlap_w // scale_factor), latent_tile_w - 1))
@@ -1708,7 +1718,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
         if split_size is not None:
             self.enable_slicing()
             self.slicing_sample_min_size = split_size
-            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
+            self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor)
         else:
             self.disable_slicing()
         for module in self.modules():

src/models/video_vae_v3/modules/causal_inflation_lib.py

@@ -91,18 +91,26 @@ class InflatedCausalConv3d(Conv3d):
         Workaround: Call torch.cudnn_convolution directly to bypass buggy layer.
         Status is logged at startup in compatibility.py.
         """
-        if NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND and weight.dtype in (torch.float16, torch.bfloat16):
+        if (NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND and 
-            # Direct cuDNN call bypasses buggy PyTorch dispatch layer
+            weight.dtype in (torch.float16, torch.bfloat16) and 
-            out = torch.cudnn_convolution(
+            hasattr(torch.backends.cudnn, 'is_available') and
-                input, weight, self.padding, self.stride, self.dilation, self.groups,
+            torch.backends.cudnn.is_available() and
-                benchmark=False, deterministic=False, allow_tf32=True
+            getattr(torch.backends.cudnn, 'enabled', True)):
-            )
+            try:
-            if bias is not None:
+                # Direct cuDNN call bypasses buggy PyTorch dispatch layer (NVIDIA only)
-                out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))
+                out = torch.cudnn_convolution(
-            return out
+                    input, weight, self.padding, self.stride, self.dilation, self.groups,
-        else:
+                    benchmark=False, deterministic=False, allow_tf32=True
-            # Use standard path for unaffected configurations
+                )
-            return super()._conv_forward(input, weight, bias, *args, **kwargs)
+                if bias is not None:
+                    out += bias.reshape((1, -1) + (1,) * (out.ndim - 2))
+                return out
+            except RuntimeError:
+                # Fallback if direct cuDNN call fails (dev builds, edge cases)
+                pass
+        
+        # Use standard path for unaffected configurations or if workaround failed
+        return super()._conv_forward(input, weight, bias, *args, **kwargs)
 
     def memory_limit_conv(
         self,

src/models/video_vae_v3/modules/video_vae.py

@@ -733,7 +733,7 @@ class VideoAutoencoderKL(nn.Module):
         if slicing_sample_min_size is None:
             slicing_sample_min_size = temporal_downsample_factor
         self.slicing_sample_min_size = slicing_sample_min_size
-        self.slicing_latent_min_size = slicing_sample_min_size // (2**temporal_scale_num)
+        self.slicing_latent_min_size = max(1, slicing_sample_min_size // (2**temporal_scale_num))
 
         # pass init params to Encoder
         self.encoder = Encoder3D(
@@ -886,7 +886,7 @@ class VideoAutoencoderKL(nn.Module):
         if split_size is not None:
             self.enable_slicing()
             self.slicing_sample_min_size = split_size
-            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
+            self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor)
         else:
             self.disable_slicing()
         for module in self.modules():
@@ -950,7 +950,7 @@ class VideoAutoencoderKLWrapper(VideoAutoencoderKL):
             self.disable_slicing()
         self.slicing_sample_min_size = split_size
         if split_size is not None:
-            self.slicing_latent_min_size = split_size // self.temporal_downsample_factor
+            self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor)
         for module in self.modules():
             if isinstance(module, InflatedCausalConv3d):
                 module.set_memory_device(memory_device)

src/optimization/blockswap.py

@@ -47,6 +47,78 @@ def is_blockswap_enabled(config: Optional[Dict[str, Any]]) -> bool:
     return blocks_to_swap > 0 or swap_io_components
 
 
+def validate_blockswap_config(
+    block_swap_config: Optional[Dict[str, Any]],
+    dit_device: 'torch.device',
+    dit_offload_device: Optional['torch.device'],
+    debug: 'Debug'
+) -> Optional[Dict[str, Any]]:
+    """
+    Validate and potentially modify BlockSwap configuration.
+    
+    Performs platform-specific validation and configuration adjustment:
+    - On macOS (MPS): Auto-disables BlockSwap since unified memory makes it meaningless
+    - On other platforms: Validates that offload_device is properly configured
+    
+    This is the single authoritative validation point for BlockSwap configuration,
+    called early in configure_runner() before any model loading.
+    
+    Args:
+        block_swap_config: BlockSwap configuration dictionary (may be None)
+        dit_device: Target device for DiT model inference
+        dit_offload_device: Device for offloading DiT blocks (may be None)
+        debug: Debug instance for logging warnings/errors
+        
+    Returns:
+        Validated/modified block_swap_config (may be None or modified copy)
+        
+    Raises:
+        ValueError: If BlockSwap is enabled but offload_device is invalid (non-MPS only)
+    """
+    if not is_blockswap_enabled(block_swap_config):
+        return block_swap_config
+    
+    blocks_to_swap = block_swap_config.get("blocks_to_swap", 0)
+    swap_io_components = block_swap_config.get("swap_io_components", False)
+    
+    # Check for macOS unified memory - BlockSwap is meaningless there
+    if dit_device.type == "mps":
+        debug.log(
+            f"BlockSwap disabled: macOS uses unified memory (no separate VRAM/RAM). "
+            f"Ignoring blocks_to_swap={blocks_to_swap}, swap_io_components={swap_io_components}",
+            level="WARNING", category="blockswap", force=True
+        )
+        # Return disabled config
+        return {
+            **block_swap_config,
+            "blocks_to_swap": 0,
+            "swap_io_components": False
+        }
+    
+    # Validate offload_device is set and different from dit_device
+    offload_device_valid = (
+        dit_offload_device is not None and 
+        str(dit_offload_device) != str(dit_device)
+    )
+    
+    if not offload_device_valid:
+        config_details = []
+        if blocks_to_swap > 0:
+            config_details.append(f"blocks_to_swap={blocks_to_swap}")
+        if swap_io_components:
+            config_details.append("swap_io_components=True")
+        
+        offload_str = str(dit_offload_device) if dit_offload_device else "none"
+        raise ValueError(
+            f"BlockSwap enabled ({', '.join(config_details)}) but dit_offload_device is invalid. "
+            f"Current: device='{dit_device}', dit_offload_device='{offload_str}'. "
+            f"BlockSwap requires offload_device on the DiT Model to be set and different from device. "
+            f"Set --dit_offload_device cpu or disable BlockSwap."
+        )
+    
+    return block_swap_config
+
+
 # Timing helpers marked to skip torch.compile tracing
 # These functions are excluded from Dynamo's graph tracing to avoid warnings
 # about non-traceable builtins like time.time(), but they still execute normally
@@ -125,7 +197,7 @@ def apply_block_swap_to_dit(
     
     debug.start_timer("apply_blockswap")
 
-    # Get the actual model (handle FP8CompatibleDiT wrapper)
+    # Get the actual model (handle CompatibleDiT wrapper)
     model = runner.dit
     if hasattr(model, "dit_model"):
         model = model.dit_model
@@ -144,22 +216,21 @@ def apply_block_swap_to_dit(
 
     total_blocks = len(model.blocks)
     
-    # Log configuration clearly based on what's enabled
+    # Clamp blocks_to_swap to available blocks BEFORE logging
-    block_text = "block" if blocks_to_swap <= 1 else "blocks"
+    effective_blocks = min(blocks_to_swap, total_blocks) if blocks_to_swap > 0 else 0
-    if blocks_to_swap > 0 and swap_io_components:
-        debug.log(f"BlockSwap: {blocks_to_swap} transformer {block_text} + I/O components offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
-    elif blocks_to_swap > 0:
-        debug.log(f"BlockSwap: {blocks_to_swap} transformer {block_text} offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
-    elif swap_io_components:
-        debug.log(f"BlockSwap: I/O components offloaded to {str(offload_device).upper()} (blocks remain on GPU)", category="blockswap", force=True)
     
-    debug.log(f"Model has {total_blocks} transformer blocks", category="blockswap")
+    # Log configuration clearly based on what's enabled
+    block_text = "block" if effective_blocks <= 1 else "blocks"
+    if effective_blocks > 0 and swap_io_components:
+        debug.log(f"BlockSwap: {effective_blocks}/{total_blocks} transformer {block_text} + I/O components offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
+    elif effective_blocks > 0:
+        debug.log(f"BlockSwap: {effective_blocks}/{total_blocks} transformer {block_text} offloaded to {str(offload_device).upper()}", category="blockswap", force=True)
+    elif swap_io_components:
+        debug.log(f"BlockSwap: I/O components offloaded to {str(offload_device).upper()} (0/{total_blocks} blocks swapped)", category="blockswap", force=True)
     
     # Configure model with blockswap attributes
     if blocks_to_swap > 0:
-        blocks_to_swap = min(blocks_to_swap, total_blocks)
+        model.blocks_to_swap = effective_blocks - 1  # Convert to 0-indexed
-        model.blocks_to_swap = blocks_to_swap - 1  # Convert to 0-indexed
-        debug.log(f"Transformer blocks to swap: {blocks_to_swap}/{total_blocks}", category="blockswap")
     else:
         # No block swapping, set to -1 so no blocks match the swap condition
         model.blocks_to_swap = -1
@@ -193,7 +264,7 @@ def apply_block_swap_to_dit(
     runner._blockswap_active = True
 
     # Store configuration for debugging and cleanup
-    runner._block_swap_config = {
+    model._block_swap_config = {
         "blocks_swapped": blocks_to_swap,
         "swap_io_components": swap_io_components,
         "total_blocks": total_blocks,
@@ -650,11 +721,11 @@ def _protect_model_from_move(
     
     Wraps model.to() method to prevent other code from accidentally moving
     the entire model to GPU, which would defeat BlockSwap's memory savings.
-    Allows movement only when explicitly bypassed via runner flag.
+    Allows movement only when explicitly bypassed via model flag.
     
     Args:
         model: DiT model to protect
-        runner: VideoDiffusionInfer instance (stored as weak reference)
+        runner: VideoDiffusionInfer instance (for active status check)
         debug: Debug instance for logging (required)
     """
     if not hasattr(model, '_original_to'):
@@ -665,34 +736,46 @@ def _protect_model_from_move(
         # Define the protected method without closures
         def protected_model_to(self, device, *args, **kwargs):
             # Check if protection is temporarily bypassed for offloading
+            # Flag is stored on model itself (not runner) to survive runner recreation
+            if getattr(self, "_blockswap_bypass_protection", False):
+                # Protection bypassed, allow movement
+                if hasattr(self, '_original_to'):
+                    return self._original_to(device, *args, **kwargs)
+            
+            # Get configured offload device directly from model
+            blockswap_offload_device = "cpu"  # default
+            if hasattr(self, "_block_swap_config"):
+                blockswap_offload_device = self._block_swap_config.get("offload_device", "cpu")
+            
+            # Check if BlockSwap is currently active via runner weak reference
             runner_ref = getattr(self, '_blockswap_runner_ref', None)
+            blockswap_is_active = False
             if runner_ref:
                 runner_obj = runner_ref()
-                if runner_obj and getattr(runner_obj, "_blockswap_bypass_protection", False):
+                if runner_obj and hasattr(runner_obj, "_blockswap_active"):
-                    # Protection bypassed, allow movement
+                    blockswap_is_active = runner_obj._blockswap_active
-                    if hasattr(self, '_original_to'):
-                        return self._original_to(device, *args, **kwargs)
             
-            # Check blockswap status using weak reference
+            # Block attempts to move model away from configured offload device when active
-            # Get configured offload device from runner
+            if blockswap_is_active and str(device) != str(blockswap_offload_device):
-            blockswap_offload_device = "cpu" # default
+                # Get debug instance from runner if available
-            if runner_ref:
+                debug_instance = None
-                runner_obj = runner_ref()
+                if runner_ref:
-                if runner_obj and hasattr(runner_obj, "_block_swap_config"):
+                    runner_obj = runner_ref()
-                    blockswap_offload_device = runner_obj._block_swap_config.get("offload_device", "cpu")
+                    if runner_obj and hasattr(runner_obj, 'debug'):
+                        debug_instance = runner_obj.debug
                 
-                # Block attempts to move model away from configured offload device
+                if debug_instance:
-                if str(device) != str(blockswap_offload_device):
+                    debug_instance.log(
-                    if runner_obj and hasattr(runner_obj, "_blockswap_active") and runner_obj._blockswap_active:
+                        f"Blocked attempt to move BlockSwap model from {blockswap_offload_device} to {device}",
-                        debug.log(f"Blocked attempt to move blockswapped model from {blockswap_offload_device} to {device}", 
+                        level="WARNING", category="blockswap", force=True
-                                level="WARNING", category="blockswap", force=True)
+                    )
-                        return self
+                return self
             
-            # Use original method stored as attribute
+            # Allow movement (either bypass is enabled or target is offload device)
             if hasattr(self, '_original_to'):
                 return self._original_to(device, *args, **kwargs)
             else:
-                # This shouldn't happen, but fallback to super().to()
+                # Fallback - shouldn't happen
                 return super(type(self), self).to(device, *args, **kwargs)
         
         # Bind as a method to the model instance
@@ -712,7 +795,13 @@ def set_blockswap_bypass(runner, bypass: bool, debug):
     if not hasattr(runner, "_blockswap_active") or not runner._blockswap_active:
         return
     
-    runner._blockswap_bypass_protection = bypass
+    # Get the actual model (handle CompatibleDiT wrapper)
+    model = runner.dit
+    if hasattr(model, "dit_model"):
+        model = model.dit_model
+    
+    # Store on model so it survives runner recreation during caching
+    model._blockswap_bypass_protection = bypass
     
     if bypass:
         debug.log("BlockSwap protection disabled to allow model DiT offloading", category="success")
@@ -741,11 +830,16 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
     
     debug = runner.debug
     
-    # Check if there's any BlockSwap state to clean up
+    # Get the actual model (handle CompatibleDiT wrapper)
+    model = runner.dit
+    if hasattr(model, "dit_model"):
+        model = model.dit_model
+    
+    # Check if there's any BlockSwap state to clean up (check both runner and model)
     has_blockswap_state = (
         hasattr(runner, "_blockswap_active") or 
-        hasattr(runner, "_block_swap_config") or
+        hasattr(model, "_block_swap_config") or
-        hasattr(runner, "_blockswap_bypass_protection")
+        hasattr(model, "_blockswap_bypass_protection")
     )
     
     if not has_blockswap_state:
@@ -757,14 +851,14 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
         # Minimal cleanup for caching - just mark as inactive and allow offloading
         # Everything else stays intact for fast reactivation
         if hasattr(runner, "_blockswap_active") and runner._blockswap_active:
-            if not getattr(runner, "_blockswap_bypass_protection", False):
+            if not getattr(model, "_blockswap_bypass_protection", False):
                 set_blockswap_bypass(runner=runner, bypass=True, debug=debug)
             runner._blockswap_active = False
         debug.log("BlockSwap deactivated for caching (configuration preserved)", category="success")
         return
 
     # Full cleanup when not caching
-    # Get the actual model (handle FP8CompatibleDiT wrapper)
+    # Get the actual model (handle CompatibleDiT wrapper)
     model = runner.dit
     if hasattr(model, "dit_model"):
         model = model.dit_model
@@ -829,7 +923,7 @@ def cleanup_blockswap(runner, keep_state_for_cache=False):
 
     # 5. Clean up BlockSwap-specific attributes
     for attr in ['_blockswap_runner_ref', 'blocks_to_swap', 'main_device', 
-                 'offload_device', '_blockswap_configured']:
+                 'offload_device']:
         if hasattr(model, attr):
             delattr(model, attr)
 

src/optimization/compatibility.py

@@ -5,51 +5,546 @@ Contains FP8/FP16 compatibility layers and wrappers for different model architec
 Extracted from: seedvr2.py (lines 1045-1630)
 """
 
-import torch
+# Compatibility shims - Must run before any torch/diffusers import
+import sys
 import types
+import importlib.machinery
+
+
+def ensure_triton_compat():
+    """Create minimal triton.ops stubs only if missing, to allow bitsandbytes import."""
+    if 'triton.ops.matmul_perf_model' in sys.modules:
+        return
+    
+    try:
+        from triton.ops.matmul_perf_model import early_config_prune  # noqa: F401
+        return
+    except (ImportError, ModuleNotFoundError, AttributeError):
+        pass
+    
+    if 'triton.ops' not in sys.modules:
+        sys.modules['triton.ops'] = types.ModuleType('triton.ops')
+    
+    matmul_perf = types.ModuleType('triton.ops.matmul_perf_model')
+    matmul_perf.early_config_prune = lambda configs, *a, **kw: configs
+    matmul_perf.estimate_matmul_time = lambda *a, **kw: 0.0
+    
+    sys.modules['triton.ops'].matmul_perf_model = matmul_perf
+    sys.modules['triton.ops.matmul_perf_model'] = matmul_perf
+
+
+def ensure_flash_attn_safe():
+    """
+    Pre-test flash_attn package; stub if DLL is broken.
+    Prevents diffusers from crashing when flash_attn has broken DLLs.
+    """
+    if 'flash_attn' in sys.modules:
+        return  # Already loaded
+    
+    try:
+        import flash_attn
+    except (ImportError, OSError):
+        # DLL broken or not installed - create stub with proper __spec__
+        stub = types.ModuleType('flash_attn')
+        stub.__spec__ = importlib.machinery.ModuleSpec('flash_attn', None)
+        stub.__file__ = None
+        stub.__path__ = []
+        stub.__loader__ = None
+        # Provide attributes that diffusers/transformers import
+        stub.flash_attn_func = None
+        stub.flash_attn_varlen_func = None
+        sys.modules['flash_attn'] = stub
+
+
+def ensure_xformers_flash_compat():
+    """
+    Pre-test xformers._C_flashattention; stub if DLL is broken.
+    Prevents xformers.ops.fmha.flash from crashing on import.
+    """
+    if 'xformers._C_flashattention' in sys.modules:
+        return  # Already loaded
+    
+    try:
+        from xformers import _C_flashattention  # noqa: F401
+    except (ImportError, OSError):
+        # DLL broken or not installed - create stub with proper __spec__
+        class _FailingStub(types.ModuleType):
+            """Stub that lets xformers gracefully disable its flash backend."""
+            def __getattr__(self, name):
+                raise ImportError("_C_flashattention unavailable")
+        
+        stub = _FailingStub('xformers._C_flashattention')
+        stub.__spec__ = importlib.machinery.ModuleSpec('xformers._C_flashattention', None)
+        stub.__file__ = None
+        stub.__path__ = []
+        stub.__loader__ = None
+        sys.modules['xformers._C_flashattention'] = stub
+
+
+def ensure_bitsandbytes_safe():
+    """
+    Pre-test bitsandbytes; stub if broken to prevent import conflicts.
+    
+    On some systems (e.g., ROCm without proper binaries), bitsandbytes registers
+    PyTorch kernels during import then fails. If another node already triggered 
+    this partial load, re-importing causes kernel registration conflicts.
+    
+    This shim catches such failures and stubs the module so diffusers can load
+    gracefully without bitsandbytes quantization support.
+    """
+    if 'bitsandbytes' in sys.modules:
+        return  # Already loaded or stubbed
+    
+    try:
+        import bitsandbytes
+        # Success - bitsandbytes works, other nodes can use it
+    except (ImportError, OSError, RuntimeError, ValueError):
+        # Installation broken, not present, or version detection failed - create stub
+        stub = types.ModuleType('bitsandbytes')
+        stub.__spec__ = importlib.machinery.ModuleSpec('bitsandbytes', None)
+        stub.__file__ = None
+        stub.__path__ = []
+        stub.__version__ = "0.0.0"
+        sys.modules['bitsandbytes'] = stub
+
+
+# Run all shims immediately on import, before torch/diffusers
+ensure_triton_compat()
+ensure_flash_attn_safe()
+ensure_xformers_flash_compat()
+ensure_bitsandbytes_safe()
+
+
+import torch
 import os
 
-# Flash Attention & Triton Compatibility Layer
-# 1. Flash Attention - speedup for attention operations
-try:
-    from flash_attn import flash_attn_varlen_func
-    FLASH_ATTN_AVAILABLE = True
-except ImportError:
-    flash_attn_varlen_func = None
-    FLASH_ATTN_AVAILABLE = False
 
-def validate_flash_attention_availability(requested_mode: str, debug=None) -> str:
+# Flash/Sage Attention & Triton Compatibility Layer
+
+# 1. Flash Attention 3 (Hopper+, faster, no dropout/window support)
+flash_attn_3_varlen_func = None
+FLASH_ATTN_3_AVAILABLE = False
+try:
+    import flash_attn_interface
+    flash_attn_3_varlen_func = flash_attn_interface.flash_attn_varlen_func
+    FLASH_ATTN_3_AVAILABLE = True
+except (ImportError, AttributeError, OSError):
+    pass
+
+# 2. Flash Attention 2 (wider compatibility, supports dropout/window)
+flash_attn_2_varlen_func = None
+FLASH_ATTN_2_AVAILABLE = False
+try:
+    from flash_attn import flash_attn_varlen_func as _fa2_varlen
+    import flash_attn_2_cuda  # noqa: F401
+    flash_attn_2_varlen_func = _fa2_varlen
+    FLASH_ATTN_2_AVAILABLE = True
+except (ImportError, AttributeError, OSError):
+    pass
+
+FLASH_ATTN_AVAILABLE = FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE
+
+# 3. SageAttention 2 (varlen support)
+sageattn_varlen = None
+SAGE_ATTN_2_AVAILABLE = False
+try:
+    from sageattention import sageattn_varlen as _sa2_varlen
+    sageattn_varlen = _sa2_varlen
+    SAGE_ATTN_2_AVAILABLE = True
+except (ImportError, AttributeError, OSError):
+    pass
+
+# 4. SageAttention 3 / Blackwell (RTX 50xx only, batched attention)
+sageattn_blackwell = None
+SAGE_ATTN_3_AVAILABLE = False
+try:
+    from sageattn3 import sageattn3_blackwell as _sa3_blackwell
+    sageattn_blackwell = _sa3_blackwell
+    SAGE_ATTN_3_AVAILABLE = True
+except (ImportError, AttributeError, OSError):
+    try:
+        from sageattention import sageattn_blackwell as _sa3_blackwell
+        sageattn_blackwell = _sa3_blackwell
+        SAGE_ATTN_3_AVAILABLE = True
+    except (ImportError, AttributeError, OSError):
+        pass
+
+SAGE_ATTN_AVAILABLE = SAGE_ATTN_2_AVAILABLE or SAGE_ATTN_3_AVAILABLE
+
+
+def validate_attention_mode(requested_mode: str, debug=None) -> str:
     """
-    Validate Flash Attention availability and warn if fallback needed.
+    Validate attention mode availability with automatic fallback.
     
     Args:
-        requested_mode: Either 'flash_attn' or 'sdpa'
+        requested_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3'
         debug: Optional debug instance for logging
         
     Returns:
-        Validated mode ('flash_attn' or 'sdpa')
+        Validated mode that is available
     """
-    if requested_mode == 'flash_attn' and not FLASH_ATTN_AVAILABLE:
+    # Flash Attention 3
+    if requested_mode == 'flash_attn_3':
+        if FLASH_ATTN_3_AVAILABLE:
+            return requested_mode
+        if FLASH_ATTN_2_AVAILABLE:
+            if debug:
+                debug.log(
+                    "Flash Attention 3 not available (requires Hopper+ GPU and flash-attn with FA3 support).\n"
+                    "Falling back to Flash Attention 2.",
+                    level="WARNING", category="setup", force=True
+                )
+            return 'flash_attn_2'
         error_msg = (
-            f"Cannot use 'flash_attn' attention mode: Flash Attention is not installed.\n"
+            "Cannot use 'flash_attn_3' attention mode: Flash Attention is not installed.\n"
-            f"\n"
+            "\n"
-            f"Flash Attention provides speedup on some hardware through optimized CUDA kernels.\n"
+            "Flash Attention 3 provides maximum speedup on Hopper+ GPUs through optimized CUDA kernels.\n"
-            f"Falling back to PyTorch SDPA (scaled dot-product attention).\n"
+            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
-            f"\n"
+            "\n"
-            f"To fix this issue:\n"
+            "To fix this issue:\n"
-            f"  1. Install Flash Attention: pip install flash-attn\n"
+            "  1. Install Flash Attention: pip install flash-attn\n"
-            f"  2. OR change attention_mode to 'sdpa' (default, always available)\n"
+            "  2. OR change attention_mode to 'sdpa' (default, always available)\n"
-            f"\n"
+            "\n"
-            f"For more info: https://github.com/Dao-AILab/flash-attention"
+            "For more info: https://github.com/Dao-AILab/flash-attention"
+        )
+        if debug:
+            debug.log(error_msg, level="WARNING", category="setup", force=True)
+        return 'sdpa'
+    
+    # Flash Attention 2
+    if requested_mode == 'flash_attn_2':
+        if FLASH_ATTN_2_AVAILABLE:
+            return requested_mode
+        error_msg = (
+            "Cannot use 'flash_attn_2' attention mode: Flash Attention 2 is not installed.\n"
+            "\n"
+            "Flash Attention 2 provides speedup on Ampere+ GPUs through optimized CUDA kernels.\n"
+            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
+            "\n"
+            "To fix this issue:\n"
+            "  1. Install Flash Attention: pip install flash-attn\n"
+            "  2. OR change attention_mode to 'sdpa' (default, always available)\n"
+            "\n"
+            "For more info: https://github.com/Dao-AILab/flash-attention"
+        )
+        if debug:
+            debug.log(error_msg, level="WARNING", category="setup", force=True)
+        return 'sdpa'
+    
+    # SageAttention 3 (Blackwell)
+    if requested_mode == 'sageattn_3':
+        if SAGE_ATTN_3_AVAILABLE:
+            return requested_mode
+        if SAGE_ATTN_2_AVAILABLE:
+            if debug:
+                debug.log(
+                    "SageAttention 3 (Blackwell) not available (requires RTX 50xx GPU and sageattn3 package).\n"
+                    "Falling back to SageAttention 2.",
+                    level="WARNING", category="setup", force=True
+                )
+            return 'sageattn_2'
+        error_msg = (
+            "Cannot use 'sageattn_3' attention mode: SageAttention is not installed.\n"
+            "\n"
+            "SageAttention 3 provides maximum speedup on Blackwell (RTX 50xx) GPUs.\n"
+            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
+            "\n"
+            "To fix this issue:\n"
+            "  1. Install SageAttention: pip install sageattention\n"
+            "  2. For SA3 Blackwell support: pip install sageattn3\n"
+            "  3. OR change attention_mode to 'flash_attn_2' or 'sdpa'\n"
+            "\n"
+            "For more info: https://github.com/thu-ml/SageAttention"
+        )
+        if debug:
+            debug.log(error_msg, level="WARNING", category="setup", force=True)
+        return 'sdpa'
+    
+    # SageAttention 2
+    if requested_mode == 'sageattn_2':
+        if SAGE_ATTN_2_AVAILABLE:
+            return requested_mode
+        error_msg = (
+            "Cannot use 'sageattn_2' attention mode: SageAttention is not installed.\n"
+            "\n"
+            "SageAttention provides speedup on NVIDIA GPUs through optimized CUDA kernels.\n"
+            "Falling back to PyTorch SDPA (scaled dot-product attention).\n"
+            "\n"
+            "To fix this issue:\n"
+            "  1. Install SageAttention: pip install sageattention\n"
+            "  2. OR change attention_mode to 'flash_attn_2' or 'sdpa'\n"
+            "\n"
+            "For more info: https://github.com/thu-ml/SageAttention"
         )
         if debug:
             debug.log(error_msg, level="WARNING", category="setup", force=True)
-        
         return 'sdpa'
     
     return requested_mode
 
 
+@torch._dynamo.disable
+def call_flash_attn_2_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
+    """
+    Wrapper for Flash Attention 2 flash_attn_varlen_func that handles tensor-to-scalar conversion.
+    
+    Flash Attention 2 supports dropout_p and window_size parameters.
+    Works on Ampere+ GPUs (RTX 30xx, 40xx, A100, etc.).
+    
+    This function is excluded from torch.compile because:
+    1. flash_attn is a C++ extension that can't be compiled anyway
+    2. It requires Python int scalars for max_seqlen parameters
+    3. Disabling compilation here keeps the rest of the model compilable
+    
+    Args:
+        q: Query tensor (total_seq, heads, head_dim)
+        k: Key tensor (total_seq, heads, head_dim)
+        v: Value tensor (total_seq, heads, head_dim)
+        cu_seqlens_q: Cumulative sequence lengths for queries
+        cu_seqlens_k: Cumulative sequence lengths for keys
+        max_seqlen_q: Maximum query sequence length (can be tensor or int)
+        max_seqlen_k: Maximum key sequence length (can be tensor or int)
+        **kwargs: Additional arguments (dropout_p, softmax_scale, causal, window_size, deterministic)
+        
+    Returns:
+        Attention output tensor (total_seq, heads, head_dim)
+    """
+    if not FLASH_ATTN_2_AVAILABLE:
+        raise ImportError("Flash Attention 2 is not available")
+    
+    # Convert tensor max_seqlen to Python int if needed
+    if torch.is_tensor(max_seqlen_q):
+        max_seqlen_q = int(max_seqlen_q.item())
+    if torch.is_tensor(max_seqlen_k):
+        max_seqlen_k = int(max_seqlen_k.item())
+    
+    return flash_attn_2_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        **kwargs
+    )
+
+
+@torch._dynamo.disable
+def call_flash_attn_3_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
+    """
+    Wrapper for Flash Attention 3 flash_attn_varlen_func that handles tensor-to-scalar conversion.
+    
+    Flash Attention 3 is faster than FA2 but does NOT support dropout_p and window_size.
+    Works on Hopper+ GPUs (H100, etc.) - requires flash_attn_interface package.
+    
+    This function is excluded from torch.compile because:
+    1. flash_attn is a C++ extension that can't be compiled anyway
+    2. It requires Python int scalars for max_seqlen parameters
+    3. Disabling compilation here keeps the rest of the model compilable
+    
+    Args:
+        q: Query tensor (total_seq, heads, head_dim)
+        k: Key tensor (total_seq, heads, head_dim)
+        v: Value tensor (total_seq, heads, head_dim)
+        cu_seqlens_q: Cumulative sequence lengths for queries
+        cu_seqlens_k: Cumulative sequence lengths for keys
+        max_seqlen_q: Maximum query sequence length (can be tensor or int)
+        max_seqlen_k: Maximum key sequence length (can be tensor or int)
+        **kwargs: Additional arguments (softmax_scale, causal, deterministic)
+                  Note: dropout_p and window_size are ignored (not supported by FA3)
+        
+    Returns:
+        Attention output tensor (total_seq, heads, head_dim)
+    """
+    if not FLASH_ATTN_3_AVAILABLE:
+        raise ImportError("Flash Attention 3 is not available")
+    
+    # Convert tensor max_seqlen to Python int if needed
+    if torch.is_tensor(max_seqlen_q):
+        max_seqlen_q = int(max_seqlen_q.item())
+    if torch.is_tensor(max_seqlen_k):
+        max_seqlen_k = int(max_seqlen_k.item())
+    
+    # FA3 doesn't support dropout_p and window_size - filter them out
+    fa3_kwargs = {key: val for key, val in kwargs.items() if key not in ('dropout_p', 'window_size')}
+    
+    # FA3 returns a tuple (output, softmax_lse), we only need output
+    return flash_attn_3_varlen_func(
+        q=q,
+        k=k,
+        v=v,
+        cu_seqlens_q=cu_seqlens_q,
+        cu_seqlens_k=cu_seqlens_k,
+        max_seqlen_q=max_seqlen_q,
+        max_seqlen_k=max_seqlen_k,
+        seqused_q=None,
+        seqused_k=None,
+        **fa3_kwargs
+    )[0]
+
+
+@torch._dynamo.disable
+def call_sage_attn_2_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
+    """
+    Wrapper for SageAttention 2 sageattn_varlen that handles tensor-to-scalar conversion.
+    
+    SageAttention 2 provides optimized attention for NVIDIA GPUs with native varlen support.
+    Works on most modern NVIDIA GPUs.
+    
+    This function is excluded from torch.compile because:
+    1. SageAttention is a C++ extension that can't be compiled anyway
+    2. It requires Python int scalars for max_seqlen parameters
+    3. Disabling compilation here keeps the rest of the model compilable
+    
+    Args:
+        q: Query tensor (total_seq, heads, head_dim)
+        k: Key tensor (total_seq, heads, head_dim)
+        v: Value tensor (total_seq, heads, head_dim)
+        cu_seqlens_q: Cumulative sequence lengths for queries
+        cu_seqlens_k: Cumulative sequence lengths for keys
+        max_seqlen_q: Maximum query sequence length (can be tensor or int)
+        max_seqlen_k: Maximum key sequence length (can be tensor or int)
+        **kwargs: Additional arguments (causal supported, others ignored)
+        
+    Returns:
+        Attention output tensor (total_seq, heads, head_dim)
+    """
+    if not SAGE_ATTN_2_AVAILABLE:
+        raise ImportError("SageAttention 2 is not available")
+    
+    # Convert tensor max_seqlen to Python int if needed
+    if torch.is_tensor(max_seqlen_q):
+        max_seqlen_q = int(max_seqlen_q.item())
+    if torch.is_tensor(max_seqlen_k):
+        max_seqlen_k = int(max_seqlen_k.item())
+    
+    # SageAttention requires half precision (fp16/bf16)
+    out_dtype = q.dtype
+    half_dtypes = (torch.float16, torch.bfloat16)
+    
+    if not (q.dtype == k.dtype == v.dtype):
+        k = k.to(q.dtype)
+        v = v.to(q.dtype)
+    
+    if q.dtype not in half_dtypes:
+        q = q.to(torch.bfloat16)
+        k = k.to(torch.bfloat16)
+        v = v.to(torch.bfloat16)
+    
+    is_causal = kwargs.get('causal', False)
+    sm_scale = 1.0 / (q.shape[-1] ** 0.5)
+    
+    out = sageattn_varlen(
+        q, k, v,
+        cu_seqlens_q, cu_seqlens_k,
+        max_seqlen_q, max_seqlen_k,
+        is_causal, sm_scale
+    )
+    
+    return out.to(out_dtype) if out.dtype != out_dtype else out
+
+
+@torch._dynamo.disable
+def call_sage_attn_3_varlen(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs):
+    """
+    Wrapper for SageAttention 3 (Blackwell) that converts varlen format to batched format.
+    
+    SageAttention 3 / Blackwell provides maximum performance on RTX 50xx series GPUs.
+    However, it only supports batched attention (uniform sequence lengths), not varlen.
+    
+    This wrapper detects uniform-length batches and reshapes accordingly.
+    For variable-length sequences, it automatically falls back to SageAttention 2.
+    
+    This function is excluded from torch.compile because:
+    1. SageAttention is a C++ extension that can't be compiled anyway
+    2. It requires Python int scalars for max_seqlen parameters
+    3. The varlen-to-batched conversion involves dynamic shapes
+    4. Disabling compilation here keeps the rest of the model compilable
+    
+    Args:
+        q: Query tensor (total_seq, heads, head_dim)
+        k: Key tensor (total_seq, heads, head_dim)
+        v: Value tensor (total_seq, heads, head_dim)
+        cu_seqlens_q: Cumulative sequence lengths for queries
+        cu_seqlens_k: Cumulative sequence lengths for keys
+        max_seqlen_q: Maximum query sequence length (can be tensor or int)
+        max_seqlen_k: Maximum key sequence length (can be tensor or int)
+        **kwargs: Additional arguments (passed to SA2 fallback if needed)
+        
+    Returns:
+        Attention output tensor (total_seq, heads, head_dim)
+    """
+    if not SAGE_ATTN_3_AVAILABLE:
+        raise ImportError("SageAttention 3 (Blackwell) is not available")
+    
+    # Convert tensor max_seqlen to Python int if needed
+    if torch.is_tensor(max_seqlen_q):
+        max_seqlen_q = int(max_seqlen_q.item())
+    if torch.is_tensor(max_seqlen_k):
+        max_seqlen_k = int(max_seqlen_k.item())
+    
+    # Check if all sequences have uniform length (required for SA3 batched API)
+    # SA3/Blackwell uses batched attention, not varlen, so we need uniform lengths
+    seq_lens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
+    seq_lens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
+    
+    uniform_q = (seq_lens_q == seq_lens_q[0]).all()
+    uniform_k = (seq_lens_k == seq_lens_k[0]).all()
+    
+    if not (uniform_q and uniform_k):
+        # Fall back to SA2 for variable-length sequences
+        # This is expected behavior - SA3 Blackwell doesn't support varlen natively
+        if SAGE_ATTN_2_AVAILABLE:
+            return call_sage_attn_2_varlen(
+                q, k, v, cu_seqlens_q, cu_seqlens_k,
+                max_seqlen_q, max_seqlen_k, **kwargs
+            )
+        raise RuntimeError(
+            "SageAttention 3 (Blackwell) requires uniform sequence lengths, "
+            "and SageAttention 2 is not available as fallback. "
+            "Please install sageattention package or use flash_attn/sdpa instead."
+        )
+    
+    # Extract batch dimensions
+    batch_size = len(cu_seqlens_q) - 1
+    seq_len_q = int(seq_lens_q[0].item())
+    seq_len_k = int(seq_lens_k[0].item())
+    heads = q.shape[1]
+    dim = q.shape[2]
+    
+    # SageAttention requires half precision (fp16/bf16)
+    out_dtype = q.dtype
+    half_dtypes = (torch.float16, torch.bfloat16)
+    
+    if not (q.dtype == k.dtype == v.dtype):
+        k = k.to(q.dtype)
+        v = v.to(q.dtype)
+    
+    if q.dtype not in half_dtypes:
+        q = q.to(torch.bfloat16)
+        k = k.to(torch.bfloat16)
+        v = v.to(torch.bfloat16)
+    
+    # Reshape varlen (total_seq, heads, dim) -> batched (batch, seq, heads, dim)
+    q_batched = q.view(batch_size, seq_len_q, heads, dim)
+    k_batched = k.view(batch_size, seq_len_k, heads, dim)
+    v_batched = v.view(batch_size, seq_len_k, heads, dim)
+    
+    # SA3/Blackwell expects (batch, heads, seq, dim) layout
+    q_batched = q_batched.transpose(1, 2)  # (batch, heads, seq, dim)
+    k_batched = k_batched.transpose(1, 2)
+    v_batched = v_batched.transpose(1, 2)
+    
+    # Call SA3 Blackwell
+    out = sageattn_blackwell(q_batched, k_batched, v_batched, per_block_mean=False)
+    
+    # Reshape back to varlen format (total_seq, heads, dim)
+    out = out.transpose(1, 2).reshape(-1, heads, dim).contiguous()
+    
+    return out.to(out_dtype) if out.dtype != out_dtype else out
+
+
 # 2. Triton - Required for torch.compile with inductor backend
 try:
     import triton
@@ -97,15 +592,29 @@ def validate_gguf_availability(operation: str = "load GGUF model", debug=None) -
         raise RuntimeError(f"GGUF library required to {operation}")
 
 
-# 4. NVIDIA Conv3d Memory Bug - Workaround for PyTorch 2.9-2.10 + cuDNN >= 91002
+# 4. NVIDIA Conv3d Memory Bug - Workaround for PyTorch >= 2.9 + cuDNN >= 91002
 def _check_conv3d_memory_bug():
     """
     Check if Conv3d memory bug workaround needed.
-    Bug: PyTorch 2.9-2.10 with cuDNN >= 91002 uses 3x memory for Conv3d 
+    Bug: PyTorch 2.9+ with cuDNN >= 91002 uses 3x memory for Conv3d 
     with fp16/bfloat16 due to buggy dispatch layer.
     """
     try:
-        if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 3):
+        # Exclude AMD ROCm/HIP builds (they use MIOpen, not cuDNN)
+        if hasattr(torch.version, 'hip') and torch.version.hip is not None:
+            return False
+        
+        # Must have CUDA available
+        if not (hasattr(torch, 'cuda') and torch.cuda.is_available()):
+            return False
+        
+        # Must have cuDNN actually available (not just the attribute)
+        if not (hasattr(torch.backends.cudnn, 'is_available') and 
+                torch.backends.cudnn.is_available()):
+            return False
+        
+        # Check device capability (NVIDIA GPUs)
+        if torch.cuda.get_device_capability()[0] < 3:
             return False
         
         # Parse torch version
@@ -113,7 +622,8 @@ def _check_conv3d_memory_bug():
         parts = version_str.split('.')
         torch_version = tuple(int(p) for p in parts[:2])
         
-        if not ((2, 9) <= torch_version <= (2, 10)):
+        # Bug affects PyTorch 2.9 and later versions
+        if torch_version < (2, 9):
             return False
         
         if not hasattr(torch.backends.cudnn, 'version'):
@@ -134,21 +644,34 @@ NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND = _check_conv3d_memory_bug()
 if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"):
     os.environ["SEEDVR2_OPTIMIZATIONS_LOGGED"] = "1"
     
-    # Flash Attention & Triton status
+    # Build status strings
-    has_both = FLASH_ATTN_AVAILABLE and TRITON_AVAILABLE
+    sage_status = "✅" if SAGE_ATTN_AVAILABLE else "❌"
-    has_neither = not FLASH_ATTN_AVAILABLE and not TRITON_AVAILABLE
+    flash_status = "✅" if FLASH_ATTN_AVAILABLE else "❌"
+    triton_status = "✅" if TRITON_AVAILABLE else "❌"
     
-    if has_both:
+    # Count available optimizations
-        print("⚡ SeedVR2 optimizations check: Flash Attention ✅ | Triton ✅")
+    available = [SAGE_ATTN_AVAILABLE, FLASH_ATTN_AVAILABLE, TRITON_AVAILABLE]
-    elif has_neither:
+    num_available = sum(available)
-        print("⚠️  SeedVR2 optimizations check: Flash Attention ❌ | Triton ❌")
+    
-        print("💡 For best performance: pip install flash-attn triton")
+    if num_available == 3:
-    elif FLASH_ATTN_AVAILABLE:
+        print(f"⚡ SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
-        print("⚡ SeedVR2 optimizations check: Flash Attention ✅ | Triton ❌")
+    elif num_available == 0:
-        print("💡 Install Triton for torch.compile: pip install triton")
+        print(f"⚠️  SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
-    else:  # TRITON_AVAILABLE only
+        print("💡 For best performance: pip install sageattention flash-attn triton")
-        print("⚠️  SeedVR2 optimizations check: Flash Attention ❌ | Triton ✅")
+    else:
-        print("💡 Install Flash Attention for faster inference: pip install flash-attn")
+        icon = "⚡" if num_available >= 2 else "⚠️ "
+        print(f"{icon} SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}")
+        
+        # Build install suggestions for missing packages
+        missing = []
+        if not SAGE_ATTN_AVAILABLE:
+            missing.append("sageattention")
+        if not FLASH_ATTN_AVAILABLE:
+            missing.append("flash-attn")
+        if not TRITON_AVAILABLE:
+            missing.append("triton")
+        if missing:
+            print(f"💡 Optional: pip install {' '.join(missing)}")
     
     # Conv3d workaround status (if applicable)
     if NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND:
@@ -157,33 +680,55 @@ if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"):
         print(f"🔧 Conv3d workaround active: PyTorch {torch_ver}, cuDNN {cudnn_ver} (fixing VAE 3x memory bug)")
 
 
+# Bfloat16 CUBLAS support
+def _probe_bfloat16_support() -> bool:
+    if not torch.cuda.is_available():
+        return True
+    try:
+        a = torch.randn(8, 8, dtype=torch.bfloat16, device='cuda:0')
+        _ = torch.matmul(a, a)
+        del a
+        return True
+    except RuntimeError as e:
+        if "CUBLAS_STATUS_NOT_SUPPORTED" in str(e):
+            return False
+        raise
+
+BFLOAT16_SUPPORTED = _probe_bfloat16_support()
+COMPUTE_DTYPE = torch.bfloat16 if BFLOAT16_SUPPORTED else torch.float16
+
+
 def call_rope_with_stability(method, *args, **kwargs):
     """
     Call RoPE method with stability fixes:
     1. Clear cache if available
-    2. Disable autocast to prevent numerical issues
+    2. Disable autocast to prevent numerical issues (CUDA only)
     This prevents artifacts in FP8/mixed precision models.
     """
     if hasattr(method, 'cache_clear'):
         method.cache_clear()
     
-    with torch.cuda.amp.autocast(enabled=False):
+    # Only use CUDA autocast context on CUDA devices
+    # MPS has no CUDA autocast to disable
+    if torch.cuda.is_available():
+        with torch.cuda.amp.autocast(enabled=False):
+            return method(*args, **kwargs)
+    else:
         return method(*args, **kwargs)
     
     
-class FP8CompatibleDiT(torch.nn.Module):
+class CompatibleDiT(torch.nn.Module):
     """
     Wrapper for DiT models with automatic compatibility management + advanced optimizations
     
     Precision Handling:
     - FP8: Keeps native FP8 parameters (memory efficient), converts inputs/outputs to compute_dtype for arithmetic
-    - FP16: Uses native FP16 precision throughout
+    - FP16/BFloat16/Float32: Uses native precision throughout
-    - BFloat16: Uses native BFloat16 precision throughout
+    - GGUF: On-the-fly dequantization to compute_dtype
-    - Float32: Uses full precision for maximum quality
+    - MPS: Forces all parameters to compute_dtype (unified memory requires dtype consistency)
     - RoPE: Converted from FP8 to compute_dtype for numerical consistency
     
     Optimizations:
-    - Flash Attention: Automatic optimization of attention layers
     - RoPE Stabilization: Error handling for numerical stability in mixed precision
     - MPS Compatibility: Unified dtype conversion for Apple Silicon backends
     """
@@ -206,9 +751,12 @@ class FP8CompatibleDiT(torch.nn.Module):
             self.debug.start_timer("_convert_rope_freqs")
             self._convert_rope_freqs(target_dtype=self.compute_dtype)
             self.debug.end_timer("_convert_rope_freqs", "RoPE freqs conversion")
-            
+        
-            if torch.mps.is_available():
+        # MPS requires unified dtype for all parameters/buffers (no autocast fallback)
-                self.debug.log(f"Also converting NaDiT parameters/buffers for MPS backend", category="setup", force=True)
+        # Apply to ALL model types (FP8, FP16, GGUF) when dtype differs from compute_dtype
+        if not skip_conversion and hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+            if self.model_dtype != self.compute_dtype:
+                self.debug.log(f"Converting NaDiT parameters/buffers to {self.compute_dtype} for MPS backend", category="setup", force=True)
                 self.debug.start_timer("_force_nadit_precision")
                 self._force_nadit_precision(target_dtype=self.compute_dtype)
                 self.debug.end_timer("_force_nadit_precision", "NaDiT parameters/buffers conversion")
@@ -218,11 +766,6 @@ class FP8CompatibleDiT(torch.nn.Module):
         self.debug.start_timer("_stabilize_rope_computations")
         self._stabilize_rope_computations()
         self.debug.end_timer("_stabilize_rope_computations", "RoPE stabilization")
-
-        # 🚀 FLASH ATTENTION OPTIMIZATION (Phase 2)
-        self.debug.start_timer("_apply_flash_attention_optimization")
-        self._apply_flash_attention_optimization()
-        self.debug.end_timer("_apply_flash_attention_optimization", "Flash Attention application")
     
     def _detect_model_dtype(self) -> torch.dtype:
         """Detect main model dtype"""
@@ -284,8 +827,11 @@ class FP8CompatibleDiT(torch.nn.Module):
                     param.data = param.data.to(target_dtype)
                 converted_count += 1
                 
-        # Also convert buffers
+        # Also convert buffers (skip GGUF quantized buffers - they have tensor_type attribute)
         for name, buffer in self.dit_model.named_buffers():
+            # Skip GGUF quantized buffers - these must stay in packed format for on-the-fly dequantization
+            if hasattr(buffer, 'tensor_type'):
+                continue
             if buffer.dtype != target_dtype:
                 if buffer.device.type == "mps":
                     temp_cpu = buffer.data.to("cpu")
@@ -345,205 +891,6 @@ class FP8CompatibleDiT(torch.nn.Module):
         
         if rope_count > 0:
             self.debug.log(f"Stabilized {rope_count} RoPE modules", category="success")
-
-    def _apply_flash_attention_optimization(self) -> None:
-        """🚀 FLASH ATTENTION OPTIMIZATION - 30-50% speedup of attention layers"""
-        attention_layers_optimized = 0
-        flash_attention_available = self._check_flash_attention_support()
-        
-        for name, module in self.dit_model.named_modules():
-            # Identify all attention layers
-            if self._is_attention_layer(name, module):
-                # Apply optimization based on availability
-                if self._optimize_attention_layer(name, module, flash_attention_available):
-                    attention_layers_optimized += 1
-        
-        if not flash_attention_available:
-            self.debug.log("Flash Attention not available, using PyTorch SDPA as fallback", category="info", force=True)
-    
-    def _check_flash_attention_support(self) -> bool:
-        """Check if Flash Attention is available"""
-        # Check PyTorch SDPA (includes Flash Attention on H100/A100)
-        if hasattr(torch.nn.functional, 'scaled_dot_product_attention'):
-            return True
-        
-        # Check flash-attn package (uses module-level check from top of file)
-        return FLASH_ATTN_AVAILABLE
-    
-    def _is_attention_layer(self, name: str, module: torch.nn.Module) -> bool:
-        """Identify if a module is an attention layer"""
-        attention_keywords = [
-            'attention', 'attn', 'self_attn', 'cross_attn', 'mhattn', 'multihead',
-            'transformer_block', 'dit_block'
-        ]
-        
-        # Check by name
-        if any(keyword in name.lower() for keyword in attention_keywords):
-            return True
-        
-        # Check by module type
-        module_type = type(module).__name__.lower()
-        if any(keyword in module_type for keyword in attention_keywords):
-            return True
-        
-        # Check by attributes (modules with q, k, v projections)
-        if hasattr(module, 'q_proj') or hasattr(module, 'qkv') or hasattr(module, 'to_q'):
-            return True
-        
-        return False
-    
-    def _optimize_attention_layer(self, name: str, module: torch.nn.Module, flash_attention_available: bool) -> bool:
-        """Optimize a specific attention layer"""
-        try:
-            # Save original forward method
-            if not hasattr(module, '_original_forward'):
-                module._original_forward = module.forward
-            
-            # Create new optimized forward method
-            if flash_attention_available:
-                optimized_forward = self._create_flash_attention_forward(module, name)
-            else:
-                optimized_forward = self._create_sdpa_forward(module, name)
-            
-            # Replace forward method
-            module.forward = optimized_forward
-            return True
-            
-        except Exception as e:
-            self.debug.log(f"Failed to optimize attention layer '{name}': {e}", level="WARNING", category="dit", force=True)
-            return False
-    
-    def _create_flash_attention_forward(self, module: torch.nn.Module, layer_name: str):
-        """Create optimized forward with Flash Attention"""
-        original_forward = module._original_forward
-        
-        def flash_attention_forward(*args, **kwargs):
-            try:
-                # Try to use Flash Attention via SDPA
-                return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
-            except Exception as e:
-                # Fallback to original implementation
-                self.debug.log(f"Flash Attention failed for {layer_name}, using original: {e}", level="WARNING", category="dit", force=True)
-                return original_forward(*args, **kwargs)
-        
-        return flash_attention_forward
-    
-    def _create_sdpa_forward(self, module: torch.nn.Module, layer_name: str):
-        """Create optimized forward with PyTorch SDPA"""
-        original_forward = module._original_forward
-        
-        def sdpa_forward(*args, **kwargs):
-            try:
-                return self._sdpa_attention_forward(original_forward, module, *args, **kwargs)
-            except Exception as e:
-                # Fallback to original implementation
-                return original_forward(*args, **kwargs)
-        
-        return sdpa_forward
-    
-    def _sdpa_attention_forward(self, original_forward, module: torch.nn.Module, *args, **kwargs):
-        """Optimized forward pass using SDPA (Scaled Dot Product Attention)"""
-        # Detect if we can intercept and optimize this layer
-        if len(args) >= 1 and isinstance(args[0], torch.Tensor):
-            input_tensor = args[0]
-            
-            # Check dimensions to ensure it's standard attention
-            if len(input_tensor.shape) >= 3:  # [batch, seq_len, hidden_dim] or similar
-                try:
-                    return self._optimized_attention_computation(module, input_tensor, *args[1:], **kwargs)
-                except:
-                    pass
-        
-        # Fallback to original implementation
-        return original_forward(*args, **kwargs)
-    
-    def _optimized_attention_computation(self, module: torch.nn.Module, input_tensor: torch.Tensor, *args, **kwargs):
-        """Optimized attention computation with SDPA"""
-        # Try to detect standard attention format
-        batch_size, seq_len = input_tensor.shape[:2]
-        
-        # Check if module has standard Q, K, V projections
-        if hasattr(module, 'qkv') or (hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj')):
-            return self._compute_sdpa_attention(module, input_tensor, *args, **kwargs)
-        
-        # If no standard format detected, use original
-        return module._original_forward(input_tensor, *args, **kwargs)
-    
-    def _compute_sdpa_attention(self, module: torch.nn.Module, x: torch.Tensor, *args, **kwargs):
-        """Optimized SDPA computation for standard attention modules"""
-        try:
-            # Case 1: Module with combined QKV projection
-            if hasattr(module, 'qkv'):
-                qkv = module.qkv(x)
-                # Reshape to separate Q, K, V
-                batch_size, seq_len, _ = qkv.shape
-                qkv = qkv.reshape(batch_size, seq_len, 3, -1)
-                q, k, v = qkv.unbind(dim=2)
-                
-            # Case 2: Separate Q, K, V projections
-            elif hasattr(module, 'q_proj') and hasattr(module, 'k_proj') and hasattr(module, 'v_proj'):
-                q = module.q_proj(x)
-                k = module.k_proj(x)
-                v = module.v_proj(x)
-            else:
-                # Unsupported format, use original
-                return module._original_forward(x, *args, **kwargs)
-            
-            # Detect number of heads
-            head_dim = getattr(module, 'head_dim', None)
-            num_heads = getattr(module, 'num_heads', None)
-            
-            if head_dim is None or num_heads is None:
-                # Try to guess from dimensions
-                hidden_dim = q.shape[-1]
-                if hasattr(module, 'num_heads'):
-                    num_heads = module.num_heads
-                    head_dim = hidden_dim // num_heads
-                else:
-                    # Reasonable defaults
-                    head_dim = 64
-                    num_heads = hidden_dim // head_dim
-            
-            # Reshape for multi-head attention
-            batch_size, seq_len = q.shape[:2]
-            q = q.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
-            k = k.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
-            v = v.view(batch_size, seq_len, num_heads, head_dim).transpose(1, 2)
-            
-            if torch.mps.is_available():
-                attn_output = torch.nn.functional.scaled_dot_product_attention(
-                    q, k, v,
-                    dropout_p=0.0,
-                    is_causal=False
-                )
-            else:
-                # Use optimized SDPA
-                with torch.backends.cuda.sdp_kernel(
-                    enable_flash=True,
-                    enable_math=True,
-                    enable_mem_efficient=True
-                ):
-                    attn_output = torch.nn.functional.scaled_dot_product_attention(
-                        q, k, v,
-                        dropout_p=0.0,
-                        is_causal=False
-                    )
-            
-            # Reshape back
-            attn_output = attn_output.transpose(1, 2).contiguous().view(
-                batch_size, seq_len, num_heads * head_dim
-            )
-            
-            # Output projection if it exists
-            if hasattr(module, 'out_proj') or hasattr(module, 'o_proj'):
-                proj = getattr(module, 'out_proj', None) or getattr(module, 'o_proj', None)
-                attn_output = proj(attn_output)
-            
-            return attn_output
-            
-        except Exception as e:
-            # In case of error, use original implementation
-            return module._original_forward(x, *args, **kwargs)
     
     def forward(self, *args, **kwargs):
         """

src/optimization/memory_manager.py

@@ -10,8 +10,40 @@ import gc
 import sys
 import time
 import psutil
+import platform
 from typing import Tuple, Dict, Any, Optional, List, Union
+
+
+def _device_str(device: Union[torch.device, str]) -> str:
+    """Normalized uppercase device string for comparison and logging. MPS variants → 'MPS'."""
+    s = str(device).upper()
+    return 'MPS' if s.startswith('MPS') else s
+
+
+def is_mps_available() -> bool:
+    """Check if MPS (Apple Metal) backend is available."""
+    return hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()
+
+
+def is_cuda_available() -> bool:
+    """Check if CUDA backend is available."""
+    return torch.cuda.is_available()
+
+
+def get_gpu_backend() -> str:
+    """Get the active GPU backend type.
     
+    Returns:
+        'cuda': NVIDIA CUDA
+        'mps': Apple Metal Performance Shaders
+        'cpu': No GPU backend available
+    """
+    if is_cuda_available():
+        return 'cuda'
+    if is_mps_available():
+        return 'mps'
+    return 'cpu'
+
 
 def get_device_list(include_none: bool = False, include_cpu: bool = False) -> List[str]:
     """
@@ -31,15 +63,15 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
     has_mps = False
     
     try:
-        if hasattr(torch, "cuda") and hasattr(torch.cuda, "is_available") and torch.cuda.is_available():
+        if is_cuda_available():
             devs += [f"cuda:{i}" for i in range(torch.cuda.device_count())]
             has_cuda = True
     except Exception:
         pass
     
     try:
-        if hasattr(torch, "mps") and hasattr(torch.mps, "is_available") and torch.mps.is_available():
+        if is_mps_available():
-            devs += [f"mps:{i}" for i in range(torch.mps.device_count())]
+            devs.append("mps")  # MPS doesn't use device indices
             has_mps = True
     except Exception:
         pass
@@ -60,7 +92,7 @@ def get_device_list(include_none: bool = False, include_cpu: bool = False) -> Li
     result.extend(devs)
     
     return result if result else []
-    
+
 
 def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]:
     """
@@ -74,13 +106,13 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
         dict: {"free_gb": float, "total_gb": float} or {"error": str}
     """
     try:
-        if torch.cuda.is_available():
+        if is_cuda_available():
             if device is None:
                 device = torch.device("cuda:0")
             elif not isinstance(device, torch.device):
                 device = torch.device(device)
             free_memory, total_memory = torch.cuda.mem_get_info(device)
-        elif torch.mps.is_available():
+        elif is_mps_available():
             # MPS doesn't support per-device queries or mem_get_info
             # Use system memory as proxy
             mem = psutil.virtual_memory()
@@ -100,13 +132,13 @@ def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]
 # Initial VRAM check at module load
 vram_info = get_basic_vram_info(device=None)
 if "error" not in vram_info:
-    backend = "MPS" if torch.mps.is_available() else "CUDA"
+    backend = "MPS" if is_mps_available() else "CUDA"
     print(f"📊 Initial {backend} memory: {vram_info['free_gb']:.2f}GB free / {vram_info['total_gb']:.2f}GB total")
 else:
     print(f"⚠️ Memory check failed: {vram_info['error']} - No available backend!")
 
 
-def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> Tuple[float, float, float]:
+def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]:
     """
     Get current VRAM usage metrics for monitoring.
     Used for tracking memory consumption during processing.
@@ -116,29 +148,30 @@ def get_vram_usage(device: Optional[torch.device] = None, debug: Optional['Debug
         debug: Optional debug instance for logging
     
     Returns:
-        tuple: (allocated_gb, reserved_gb, max_allocated_gb)
+        tuple: (allocated_gb, reserved_gb, peak_allocated_gb, peak_reserved_gb)
-               Returns (0, 0, 0) if no GPU available
+               Returns (0, 0, 0, 0) if no GPU available
     """
     try:
-        if torch.cuda.is_available():
+        if is_cuda_available():
             if device is None:
                 device = torch.device("cuda:0")
             elif not isinstance(device, torch.device):
                 device = torch.device(device)
             allocated = torch.cuda.memory_allocated(device) / (1024**3)
             reserved = torch.cuda.memory_reserved(device) / (1024**3)
-            max_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
+            peak_allocated = torch.cuda.max_memory_allocated(device) / (1024**3)
-            return allocated, reserved, max_allocated
+            peak_reserved = torch.cuda.max_memory_reserved(device) / (1024**3)
-        elif torch.mps.is_available():
+            return allocated, reserved, peak_allocated, peak_reserved
+        elif is_mps_available():
             # MPS doesn't support per-device queries - uses global memory tracking
             allocated = torch.mps.current_allocated_memory() / (1024**3)
             reserved = torch.mps.driver_allocated_memory() / (1024**3)
-            max_allocated = allocated  # MPS doesn't track peak separately
+            # MPS doesn't track peak separately
-            return allocated, reserved, max_allocated
+            return allocated, reserved, allocated, reserved
     except Exception as e:
         if debug:
             debug.log(f"Failed to get VRAM usage: {e}", level="WARNING", category="memory", force=True)
-    return 0.0, 0.0, 0.0
+    return 0.0, 0.0, 0.0, 0.0
 
 
 def get_ram_usage(debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]:
@@ -229,17 +262,17 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
         # Use existing function for memory info
         mem_info = get_basic_vram_info(device=None)
         
-        if "error" not in mem_info:
+        if "error" not in mem_info and mem_info["total_gb"] > 0:
             # Check VRAM/MPS memory pressure (5% free threshold)
             free_ratio = mem_info["free_gb"] / mem_info["total_gb"]
             if free_ratio < 0.05:
                 should_clear = True
                 if debug:
-                    backend = "MPS" if torch.mps.is_available() else "VRAM"
+                    backend = "Unified Memory" if is_mps_available() else "VRAM"
                     debug.log(f"{backend} pressure: {mem_info['free_gb']:.2f}GB free of {mem_info['total_gb']:.2f}GB", category="memory")
         
         # For non-MPS systems, also check system RAM separately
-        if not should_clear and not torch.mps.is_available():
+        if not should_clear and not is_mps_available():
             mem = psutil.virtual_memory()
             if mem.available < mem.total * 0.05:
                 should_clear = True
@@ -262,10 +295,10 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
     if debug:
         debug.start_timer(gpu_timer)
     
-    if torch.cuda.is_available():
+    if is_cuda_available():
         torch.cuda.empty_cache()
         torch.cuda.ipc_collect()
-    elif torch.mps.is_available():
+    elif is_mps_available():
         torch.mps.empty_cache()
     
     if debug:
@@ -302,7 +335,7 @@ def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: boo
                 handle = _os_memory_lib.GetCurrentProcess()
                 _os_memory_lib.SetProcessWorkingSetSize(handle, -1, -1)
                 
-            elif torch.mps.is_available():
+            elif is_mps_available():
                 # macOS with MPS
                 import ctypes  # Import only when needed
                 import ctypes.util
@@ -379,7 +412,7 @@ def reset_vram_peak(device: Optional[torch.device] = None, debug: Optional['Debu
     if debug and debug.enabled:
         debug.log("Resetting VRAM peak memory statistics", category="memory")
     try:
-        if torch.cuda.is_available():
+        if is_cuda_available():
             if device is None:
                 device = torch.device("cuda:0")
             elif not isinstance(device, torch.device):
@@ -423,12 +456,11 @@ def clear_rope_lru_caches(model: Optional[torch.nn.Module], debug: Optional['Deb
 
 
 def release_tensor_memory(tensor: Optional[torch.Tensor]) -> None:
-    """Release tensor memory properly without CPU allocation"""
+    """Release tensor memory from any device (CPU/CUDA/MPS)"""
     if tensor is not None and torch.is_tensor(tensor):
-        if tensor.is_cuda or tensor.is_mps:
+        # Release storage for all devices (CPU, CUDA, MPS)
-            # Release GPU memory directly without CPU transfer
+        if tensor.numel() > 0:
-            if tensor.numel() > 0:
+            tensor.data.set_()
-                tensor.data.set_()
         tensor.grad = None
 
 
@@ -592,7 +624,7 @@ def manage_tensor(
     target_dtype = dtype if dtype is not None else current_dtype
     
     # Check if movement is actually needed
-    needs_device_move = current_device != target_device
+    needs_device_move = _device_str(current_device) != _device_str(target_device)
     needs_dtype_change = dtype is not None and current_dtype != target_dtype
     
     if not needs_device_move and not needs_dtype_change:
@@ -610,8 +642,8 @@ def manage_tensor(
     
     # Log the movement
     if debug:
-        current_device_str = str(current_device).upper()
+        current_device_str = _device_str(current_device)
-        target_device_str = str(target_device).upper()
+        target_device_str = _device_str(target_device)
         
         dtype_info = ""
         if needs_dtype_change:
@@ -670,7 +702,7 @@ def manage_model_device(model: torch.nn.Module, target_device: torch.device, mod
         
         if has_blockswap_config:
             is_blockswap_model = True
-            # Get the actual model (handle FP8CompatibleDiT wrapper)
+            # Get the actual model (handle CompatibleDiT wrapper)
             if hasattr(model, "dit_model"):
                 actual_model = model.dit_model
 
@@ -682,8 +714,8 @@ def manage_model_device(model: torch.nn.Module, target_device: torch.device, mod
     
     # Extract device type for comparison (both are torch.device objects)
     target_type = target_device.type
-    current_device_upper = str(current_device).upper()
+    current_device_upper = _device_str(current_device)
-    target_device_upper = str(target_device).upper()
+    target_device_upper = _device_str(target_device)
 
     # Compare normalized device types
     if current_device_upper == target_device_upper and not is_blockswap_model:
@@ -738,10 +770,10 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
                 actual_source_device = param.device
                 break
         
-        source_device_desc = str(actual_source_device).upper() if actual_source_device else str(target_device).upper()
+        source_device_desc = _device_str(actual_source_device) if actual_source_device else _device_str(target_device)
         
         if debug:
-            debug.log(f"Moving {model_name} from {source_device_desc} to {str(target_device).upper()} ({reason or 'model caching'})", category="general")
+            debug.log(f"Moving {model_name} from {source_device_desc} to {_device_str(target_device)} ({reason or 'model caching'})", category="general")
         
         # Enable bypass to allow movement
         set_blockswap_bypass(runner=runner, bypass=True, debug=debug)
@@ -756,14 +788,14 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
         model.zero_grad(set_to_none=True)
         
         if debug:
-            debug.end_timer(timer_name, f"BlockSwap model offloaded to {str(target_device).upper()}")
+            debug.end_timer(timer_name, f"BlockSwap model offloaded to {_device_str(target_device)}")
         
         return True
         
     else:
         # Moving to GPU (reload)
         # Check if we're in bypass mode (coming from offload)
-        if not getattr(runner, "_blockswap_bypass_protection", False):
+        if not getattr(model, "_blockswap_bypass_protection", False):
             # Not in bypass mode, blocks are already configured
             if debug:
                 debug.log(f"{model_name} with BlockSwap active - blocks already distributed across devices, skipping movement", category="general")
@@ -776,10 +808,10 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
                 actual_current_device = param.device
                 break
         
-        current_device_desc = str(actual_current_device).upper() if actual_current_device else "OFFLOAD"
+        current_device_desc = _device_str(actual_current_device) if actual_current_device else "OFFLOAD"
         
         if debug:
-            debug.log(f"Moving {model_name} from {current_device_desc} to {str(target_device).upper()} ({reason or 'inference requirement'})", category="general")
+            debug.log(f"Moving {model_name} from {current_device_desc} to {_device_str(target_device)} ({reason or 'inference requirement'})", category="general")
         
         timer_name = f"{model_name.lower()}_to_gpu"
         if debug:
@@ -788,7 +820,7 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
         # Restore blocks to their configured devices
         if hasattr(model, "blocks") and hasattr(model, "blocks_to_swap"):
             # Use configured offload_device from BlockSwap config
-            offload_device = runner._block_swap_config.get("offload_device")
+            offload_device = model._block_swap_config.get("offload_device")
             if not offload_device:
                 raise ValueError("BlockSwap config missing offload_device")
             
@@ -802,7 +834,7 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
                     block.to(offload_device)
             
             # Handle I/O components
-            if not runner._block_swap_config.get("swap_io_components", False):
+            if not model._block_swap_config.get("swap_io_components", False):
                 # I/O components should be on GPU if not offloaded
                 for name, module in model.named_children():
                     if name != "blocks":
@@ -815,11 +847,11 @@ def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module,
             
             if debug:
                 # Get actual configuration from runner
-                if hasattr(runner, '_block_swap_config'):
+                if hasattr(model, '_block_swap_config'):
-                    blocks_on_gpu = runner._block_swap_config.get('total_blocks', 32) - runner._block_swap_config.get('blocks_swapped', 16)
+                    blocks_on_gpu = model._block_swap_config.get('total_blocks', 32) - model._block_swap_config.get('blocks_swapped', 16)
-                    total_blocks = runner._block_swap_config.get('total_blocks', 32)
+                    total_blocks = model._block_swap_config.get('total_blocks', 32)
-                    main_device = runner._block_swap_config.get('main_device', 'GPU')
+                    main_device = model._block_swap_config.get('main_device', 'GPU')
-                    debug.log(f"BlockSwap blocks restored to configured devices ({blocks_on_gpu}/{total_blocks} blocks on {str(main_device).upper()})", category="success")
+                    debug.log(f"BlockSwap blocks restored to configured devices ({blocks_on_gpu}/{total_blocks} blocks on {_device_str(main_device)})", category="success")
                 else:
                     debug.log("BlockSwap blocks restored to configured devices", category="success")
 
@@ -866,8 +898,8 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic
     
     # Log the movement with full device strings
     if debug:
-        current_device_str = str(current_device).upper()
+        current_device_str = _device_str(current_device)
-        target_device_str = str(target_device).upper()
+        target_device_str = _device_str(target_device)
         debug.log(f"Moving {model_name} from {current_device_str} to {target_device_str} ({reason})", category="general")
 
     # Start timer based on direction
@@ -892,7 +924,7 @@ def _standard_model_movement(model: torch.nn.Module, current_device: torch.devic
     
     # End timer
     if debug:
-        debug.end_timer(timer_name, f"{model_name} moved to {str(target_device).upper()}")
+        debug.end_timer(timer_name, f"{model_name} moved to {_device_str(target_device)}")
     
     return True
 
@@ -1018,15 +1050,17 @@ def cleanup_dit(runner: Any, debug: Optional['Debug'] = None, cache_model: bool
         
         # Move model off GPU if needed
         if param_device.type not in ['meta', 'cpu']:
-            # Get offload target - default to 'cpu' if not configured or set to 'none'
+            # MPS: skip CPU movement before deletion (unified memory, just causes sync)
-            offload_target = getattr(runner, '_dit_offload_device', None)
+            if param_device.type == 'mps' and not cache_model:
-            if offload_target is None or offload_target == 'none':
+                if debug:
-                offload_target = torch.device('cpu')
+                    debug.log("DiT on MPS - skipping CPU movement before deletion", category="cleanup")
-            
+            else:
-            # Move model off GPU (either for caching or before deletion)
+                offload_target = getattr(runner, '_dit_offload_device', None)
-            reason = "model caching" if cache_model else "releasing GPU memory"
+                if offload_target is None or offload_target == 'none':
-            manage_model_device(model=runner.dit, target_device=offload_target, model_name="DiT", 
+                    offload_target = torch.device('cpu')
-                               debug=debug, reason=reason, runner=runner)
+                reason = "model caching" if cache_model else "releasing GPU memory"
+                manage_model_device(model=runner.dit, target_device=offload_target, model_name="DiT", 
+                                   debug=debug, reason=reason, runner=runner)
         elif param_device.type == 'meta' and debug:
             debug.log("DiT on meta device - keeping structure for cache", category="cleanup")
     except StopIteration:
@@ -1094,15 +1128,17 @@ def cleanup_vae(runner: Any, debug: Optional['Debug'] = None, cache_model: bool
         
         # Move model off GPU if needed
         if param_device.type not in ['meta', 'cpu']:
-            # Get offload target - default to 'cpu' if not configured or set to 'none'
+            # MPS: skip CPU movement before deletion (unified memory, just causes sync)
-            offload_target = getattr(runner, '_vae_offload_device', None)
+            if param_device.type == 'mps' and not cache_model:
-            if offload_target is None or offload_target == 'none':
+                if debug:
-                offload_target = torch.device('cpu')
+                    debug.log("VAE on MPS - skipping CPU movement before deletion", category="cleanup")
-            
+            else:
-            # Move model off GPU (either for caching or before deletion)
+                offload_target = getattr(runner, '_vae_offload_device', None)
-            reason = "model caching" if cache_model else "releasing GPU memory"
+                if offload_target is None or offload_target == 'none':
-            manage_model_device(model=runner.vae, target_device=offload_target, model_name="VAE", 
+                    offload_target = torch.device('cpu')
-                               debug=debug, reason=reason, runner=runner)
+                reason = "model caching" if cache_model else "releasing GPU memory"
+                manage_model_device(model=runner.vae, target_device=offload_target, model_name="VAE", 
+                                   debug=debug, reason=reason, runner=runner)
         elif param_device.type == 'meta' and debug:
             debug.log("VAE on meta device - keeping structure for cache", category="cleanup")
     except StopIteration:

src/optimization/performance.py

@@ -9,6 +9,40 @@ import torch
 from typing import List
 
 
+def optimized_channels_to_last(tensor):
+    """🚀 Optimized replacement for rearrange(tensor, 'b c ... -> b ... c')
+    Moves channels from position 1 to last position using PyTorch native operations.
+    """
+    if tensor.ndim == 3:  # [batch, channels, spatial]
+        return tensor.permute(0, 2, 1)
+    elif tensor.ndim == 4:  # [batch, channels, height, width]
+        return tensor.permute(0, 2, 3, 1)
+    elif tensor.ndim == 5:  # [batch, channels, depth, height, width]
+        return tensor.permute(0, 2, 3, 4, 1)
+    else:
+        # Fallback for other dimensions - move channel (dim=1) to last
+        dims = list(range(tensor.ndim))
+        dims = [dims[0]] + dims[2:] + [dims[1]]  # [0, 2, 3, ..., 1]
+        return tensor.permute(*dims)
+
+
+def optimized_channels_to_second(tensor):
+    """🚀 Optimized replacement for rearrange(tensor, 'b ... c -> b c ...')
+    Moves channels from last position to position 1 using PyTorch native operations.
+    """
+    if tensor.ndim == 3:  # [batch, spatial, channels]
+        return tensor.permute(0, 2, 1)
+    elif tensor.ndim == 4:  # [batch, height, width, channels]
+        return tensor.permute(0, 3, 1, 2)
+    elif tensor.ndim == 5:  # [batch, depth, height, width, channels]
+        return tensor.permute(0, 4, 1, 2, 3)
+    else:
+        # Fallback for other dimensions - move last dim to position 1
+        dims = list(range(tensor.ndim))
+        dims = [dims[0], dims[-1]] + dims[1:-1]  # [0, -1, 1, 2, ..., -2]
+        return tensor.permute(*dims)
+
+
 def optimized_video_rearrange(video_tensors: List[torch.Tensor]) -> List[torch.Tensor]:
     """
     🚀 OPTIMIZED version of video rearrangement
@@ -132,34 +166,4 @@ def optimized_sample_to_image_format(sample: torch.Tensor) -> torch.Tensor:
         return sample.permute(0, 2, 3, 1)
 
 
-def temporal_latent_blending(latents1: torch.Tensor, latents2: torch.Tensor, blend_frames: int) -> torch.Tensor:
-    """
-    🎨 Temporal blending in latent space to avoid discontinuities
-    
-    Args:
-        latents1: Latents from previous batch (end frames)
-        latents2: Latents from current batch (start frames)
-        blend_frames: Number of frames to blend
-    
-    Returns:
-        Blended latents for smooth transition
-    """
-    if latents1.shape[0] != latents2.shape[0]:
-        # Adjust dimensions if necessary
-        min_frames = min(latents1.shape[0], latents2.shape[0])
-        latents1 = latents1[:min_frames]
-        latents2 = latents2[:min_frames]
-    
-    # Create linear blending weights
-    # Frame 0: 100% latents1, 0% latents2
-    # Frame n: 0% latents1, 100% latents2
-    weights1 = torch.linspace(1.0, 0.0, blend_frames).view(-1, 1, 1, 1).to(latents1.device)
-    weights2 = torch.linspace(0.0, 1.0, blend_frames).view(-1, 1, 1, 1).to(latents2.device)
-    
-    # Apply blending
-    blended_latents = weights1 * latents1 + weights2 * latents2
-    
-    return blended_latents
-
-
 

src/utils/color_fix.py

@@ -84,9 +84,9 @@ def calc_mean_std(feat: Tensor, eps: float = 1e-5) -> tuple[Tensor, Tensor]:
     assert len(size) == 4, 'The input feature should be 4D tensor.'
     b, c = size[:2]
     
-    feat_var = feat.view(b, c, -1).var(dim=2) + eps
+    feat_var = feat.reshape(b, c, -1).var(dim=2) + eps
-    feat_std = feat_var.sqrt().view(b, c, 1, 1)
+    feat_std = feat_var.sqrt().reshape(b, c, 1, 1)
-    feat_mean = feat.view(b, c, -1).mean(dim=2).view(b, c, 1, 1)
+    feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1)
     
     return feat_mean, feat_std
 
@@ -178,7 +178,7 @@ def wavelet_decomposition(image: Tensor, levels: int = 5) -> tuple[Tensor, Tenso
     for i in range(levels):
         radius = 2 ** i
         low_freq = wavelet_blur(image, radius)
-        high_freq += (image - low_freq)
+        high_freq.add_(image).sub_(low_freq)
         image = low_freq
     
     return high_freq, low_freq
@@ -237,13 +237,13 @@ def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor, debug: Opti
             align_corners=False
         )
     
-    # Reconstruct: content details + style color
+    # Reconstruct: content details + style color (in-place on content_high_freq)
-    result = content_high_freq + style_low_freq
+    content_high_freq.add_(style_low_freq)
     
-    # Safety clamp for normalized SDR range
+    # Safety clamp for normalized SDR range (in-place)
     # This prevents numerical errors from propagating
     # Note: For HDR support, this would need to be removed
-    return torch.clamp(result, -1.0, 1.0)
+    return content_high_freq.clamp_(-1.0, 1.0)
 
 
 def lab_color_transfer(
@@ -255,17 +255,15 @@ def lab_color_transfer(
     """
     Perceptually-accurate color transfer using CIELAB color space.
     
-    LAB provides superior perceptual uniformity compared to RGB/HSV, enabling
+    Combines wavelet reconstruction (for spatial continuity) with LAB histogram
-    highly accurate color matching that preserves the original image's appearance.
+    matching (for precise color matching). This eliminates tile artifacts while
-    This is the RECOMMENDED method for color correction.
+    providing superior color accuracy.
     
     Algorithm:
-    1. Convert both images to LAB color space (D65 illuminant)
+    1. Apply wavelet reconstruction to get artifact-free base
-    2. Apply histogram matching to all LAB channels:
+    2. Convert both wavelet result and style to LAB color space
-       - L* (luminance): Weighted blend to preserve detail
+    3. Apply histogram matching to LAB channels
-       - a* (green-red): Full histogram matching
+    4. Convert back to RGB
-       - b* (blue-yellow): Full histogram matching
-    3. Convert back to RGB
     
     Args:
         content_feat: Target tensor [B, C, H, W] in [-1, 1] with upscaled details
@@ -278,7 +276,10 @@ def lab_color_transfer(
     Returns:
         Color-corrected tensor [B, C, H, W] in [-1, 1]
     """
-    # Handle spatial dimension mismatch
+    # Step 1: Apply wavelet to get artifact-free base with correct spatial structure
+    content_feat = wavelet_reconstruction(content_feat, style_feat, debug=None)
+    
+    # Handle spatial dimension mismatch (should already match after wavelet)
     if content_feat.shape != style_feat.shape:
         debug.log(
             f"LAB: Resizing style {style_feat.shape} to match content {content_feat.shape}",
@@ -291,14 +292,14 @@ def lab_color_transfer(
             align_corners=False
         )
     
-    # Store device
+    # Store device and convert to float32
     device = content_feat.device
     
     # Convert to float32 for accurate color space conversion
     content_feat, original_dtype = ensure_float32_precision(content_feat)
     style_feat, _ = ensure_float32_precision(style_feat)
     
-    # Precompute color space conversion matrices (once per batch, not per pixel)
+    # Precompute color space conversion matrices
     rgb_to_xyz_matrix = torch.tensor([
         [0.4124564, 0.3575761, 0.1804375],
         [0.2126729, 0.7151522, 0.0721750],
@@ -315,17 +316,16 @@ def lab_color_transfer(
     epsilon = 6.0 / 29.0
     kappa = (29.0 / 3.0) ** 3
     
-    # Convert from [-1, 1] to [0, 1] range (in-place where possible)
+    # Convert from [-1, 1] to [0, 1] range (in-place)
-    content_rgb = content_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0)
+    content_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
-    style_rgb = style_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0)
+    style_feat.add_(1.0).mul_(0.5).clamp_(0.0, 1.0)
-    del content_feat, style_feat
     
     # Convert to LAB color space
-    content_lab = _rgb_to_lab_batch(content_rgb, device, rgb_to_xyz_matrix, epsilon, kappa)
+    content_lab = _rgb_to_lab_batch(content_feat, device, rgb_to_xyz_matrix, epsilon, kappa)
-    del content_rgb
+    del content_feat
     
-    style_lab = _rgb_to_lab_batch(style_rgb, device, rgb_to_xyz_matrix, epsilon, kappa)
+    style_lab = _rgb_to_lab_batch(style_feat, device, rgb_to_xyz_matrix, epsilon, kappa)
-    del style_rgb, rgb_to_xyz_matrix
+    del style_feat, rgb_to_xyz_matrix
     
     # Match chrominance channels (a*, b*) for accurate color transfer
     matched_a = _histogram_matching_channel(content_lab[:, 1], style_lab[:, 1], device)
@@ -381,6 +381,8 @@ def _rgb_to_lab_batch(rgb: Tensor, device: torch.device, matrix: Tensor, epsilon
     rgb_flat = rgb_linear.permute(0, 2, 3, 1).reshape(-1, 3)
     del rgb_linear
     
+    # Ensure dtype consistency for matrix multiplication
+    rgb_flat = rgb_flat.to(dtype=matrix.dtype)
     xyz_flat = torch.matmul(rgb_flat, matrix.T)
     del rgb_flat
     
@@ -452,6 +454,8 @@ def _lab_to_rgb_batch(lab: Tensor, device: torch.device, matrix_inv: Tensor, eps
     xyz_flat = xyz.permute(0, 2, 3, 1).reshape(-1, 3)
     del xyz
     
+    # Ensure dtype consistency for matrix multiplication
+    xyz_flat = xyz_flat.to(dtype=matrix_inv.dtype)
     rgb_linear_flat = torch.matmul(xyz_flat, matrix_inv.T)
     del xyz_flat
     
@@ -490,6 +494,7 @@ def _histogram_matching_channel(source: Tensor, reference: Tensor, device: torch
     # Sort both arrays
     source_sorted, source_indices = torch.sort(source_flat)
     reference_sorted, _ = torch.sort(reference_flat)
+    del reference_flat
     
     # Quantile mapping
     n_source = len(source_sorted)
@@ -503,12 +508,15 @@ def _histogram_matching_channel(source: Tensor, reference: Tensor, device: torch
         ref_indices = (source_quantiles * (n_reference - 1)).long()
         ref_indices.clamp_(0, n_reference - 1)
         matched_sorted = reference_sorted[ref_indices]
-        del source_quantiles, ref_indices
+        del source_quantiles, ref_indices, reference_sorted
     
-    # Reconstruct with matched values
+    del source_sorted, source_flat
-    matched_flat = torch.empty_like(source_flat)
+    
-    matched_flat.scatter_(0, source_indices, matched_sorted)
+    # Reconstruct using argsort (portable across CUDA/ROCm/MPS)
-    del source_flat, reference_flat, source_sorted, source_indices, reference_sorted, matched_sorted
+    inverse_indices = torch.argsort(source_indices)
+    del source_indices
+    matched_flat = matched_sorted[inverse_indices]
+    del matched_sorted, inverse_indices
     
     return matched_flat.reshape(original_shape)
 
@@ -748,11 +756,15 @@ def _histogram_match_1d(source: Tensor, reference: Tensor, device: torch.device)
         ref_indices = (source_quantiles * (n_reference - 1)).long()
         ref_indices.clamp_(0, n_reference - 1)
         matched_sorted = reference_sorted[ref_indices]
-        del source_quantiles, ref_indices
+        del source_quantiles, ref_indices, reference_sorted
     
-    matched = torch.empty_like(source)
+    del source_sorted
-    matched.scatter_(0, source_indices, matched_sorted)
+    
-    del source_sorted, source_indices, reference_sorted, matched_sorted
+    # Reconstruct using argsort (portable across CUDA/ROCm/MPS)
+    inverse_indices = torch.argsort(source_indices)
+    del source_indices
+    matched = matched_sorted[inverse_indices]
+    del matched_sorted, inverse_indices
     
     return matched
 

src/utils/constants.py

@@ -3,6 +3,9 @@ Shared constants and utilities for SeedVR2
 Only includes constants actually used in the codebase
 """
 
+# Version information
+__version__ = "2.5.24"
+
 import os
 import warnings
 import inspect
@@ -52,14 +55,33 @@ def get_base_cache_dir() -> str:
 
 
 def get_all_model_paths() -> list:
-    """Get all registered model paths including those from extra_model_paths.yaml"""
+    """Get all registered model paths including those from extra_model_paths.yaml (case-insensitive)"""
     try:
         import folder_paths
         # Ensure default path is registered first
         get_base_cache_dir()
-        # Get all paths registered for seedvr2 model type
+        
-        paths = folder_paths.get_folder_paths(SEEDVR2_MODEL_TYPE)
+        # Case-insensitive lookup: search through all registered folder types
-        return paths if paths else [get_base_cache_dir()]
+        # This handles any case variation users might use in extra_model_paths.yaml
+        all_paths = []
+        target_lower = SEEDVR2_MODEL_TYPE.lower()
+        
+        # folder_paths.folder_names_and_paths is the underlying dict: {type: ([paths], extensions)}
+        if hasattr(folder_paths, 'folder_names_and_paths'):
+            for folder_type, (paths, _) in folder_paths.folder_names_and_paths.items():
+                if folder_type.lower() == target_lower:
+                    all_paths.extend(paths)
+        
+        # Remove duplicates while preserving order (os.path.normpath handles Windows/Linux path differences)
+        seen = set()
+        unique_paths = []
+        for path in all_paths:
+            normalized = os.path.normpath(path.lower())
+            if normalized not in seen:
+                seen.add(normalized)
+                unique_paths.append(path)
+        
+        return unique_paths if unique_paths else [get_base_cache_dir()]
     except:
         return [get_base_cache_dir()]
 

src/utils/debug.py

@@ -10,7 +10,33 @@ import torch
 import gc
 from typing import Optional, List, Dict, Any, Union
 from datetime import datetime
-from ..optimization.memory_manager import get_vram_usage, get_basic_vram_info, get_ram_usage, reset_vram_peak
+import platform
+from ..optimization.memory_manager import (
+    get_vram_usage, 
+    get_basic_vram_info, 
+    get_ram_usage, 
+    reset_vram_peak,
+    is_mps_available,
+    is_cuda_available
+)
+from ..utils.constants import __version__
+
+
+def _format_peak_with_overflow(peak_gb: float, total_vram_gb: float) -> str:
+    """Format peak reserved memory, showing overflow breakdown on Windows.
+    
+    Args:
+        peak_gb: Peak reserved memory from PyTorch
+        total_vram_gb: Physical GPU VRAM capacity
+    """
+    if total_vram_gb <= 0:
+        return f"{peak_gb:.2f}GB reserved"
+    
+    overflow_gb = peak_gb - total_vram_gb
+    if overflow_gb <= 0 or platform.system() != 'Windows':
+        return f"{peak_gb:.2f}GB reserved"
+    
+    return f"{peak_gb:.2f}GB reserved ({total_vram_gb:.0f}GB GPU + {overflow_gb:.2f}GB overflow)"
 
 
 class Debug:
@@ -52,7 +78,7 @@ class Debug:
         "device": "🖥️",       # Device info
         "file": "📂",         # File operations
         "alpha": "👻",        # Alpha operations
-        "star": "⭐",         # Star
+        "starlove": "⭐💝",   # Star + love
         "dialogue": "💬",     # Dialogue
         "none" : "",
     }
@@ -71,7 +97,9 @@ class Debug:
         self.vram_history: List[float] = []
         self.active_timer_stack: List[str] = [] 
         self.timer_namespace: str = ""
-        self.phase_peaks: Dict[str, float] = {}
+        self.phase_vram_peaks_alloc: Dict[str, float] = {}
+        self.phase_vram_peaks_rsv: Dict[str, float] = {}
+        self.phase_ram_peaks: Dict[str, float] = {}
         
     @torch._dynamo.disable  # Skip tracing to avoid datetime.now() warnings
     def log(self, message: str, level: str = "INFO", category: str = "general", force: bool = False, indent_level: int = 0) -> None:
@@ -111,32 +139,129 @@ class Debug:
         # Add indentation
         indent = " " * (indent_level * 2)
         
-        print(f"{prefix} {indent}{message}")
+        print(f"{prefix} {indent}{message}", flush=True)
 
     def print_header(self, cli: bool = False) -> None:
         """Print the header with banner - always displayed"""
-        # Intro logo
+        # Temporarily disable timestamps for clean header display
+        original_timestamps = self.show_timestamps
+        self.show_timestamps = False
+        
+        # ASCII art logo
         self.log("", category="none", force=True)
-        self.log(" ╔══════════════════════════════════════════════════════════╗", category="none", force=True)
+        self.log("", category="none", force=True)
-        self.log(" ║ ███████ ███████ ███████ ██████  ██    ██ ██████  ███████ ║", category="none", force=True)
+        self.log("███████╗███████╗███████╗██████╗ ██╗   ██╗██████╗     ██████╗       ███████╗", category="none", force=True, indent_level=1)
-        self.log(" ║ ██      ██      ██      ██   ██ ██    ██ ██   ██      ██ ║", category="none", force=True)
+        self.log("██╔════╝██╔════╝██╔════╝██╔══██╗██║   ██║██╔══██╗    ╚════██╗      ██╔════╝", category="none", force=True, indent_level=1)
-        self.log(" ║ ███████ █████   █████   ██   ██ ██    ██ ██████  █████   ║", category="none", force=True)
+        self.log("███████╗█████╗  █████╗  ██║  ██║██║   ██║██████╔╝     █████╔╝      ███████╗", category="none", force=True, indent_level=1)
-        self.log(" ║      ██ ██      ██      ██   ██  ██  ██  ██   ██ ██      ║", category="none", force=True)
+        self.log("╚════██║██╔══╝  ██╔══╝  ██║  ██║╚██╗ ██╔╝██╔══██╗    ██╔═══╝       ╚════██║", category="none", force=True, indent_level=1)
-        self.log(" ║ ███████ ███████ ███████ ██████    ████   ██   ██ ███████ ║", category="none", force=True)
+        self.log("███████║███████╗███████╗██████╔╝ ╚████╔╝ ██║  ██║    ███████╗  ██╗ ███████║", category="none", force=True, indent_level=1)
-        if cli:
+        self.log("╚══════╝╚══════╝╚══════╝╚═════╝   ╚═══╝  ╚═╝  ╚═╝    ╚══════╝  ╚═╝ ╚══════╝", category="none", force=True, indent_level=1)
-            self.log(" ║ 💻 CLI mode             © ByteDance Seed · NumZ · AInVFX ║", category="none", force=True)
+        # Version and credits - left/right aligned to logo width
+        version_text = f"v{__version__}"
+        cli_indicator = "💻 CLI · " if cli else ""
+        left_part = f"{cli_indicator}{version_text}"
+        right_part = "© ByteDance Seed · NumZ · AInVFX"
+        logo_width = 75
+        emoji_compensation = 1 if cli else 0
+        padding = logo_width - len(left_part) - len(right_part) - emoji_compensation
+        self.log(f"{left_part}{' ' * max(1, padding)}{right_part}", category="none", force=True, indent_level=1)
+        self.log("━" * logo_width, category="none", force=True, indent_level=1)
+        self.log("", category="none", force=True)
+        
+        # Restore timestamps setting
+        self.show_timestamps = original_timestamps
+        
+        # Environment info - only in debug mode
+        if self.enabled:
+            self._print_environment_info(cli)
+
+    def _print_environment_info(self, cli: bool = False) -> None:
+        """Print concise environment info for bug reports - zero cost when debug disabled"""
+        import platform
+        import sys
+        
+        # OS
+        os_name = platform.system()
+        if os_name == "Windows":
+            os_str = f"Windows ({platform.version()})"
+        elif os_name == "Darwin":
+            os_str = f"macOS {platform.mac_ver()[0]}"
         else:
-            self.log(" ║                         © ByteDance Seed · NumZ · AInVFX ║", category="none", force=True)
+            try:
-        self.log(" ╚══════════════════════════════════════════════════════════╝", category="none", force=True)
+                distro = platform.freedesktop_os_release()
-        self.log("", category="none", force=True)
+                os_str = f"{distro.get('NAME', 'Linux')} {distro.get('VERSION_ID', '')}"
+            except (OSError, AttributeError):
+                os_str = f"Linux {platform.release()}"
+        
+        # Python & PyTorch & CUDA
+        py_ver = f"{sys.version_info.major}.{sys.version_info.minor}.{sys.version_info.micro}"
+        torch_ver = torch.__version__
+        cuda_ver = getattr(torch.version, 'cuda', None) or "N/A"
+        
+        # GPU
+        if is_cuda_available():
+            try:
+                props = torch.cuda.get_device_properties(0)
+                gpu_str = f"{props.name} ({round(props.total_memory / (1024**3))}GB)"
+                cudnn_ver = str(torch.backends.cudnn.version()) if torch.backends.cudnn.is_available() else "N/A"
+            except Exception:
+                gpu_str = "CUDA"
+                cudnn_ver = "N/A"
+        elif is_mps_available():
+            gpu_str = "Apple Silicon (MPS)"
+            cudnn_ver = "N/A"
+        else:
+            gpu_str = "CPU"
+            cudnn_ver = "N/A"
+        
+        # Flash Attn, SageAttn & Triton - reuse existing module constants
+        try:
+            from ..optimization.compatibility import (
+                FLASH_ATTN_2_AVAILABLE, FLASH_ATTN_3_AVAILABLE,
+                SAGE_ATTN_2_AVAILABLE, SAGE_ATTN_3_AVAILABLE,
+                TRITON_AVAILABLE
+            )
+            fa_parts = []
+            if FLASH_ATTN_3_AVAILABLE:
+                fa_parts.append("3")
+            if FLASH_ATTN_2_AVAILABLE:
+                fa_parts.append("2")
+            flash_str = f"v{','.join(fa_parts)} ✓" if fa_parts else "✗"
+            
+            sa_parts = []
+            if SAGE_ATTN_3_AVAILABLE:
+                sa_parts.append("3")
+            if SAGE_ATTN_2_AVAILABLE:
+                sa_parts.append("2")
+            sage_str = f"v{','.join(sa_parts)} ✓" if sa_parts else "✗"
+            
+            triton_str = "✓" if TRITON_AVAILABLE else "✗"
+        except ImportError:
+            flash_str = sage_str = triton_str = "?"
+        
+        # ComfyUI version
+        comfy_str = None
+        if not cli:
+            try:
+                from comfyui_version import __version__ as comfy_ver
+                comfy_str = comfy_ver
+            except ImportError:
+                pass
+        
+        # Print
+        self.log(f"OS: {os_str} | GPU: {gpu_str}", category="info")
+        self.log(f"Python: {py_ver} | PyTorch: {torch_ver} | FlashAttn: {flash_str} | SageAttn: {sage_str} | Triton: {triton_str}", category="info")
+        cuda_line = f"CUDA: {cuda_ver} | cuDNN: {cudnn_ver}"
+        self.log(f"{cuda_line} | ComfyUI: {comfy_str}" if comfy_str else cuda_line, category="info")
+        self.log("", category="none")
 
     def print_footer(self) -> None:
         """Print the footer with links - always displayed"""
         self.log("", category="none", force=True)
         self.log("────────────────────────", category="none", force=True)
-        self.log("Questions? Updates? Watch the videos, star the repo & join us!", category="dialogue", force=True)
+        self.log("Questions? Updates? Watch, star & sponsor if you can!", category="dialogue", force=True)
         self.log("https://www.youtube.com/@AInVFX", category="generation", force=True)
-        self.log("https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler", category="star", force=True)
+        self.log("https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler", category="starlove", force=True)
     
     @torch._dynamo.disable  # Skip tracing to avoid time.time() warnings
     def start_timer(self, name: str, force: bool = False) -> None:
@@ -299,31 +424,55 @@ class Debug:
         if show_diff and self.memory_checkpoints:
             self._log_memory_diff(current_metrics=memory_info, force=force)
 
-       # Log detailed analysis if requested
+        # Overflow warning (Windows only - WDDM can page to system RAM)
+        overflow = memory_info.get('vram_overflow', 0.0)
+        
+        if overflow > 0 and platform.system() == 'Windows':
+            self.log(f"VRAM overflow: {overflow:.2f}GB paged to system RAM - severe slowdown expected. "
+                     "Consider optimizing (e.g., reduce resolution, batch size, enable BlockSwap, VAE tiling...).",
+                     level="WARNING", category="memory", force=True)
+
+        # Log detailed analysis if requested
         if detailed_tensors and tensor_stats.get('details'):
             self._log_detailed_tensor_analysis(details=tensor_stats['details'], force=force)
                 
         # Store checkpoint with memory limit
         self._store_checkpoint(label, memory_info)
 
-        # Update phase peak if we're in an active phase
+        # Update phase peaks if we're in an active phase
-        if self.current_phase and memory_info['vram_peak_since_last'] > 0:
+        if self.current_phase:
-            self.phase_peaks[self.current_phase] = max(
+            if memory_info['vram_peak_alloc'] > 0:
-                self.phase_peaks.get(self.current_phase, 0),
+                self.phase_vram_peaks_alloc[self.current_phase] = max(
-                memory_info['vram_peak_since_last']
+                    self.phase_vram_peaks_alloc.get(self.current_phase, 0),
-            )
+                    memory_info['vram_peak_alloc']
+                )
+            if memory_info['vram_peak_rsv'] > 0:
+                self.phase_vram_peaks_rsv[self.current_phase] = max(
+                    self.phase_vram_peaks_rsv.get(self.current_phase, 0),
+                    memory_info['vram_peak_rsv']
+                )
+            if memory_info['ram_process'] > 0:
+                self.phase_ram_peaks[self.current_phase] = max(
+                    self.phase_ram_peaks.get(self.current_phase, 0),
+                    memory_info['ram_process']
+                )
 
         # Reset PyTorch's peak memory stats for next interval
         reset_vram_peak(device=None, debug=self)
     
     def _collect_memory_metrics(self) -> Dict[str, Any]:
-        """Collect current memory metrics efficiently."""
+        """Collect current memory metrics."""
+        is_mps = is_mps_available()
+        has_gpu = is_mps or is_cuda_available()
+        
         metrics = {
             'vram_allocated': 0.0,
             'vram_reserved': 0.0,
             'vram_free': 0.0,
             'vram_total': 0.0,
-            'vram_peak_since_last': 0.0,
+            'vram_peak_alloc': 0.0,
+            'vram_peak_rsv': 0.0,
+            'vram_overflow': 0.0,
             'ram_process': 0.0,
             'ram_available': 0.0,
             'ram_total': 0.0,
@@ -332,45 +481,36 @@ class Debug:
             'summary_ram': ""
         }
         
-        # VRAM metrics
+        if has_gpu:
-        if torch.cuda.is_available() or torch.mps.is_available():
+            metrics['vram_allocated'], metrics['vram_reserved'], metrics['vram_peak_alloc'], metrics['vram_peak_rsv'] = get_vram_usage(device=None, debug=self)
-            metrics['vram_allocated'], metrics['vram_reserved'], current_global_peak = get_vram_usage(device=None, debug=self)
-
-            # Calculate peak since last log_memory_state
-            # This captures the actual peak that occurred between calls
-            metrics['vram_peak_since_last'] = current_global_peak
-            
             vram_info = get_basic_vram_info(device=None)
             
-            if "error" not in vram_info:
+            if "error" not in vram_info and vram_info["total_gb"] > 0:
                 metrics['vram_free'] = vram_info["free_gb"]
                 metrics['vram_total'] = vram_info["total_gb"]
+                metrics['vram_overflow'] = max(0.0, metrics['vram_peak_rsv'] - metrics['vram_total'])
                 
-                backend = "MPS" if torch.mps.is_available() else "VRAM"
+                backend = "Unified Memory" if is_mps else "VRAM"
-                metrics['summary_vram'] = (f"  [{backend}] {metrics['vram_allocated']:.2f}GB allocated / "
+                metrics['summary_vram'] = (
-                        f"{metrics['vram_reserved']:.2f}GB reserved / "
+                    f"  [{backend}] {metrics['vram_allocated']:.2f}GB allocated / "
-                        f"Peak: {metrics['vram_peak_since_last']:.2f}GB / "
+                    f"{metrics['vram_reserved']:.2f}GB reserved / "
-                        f"{metrics['vram_free']:.2f}GB free / "
+                    f"Peak: {metrics['vram_peak_alloc']:.2f}GB / "
-                        f"{metrics['vram_total']:.2f}GB total")
+                    f"{metrics['vram_free']:.2f}GB free / "
-            else:
+                    f"{metrics['vram_total']:.2f}GB total"
-                metrics['summary_vram'] = ""
+                )
-        else:
+            
-            metrics['summary_vram'] = ""
+            self.vram_history.append(metrics['vram_reserved'])
         
-        # RAM metrics using new function
+        # RAM metrics
         metrics['ram_process'], metrics['ram_available'], metrics['ram_total'], metrics['ram_others'] = get_ram_usage(debug=self)
         
         if metrics['ram_total'] > 0:
-            metrics['summary_ram'] = (f"  [RAM] {metrics['ram_process']:.2f}GB process / "
+            metrics['summary_ram'] = (
-                      f"{metrics['ram_others']:.2f}GB others / "
+                f"  [RAM] {metrics['ram_process']:.2f}GB process / "
-                      f"{metrics['ram_available']:.2f}GB free / "
+                f"{metrics['ram_others']:.2f}GB others / "
-                      f"{metrics['ram_total']:.2f}GB total")
+                f"{metrics['ram_available']:.2f}GB free / "
-        else:
+                f"{metrics['ram_total']:.2f}GB total"
-            metrics['summary_ram'] = ""
+            )
-        
-        # Update VRAM history for tracking
-        if torch.cuda.is_available() or torch.mps.is_available():
-            self.vram_history.append(metrics['vram_allocated'])
         
         return metrics
     
@@ -497,9 +637,9 @@ class Debug:
         if diffs:
             self.log(f"Memory changes: {', '.join(diffs)}", category="memory", force=force, indent_level=1)
     
-    def log_peak_vram_summary(self, force: bool = True) -> None:
+    def log_peak_memory_summary(self, force: bool = True) -> None:
-        """Display peak VRAM usage across all phases"""
+        """Display peak memory usage across all phases."""
-        if not self.phase_peaks:
+        if not self.phase_vram_peaks_alloc and not self.phase_ram_peaks:
             return
         
         phase_names = {
@@ -509,17 +649,42 @@ class Debug:
             'phase4': 'Post-processing'
         }
         
+        is_mps = is_mps_available()
+        
+        # Get total VRAM for overflow formatting (Windows only)
+        total_vram_gb = 0.0
+        if not is_mps:
+            vram_info = get_basic_vram_info(device=None)
+            if "error" not in vram_info:
+                total_vram_gb = vram_info["total_gb"]
+        
         self.log("", category="none", force=force)
         self.log("────────────────────────", category="none", force=force)
-        self.log("Peak VRAM by Phase:", category="memory", force=force)
+        self.log("Peak memory by phase:", category="memory", force=force)
         
-        for phase_key in sorted(self.phase_peaks.keys()):
+        all_phases = sorted(set(self.phase_vram_peaks_alloc.keys()) | set(self.phase_ram_peaks.keys()))
+        for phase_key in all_phases:
             phase_num = phase_key[-1]
-            self.log(f"  Phase {phase_num}: {phase_names[phase_key]}: {self.phase_peaks[phase_key]:.2f}GB", 
+            phase_name = phase_names.get(phase_key, phase_key)
-                    category="memory", force=force)
+            alloc = self.phase_vram_peaks_alloc.get(phase_key, 0)
+            rsv = self.phase_vram_peaks_rsv.get(phase_key, 0)
+            ram = self.phase_ram_peaks.get(phase_key, 0)
+            
+            if is_mps:
+                self.log(f"{phase_num}. {phase_name}: {alloc:.2f}GB", category="memory", indent_level=1, force=force)
+            else:
+                rsv_str = _format_peak_with_overflow(rsv, total_vram_gb)
+                self.log(f"{phase_num}. {phase_name}: VRAM {alloc:.2f}GB allocated, {rsv_str} | RAM {ram:.2f}GB", category="memory", indent_level=1, force=force)
         
-        overall_max = max(self.phase_peaks.values())
+        overall_alloc = max(self.phase_vram_peaks_alloc.values()) if self.phase_vram_peaks_alloc else 0
-        self.log(f"Overall Peak VRAM: {overall_max:.2f}GB", category="memory", force=force)
+        overall_rsv = max(self.phase_vram_peaks_rsv.values()) if self.phase_vram_peaks_rsv else 0
+        overall_ram = max(self.phase_ram_peaks.values()) if self.phase_ram_peaks else 0
+        
+        if is_mps:
+            self.log(f"Overall peak: {overall_alloc:.2f}GB", category="memory", force=force)
+        else:
+            overall_rsv_str = _format_peak_with_overflow(overall_rsv, total_vram_gb)
+            self.log(f"Overall peak: VRAM {overall_alloc:.2f}GB allocated, {overall_rsv_str} | RAM {overall_ram:.2f}GB", category="memory", force=force)
     
     @torch._dynamo.disable  # Skip tracing to avoid time.time() warnings
     def _store_checkpoint(self, label: str, metrics: Dict[str, Any]) -> None:
@@ -629,5 +794,7 @@ class Debug:
         self.timer_durations.clear()
         self.timer_messages.clear()
         self.active_timer_stack.clear()
-        self.phase_peaks.clear()
+        self.phase_vram_peaks_alloc.clear()
+        self.phase_vram_peaks_rsv.clear()
+        self.phase_ram_peaks.clear()
         self.current_phase = None

src/utils/downloads.py

@@ -14,6 +14,7 @@ import time
 from .model_registry import MODEL_REGISTRY, DEFAULT_VAE
 from .constants import (
     get_base_cache_dir,
+    get_all_model_paths,
     find_model_file,
     get_validation_cache_path,
     HUGGINGFACE_BASE_URL,
@@ -179,20 +180,36 @@ def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = N
                          level="WARNING", category="setup")
             continue
         
+        # Get model type for logging
+        model_type = "VAE" if model_info.category == "vae" else "DiT"
+        
         # Check if file exists in any registered path first
         existing_filepath = find_model_file(filename, fallback_dir=cache_dir)
         
         # Use existing file path if it exists, otherwise download to cache_dir
         if os.path.exists(existing_filepath):
             filepath = existing_filepath
+            # Debug log: Model found
+            if debug:
+                debug.log(f"{model_type} model found: {filepath}", category="setup")
         else:
             filepath = os.path.join(cache_dir, filename)
+            # Debug log: Model not found, will need to download
+            if debug:
+                searched_paths = get_all_model_paths()
+                debug.log(f"{model_type} model not found: {filename}", category="setup")
+                debug.log(f"Searched in {len(searched_paths)} location(s):", category="setup")
+                for i, path in enumerate(searched_paths, 1):
+                    debug.log(f"[{i}] {path}", category="setup", indent_level=1)
             
         expected_hash = model_info.sha256
         repo = model_info.repo
         
         ## Quick cache check first
         if is_file_validated_cached(filepath, cache_dir):
+            # Debug log: Model already validated (using cache)
+            if debug:
+                debug.log(f"{model_type} model already validated (cache): {filepath}", category="setup")
             continue
         
         # File exists - validate it
@@ -201,6 +218,9 @@ def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = N
                 debug.log(f"Validating {filename}...", category="setup", force=True)
             
             if validate_file(filepath, expected_hash, cache_dir):
+                # Debug log: Model validated successfully
+                if debug:
+                    debug.log(f"{model_type} model validated successfully: {filepath}", category="setup")
                 continue
             else:
                 # File is corrupted