diff --git a/.gitattributes b/.gitattributes index 7f1e78a3633bb6b5f7efb6172e3fe90ce791c5ce..b8833f46d66111051b21264ed307416050d78580 100644 --- a/.gitattributes +++ b/.gitattributes @@ -296,3 +296,9 @@ v3-nodes/ComfyUI-RMBG/example_workflows/V3.0.0_nodes.jpg filter=lfs diff=lfs mer v3-nodes/ComfyUI-RMBG/example_workflows/YOLO_Node.jpg filter=lfs diff=lfs merge=lfs -text v3-nodes/ComfyUI-RMBG/example_workflows/florence2_node.jpg filter=lfs diff=lfs merge=lfs -text v3-nodes/ComfyUI-RMBG/models/sam3/perflib/tests/assets/masks.tiff filter=lfs diff=lfs merge=lfs -text +v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_01.jpg filter=lfs diff=lfs merge=lfs -text +v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_02.jpg filter=lfs diff=lfs merge=lfs -text +v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_01.png filter=lfs diff=lfs merge=lfs -text +v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_02.png filter=lfs diff=lfs merge=lfs -text +v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Mustache_640x360.mp4 filter=lfs diff=lfs merge=lfs -text +v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Sadhu_320x478.png filter=lfs diff=lfs merge=lfs -text diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.github/FUNDING.yml b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.github/FUNDING.yml new file mode 100644 index 0000000000000000000000000000000000000000..30b15aa2a72d112f87534655c707130bba80133b --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.github/FUNDING.yml @@ -0,0 +1 @@ +github: [adrientoupet, numz] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.github/workflows/publish.yml b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.github/workflows/publish.yml new file mode 100644 index 0000000000000000000000000000000000000000..dd44ae2363bbcd575636a8e1ab1a8fd545dcdc38 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.github/workflows/publish.yml @@ -0,0 +1,27 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'numz' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + with: + submodules: true + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@v1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.gitignore b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.gitignore new file mode 100644 index 0000000000000000000000000000000000000000..7b478a87af1b3474697d0f55cdb4770fb35c1ef9 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/.gitignore @@ -0,0 +1,26 @@ +.git/* +**/__pycache__/ +tests/ +.vscode/ +.cursor/ +benchmark/ +advanced_optimizations.py +BENCHMARK_* +environment.yml +install_safetensors.py +manage_quantized_models.py +quantization_config.py +quantize_model.py +quantize_vae.py +README_*.md +run_*.py +vram_diagnostic.py +test_*.py +VRAM_OPTIMIZATIONS_SUMMARY.md +seedvr2.py +src/core/isolated_generation.py +src/core/subprocess_runner.py +models/video_vae_v3_mine_bad/ +src/processing/ +TILE_VAE* +.DS_Store \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/CONTRIBUTING.md b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/CONTRIBUTING.md new file mode 100644 index 0000000000000000000000000000000000000000..87020696b5bcb680026025ae960ad012cc2eaa26 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/CONTRIBUTING.md @@ -0,0 +1,34 @@ +# Contributing to ComfyUI-SeedVR2_VideoUpscaler + +Thank you for your interest in contributing to ComfyUI-SeedVR2_VideoUpscaler! We appreciate your effort and to help us incorporate your contribution in the best way possible, please follow the following contribution guidelines. + +## Reporting Bugs + +If you find a bug in the project, we encourage you to report it. Here's how: + +1. First, check the [existing Issues](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues) to see if the issue has already been reported. If it has, please add a comment to the existing issue rather than creating a new one. +2. If you can't find an existing issue that matches your bug, create a new issue. Make sure to include as many details as possible so we can understand and reproduce the problem. + +## Proposing Changes + +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 the **`main`** branch. + +We use the GitHub Flow workflow. + +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. + +## Contact + +If you have any questions or need help, please reach out to the developers: + +- **NumZ**: Discord NumZ#7184 +- **adrientoupet** from AInVFX: [YouTube Channel](https://www.youtube.com/@AInVFX) + +You can also open an issue on GitHub for general questions and discussions. + +Thank you again for your contribution ! diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/LICENSE b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..9b0726cabaaf26dcc32cfe2f95c7780039e82a2c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "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. diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/README.md b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/README.md new file mode 100644 index 0000000000000000000000000000000000000000..197ad2d8880c44478beb72f7b21578ecfbc7028c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/README.md @@ -0,0 +1,1053 @@ +# ComfyUI-SeedVR2_VideoUpscaler + +[![View Code](https://img.shields.io/badge/๐Ÿ“‚_View_Code-GitHub-181717?style=for-the-badge&logo=github)](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler) + +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. + +[![SeedVR2 v2.5 Deep Dive Tutorial](https://img.youtube.com/vi/MBtWYXq_r60/maxresdefault.jpg)](https://youtu.be/MBtWYXq_r60) + +![Usage Example](docs/usage_01.png) + +![Usage Example](docs/usage_02.png) + +## ๐Ÿ“‹ Quick Access + +- [๐Ÿ†™ Future Work](#-future-work) +- [๐Ÿš€ Release Notes](#-release-notes) +- [๐ŸŽฏ Features](#-features) +- [๐Ÿ”ง Requirements](#-requirements) +- [๐Ÿ“ฆ Installation](#-installation) +- [๐Ÿ“– Usage](#-usage) +- [๐Ÿ–ฅ๏ธ Run as Standalone](#๏ธ-run-as-standalone-cli) +- [โš ๏ธ Limitations](#๏ธ-limitations) +- [๐Ÿค Contributing](#-contributing) +- [๐Ÿ™ Credits](#-credits) +- [๐Ÿ“œ License](#-license) + +## ๐Ÿ†™ Future Work + +We're actively working on improvements and new features. To stay informed: + +- **๐Ÿ“Œ Track Active Development**: Visit [Issues](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues) to see active development, report bugs, and request new features +- **๐Ÿ’ฌ 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) + +## ๐Ÿš€ 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** ๐ŸŽ‰ + +โš ๏ธ **BREAKING CHANGE**: This is a major update requiring workflow recreation. All nodes and CLI parameters have been redesigned for better usability and consistency. Watch the latest video from [AInVFX](https://www.youtube.com/@AInVFX) for a deep dive and check out the [usage](#-usage) section. + +**๐Ÿ“ฆ Official Release**: Now available on main branch with ComfyUI Manager support for easy installation and automatic version tracking. Updated dependencies and local imports prevent conflicts with other ComfyUI custom nodes. + +### ๐ŸŽจ ComfyUI Improvements + +- **Four-Node Modular Architecture**: Split into dedicated nodes for DiT model, VAE model, torch.compile settings, and main upscaler for granular control +- **Global Model Cache**: Models now shared across multiple upscaler instances with automatic config updates - no more redundant loading +- **ComfyUI V3 Migration**: Full compatibility with ComfyUI V3 stateless node design +- **RGBA Support**: Native alpha channel processing with edge-guided upscaling for clean transparency +- **Improved Memory Management**: Streaming architecture prevents VRAM spikes regardless of video length +- **Flexible Resolution Support**: Upscale to any resolution divisible by 2 with lossless padding approach (replaced restrictive cropping) +- **Enhanced Parameters**: Added `uniform_batch_size`, `temporal_overlap`, `prepend_frames`, and `max_resolution` for better control + +### ๐Ÿ–ฅ๏ธ CLI Enhancements + +- **Batch Directory Processing**: Process entire folders of videos/images with model caching for efficiency +- **Single Image Support**: Direct image upscaling without video conversion +- **Smart Output Detection**: Auto-detects output format (MP4/PNG) based on input type +- **Enhanced Multi-GPU**: Improved workload distribution with temporal overlap blending +- **Unified Parameters**: CLI and ComfyUI now use identical parameter names for consistency +- **Better UX**: Auto-display help, validation improvements, progress tracking, and cleaner output + +### โšก Performance & Optimization + +- **torch.compile Support**: 20-40% DiT speedup and 15-25% VAE speedup with full graph compilation +- **Optimized BlockSwap**: Adaptive memory clearing (5% threshold), separate I/O component handling, reduced overhead +- **Enhanced VAE Tiling**: Tensor offload support for accumulation buffers, separate encode/decode configuration +- **Native Dtype Pipeline**: Eliminated unnecessary conversions, maintains bfloat16 precision throughout for speed and quality +- **Optimized Tensor Operations**: Replaced einops rearrange with native PyTorch ops for 2-5x faster transforms + +### ๐ŸŽฏ Quality Improvements + +- **LAB Color Correction**: New perceptual color transfer method with superior color accuracy (now default) +- **Additional Color Methods**: HSV saturation matching, wavelet adaptive, and hybrid approaches +- **Deterministic Generation**: Seed-based reproducibility with phase-specific seeding strategy +- **Better Temporal Consistency**: Hann window blending for smooth transitions between batches + +### ๐Ÿ’พ Memory Management + +- **Smarter Offloading**: Independent device configuration for DiT, VAE, and tensors (CPU/GPU/none) +- **Four-Phase Pipeline**: Completes each phase (encodeโ†’upscaleโ†’decodeโ†’postprocess) for all batches before moving to next, minimizing model swaps +- **Better Cleanup**: Phase-specific resource management with proper tensor memory release +- **Peak VRAM Tracking**: Per-phase memory monitoring with summary display + +### ๐Ÿ”ง Technical Improvements + +- **GGUF Quantization Support**: Added full GGUF support for 4-bit/8-bit inference on low-VRAM systems +- **Improved GGUF Handling**: Fixed VRAM leaks, torch.compile compatibility, non-persistent buffers +- **Apple Silicon Support**: Full MPS (Metal Performance Shaders) support for Apple Silicon Macs +- **AMD ROCm Compatibility**: Conditional FSDP imports for PyTorch ROCm 7+ support +- **Conv3d Memory Workaround**: Fixes PyTorch 2.9+ cuDNN memory bug (3x usage reduction) +- **Flash Attention Optional**: Graceful fallback to SDPA when flash-attn unavailable + +### ๐Ÿ“š Code Quality + +- **Modular Architecture**: Split monolithic files into focused modules (generation_phases, model_configuration, etc.) +- **Comprehensive Documentation**: Extensive docstrings with type hints across all modules +- **Better Error Handling**: Early validation, clear error messages, installation instructions +- **Consistent Logging**: Unified indentation, better categorization, concise messages + +**2025.08.07** + +- ๐ŸŽฏ **Unified Debug System**: New structured logging with categories, timers, and memory tracking. `enable_debug` now available on main node +- โšก **Smart FP8 Optimization**: FP8 models now keep native FP8 storage, converting to BFloat16 only for arithmetic - faster and more memory efficient than FP16 +- ๐Ÿ“ฆ **Model Registry**: Multi-repo support (numz/ & AInVFX/), auto-discovery of user models, added mixed FP8 variants to fix 7B artifacts +- ๐Ÿ’พ **Model Caching**: `cache_model` moved to main node, fixed memory leaks with proper RoPE/wrapper cleanup +- ๐Ÿงน **Code Cleanup**: New modular structure (`constants.py`, `model_registry.py`, `debug.py`), removed legacy code +- ๐Ÿš€ **Performance**: Better memory management with `torch.cuda.ipc_collect()`, improved RoPE handling + +**2025.07.17** + +- ๐Ÿ› ๏ธ Add 7B sharp Models: add 2 new 7B models with sharpen output + +**2025.07.11** + +- ๐ŸŽฌ Complete tutorial released: Adrien from [AInVFX](https://www.youtube.com/@AInVFX) created an in-depth ComfyUI SeedVR2 guide covering everything from basic setup to advanced BlockSwap techniques for running on consumer GPUs. Perfect for understanding memory optimization and upscaling of image sequences with alpha channel! [Watch the tutorial](#-usage) + +**2025.09.07** + +- ๐Ÿ› ๏ธ Blockswap Integration: Big thanks to [Adrien Toupet](https://github.com/adrientoupet) from [AInVFX](https://www.youtube.com/@AInVFX) for this :), useful for low VRAM users (see [usage](#-usage) section) + +**2025.07.03** + +- ๐Ÿ› ๏ธ Can run as **standalone mode** with **Multi GPU** see [๐Ÿ–ฅ๏ธ Run as Standalone](#run-as-standalone-cli) + +**2025.06.30** + +- ๐Ÿš€ Speed Up the process and less VRAM used +- ๐Ÿ› ๏ธ Fixed memory leak on 3B models +- โŒ Can now interrupt process if needed +- โœ… Refactored the code for better sharing with the community, feel free to propose pull requests +- ๐Ÿ› ๏ธ Removed flash attention dependency (thanks to [luke2642](https://github.com/Luke2642) !!) + +**2025.06.24** + +- ๐Ÿš€ Speed up the process until x4 + +**2025.06.22** + +- ๐Ÿ’ช FP8 compatibility ! +- ๐Ÿš€ Speed Up all Process +- ๐Ÿš€ less VRAM consumption (Stay high, batch_size=1 for RTX4090 max, I'm trying to fix that) +- ๐Ÿ› ๏ธ Better benchmark coming soon + +**2025.06.20** + +- ๐Ÿ› ๏ธ Initial push + +## ๐ŸŽฏ Features + +### Core Capabilities +- **High-Quality Diffusion-Based Upscaling**: One-step diffusion model for video and image enhancement +- **Temporal Consistency**: Maintains coherence across video frames with configurable batch processing +- **Multi-Format Support**: Handles RGB and RGBA (alpha channel) for both videos and images +- **Any Video Length**: Suitable for any video length + +### Model Support +- **Multiple Model Variants**: 3B and 7B parameter models with different precision options +- **FP16, FP8, and GGUF Quantization**: Choose between full precision (FP16), mixed precision (FP8), or heavily quantized GGUF models for different VRAM requirements +- **Automatic Model Downloads**: Models are automatically downloaded from HuggingFace on first use + +### Memory Optimization +- **BlockSwap Technology**: Dynamically swap transformer blocks between GPU and CPU memory to run large models on limited VRAM +- **VAE Tiling**: Process large resolutions with tiled encoding/decoding to reduce VRAM usage +- **Intelligent Offloading**: Offload models and intermediate tensors to CPU or secondary GPUs between processing phases +- **GGUF Quantization Support**: Run models with 4-bit or 8-bit quantization for extreme VRAM savings + +### 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 between generations for single-GPU directory processing or multi-GPU streaming +- **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 +- **Noise Injection Controls**: Fine-tune input and latent noise scales for artifact reduction at high resolutions +- **Configurable Resolution Limits**: Set target and maximum resolutions with automatic aspect ratio preservation + +### Workflow Features +- **ComfyUI Integration**: Four dedicated nodes for complete control over the upscaling pipeline +- **Standalone CLI**: Command-line interface for batch processing and automation +- **Debug Logging**: Comprehensive debug mode with memory tracking, timing information, and processing details +- **Progress Reporting**: Real-time progress updates during processing + +## ๐Ÿ”ง Requirements + +### Hardware + +With the current optimizations (tiling, BlockSwap, GGUF quantization), SeedVR2 can run on a wide range of hardware: + +- **Minimal VRAM** (8GB or less): Use GGUF Q4_K_M models with BlockSwap and VAE tiling enabled +- **Moderate VRAM** (12-16GB): Use FP8 models with BlockSwap or VAE tiling as needed +- **High VRAM** (24GB+): Use FP16 models for best quality and speed without memory optimizations + +### Software + +- **ComfyUI**: Latest version recommended +- **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 / 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 + +### Option 1: ComfyUI Manager (Recommended) + +1. Open ComfyUI Manager in your ComfyUI interface +2. Click "Custom Nodes Manager" +3. Search for "ComfyUI-SeedVR2_VideoUpscaler" +4. Click "Install" and restart ComfyUI + +**Registry Link**: [ComfyUI Registry - SeedVR2 Video Upscaler](https://registry.comfy.org/nodes/seedvr2_videoupscaler) + +### Option 2: Manual Installation + +1. **Clone the repository** into your ComfyUI custom nodes directory: +```bash +cd ComfyUI +git clone https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler.git custom_nodes/seedvr2_videoupscaler +``` + +2. **Install dependencies using standalone Python**: +```bash +# Install requirements (from same ComfyUI directory) +# Windows: +.venv\Scripts\python.exe -m pip install -r custom_nodes\seedvr2_videoupscaler\requirements.txt +# Linux/macOS: +.venv/bin/python -m pip install -r custom_nodes/seedvr2_videoupscaler/requirements.txt +``` + +3. **Restart ComfyUI** + +### Model Installation + +Models will be **automatically downloaded** on first use and saved to `ComfyUI/models/SEEDVR2`. + +You can also manually download models from: +- Main models available at [numz/SeedVR2_comfyUI](https://huggingface.co/numz/SeedVR2_comfyUI/tree/main) and [AInVFX/SeedVR2_comfyUI](https://huggingface.co/AInVFX/SeedVR2_comfyUI/tree/main) +- Additional GGUF models available at [cmeka/SeedVR2-GGUF](https://huggingface.co/cmeka/SeedVR2-GGUF/tree/main) + +## ๐Ÿ“– Usage + +### ๐ŸŽฌ Video Tutorials + +#### Latest Version Deep Dive (Recommended) + +Complete walkthrough of version 2.5 by Adrien from [AInVFX](https://www.youtube.com/@AInVFX), covering the new 4-node architecture, GGUF support, memory optimizations, and production workflows: + +[![SeedVR2 v2.5 Deep Dive Tutorial](https://img.youtube.com/vi/MBtWYXq_r60/maxresdefault.jpg)](https://youtu.be/MBtWYXq_r60) + +This comprehensive tutorial covers: +- Installing v2.5 through ComfyUI Manager and troubleshooting conflicts +- Understanding the new 4-node modular architecture and why we rebuilt it +- Running 7B models on 8GB VRAM with GGUF quantization +- Configuring BlockSwap, VAE tiling, and torch.compile for your hardware +- Image and video upscaling workflows with alpha channel support +- CLI for batch processing and multi-GPU rendering +- Memory optimization strategies for different VRAM levels +- Real production tips and the critical batch_size formula (4n+1) + +#### Previous Version Tutorial + +For reference, here's the original tutorial covering the initial release: + +[![SeedVR2 Deep Dive Tutorial](https://img.youtube.com/vi/I0sl45GMqNg/maxresdefault.jpg)](https://youtu.be/I0sl45GMqNg) + +*Note: This tutorial covers the previous single-node architecture. While the UI has changed significantly in v2.5, the core concepts about BlockSwap and memory management remain valuable.* + +### Node Setup + +SeedVR2 uses a modular node architecture with four specialized nodes: + +#### 1. SeedVR2 (Down)Load DiT Model + +![SeedVR2 (Down)Load DiT Model](docs/dit_model_loader.png) + +Configure the DiT (Diffusion Transformer) model for video upscaling. + +**Parameters:** + +- **model**: Choose your DiT model + - **3B Models**: Faster, lower VRAM requirements + - `seedvr2_ema_3b_fp16.safetensors`: FP16 (best quality) + - `seedvr2_ema_3b_fp8_e4m3fn.safetensors`: FP8 8-bit (good quality) + - `seedvr2_ema_3b-Q4_K_M.gguf`: GGUF 4-bit quantized (acceptable quality) + - `seedvr2_ema_3b-Q8_0.gguf`: GGUF 8-bit quantized (good quality) + - **7B Models**: Higher quality, higher VRAM requirements + - `seedvr2_ema_7b_fp16.safetensors`: FP16 (best quality) + - `seedvr2_ema_7b_fp8_e4m3fn_mixed_block35_fp16.safetensors`: FP8 with last block in FP16 to reduce artifacts (good quality) + - `seedvr2_ema_7b-Q4_K_M.gguf`: GGUF 4-bit quantized (acceptable quality) + - `seedvr2_ema_7b_sharp_*`: Sharp variants for enhanced detail + +- **device**: GPU device for DiT inference (e.g., `cuda:0`) + +- **offload_device**: Device to offload DiT model when not actively processing + - `none`: Keep model on inference device (fastest, highest VRAM) + - `cpu`: Offload to system RAM (reduces VRAM) + - `cuda:X`: Offload to another GPU (good balance if available) + +- **cache_model**: Keep DiT model loaded on offload_device between workflow runs + - Useful for batch processing to avoid repeated loading + - Requires offload_device to be set + +- **blocks_to_swap**: BlockSwap memory optimization + - `0`: Disabled (default) + - `1-32`: Number of transformer blocks to swap for 3B model + - `1-36`: Number of transformer blocks to swap for 7B model + - Higher values = more VRAM savings but slower processing + - Requires offload_device to be set and different from device + +- **swap_io_components**: Offload input/output embeddings and normalization layers + - Additional VRAM savings when combined with blocks_to_swap + - Requires offload_device to be set and different from device + +- **attention_mode**: Attention computation backend + - `sdpa`: PyTorch scaled_dot_product_attention (default, always available) + - `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. + +> **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 +- **How to configure**: + 1. Set `offload_device` to `cpu` or another GPU + 2. Start with `blocks_to_swap=16` (half the blocks) + 3. If still getting OOM, increase to 24 or 32 (3B) / 36 (7B) + 4. Enable `swap_io_components` for maximum VRAM savings + 5. If you have plenty of VRAM, decrease or set to 0 for faster processing + +**Example Configuration for Low VRAM (8GB)**: +- model: `seedvr2_ema_3b-Q8_0.gguf` +- device: `cuda:0` +- offload_device: `cpu` +- blocks_to_swap: `32` +- swap_io_components: `True` + +#### 2. SeedVR2 (Down)Load VAE Model + +![SeedVR2 (Down)Load VAE Model](docs/vae_model_loader.png) + +Configure the VAE (Variational Autoencoder) model for encoding/decoding video frames. + +**Parameters:** + +- **model**: VAE model selection + - `ema_vae_fp16.safetensors`: Default and recommended + +- **device**: GPU device for VAE inference (e.g., `cuda:0`) + +- **offload_device**: Device to offload VAE model when not actively processing + - `none`: Keep model on inference device (default, fastest) + - `cpu`: Offload to system RAM (reduces VRAM) + - `cuda:X`: Offload to another GPU (good balance if available) + +- **cache_model**: Keep VAE model loaded on offload_device between workflow runs + - Requires offload_device to be set + +- **encode_tiled**: Enable tiled encoding to reduce VRAM usage during encoding phase + - Enable if you see OOM errors during the "Encoding" phase in debug logs + +- **encode_tile_size**: Encoding tile size in pixels (default: 1024) + - Applied to both height and width + - Lower values reduce VRAM but may increase processing time + +- **encode_tile_overlap**: Encoding tile overlap in pixels (default: 128) + - Reduces visible seams between tiles + +- **decode_tiled**: Enable tiled decoding to reduce VRAM usage during decoding phase + - Enable if you see OOM errors during the "Decoding" phase in debug logs + +- **decode_tile_size**: Decoding tile size in pixels (default: 1024) + +- **decode_tile_overlap**: Decoding tile overlap in pixels (default: 128) + +- **torch_compile_args**: Connect to SeedVR2 Torch Compile Settings node for 15-25% speedup + +**VAE Tiling Explained:** + +VAE tiling processes large resolutions in smaller tiles to reduce VRAM requirements. Here's how to use it: + +1. **Run without tiling first** and monitor the debug logs (enable `enable_debug` on main node) +2. **If OOM during "Encoding" phase**: + - Enable `encode_tiled` + - If still OOM, reduce `encode_tile_size` (try 768, 512, etc.) +3. **If OOM during "Decoding" phase**: + - Enable `decode_tiled` + - If still OOM, reduce `decode_tile_size` +4. **Adjust overlap** (default 128) if you see visible seams in output (increase it) or processing times are too slow (decrease it). + +**Example Configuration for High Resolution (4K)**: +- encode_tiled: `True` +- encode_tile_size: `1024` +- encode_tile_overlap: `128` +- decode_tiled: `True` +- decode_tile_size: `1024` +- decode_tile_overlap: `128` + +#### 3. SeedVR2 Torch Compile Settings (Optional) + +![SeedVR2 Torch Compile Settings](docs/torch_compile_settings.png) + +Configure torch.compile optimization for 20-40% DiT speedup and 15-25% VAE speedup. + +**Requirements:** +- PyTorch 2.0+ +- Triton (for inductor backend) + +**Parameters:** + +- **backend**: Compilation backend + - `inductor`: Full optimization with Triton kernel generation and fusion (recommended) + - `cudagraphs`: Lightweight wrapper using CUDA graphs, no kernel optimization + +- **mode**: Optimization level (compilation time vs runtime performance) + - `default`: Fast compilation with good speedup (recommended for development) + - `reduce-overhead`: Lower overhead, optimized for smaller models + - `max-autotune`: Slowest compilation, best runtime performance (recommended for production) + - `max-autotune-no-cudagraphs`: Like max-autotune but without CUDA graphs + +- **fullgraph**: Compile entire model as single graph without breaks + - `False`: Allow graph breaks for better compatibility (default, recommended) + - `True`: Enforce no breaks for maximum optimization (may fail with dynamic shapes) + +- **dynamic**: Handle varying input shapes without recompilation + - `False`: Specialize for exact input shapes (default) + - `True`: Create dynamic kernels that adapt to shape variations (enable when processing different resolutions or batch sizes) + +- **dynamo_cache_size_limit**: Max cached compiled versions per function (default: 64) + - Higher = more memory, lower = more recompilation + +- **dynamo_recompile_limit**: Max recompilation attempts before falling back to eager mode (default: 128) + - Safety limit to prevent compilation loops + +**Usage:** +1. Add this node to your workflow +2. Connect its output to the `torch_compile_args` input of DiT and/or VAE loader nodes +3. First run will be slow (compilation), subsequent runs will be much faster + +**When to use:** +- torch.compile only makes sense when processing **multiple batches, long videos, or many tiles** +- For single images or short clips, the compilation time outweighs the speed improvement +- Best suited for batch processing workflows or long videos + +**Recommended Settings:** +- For development/testing: `mode=default`, `backend=inductor`, `fullgraph=False` +- For production: `mode=max-autotune`, `backend=inductor`, `fullgraph=False` + +#### 4. SeedVR2 Video Upscaler (Main Node) + +![SeedVR2 Video Upscaler](docs/video_upscaler.png) + +Main upscaling node that processes video frames using DiT and VAE models. + +**Required Inputs:** + +- **image**: Input video frames as image batch (RGB or RGBA format) +- **dit**: DiT model configuration from SeedVR2 (Down)Load DiT Model node +- **vae**: VAE model configuration from SeedVR2 (Down)Load VAE Model node + +**Parameters:** + +- **seed**: Random seed for reproducible generation (default: 42) + - Same seed with same inputs produces identical output + +- **resolution**: Target resolution for shortest edge in pixels (default: 1080) + - Maintains aspect ratio automatically + +- **max_resolution**: Maximum resolution for any edge (default: 0 = no limit) + - Automatically scales down if exceeded to prevent OOM + +- **batch_size**: Frames per batch (default: 5) + - **CRITICAL REQUIREMENT**: Must follow the **4n+1 formula** (1, 5, 9, 13, 17, 21, 25, ...) + - **Why this matters**: The model uses these frames for temporal consistency calculations + - **Minimum 5 for temporal consistency**: Use 1 only for single images or when temporal consistency isn't needed + - **Match shot length ideally**: For best results, set batch_size to match your shot length (e.g., batch_size=21 for a 20-frame shot) + - **VRAM impact**: Higher batch_size = better quality and speed but requires more VRAM + - **If you get OOM with batch_size=5**: Try optimization techniques first (model offloading, BlockSwap, GGUF models...) before reducing batch_size or input resolution, as these directly impact quality + +**uniform_batch_size** (default: False) + - Pads the final batch to match `batch_size` for uniform processing + - Prevents temporal artifacts when the last batch is significantly smaller than others + - Example: 45 frames with `batch_size=33` creates [33, 33] instead of [33, 12] + - Recommended when using large batch sizes and video length is not a multiple of `batch_size` + - Increases VRAM usage slightly but ensures consistent temporal coherence across all batches + +- **temporal_overlap**: Overlapping frames between batches (default: 0) + - Used for blending between batches to reduce temporal artifacts + - Range: 0-16 frames + +- **prepend_frames**: Frames to prepend (default: 0) + - Prepends reversed frames to reduce artifacts at video start + - Automatically removed after processing + - Range: 0-32 frames + +- **color_correction**: Color correction method (default: "wavelet") + - **`lab`**: Full perceptual color matching with detail preservation (recommended for highest fidelity to original) + - **`wavelet`**: Frequency-based natural colors, preserves details well + - **`wavelet_adaptive`**: Wavelet base + targeted saturation correction + - **`hsv`**: Hue-conditional saturation matching + - **`adain`**: Statistical style transfer + - **`none`**: No color correction + +- **input_noise_scale**: Input noise injection scale 0.0-1.0 (default: 0.0) + - Adds noise to input frames to reduce artifacts at very high resolutions + - Try 0.1-0.3 if you see artifacts with high output resolutions + +- **latent_noise_scale**: Latent space noise scale 0.0-1.0 (default: 0.0) + - Adds noise during diffusion process, can soften excessive detail + - Use if input_noise doesn't help, try 0.05-0.15 + +- **offload_device**: Device for storing intermediate tensors between processing phases (default: "cpu") + - `none`: Keep all tensors on inference device (fastest but highest VRAM) + - `cpu`: Offload to system RAM (recommended for long videos, slower transfers) + - `cuda:X`: Offload to another GPU (good balance if available, faster than CPU) + +- **enable_debug**: Enable detailed debug logging (default: False) + - Shows memory usage, timing information, and processing details + - **Highly recommended** for troubleshooting OOM issues + +**Output:** +- Upscaled video frames with color correction applied +- Format (RGB/RGBA) matches input +- Range [0, 1] normalized for ComfyUI compatibility + +### Typical Workflow Setup + +**Basic Workflow (High VRAM - 24GB+)**: +``` +Load Video Frames + โ†“ +SeedVR2 Load DiT Model + โ”œโ”€ model: seedvr2_ema_3b_fp16.safetensors + โ””โ”€ device: cuda:0 + โ†“ +SeedVR2 Load VAE Model + โ”œโ”€ model: ema_vae_fp16.safetensors + โ””โ”€ device: cuda:0 + โ†“ +SeedVR2 Video Upscaler + โ”œโ”€ batch_size: 21 + โ””โ”€ resolution: 1080 + โ†“ +Save Video/Frames +``` + +**Low VRAM Workflow (8-12GB)**: +``` +Load Video Frames + โ†“ +SeedVR2 Load DiT Model + โ”œโ”€ model: seedvr2_ema_3b-Q8_0.gguf + โ”œโ”€ device: cuda:0 + โ”œโ”€ offload_device: cpu + โ”œโ”€ blocks_to_swap: 32 + โ””โ”€ swap_io_components: True + โ†“ +SeedVR2 Load VAE Model + โ”œโ”€ model: ema_vae_fp16.safetensors + โ”œโ”€ device: cuda:0 + โ”œโ”€ encode_tiled: True + โ””โ”€ decode_tiled: True + โ†“ +SeedVR2 Video Upscaler + โ”œโ”€ batch_size: 5 + โ””โ”€ resolution: 720 + โ†“ +Save Video/Frames +``` + +**High Performance Workflow (24GB+ with torch.compile)**: +``` +Load Video Frames + โ†“ +SeedVR2 Torch Compile Settings + โ”œโ”€ mode: max-autotune + โ””โ”€ backend: inductor + โ†“ +SeedVR2 Load DiT Model + โ”œโ”€ model: seedvr2_ema_7b_sharp_fp16.safetensors + โ”œโ”€ device: cuda:0 + โ””โ”€ torch_compile_args: connected + โ†“ +SeedVR2 Load VAE Model + โ”œโ”€ model: ema_vae_fp16.safetensors + โ”œโ”€ device: cuda:0 + โ””โ”€ torch_compile_args: connected + โ†“ +SeedVR2 Video Upscaler + โ”œโ”€ batch_size: 81 + โ””โ”€ resolution: 1080 + โ†“ +Save Video/Frames +``` + +## ๐Ÿ–ฅ๏ธ Run as Standalone (CLI) + +The standalone CLI provides powerful batch processing capabilities with multi-GPU support and sophisticated optimization options. + +### Prerequisites + +Choose the appropriate setup based on your installation: + +#### Option 1: Already Have ComfyUI with SeedVR2 Installed + +If you've already installed SeedVR2 as part of ComfyUI (via [ComfyUI installation](#-installation)), you can use the CLI directly: + +```bash +# Navigate to your ComfyUI directory +cd ComfyUI + +# Run the CLI using standalone Python (display help message) +# Windows: +.venv\Scripts\python.exe custom_nodes\seedvr2_videoupscaler\inference_cli.py --help +# Linux/macOS: +.venv/bin/python custom_nodes/seedvr2_videoupscaler/inference_cli.py --help +``` + +**Skip to [Command Line Usage](#command-line-usage) below.** + +#### Option 2: Standalone Installation (Without ComfyUI) + +If you want to use the CLI without ComfyUI installation, follow these steps: + +1. **Install [uv](https://docs.astral.sh/uv/getting-started/installation/)** (modern Python package manager): +```bash +# Windows +powershell -ExecutionPolicy ByPass -c "irm https://astral.sh/uv/install.ps1 | iex" + +# macOS and Linux +curl -LsSf https://astral.sh/uv/install.sh | sh +``` + +2. **Clone the repository**: +```bash +git clone https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler.git seedvr2_videoupscaler +cd seedvr2_videoupscaler +``` + +3. **Create virtual environment and install dependencies**: +```bash +# Create virtual environment with Python 3.13 +uv venv --python 3.13 + +# Activate virtual environment +# Windows: +.venv\Scripts\activate +# Linux/macOS: +source .venv/bin/activate + +# Install PyTorch with CUDA support +# Check command line based on your environment: https://pytorch.org/get-started/locally/ +uv pip install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu130 + +# Install SeedVR2 requirements +uv pip install -r requirements.txt + +# Run the CLI (display help message) +# Windows: +.venv\Scripts\python.exe inference_cli.py --help +# Linux/macOS: +.venv/bin/python inference_cli.py --help +``` + +### Command Line Usage + +The CLI provides comprehensive options for single-GPU, multi-GPU, and batch processing workflows. + +**Basic Usage Examples:** + +```bash +# Basic image upscaling +python inference_cli.py image.jpg + +# 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 \ + --resolution 1080 \ + --batch_size 81 \ + --uniform_batch_size \ + --temporal_overlap 3 \ + --prepend_frames 4 + +# Memory-optimized for low VRAM (8GB) +python inference_cli.py image.png \ + --dit_model seedvr2_ema_3b-Q8_0.gguf \ + --resolution 1080 \ + --blocks_to_swap 32 \ + --swap_io_components \ + --dit_offload_device cpu \ + --vae_offload_device cpu + +# High resolution with VAE tiling +python inference_cli.py video.mp4 \ + --resolution 1440 \ + --batch_size 31 \ + --uniform_batch_size \ + --temporal_overlap 3 \ + --vae_encode_tiled \ + --vae_decode_tiled + +# Batch directory processing with model caching +python inference_cli.py media_folder/ \ + --output processed/ \ + --cuda_device 0 \ + --cache_dit \ + --cache_vae \ + --dit_offload_device cpu \ + --vae_offload_device cpu \ + --resolution 1080 \ + --max_resolution 1920 +``` + +### Command Line Arguments + +**Input/Output:** +- ``: 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:** +- `--dit_model`: DiT model to use. Options: 3B/7B with fp16/fp8/GGUF variants (default: 3B FP8) + +**Processing Parameters:** +- `--resolution`: Target short-side resolution in pixels (default: 1080) +- `--max_resolution`: Maximum resolution for any edge. Scales down if exceeded. 0 = no limit (default: 0) +- `--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`: 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) + +**Quality Control:** +- `--color_correction`: Color correction method: 'lab' (perceptual, recommended), 'wavelet', 'wavelet_adaptive', 'hsv', 'adain', or 'none' (default: lab) +- `--input_noise_scale`: Input noise injection scale (0.0-1.0). Reduces artifacts at high resolutions (default: 0.0) +- `--latent_noise_scale`: Latent space noise scale (0.0-1.0). Softens details if needed (default: 0.0) + +**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). Not available on macOS. +- `--swap_io_components`: Offload I/O components for additional VRAM savings. Requires dit_offload_device. Not available on macOS. + +**VAE Tiling:** +- `--vae_encode_tiled`: Enable VAE encode tiling to reduce VRAM during encoding +- `--vae_encode_tile_size`: VAE encode tile size in pixels (default: 1024) +- `--vae_encode_tile_overlap`: VAE encode tile overlap in pixels (default: 128) +- `--vae_decode_tiled`: Enable VAE decode tiling to reduce VRAM during decoding +- `--vae_decode_tile_size`: VAE decode tile size in pixels (default: 1024) +- `--vae_decode_tile_overlap`: VAE decode tile overlap in pixels (default: 128) +- `--tile_debug`: Visualize tiles: 'false' (default), 'encode', or 'decode' + +**Performance Optimization:** +- `--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) +- `--compile_mode`: Optimization level: 'default', 'reduce-overhead', 'max-autotune', 'max-autotune-no-cudagraphs' (default: default) +- `--compile_fullgraph`: Compile entire model as single graph (faster but less flexible) (default: False) +- `--compile_dynamic`: Handle varying input shapes without recompilation (default: False) +- `--compile_dynamo_cache_size_limit`: Max cached compiled versions per function (default: 64) +- `--compile_dynamo_recompile_limit`: Max recompilation attempts before fallback (default: 128) + +**Model Caching (batch 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`: 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 + +**Debugging:** +- `--debug`: Enable verbose debug logging + +### Multi-GPU Processing Explained + +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 frames are split evenly across GPUs (e.g., 100 frames on 2 GPUs โ†’ 50 frames each) +2. Each GPU loads its own copy of the models and processes its chunk independently +3. When `--temporal_overlap` is set, chunks include overlapping frames for seamless blending +4. Results are concatenated (and blended at overlap regions) into the final video + +**Example for 100 frames on 2 GPUs with temporal_overlap=4:** +``` +GPU 0: Frames 0-53 (50 base + 4 overlap at end, processed as independent video) +GPU 1: Frames 50-99 (50 frames, 4 overlap at start, processed as independent video) +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-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 +- For optimal quality on short videos, use single GPU with `batch_size` matching your shot length + +## โš ๏ธ Limitations + +### Model Limitations + +**Batch Size Constraint**: The model requires batch_size to follow the **4n+1 formula** (1, 5, 9, 13, 17, 21, 25, ...) due to temporal consistency architecture. All frames in a batch are processed together for temporal coherence, then batches can be blended using temporal_overlap. Ideally, set batch_size to match your shot length for optimal quality. + +### Performance Considerations + +**VAE Bottleneck**: Even with optimized DiT upscaling (BlockSwap, GGUF, torch.compile), the VAE encoding/decoding stages can be the bottleneck, especially for high resolutions. The VAE is slow. Use large batch_size to mitigate this. + +**VRAM Usage**: While the integration now supports low VRAM systems (8GB or less with proper optimization), VRAM usage varies based on: +- Input/output resolution (larger = more VRAM) +- Batch size (higher = more VRAM but better temporal consistency and speed) +- Model choice (FP16 > FP8 > GGUF in VRAM usage) +- Optimization settings (BlockSwap, VAE tiling significantly reduce VRAM) + +**Speed**: Processing speed depends on: +- GPU capabilities (compute performance, VRAM bandwidth, and architecture generation) +- Model size (3B faster than 7B) +- Batch size (larger batch sizes are faster per frame due to better GPU utilization) +- Optimization settings (torch.compile provides significant speedup) +- Resolution (higher resolutions are slower) + +### Best Practices + +1. **Start with debug enabled** to understand where VRAM is being used +2. **For OOM errors during encoding**: Enable VAE encode tiling and reduce tile size +3. **For OOM errors during upscaling**: Enable BlockSwap and increase blocks_to_swap +4. **For OOM errors during decoding**: Enable VAE decode tiling and reduce tile size + - **If still getting OOM after trying all above**: Reduce batch_size or resolution +5. **For best quality**: Use higher batch_size matching your shot length, FP16 models, and LAB color correction +6. **For speed**: Use FP8/GGUF models, enable torch.compile, and use Flash Attention if available +7. **Test settings with a short clip first** before processing long videos + +## ๐Ÿค Contributing + +Contributions are welcome! We value community input and improvements. + +For detailed contribution guidelines, see [CONTRIBUTING.md](CONTRIBUTING.md). + +**Quick Start:** + +1. Fork the repository +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 the **main** branch + +**Get Help:** +- YouTube: [AInVFX Channel](https://www.youtube.com/@AInVFX) +- GitHub [Issues](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues): For bug reports and feature requests +- GitHub [Discussions](https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions): For questions and community support +- Discord: adrientoupet & NumZ#7184 + +## ๐Ÿ™ Credits + +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 [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 Apache 2.0 license as found in the [LICENSE](LICENSE) file. \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c38b50c22f98997cd857b80614ddd2ae80144551 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/__init__.py @@ -0,0 +1,9 @@ +""" +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"] \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/configs_3b/main.yaml b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/configs_3b/main.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ebd5adf8b32446d4b5742e4cc406a0d84f48679d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/configs_3b/main.yaml @@ -0,0 +1,91 @@ +__object__: + path: projects.video_diffusion_sr.train + name: VideoDiffusionTrainer + +dit: + model: + __object__: + path: "dit_3b.nadit" + name: "NaDiT" + args: "as_params" + vid_in_channels: 33 + vid_out_channels: 16 + vid_dim: 2560 + vid_out_norm: fusedrms + txt_in_dim: 5120 + txt_in_norm: fusedln + txt_dim: ${.vid_dim} + emb_dim: ${eval:'6 * ${.vid_dim}'} + heads: 20 + head_dim: 128 # llm-like + expand_ratio: 4 + norm: fusedrms + norm_eps: 1.0e-05 + ada: single + qk_bias: False + qk_norm: fusedrms + patch_size: [1, 2, 2] + num_layers: 32 # llm-like + mm_layers: 10 + mlp_type: swiglu + msa_type: None + block_type: ${eval:'${.num_layers} * ["mmdit_sr"]'} # space-full + window: ${eval:'${.num_layers} * [(4,3,3)]'} # space-full + window_method: ${eval:'${.num_layers} // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"]'} # space-full + rope_type: mmrope3d + rope_dim: 128 + compile: False + gradient_checkpoint: True + fsdp: + sharding_strategy: _HYBRID_SHARD_ZERO2 + +ema: + decay: 0.9998 + +vae: + model: + __object__: + path: "video_vae_v3.modules.attn_video_vae" + name: "VideoAutoencoderKLWrapper" + args: "as_params" + freeze_encoder: False + gradient_checkpoint: True # Disabled to prevent VRAM leaks in inference + slicing: + split_size: 4 + memory_device: same + memory_limit: + conv_max_mem: 0.5 + norm_max_mem: 0.5 + checkpoint: ema_vae_fp16.safetensors + scaling_factor: 0.9152 + compile: False + grouping: False + dtype: float16 + +diffusion: + schedule: + type: lerp + T: 1000.0 + sampler: + type: euler + prediction_type: v_lerp + timesteps: + training: + type: logitnormal + loc: 0.0 + scale: 1.0 + sampling: + type: uniform_trailing + steps: 50 + transform: True + loss: + type: v_lerp + cfg: + scale: 7.5 + rescale: 0 + +condition: + i2v: 0.0 + v2v: 0.0 + sr: 1.0 + noise_scale: 0.25 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/configs_7b/main.yaml b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/configs_7b/main.yaml new file mode 100644 index 0000000000000000000000000000000000000000..ac66e9a016b9849b5fded7313bf01e52f2b69aed --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/configs_7b/main.yaml @@ -0,0 +1,88 @@ +__object__: + path: projects.video_diffusion_sr.train + name: VideoDiffusionTrainer + +dit: + model: + __object__: + path: "dit_7b.nadit" + name: "NaDiT" + args: "as_params" + vid_in_channels: 33 + vid_out_channels: 16 + vid_dim: 3072 + txt_in_dim: 5120 + txt_dim: ${.vid_dim} + emb_dim: ${eval:'6 * ${.vid_dim}'} + heads: 24 + head_dim: 128 # llm-like + expand_ratio: 4 + norm: fusedrms + norm_eps: 1e-5 + ada: single + qk_bias: False + qk_rope: True + qk_norm: fusedrms + patch_size: [1, 2, 2] + num_layers: 36 # llm-like + shared_mlp: False + shared_qkv: False + mlp_type: normal + block_type: ${eval:'${.num_layers} * ["mmdit_sr"]'} # space-full + window: ${eval:'${.num_layers} * [(4,3,3)]'} # space-full + window_method: ${eval:'${.num_layers} // 2 * ["720pwin_by_size_bysize","720pswin_by_size_bysize"]'} # space-full + compile: False + gradient_checkpoint: True + fsdp: + sharding_strategy: _HYBRID_SHARD_ZERO2 + +ema: + decay: 0.9998 + +vae: + model: + __object__: + path: "video_vae_v3.modules.attn_video_vae" + name: "VideoAutoencoderKLWrapper" + args: "as_params" + freeze_encoder: False + # gradient_checkpoint: True + slicing: + split_size: 4 + memory_device: same + memory_limit: + conv_max_mem: 0.5 + norm_max_mem: 0.5 + checkpoint: ema_vae_fp16.safetensors + scaling_factor: 0.9152 + compile: False + grouping: False + dtype: float16 + +diffusion: + schedule: + type: lerp + T: 1000.0 + sampler: + type: euler + prediction_type: v_lerp + timesteps: + training: + type: logitnormal + loc: 0.0 + scale: 1.0 + sampling: + type: uniform_trailing + steps: 50 + transform: True + loss: + type: v_lerp + cfg: + scale: 7.5 + rescale: 0 + +condition: + i2v: 0.0 + v2v: 0.0 + sr: 1.0 + noise_scale: 0.25 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_01.jpg b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_01.jpg new file mode 100644 index 0000000000000000000000000000000000000000..4e77cec2f0d5cc8e530f20c506790052008a5131 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_01.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:356496e9421e5af83c40916cc45a44122f935b770dd5b7c62f8ca6c2f3065b5a +size 453418 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_02.jpg b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_02.jpg new file mode 100644 index 0000000000000000000000000000000000000000..c32e719d0f9181fa10ede9f2e4b1446e3338dbf4 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/demo_02.jpg @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:50bcd1fac36961f979c01fea565dc6c55babcc16c0ef5b8bb1b4c8dfbd75edaf +size 152095 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/dit_model_loader.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/dit_model_loader.png new file mode 100644 index 0000000000000000000000000000000000000000..837d356c8ab264ce219d4c6cc542644d05479d4d Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/dit_model_loader.png differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/seedvr_logo.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/seedvr_logo.png new file mode 100644 index 0000000000000000000000000000000000000000..5582b98fa9e29073c94071a98d7e6ce97833630b Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/seedvr_logo.png differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/torch_compile_settings.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/torch_compile_settings.png new file mode 100644 index 0000000000000000000000000000000000000000..73555e746e5a54dcdb5ea7b8bceb242711a2628e Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/torch_compile_settings.png differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_01.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_01.png new file mode 100644 index 0000000000000000000000000000000000000000..d290c94c539d6c8645bbf01ac25fafb0aa0902d9 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_01.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:927198f9c5119df8d4b126a8ec4cd367a8b0ab2f471f755636843ecac5ebcfff +size 1366882 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_02.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_02.png new file mode 100644 index 0000000000000000000000000000000000000000..63ac40679ad2fd64390844f2537278a9167c8996 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/usage_02.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:73266574b70c31fc06cc990e6229471ddba373ea5e808f92a1214c422bfacc37 +size 686123 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/vae_model_loader.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/vae_model_loader.png new file mode 100644 index 0000000000000000000000000000000000000000..349d8b278e85721c57c175b1e7d165b9f3793977 Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/vae_model_loader.png differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/video_upscaler.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/video_upscaler.png new file mode 100644 index 0000000000000000000000000000000000000000..ad5e3d1eb4ce1c86832114efee2212b1b7706925 Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/docs/video_upscaler.png differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.jpg b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.jpg new file mode 100644 index 0000000000000000000000000000000000000000..ab06d99891ce3e6cb0c58fca9bad1034fe8cf011 Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.jpg differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.json b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.json new file mode 100644 index 0000000000000000000000000000000000000000..54a1d42349fa514a8ff632d570a56c7e894630e2 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_4K_image_upscale.json @@ -0,0 +1 @@ +{"id":"e2c8e6ec-9e8a-4a81-a240-7e7d2befc2bf","revision":0,"last_node_id":20,"last_link_id":19,"nodes":[{"id":19,"type":"SeedVR2TorchCompileSettings","pos":[-366.25065217175177,401.8303561599628],"size":[307.671484375,178],"flags":{},"order":0,"mode":4,"inputs":[{"localized_name":"backend","name":"backend","type":"COMBO","widget":{"name":"backend"},"link":null},{"localized_name":"mode","name":"mode","type":"COMBO","widget":{"name":"mode"},"link":null},{"localized_name":"fullgraph","name":"fullgraph","type":"BOOLEAN","widget":{"name":"fullgraph"},"link":null},{"localized_name":"dynamic","name":"dynamic","type":"BOOLEAN","widget":{"name":"dynamic"},"link":null},{"localized_name":"dynamo_cache_size_limit","name":"dynamo_cache_size_limit","type":"INT","widget":{"name":"dynamo_cache_size_limit"},"link":null},{"localized_name":"dynamo_recompile_limit","name":"dynamo_recompile_limit","type":"INT","widget":{"name":"dynamo_recompile_limit"},"link":null}],"outputs":[{"localized_name":"TORCH_COMPILE_ARGS","name":"TORCH_COMPILE_ARGS","type":"TORCH_COMPILE_ARGS","links":[18,19]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2TorchCompileSettings"},"widgets_values":["inductor","default",false,false,64,128]},{"id":18,"type":"Note","pos":[420.71483854947945,-81.0108307454865],"size":[210,94.05179298404067],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Enable to upscale alpha/mask channel along with RGB channel (for RGBA inputs)"],"color":"#432","bgcolor":"#653"},{"id":16,"type":"LoadImage","pos":[80.5773583147188,-182.2697173004957],"size":[274.080078125,314],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[14]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":[17]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"LoadImage"},"widgets_values":["Sadhu_320x478.png","image"]},{"id":14,"type":"SeedVR2LoadDiTModel","pos":[64.4234520647185,212.7072358245033],"size":[307.6646484375,202],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"torch_compile_args","name":"torch_compile_args","shape":7,"type":"TORCH_COMPILE_ARGS","link":18},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"blocks_to_swap","name":"blocks_to_swap","shape":7,"type":"INT","widget":{"name":"blocks_to_swap"},"link":null},{"localized_name":"swap_io_components","name":"swap_io_components","shape":7,"type":"BOOLEAN","widget":{"name":"swap_io_components"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"cache_model","name":"cache_model","shape":7,"type":"BOOLEAN","widget":{"name":"cache_model"},"link":null},{"localized_name":"attention_mode","name":"attention_mode","shape":7,"type":"COMBO","widget":{"name":"attention_mode"},"link":null}],"outputs":[{"localized_name":"SEEDVR2_DIT","name":"SEEDVR2_DIT","type":"SEEDVR2_DIT","links":[11]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2LoadDiTModel"},"widgets_values":["seedvr2_ema_7b_sharp_fp16.safetensors","cuda:0",36,false,"cpu",false,"sdpa"]},{"id":13,"type":"SeedVR2LoadVAEModel","pos":[60.33985831471847,499.8798920745057],"size":[312.866796875,298],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"torch_compile_args","name":"torch_compile_args","shape":7,"type":"TORCH_COMPILE_ARGS","link":19},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"encode_tiled","name":"encode_tiled","shape":7,"type":"BOOLEAN","widget":{"name":"encode_tiled"},"link":null},{"localized_name":"encode_tile_size","name":"encode_tile_size","shape":7,"type":"INT","widget":{"name":"encode_tile_size"},"link":null},{"localized_name":"encode_tile_overlap","name":"encode_tile_overlap","shape":7,"type":"INT","widget":{"name":"encode_tile_overlap"},"link":null},{"localized_name":"decode_tiled","name":"decode_tiled","shape":7,"type":"BOOLEAN","widget":{"name":"decode_tiled"},"link":null},{"localized_name":"decode_tile_size","name":"decode_tile_size","shape":7,"type":"INT","widget":{"name":"decode_tile_size"},"link":null},{"localized_name":"decode_tile_overlap","name":"decode_tile_overlap","shape":7,"type":"INT","widget":{"name":"decode_tile_overlap"},"link":null},{"localized_name":"tile_debug","name":"tile_debug","shape":7,"type":"COMBO","widget":{"name":"tile_debug"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"cache_model","name":"cache_model","shape":7,"type":"BOOLEAN","widget":{"name":"cache_model"},"link":null}],"outputs":[{"localized_name":"SEEDVR2_VAE","name":"SEEDVR2_VAE","type":"SEEDVR2_VAE","links":[10]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2LoadVAEModel"},"widgets_values":["ema_vae_fp16.safetensors","cuda:0",true,1024,128,true,1024,128,"false","cpu",false]},{"id":10,"type":"SeedVR2VideoUpscaler","pos":[843.8992333147189,191.3224701995034],"size":[270,386],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":16},{"localized_name":"dit","name":"dit","type":"SEEDVR2_DIT","link":11},{"localized_name":"vae","name":"vae","type":"SEEDVR2_VAE","link":10},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"localized_name":"resolution","name":"resolution","type":"INT","widget":{"name":"resolution"},"link":null},{"localized_name":"max_resolution","name":"max_resolution","type":"INT","widget":{"name":"max_resolution"},"link":null},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null},{"localized_name":"uniform_batch_size","name":"uniform_batch_size","type":"BOOLEAN","widget":{"name":"uniform_batch_size"},"link":null},{"localized_name":"color_correction","name":"color_correction","type":"COMBO","widget":{"name":"color_correction"},"link":null},{"localized_name":"temporal_overlap","name":"temporal_overlap","shape":7,"type":"INT","widget":{"name":"temporal_overlap"},"link":null},{"localized_name":"prepend_frames","name":"prepend_frames","shape":7,"type":"INT","widget":{"name":"prepend_frames"},"link":null},{"localized_name":"input_noise_scale","name":"input_noise_scale","shape":7,"type":"FLOAT","widget":{"name":"input_noise_scale"},"link":null},{"localized_name":"latent_noise_scale","name":"latent_noise_scale","shape":7,"type":"FLOAT","widget":{"name":"latent_noise_scale"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"enable_debug","name":"enable_debug","shape":7,"type":"BOOLEAN","widget":{"name":"enable_debug"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[12]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2VideoUpscaler"},"widgets_values":[42,"fixed",4096,4096,1,false,"lab",0,0,0,0,"cpu",false]},{"id":17,"type":"JoinImageWithAlpha","pos":[449.4979840435381,-180.32741340177097],"size":[176.86484375,46],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":14},{"localized_name":"alpha","name":"alpha","type":"MASK","link":17}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[16]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"JoinImageWithAlpha"},"widgets_values":[]},{"id":15,"type":"SaveImage","pos":[1167.7234520647182,192.0920014495034],"size":[270,270],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":12},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.68"},"widgets_values":["ComfyUI"]},{"id":20,"type":"Note","pos":[-353.44461861972593,641.0883773827939],"size":[290.1922351462714,97.95592567265356],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Enable torch.compile optimization. First batch has compilation overhead, but speeds up subsequent batches. Best for long videos, not single images."],"color":"#432","bgcolor":"#653"}],"links":[[10,13,0,10,2,"SEEDVR2_VAE"],[11,14,0,10,1,"SEEDVR2_DIT"],[12,10,0,15,0,"IMAGE"],[14,16,0,17,0,"IMAGE"],[16,17,0,10,0,"IMAGE"],[17,16,1,17,1,"MASK"],[18,19,0,14,0,"TORCH_COMPILE_ARGS"],[19,19,0,13,0,"TORCH_COMPILE_ARGS"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.6934334949441353,"offset":[1011.3669791186508,462.16171611318566]}},"version":0.4} \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.jpg b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.jpg new file mode 100644 index 0000000000000000000000000000000000000000..7bc428787e7efce6b2684f31cbd76cb794227915 Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.jpg differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.json b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.json new file mode 100644 index 0000000000000000000000000000000000000000..24b46f9a613482c741e437f66c62b9faa5fed84e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_HD_video_upscale.json @@ -0,0 +1 @@ +{"id":"e2c8e6ec-9e8a-4a81-a240-7e7d2befc2bf","revision":0,"last_node_id":24,"last_link_id":25,"nodes":[{"id":17,"type":"JoinImageWithAlpha","pos":[415.46463925106394,-269.3527671179127],"size":[176.86484375,46],"flags":{},"order":7,"mode":4,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":21},{"localized_name":"alpha","name":"alpha","type":"MASK","link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[16]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"JoinImageWithAlpha"},"widgets_values":[]},{"id":19,"type":"SeedVR2TorchCompileSettings","pos":[-369.50847923647007,315.0949953354577],"size":[307.671484375,178],"flags":{},"order":0,"mode":4,"inputs":[{"localized_name":"backend","name":"backend","type":"COMBO","widget":{"name":"backend"},"link":null},{"localized_name":"mode","name":"mode","type":"COMBO","widget":{"name":"mode"},"link":null},{"localized_name":"fullgraph","name":"fullgraph","type":"BOOLEAN","widget":{"name":"fullgraph"},"link":null},{"localized_name":"dynamic","name":"dynamic","type":"BOOLEAN","widget":{"name":"dynamic"},"link":null},{"localized_name":"dynamo_cache_size_limit","name":"dynamo_cache_size_limit","type":"INT","widget":{"name":"dynamo_cache_size_limit"},"link":null},{"localized_name":"dynamo_recompile_limit","name":"dynamo_recompile_limit","type":"INT","widget":{"name":"dynamo_recompile_limit"},"link":null}],"outputs":[{"localized_name":"TORCH_COMPILE_ARGS","name":"TORCH_COMPILE_ARGS","type":"TORCH_COMPILE_ARGS","links":[18,19]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2TorchCompileSettings"},"widgets_values":["inductor","default",false,false,64,128]},{"id":20,"type":"Note","pos":[-358.9923959032341,553.3187800066305],"size":[290.1922351462714,97.95592567265356],"flags":{},"order":1,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Enable torch.compile optimization. First batch has compilation overhead, but speeds up subsequent batches. Best for long videos, not single images."],"color":"#432","bgcolor":"#653"},{"id":21,"type":"LoadVideo","pos":[-201.66033140975628,-266.86977401462417],"size":[274.080078125,228.1700439453125],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[20]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"LoadVideo"},"widgets_values":["Mustache_640x360.mp4","image"]},{"id":18,"type":"Note","pos":[417.45701148476115,-167.7461915699906],"size":[210,94.05179298404067],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Enable + connect input alpha to upscale alpha/mask channel along with RGB channel (for RGBA inputs)"],"color":"#432","bgcolor":"#653"},{"id":24,"type":"CreateVideo","pos":[1171.8133710326213,102.33318716823825],"size":[270,78],"flags":{},"order":9,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":25},{"localized_name":"fps","name":"fps","type":"FLOAT","widget":{"name":"fps"},"link":24}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[23]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"CreateVideo"},"widgets_values":[30]},{"id":23,"type":"SaveVideo","pos":[1499.476781650457,101.92559527677722],"size":[270,249.875],"flags":{},"order":10,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":23},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null},{"localized_name":"format","name":"format","type":"COMBO","widget":{"name":"format"},"link":null},{"localized_name":"codec","name":"codec","type":"COMBO","widget":{"name":"codec"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"SaveVideo"},"widgets_values":["video/ComfyUI","auto","auto"]},{"id":22,"type":"GetVideoComponents","pos":[149.28727188425574,-267.41033914727126],"size":[185.17734375,66],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"video","name":"video","type":"VIDEO","link":20}],"outputs":[{"localized_name":"images","name":"images","type":"IMAGE","links":[21]},{"localized_name":"audio","name":"audio","type":"AUDIO","links":[25]},{"localized_name":"fps","name":"fps","type":"FLOAT","links":[24]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"GetVideoComponents"},"widgets_values":[]},{"id":13,"type":"SeedVR2LoadVAEModel","pos":[57.08203125,413.14453125],"size":[312.866796875,298],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"torch_compile_args","name":"torch_compile_args","shape":7,"type":"TORCH_COMPILE_ARGS","link":19},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"encode_tiled","name":"encode_tiled","shape":7,"type":"BOOLEAN","widget":{"name":"encode_tiled"},"link":null},{"localized_name":"encode_tile_size","name":"encode_tile_size","shape":7,"type":"INT","widget":{"name":"encode_tile_size"},"link":null},{"localized_name":"encode_tile_overlap","name":"encode_tile_overlap","shape":7,"type":"INT","widget":{"name":"encode_tile_overlap"},"link":null},{"localized_name":"decode_tiled","name":"decode_tiled","shape":7,"type":"BOOLEAN","widget":{"name":"decode_tiled"},"link":null},{"localized_name":"decode_tile_size","name":"decode_tile_size","shape":7,"type":"INT","widget":{"name":"decode_tile_size"},"link":null},{"localized_name":"decode_tile_overlap","name":"decode_tile_overlap","shape":7,"type":"INT","widget":{"name":"decode_tile_overlap"},"link":null},{"localized_name":"tile_debug","name":"tile_debug","shape":7,"type":"COMBO","widget":{"name":"tile_debug"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"cache_model","name":"cache_model","shape":7,"type":"BOOLEAN","widget":{"name":"cache_model"},"link":null}],"outputs":[{"localized_name":"SEEDVR2_VAE","name":"SEEDVR2_VAE","type":"SEEDVR2_VAE","links":[10]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2LoadVAEModel"},"widgets_values":["ema_vae_fp16.safetensors","cuda:0",true,1024,128,true,768,128,"false","cpu",false]},{"id":10,"type":"SeedVR2VideoUpscaler","pos":[840.6414062499999,104.58710937499997],"size":[270,386],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":16},{"localized_name":"dit","name":"dit","type":"SEEDVR2_DIT","link":11},{"localized_name":"vae","name":"vae","type":"SEEDVR2_VAE","link":10},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"localized_name":"resolution","name":"resolution","type":"INT","widget":{"name":"resolution"},"link":null},{"localized_name":"max_resolution","name":"max_resolution","type":"INT","widget":{"name":"max_resolution"},"link":null},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null},{"localized_name":"uniform_batch_size","name":"uniform_batch_size","type":"BOOLEAN","widget":{"name":"uniform_batch_size"},"link":null},{"localized_name":"color_correction","name":"color_correction","type":"COMBO","widget":{"name":"color_correction"},"link":null},{"localized_name":"temporal_overlap","name":"temporal_overlap","shape":7,"type":"INT","widget":{"name":"temporal_overlap"},"link":null},{"localized_name":"prepend_frames","name":"prepend_frames","shape":7,"type":"INT","widget":{"name":"prepend_frames"},"link":null},{"localized_name":"input_noise_scale","name":"input_noise_scale","shape":7,"type":"FLOAT","widget":{"name":"input_noise_scale"},"link":null},{"localized_name":"latent_noise_scale","name":"latent_noise_scale","shape":7,"type":"FLOAT","widget":{"name":"latent_noise_scale"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"enable_debug","name":"enable_debug","shape":7,"type":"BOOLEAN","widget":{"name":"enable_debug"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[22]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2VideoUpscaler"},"widgets_values":[42,"fixed",1080,0,33,true,"lab",3,0,0,0,"cpu",false]},{"id":14,"type":"SeedVR2LoadDiTModel","pos":[61.165625,125.97187500000004],"size":[307.6646484375,202],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"torch_compile_args","name":"torch_compile_args","shape":7,"type":"TORCH_COMPILE_ARGS","link":18},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"blocks_to_swap","name":"blocks_to_swap","shape":7,"type":"INT","widget":{"name":"blocks_to_swap"},"link":null},{"localized_name":"swap_io_components","name":"swap_io_components","shape":7,"type":"BOOLEAN","widget":{"name":"swap_io_components"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"cache_model","name":"cache_model","shape":7,"type":"BOOLEAN","widget":{"name":"cache_model"},"link":null},{"localized_name":"attention_mode","name":"attention_mode","shape":7,"type":"COMBO","widget":{"name":"attention_mode"},"link":null}],"outputs":[{"localized_name":"SEEDVR2_DIT","name":"SEEDVR2_DIT","type":"SEEDVR2_DIT","links":[11]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2LoadDiTModel"},"widgets_values":["seedvr2_ema_3b_fp16.safetensors","cuda:0",32,false,"cpu",false,"sdpa"]}],"links":[[10,13,0,10,2,"SEEDVR2_VAE"],[11,14,0,10,1,"SEEDVR2_DIT"],[16,17,0,10,0,"IMAGE"],[18,19,0,14,0,"TORCH_COMPILE_ARGS"],[19,19,0,13,0,"TORCH_COMPILE_ARGS"],[20,21,0,22,0,"VIDEO"],[21,22,0,17,0,"IMAGE"],[22,10,0,24,0,"IMAGE"],[23,24,0,23,0,"VIDEO"],[24,22,2,24,2,"FLOAT"],[25,22,1,24,1,"AUDIO"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.8441465687687482,"offset":[739.339356050958,625.120941260255]}},"version":0.4} \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.jpg b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.jpg new file mode 100644 index 0000000000000000000000000000000000000000..510101e3a02623d6d139831c95596521cddc1178 Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.jpg differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.json b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.json new file mode 100644 index 0000000000000000000000000000000000000000..adfbeb4f2dfef021a3fc960af58dc4d700ca0143 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/SeedVR2_simple_image_upscale.json @@ -0,0 +1 @@ +{"id":"e2c8e6ec-9e8a-4a81-a240-7e7d2befc2bf","revision":0,"last_node_id":20,"last_link_id":19,"nodes":[{"id":16,"type":"LoadImage","pos":[529.6515751584279,98.66389637266899],"size":[274.080078125,314],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[14]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":[17]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":17,"type":"JoinImageWithAlpha","pos":[867.7966831594896,98.31620737975611],"size":[176.86484375,46],"flags":{},"order":4,"mode":4,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":14},{"localized_name":"alpha","name":"alpha","type":"MASK","link":17}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[16]}],"properties":{"cnr_id":"comfy-core","ver":"0.3.68","Node name for S&R":"JoinImageWithAlpha"},"widgets_values":[]},{"id":19,"type":"SeedVR2TorchCompileSettings","pos":[82.82356467195704,682.7639698331246],"size":[307.671484375,178],"flags":{},"order":1,"mode":4,"inputs":[{"localized_name":"backend","name":"backend","type":"COMBO","widget":{"name":"backend"},"link":null},{"localized_name":"mode","name":"mode","type":"COMBO","widget":{"name":"mode"},"link":null},{"localized_name":"fullgraph","name":"fullgraph","type":"BOOLEAN","widget":{"name":"fullgraph"},"link":null},{"localized_name":"dynamic","name":"dynamic","type":"BOOLEAN","widget":{"name":"dynamic"},"link":null},{"localized_name":"dynamo_cache_size_limit","name":"dynamo_cache_size_limit","type":"INT","widget":{"name":"dynamo_cache_size_limit"},"link":null},{"localized_name":"dynamo_recompile_limit","name":"dynamo_recompile_limit","type":"INT","widget":{"name":"dynamo_recompile_limit"},"link":null}],"outputs":[{"localized_name":"TORCH_COMPILE_ARGS","name":"TORCH_COMPILE_ARGS","type":"TORCH_COMPILE_ARGS","links":[18,19]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2TorchCompileSettings"},"widgets_values":["inductor","default",false,false,64,128]},{"id":13,"type":"SeedVR2LoadVAEModel","pos":[509.41407515842775,780.813505747667],"size":[312.866796875,298],"flags":{},"order":6,"mode":0,"inputs":[{"localized_name":"torch_compile_args","name":"torch_compile_args","shape":7,"type":"TORCH_COMPILE_ARGS","link":19},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"encode_tiled","name":"encode_tiled","shape":7,"type":"BOOLEAN","widget":{"name":"encode_tiled"},"link":null},{"localized_name":"encode_tile_size","name":"encode_tile_size","shape":7,"type":"INT","widget":{"name":"encode_tile_size"},"link":null},{"localized_name":"encode_tile_overlap","name":"encode_tile_overlap","shape":7,"type":"INT","widget":{"name":"encode_tile_overlap"},"link":null},{"localized_name":"decode_tiled","name":"decode_tiled","shape":7,"type":"BOOLEAN","widget":{"name":"decode_tiled"},"link":null},{"localized_name":"decode_tile_size","name":"decode_tile_size","shape":7,"type":"INT","widget":{"name":"decode_tile_size"},"link":null},{"localized_name":"decode_tile_overlap","name":"decode_tile_overlap","shape":7,"type":"INT","widget":{"name":"decode_tile_overlap"},"link":null},{"localized_name":"tile_debug","name":"tile_debug","shape":7,"type":"COMBO","widget":{"name":"tile_debug"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"cache_model","name":"cache_model","shape":7,"type":"BOOLEAN","widget":{"name":"cache_model"},"link":null}],"outputs":[{"localized_name":"SEEDVR2_VAE","name":"SEEDVR2_VAE","type":"SEEDVR2_VAE","links":[10]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2LoadVAEModel"},"widgets_values":["ema_vae_fp16.safetensors","cuda:0",false,1024,128,false,1024,128,"false","none",false]},{"id":14,"type":"SeedVR2LoadDiTModel","pos":[513.4976689084277,493.64084949766976],"size":[307.6646484375,202],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"torch_compile_args","name":"torch_compile_args","shape":7,"type":"TORCH_COMPILE_ARGS","link":18},{"localized_name":"model","name":"model","type":"COMBO","widget":{"name":"model"},"link":null},{"localized_name":"device","name":"device","type":"COMBO","widget":{"name":"device"},"link":null},{"localized_name":"blocks_to_swap","name":"blocks_to_swap","shape":7,"type":"INT","widget":{"name":"blocks_to_swap"},"link":null},{"localized_name":"swap_io_components","name":"swap_io_components","shape":7,"type":"BOOLEAN","widget":{"name":"swap_io_components"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"cache_model","name":"cache_model","shape":7,"type":"BOOLEAN","widget":{"name":"cache_model"},"link":null},{"localized_name":"attention_mode","name":"attention_mode","shape":7,"type":"COMBO","widget":{"name":"attention_mode"},"link":null}],"outputs":[{"localized_name":"SEEDVR2_DIT","name":"SEEDVR2_DIT","type":"SEEDVR2_DIT","links":[11]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2LoadDiTModel"},"widgets_values":["seedvr2_ema_3b_fp8_e4m3fn.safetensors","cuda:0",0,false,"none",false,"sdpa"]},{"id":18,"type":"Note","pos":[869.7890553931869,199.9227829276786],"size":[210,94.05179298404067],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Enable to upscale alpha/mask channel along with RGB channel (for RGBA inputs)"],"color":"#432","bgcolor":"#653"},{"id":20,"type":"Note","pos":[93.33964800519301,920.9877545042978],"size":[290.1922351462714,97.95592567265356],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"properties":{},"widgets_values":["Enable torch.compile optimization. First batch has compilation overhead, but speeds up subsequent batches. Best for long videos, not single images."],"color":"#432","bgcolor":"#653"},{"id":10,"type":"SeedVR2VideoUpscaler","pos":[1292.9734501584262,472.2560838726697],"size":[270,386],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"IMAGE","link":16},{"localized_name":"dit","name":"dit","type":"SEEDVR2_DIT","link":11},{"localized_name":"vae","name":"vae","type":"SEEDVR2_VAE","link":10},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null},{"localized_name":"resolution","name":"resolution","type":"INT","widget":{"name":"resolution"},"link":null},{"localized_name":"max_resolution","name":"max_resolution","type":"INT","widget":{"name":"max_resolution"},"link":null},{"localized_name":"batch_size","name":"batch_size","type":"INT","widget":{"name":"batch_size"},"link":null},{"localized_name":"uniform_batch_size","name":"uniform_batch_size","type":"BOOLEAN","widget":{"name":"uniform_batch_size"},"link":null},{"localized_name":"color_correction","name":"color_correction","type":"COMBO","widget":{"name":"color_correction"},"link":null},{"localized_name":"temporal_overlap","name":"temporal_overlap","shape":7,"type":"INT","widget":{"name":"temporal_overlap"},"link":null},{"localized_name":"prepend_frames","name":"prepend_frames","shape":7,"type":"INT","widget":{"name":"prepend_frames"},"link":null},{"localized_name":"input_noise_scale","name":"input_noise_scale","shape":7,"type":"FLOAT","widget":{"name":"input_noise_scale"},"link":null},{"localized_name":"latent_noise_scale","name":"latent_noise_scale","shape":7,"type":"FLOAT","widget":{"name":"latent_noise_scale"},"link":null},{"localized_name":"offload_device","name":"offload_device","shape":7,"type":"COMBO","widget":{"name":"offload_device"},"link":null},{"localized_name":"enable_debug","name":"enable_debug","shape":7,"type":"BOOLEAN","widget":{"name":"enable_debug"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[12]}],"properties":{"aux_id":"ainvfx/ComfyUI-SeedVR2_VideoUpscaler","ver":"690cc39379c1481159ddd451368dbf2295930fc6","Node name for S&R":"SeedVR2VideoUpscaler"},"widgets_values":[42,"randomize",1080,0,1,false,"lab",0,0,0,0,"cpu",false]},{"id":15,"type":"SaveImage","pos":[1616.797668908428,473.0256151226697],"size":[270,270],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":12},{"localized_name":"filename_prefix","name":"filename_prefix","type":"STRING","widget":{"name":"filename_prefix"},"link":null}],"outputs":[],"properties":{"cnr_id":"comfy-core","ver":"0.3.68"},"widgets_values":["ComfyUI"]}],"links":[[10,13,0,10,2,"SEEDVR2_VAE"],[11,14,0,10,1,"SEEDVR2_DIT"],[12,10,0,15,0,"IMAGE"],[14,16,0,17,0,"IMAGE"],[16,17,0,10,0,"IMAGE"],[17,16,1,17,1,"MASK"],[18,19,0,14,0,"TORCH_COMPILE_ARGS"],[19,19,0,13,0,"TORCH_COMPILE_ARGS"]],"groups":[],"config":{},"extra":{"ds":{"scale":0.9229599817706443,"offset":[268.65279462308837,183.8749573202125]}},"version":0.4} \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Eyes_212x120.mp4 b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Eyes_212x120.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..eb87d4522680659f2b0e9dbcde63fe23350f09e5 Binary files /dev/null and b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Eyes_212x120.mp4 differ diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Mustache_640x360.mp4 b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Mustache_640x360.mp4 new file mode 100644 index 0000000000000000000000000000000000000000..3b6d61890056c867003f2b7ba2be40406124b903 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Mustache_640x360.mp4 @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:21581bc8454e234e3d3f833172bb358215f47d33cfcf5fc12f3dd0dff3319d1d +size 115496 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Sadhu_320x478.png b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Sadhu_320x478.png new file mode 100644 index 0000000000000000000000000000000000000000..e593f5104bbc3392522378bc09dc2c8a5639917f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/example_workflows/example_inputs/Sadhu_320x478.png @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b7faf22b8911a49742b1577b461f2187e78942d5f15c3998254edadfee970084 +size 417158 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/inference_cli.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/inference_cli.py new file mode 100644 index 0000000000000000000000000000000000000000..2d4fff18667e12a7d35ec8a43b952820a6b73ee9 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/inference_cli.py @@ -0,0 +1,1712 @@ +#!/usr/bin/env python3 +""" +SeedVR2 Video Upscaler - Standalone CLI Interface + +Command-line interface for high-quality upscaling using SeedVR2 diffusion models. +Supports single and multi-GPU processing with advanced memory optimization. + +Key Features: + โ€ข Multi-GPU Processing: Automatic workload distribution across multiple GPUs with + temporal overlap blending for seamless transitions + โ€ข Streaming Mode: Memory-efficient processing of long videos in chunks, avoiding + full video loading into RAM while maintaining temporal consistency + โ€ข Memory Optimization: BlockSwap for limited VRAM, VAE tiling for large resolutions, + intelligent tensor offloading between processing phases + โ€ข Performance: Torch.compile integration, BFloat16 compute pipeline, + efficient model caching for batch and streaming processing + โ€ข Flexibility: Multiple output formats (MP4/PNG), advanced color correction methods, + directory batch processing with auto-format detection + โ€ข Quality Control: Temporal overlap blending, frame prepending for artifact reduction, + configurable noise scales for detail preservation + +Architecture: + The CLI implements a 4-phase processing pipeline: + 1. Encode: VAE encoding with optional input noise and tiling + 2. Upscale: DiT transformer upscaling with latent space diffusion + 3. Decode: VAE decoding with optional tiling + 4. Postprocess: Color correction and temporal blending + +Usage: + python inference_cli.py video.mp4 --resolution 1080 + For complete usage examples, run: python inference_cli.py --help + +Requirements: + โ€ข Python 3.10+ + โ€ข PyTorch 2.4+ with CUDA 12.1+ (NVIDIA) or MPS (Apple Silicon) + โ€ข 16GB+ VRAM recommended (8GB minimum with BlockSwap) + โ€ข OpenCV, NumPy for video I/O + +Model Support: + โ€ข 3B models: seedvr2_ema_3b_fp16.safetensors (default), _fp8_e4m3fn/GGUF variants + โ€ข 7B models: seedvr2_ema_7b_fp16.safetensors, _fp8_e4m3fn/GGUF variants + โ€ข VAE: ema_vae_fp16.safetensors (shared across all models) + โ€ข Auto-downloads from HuggingFace on first run with SHA256 validation +""" + +# Standard library imports +import sys +import os +import argparse +import time +import platform +import multiprocessing as mp +from typing import Dict, Any, List, Optional, Tuple, Literal, Generator +from datetime import datetime +from pathlib import Path + +# Set up path before any other imports to fix module resolution +script_dir = os.path.dirname(os.path.abspath(__file__)) +if script_dir not in sys.path: + sys.path.insert(0, script_dir) + +# Set environment variable so all spawned processes can find modules +os.environ['PYTHONPATH'] = script_dir + ':' + os.environ.get('PYTHONPATH', '') + +# Ensure safe CUDA usage with multiprocessing +if mp.get_start_method(allow_none=True) != 'spawn': + mp.set_start_method('spawn', force=True) + +# Configure platform-specific memory management before heavy imports +# Must be set BEFORE import torch +if platform.system() == "Darwin": + # MPS allocator requires: low_watermark <= high_watermark + # Setting both to 0.0 disables PyTorch memory limits, letting macOS manage memory + os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0") + os.environ.setdefault("PYTORCH_MPS_LOW_WATERMARK_RATIO", "0.0") +else: + os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync") + + # Pre-parse arguments that must be handled before torch import + _pre_parser = argparse.ArgumentParser(add_help=False) + _pre_parser.add_argument("--cuda_device", type=str, default=None) + _pre_args, _ = _pre_parser.parse_known_args() + + if _pre_args.cuda_device is not None: + device_list_env = [x.strip() for x in _pre_args.cuda_device.split(',') if x.strip()!=''] + + # Skip validation if CUDA_VISIBLE_DEVICES is already set (worker process) + if os.environ.get("CUDA_VISIBLE_DEVICES") is None: + # Temporary torch import for CUDA device validation only + # Must happen before setting CUDA_VISIBLE_DEVICES and before main torch import + import torch as _torch_check + if _torch_check.cuda.is_available(): + available_count = _torch_check.cuda.device_count() + invalid_devices = [d for d in device_list_env if not d.isdigit() or int(d) >= available_count] + if invalid_devices: + print(f"โŒ [ERROR] Invalid CUDA device ID(s): {', '.join(invalid_devices)}. " + f"Available devices: 0-{available_count-1} (total: {available_count})") + sys.exit(1) + else: + print("โŒ [ERROR] CUDA is not available on this system. Cannot use --cuda_device argument.") + sys.exit(1) + + # Set CUDA_VISIBLE_DEVICES for single GPU after validation + if len(device_list_env) == 1: + os.environ["CUDA_VISIBLE_DEVICES"] = device_list_env[0] + +# Heavy dependency imports after environment configuration +import torch +import cv2 +import numpy as np +import subprocess +import shutil + +# Project imports +from src.utils.downloads import download_weight +from src.utils.model_registry import get_available_dit_models, DEFAULT_DIT, DEFAULT_VAE +from src.utils.constants import SEEDVR2_FOLDER_NAME +from src.core.generation_utils import ( + setup_generation_context, + prepare_runner, + compute_generation_info, + log_generation_start, + blend_overlapping_frames, + load_text_embeddings, + script_directory +) +from src.core.generation_phases import ( + encode_all_batches, + upscale_all_batches, + decode_all_batches, + postprocess_all_batches +) +from src.utils.debug import Debug +from src.optimization.memory_manager import clear_memory, get_gpu_backend, is_cuda_available +debug = Debug(enabled=False) # Will be enabled via --debug CLI flag + + +# ============================================================================= +# FFMPEG Class +# ============================================================================= + +class FFMPEGVideoWriter: + """ + Video writer using ffmpeg subprocess for encoding with 10-bit support. + + Provides cv2.VideoWriter-compatible interface (write, isOpened, release) while + using ffmpeg for encoding. Enables 10-bit output (yuv420p10le with x265) which + reduces banding artifacts in gradients compared to 8-bit opencv output. + + Args: + path: Output video file path + width: Frame width in pixels + height: Frame height in pixels + fps: Frames per second + use_10bit: If True, uses x265 codec with yuv420p10le pixel format. + If False, uses x264 with yuv420p (default: False) + + Raises: + RuntimeError: If ffmpeg is not found in system PATH + + Note: + Frames must be passed to write() in BGR format (same as cv2.VideoWriter). + Internally converts to RGB for ffmpeg rawvideo input. + """ + + def __init__(self, path: str, width: int, height: int, fps: float, use_10bit: bool = False): + pix_fmt = 'yuv420p10le' if use_10bit else 'yuv420p' + codec = 'libx265' if use_10bit else 'libx264' + + self.proc = subprocess.Popen( + ['ffmpeg', '-y', '-f', 'rawvideo', '-pix_fmt', 'rgb24', + '-s', f'{width}x{height}', '-r', str(fps), '-i', '-', + '-c:v', codec, '-pix_fmt', pix_fmt, '-preset', 'medium', '-crf', '12', path], + stdin=subprocess.PIPE, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL + ) + + def write(self, frame_bgr: np.ndarray): + if not self.isOpened(): + raise RuntimeError("FFMPEGVideoWriter: ffmpeg process is not running") + + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + try: + self.proc.stdin.write(frame_rgb.astype(np.uint8).tobytes()) + self.proc.stdin.flush() # Critical: prevent buffering issues + except BrokenPipeError: + raise RuntimeError( + "FFMPEGVideoWriter: ffmpeg process terminated unexpectedly. " + "Check video path, codec support, and disk space." + ) + + def isOpened(self) -> bool: + return self.proc is not None and self.proc.poll() is None + + def release(self): + if self.proc: + try: + self.proc.stdin.close() + except Exception: + pass # Ignore errors on close + + self.proc.wait() + + if self.proc.returncode != 0: + debug.log( + f"ffmpeg exited with code {self.proc.returncode}. " + "Check output file for corruption.", + level="WARNING", force=True, category="file" + ) + self.proc = None + + +# ============================================================================= +# Device Management Helpers +# ============================================================================= + +def _device_id_to_name(device_id: str, platform_type: str = None) -> str: + """ + Convert device ID to full device name. + + Args: + device_id: Device ID ("0", "1") or special value ("cpu", "none") + platform_type: Override platform type ("cuda", "mps", "cpu") + + Returns: + Full device name ("cuda:0", "mps:0", "cpu", "none") + """ + if device_id in ("cpu", "none"): + return device_id + + if platform_type is None: + platform_type = get_gpu_backend() + + # MPS typically doesn't use indices + if platform_type == "mps": + return "mps" + + return f"{platform_type}:{device_id}" + + +def _parse_offload_device(offload_arg: str, platform_type: str = None, cache_enabled: bool = False) -> Optional[str]: + """ + Parse offload device argument to full device name. + + Args: + offload_arg: Offload device argument ("none", "cpu", "0", "1", or "cuda:1") + platform_type: Override platform type + cache_enabled: If True and offload_arg is "none", default to "cpu" + + Returns: + Full device name or None + """ + if offload_arg == "none": + # If caching enabled but no offload device specified, default to CPU + return "cpu" if cache_enabled else None + + if offload_arg == "cpu": + return "cpu" + + # If already a full device name (cuda:1, mps:0), return as-is + if ":" in offload_arg: + return offload_arg + + # Otherwise treat as device ID + return _device_id_to_name(offload_arg, platform_type) + + +# ============================================================================= +# Constants +# ============================================================================= + +# Supported file extensions +VIDEO_EXTENSIONS = {'.mp4', '.avi', '.mov', '.mkv', '.webm', '.flv', '.wmv', '.m4v'} +IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.bmp', '.tiff', '.tif', '.webp'} + + +# ============================================================================= +# Video I/O Functions +# ============================================================================= + +def get_media_files(directory: str) -> List[str]: + """ + Get all video and image files from directory, sorted alphabetically. + + Args: + directory: Path to directory to scan + + Returns: + Sorted list of file paths (strings) matching video or image extensions + """ + valid_extensions = VIDEO_EXTENSIONS | IMAGE_EXTENSIONS + path = Path(directory) + + # Get all files and filter by extension (case-insensitive) + files = [f for f in path.iterdir() if f.is_file() and f.suffix.lower() in valid_extensions] + + return sorted([str(f) for f in files]) + + +def extract_frames_from_image(image_path: str) -> Tuple[torch.Tensor, float]: + """ + Extract single frame from image file and convert to tensor format. + + Reads image using OpenCV, converts BGR to RGB, normalizes to [0,1] range, + and formats as single-frame video tensor for consistent processing. + + Args: + image_path: Path to input image file + + Returns: + Tuple containing: + - frames_tensor: Single frame as tensor [1, H, W, C], Float16, range [0,1] (C=3 for RGB, C=4 for RGBA) + - fps: Default FPS value (30.0) for image-to-video conversion + + Raises: + FileNotFoundError: If image file doesn't exist + ValueError: If image cannot be opened + """ + debug.log(f"Loading image: {image_path}", category="file") + + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + + # Read image with alpha channel preserved + frame = cv2.imread(image_path, cv2.IMREAD_UNCHANGED) + if frame is None: + raise ValueError(f"Cannot open image file: {image_path}") + + # Convert BGR(A) to RGB(A) based on channel count + if frame.shape[2] == 4: + frame = cv2.cvtColor(frame, cv2.COLOR_BGRA2RGBA) + debug.log(f"Detected RGBA image (alpha channel preserved)", category="file") + else: + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + + # Convert to float32 and normalize + frame = frame.astype(np.float32) / 255.0 + + # Convert to tensor [1, H, W, C] + frames_tensor = torch.from_numpy(frame[None, ...]).to(torch.float16) + + debug.log(f"Image tensor shape: {frames_tensor.shape}, dtype: {frames_tensor.dtype}", category="memory") + + return frames_tensor, 30.0 # Default FPS for images + + +def get_input_type(input_path: str) -> Literal['video', 'image', 'directory', 'unknown']: + """ + Determine input type from file path. + + Args: + input_path: Path to input file or directory + + Returns: + Input type: 'video', 'image', 'directory', or 'unknown' + + Raises: + FileNotFoundError: If input path doesn't exist + """ + path = Path(input_path) + + if not path.exists(): + raise FileNotFoundError(f"Input path not found: {input_path}") + + if path.is_dir(): + return 'directory' + + ext = path.suffix.lower() + if ext in VIDEO_EXTENSIONS: + return "video" + elif ext in IMAGE_EXTENSIONS: + return "image" + else: + return "unknown" + + +def generate_output_path(input_path: str, output_format: str, output_dir: Optional[str] = None, + input_type: Optional[str] = None, from_directory: bool = False) -> str: + """ + Generate output path based on input path and format. + + Args: + input_path: Source file path + output_format: "mp4" or "png" + output_dir: Optional output directory (overrides default behavior) + input_type: Optional input type ("image", "video", "directory") + from_directory: True if processing files from a directory (batch mode) + + Returns: + Absolute output path (file for single image/video, directory for sequences) + """ + input_path_obj = Path(input_path) + input_name = input_path_obj.stem + + # Determine base directory and whether to add suffix + if output_dir: + # User specified output directory - use as-is, no suffix + base_dir = Path(output_dir) + add_suffix = False + elif from_directory: + # Batch mode: create sibling folder with _upscaled, keep original filenames + original_dir = input_path_obj.parent + base_dir = original_dir.parent / f"{original_dir.name}_upscaled" + add_suffix = False + else: + # Single file mode: output to same directory with _upscaled suffix + base_dir = input_path_obj.parent + add_suffix = True + + # Build filename with optional suffix + file_suffix = "_upscaled" if add_suffix else "" + + # Generate output path based on format + if output_format == "png": + if input_type == "image": + output_path = base_dir / f"{input_name}{file_suffix}.png" + else: + output_path = base_dir / f"{input_name}{file_suffix}" + else: + output_path = base_dir / f"{input_name}{file_suffix}.mp4" + + return str(output_path.resolve()) + + +def process_single_file(input_path: str, args: argparse.Namespace, device_list: List[str], + output_path: Optional[str] = None, format_auto_detected: bool = False, + runner_cache: Optional[Dict[str, Any]] = None) -> int: + """ + Process a single video or image file with optional model caching. + + For videos, supports streaming mode (chunk_size > 0) which processes in memory-bounded + chunks with temporal overlap for seamless transitions between chunks. + + Args: + input_path: Path to input file + args: Command-line arguments with all processing settings + device_list: List of GPU device IDs as strings + output_path: Optional explicit output path (auto-generated if None) + format_auto_detected: Whether output format was auto-detected + runner_cache: Optional cache dict for model reuse across multiple files + + Returns: + Number of frames written to output + """ + input_type = get_input_type(input_path) + + if input_type == "unknown": + debug.log(f"Skipping unsupported file: {input_path}", level="WARNING", category="file", force=True) + return 0 + + debug.log(f"Processing {input_type}: {Path(input_path).name}", category="generation", force=True) + + # Generate or validate output path + if output_path is None: + output_path = generate_output_path(input_path, args.output_format, input_type=input_type) + elif not Path(output_path).suffix or (args.output_format == "png" and input_type != "image"): + # No extension or PNG sequence โ†’ treat as directory, generate filename + output_path = generate_output_path(input_path, args.output_format, + output_dir=output_path, input_type=input_type) + + # Show format with auto-detection indicator + format_prefix = "Auto-detected" if format_auto_detected else "Requested" + debug.log(f"{format_prefix} output format: {args.output_format}", category="info", force=True, indent_level=1) + + # === VIDEO PROCESSING === + if input_type == "video": + if not os.path.exists(input_path): + raise FileNotFoundError(f"Video file not found: {input_path}") + + cap = cv2.VideoCapture(input_path) + if not cap.isOpened(): + raise ValueError(f"Cannot open video file: {input_path}") + + fps = cap.get(cv2.CAP_PROP_FPS) or 30.0 + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + debug.log(f"Video info: {total_frames} frames, {width}x{height}, {fps:.2f} FPS", category="info") + + # Skip initial frames + if args.skip_first_frames > 0: + debug.log(f"Skipping first {args.skip_first_frames} frames", category="info") + cap.set(cv2.CAP_PROP_POS_FRAMES, args.skip_first_frames) + + # Calculate frames to process (apply load_cap if set) + frames_to_process = total_frames - args.skip_first_frames + if args.load_cap > 0: + frames_to_process = min(frames_to_process, args.load_cap) + + # Early exit for empty/exhausted video + if frames_to_process <= 0: + debug.log(f"No frames to process after skipping {args.skip_first_frames} of {total_frames}", + level="WARNING", category="file", force=True) + cap.release() + return 0 + + # Streaming mode: process in chunks + chunk_size = args.chunk_size if args.chunk_size > 0 else frames_to_process + streaming = args.chunk_size > 0 + total_chunks = (frames_to_process + chunk_size - 1) // chunk_size # ceiling division + + if streaming: + debug.log(f"Streaming mode: chunks of {chunk_size} frames, overlap={args.temporal_overlap}", + category="info", force=True, indent_level=1) + + is_png = args.output_format == "png" + video_writer = None + overlap = args.temporal_overlap + frames_written = 0 + chunk_idx = 0 + base_name = Path(input_path).stem + + # Multi-GPU: workers stream their own segments + if len(device_list) > 1: + cap.release() # Workers will reopen + video_info = { + 'video_path': input_path, + 'start_frame': args.skip_first_frames, + 'frames_to_process': frames_to_process, + } + result = _gpu_processing(None, device_list, args, video_info=video_info) + + # Save result + if is_png: + save_frames_to_image(result, output_path, base_name) + else: + video_writer = save_frames_to_video(result, output_path, fps, + video_backend=args.video_backend, use_10bit=args.use_10bit) + if video_writer is not None: + video_writer.release() + + frames_written = result.shape[0] + + # Single GPU: stream in main process + else: + chunk_count = 0 + for result in _stream_video_chunks( + cap=cap, + frames_to_process=frames_to_process, + chunk_size=chunk_size, + overlap=overlap, + args=args, + device_id=device_list[0], + debug=debug, + runner_cache=runner_cache, + log_progress=streaming, + total_chunks=total_chunks, + cleanup_timer_name="chunk_cleanup" + ): + chunk_count += 1 + + # Save output + if is_png: + save_frames_to_image(result, output_path, base_name, start_index=frames_written) + else: + video_writer = save_frames_to_video(result, output_path, fps, writer=video_writer, + video_backend=args.video_backend, use_10bit=args.use_10bit) + + frames_written += result.shape[0] + del result + + chunk_idx = chunk_count + cap.release() + if video_writer is not None: + video_writer.release() + + if streaming: + debug.log("", category="none", force=True) + if len(device_list) > 1: + debug.log(f"Streaming complete: {frames_written} frames across {len(device_list)} GPUs", category="success", force=True) + else: + debug.log(f"Streaming complete: {frames_written} frames in {chunk_idx} chunks", category="success", force=True) + + debug.log(f"Output saved to: {output_path}", category="file", force=True) + return frames_written + + # === IMAGE PROCESSING === + frames_tensor, _ = extract_frames_from_image(input_path) + + processing_start = time.time() + # Process frames (multiprocessing only for multi-GPU) + if len(device_list) > 1: + result = _gpu_processing(frames_tensor, device_list, args) + else: + result = _single_gpu_direct_processing(frames_tensor, args, device_list[0], runner_cache) + debug.log(f"Processing time: {time.time() - processing_start:.2f}s", category="timing") + + # Save single image + os.makedirs(Path(output_path).parent, exist_ok=True) + frame_np = (result[0].cpu().numpy() * 255.0).astype(np.uint8) + _save_image_bgr(frame_np, output_path) + + debug.log(f"Output saved to: {output_path}", category="file", force=True) + return 1 + + +def _read_frames_from_cap(cap: cv2.VideoCapture, max_frames: int) -> Optional[torch.Tensor]: + """ + Read up to max_frames from an already-open VideoCapture. + + Args: + cap: An already opened cv2.VideoCapture instance + max_frames: Maximum number of frames to read in this call + + Returns: + Tensor [T, H, W, C] float32 [0,1], or None if no frames available + """ + frames = [] + for _ in range(max_frames): + ret, frame = cap.read() + if not ret: + break + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0 + frames.append(frame) + + if not frames: + return None + return torch.from_numpy(np.stack(frames)).to(torch.float32) + + +def _stream_video_chunks( + cap: cv2.VideoCapture, + frames_to_process: int, + chunk_size: int, + overlap: int, + args: argparse.Namespace, + device_id: str, + debug: 'Debug', + runner_cache: Optional[Dict[str, Any]], + log_progress: bool = False, + total_chunks: int = 0, + cleanup_timer_name: Optional[str] = None, + log_prefix: str = "" +) -> Generator[torch.Tensor, None, None]: + """ + Generator that streams and processes video chunks. + + Handles frame reading, temporal context prepending, processing via + _process_frames_core, context removal from output, and memory cleanup. + Caller is responsible for VideoCapture lifecycle and result handling. + + Args: + cap: Open VideoCapture positioned at start frame + frames_to_process: Total frames to read and process + chunk_size: Frames per chunk (use frames_to_process for single chunk) + overlap: Temporal overlap frames between chunks for blending + args: Processing arguments (copied internally, prepend_frames zeroed after first chunk) + device_id: GPU device ID for processing + debug: Debug instance for logging + runner_cache: Optional model cache dict for reuse across chunks + log_progress: If True, log chunk progress with separators + total_chunks: Total chunks for progress display (used if log_progress=True) + cleanup_timer_name: Optional timer name for memory cleanup logging + log_prefix: Optional prefix for log messages (e.g., "[GPU 0] " for worker identification) + + Yields: + Processed frames tensor [T, H, W, C] for each chunk, context frames removed + """ + chunk_args = argparse.Namespace(**vars(args)) + frames_read = 0 + prev_raw_tail = None + chunk_idx = 0 + streaming = chunk_size < frames_to_process + + while frames_read < frames_to_process: + read_count = min(chunk_size, frames_to_process - frames_read) + new_frames = _read_frames_from_cap(cap, read_count) + if new_frames is None: + break + frames_read += new_frames.shape[0] + chunk_idx += 1 + + # Disable prepend_frames after first chunk + if chunk_idx > 1: + chunk_args.prepend_frames = 0 + + # Prepend context from previous chunk + if prev_raw_tail is not None and overlap > 0: + context_count = min(overlap, prev_raw_tail.shape[0]) + frames = torch.cat([prev_raw_tail[-context_count:], new_frames], dim=0) + else: + frames = new_frames + context_count = 0 + + # Log progress if enabled + if log_progress and streaming: + if chunk_idx > 1: + debug.log("", category="none", force=True) + debug.log("โ”" * 60, category="none", force=True) + debug.log("", category="none", force=True) + debug.log(f"{log_prefix}Chunk {chunk_idx}/{total_chunks}: {new_frames.shape[0]} new + {context_count} context frames", + category="generation", force=True) + debug.log("", category="none", force=True) + + # Process chunk + result = _process_frames_core( + frames_tensor=frames.to(torch.float16), + args=chunk_args, + device_id=device_id, + debug=debug, + runner_cache=runner_cache + ) + + # Remove context frames from output + if context_count > 0: + result = result[context_count:] + + # Save tail for next chunk context + prev_raw_tail = new_frames[-overlap:].clone() if overlap > 0 else None + + # Cleanup before yield + del frames + + yield result + + # Memory cleanup between chunks + if streaming: + clear_memory(debug=debug, deep=True, force=True, timer_name=cleanup_timer_name) + + +def _save_image_bgr(frame_np: np.ndarray, file_path: str) -> None: + """ + Save a single RGB(A) uint8 frame to disk, converting to BGR(A) for OpenCV. + + Args: + frame_np: Frame as uint8 numpy array [H, W, C] where C is 3 (RGB) or 4 (RGBA) + file_path: Output file path + """ + if frame_np.shape[2] == 4: + frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGBA2BGRA) + else: + frame_bgr = cv2.cvtColor(frame_np, cv2.COLOR_RGB2BGR) + cv2.imwrite(file_path, frame_bgr) + + +def save_frames_to_video( + frames_tensor: torch.Tensor, + output_path: str, + fps: float = 30.0, + writer: Optional[cv2.VideoWriter] = None, + video_backend: str = "opencv", + use_10bit: bool = False +) -> Optional[cv2.VideoWriter]: + """ + Save frames tensor to MP4 video file. + + Converts tensor from Float32 [0,1] to uint8 [0,255], RGB to BGR for OpenCV, + and writes to video file using mp4v codec. Supports streaming mode where + an existing writer is passed and kept open for subsequent chunks. + + Args: + frames_tensor: Frames in format [T, H, W, C], Float32, range [0,1] + output_path: Output video file path (directory created if doesn't exist) + fps: Frames per second for output video (default: 30.0) + writer: Existing VideoWriter for streaming (if None, creates new one) + + Returns: + VideoWriter if streaming mode (caller must close), None if standalone mode + + Raises: + ValueError: If video writer cannot be initialized + """ + frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8) + T, H, W, C = frames_np.shape + + if writer is None: + debug.log(f"Saving {T} frames to video: {output_path} (backend={video_backend})", category="file") + os.makedirs(Path(output_path).parent, exist_ok=True) + if video_backend == "ffmpeg": + writer = FFMPEGVideoWriter(output_path, W, H, fps, use_10bit) + else: + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + writer = cv2.VideoWriter(output_path, fourcc, fps, (W, H)) + if not writer.isOpened(): + raise ValueError(f"Cannot create video writer for: {output_path}") + + for i, frame in enumerate(frames_np): + frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) + writer.write(frame_bgr) + if debug.enabled and (i + 1) % 100 == 0: + debug.log(f"Written {i + 1}/{T} frames", category="file") + + return writer # Caller always closes + + +def save_frames_to_image( + frames_tensor: torch.Tensor, + output_dir: str, + base_name: str, + start_index: int = 0 +) -> int: + """ + Save frames tensor as sequential PNG image files. + + Each frame saved as {base_name}_{index:0Nd}.png with zero-padded indices. + Converts Float32 [0,1] to uint8 [0,255] and RGB(A) to BGR(A) for OpenCV. + + Args: + frames_tensor: Frames in format [T, H, W, C], Float32, range [0,1] + output_dir: Directory to save PNG files (created if doesn't exist) + base_name: Base name for output files (e.g., "frame" โ†’ "frame_00000.png") + start_index: Starting index for filenames (for streaming continuation) + + Returns: + Number of frames saved + """ + os.makedirs(output_dir, exist_ok=True) + + frames_np = (frames_tensor.cpu().numpy() * 255.0).astype(np.uint8) + total = frames_np.shape[0] + + if start_index == 0: + debug.log(f"Saving {total} frames as PNGs to directory: {output_dir}", category="file") + digits = 6 # Supports up to 999,999 frames (~11.5 hours at 24fps) + + for idx, frame in enumerate(frames_np): + filename = f"{base_name}_{start_index + idx:0{digits}d}.png" + file_path = os.path.join(output_dir, filename) + _save_image_bgr(frame, file_path) + if debug.enabled and (idx + 1) % 100 == 0: + debug.log(f"Saved {idx + 1}/{total} images", category="file") + + debug.log(f"Saved {total} images to '{output_dir}'", category="success") + return total + + +# ============================================================================= +# Core Processing Logic +# ============================================================================= + +def _process_frames_core( + frames_tensor: torch.Tensor, + args: argparse.Namespace, + device_id: str, + debug: Debug, + runner_cache: Optional[Dict[str, Any]] = None +) -> torch.Tensor: + """ + Core frame processing logic shared between worker and direct processing. + + Executes the complete 4-phase pipeline: encode โ†’ upscale โ†’ decode โ†’ postprocess. + Supports both cached (direct) and non-cached (worker) execution modes. + + Args: + frames_tensor: Input frames [T, H, W, C], Float16/Float32, range [0,1] + args: Command-line arguments with all processing settings + device_id: Device ID for inference ("0", "1", etc.) + debug: Debug instance for logging + runner_cache: Optional cache dict for model reuse (direct mode only) + + Returns: + Upscaled frames tensor [T', H', W', C], Float32, range [0,1] + """ + # Determine platform and convert device IDs to full names + platform_type = get_gpu_backend() + inference_device = _device_id_to_name(device_id, platform_type) + + # Parse offload devices (with caching defaults) + cache_dit = args.cache_dit if runner_cache is not None else False + cache_vae = args.cache_vae if runner_cache is not None else False + + dit_offload = _parse_offload_device(args.dit_offload_device, platform_type, cache_dit) + vae_offload = _parse_offload_device(args.vae_offload_device, platform_type, cache_vae) + tensor_offload = _parse_offload_device(args.tensor_offload_device, platform_type, False) + + # Setup or reuse generation context + if runner_cache is not None and 'ctx' in runner_cache: + ctx = runner_cache['ctx'] + # Clear previous run data but keep device config + keys_to_keep = {'dit_device', 'vae_device', 'dit_offload_device', + 'vae_offload_device', 'tensor_offload_device', 'compute_dtype'} + for key in list(ctx.keys()): + if key not in keys_to_keep: + del ctx[key] + else: + ctx = setup_generation_context( + dit_device=inference_device, + vae_device=inference_device, + dit_offload_device=dit_offload, + vae_offload_device=vae_offload, + tensor_offload_device=tensor_offload, + debug=debug + ) + if runner_cache is not None: + runner_cache['ctx'] = ctx + + # Build torch compile args + torch_compile_args_dit = None + torch_compile_args_vae = None + if args.compile_dit: + torch_compile_args_dit = { + "backend": args.compile_backend, + "mode": args.compile_mode, + "fullgraph": args.compile_fullgraph, + "dynamic": args.compile_dynamic, + "dynamo_cache_size_limit": args.compile_dynamo_cache_size_limit, + "dynamo_recompile_limit": args.compile_dynamo_recompile_limit, + } + if args.compile_vae: + torch_compile_args_vae = { + "backend": args.compile_backend, + "mode": args.compile_mode, + "fullgraph": args.compile_fullgraph, + "dynamic": args.compile_dynamic, + "dynamo_cache_size_limit": args.compile_dynamo_cache_size_limit, + "dynamo_recompile_limit": args.compile_dynamo_recompile_limit, + } + + # Prepare runner with caching support + model_dir = args.model_dir if args.model_dir is not None else f"./models/{SEEDVR2_FOLDER_NAME}" + + # Use fixed IDs for CLI caching when enabled + dit_id = "cli_dit" if cache_dit else None + vae_id = "cli_vae" if cache_vae else None + + runner, cache_context = prepare_runner( + dit_model=args.dit_model, + vae_model=DEFAULT_VAE, + model_dir=model_dir, + debug=debug, + ctx=ctx, + dit_cache=cache_dit, + vae_cache=cache_vae, + dit_id=dit_id, + vae_id=vae_id, + block_swap_config={ + 'blocks_to_swap': args.blocks_to_swap, + 'swap_io_components': args.swap_io_components, + 'offload_device': dit_offload, + }, + encode_tiled=args.vae_encode_tiled, + encode_tile_size=(args.vae_encode_tile_size, args.vae_encode_tile_size), + encode_tile_overlap=(args.vae_encode_tile_overlap, args.vae_encode_tile_overlap), + decode_tiled=args.vae_decode_tiled, + decode_tile_size=(args.vae_decode_tile_size, args.vae_decode_tile_size), + decode_tile_overlap=(args.vae_decode_tile_overlap, args.vae_decode_tile_overlap), + tile_debug=args.tile_debug.lower() if args.tile_debug else "false", + attention_mode=args.attention_mode, + torch_compile_args_dit=torch_compile_args_dit, + torch_compile_args_vae=torch_compile_args_vae + ) + + ctx['cache_context'] = cache_context + if runner_cache is not None: + runner_cache['runner'] = runner + + # 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") + + # Compute generation info and log start (handles prepending internally) + frames_tensor, gen_info = compute_generation_info( + ctx=ctx, + images=frames_tensor, + resolution=args.resolution, + max_resolution=args.max_resolution, + batch_size=args.batch_size, + uniform_batch_size=args.uniform_batch_size, + seed=args.seed, + prepend_frames=args.prepend_frames, + temporal_overlap=args.temporal_overlap, + debug=debug + ) + log_generation_start(gen_info, debug) + + # Phase 1: Encode + ctx = encode_all_batches( + runner, ctx=ctx, images=frames_tensor, + debug=debug, + batch_size=args.batch_size, + uniform_batch_size=args.uniform_batch_size, + seed=args.seed, + progress_callback=None, + temporal_overlap=args.temporal_overlap, + resolution=args.resolution, + max_resolution=args.max_resolution, + input_noise_scale=args.input_noise_scale, + color_correction=args.color_correction + ) + + # Phase 2: Upscale + ctx = upscale_all_batches( + runner, ctx=ctx, debug=debug, progress_callback=None, + seed=args.seed, + latent_noise_scale=args.latent_noise_scale, + cache_model=cache_dit + ) + + # Phase 3: Decode + ctx = decode_all_batches( + runner, ctx=ctx, debug=debug, progress_callback=None, + cache_model=cache_vae + ) + + # Phase 4: Post-process + ctx = postprocess_all_batches( + ctx=ctx, debug=debug, progress_callback=None, + color_correction=args.color_correction, + prepend_frames=0, # Worker mode handles this in main process + temporal_overlap=args.temporal_overlap, + batch_size=args.batch_size + ) + + result_tensor = ctx['final_video'] + + # Convert to CPU and compatible dtype + if result_tensor.is_cuda or result_tensor.is_mps: + result_tensor = result_tensor.cpu() + if result_tensor.dtype in (torch.bfloat16, torch.float8_e4m3fn, torch.float8_e5m2): + result_tensor = result_tensor.to(torch.float32) + + return result_tensor + + +def _worker_process( + proc_idx: int, + device_id: str, + frames_np: Optional[np.ndarray], + shared_args: Dict[str, Any], + return_queue: mp.Queue, + done_barrier: mp.Barrier, + video_info: Optional[Dict[str, Any]] = None +) -> None: + """ + Worker process for multi-GPU upscaling. + + Supports two modes: + 1. frames_np provided: Process pre-loaded frames (for images) + 2. video_info provided: Stream video segment internally (for videos) + - Each worker opens the video, seeks to its assigned range, and streams + with internal chunking and model caching for memory efficiency + + Args: + proc_idx: Worker index for result ordering + device_id: GPU device ID (used for CUDA_VISIBLE_DEVICES inheritance) + frames_np: Pre-loaded frames as numpy array, or None for video streaming + shared_args: Serialized args namespace as dict + return_queue: Queue for returning results to parent + done_barrier: Barrier for synchronizing shared memory handoff + video_info: Optional dict with 'video_path', 'start_frame', 'end_frame' + for video streaming mode + """ + # Create debug instance for this worker + worker_debug = Debug(enabled=shared_args["debug"]) + + args = argparse.Namespace(**shared_args) + + # Video streaming mode: worker reads and processes its assigned segment + if video_info is not None: + cap = cv2.VideoCapture(video_info['video_path']) + cap.set(cv2.CAP_PROP_POS_FRAMES, video_info['start_frame']) + + segment_frames = video_info['end_frame'] - video_info['start_frame'] + chunk_size = args.chunk_size if args.chunk_size > 0 else segment_frames + + worker_debug.log(f"GPU {proc_idx}: frames {video_info['start_frame']}-{video_info['end_frame']} " + f"({segment_frames} frames, chunks of {chunk_size})", + category="generation", force=True) + + # Only GPU 0 uses prepend_frames (applies to video start only) + worker_args = argparse.Namespace(**vars(args)) + if proc_idx != 0: + worker_args.prepend_frames = 0 + + # Enable model caching within worker only if requested + runner_cache = {} if (args.cache_dit or args.cache_vae) else None + + total_chunks = (segment_frames + chunk_size - 1) // chunk_size + results = [] + for result in _stream_video_chunks( + cap=cap, + frames_to_process=segment_frames, + chunk_size=chunk_size, + overlap=args.temporal_overlap, + args=worker_args, + device_id="0", + debug=worker_debug, + runner_cache=runner_cache, + log_progress=total_chunks > 1, + total_chunks=total_chunks, + log_prefix=f"[GPU {proc_idx}] " + ): + results.append(result.cpu()) + + cap.release() + result_tensor = torch.cat(results, dim=0) if results else torch.empty(0, dtype=torch.float32) + + # Pre-loaded frames mode (original behavior) + else: + frames_tensor = torch.from_numpy(frames_np).to(torch.float16) + result_tensor = _process_frames_core( + frames_tensor=frames_tensor, + args=args, + device_id="0", + debug=worker_debug, + runner_cache=None + ) + + # Share tensor memory for efficient cross-process transfer (avoids pickling large arrays) + return_queue.put((proc_idx, result_tensor.share_memory_())) + + # Wait for parent to copy shared tensors before exiting + # (shared memory requires creating process to stay alive during access) + done_barrier.wait() + + +def _single_gpu_direct_processing( + frames_tensor: torch.Tensor, + args: argparse.Namespace, + device_id: str, + runner_cache: Optional[Dict[str, Any]] +) -> torch.Tensor: + """ + Direct single-GPU processing with model caching support. + + Uses main process and shared runner cache for efficient multi-file processing. + """ + return _process_frames_core( + frames_tensor=frames_tensor, + args=args, + device_id=device_id, + debug=debug, + runner_cache=runner_cache + ) + + +def _gpu_processing( + frames_tensor: Optional[torch.Tensor], + device_list: List[str], + args: argparse.Namespace, + video_info: Optional[Dict[str, Any]] = None +) -> torch.Tensor: + """ + Orchestrate multi-GPU parallel video upscaling with temporal overlap blending. + + Supports two modes: + 1. video_info provided: Workers stream their assigned video segments internally + (each GPU reads and processes its frame range with internal chunking) + 2. frames_tensor provided: Workers process pre-loaded frame chunks + (non streaming behavior for images or pre-loaded videos) + + Args: + frames_tensor: Input frames [T, H, W, C] or None if using video_info mode + device_list: List of device IDs as strings (e.g., ["0", "1"]) + args: Parsed command-line arguments containing all processing settings + video_info: Optional dict with 'video_path', 'start_frame', 'frames_to_process' + for streaming mode where workers read video directly + + Returns: + Upscaled frames tensor [T', H', W', C], Float32, range [0,1] + """ + num_devices = len(device_list) + overlap = args.temporal_overlap + + return_queue = mp.Queue(maxsize=0) + done_barrier = mp.Barrier(num_devices + 1) + workers = [] + shared_args = vars(args).copy() + + # Video streaming mode: distribute frame ranges to workers + if video_info is not None: + total_frames = video_info['frames_to_process'] + start_frame = video_info['start_frame'] + video_path = video_info['video_path'] + + base_per_gpu = total_frames // num_devices + remainder = total_frames % num_devices + + current_start = start_frame + for idx, device_id in enumerate(device_list): + gpu_frames = base_per_gpu + (1 if idx < remainder else 0) + gpu_end = current_start + gpu_frames + + # Add overlap frames for blending (except last GPU) + if idx < num_devices - 1 and overlap > 0: + gpu_end = min(gpu_end + overlap, start_frame + total_frames) + + worker_video_info = { + 'video_path': video_path, + 'start_frame': current_start, + 'end_frame': gpu_end, + } + + os.environ["CUDA_VISIBLE_DEVICES"] = device_id + p = mp.Process( + target=_worker_process, + args=(idx, device_id, None, shared_args, return_queue, done_barrier), + kwargs={'video_info': worker_video_info} + ) + p.start() + workers.append(p) + + current_start += gpu_frames + + # Pre-loaded frames mode (original behavior for images or non-streaming) + else: + total_frames = frames_tensor.shape[0] + + if overlap > 0 and num_devices > 1: + chunk_with_overlap = total_frames // num_devices + overlap + if args.batch_size > 1: + chunk_with_overlap = ((chunk_with_overlap + args.batch_size - 1) // args.batch_size) * args.batch_size + base_chunk_size = chunk_with_overlap - overlap + + chunks = [] + for i in range(num_devices): + start_idx = i * base_chunk_size + if i == num_devices - 1: + end_idx = total_frames + else: + end_idx = min(start_idx + chunk_with_overlap, total_frames) + chunks.append(frames_tensor[start_idx:end_idx]) + else: + chunks = torch.chunk(frames_tensor, num_devices, dim=0) + + for idx, (device_id, chunk_tensor) in enumerate(zip(device_list, chunks)): + os.environ["CUDA_VISIBLE_DEVICES"] = device_id + p = mp.Process( + target=_worker_process, + args=(idx, device_id, chunk_tensor.cpu().numpy(), shared_args, return_queue, done_barrier), + ) + p.start() + workers.append(p) + + # Collect results before joining to prevent deadlock + # Tensors arrive via shared memory - copy to numpy while workers still alive + results_np = [None] * num_devices + collected = 0 + while collected < num_devices: + proc_idx, result_tensor = return_queue.get() + results_np[proc_idx] = result_tensor.numpy() + collected += 1 + + # Release workers now that shared tensors are copied + done_barrier.wait() + + # Now safe to join + for p in workers: + p.join() + + # Concatenate results with overlap blending using shared function + if args.temporal_overlap > 0 and num_devices > 1: + overlap = args.temporal_overlap + result_tensor = None + + for idx, res_np in enumerate(results_np): + chunk_tensor = torch.from_numpy(res_np).to(torch.float32) + + if idx == 0: + # First chunk: keep all frames + result_tensor = chunk_tensor + else: + # Subsequent chunks: blend overlapping region with accumulated result + if chunk_tensor.shape[0] > overlap and result_tensor.shape[0] >= overlap: + # Get overlapping regions + prev_tail = result_tensor[-overlap:] # Last N frames from accumulated result + cur_head = chunk_tensor[:overlap] # First N frames from current chunk + + # Blend using shared function + blended = blend_overlapping_frames(prev_tail, cur_head, overlap) + + # Replace tail of result with blended frames, then append rest of chunk + result_tensor = torch.cat([ + result_tensor[:-overlap], # Everything except the tail + blended, # Blended overlapping frames + chunk_tensor[overlap:] # Non-overlapping part of current chunk + ], dim=0) + else: + # Edge case: chunk too small, just append non-overlapping part + if chunk_tensor.shape[0] > overlap: + result_tensor = torch.cat([result_tensor, chunk_tensor[overlap:]], dim=0) + + if result_tensor is None: + result_tensor = torch.from_numpy(results_np[0]).to(torch.float32) + else: + # Simple concatenation without overlap + result_tensor = torch.from_numpy(np.concatenate(results_np, axis=0)).to(torch.float32) + + # Handle prepend_frames removal (multi-GPU safe - done after all workers complete) + if args.prepend_frames > 0: + if args.prepend_frames < result_tensor.shape[0]: + debug.log(f"Removing {args.prepend_frames} prepended frames from output", category="generation") + result_tensor = result_tensor[args.prepend_frames:] + else: + debug.log(f"prepend_frames ({args.prepend_frames}) >= total frames ({result_tensor.shape[0]}), skipping removal", + level="WARNING", category="generation", force=True) + + return result_tensor + + +# ============================================================================= +# Argument Parsing +# ============================================================================= + +def parse_arguments() -> argparse.Namespace: + """ + Parse and validate command-line arguments for SeedVR2 CLI. + + Configures all available options including model selection, processing parameters, + memory optimization settings, and output configuration. + + Returns: + Parsed arguments namespace with all CLI parameters + + Note: + - cuda_device argument only available on non-macOS systems + - Default model directory resolves to "models/SEEDVR2" if not specified + """ + + # Get the actual invocation path for usage examples + invocation = sys.argv[0] + + # Multi-line usage examples for --help + usage_examples = f""" +Examples: + + Basic image upscaling: + python {invocation} image.jpg + + Basic video upscaling with temporal consistency: + python {invocation} video.mp4 --resolution 720 --batch_size 33 + + Streaming mode for long videos with 10-bit video output (requires FFMPEG): + python {invocation} 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 {invocation} video.mp4 --cuda_device 0,1 --resolution 1080 --batch_size 81 --uniform_batch_size --temporal_overlap 3 --prepend_frames 4 + + Memory-optimized for low VRAM (8GB): + python {invocation} image.png --dit_model seedvr2_ema_3b-Q8_0.gguf --blocks_to_swap 32 --swap_io_components --dit_offload_device cpu --vae_offload_device cpu + + High resolution with VAE tiling: + python {invocation} video.mp4 --resolution 1440 --batch_size 31 --uniform_batch_size --temporal_overlap 3 --vae_encode_tiled --vae_decode_tiled + + Batch directory processing: + python {invocation} media_folder/ --output processed/ --cuda_device 0 --cache_dit --cache_vae --dit_offload_device cpu --vae_offload_device cpu --resolution 1080 --max_resolution 1920 +""" + + parser = argparse.ArgumentParser( + description="SeedVR2 Video Upscaler - CLI for high-quality image/video upscaling and batch processing", + epilog=usage_examples, + formatter_class=argparse.RawDescriptionHelpFormatter, + allow_abbrev=False + ) + + # Input/Output + io_group = parser.add_argument_group('Input/Output options') + io_group.add_argument("input", type=str, + help="Input: video file (.mp4, .avi, etc.), image file (.png, .jpg, etc.), or directory") + io_group.add_argument("--output", type=str, default=None, + help="Output path (default: auto-generated in 'output/' directory)") + io_group.add_argument("--output_format", type=str, default=None, choices=["mp4", "png", None], + help="Output format: 'mp4' (video) or 'png' (image sequence). Default: auto-detect from input type") + io_group.add_argument("--video_backend", type=str, default="opencv", choices=["opencv", "ffmpeg"], + help="Video encoder backend: 'opencv' (default) or 'ffmpeg' (requires ffmpeg in PATH)") + io_group.add_argument("--10bit", dest="use_10bit", action="store_true", + help="Save 10-bit video with x265 codec (reduces banding). Without this flag, " + "ffmpeg uses x264 for maximum compatibility. Requires --video_backend ffmpeg") + io_group.add_argument("--model_dir", type=str, default=None, + help=f"Model directory (default: ./models/{SEEDVR2_FOLDER_NAME})") + + # Model Selection + model_group = parser.add_argument_group('Model selection') + model_group.add_argument("--dit_model", type=str, default=DEFAULT_DIT, + choices=get_available_dit_models(), + help="DiT model to use. Options: 3B (fp16/fp8/GGUF) or 7B (fp16/fp8/GGUF). Default: 3B FP8") + + # Processing Parameters + process_group = parser.add_argument_group('Processing parameters') + process_group.add_argument("--resolution", type=int, default=1080, + help="Target short-side resolution in pixels (default: 1080)") + process_group.add_argument("--max_resolution", type=int, default=0, + help="Maximum resolution for any edge. Scales down if exceeded. 0 = no limit (default: 0)") + process_group.add_argument("--batch_size", type=int, default=5, + help="Frames per batch (must follow 4n+1: 1, 5, 9, 13, 17, 21,...). " + "Ideally matches shot length for best temporal consistency. Higher values improve " + "quality and speed but require more VRAM. Default: 5") + process_group.add_argument("--uniform_batch_size", action="store_true", + help="Pad final batch to match batch_size. Prevents temporal artifacts caused by small " + "final batches. Add extra compute but recommended for optimal quality.") + process_group.add_argument("--seed", type=int, default=42, + help="Random seed for reproducibility (default: 42)") + process_group.add_argument("--skip_first_frames", type=int, default=0, + help="Skip N initial frames (default: 0)") + process_group.add_argument("--load_cap", type=int, default=0, + help="Load maximum N frames from video. 0 = load all (default: 0)") + process_group.add_argument("--chunk_size", type=int, default=0, + help="Frames per chunk for streaming mode. When > 0, processes video in " + "memory-bounded chunks of N frames. 0 = load all frames at once (default: 0)") + process_group.add_argument("--prepend_frames", type=int, default=0, + help="Prepend N reversed frames to reduce start artifacts (auto-removed). Default: 0") + process_group.add_argument("--temporal_overlap", type=int, default=0, + help="Frames to overlap between batches/GPUs for smooth blending (default: 0)") + + # Quality Control + quality_group = parser.add_argument_group('Quality control') + quality_group.add_argument("--color_correction", type=str, default="lab", + choices=["lab", "wavelet", "wavelet_adaptive", "hsv", "adain", "none"], + help="Color correction method: 'lab' (perceptual color matching, recommended), 'wavelet' (frequency-based), " + "'wavelet_adaptive' (wavelet + saturation correction), 'hsv' (hue-conditional), 'adain' (statistical transfer), " + "'none' (disabled) (default: lab)") + quality_group.add_argument("--input_noise_scale", type=float, default=0.0, + help="Input noise injection scale (0.0-1.0). Adds variation to input images (default: 0.0)") + quality_group.add_argument("--latent_noise_scale", type=float, default=0.0, + help="Latent noise injection scale (0.0-1.0). Adds variation to latent space (default: 0.0)") + + # Device Management + device_group = parser.add_argument_group('Device management') + if platform.system() != "Darwin": + device_group.add_argument("--cuda_device", type=str, default=None, + help="CUDA device(s): single '0' or multi-GPU '0,1,2'. Default: device 0") + device_group.add_argument("--dit_offload_device", type=str, default="none", + help="DiT offload device when idle: 'none' (keep on GPU), 'cpu' (offload to RAM), or GPU ID. " + "Frees VRAM between phases. Required for BlockSwap. Default: none") + device_group.add_argument("--vae_offload_device", type=str, default="none", + help="VAE offload device when idle: 'none', 'cpu', or GPU ID. Frees VRAM between phases. Default: none") + device_group.add_argument("--tensor_offload_device", type=str, default="cpu", + help="Intermediate tensor storage: 'cpu' (recommended), 'none' (keep on GPU), or GPU ID. Default: cpu") + + # Memory Optimization (BlockSwap) + blockswap_group = parser.add_argument_group('Memory optimization (BlockSwap)') + blockswap_group.add_argument("--blocks_to_swap", type=int, default=0, + help="Transformer blocks to swap for VRAM savings. 0-32 (3B) or 0-36 (7B). " + "Requires --dit_offload_device. Not available on macOS. Default: 0 (disabled)") + blockswap_group.add_argument("--swap_io_components", action="store_true", + help="Offload DiT I/O layers for extra VRAM savings. Requires --dit_offload_device. " + "Not available on macOS") + + # VAE Tiling + vae_group = parser.add_argument_group('VAE tiling (for high resolution upscale)') + vae_group.add_argument("--vae_encode_tiled", action="store_true", + help="Enable VAE encode tiling to reduce VRAM during encoding") + vae_group.add_argument("--vae_encode_tile_size", type=int, default=1024, + help="VAE encode tile size in pixels (default: 1024). Applied to both height and width. Only used if --vae_encode_tiled is set") + vae_group.add_argument("--vae_encode_tile_overlap", type=int, default=128, + help="VAE encode tile overlap in pixels (default: 128). Reduces visible seams between tiles. Only used if --vae_encode_tiled is set") + vae_group.add_argument("--vae_decode_tiled", action="store_true", + help="Enable VAE decode tiling to reduce VRAM during decoding") + vae_group.add_argument("--vae_decode_tile_size", type=int, default=1024, + help="VAE decode tile size in pixels (default: 1024). Applied to both height and width. Only used if --vae_decode_tiled is set") + vae_group.add_argument("--vae_decode_tile_overlap", type=int, default=128, + help="VAE decode tile overlap in pixels (default: 128). Reduces visible seams between tiles. Only used if --vae_decode_tiled is set") + vae_group.add_argument("--tile_debug", type=str, default="false", choices=["false", "encode", "decode"], + help="Visualize tiles: 'false' (default), 'encode', or 'decode'") + + # Performance + perf_group = parser.add_argument_group('Performance optimization') + perf_group.add_argument("--attention_mode", type=str, default="sdpa", + choices=["sdpa", "flash_attn_2", "flash_attn_3", "sageattn_2", "sageattn_3"], + help="Attention backend: 'sdpa' (default), 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3' (Blackwell GPUs)") + perf_group.add_argument("--compile_dit", action="store_true", + help="Enable torch.compile for DiT model (20-40%% speedup, requires PyTorch 2.0+ and Triton)") + perf_group.add_argument("--compile_vae", action="store_true", + help="Enable torch.compile for VAE model (15-25%% speedup, requires PyTorch 2.0+ and Triton)") + perf_group.add_argument("--compile_backend", type=str, default="inductor", choices=["inductor", "cudagraphs"], + help="Compilation backend: 'inductor' (full optimization with Triton) or 'cudagraphs' (lightweight, no kernel optimization) (default: inductor)") + perf_group.add_argument("--compile_mode", type=str, default="default", choices=["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"], + help="Optimization level: 'default' (fast compilation), 'reduce-overhead' (lower overhead), 'max-autotune' (best runtime, slow compilation), " + "'max-autotune-no-cudagraphs' (like max-autotune without cudagraphs) (default: default)") + perf_group.add_argument("--compile_fullgraph", action="store_true", + help="Compile entire model as single graph (faster but less flexible). May fail with dynamic shapes (default: False)") + perf_group.add_argument("--compile_dynamic", action="store_true", + help="Handle varying input shapes without recompilation. Useful for different resolutions/batch sizes (default: False)") + perf_group.add_argument("--compile_dynamo_cache_size_limit", type=int, default=64, + help="Max cached compiled versions per function. Increase when using many different input shapes. Higher uses more memory (default: 64)") + perf_group.add_argument("--compile_dynamo_recompile_limit", type=int, default=128, + help="Max recompilation attempts before fallback to eager mode. Safety limit to prevent compilation loops (default: 128)") + + # Model Caching (for batch processing) + cache_group = parser.add_argument_group('Model caching (batch processing)') + cache_group.add_argument("--cache_dit", action="store_true", + help="Keep DiT model in memory between generations. Works with single-GPU directory processing " + "or multi-GPU streaming (--chunk_size). Requires --dit_offload_device") + cache_group.add_argument("--cache_vae", action="store_true", + help="Keep VAE model in memory between generations. Works with single-GPU directory processing " + "or multi-GPU streaming (--chunk_size). Requires --vae_offload_device") + + # Debugging + debug_group = parser.add_argument_group('Debugging') + debug_group.add_argument("--debug", action="store_true", + help="Enable verbose debug logging") + + # Auto-show help if no arguments provided + if len(sys.argv) == 1: + sys.argv.append('--help') + + return parser.parse_args() + + +# ============================================================================= +# Main Entry Point +# ============================================================================= + +def main() -> None: + """ + Main entry point for SeedVR2 Video Upscaler CLI. + + Orchestrates the complete upscaling workflow: + 1. Parse and validate command-line arguments + 2. Extract frames from input video/image(s) + 3. Download required models if not cached + 4. Process frames on single or multiple GPUs + 5. Save results as video or PNG sequence + 6. Report timing and FPS (calculated from total wall-clock time) + + Error handling: + - Validates tile configuration before processing + - Provides detailed error messages with traceback + - Ensures proper cleanup on exit (VRAM automatically freed) + + Raises: + SystemExit: On argument validation failure or processing error + """ + # Parse arguments + args = parse_arguments() + + # Update debug instance with --debug flag + debug.enabled = args.debug + + # print header + debug.print_header(cli=True) + + debug.log("Arguments:", category="setup") + for key, value in vars(args).items(): + debug.log(f"{key}: {value}", category="none", indent_level=1) + + if args.vae_encode_tiled and args.vae_encode_tile_overlap >= args.vae_encode_tile_size: + debug.log(f"VAE encode tile overlap ({args.vae_encode_tile_overlap}) must be smaller than tile size ({args.vae_encode_tile_size})", level="ERROR", category="vae", force=True) + sys.exit(1) + + if args.vae_decode_tiled and args.vae_decode_tile_overlap >= args.vae_decode_tile_size: + debug.log(f"VAE decode tile overlap ({args.vae_decode_tile_overlap}) must be smaller than tile size ({args.vae_decode_tile_size})", level="ERROR", category="vae", force=True) + sys.exit(1) + + # Validate ffmpeg availability if selected + if args.video_backend == "ffmpeg" and shutil.which("ffmpeg") is None: + debug.log("--video_backend ffmpeg requires ffmpeg in PATH. Install ffmpeg or use --video_backend opencv", + level="ERROR", category="setup", force=True) + sys.exit(1) + + # Inform about caching defaults + if args.cache_dit and args.dit_offload_device == "none": + offload_target = "system memory (CPU)" if get_gpu_backend() != "mps" else "unified memory" + debug.log( + f"DiT caching enabled: Using default {offload_target} for offload. " + "Set --dit_offload_device explicitly to use a different device.", + category="cache", force=True + ) + + if args.cache_vae and args.vae_offload_device == "none": + offload_target = "system memory (CPU)" if get_gpu_backend() != "mps" else "unified memory" + debug.log( + f"VAE caching enabled: Using default {offload_target} for offload. " + "Set --vae_offload_device explicitly to use a different device.", + category="cache", force=True + ) + + if args.debug: + if platform.system() == "Darwin": + debug.log("You are running on macOS and will use the MPS backend!", category="info", force=True) + else: + # Show actual CUDA device visibility + debug.log(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set (all)')}", category="device") + if is_cuda_available(): + debug.log(f"torch.cuda.device_count(): {torch.cuda.device_count()}", category="device") + debug.log(f"Using device index 0 inside script (mapped to selected GPU)", category="device") + + try: + start_time = time.time() + + # Parse GPU list + if platform.system() == "Darwin": + device_list = ["0"] + else: + if args.cuda_device: + device_list = [d.strip() for d in str(args.cuda_device).split(',') if d.strip()] + else: + device_list = ["0"] + if args.debug: + debug.log(f"Using devices: {device_list}", category="device") + + # Download models once before processing + if not download_weight(dit_model=args.dit_model, vae_model=DEFAULT_VAE, model_dir=args.model_dir, debug=debug): + debug.log("Failed to download required models. Check console output above.", level="ERROR", category="download", force=True) + sys.exit(1) + + # Determine input type and process accordingly + input_type = get_input_type(args.input) + + # Track total frames for FPS calculation (time tracked via start_time) + total_frames_processed = 0 + + # Track if output format was user-specified or auto-detected + format_auto_detected = args.output_format is None + + if input_type == 'directory': + media_files = get_media_files(args.input) + if not media_files: + debug.log(f"No video or image files found in directory: {args.input}", + level="ERROR", category="file", force=True) + sys.exit(1) + + debug.log(f"Found {len(media_files)} media files to process", category="file", force=True) + + # Multi-GPU caching requires streaming (workers cache within their chunk loops) + if (args.cache_dit or args.cache_vae) and len(device_list) > 1 and args.chunk_size <= 0: + debug.log( + "Model caching requires streaming mode (--chunk_size > 0) for multi-GPU. " + "Disabling caching for this run.", + level="WARNING", category="cache", force=True + ) + args.cache_dit = False + args.cache_vae = False + + # Single-GPU: runner_cache persists across files; multi-GPU: workers cache internally + runner_cache = {} if (args.cache_dit or args.cache_vae) and len(device_list) == 1 else None + + for idx, file_path in enumerate(media_files, 1): + # Visual separation between files (except before first file) + if idx > 1: + debug.log("", category="none", force=True) + debug.log("โ”" * 60, category="none", force=True) + debug.log("", category="none", force=True) + + debug.log(f"Processing file {idx}/{len(media_files)}", category="generation", force=True) + + # Auto-detect format per file if not user-specified + if format_auto_detected: + file_type = get_input_type(file_path) + file_output_format = "mp4" if file_type == "video" else "png" + else: + file_output_format = args.output_format + + # Temporarily override args.output_format for this file + original_format = args.output_format + args.output_format = file_output_format + + # generate_output_path handles None gracefully with "outputs" default + output_path = generate_output_path(file_path, file_output_format, args.output, + input_type=get_input_type(file_path), from_directory=True) + + # Process with explicit output path and runner cache + frames = process_single_file(file_path, args, device_list, output_path, + format_auto_detected=format_auto_detected, + runner_cache=runner_cache) + total_frames_processed += frames + + # Restore original format + args.output_format = original_format + + elif input_type in ("video", "image"): + # Auto-detect output format for single file if not specified + if format_auto_detected: + args.output_format = "mp4" if input_type == "video" else "png" + + # Caching: single-GPU streaming uses runner_cache, multi-GPU streaming workers cache internally + runner_cache = None + streaming = args.chunk_size > 0 + + if args.cache_dit or args.cache_vae: + if len(device_list) > 1: + if not streaming: + debug.log( + "Model caching requires streaming mode (--chunk_size > 0) for multi-GPU. " + "Disabling caching for this run.", + level="WARNING", category="cache", force=True + ) + args.cache_dit = False + args.cache_vae = False + elif streaming: + runner_cache = {} + else: + debug.log( + "Model caching has no benefit for single file processing (only useful for directories or streaming mode). " + "Consider removing --cache_dit/--cache_vae for single files.", + category="tip", force=True + ) + + frames = process_single_file(args.input, args, device_list, args.output, + format_auto_detected=format_auto_detected, + runner_cache=runner_cache) + total_frames_processed += frames + + else: + debug.log(f"Unsupported input type: {args.input}", level="ERROR", category="file", force=True) + sys.exit(1) + + # Calculate total execution time + total_time = time.time() - start_time + + debug.log("", category="none", force=True) + debug.log(f"All upscaling processes completed successfully in {total_time:.2f}s", category="success", force=True) + + # Calculate and display FPS based on overall wall-clock time + if total_time > 0 and total_frames_processed > 0: + fps = total_frames_processed / total_time + debug.log(f"Average FPS: {fps:.2f} frames/sec", category="timing", force=True) + + except Exception as e: + debug.log(f"Error during processing: {e}", level="ERROR", category="generation", force=True) + import traceback + traceback.print_exc() + sys.exit(1) + + finally: + debug.log(f"Process {os.getpid()} terminating - VRAM will be automatically freed", category="cleanup", force=True) + + # print footer + debug.print_footer() + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/neg_emb.pt b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/neg_emb.pt new file mode 100644 index 0000000000000000000000000000000000000000..5eadd87ca26f82c5701178d4d6b2759c50b93cb2 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/neg_emb.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:6a43e5800ef2354f1c156d27535834da055cbec8248298b8923492bba2076581 +size 656540 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/pos_emb.pt b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/pos_emb.pt new file mode 100644 index 0000000000000000000000000000000000000000..4a51379217853a074b84d775d90c36985e238c2d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/pos_emb.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:fa07a14844314772266b66c3b95deb0027696d8fe7065721263db5176f45d799 +size 595100 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/pyproject.toml b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..71c12e21aac1cd6540d2aeb05149b989014d6595 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/pyproject.toml @@ -0,0 +1,40 @@ +[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.24" +authors = [ + {name = "numz"}, + {name = "adrientoupet"} +] +license = {file = "LICENSE"} +classifiers = [ + "Operating System :: OS Independent" +] +dependencies = [ + "torch", + "torchvision", + "safetensors", + "numpy", + "tqdm", + "psutil", + "einops", + "omegaconf>=2.3.0", + "diffusers>=0.33.1", + "peft>=0.17.0", + "rotary_embedding_torch>=0.5.3", + "opencv-python", + "gguf", + "matplotlib" +] + +[project.urls] +Repository = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler" +Documentation = "https://www.youtube.com/@AInVFX" +"Bug Tracker" = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/issues" +Forum = "https://github.com/numz/ComfyUI-SeedVR2_VideoUpscaler/discussions" + +[tool.comfy] +PublisherId = "ainvfx" +DisplayName = "ComfyUI-SeedVR2_VideoUpscaler" +Icon = "https://raw.githubusercontent.com/numz/ComfyUI-SeedVR2_VideoUpscaler/refs/heads/main/docs/seedvr_logo.png" +includes = [] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/requirements.txt b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..73d2db00f0b9f47eae750cafed4f0d5354542785 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/requirements.txt @@ -0,0 +1,14 @@ +torch +torchvision +safetensors +numpy +tqdm +psutil +einops +omegaconf>=2.3.0 +diffusers>=0.33.1 +peft>=0.17.0 +rotary_embedding_torch>=0.5.3 +opencv-python +gguf +matplotlib \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/cache.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/cache.py new file mode 100644 index 0000000000000000000000000000000000000000..89592fe8747a0b68b8553729abe908c6f06a5aa5 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/cache.py @@ -0,0 +1,47 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Callable + + +class Cache: + """Caching reusable args for faster inference""" + + def __init__(self, disable=False, prefix="", cache=None): + self.cache = cache if cache is not None else {} + self.disable = disable + self.prefix = prefix + + def __call__(self, key: str, fn: Callable): + if self.disable: + return fn() + + key = self.prefix + key + try: + result = self.cache[key] + except KeyError: + result = fn() + self.cache[key] = result + return result + + def namespace(self, namespace: str): + return Cache( + disable=self.disable, + prefix=self.prefix + namespace + ".", + cache=self.cache, + ) + + def get(self, key: str): + key = self.prefix + key + return self.cache[key] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/config.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c63366ab8cb3957f9278e95020b23751ce1a4c2e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/config.py @@ -0,0 +1,134 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Configuration utility functions +""" + +import importlib +from typing import Any, Callable, List, Union +from omegaconf import DictConfig, ListConfig, OmegaConf +from ..utils.model_registry import MODEL_CLASSES + +try: + OmegaConf.register_new_resolver("eval", eval) +except Exception as e: + if "already registered" not in str(e): + raise + + + +def load_config(path: str, argv: List[str] = None) -> Union[DictConfig, ListConfig]: + """ + Load a configuration. Will resolve inheritance. + """ + + #print(path) + config = OmegaConf.load(path) + if argv is not None: + config_argv = OmegaConf.from_dotlist(argv) + config = OmegaConf.merge(config, config_argv) + config = resolve_recursive(config, resolve_inheritance) + return config + + +def resolve_recursive( + config: Any, + resolver: Callable[[Union[DictConfig, ListConfig]], Union[DictConfig, ListConfig]], +) -> Any: + config = resolver(config) + if isinstance(config, DictConfig): + for k in config.keys(): + v = config.get(k) + if isinstance(v, (DictConfig, ListConfig)): + config[k] = resolve_recursive(v, resolver) + if isinstance(config, ListConfig): + for i in range(len(config)): + v = config.get(i) + if isinstance(v, (DictConfig, ListConfig)): + config[i] = resolve_recursive(v, resolver) + return config + + +def resolve_inheritance(config: Union[DictConfig, ListConfig]) -> Any: + """ + Recursively resolve inheritance if the config contains: + __inherit__: path/to/parent.yaml or a ListConfig of such paths. + """ + if isinstance(config, DictConfig): + inherit = config.pop("__inherit__", None) + + if inherit: + inherit_list = inherit if isinstance(inherit, ListConfig) else [inherit] + + parent_config = None + for parent_path in inherit_list: + assert isinstance(parent_path, str) + parent_config = ( + load_config(parent_path) + if parent_config is None + else OmegaConf.merge(parent_config, load_config(parent_path)) + ) + + if len(config.keys()) > 0: + config = OmegaConf.merge(parent_config, config) + else: + config = parent_config + return config + + +def import_item(path: str, name: str) -> Any: + """ + Import a python item, checking model registry first. + + Args: + path: Module path + name: Class/function name to import + + Returns: + Imported object + """ + # Simple lookup with path as key + if path in MODEL_CLASSES: + return MODEL_CLASSES[path] + + # Fallback to dynamic import for everything else + try: + return getattr(importlib.import_module(path), name) + except (ImportError, AttributeError) as e: + raise ImportError(f"Could not import '{name}' from '{path}': {e}") + + +def create_object(config: DictConfig) -> Any: + """ + Create an object from config. + The config is expected to contains the following: + __object__: + path: path.to.module + name: MyClass + args: as_config | as_params (default to as_config) + """ + + item = import_item( + path=config.__object__.path, + name=config.__object__.name, + ) + args = config.__object__.get("args", "as_config") + if args == "as_config": + return item(config) + if args == "as_params": + config = OmegaConf.to_object(config) + config.pop("__object__") + return item(**config) + raise NotImplementedError(f"Unknown args type: {args}") \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/decorators.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/decorators.py new file mode 100644 index 0000000000000000000000000000000000000000..30618855ec78cdbbe0e0d69e89e51dd912c1299f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/decorators.py @@ -0,0 +1,130 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Decorators. +""" + +import functools +import threading +import time +from typing import Callable +import torch + +from .distributed import barrier_if_distributed, get_global_rank, get_local_rank +from .logger import get_logger + +logger = get_logger(__name__) + + +def log_on_entry(func: Callable) -> Callable: + """ + Functions with this decorator will log the function name at entry. + When using multiple decorators, this must be applied innermost to properly capture the name. + """ + + def log_on_entry_wrapper(*args, **kwargs): + logger.info(f"Entering {func.__name__}") + return func(*args, **kwargs) + + return log_on_entry_wrapper + + +def barrier_on_entry(func: Callable) -> Callable: + """ + Functions with this decorator will start executing when all ranks are ready to enter. + """ + + def barrier_on_entry_wrapper(*args, **kwargs): + barrier_if_distributed() + return func(*args, **kwargs) + + return barrier_on_entry_wrapper + + +def _conditional_execute_wrapper_factory(execute: bool, func: Callable) -> Callable: + """ + Helper function for local_rank_zero_only and global_rank_zero_only. + """ + + def conditional_execute_wrapper(*args, **kwargs): + # Only execute if needed. + result = func(*args, **kwargs) if execute else None + # All GPUs must wait. + barrier_if_distributed() + # Return results. + return result + + return conditional_execute_wrapper + + +def _asserted_wrapper_factory(condition: bool, func: Callable, err_msg: str = "") -> Callable: + """ + Helper function for some functions with special constraints, + especially functions called by other global_rank_zero_only / local_rank_zero_only ones, + in case they are wrongly invoked in other scenarios. + """ + + def asserted_execute_wrapper(*args, **kwargs): + assert condition, err_msg + result = func(*args, **kwargs) + return result + + return asserted_execute_wrapper + + +def local_rank_zero_only(func: Callable) -> Callable: + """ + Functions with this decorator will only execute on local rank zero. + """ + return _conditional_execute_wrapper_factory(get_local_rank() == 0, func) + + +def global_rank_zero_only(func: Callable) -> Callable: + """ + Functions with this decorator will only execute on global rank zero. + """ + return _conditional_execute_wrapper_factory(get_global_rank() == 0, func) + + +def assert_only_global_rank_zero(func: Callable) -> Callable: + """ + Functions with this decorator are only accessible to processes with global rank zero. + """ + return _asserted_wrapper_factory( + get_global_rank() == 0, func, err_msg="Not accessible to processes with global_rank != 0" + ) + + +def assert_only_local_rank_zero(func: Callable) -> Callable: + """ + Functions with this decorator are only accessible to processes with local rank zero. + """ + return _asserted_wrapper_factory( + get_local_rank() == 0, func, err_msg="Not accessible to processes with local_rank != 0" + ) + + +def new_thread(func: Callable) -> Callable: + """ + Functions with this decorator will run in a new thread. + The function will return the thread, which can be joined to wait for completion. + """ + + def new_thread_wrapper(*args, **kwargs): + thread = threading.Thread(target=func, args=args, kwargs=kwargs) + thread.start() + return thread + + return new_thread_wrapper \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..034e36ef7f9eb0b3ae94280165e622a362e9fc1e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/__init__.py @@ -0,0 +1,56 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Diffusion package. +""" + +from .config import ( + create_sampler_from_config, + create_sampling_timesteps_from_config, + create_schedule_from_config, +) +from .samplers.base import Sampler +from .samplers.euler import EulerSampler +from .schedules.base import Schedule +from .schedules.lerp import LinearInterpolationSchedule +from .timesteps.base import SamplingTimesteps, Timesteps +from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps +from .types import PredictionType, SamplingDirection +from .utils import classifier_free_guidance, classifier_free_guidance_dispatcher, expand_dims + +__all__ = [ + # Configs + "create_sampler_from_config", + "create_sampling_timesteps_from_config", + "create_schedule_from_config", + # Schedules + "Schedule", + "DiscreteVariancePreservingSchedule", + "LinearInterpolationSchedule", + # Samplers + "Sampler", + "EulerSampler", + # Timesteps + "Timesteps", + "SamplingTimesteps", + # Types + "PredictionType", + "SamplingDirection", + "UniformTrailingSamplingTimesteps", + # Utils + "classifier_free_guidance", + "classifier_free_guidance_dispatcher", + "expand_dims", +] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/config.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/config.py new file mode 100644 index 0000000000000000000000000000000000000000..dabd17955c2d87d8d35fd2ceec33555daeed4ba6 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/config.py @@ -0,0 +1,75 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Utility functions for creating schedules and samplers from config. +""" + +import torch +from omegaconf import DictConfig + +from .samplers.base import Sampler +from .samplers.euler import EulerSampler +from .schedules.base import Schedule +from .schedules.lerp import LinearInterpolationSchedule +from .timesteps.base import SamplingTimesteps +from .timesteps.sampling.trailing import UniformTrailingSamplingTimesteps + + +def create_schedule_from_config( + config: DictConfig, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> Schedule: + """ + Create a schedule from configuration. + """ + if config.type == "lerp": + return LinearInterpolationSchedule(T=config.get("T", 1.0)) + + raise NotImplementedError + + +def create_sampler_from_config( + config: DictConfig, + schedule: Schedule, + timesteps: SamplingTimesteps, +) -> Sampler: + """ + Create a sampler from configuration. + """ + if config.type == "euler": + return EulerSampler( + schedule=schedule, + timesteps=timesteps, + prediction_type=config.prediction_type, + ) + raise NotImplementedError + + +def create_sampling_timesteps_from_config( + config: DictConfig, + schedule: Schedule, + device: torch.device, + dtype: torch.dtype = torch.float32, +) -> SamplingTimesteps: + if config.type == "uniform_trailing": + return UniformTrailingSamplingTimesteps( + T=schedule.T, + steps=config.steps, + shift=config.get("shift", 1.0), + device=device, + dtype=dtype, + ) + raise NotImplementedError \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/base.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/base.py new file mode 100644 index 0000000000000000000000000000000000000000..8e65f19896b6d5844e769762e76d699b96abc733 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/base.py @@ -0,0 +1,108 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Sampler base class. +""" + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import Callable +import torch +from tqdm import tqdm + +from ..schedules.base import Schedule +from ..timesteps.base import SamplingTimesteps +from ..types import PredictionType, SamplingDirection +from ..utils import assert_schedule_timesteps_compatible + + +@dataclass +class SamplerModelArgs: + x_t: torch.Tensor + t: torch.Tensor + i: int + + +class Sampler(ABC): + """ + Samplers are ODE/SDE solvers. + """ + + def __init__( + self, + schedule: Schedule, + timesteps: SamplingTimesteps, + prediction_type: PredictionType, + return_endpoint: bool = True, + ): + assert_schedule_timesteps_compatible( + schedule=schedule, + timesteps=timesteps, + ) + self.schedule = schedule + self.timesteps = timesteps + self.prediction_type = prediction_type + self.return_endpoint = return_endpoint + + @abstractmethod + def sample( + self, + x: torch.Tensor, + f: Callable[[SamplerModelArgs], torch.Tensor], + ) -> torch.Tensor: + """ + Generate a new sample given the the intial sample x and score function f. + """ + + def get_next_timestep( + self, + t: torch.Tensor, + ) -> torch.Tensor: + """ + Get the next sample timestep. + Support multiple different timesteps t in a batch. + If no more steps, return out of bound value -1 or T+1. + """ + T = self.timesteps.T + steps = len(self.timesteps) + curr_idx = self.timesteps.index(t) + next_idx = curr_idx + 1 + bound = -1 if self.timesteps.direction == SamplingDirection.backward else T + 1 + + s = self.timesteps[next_idx.clamp_max(steps - 1)] + s = s.where(next_idx < steps, bound) + return s + + def get_endpoint( + self, + pred: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + """ + Get to the endpoint of the probability flow. + """ + x_0, x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t) + return x_0 if self.timesteps.direction == SamplingDirection.backward else x_T + + def get_progress_bar(self): + """ + Get progress bar for sampling. + """ + return tqdm( + iterable=range(len(self.timesteps) - (0 if self.return_endpoint else 1)), + dynamic_ncols=True, + desc=self.__class__.__name__, + ) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/euler.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/euler.py new file mode 100644 index 0000000000000000000000000000000000000000..056b5b6aa4a67dd6c5f42224987dc58f433e116f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/samplers/euler.py @@ -0,0 +1,99 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + + +""" +Euler ODE solver. +""" + +from typing import Callable +import torch +from einops import rearrange +from torch.nn import functional as F + +from ..types import PredictionType +from ..utils import expand_dims +from .base import Sampler, SamplerModelArgs + + +class EulerSampler(Sampler): + """ + The Euler method is the simplest ODE solver. + + """ + + def sample( + self, + x: torch.Tensor, + f: Callable[[SamplerModelArgs], torch.Tensor], + ) -> torch.Tensor: + timesteps = self.timesteps.timesteps + progress = self.get_progress_bar() + i = 0 + + # Keep native dtype throughout sampling + # The DiT model already handles dtype internally via compatibility wrapper + for t, s in zip(timesteps[:-1], timesteps[1:]): + pred = f(SamplerModelArgs(x, t, i)) + + # Next step + x = self.step_to(pred, x, t, s) + + # Clean up temporary tensors + del pred + + i += 1 + progress.update() + + if self.return_endpoint: + t = timesteps[-1] + pred = f(SamplerModelArgs(x, t, i)) + x = self.get_endpoint(pred, x, t) + del pred + progress.update() + + return x + + def step( + self, + pred: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, + ) -> torch.Tensor: + """ + Step to the next timestep. + """ + return self.step_to(pred, x_t, t, self.get_next_timestep(t)) + + def step_to( + self, + pred: torch.Tensor, + x_t: torch.Tensor, + t: torch.Tensor, + s: torch.Tensor, + ) -> torch.Tensor: + """ + Steps from x_t at timestep t to x_s at timestep s. Returns x_s. + """ + t = expand_dims(t, x_t.ndim) + s = expand_dims(s, x_t.ndim) + T = self.schedule.T + # Step from x_t to x_s. + pred_x_0, pred_x_T = self.schedule.convert_from_pred(pred, self.prediction_type, x_t, t) + pred_x_s = self.schedule.forward(pred_x_0, pred_x_T, s.clamp(0, T)) + # Clamp x_s to x_0 and x_T if s is out of bound. + pred_x_s = pred_x_s.where(s >= 0, pred_x_0) + pred_x_s = pred_x_s.where(s <= T, pred_x_T) + return pred_x_s + \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/base.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/base.py new file mode 100644 index 0000000000000000000000000000000000000000..bcf6c6b6460977c6e2687e225c5c913a928bf812 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/base.py @@ -0,0 +1,131 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Schedule base class. +""" + +from abc import ABC, abstractmethod, abstractproperty +from typing import Tuple, Union +import torch + +from ..types import PredictionType +from ..utils import expand_dims + + +class Schedule(ABC): + """ + Diffusion schedules are uniquely defined by T, A, B: + + x_t = A(t) * x_0 + B(t) * x_T, where t in [0, T] + + Schedules can be continuous or discrete. + """ + + @abstractproperty + def T(self) -> Union[int, float]: + """ + Maximum timestep inclusive. + Schedule is continuous if float, discrete if int. + """ + + @abstractmethod + def A(self, t: torch.Tensor) -> torch.Tensor: + """ + Interpolation coefficient A. + Returns tensor with the same shape as t. + """ + + @abstractmethod + def B(self, t: torch.Tensor) -> torch.Tensor: + """ + Interpolation coefficient B. + Returns tensor with the same shape as t. + """ + + # ---------------------------------------------------- + + def snr(self, t: torch.Tensor) -> torch.Tensor: + """ + Signal to noise ratio. + Returns tensor with the same shape as t. + """ + return (self.A(t) ** 2) / (self.B(t) ** 2) + + def isnr(self, snr: torch.Tensor) -> torch.Tensor: + """ + Inverse signal to noise ratio. + Returns tensor with the same shape as snr. + Subclass may implement. + """ + raise NotImplementedError + + # ---------------------------------------------------- + + def is_continuous(self) -> bool: + """ + Whether the schedule is continuous. + """ + return isinstance(self.T, float) + + def forward(self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor) -> torch.Tensor: + """ + Diffusion forward function. + """ + t = expand_dims(t, x_0.ndim) + return self.A(t) * x_0 + self.B(t) * x_T + + def convert_from_pred( + self, pred: torch.Tensor, pred_type: PredictionType, x_t: torch.Tensor, t: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Convert from prediction. Return predicted x_0 and x_T. + """ + t = expand_dims(t, x_t.ndim) + A_t = self.A(t) + B_t = self.B(t) + + if pred_type == PredictionType.x_T: + pred_x_T = pred + pred_x_0 = (x_t - B_t * pred_x_T) / A_t + elif pred_type == PredictionType.x_0: + pred_x_0 = pred + pred_x_T = (x_t - A_t * pred_x_0) / B_t + elif pred_type == PredictionType.v_cos: + pred_x_0 = A_t * x_t - B_t * pred + pred_x_T = A_t * pred + B_t * x_t + elif pred_type == PredictionType.v_lerp: + pred_x_0 = (x_t - B_t * pred) / (A_t + B_t) + pred_x_T = (x_t + A_t * pred) / (A_t + B_t) + else: + raise NotImplementedError + + return pred_x_0, pred_x_T + + def convert_to_pred( + self, x_0: torch.Tensor, x_T: torch.Tensor, t: torch.Tensor, pred_type: PredictionType + ) -> torch.FloatTensor: + """ + Convert to prediction target given x_0 and x_T. + """ + if pred_type == PredictionType.x_T: + return x_T + if pred_type == PredictionType.x_0: + return x_0 + if pred_type == PredictionType.v_cos: + t = expand_dims(t, x_0.ndim) + return self.A(t) * x_T - self.B(t) * x_0 + if pred_type == PredictionType.v_lerp: + return x_T - x_0 + raise NotImplementedError diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/lerp.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/lerp.py new file mode 100644 index 0000000000000000000000000000000000000000..56b42bc17538b3217b2209234fc723ac3f58a746 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/schedules/lerp.py @@ -0,0 +1,55 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Linear interpolation schedule (lerp). +""" + +from typing import Union +import torch + +from .base import Schedule + + +class LinearInterpolationSchedule(Schedule): + """ + Linear interpolation schedule (lerp) is proposed by flow matching and rectified flow. + It leads to straighter probability flow theoretically. It is also used by Stable Diffusion 3. + + + + x_t = (1 - t) * x_0 + t * x_T + + Can be either continuous or discrete. + """ + + def __init__(self, T: Union[int, float] = 1.0): + self._T = T + + @property + def T(self) -> Union[int, float]: + return self._T + + def A(self, t: torch.Tensor) -> torch.Tensor: + return 1 - (t / self.T) + + def B(self, t: torch.Tensor) -> torch.Tensor: + return t / self.T + + # ---------------------------------------------------- + + def isnr(self, snr: torch.Tensor) -> torch.Tensor: + t = self.T / (1 + snr**0.5) + t = t if self.is_continuous() else t.round().int() + return t diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/base.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/base.py new file mode 100644 index 0000000000000000000000000000000000000000..d1a598103547694d5ef4dc5db0be1e5be2deb60c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/base.py @@ -0,0 +1,72 @@ +from abc import ABC, abstractmethod +from typing import Sequence, Union +import torch + +from ..types import SamplingDirection + + +class Timesteps(ABC): + """ + Timesteps base class. + """ + + def __init__(self, T: Union[int, float]): + assert T > 0 + self._T = T + + @property + def T(self) -> Union[int, float]: + """ + Maximum timestep inclusive. + int if discrete, float if continuous. + """ + return self._T + + def is_continuous(self) -> bool: + """ + Whether the schedule is continuous. + """ + return isinstance(self.T, float) + + +class SamplingTimesteps(Timesteps): + """ + Sampling timesteps. + It defines the discretization of sampling steps. + """ + + def __init__( + self, + T: Union[int, float], + timesteps: torch.Tensor, + direction: SamplingDirection, + ): + assert timesteps.ndim == 1 + super().__init__(T) + self.timesteps = timesteps + self.direction = direction + + def __len__(self) -> int: + """ + Number of sampling steps. + """ + return len(self.timesteps) + + def __getitem__(self, idx: Union[int, torch.IntTensor]) -> torch.Tensor: + """ + The timestep at the sampling step. + Returns a scalar tensor if idx is int, + or tensor of the same size if idx is a tensor. + """ + return self.timesteps[idx] + + def index(self, t: torch.Tensor) -> torch.Tensor: + """ + Find index by t. + Return index of the same shape as t. + Index is -1 if t not found in timesteps. + """ + i, j = t.reshape(-1, 1).eq(self.timesteps).nonzero(as_tuple=True) + idx = torch.full_like(t, fill_value=-1, dtype=torch.int) + idx.view(-1)[i] = j.int() + return idx diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/sampling/trailing.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/sampling/trailing.py new file mode 100644 index 0000000000000000000000000000000000000000..a6524b0a80d59d440f0a8f8e7a49073829cdcace --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/timesteps/sampling/trailing.py @@ -0,0 +1,50 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import torch + +from ...types import SamplingDirection +from ..base import SamplingTimesteps + + +class UniformTrailingSamplingTimesteps(SamplingTimesteps): + """ + Uniform trailing sampling timesteps. + Defined in (https://arxiv.org/abs/2305.08891) + + Shift is proposed in SD3 for RF schedule. + Defined in (https://arxiv.org/pdf/2403.03206) eq.23 + """ + + def __init__( + self, + T: int, + steps: int, + shift: float = 1.0, + device: torch.device = "cpu", + dtype: torch.dtype = torch.float32, + ): + # Create trailing timesteps with specified 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) + + # Scale to T range. + if isinstance(T, float): + timesteps = timesteps * T + else: + timesteps = timesteps.mul(T + 1).sub(1).round().int() + + super().__init__(T=T, timesteps=timesteps, direction=SamplingDirection.backward) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/types.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/types.py new file mode 100644 index 0000000000000000000000000000000000000000..076295f2be24dadc79da20a5f335b391eb9543bb --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/types.py @@ -0,0 +1,59 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Type definitions. +""" + +from enum import Enum + + +class PredictionType(str, Enum): + """ + x_0: + Predict data sample. + x_T: + Predict noise sample. + Proposed by DDPM (https://arxiv.org/abs/2006.11239) + Proved problematic by zsnr paper (https://arxiv.org/abs/2305.08891) + v_cos: + Predict velocity dx/dt based on the cosine schedule (A_t * x_T - B_t * x_0). + Proposed by progressive distillation (https://arxiv.org/abs/2202.00512) + v_lerp: + Predict velocity dx/dt based on the lerp schedule (x_T - x_0). + Proposed by rectified flow (https://arxiv.org/abs/2209.03003) + """ + + x_0 = "x_0" + x_T = "x_T" + v_cos = "v_cos" + v_lerp = "v_lerp" + + +class SamplingDirection(str, Enum): + """ + backward: Sample from x_T to x_0 for data generation. + forward: Sample from x_0 to x_T for noise inversion. + """ + + backward = "backward" + forward = "forward" + + @staticmethod + def reverse(direction): + if direction == SamplingDirection.backward: + return SamplingDirection.forward + if direction == SamplingDirection.forward: + return SamplingDirection.backward + raise NotImplementedError diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/utils.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..69d4aec34f59b293e2354744a4329008063a30e3 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/diffusion/utils.py @@ -0,0 +1,84 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Utility functions. +""" + +from typing import Callable +import torch + + +def expand_dims(tensor: torch.Tensor, ndim: int): + """ + Expand tensor to target ndim. New dims are added to the right. + For example, if the tensor shape was (8,), target ndim is 4, return (8, 1, 1, 1). + """ + shape = tensor.shape + (1,) * (ndim - tensor.ndim) + return tensor.reshape(shape) + + +def assert_schedule_timesteps_compatible(schedule, timesteps): + """ + Check if schedule and timesteps are compatible. + """ + if schedule.T != timesteps.T: + raise ValueError("Schedule and timesteps must have the same T.") + if schedule.is_continuous() != timesteps.is_continuous(): + raise ValueError("Schedule and timesteps must have the same continuity.") + + +def classifier_free_guidance( + pos: torch.Tensor, + neg: torch.Tensor, + scale: float, + rescale: float = 0.0, +): + """ + Apply classifier-free guidance. + """ + # Classifier-free guidance (https://arxiv.org/abs/2207.12598) + cfg = neg + scale * (pos - neg) + + # Classifier-free guidance rescale (https://arxiv.org/pdf/2305.08891.pdf) + if rescale != 0.0: + pos_std = pos.std(dim=list(range(1, pos.ndim)), keepdim=True) + cfg_std = cfg.std(dim=list(range(1, cfg.ndim)), keepdim=True) + factor = pos_std / cfg_std + factor = rescale * factor + (1 - rescale) + cfg *= factor + + return cfg + + +def classifier_free_guidance_dispatcher( + pos: Callable, + neg: Callable, + scale: float, + rescale: float = 0.0, +): + """ + Optionally execute models depending on classifer-free guidance scale. + """ + # If scale is 1, no need to execute neg model. + if scale == 1.0: + return pos() + + # Otherwise, execute both pos nad neg models and apply cfg. + return classifier_free_guidance( + pos=pos(), + neg=neg(), + scale=scale, + rescale=rescale, + ) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a5b4f873ae3e5524c88942bb27ec98ac98c3b5b5 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/__init__.py @@ -0,0 +1,37 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Distributed package. +""" + +from .basic import ( + barrier_if_distributed, + convert_to_ddp, + get_device, + get_global_rank, + get_local_rank, + get_world_size, + init_torch, +) + +__all__ = [ + "barrier_if_distributed", + "convert_to_ddp", + "get_device", + "get_global_rank", + "get_local_rank", + "get_world_size", + "init_torch", +] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/advanced.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/advanced.py new file mode 100644 index 0000000000000000000000000000000000000000..68d48877cf5a7a7abd2c2031d0bac54ea67b5c1c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/advanced.py @@ -0,0 +1,223 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Advanced distributed functions for sequence parallel. +""" + +from typing import Optional, List, TYPE_CHECKING +import torch +import torch.distributed as dist + +# Conditional imports for distributed training features (not needed for inference) +try: + from torch.distributed.device_mesh import DeviceMesh, init_device_mesh + from torch.distributed.fsdp import ShardingStrategy + _FSDP_AVAILABLE = True +except (ImportError, AttributeError): + # AMD ROCm and other builds may not have full FSDP support + DeviceMesh = None + init_device_mesh = None + ShardingStrategy = None + _FSDP_AVAILABLE = False + +from .basic import get_global_rank, get_world_size + + +_DATA_PARALLEL_GROUP = None +_SEQUENCE_PARALLEL_GROUP = None +_SEQUENCE_PARALLEL_CPU_GROUP = None +_MODEL_SHARD_CPU_INTER_GROUP = None +_MODEL_SHARD_CPU_INTRA_GROUP = None +_MODEL_SHARD_INTER_GROUP = None +_MODEL_SHARD_INTRA_GROUP = None +_SEQUENCE_PARALLEL_GLOBAL_RANKS = None + + +def get_data_parallel_group() -> Optional[dist.ProcessGroup]: + """ + Get data parallel process group. + """ + return _DATA_PARALLEL_GROUP + + +def get_sequence_parallel_group() -> Optional[dist.ProcessGroup]: + """ + Get sequence parallel process group. + """ + return _SEQUENCE_PARALLEL_GROUP + + +def get_sequence_parallel_cpu_group() -> Optional[dist.ProcessGroup]: + """ + Get sequence parallel CPU process group. + """ + return _SEQUENCE_PARALLEL_CPU_GROUP + + +def get_data_parallel_rank() -> int: + """ + Get data parallel rank. + """ + group = get_data_parallel_group() + return dist.get_rank(group) if group else get_global_rank() + + +def get_data_parallel_world_size() -> int: + """ + Get data parallel world size. + """ + group = get_data_parallel_group() + return dist.get_world_size(group) if group else get_world_size() + + +def get_sequence_parallel_rank() -> int: + """ + Get sequence parallel rank. + """ + group = get_sequence_parallel_group() + return dist.get_rank(group) if group else 0 + + +def get_sequence_parallel_world_size() -> int: + """ + Get sequence parallel world size. + """ + group = get_sequence_parallel_group() + return dist.get_world_size(group) if group else 1 + + +def get_model_shard_cpu_intra_group() -> Optional[dist.ProcessGroup]: + """ + Get the CPU intra process group of model sharding. + """ + return _MODEL_SHARD_CPU_INTRA_GROUP + + +def get_model_shard_cpu_inter_group() -> Optional[dist.ProcessGroup]: + """ + Get the CPU inter process group of model sharding. + """ + return _MODEL_SHARD_CPU_INTER_GROUP + + +def get_model_shard_intra_group() -> Optional[dist.ProcessGroup]: + """ + Get the GPU intra process group of model sharding. + """ + return _MODEL_SHARD_INTRA_GROUP + + +def get_model_shard_inter_group() -> Optional[dist.ProcessGroup]: + """ + Get the GPU inter process group of model sharding. + """ + return _MODEL_SHARD_INTER_GROUP + + +def init_sequence_parallel(sequence_parallel_size: int): + """ + Initialize sequence parallel. + """ + global _DATA_PARALLEL_GROUP + global _SEQUENCE_PARALLEL_GROUP + global _SEQUENCE_PARALLEL_CPU_GROUP + global _SEQUENCE_PARALLEL_GLOBAL_RANKS + assert dist.is_initialized() + world_size = dist.get_world_size() + rank = dist.get_rank() + data_parallel_size = world_size // sequence_parallel_size + for i in range(data_parallel_size): + start_rank = i * sequence_parallel_size + end_rank = (i + 1) * sequence_parallel_size + ranks = range(start_rank, end_rank) + group = dist.new_group(ranks) + cpu_group = dist.new_group(ranks, backend="gloo") + if rank in ranks: + _SEQUENCE_PARALLEL_GROUP = group + _SEQUENCE_PARALLEL_CPU_GROUP = cpu_group + _SEQUENCE_PARALLEL_GLOBAL_RANKS = list(ranks) + + +def init_model_shard_group( + *, + sharding_strategy: ShardingStrategy, + device_mesh: Optional[DeviceMesh] = None, +): + """ + Initialize process group of model sharding. + """ + if not _FSDP_AVAILABLE: + raise RuntimeError( + "FSDP features are not available in this PyTorch build. " + "Model sharding requires torch.distributed.fsdp support." + ) + global _MODEL_SHARD_INTER_GROUP + global _MODEL_SHARD_INTRA_GROUP + global _MODEL_SHARD_CPU_INTER_GROUP + global _MODEL_SHARD_CPU_INTRA_GROUP + assert dist.is_initialized() + world_size = dist.get_world_size() + if device_mesh is not None: + num_shards_per_group = device_mesh.shape[1] + elif sharding_strategy == ShardingStrategy.NO_SHARD: + num_shards_per_group = 1 + elif sharding_strategy in [ + ShardingStrategy.HYBRID_SHARD, + ShardingStrategy._HYBRID_SHARD_ZERO2, + ]: + num_shards_per_group = torch.cuda.device_count() + else: + num_shards_per_group = world_size + num_groups = world_size // num_shards_per_group + device_mesh = (num_groups, num_shards_per_group) + + gpu_mesh_2d = init_device_mesh("cuda", device_mesh, mesh_dim_names=("inter", "intra")) + cpu_mesh_2d = init_device_mesh("cpu", device_mesh, mesh_dim_names=("inter", "intra")) + + _MODEL_SHARD_INTER_GROUP = gpu_mesh_2d.get_group("inter") + _MODEL_SHARD_INTRA_GROUP = gpu_mesh_2d.get_group("intra") + _MODEL_SHARD_CPU_INTER_GROUP = cpu_mesh_2d.get_group("inter") + _MODEL_SHARD_CPU_INTRA_GROUP = cpu_mesh_2d.get_group("intra") + +def get_sequence_parallel_global_ranks() -> List[int]: + """ + Get all global ranks of the sequence parallel process group + that the caller rank belongs to. + """ + if _SEQUENCE_PARALLEL_GLOBAL_RANKS is None: + return [dist.get_rank()] + return _SEQUENCE_PARALLEL_GLOBAL_RANKS + + +def get_next_sequence_parallel_rank() -> int: + """ + Get the next global rank of the sequence parallel process group + that the caller rank belongs to. + """ + sp_global_ranks = get_sequence_parallel_global_ranks() + sp_rank = get_sequence_parallel_rank() + sp_size = get_sequence_parallel_world_size() + return sp_global_ranks[(sp_rank + 1) % sp_size] + + +def get_prev_sequence_parallel_rank() -> int: + """ + Get the previous global rank of the sequence parallel process group + that the caller rank belongs to. + """ + sp_global_ranks = get_sequence_parallel_global_ranks() + sp_rank = get_sequence_parallel_rank() + sp_size = get_sequence_parallel_world_size() + return sp_global_ranks[(sp_rank + sp_size - 1) % sp_size] \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/basic.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/basic.py new file mode 100644 index 0000000000000000000000000000000000000000..d92610eeae345a72b967633654a3e7b183c4fef4 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/basic.py @@ -0,0 +1,87 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Distributed basic functions. +""" + +import os +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: + """ + Get the global rank, the global index of the GPU. + """ + return int(os.environ.get("RANK", "0")) + + +def get_local_rank() -> int: + """ + Get the local rank, the local index of the GPU. + """ + return int(os.environ.get("LOCAL_RANK", "0")) + + +def get_world_size() -> int: + """ + Get the world size, the total amount of GPUs. + """ + return int(os.environ.get("WORLD_SIZE", "1")) + + +def get_device() -> torch.device: + """ + Get current rank device. + """ + if is_mps_available(): + return torch.device("mps") + return torch.device("cuda", get_local_rank()) + + +def barrier_if_distributed(*args, **kwargs): + """ + Synchronizes all processes if under distributed context. + """ + if dist.is_initialized(): + return dist.barrier(*args, **kwargs) + + +def init_torch(cudnn_benchmark=True, timeout=timedelta(seconds=600)): + """ + Common PyTorch initialization configuration. + """ + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + torch.backends.cudnn.benchmark = cudnn_benchmark + torch.cuda.set_device(get_local_rank()) + dist.init_process_group( + backend="nccl", + rank=get_global_rank(), + world_size=get_world_size(), + timeout=timeout, + ) + + +def convert_to_ddp(module: torch.nn.Module, **kwargs) -> DistributedDataParallel: + return DistributedDataParallel( + module=module, + device_ids=[get_local_rank()], + output_device=get_local_rank(), + **kwargs, + ) + \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/meta_init_utils.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/meta_init_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e6a92b513a995ede11a6f3aa47084395abad0c8e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/meta_init_utils.py @@ -0,0 +1,40 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import torch +from rotary_embedding_torch import RotaryEmbedding +from torch import nn + +__all__ = ["meta_non_persistent_buffer_init_fn"] + + +def meta_non_persistent_buffer_init_fn(module: nn.Module) -> nn.Module: + """ + Used for materializing `non-persistent tensor buffers` while model resuming. + + Since non-persistent tensor buffers are not saved in state_dict, + when initializing model with meta device, user should materialize those buffers manually. + + Currently, only `rope.dummy` is this special case. + """ + with torch.no_grad(): + for submodule in module.modules(): + if not isinstance(submodule, RotaryEmbedding): + continue + for buffer_name, buffer in submodule.named_buffers(recurse=False): + if buffer.is_meta and "dummy" in buffer_name: + materialized_buffer = torch.zeros_like(buffer, device="cpu") + setattr(submodule, buffer_name, materialized_buffer) + assert not any(b.is_meta for n, b in module.named_buffers()) + return module diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/ops.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/ops.py new file mode 100644 index 0000000000000000000000000000000000000000..380d95c80742b4e5a678888a342d04631c9672a9 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/distributed/ops.py @@ -0,0 +1,495 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Distributed ops for supporting sequence parallel. +""" + +from collections import defaultdict +from typing import Any, Callable, Dict, List, Optional, Tuple, Union +import torch +import torch.distributed as dist +from torch import Tensor + +from ..cache import Cache +from .advanced import ( + get_sequence_parallel_group, + get_sequence_parallel_rank, + get_sequence_parallel_world_size, +) + +from .basic import get_device + +_SEQ_DATA_BUF = defaultdict(lambda: [None, None, None]) +_SEQ_DATA_META_SHAPES = defaultdict() +_SEQ_DATA_META_DTYPES = defaultdict() +_SEQ_DATA_ASYNC_COMMS = defaultdict(list) +_SYNC_BUFFER = defaultdict(dict) + + +def single_all_to_all( + local_input: Tensor, + scatter_dim: int, + gather_dim: int, + group: dist.ProcessGroup, + async_op: bool = False, +): + """ + A function to do all-to-all on a tensor + """ + seq_world_size = dist.get_world_size(group) + prev_scatter_dim = scatter_dim + if scatter_dim != 0: + local_input = local_input.transpose(0, scatter_dim) + if gather_dim == 0: + gather_dim = scatter_dim + scatter_dim = 0 + + inp_shape = list(local_input.shape) + inp_shape[scatter_dim] = inp_shape[scatter_dim] // seq_world_size + input_t = local_input.reshape( + [seq_world_size, inp_shape[scatter_dim]] + inp_shape[scatter_dim + 1 :] + ).contiguous() + output = torch.empty_like(input_t) + comm = dist.all_to_all_single(output, input_t, group=group, async_op=async_op) + if async_op: + # let user's code transpose & reshape + return output, comm, prev_scatter_dim + + # first dim is seq_world_size, so we can split it directly + output = torch.cat(output.split(1), dim=gather_dim + 1).squeeze(0) + if prev_scatter_dim: + output = output.transpose(0, prev_scatter_dim).contiguous() + return output + + +def _all_to_all( + local_input: Tensor, + scatter_dim: int, + gather_dim: int, + group: dist.ProcessGroup, +): + seq_world_size = dist.get_world_size(group) + input_list = [ + t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim) + ] + output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)] + dist.all_to_all(output_list, input_list, group=group) + return torch.cat(output_list, dim=gather_dim).contiguous() + + +class SeqAllToAll(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + local_input: Tensor, + scatter_dim: int, + gather_dim: int, + async_op: bool, + ) -> Tensor: + ctx.group = group + ctx.scatter_dim = scatter_dim + ctx.gather_dim = gather_dim + ctx.async_op = async_op + if async_op: + output, comm, prev_scatter_dim = single_all_to_all( + local_input, scatter_dim, gather_dim, group, async_op=async_op + ) + ctx.prev_scatter_dim = prev_scatter_dim + return output, comm + + return _all_to_all(local_input, scatter_dim, gather_dim, group) + + @staticmethod + def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: + if ctx.async_op: + input_t = torch.cat(grad_output[0].split(1), dim=ctx.gather_dim + 1).squeeze(0) + if ctx.prev_scatter_dim: + input_t = input_t.transpose(0, ctx.prev_scatter_dim) + else: + input_t = grad_output[0] + return ( + None, + _all_to_all(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group), + None, + None, + None, + ) + + +class Slice(torch.autograd.Function): + @staticmethod + def forward(ctx: Any, group: dist.ProcessGroup, local_input: Tensor, dim: int) -> Tensor: + ctx.group = group + ctx.rank = dist.get_rank(group) + seq_world_size = dist.get_world_size(group) + ctx.seq_world_size = seq_world_size + ctx.dim = dim + dim_size = local_input.shape[dim] + return local_input.split(dim_size // seq_world_size, dim=dim)[ctx.rank].contiguous() + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor, None]: + dim_size = list(grad_output.size()) + split_size = dim_size[0] + dim_size[0] = dim_size[0] * ctx.seq_world_size + output = torch.empty(dim_size, dtype=grad_output.dtype, device=torch.cuda.current_device()) + dist._all_gather_base(output, grad_output, group=ctx.group) + return (None, torch.cat(output.split(split_size), dim=ctx.dim), None) + + +class Gather(torch.autograd.Function): + @staticmethod + def forward( + ctx: Any, + group: dist.ProcessGroup, + local_input: Tensor, + dim: int, + grad_scale: Optional[bool] = False, + ) -> Tensor: + ctx.group = group + ctx.rank = dist.get_rank(group) + ctx.dim = dim + ctx.grad_scale = grad_scale + seq_world_size = dist.get_world_size(group) + ctx.seq_world_size = seq_world_size + dim_size = list(local_input.size()) + split_size = dim_size[0] + ctx.part_size = dim_size[dim] + dim_size[0] = dim_size[0] * seq_world_size + output = torch.empty(dim_size, dtype=local_input.dtype, device=torch.cuda.current_device()) + dist._all_gather_base(output, local_input.contiguous(), group=ctx.group) + return torch.cat(output.split(split_size), dim=dim) + + @staticmethod + def backward(ctx: Any, grad_output: Tensor) -> Tuple[None, Tensor]: + if ctx.grad_scale: + grad_output = grad_output * ctx.seq_world_size + return ( + None, + grad_output.split(ctx.part_size, dim=ctx.dim)[ctx.rank].contiguous(), + None, + None, + ) + + +def gather_seq_scatter_heads_qkv( + qkv_tensor: Tensor, + *, + seq_dim: int, + qkv_shape: Optional[Tensor] = None, + cache: Cache = Cache(disable=True), + restore_shape: bool = True, +): + """ + A func to sync splited qkv tensor + qkv_tensor: the tensor we want to do alltoall with. The last dim must + be the projection_idx, which we will split into 3 part. After + spliting, the gather idx will be projecttion_idx + 1 + seq_dim: gather_dim for all2all comm + restore_shape: if True, output will has the same shape length as input + """ + group = get_sequence_parallel_group() + if not group: + return qkv_tensor + world = get_sequence_parallel_world_size() + orig_shape = qkv_tensor.shape + scatter_dim = qkv_tensor.dim() + bef_all2all_shape = list(orig_shape) + qkv_proj_dim = bef_all2all_shape[-1] + bef_all2all_shape = bef_all2all_shape[:-1] + [3, qkv_proj_dim // 3] + qkv_tensor = qkv_tensor.view(bef_all2all_shape) + qkv_tensor = SeqAllToAll.apply(group, qkv_tensor, scatter_dim, seq_dim, False) + if restore_shape: + out_shape = list(orig_shape) + out_shape[seq_dim] *= world + out_shape[-1] = qkv_proj_dim // world + qkv_tensor = qkv_tensor.view(out_shape) + + # remove padding + if qkv_shape is not None: + unpad_dim_size = cache( + "unpad_dim_size", lambda: torch.sum(torch.prod(qkv_shape, dim=-1)).item() + ) + if unpad_dim_size % world != 0: + padding_size = qkv_tensor.size(seq_dim) - unpad_dim_size + qkv_tensor = _unpad_tensor(qkv_tensor, seq_dim, padding_size) + return qkv_tensor + + +def slice_inputs(x: Tensor, dim: int, padding: bool = True): + """ + A func to slice the input sequence in sequence parallel + """ + group = get_sequence_parallel_group() + if group is None: + return x + sp_rank = get_sequence_parallel_rank() + sp_world = get_sequence_parallel_world_size() + dim_size = x.shape[dim] + unit = (dim_size + sp_world - 1) // sp_world + if padding and dim_size % sp_world: + padding_size = sp_world - (dim_size % sp_world) + x = _pad_tensor(x, dim, padding_size) + slc = [slice(None)] * len(x.shape) + slc[dim] = slice(unit * sp_rank, unit * (sp_rank + 1)) + return x[slc] + + +def remove_seqeunce_parallel_padding(x: Tensor, dim: int, unpad_dim_size: int): + """ + A func to remove the padding part of the tensor based on its original shape + """ + group = get_sequence_parallel_group() + if group is None: + return x + sp_world = get_sequence_parallel_world_size() + if unpad_dim_size % sp_world == 0: + return x + padding_size = sp_world - (unpad_dim_size % sp_world) + assert (padding_size + unpad_dim_size) % sp_world == 0 + return _unpad_tensor(x, dim=dim, padding_size=padding_size) + + +def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int) -> Tensor: + """ + A func to sync attention result with alltoall in sequence parallel + """ + group = get_sequence_parallel_group() + if not group: + return x + dim_size = x.size(seq_dim) + sp_world = get_sequence_parallel_world_size() + if dim_size % sp_world != 0: + padding_size = sp_world - (dim_size % sp_world) + x = _pad_tensor(x, seq_dim, padding_size) + return SeqAllToAll.apply(group, x, seq_dim, head_dim, False) + + +def gather_seq_scatter_heads(x: Tensor, seq_dim: int, head_dim: int) -> Tensor: + """ + A func to sync embedding input with alltoall in sequence parallel + """ + group = get_sequence_parallel_group() + if not group: + return x + return SeqAllToAll.apply(group, x, head_dim, seq_dim, False) + + +def scatter_heads(x: Tensor, dim: int) -> Tensor: + """ + A func to split heads before attention in sequence parallel + """ + group = get_sequence_parallel_group() + if not group: + return x + return Slice.apply(group, x, dim) + + +def gather_heads(x: Tensor, dim: int, grad_scale: Optional[bool] = False) -> Tensor: + """ + A func to gather heads for the attention result in sequence parallel + """ + group = get_sequence_parallel_group() + if not group: + return x + return Gather.apply(group, x, dim, grad_scale) + + +def gather_outputs( + x: Tensor, + *, + gather_dim: int, + padding_dim: Optional[int] = None, + unpad_shape: Optional[Tensor] = None, + cache: Cache = Cache(disable=True), + scale_grad=True, +): + """ + A func to gather the outputs for the model result in sequence parallel + """ + group = get_sequence_parallel_group() + if not group: + return x + x = Gather.apply(group, x, gather_dim, scale_grad) + if padding_dim is not None: + unpad_dim_size = cache( + "unpad_dim_size", lambda: torch.sum(torch.prod(unpad_shape, dim=1)).item() + ) + x = remove_seqeunce_parallel_padding(x, padding_dim, unpad_dim_size) + return x + + +def _pad_tensor(x: Tensor, dim: int, padding_size: int): + shape = list(x.shape) + shape[dim] = padding_size + pad = torch.zeros(shape, dtype=x.dtype, device=x.device) + return torch.cat([x, pad], dim=dim) + + +def _unpad_tensor(x: Tensor, dim: int, padding_size): + slc = [slice(None)] * len(x.shape) + slc[dim] = slice(0, -padding_size) + return x[slc] + + +def _broadcast_data(data, shape, dtype, src, group, async_op): + comms = [] + if isinstance(data, (list, tuple)): + for i, sub_shape in enumerate(shape): + comms += _broadcast_data(data[i], sub_shape, dtype[i], src, group, async_op) + elif isinstance(data, dict): + for key, sub_data in data.items(): + comms += _broadcast_data(sub_data, shape[key], dtype[key], src, group, async_op) + elif isinstance(data, Tensor): + comms.append(dist.broadcast(data, src=src, group=group, async_op=async_op)) + return comms + + +def _traverse(data: Any, op: Callable) -> Union[None, List, Dict, Any]: + if isinstance(data, (list, tuple)): + return [_traverse(sub_data, op) for sub_data in data] + elif isinstance(data, dict): + return {key: _traverse(sub_data, op) for key, sub_data in data.items()} + elif isinstance(data, Tensor): + return op(data) + else: + return None + + +def _get_shapes(data): + return _traverse(data, op=lambda x: x.shape) + + +def _get_dtypes(data): + return _traverse(data, op=lambda x: x.dtype) + + +def _construct_broadcast_buffer(shapes, dtypes, device): + if isinstance(shapes, torch.Size): + return torch.empty(shapes, dtype=dtypes, device=device) + + if isinstance(shapes, (list, tuple)): + buffer = [] + for i, sub_shape in enumerate(shapes): + buffer.append(_construct_broadcast_buffer(sub_shape, dtypes[i], device)) + elif isinstance(shapes, dict): + buffer = {} + for key, sub_shape in shapes.items(): + buffer[key] = _construct_broadcast_buffer(sub_shape, dtypes[key], device) + else: + return None + return buffer + + +class SPDistForward: + """A forward tool to sync different result across sp group + + Args: + module: a function or module to process users input + sp_step: current training step to judge which rank to broadcast its result to all + name: a distinct str to save meta and async comm + comm_shape: if different ranks have different shape, mark this arg to True + device: the device for current rank, can be empty + """ + + def __init__( + self, + name: str, + comm_shape: bool, + device: torch.device = None, + ): + self.name = name + self.comm_shape = comm_shape + if device: + self.device = device if isinstance(device, torch.device) else torch.device(device) + else: + # Fallback to standard device detection + self.device = get_device() + + def __call__(self, inputs) -> Any: + group = get_sequence_parallel_group() + if not group: + yield inputs + else: + device = self.device + sp_world = get_sequence_parallel_world_size() + sp_rank = get_sequence_parallel_rank() + for local_step in range(sp_world): + src_rank = dist.get_global_rank(group, local_step) + is_src = sp_rank == local_step + local_shapes = [] + local_dtypes = [] + if local_step == 0: + local_result = inputs + _SEQ_DATA_BUF[self.name][-1] = local_result + local_shapes = _get_shapes(local_result) + local_dtypes = _get_dtypes(local_result) + if self.comm_shape: + group_shapes_lists = [None] * sp_world + dist.all_gather_object(group_shapes_lists, local_shapes, group=group) + _SEQ_DATA_META_SHAPES[self.name] = group_shapes_lists + else: + _SEQ_DATA_META_SHAPES[self.name] = [local_shapes] * sp_world + _SEQ_DATA_META_DTYPES[self.name] = local_dtypes + shapes = _SEQ_DATA_META_SHAPES[self.name][local_step] + dtypes = _SEQ_DATA_META_DTYPES[self.name] + buf_id = local_step % 2 + if local_step == 0: + sync_data = ( + local_result + if is_src + else _construct_broadcast_buffer(shapes, dtypes, device) + ) + _broadcast_data(sync_data, shapes, dtypes, src_rank, group, False) + _SEQ_DATA_BUF[self.name][buf_id] = sync_data + + # wait for async comm ops + if _SEQ_DATA_ASYNC_COMMS[self.name]: + for comm in _SEQ_DATA_ASYNC_COMMS[self.name]: + comm.wait() + # before return the sync result, do async broadcast for next batch + if local_step < sp_world - 1: + next_buf_id = 1 - buf_id + shapes = _SEQ_DATA_META_SHAPES[self.name][local_step + 1] + src_rank = dist.get_global_rank(group, local_step + 1) + is_src = sp_rank == local_step + 1 + next_sync_data = ( + _SEQ_DATA_BUF[self.name][-1] + if is_src + else _construct_broadcast_buffer(shapes, dtypes, device) + ) + _SEQ_DATA_ASYNC_COMMS[self.name] = _broadcast_data( + next_sync_data, shapes, dtypes, src_rank, group, True + ) + _SEQ_DATA_BUF[self.name][next_buf_id] = next_sync_data + yield _SEQ_DATA_BUF[self.name][buf_id] + + +sync_inputs = SPDistForward(name="bef_fwd", comm_shape=True) + + +def sync_data(data, sp_idx, name="tmp"): + group = get_sequence_parallel_group() + if group is None: + return data + # if sp_idx in _SYNC_BUFFER[name]: + # return _SYNC_BUFFER[name][sp_idx] + sp_rank = get_sequence_parallel_rank() + src_rank = dist.get_global_rank(group, sp_idx) + objects = [data] if sp_rank == sp_idx else [None] + dist.broadcast_object_list(objects, src=src_rank, group=group) + # _SYNC_BUFFER[name] = {sp_idx: objects[0]} + return objects[0] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/half_precision_fixes.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/half_precision_fixes.py new file mode 100644 index 0000000000000000000000000000000000000000..961f9f96e6c056293dc88af57f663ec7520435e0 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/half_precision_fixes.py @@ -0,0 +1,160 @@ +""" +Half Precision Compatibility Fixes for SeedVR2 + +Provides safe wrappers around PyTorch operations that may fail with reduced +precision dtypes (float16, bfloat16, float8). These utilities automatically +detect precision issues and apply temporary conversions when needed, then +restore original dtypes. +""" + +import torch +import torch.nn.functional as F +from typing import Optional, Tuple, Union + + +def safe_pad_operation( + x: torch.Tensor, + padding: Union[Tuple[int, ...], int], + mode: str = 'constant', + value: float = 0.0 +) -> torch.Tensor: + """ + Safe padding operation with automatic float16 compatibility handling. + + Certain padding modes ('replicate', 'reflect', 'circular') are not implemented + for float16 on some backends. This function automatically detects failures and + applies a temporary float32 conversion, then restores the original dtype. + + Args: + x: Input tensor of any dtype + padding: Padding specification (left, right, top, bottom, front, back) + mode: Padding mode - 'constant', 'replicate', 'reflect', or 'circular' + value: Fill value for 'constant' mode (default: 0.0) + + Returns: + Padded tensor in original dtype + """ + # Modes that may require float16 compatibility fixes + problematic_modes = ['replicate', 'reflect', 'circular'] + + if mode in problematic_modes: + try: + return F.pad(x, padding, mode=mode, value=value) + except RuntimeError as e: + if "not implemented for 'Half'" in str(e): + original_dtype = x.dtype + result = F.pad(x.float(), padding, mode=mode, value=value) + return result.to(original_dtype) + else: + raise + else: + # 'constant' and other compatible modes work natively + return F.pad(x, padding, mode=mode, value=value) + + +def safe_interpolate_operation( + x: torch.Tensor, + size: Optional[Union[int, Tuple[int, ...]]] = None, + scale_factor: Optional[Union[float, Tuple[float, ...]]] = None, + mode: str = 'nearest', + align_corners: Optional[bool] = None, + recompute_scale_factor: Optional[bool] = None +) -> torch.Tensor: + """ + Safe interpolation operation with automatic float16 compatibility handling. + + Interpolation modes like 'bilinear', 'bicubic', and 'trilinear' may not be + implemented for float16 on some backends. This function automatically detects + failures and applies a temporary float32 conversion, then restores the original dtype. + + Args: + x: Input tensor of any dtype + size: Target output size (height, width) or (depth, height, width) + scale_factor: Multiplier for spatial size (alternative to size) + mode: Interpolation mode - 'nearest', 'linear', 'bilinear', 'bicubic', 'trilinear' + align_corners: If True, align corner pixels of input and output + recompute_scale_factor: Recompute scale_factor for use in interpolation + + Returns: + Interpolated tensor in original dtype + """ + # Modes that may require float16 compatibility fixes + problematic_modes = ['bilinear', 'bicubic', 'trilinear'] + + if mode in problematic_modes: + try: + return F.interpolate( + x, + size=size, + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners, + recompute_scale_factor=recompute_scale_factor + ) + except RuntimeError as e: + # Check for float16 incompatibility errors + if ("not implemented for 'Half'" in str(e) or + "compute_indices_weights" in str(e)): + original_dtype = x.dtype + result = F.interpolate( + x.float(), + size=size, + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners, + recompute_scale_factor=recompute_scale_factor + ) + return result.to(original_dtype) + else: + raise + else: + # 'nearest' and other compatible modes work natively + return F.interpolate( + x, + size=size, + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners, + recompute_scale_factor=recompute_scale_factor + ) + + +def ensure_float32_precision( + tensor: torch.Tensor, + force_float32: bool = True +) -> Tuple[torch.Tensor, torch.dtype]: + """ + Ensure tensor is in float32 for precision-sensitive operations. + + Many numerical operations require full precision to avoid accumulated errors: + - Color space conversions (RGBโ†”LAB, RGBโ†”HSV) with matrix multiplications + - Statistical operations (mean, variance, standard deviation) + - Edge detection using derivative filters (Sobel, Canny) + - Histogram matching and CDF computations + - Guided filtering with covariance calculations + + This function upgrades reduced precision dtypes (float16, bfloat16, float8) + to float32, while preserving the original dtype for restoration after computation. + + Args: + tensor: Input tensor of any dtype + force_float32: If True, convert reduced precision dtypes to float32. + If False, return tensor unchanged (useful for disabling conversion) + + Returns: + Tuple of (converted_tensor, original_dtype) for easy restoration: + - converted_tensor: Tensor in float32 (or original dtype if already full precision) + - original_dtype: Original tensor dtype for restoration via .to(original_dtype) + """ + original_dtype = tensor.dtype + + # Skip conversion if disabled + if not force_float32: + return tensor, original_dtype + + # Convert reduced precision dtypes to float32 + if original_dtype not in (torch.float32, torch.float64): + return tensor.float(), original_dtype + + # Already full precision - return as-is (no copy) + return tensor, original_dtype \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/logger.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/logger.py new file mode 100644 index 0000000000000000000000000000000000000000..ada368e79c38d3300ebfe703eae593ac22572467 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/logger.py @@ -0,0 +1,44 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Logging utility functions. +""" + +import logging +import sys +from typing import Optional + +from .distributed import get_global_rank, get_local_rank, get_world_size + +_default_handler = logging.StreamHandler(sys.stdout) +_default_handler.setFormatter( + logging.Formatter( + "%(asctime)s " + + (f"[Rank:{get_global_rank()}]" if get_world_size() > 1 else "") + + (f"[LocalRank:{get_local_rank()}]" if get_world_size() > 1 else "") + + "[%(threadName).12s][%(name)s][%(levelname).5s] " + + "%(message)s" + ) +) + + +def get_logger(name: Optional[str] = None) -> logging.Logger: + """ + Get a logger. + """ + logger = logging.getLogger(name) + logger.addHandler(_default_handler) + logger.setLevel(logging.INFO) + return logger diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/partition.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/partition.py new file mode 100644 index 0000000000000000000000000000000000000000..648c87fe2a61294c09704b9af3e47f5a8570c215 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/partition.py @@ -0,0 +1,59 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Partition utility functions. +""" + +from typing import Any, List + + +def partition_by_size(data: List[Any], size: int) -> List[List[Any]]: + """ + Partition a list by size. + When indivisible, the last group contains fewer items than the target size. + + Examples: + - data: [1,2,3,4,5] + - size: 2 + - return: [[1,2], [3,4], [5]] + """ + assert size > 0 + return [data[i : (i + size)] for i in range(0, len(data), size)] + + +def partition_by_groups(data: List[Any], groups: int) -> List[List[Any]]: + """ + Partition a list by groups. + When indivisible, some groups may have more items than others. + + Examples: + - data: [1,2,3,4,5] + - groups: 2 + - return: [[1,3,5], [2,4]] + """ + assert groups > 0 + return [data[i::groups] for i in range(groups)] + + +def shift_list(data: List[Any], n: int) -> List[Any]: + """ + Rotate a list by n elements. + + Examples: + - data: [1,2,3,4,5] + - n: 3 + - return: [4,5,1,2,3] + """ + return data[(n % len(data)) :] + data[: (n % len(data))] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/seed.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/seed.py new file mode 100644 index 0000000000000000000000000000000000000000..f00e33db6ceed7fe96c8024586dc559910319d3d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/common/seed.py @@ -0,0 +1,30 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import random +from typing import Optional +import numpy as np +import torch + +from .distributed import get_global_rank + + +def set_seed(seed: Optional[int], same_across_ranks: bool = False): + """Function that sets the seed for pseudo-random number generators.""" + if seed is not None: + seed += get_global_rank() if not same_across_ranks else 0 + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/alpha_upscaling.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/alpha_upscaling.py new file mode 100644 index 0000000000000000000000000000000000000000..6059c7da52213f22d86cad8616c0d7169e59d649 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/alpha_upscaling.py @@ -0,0 +1,438 @@ +""" +Alpha Channel Upscaling Module + +Provides edge-guided Alpha upscaling that matches RGB structural changes +without diffusion artifacts. Uses joint bilateral filtering and edge detection +to maintain alignment between upscaled RGB and Alpha channels. +""" + +import torch +import torch.nn.functional as F +import cv2 +import numpy as np +from typing import Optional, Any, List +from ..common.half_precision_fixes import ensure_float32_precision +from ..optimization.memory_manager import manage_tensor + + +def process_alpha_for_batch( + rgb_samples: List[torch.Tensor], + alpha_original: torch.Tensor, + rgb_original: torch.Tensor, + device: torch.device, + compute_dtype: torch.dtype, + debug: Optional['Debug'] = None +) -> List[torch.Tensor]: + """ + Process Alpha channel for an entire batch with temporal consistency. + Called during postprocess phase when VRAM is available. + + Args: + rgb_samples: List of decoded RGB samples from VAE + alpha_original: Original Alpha channel (C, T, H, W) + rgb_original: Original RGB for guidance (C, T, H, W) + device: Target device for processing (typically CUDA) + compute_dtype: Pipeline compute dtype (e.g., bfloat16) for final output. + debug: Debug instance for logging + + Returns: + List of RGBA samples with Alpha merged, in compute_dtype + """ + # Move alpha and RGB guidance tensors to processing device (GPU) + alpha_original = manage_tensor( + tensor=alpha_original, + target_device=device, + tensor_name="alpha_original", + debug=debug, + reason="Alpha processing", + indent_level=1 + ) + rgb_original = manage_tensor( + tensor=rgb_original, + target_device=device, + tensor_name="rgb_original", + debug=debug, + reason="Alpha processing", + indent_level=1 + ) + + # Process each RGB sample individually to handle variable frame counts + processed_samples = [] + + for rgb_sample in rgb_samples: + # Move RGB sample to processing device before use + rgb_sample = manage_tensor( + tensor=rgb_sample, + target_device=device, + tensor_name="rgb_sample", + debug=debug, + reason="Alpha processing", + indent_level=1 + ) + # Handle dimension variations between single frames and video sequences + if rgb_sample.ndim == 3: + # Single frame: (C, H, W) -> need to add T dimension + rgb_sample_4d = rgb_sample.unsqueeze(0) # (C, H, W) -> (1, C, H, W) + alpha_4d = alpha_original.unsqueeze(1) if alpha_original.ndim == 3 else alpha_original + rgb_guide_4d = rgb_original.unsqueeze(1) if rgb_original.ndim == 3 else rgb_original + was_single_frame = True + else: + # Video: (T, C, H, W) -> already correct + rgb_sample_4d = rgb_sample + alpha_4d = alpha_original + rgb_guide_4d = rgb_original + was_single_frame = False + + # Apply edge-guided upscaling to match RGB resolution while preserving alpha structure + alpha_upscaled = edge_guided_alpha_upscale( + input_alpha=alpha_4d[:, :1, :, :] if alpha_4d.shape[1] > 1 else alpha_4d, # Ensure (T, 1, H, W) + input_rgb=rgb_guide_4d[:, :3, :, :], # Ensure (T, 3, H, W) + upscaled_rgb=rgb_sample_4d[:, :3, :, :], # Ensure (T, 3, H, W) + method='guided', + debug=debug + ) + + # Convert Alpha from float32 to compute dtype + if alpha_upscaled.dtype != compute_dtype: + alpha_upscaled = manage_tensor( + tensor=alpha_upscaled, + target_device=alpha_upscaled.device, + tensor_name="alpha_upscaled", + dtype=compute_dtype, + debug=debug, + reason="dtype alignment for RGBA concatenation", + indent_level=1 + ) + + # Concatenate RGB and upscaled alpha to create RGBA output (T, 4, H, W) + rgba_sample = torch.cat([rgb_sample_4d[:, :3, :, :], alpha_upscaled], dim=1) + + # Restore original format + if was_single_frame: + rgba_sample = rgba_sample.squeeze(0) # (1, 4, H, W) -> (4, H, W) + + processed_samples.append(rgba_sample) + + # Release memory immediately after processing each sample + del alpha_upscaled + + # Clean up + del alpha_original, rgb_original + + return processed_samples + + +def detect_edges_batch( + images: torch.Tensor, + method: str = 'sobel', + debug: Optional['Debug'] = None + ) -> torch.Tensor: + """ + Detect edges in a batch of images using Sobel or Canny. + + Args: + images: Tensor of shape (T, C, H, W) in range [-1, 1] or [0, 1] + method: 'sobel' or 'canny' + debug: Optional debug instance for logging + + Returns: + Edge map tensor of shape (T, 1, H, W) in range [0, 1] + """ + # Convert to float32 for OpenCV processing (will be converted to numpy anyway) + images, images_dtype = ensure_float32_precision(images, force_float32=True) + + images_np = images.cpu().numpy() + T, C, H, W = images_np.shape + + # Denormalize if needed + if images_np.min() < 0: + images_np = (images_np + 1) / 2 + + # Convert to 0-255 uint8 + images_np = (images_np * 255).clip(0, 255).astype(np.uint8) + + edges = [] + for t in range(T): + frame = images_np[t].transpose(1, 2, 0) + + if C == 3: + gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) + else: + gray = frame[..., 0] + + if method == 'sobel': + sobelx = cv2.Sobel(gray, cv2.CV_64F, 1, 0, ksize=3) + sobely = cv2.Sobel(gray, cv2.CV_64F, 0, 1, ksize=3) + edge = np.sqrt(sobelx**2 + sobely**2) + edge = (edge / edge.max() * 255).astype(np.uint8) + else: + edge = cv2.Canny(gray, 50, 150) + + edges.append(edge) + + edges = np.stack(edges) + edges = torch.from_numpy(edges).float() / 255.0 + edges = edges.unsqueeze(1) + + # Move back to images device after CPU numpy processing + edges = manage_tensor( + tensor=edges, + target_device=images.device, + tensor_name="edge_map", + dtype=images_dtype, + debug=debug, + reason="edge detection result", + indent_level=1 + ) + + return edges + + +def guided_filter_pytorch(guide: torch.Tensor, src: torch.Tensor, + radius: int = 8, eps: float = 0.01) -> torch.Tensor: + """ + Apply guided filter for edge-preserving smoothing. + + Args: + guide: Guidance image (T, C, H, W) + src: Input to filter (T, 1, H, W) + radius: Filter radius + eps: Regularization parameter + + Returns: + Filtered output (T, 1, H, W) + """ + # Convert to float32 for numerical stability in statistical operations + guide, guide_dtype = ensure_float32_precision(guide, force_float32=True) + src, _ = ensure_float32_precision(src, force_float32=True) + + T, C, H, W = guide.shape + + # Convert to grayscale if RGB + if C == 3: + guide_gray = guide.mean(dim=1, keepdim=True) + else: + guide_gray = guide + + # Apply guided filter + output = _apply_guided_filter(guide_gray, src, radius, eps) + + # Restore original dtype after float32 processing + if output.dtype != guide_dtype: + output = manage_tensor( + tensor=output, + target_device=output.device, + tensor_name="filtered_output", + dtype=guide_dtype, + debug=debug, + reason="dtype restoration" + ) + + return output + + +def _apply_guided_filter(guide_gray: torch.Tensor, src: torch.Tensor, + radius: int, eps: float) -> torch.Tensor: + """ + Apply guided filter algorithm. + + Args: + guide_gray: Grayscale guidance image (T, 1, H, W) + src: Source image to filter (T, 1, H, W) + radius: Filter radius in pixels + eps: Regularization parameter for edge preservation + + Returns: + Filtered output tensor (T, 1, H, W) + + Note: + Uses box filtering for efficient implementation + """ + + def box_filter(x, r): + return F.avg_pool2d(x, kernel_size=2*r+1, stride=1, padding=r) + + # Compute means + mean_guide = box_filter(guide_gray, radius) + mean_src = box_filter(src, radius) + + # Compute correlations + corr_guide = box_filter(guide_gray * guide_gray, radius) + corr_guide_src = box_filter(guide_gray * src, radius) + + # Compute variance and covariance + var_guide = corr_guide - mean_guide * mean_guide + cov_guide_src = corr_guide_src - mean_guide * mean_src + + del corr_guide, corr_guide_src + + # Linear coefficients + a = cov_guide_src / (var_guide + eps) + b = mean_src - a * mean_guide + + del cov_guide_src, var_guide, mean_guide, mean_src + + # Average coefficients + mean_a = box_filter(a, radius) + mean_b = box_filter(b, radius) + + del a, b + + # Apply filter + output = mean_a * guide_gray + mean_b + + del mean_a, mean_b + + return output + + +def edge_guided_alpha_upscale( + input_alpha: torch.Tensor, + input_rgb: torch.Tensor, + upscaled_rgb: torch.Tensor, + method: str = 'guided', + debug: Optional['Debug'] = None +) -> torch.Tensor: + """ + Upscale Alpha channel using RGB edge structure as guidance. + Creates tight, smooth transitions aligned with RGB edges. + + Args: + input_alpha: Original Alpha (T, 1, H_in, W_in) in [0, 1] + input_rgb: Original RGB (T, 3, H_in, W_in) in [-1, 1] or [0, 1] + upscaled_rgb: Upscaled RGB (T, 3, H_out, W_out) in [-1, 1] or [0, 1] + method: 'guided' or 'bilateral' + debug: Debug instance + + Returns: + Upscaled Alpha (T, 1, H_out, W_out) in [0, 1] + """ + T, _, H_out, W_out = upscaled_rgb.shape + device = input_alpha.device + dtype = input_alpha.dtype + + # Convert to float32 for numerical stability in edge detection and filtering + input_alpha, alpha_dtype = ensure_float32_precision(input_alpha, force_float32=True) + upscaled_rgb, _ = ensure_float32_precision(upscaled_rgb, force_float32=True) + input_rgb, _ = ensure_float32_precision(input_rgb, force_float32=True) + + # Analyze alpha distribution to detect binary masks vs gradient alphas + alpha_flat = input_alpha.flatten() + near_zero = (alpha_flat < 0.1).sum().float() + near_one = (alpha_flat > 0.9).sum().float() + binary_ratio = (near_zero + near_one) / alpha_flat.numel() + is_binary_mask = binary_ratio > 0.95 + + if debug: + debug.log(f"Alpha type: {'binary mask' if is_binary_mask else 'gradient alpha'}", category="alpha", indent_level=1) + debug.log(f"Binary ratio: {binary_ratio:.2%}", category="alpha", indent_level=1) + debug.log("Creating tight edge-aligned transitions", category="alpha", indent_level=1) + + # Normalize RGB from [-1, 1] to [0, 1] for edge detection + rgb_normalized = upscaled_rgb.clone() + if rgb_normalized.min() < 0: + rgb_normalized = (rgb_normalized + 1) / 2 + + # Detect edges in upscaled RGB + rgb_edges = detect_edges_batch(images=rgb_normalized, method='sobel', debug=debug) + + # Step 1: Initial bicubic upscale provides smooth base before edge refinement + # MPS on PyTorch < 2.8 doesn't support bicubic+antialias - use CPU fallback + try: + alpha_upscaled = F.interpolate( + input_alpha, + size=(H_out, W_out), + mode='bicubic', + 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: + debug.log("Applying tight edge-aware refinement", category="alpha", indent_level=1) + + # Step 2: Single-pass guided filter with small radius for tight edges + alpha_refined = guided_filter_pytorch( + guide=rgb_normalized, + src=alpha_upscaled, + radius=2, # Reduced from 3 for tighter edges + eps=0.002 + ) + + # Step 3: Create tight transition zone using 3x3 max pooling on edge map + edge_map = rgb_edges + transition_zone = F.max_pool2d(edge_map, kernel_size=3, stride=1, padding=1) + + # Step 4: Identify solid regions (far from edges) vs transition regions (near edges) + solid_threshold = 0.05 # Reduced from 0.08 + is_solid = transition_zone < solid_threshold + + # Convert solid regions to binary (0 or 1) based on 0.5 threshold + alpha_binary = (alpha_refined > 0.5).float() + + # Step 5: Apply strong contrast enhancement in narrow transition zones to complete gradients + in_edge_region = (transition_zone >= solid_threshold) & (transition_zone < 0.25) # Narrow band + + # Apply strong sigmoid (strength=12.0) to push alpha values toward binary endpoints + contrast_enhanced = torch.sigmoid((alpha_refined - 0.5) * 12.0) # Increased from 7.0 + + # Blend original and enhanced alpha based on edge strength for smooth transitions + edge_strength = torch.clamp(edge_map / 0.25, 0, 1) + alpha_in_edges = alpha_refined * (1 - edge_strength) + contrast_enhanced * edge_strength + + # Step 6: Combine solid binary regions with enhanced transition regions + alpha_combined = torch.where(is_solid, alpha_binary, alpha_in_edges) + + # Step 7: Snap remaining mid-gray values to binary outside tightest edge zone + very_solid = transition_zone < 0.03 # Even tighter threshold + final_binary = (alpha_combined > 0.5).float() + alpha_final = torch.where(very_solid, final_binary, alpha_combined) + + # Step 8: Final cleanup - eliminate mid-gray artifacts, preserving only smooth edges in tight regions + tightest_edges = edge_map > 0.15 + mid_gray = (alpha_final > 0.3) & (alpha_final < 0.7) + should_be_binary = mid_gray & ~tightest_edges + + alpha_final = torch.where( + should_be_binary, + (alpha_final > 0.5).float(), + alpha_final + ) + + else: + # For gradient alphas: single guided filter pass preserves smooth transitions + if debug: + debug.log("Applying guided filter for gradient alpha", category="alpha", indent_level=1) + + alpha_final = alpha_upscaled.clone() + + # Single pass with optimized parameters for gradients + alpha_final = guided_filter_pytorch( + guide=rgb_normalized, + src=alpha_final, + radius=3, + eps=0.002 + ) + + # Clamp output to valid alpha range [0, 1] + alpha_final = alpha_final.clamp(0, 1) + + if debug: + final_values = alpha_final.flatten() + exact_zeros = (final_values < 0.02).sum().item() + exact_ones = (final_values > 0.98).sum().item() + smooth_edges = ((final_values >= 0.02) & (final_values <= 0.98)).sum().item() + mid_grays = ((final_values >= 0.35) & (final_values <= 0.65)).sum().item() + total = final_values.numel() + + debug.log(f"Pure 0s: {exact_zeros/total*100:.1f}% | Pure 1s: {exact_ones/total*100:.1f}% | Smooth edges: {smooth_edges/total*100:.1f}% | Mid-grays: {mid_grays/total*100:.1f}%", category="alpha", indent_level=1) + + return alpha_final \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_phases.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_phases.py new file mode 100644 index 0000000000000000000000000000000000000000..3b7e6ea01537d0932800108a3ead8df4180e6a43 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_phases.py @@ -0,0 +1,1478 @@ +""" +Generation Logic Module for SeedVR2 + +This module implements a four-phase batch processing pipeline for video upscaling: +- Phase 1: Batch VAE encoding of all input frames +- Phase 2: Batch DiT upscaling of all encoded latents +- Phase 3: Batch VAE decoding of all upscaled latents +- Phase 4: Post-processing and final video assembly + +This architecture minimizes model swapping overhead by completing each phase +for all batches before moving to the next phase, significantly improving +performance especially when using model offloading. + +Key Features: +- Four-phase pipeline (encode-all โ†’ upscale-all โ†’ decode-all โ†’ postprocess-all) for efficiency +- Native FP8 pipeline support for 2x speedup and 50% VRAM reduction +- Temporal overlap support for smooth transitions between batches +- Adaptive dtype detection and configuration +- Memory-efficient pre-allocated batch processing +- Stream-based assembly eliminates memory spikes for long videos +- Advanced video format handling (4n+1 constraint) +- Clean separation of concerns with phase-specific resource management +- Each phase handles its own cleanup in finally blocks +""" + +import os +import torch +from typing import Dict, List, Optional, Tuple, Any, Callable + +from .generation_utils import ( + setup_video_transform, + pad_video_temporal, + check_interrupt, + ensure_precision_initialized, + _draw_tile_boundaries, + load_text_embeddings, + blend_overlapping_frames, + calculate_optimal_batch_params, + script_directory +) +from .model_configuration import apply_model_specific_config +from .model_loader import materialize_model +from .alpha_upscaling import process_alpha_for_batch +from .infer import VideoDiffusionInfer +from ..common.seed import set_seed +from ..optimization.memory_manager import ( + cleanup_dit, + cleanup_vae, + cleanup_text_embeddings, + manage_tensor, + manage_model_device, + release_tensor_memory, + release_tensor_collection +) +from ..optimization.performance import ( + optimized_video_rearrange, + optimized_single_video_rearrange, + optimized_sample_to_image_format +) +from ..utils.color_fix import ( + lab_color_transfer, + wavelet_adaptive_color_correction, + hsv_saturation_histogram_match, + wavelet_reconstruction, + adaptive_instance_normalization +) + + +def _prepare_video_batch( + images: torch.Tensor, + start_idx: int, + end_idx: int, + uniform_padding: int = 0, + debug: Optional['Debug'] = None, + log_info: bool = False +) -> torch.Tensor: + """ + Extract and prepare video batch with uniform padding and permutation. + + Args: + images: Source video frames [T, H, W, C] + start_idx: Start frame index + end_idx: End frame index (exclusive) + uniform_padding: Number of frames to pad (0 = no padding) + debug: Debug instance for optional logging + log_info: If True, log padding operations (used during encoding only) + + Returns: + Prepared video in TCHW format + """ + # Extract frames (view/slice, not copy) + video = images[start_idx:end_idx] + + # Apply uniform padding if needed + if uniform_padding > 0: + if log_info and debug: + current_frames = end_idx - start_idx + debug.log(f"Sequence of {current_frames} frames", category="video", force=True, indent_level=1) + debug.log(f"Padding batch: {uniform_padding} frame{'s' if uniform_padding != 1 else ''} added ({current_frames} โ†’ {current_frames + uniform_padding}) for uniform batches", + category="video", force=True, indent_level=1) + video = pad_video_temporal(video, count=uniform_padding, temporal_dim=0, prepend=False, debug=None) + + # Permute to TCHW format + video = video.permute(0, 3, 1, 2) + + return video + + +def _apply_4n1_padding(video: torch.Tensor) -> torch.Tensor: + """ + Apply 4n+1 temporal padding constraint required by VAE. + + Args: + video: Video tensor in TCHW format + + Returns: + Padded video in TCHW format + """ + t = video.size(0) + if t % 4 != 1: + video = optimized_single_video_rearrange(video) # TCHW -> CTHW + video = pad_video_temporal(video, temporal_dim=1, prepend=False, debug=None) + video = optimized_single_video_rearrange(video) # CTHW -> TCHW + return video + + +def _reconstruct_and_transform_batch( + ctx: Dict[str, Any], + batch_idx: int, + debug: Optional['Debug'] = None +) -> torch.Tensor: + """ + Reconstruct and transform a video batch for color correction (Phase 4). + + Args: + ctx: Context with input_images, batch_metadata, video_transform + batch_idx: Index of batch to reconstruct + debug: Debug instance for logging + + Returns: + Transformed video in CTHW format, ready for color correction + """ + start_idx, end_idx, uniform_padding = ctx['batch_metadata'][batch_idx] + + # Prepare video batch + video = _prepare_video_batch( + images=ctx['input_images'], + start_idx=start_idx, + end_idx=end_idx, + uniform_padding=uniform_padding, + debug=None, + log_info=False + ) + + # Apply 4n+1 padding using shared helper + video = _apply_4n1_padding(video) + + # Extract RGB and transform + if ctx.get('is_rgba', False): + rgb_video = video[:, :3, :, :] + else: + rgb_video = video + + transformed_video = ctx['video_transform'](rgb_video) + + del video + + return transformed_video + + +def encode_all_batches( + runner: 'VideoDiffusionInfer', + ctx: Dict[str, Any], + images: torch.Tensor, + debug: 'Debug', + batch_size: int = 5, + uniform_batch_size: bool = False, + seed: int = 42, + progress_callback: Optional[Callable[[int, int, int, str], None]] = None, + temporal_overlap: int = 0, + resolution: int = 1080, + max_resolution: int = 0, + input_noise_scale: float = 0.0, + color_correction: str = "wavelet" +) -> Dict[str, Any]: + """ + Phase 1: VAE Encoding for all batches + + Encodes video frames to latents in batches, handling temporal overlap and + memory optimization. Creates context automatically if not provided. + + Args: + runner: VideoDiffusionInfer instance with loaded models (required) + ctx: Generation context from setup_generation_context (required) + images: Input frames tensor [T, H, W, C] range [0,1] (required) + debug: Debug instance for logging (required) + batch_size: Frames per batch (4n+1 format: 1, 5, 9, 13...) + uniform_batch_size: Pad final batch to match batch_size for uniform batches + seed: Random seed for deterministic VAE sampling (default: 42) + progress_callback: Optional callback(current, total, frames, phase_name) + temporal_overlap: Overlapping frames between batches for continuity + resolution: Target resolution for shortest edge + max_resolution: Maximum resolution for any edge (0 = no limit) + input_noise_scale: Scale for input noise (0.0-1.0). Adds noise to input images + before VAE encoding to reduce artifacts at high resolutions. + color_correction: Color correction method - "wavelet", "adain", or "none" (default: "wavelet") + Determines if transformed videos need to be stored for later use. + + Returns: + dict: Context containing: + - batch_metadata: Lightweight indices for on-demand transform reconstruction + - all_latents: List of encoded latents ready for upscaling + - Other state for subsequent phases + + Raises: + ValueError: If required inputs are missing or invalid + RuntimeError: If encoding fails + """ + if debug is None: + raise ValueError("Debug instance must be provided to encode_all_batches") + + debug.log("", category="none", force=True) + debug.log("โ”โ”โ”โ”โ”โ”โ”โ” Phase 1: VAE encoding โ”โ”โ”โ”โ”โ”โ”โ”", category="none", force=True) + debug.start_timer("phase1_encoding") + + # Context must be provided + if ctx is None: + raise ValueError("Generation context must be provided to encode_all_batches") + + # Validate and store inputs + if images is None: + raise ValueError("Images to encode must be provided") + else: + # MPS: keep on device to avoid sync overhead in Phase 4 color correction + if ctx['vae_device'].type == 'mps' and images.device.type != 'mps': + ctx['input_images'] = images.to(ctx['vae_device']) + else: + ctx['input_images'] = images + + # Get total frame count from context (set in video_upscaler before encoding) + total_frames = ctx.get('total_frames', len(images)) + + # Set it if not already set (for standalone/CLI usage) + if 'total_frames' not in ctx: + ctx['total_frames'] = total_frames + + if total_frames == 0: + raise ValueError("No frames to process") + + # Setup video transformation pipeline and compute dimensions if not already done + if 'true_target_dims' not in ctx: + sample_frame = images[0].permute(2, 0, 1).unsqueeze(0) + setup_video_transform(ctx, resolution, max_resolution, debug, sample_frame) + del sample_frame + else: + setup_video_transform(ctx, resolution, max_resolution, debug) + + # Detect if input is RGBA (4 channels) + ctx['is_rgba'] = images[0].shape[-1] == 4 + + # Display batch optimization tip if applicable + if total_frames > 0: + batch_params = calculate_optimal_batch_params(total_frames, batch_size, temporal_overlap) + if batch_params['best_batch'] != batch_size and batch_params['best_batch'] <= total_frames: + debug.log("", category="none", force=True) + debug.log(f"Tip: For {total_frames} frames, batch_size={batch_params['best_batch']} matches video length optimally", category="tip", force=True) + debug.log(f"Matching batch_size to shot length improves temporal coherence", category="tip", force=True, indent_level=1) + debug.log("", category="none", force=True) + + # Calculate batching parameters + step = batch_size - temporal_overlap if temporal_overlap > 0 else batch_size + if step <= 0: + step = batch_size + temporal_overlap = 0 + debug.log(f"temporal_overlap >= batch_size, resetting to 0", level="WARNING", category="setup", force=True) + + # Store actual temporal overlap used (may differ from parameter if reset) + ctx['actual_temporal_overlap'] = temporal_overlap + + # Calculate number of batches + num_encode_batches = 0 + for idx in range(0, total_frames, step): + end_idx = min(idx + batch_size, total_frames) + if idx > 0 and end_idx - idx <= temporal_overlap: + break + num_encode_batches += 1 + + # Pre-allocate lists for memory efficiency + ctx['all_latents'] = [None] * num_encode_batches + ctx['all_ori_lengths'] = [None] * num_encode_batches + if color_correction != "none": + ctx['batch_metadata'] = [None] * num_encode_batches + + encode_idx = 0 + + try: + # Materialize VAE if still on meta device + if runner.vae and next(runner.vae.parameters()).device.type == 'meta': + materialize_model(runner, "vae", ctx['vae_device'], runner.config, debug) + else: + # Model already materialized (cached) - apply any pending configs if needed + if getattr(runner, '_vae_config_needs_application', False): + debug.log("Applying updated VAE configuration", category="vae", force=True) + apply_model_specific_config(runner.vae, runner, runner.config, False, debug) + + # Initialize precision after VAE is materialized with actual weights + ensure_precision_initialized(ctx, runner, debug) + + # Cache VAE now that it's fully configured and ready for inference + if ctx['cache_context']['vae_cache'] and not ctx['cache_context']['cached_vae']: + runner.vae._model_name = ctx['cache_context']['vae_model'] + ctx['cache_context']['global_cache'].set_vae( + {'node_id': ctx['cache_context']['vae_id'], 'cache_model': True}, + runner.vae, ctx['cache_context']['vae_model'], debug + ) + ctx['cache_context']['vae_newly_cached'] = True + + # If both models now cached, cache runner template + dit_is_cached = ctx['cache_context']['cached_dit'] or ctx['cache_context']['dit_newly_cached'] + if dit_is_cached: + ctx['cache_context']['global_cache'].set_runner( + ctx['cache_context']['dit_id'], ctx['cache_context']['vae_id'], + runner, debug + ) + + # Set deterministic seed for VAE encoding (separate from diffusion noise) + # Uses seed + 1,000,000 to avoid collision with upscaling batch seeds + # This ensures VAE sampling is deterministic while maintaining quality + seed_vae = seed + 1000000 + set_seed(seed_vae) + debug.log(f"Using seed: {seed_vae} (VAE uses seed+1000000 for deterministic sampling)", category="vae") + + # Move VAE to GPU for encoding (no-op if already there) + manage_model_device(model=runner.vae, target_device=ctx['vae_device'], + model_name="VAE", debug=debug, runner=runner) + + debug.log_memory_state("After VAE loading for encoding", detailed_tensors=False) + + # Initialize tile_boundaries for encoding debug + if runner.tile_debug == "encode" and runner.encode_tiled: + debug.encode_tile_boundaries = [] + debug.log("Tile debug enabled: encode tile boundaries will be visualized", category="vae", force=True) + debug.log("Remember to disable --tile_debug in production to remove overlay visualization", category="tip", indent_level=1, force=True) + + # Process encoding + for batch_idx in range(0, total_frames, step): + check_interrupt(ctx) + + # Calculate indices with temporal overlap + if batch_idx == 0: + start_idx = 0 + end_idx = min(batch_size, total_frames) + else: + start_idx = batch_idx + end_idx = min(start_idx + batch_size, total_frames) + if end_idx - start_idx <= temporal_overlap: + break + + current_frames = end_idx - start_idx + is_uniform_padding = uniform_batch_size and current_frames < batch_size + + debug.log(f"Encoding batch {encode_idx+1}/{num_encode_batches}", category="vae", force=True) + debug.start_timer(f"encode_batch_{encode_idx+1}") + + # Save original length before any padding + ori_length = current_frames + + # Prepare video batch with uniform padding + video = _prepare_video_batch( + images=images, + start_idx=start_idx, + end_idx=end_idx, + uniform_padding=batch_size - current_frames if is_uniform_padding else 0, + debug=debug, + log_info=True + ) + if is_uniform_padding: + current_frames = batch_size + + video = manage_tensor( + tensor=video, + target_device=ctx['vae_device'], + tensor_name=f"video_batch_{encode_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="VAE encoding", + indent_level=1 + ) + + # Check temporal dimension for 4n+1 padding + t = video.size(0) + + # Log sequence size if not already logged (for non-uniform batches) + if not is_uniform_padding: + debug.log(f"Sequence of {t} frames", category="video", force=True, indent_level=1) + + # Apply 4n+1 padding using shared helper + if t % 4 != 1: + target = ((t-1)//4+1)*4+1 + padding_frames = target - t + debug.log(f"Padding batch: {padding_frames} frame{'s' if padding_frames != 1 else ''} added ({t} โ†’ {target}) to meet 4n+1 constraint", + category="video", force=True, indent_level=1) + # Apply 4n+1 padding to match exact frame count from encoding + video = _apply_4n1_padding(video) + + # Apply transformations (matches reconstruction logic) + if ctx.get('is_rgba', False): + debug.log(f"Extracted Alpha channel for edge-guided upscaling", category="alpha", indent_level=1) + rgb_video = video[:, :3, :, :] + else: + rgb_video = video + + transformed_video = ctx['video_transform'](rgb_video) + + # Apply input noise if requested (to reduce artifacts at high resolutions) + if input_noise_scale > 0: + debug.log(f"Applying input noise (scale: {input_noise_scale:.2f})", category="video", indent_level=1) + + # Generate noise matching the video shape + noise = torch.randn_like(transformed_video) + + # Subtle noise amplitude + noise = noise * 0.05 + + # Linear blend factor: 0 at scale=0, 0.5 at scale=1 + blend_factor = input_noise_scale * 0.5 + + # Apply blend + transformed_video = transformed_video * (1 - blend_factor) + (transformed_video + noise) * blend_factor + + del noise + + # Store original length for proper trimming later + ctx['all_ori_lengths'][encode_idx] = ori_length + + # Store batch frame indices for on-demand reconstruction + if color_correction != "none": + ctx['batch_metadata'][encode_idx] = (start_idx, end_idx, batch_size - ori_length if is_uniform_padding else 0) + + # Extract and store Alpha and RGB from padded original video (before encoding) + if ctx.get('is_rgba', False): + if 'all_alpha_channels' not in ctx: + ctx['all_alpha_channels'] = [None] * num_encode_batches + if 'all_input_rgb' not in ctx: + ctx['all_input_rgb'] = [None] * num_encode_batches + + # Extract from padded RGBA video (format: T, 4, H, W) + alpha_channel = video[:, 3:4, :, :] + rgb_video_original = video[:, :3, :, :] + + # Store on tensor_offload_device to save VRAM (or keep on device if none) + if ctx['tensor_offload_device'] is not None: + ctx['all_alpha_channels'][encode_idx] = manage_tensor( + tensor=alpha_channel, + target_device=ctx['tensor_offload_device'], + tensor_name=f"alpha_channel_{encode_idx+1}", + debug=debug, + reason="storing Alpha channel for upscaling", + indent_level=1 + ) + ctx['all_input_rgb'][encode_idx] = manage_tensor( + tensor=rgb_video_original, + target_device=ctx['tensor_offload_device'], + tensor_name=f"rgb_original_{encode_idx+1}", + debug=debug, + reason="storing RGB edge guidance for Alpha upscaling", + indent_level=1 + ) + else: + ctx['all_alpha_channels'][encode_idx] = alpha_channel + ctx['all_input_rgb'][encode_idx] = rgb_video_original + + del alpha_channel, rgb_video_original + + del video + + # Move to VAE device with correct dtype for encoding + transformed_video = manage_tensor( + tensor=transformed_video, + target_device=ctx['vae_device'], + tensor_name=f"transformed_video_{encode_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="VAE encoding", + indent_level=1 + ) + + # Encode to latents + cond_latents = runner.vae_encode([transformed_video]) + + # Don't store transformed_video - will reconstruct on-demand in Phase 4 + del transformed_video, rgb_video + + # Convert from VAE dtype to compute dtype and offload to avoid VRAM accumulation + if ctx['tensor_offload_device'] is not None and (cond_latents[0].is_cuda or cond_latents[0].is_mps): + ctx['all_latents'][encode_idx] = manage_tensor( + tensor=cond_latents[0], + target_device=ctx['tensor_offload_device'], + tensor_name=f"latent_{encode_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="storing encoded latents for upscaling", + indent_level=1 + ) + else: + # Stay on current device but convert to compute dtype + ctx['all_latents'][encode_idx] = manage_tensor( + tensor=cond_latents[0], + target_device=cond_latents[0].device, + tensor_name=f"latent_{encode_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="VAE dtype โ†’ compute dtype", + indent_level=1 + ) + + del cond_latents + + debug.end_timer(f"encode_batch_{encode_idx+1}", f"Encoded batch {encode_idx+1}") + + if progress_callback: + progress_callback(encode_idx+1, num_encode_batches, + current_frames, "Phase 1: Encoding") + + encode_idx += 1 + + except Exception as e: + debug.log(f"Error in Phase 1 (Encoding): {e}", level="ERROR", category="error", force=True) + raise + finally: + # Offload VAE to configured offload device if specified + if ctx['vae_offload_device'] is not None: + manage_model_device(model=runner.vae, target_device=ctx['vae_offload_device'], + model_name="VAE", debug=debug, reason="VAE offload", runner=runner) + + debug.end_timer("phase1_encoding", "Phase 1: VAE encoding complete", show_breakdown=True) + debug.log_memory_state("After phase 1 (VAE encoding)", show_tensors=False) + + return ctx + + +def upscale_all_batches( + runner: 'VideoDiffusionInfer', + ctx: Dict[str, Any], + debug: 'Debug', + progress_callback: Optional[Callable[[int, int, int, str], None]] = None, + seed: int = 42, + latent_noise_scale: float = 0.0, + cache_model: bool = False +) -> Dict[str, Any]: + """ + Phase 2: DiT Upscaling for all encoded batches. + + Processes all encoded latents through the diffusion model for upscaling. + Requires context from encode_all_batches with encoded latents. + + Args: + runner: VideoDiffusionInfer instance with loaded models (required) + ctx: Context from encode_all_batches containing latents (required) + debug: Debug instance for logging (required) + progress_callback: Optional callback(current, total, frames, phase_name) + seed: Random seed for reproducible generation + latent_noise_scale: Noise scale for latent space augmentation (0.0-1.0). + Adds noise during diffusion conditioning. Can soften details + but may help with certain artifacts. 0.0 = no noise (crisp), + 1.0 = maximum noise (softer) + cache_model: If True, keep DiT model for reuse instead of deleting it + + Returns: + dict: Updated context containing: + - all_upscaled_latents: List of upscaled latents ready for decoding + - Preserved state from encoding phase + + Raises: + ValueError: If context is missing or has no encoded latents + RuntimeError: If upscaling fails + """ + if debug is None: + raise ValueError("Debug instance must be provided to upscale_all_batches") + + if ctx is None: + raise ValueError("Context is required for upscale_all_batches. Run encode_all_batches first.") + + # Validate we have encoded latents + if 'all_latents' not in ctx or not ctx['all_latents']: + raise ValueError("No encoded latents found. Run encode_all_batches first.") + + debug.log("", category="none", force=True) + debug.log("โ”โ”โ”โ”โ”โ”โ”โ” Phase 2: DiT upscaling โ”โ”โ”โ”โ”โ”โ”โ”", category="none", force=True) + debug.start_timer("phase2_upscaling") + + # Load text embeddings if not already loaded + if ctx.get('text_embeds') is None: + ctx['text_embeds'] = load_text_embeddings(script_directory, ctx['dit_device'], ctx['compute_dtype'], debug) + debug.log("Loaded text embeddings for DiT", category="dit") + + # Configure diffusion parameters + # Force cfg_scale = 1.0 for one-step distilled models (CFG is incompatible with distillation) + runner.config.diffusion.cfg.scale = 1.0 + runner.config.diffusion.cfg.rescale = 0.0 + runner.config.diffusion.timesteps.sampling.steps = 1 + runner.configure_diffusion(device=ctx['dit_device'], dtype=ctx['compute_dtype']) + + # Count valid latents + num_valid_latents = len([l for l in ctx['all_latents'] if l is not None]) + + # Safety check for empty latents + if num_valid_latents == 0: + debug.log("No valid latents to upscale", level="WARNING", category="dit", force=True) + ctx['all_upscaled_latents'] = [] + return ctx + + # Pre-allocate list for upscaled latents + ctx['all_upscaled_latents'] = [None] * num_valid_latents + + upscale_idx = 0 + + try: + # Materialize DiT if still on meta device + if runner.dit and next(runner.dit.parameters()).device.type == 'meta': + materialize_model(runner, "dit", ctx['dit_device'], runner.config, debug) + else: + # Model already materialized (cached) - apply any pending configs if needed + if getattr(runner, '_dit_config_needs_application', False): + debug.log("Applying updated DiT configuration", category="dit", force=True) + apply_model_specific_config(runner.dit, runner, runner.config, True, debug) + + # Initialize precision after DiT is materialized with actual weights + ensure_precision_initialized(ctx, runner, debug) + + # Cache DiT now that it's fully configured and ready for inference + if ctx['cache_context']['dit_cache'] and not ctx['cache_context']['cached_dit']: + runner.dit._model_name = ctx['cache_context']['dit_model'] + ctx['cache_context']['global_cache'].set_dit( + {'node_id': ctx['cache_context']['dit_id'], 'cache_model': True}, + runner.dit, ctx['cache_context']['dit_model'], debug + ) + ctx['cache_context']['dit_newly_cached'] = True + + # If both models now cached, cache runner template + vae_is_cached = ctx['cache_context']['cached_vae'] or ctx['cache_context']['vae_newly_cached'] + if vae_is_cached: + ctx['cache_context']['global_cache'].set_runner( + ctx['cache_context']['dit_id'], ctx['cache_context']['vae_id'], + runner, debug + ) + + # Move DiT to GPU for upscaling (no-op if already there) + manage_model_device(model=runner.dit, target_device=ctx['dit_device'], + model_name="DiT", debug=debug, runner=runner) + + debug.log_memory_state("After DiT loading for upscaling", detailed_tensors=False) + + for batch_idx, latent in enumerate(ctx['all_latents']): + if latent is None: + continue + + check_interrupt(ctx) + + debug.log(f"Upscaling batch {upscale_idx+1}/{num_valid_latents}", category="generation", force=True) + # Reset seed for each batch to ensure identical RNG state + # This ensures identical inputs produce identical outputs regardless of batch position + set_seed(seed) + debug.log(f"Using seed: {seed} for deterministic generation", category="dit") + + debug.start_timer(f"upscale_batch_{upscale_idx+1}") + + # Move to DiT device with correct dtype for upscaling (no-op if already there) + latent = manage_tensor( + tensor=latent, + target_device=ctx['dit_device'], + tensor_name=f"latent_{upscale_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="DiT upscaling", + indent_level=1 + ) + + # Generate noise (randn_like automatically uses latent's device) + base_noise = torch.randn_like(latent, dtype=ctx['compute_dtype']) + + noises = [base_noise] + aug_noises = [base_noise * 0.1 + torch.randn_like(base_noise) * 0.05] + + # Log latent noise application if enabled + if latent_noise_scale > 0: + debug.log(f"Applying latent noise (scale: {latent_noise_scale:.3f})", category="generation") + + def _add_noise(x, aug_noise): + if latent_noise_scale == 0.0: + return x + t = torch.tensor([1000.0], device=ctx['dit_device'], dtype=ctx['compute_dtype']) * latent_noise_scale + shape = torch.tensor(x.shape[1:], device=ctx['dit_device'])[None] + t = runner.timestep_transform(t, shape) + x = runner.schedule.forward(x, aug_noise, t) + del t, shape + return x + + # Generate condition + condition = runner.get_condition( + noises[0], + task="sr", + latent_blur=_add_noise(latent, aug_noises[0]), + ) + conditions = [condition] + + # Detect DiT model dtype (handle CompatibleDiT wrapper) + dit_model = runner.dit.dit_model if hasattr(runner.dit, 'dit_model') else runner.dit + try: + dit_dtype = next(dit_model.parameters()).dtype + except StopIteration: + dit_dtype = ctx['compute_dtype'] # Fallback for meta device or empty model + + # Use autocast if DiT dtype differs from compute dtype + # Skip autocast on MPS (CompatibleDiT already handles dtype conversion) + debug.start_timer(f"dit_inference_{upscale_idx+1}") + with torch.no_grad(): + if dit_dtype != ctx['compute_dtype'] and ctx['dit_device'].type != 'mps': + with torch.autocast(ctx['dit_device'].type, ctx['compute_dtype'], enabled=True): + upscaled_latents = runner.inference( + noises=noises, + conditions=conditions, + **ctx['text_embeds'], + ) + else: + upscaled_latents = runner.inference( + noises=noises, + conditions=conditions, + **ctx['text_embeds'], + ) + debug.end_timer(f"dit_inference_{upscale_idx+1}", f"DiT inference {upscale_idx+1}") + + # Offload upscaled latents to avoid VRAM accumulation + if ctx['tensor_offload_device'] is not None and (upscaled_latents[0].is_cuda or upscaled_latents[0].is_mps): + ctx['all_upscaled_latents'][upscale_idx] = manage_tensor( + tensor=upscaled_latents[0], + target_device=ctx['tensor_offload_device'], + tensor_name=f"upscaled_latent_{upscale_idx+1}", + debug=debug, + reason="storing upscaled latents for decoding", + indent_level=1 + ) + else: + ctx['all_upscaled_latents'][upscale_idx] = upscaled_latents[0] + + # Free original latent - release tensor memory first + release_tensor_memory(ctx['all_latents'][batch_idx]) + ctx['all_latents'][batch_idx] = None + + del noises, aug_noises, latent, conditions, condition, base_noise, upscaled_latents + + debug.end_timer(f"upscale_batch_{upscale_idx+1}", f"Upscaled batch {upscale_idx+1}") + + if progress_callback: + progress_callback(upscale_idx+1, num_valid_latents, + 1, "Phase 2: Upscaling") + + upscale_idx += 1 + + except Exception as e: + debug.log(f"Error in Phase 2 (Upscaling): {e}", level="ERROR", category="error", force=True) + raise + finally: + # Log BlockSwap summary if it was used + if hasattr(runner, '_blockswap_active') and runner._blockswap_active: + swap_summary = debug.get_swap_summary() + if swap_summary and swap_summary.get('total_swaps', 0) > 0: + total_time = swap_summary.get('block_total_ms', 0) + swap_summary.get('io_total_ms', 0) + debug.log("BlockSwap Summary", category="blockswap") + debug.log(f"BlockSwap overhead: {total_time:.2f}ms", category="blockswap", indent_level=1) + debug.log(f"Total swaps: {swap_summary['total_swaps']}", category="blockswap", indent_level=1) + + # Show block swap details + if 'block_swaps' in swap_summary and swap_summary['block_swaps'] > 0: + avg_ms = swap_summary.get('block_avg_ms', 0) + total_ms = swap_summary.get('block_total_ms', 0) + min_ms = swap_summary.get('block_min_ms', 0) + max_ms = swap_summary.get('block_max_ms', 0) + + debug.log(f"Block swaps: {swap_summary['block_swaps']} " + f"(avg: {avg_ms:.2f}ms, min: {min_ms:.2f}ms, max: {max_ms:.2f}ms, total: {total_ms:.2f}ms)", + category="blockswap", indent_level=1) + + # Show most frequently swapped block + if 'most_swapped_block' in swap_summary: + debug.log(f"Most swapped: Block {swap_summary['most_swapped_block']} " + f"({swap_summary['most_swapped_count']} times)", category="blockswap", indent_level=1) + + # Show I/O swap details if present + if 'io_swaps' in swap_summary and swap_summary['io_swaps'] > 0: + debug.log(f"I/O swaps: {swap_summary['io_swaps']} " + f"(avg: {swap_summary.get('io_avg_ms', 0):.2f}ms, total: {swap_summary.get('io_total_ms', 0):.2f}ms)", + category="blockswap", indent_level=1) + + # Cleanup DiT as it's no longer needed after upscaling + cleanup_dit(runner=runner, debug=debug, cache_model=cache_model) + + # Cleanup text embeddings as they're no longer needed after upscaling + cleanup_text_embeddings(ctx, debug) + + debug.end_timer("phase2_upscaling", "Phase 2: DiT upscaling complete", show_breakdown=True) + debug.log_memory_state("After phase 2 (DiT upscaling)", show_tensors=False) + + return ctx + + +def decode_all_batches( + runner: 'VideoDiffusionInfer', + ctx: Dict[str, Any], + debug: 'Debug', + progress_callback: Optional[Callable[[int, int, int, str], None]] = None, + cache_model: bool = False +) -> Dict[str, Any]: + """ + Phase 3: VAE Decoding. + + Decodes all upscaled latents back to pixel space and writes directly to + pre-allocated final_video tensor. This avoids memory duplication by not + storing intermediate batch_samples. + + Requires context from upscale_all_batches with upscaled latents. + + Args: + runner: VideoDiffusionInfer instance with loaded models (required) + ctx: Context from upscale_all_batches containing upscaled latents (required) + debug: Debug instance for logging (required) + progress_callback: Optional callback(current, total, frames, phase_name) + cache_model: If True, keep VAE model for reuse instead of deleting it + + Returns: + dict: Updated context containing: + - final_video: Pre-allocated tensor with decoded samples (unnormalized, in [-1,1]) + - decode_batch_info: List of (start_idx, end_idx, ori_length) for Phase 4 processing + - VAE cleanup completed + + Raises: + ValueError: If context is missing or has no upscaled latents + RuntimeError: If decoding fails + """ + if debug is None: + raise ValueError("Debug instance must be provided to decode_all_batches") + + if ctx is None: + raise ValueError("Context is required for decode_all_batches. Run upscale_all_batches first.") + + # Validate we have upscaled latents + if 'all_upscaled_latents' not in ctx or not ctx['all_upscaled_latents']: + raise ValueError("No upscaled latents found. Run upscale_all_batches first.") + + debug.log("", category="none", force=True) + debug.log("โ”โ”โ”โ”โ”โ”โ”โ” Phase 3: VAE decoding โ”โ”โ”โ”โ”โ”โ”โ”", category="none", force=True) + debug.start_timer("phase3_decoding") + + # Count valid latents + num_valid_latents = len([l for l in ctx['all_upscaled_latents'] if l is not None]) + num_batches = len([l for l in ctx['all_ori_lengths'] if l is not None]) + + # Get output dimensions from context (set during Phase 1) + if 'true_target_dims' not in ctx: + raise ValueError("true_target_dims not found in context. Run encode_all_batches first.") + true_h, true_w = ctx['true_target_dims'] + total_frames = ctx.get('total_frames', 0) + C = 4 if ctx.get('is_rgba', False) else 3 + + # Pre-allocate final_video at the START of decode phase (before any batch processing) + # This ensures we only need memory for final_video + 1 batch, not final_video + all batch_samples + # MPS: keep on device (unified memory, no benefit to CPU offload) + if ctx['tensor_offload_device'] is not None: + target_device = ctx['tensor_offload_device'] + elif ctx['vae_device'].type == 'mps': + target_device = ctx['vae_device'] + else: + target_device = 'cpu' + channels_str = "RGBA" if C == 4 else "RGB" + required_gb = (total_frames * true_h * true_w * C * 2) / (1024**3) + debug.log(f"Pre-allocating output tensor: {total_frames} frames, {true_w}x{true_h}px, {channels_str} ({required_gb:.2f}GB)", + category="setup", force=True) + + ctx['final_video'] = torch.empty((total_frames, true_h, true_w, C), dtype=ctx['compute_dtype'], device=target_device) + + # Track batch write positions for Phase 4 processing + # Each entry: (write_start, write_end, batch_idx, ori_length) + ctx['decode_batch_info'] = [] + + # Get temporal overlap from context (set during Phase 1) + temporal_overlap = ctx.get('actual_temporal_overlap', 0) + + # Track padding removed for final summary + total_padding_removed = 0 + + current_write_idx = 0 + decode_idx = 0 + + try: + # VAE should already be materialized from encoding phase + if runner.vae and next(runner.vae.parameters()).device.type == 'meta': + materialize_model(runner, "vae", ctx['vae_device'], runner.config, debug) + + # Precision should already be initialized from encoding phase + ensure_precision_initialized(ctx, runner, debug) + + # Move VAE to GPU for decoding (no-op if already there) + manage_model_device(model=runner.vae, target_device=ctx['vae_device'], + model_name="VAE", debug=debug, runner=runner) + + debug.log_memory_state("After VAE loading for decoding", detailed_tensors=False) + + # Initialize tile_boundaries for decoding debug + if runner.tile_debug == "decode" and runner.decode_tiled: + debug.decode_tile_boundaries = [] + debug.log("Tile debug enabled: decode tile boundaries will be visualized", category="vae", force=True) + debug.log("Remember to disable --tile_debug in production to remove overlay visualization", category="tip", indent_level=1, force=True) + + # Process decoding + for batch_idx, upscaled_latent in enumerate(ctx['all_upscaled_latents']): + if upscaled_latent is None: + continue + + check_interrupt(ctx) + + debug.log(f"Decoding batch {decode_idx+1}/{num_valid_latents}", category="vae", force=True) + debug.start_timer(f"decode_batch_{decode_idx+1}") + + # Move to VAE device with correct dtype for decoding (no-op if already there) + upscaled_latent = manage_tensor( + tensor=upscaled_latent, + target_device=ctx['vae_device'], + tensor_name=f"upscaled_latent_{decode_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="VAE decoding", + indent_level=1 + ) + + # Decode latent + debug.start_timer("vae_decode") + samples = runner.vae_decode([upscaled_latent]) + debug.end_timer("vae_decode", "VAE decode") + + # Process samples - get the single decoded sample + debug.start_timer("optimized_video_rearrange") + samples = optimized_video_rearrange(samples) + debug.end_timer("optimized_video_rearrange", "Video rearrange") + + # Get the decoded sample (always single-element list) + sample = samples[0] + del samples + + # Get original length for this batch (before any padding was added) + ori_length = ctx['all_ori_lengths'][decode_idx] if decode_idx < len(ctx['all_ori_lengths']) else sample.shape[0] + + # Trim temporal padding: sample is in [T, C, H, W] format after rearrange + if ori_length < sample.shape[0]: + padding_removed = sample.shape[0] - ori_length + debug.log(f"Trimming temporal padding: {padding_removed} frames removed ({sample.shape[0]} โ†’ {ori_length})", + category="video", indent_level=1) + sample = sample[:ori_length] + total_padding_removed += padding_removed + + # Trim spatial padding to true target dimensions + current_h, current_w = sample.shape[-2:] + if current_h != true_h or current_w != true_w: + debug.log(f"Trimming spatial padding: {current_w}x{current_h} โ†’ {true_w}x{true_h}", + category="video", indent_level=1) + sample = sample[:, :, :true_h, :true_w] + + # Convert to output format: [T, C, H, W] โ†’ [T, H, W, C] + # Note: We keep values in [-1, 1] range - normalization happens in Phase 4 + sample = optimized_sample_to_image_format(sample) # T, C, H, W โ†’ T, H, W, C + + # Calculate write position with temporal overlap handling + batch_frames = sample.shape[0] + if decode_idx == 0 or temporal_overlap == 0: + # First batch or no overlap: write all frames + write_start = current_write_idx + write_end = current_write_idx + batch_frames + else: + # Subsequent batches with overlap: blend overlapping region + if temporal_overlap < batch_frames and current_write_idx >= temporal_overlap: + # Blend overlapping region in-place on final_video + prev_tail = ctx['final_video'][current_write_idx - temporal_overlap:current_write_idx] + cur_head = sample[:temporal_overlap] + + # Move to same device for blending if needed + if prev_tail.device != cur_head.device: + cur_head = cur_head.to(prev_tail.device) + + blended = blend_overlapping_frames(prev_tail, cur_head, temporal_overlap) + ctx['final_video'][current_write_idx - temporal_overlap:current_write_idx] = blended + + debug.log(f"Blended {temporal_overlap} overlapping frames at positions {current_write_idx - temporal_overlap}-{current_write_idx}", + category="video", indent_level=1) + + # Write only non-overlapping part + sample = sample[temporal_overlap:] + batch_frames = sample.shape[0] + del prev_tail, cur_head, blended + + write_start = current_write_idx + write_end = current_write_idx + batch_frames + + # Move sample to target device and write directly to final_video + sample = manage_tensor( + tensor=sample, + target_device=target_device, + tensor_name=f"sample_{decode_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="writing to final_video", + indent_level=1 + ) + + # Write to final_video - for RGBA, write only RGB channels (VAE outputs 3 channels) + if ctx.get('is_rgba', False): + ctx['final_video'][write_start:write_end, :, :, :3] = sample + else: + ctx['final_video'][write_start:write_end] = sample + + # Store batch info for Phase 4 processing + ctx['decode_batch_info'].append((write_start, write_end, decode_idx, ori_length)) + current_write_idx = write_end + + debug.log(f"Wrote {batch_frames} frames to positions {write_start}-{write_end}", + category="video", indent_level=1) + + # Free memory immediately - no batch_samples storage + release_tensor_memory(ctx['all_upscaled_latents'][batch_idx]) + ctx['all_upscaled_latents'][batch_idx] = None + del upscaled_latent, sample + + debug.end_timer(f"decode_batch_{decode_idx+1}", f"Decoded batch {decode_idx+1}") + + if progress_callback: + progress_callback(decode_idx+1, num_valid_latents, + 1, "Phase 3: Decoding") + + decode_idx += 1 + + # Store padding stats for Phase 4 final summary + ctx['total_padding_removed'] = total_padding_removed + + except Exception as e: + debug.log(f"Error in Phase 3 (Decoding): {e}", level="ERROR", category="error", force=True) + raise + finally: + # Cleanup VAE as it's no longer needed + cleanup_vae(runner=runner, debug=debug, cache_model=cache_model) + + # Clean up upscaled latents storage + if 'all_upscaled_latents' in ctx: + release_tensor_collection(ctx['all_upscaled_latents']) + del ctx['all_upscaled_latents'] + + debug.end_timer("phase3_decoding", "Phase 3: VAE decoding complete", show_breakdown=True) + debug.log_memory_state("After phase 3 (VAE decoding)", show_tensors=False) + + return ctx + + +def postprocess_all_batches( + ctx: Dict[str, Any], + debug: 'Debug', + progress_callback: Optional[Callable[[int, int, int, str], None]] = None, + color_correction: str = "wavelet", + prepend_frames: int = 0, + temporal_overlap: int = 0, + batch_size: int = 5 +) -> Dict[str, Any]: + """ + Phase 4: Post-processing and Final Assembly. + + Processes final_video slices in-place: applies alpha upscaling, color correction, + and normalization. Reads from and writes back to the same final_video tensor + to avoid memory duplication. + + Args: + ctx: Context from decode_all_batches containing final_video (required) + debug: Debug instance for logging (required) + progress_callback: Optional callback(current, total, frames, phase_name) + color_correction: Color correction method - "wavelet", "adain", or "none" (default: "wavelet") + prepend_frames: Number of prepended frames to remove from final output (default: 0) + temporal_overlap: Number of overlapping frames between batches for blending (default: 0) + batch_size: Frames per batch used during encoding for overlap calculation (default: 5) + + Returns: + dict: Updated context containing: + - final_video: Assembled video tensor [T, H, W, C] range [0,1] with overlap blended and prepended frames removed + - All intermediate storage cleared for memory efficiency + + Raises: + ValueError: If context is missing or has no final_video + """ + if debug is None: + raise ValueError("Debug instance must be provided to postprocess_all_batches") + + if ctx is None: + raise ValueError("Context is required for postprocess_all_batches. Run decode_all_batches first.") + + # Validate we have final_video (pre-allocated in decode_all_batches) + if 'final_video' not in ctx or ctx['final_video'] is None: + raise ValueError("final_video not found. Run decode_all_batches first.") + + # Validate we have batch info for processing + if 'decode_batch_info' not in ctx or not ctx['decode_batch_info']: + raise ValueError("decode_batch_info not found. Run decode_all_batches first.") + + debug.log("", category="none", force=True) + debug.log("โ”โ”โ”โ”โ”โ”โ”โ” Phase 4: Post-processing โ”โ”โ”โ”โ”โ”โ”โ”", category="none", force=True) + debug.start_timer("phase4_postprocessing") + + # Total_frames represents the original input frame count (set in Phase 1) + total_frames = ctx.get('total_frames', 0) + + # Early exit if no frames to process + if total_frames == 0: + ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=ctx['compute_dtype']) + debug.log("No frames to process", level="WARNING", category="generation", force=True) + return ctx + + # Get batch info from decode phase + batch_info_list = ctx['decode_batch_info'] + num_valid_samples = len(batch_info_list) + + # Calculate total post-processing work units + # For RGBA: each batch needs 2 steps (alpha processing + color correction/assembly) + # For RGB: each batch needs 1 step (color correction/assembly only) + has_alpha_processing = (ctx.get('is_rgba', False) and + 'all_alpha_channels' in ctx and + 'all_input_rgb' in ctx and + isinstance(ctx.get('all_alpha_channels'), list)) + + if has_alpha_processing: + total_postprocessing_steps = num_valid_samples * 2 # Alpha + main processing + else: + total_postprocessing_steps = num_valid_samples # Main processing only + + current_postprocessing_step = 0 + + # Get padding stats from Phase 3 + total_padding_removed = ctx.get('total_padding_removed', 0) + + # Alpha processing - handle RGBA inputs with edge-guided upscaling + # Process alpha on final_video slices in-place + if has_alpha_processing: + debug.log("Processing Alpha channel with edge-guided upscaling...", category="alpha") + + # Validate alpha channel data exists + if not isinstance(ctx.get('all_alpha_channels'), list) or not isinstance(ctx.get('all_input_rgb'), list): + debug.log("WARNING: Alpha channel data malformed, skipping alpha processing", + level="WARNING", category="alpha", force=True) + else: + for write_start, write_end, batch_idx, ori_length in batch_info_list: + # Bounds checking for alpha channel lists + if batch_idx >= len(ctx['all_alpha_channels']) or ctx['all_alpha_channels'][batch_idx] is None: + continue + + # Validate alpha channel tensor integrity + if not isinstance(ctx['all_alpha_channels'][batch_idx], torch.Tensor): + debug.log(f"WARNING: Alpha channel {batch_idx} is not a tensor, skipping", + level="WARNING", category="alpha", force=True) + continue + + debug.log(f"Processing Alpha batch {batch_idx+1}/{num_valid_samples}", category="alpha", force=True) + debug.start_timer(f"alpha_batch_{batch_idx+1}") + + # Get RGB slice from final_video for alpha processing + # final_video is [T, H, W, C], process_alpha_for_batch expects list of [T, C, H, W] + rgb_slice = ctx['final_video'][write_start:write_end, :, :, :3] # Only RGB + rgb_tchw = rgb_slice.permute(0, 3, 1, 2) # [T, H, W, 3] โ†’ [T, 3, H, W] + + # Process Alpha and merge with RGB + processed_samples = process_alpha_for_batch( + rgb_samples=[rgb_tchw], + alpha_original=ctx['all_alpha_channels'][batch_idx], + rgb_original=ctx['all_input_rgb'][batch_idx], + device=ctx['vae_device'], + compute_dtype=ctx['compute_dtype'], + debug=debug + ) + + # processed_samples[0] is [T, 4, H, W] (RGBA) + # Extract only the alpha channel and write to final_video's alpha slot + processed_rgba = processed_samples[0] # [T, 4, H, W] + alpha_channel = processed_rgba[:, 3:4, :, :] # [T, 1, H, W] + alpha_thwc = alpha_channel.permute(0, 2, 3, 1) # [T, 1, H, W] โ†’ [T, H, W, 1] + + alpha_thwc = manage_tensor( + tensor=alpha_thwc, + target_device=ctx['final_video'].device, + tensor_name=f"alpha_channel_{batch_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="writing alpha channel to final_video", + indent_level=1 + ) + + # Write only the alpha channel to the 4th channel slot + ctx['final_video'][write_start:write_end, :, :, 3:4] = alpha_thwc + + del rgb_slice, rgb_tchw, processed_samples, processed_rgba, alpha_channel, alpha_thwc + + # Free memory immediately + release_tensor_memory(ctx['all_alpha_channels'][batch_idx]) + ctx['all_alpha_channels'][batch_idx] = None + + release_tensor_memory(ctx['all_input_rgb'][batch_idx]) + ctx['all_input_rgb'][batch_idx] = None + + debug.end_timer(f"alpha_batch_{batch_idx+1}", f"Alpha batch {batch_idx+1}") + + # Update progress for alpha processing step + current_postprocessing_step += 1 + if progress_callback: + progress_callback(current_postprocessing_step, total_postprocessing_steps, + 1, "Phase 4: Post-processing") + + debug.log("Alpha processing complete for all batches", category="alpha") + + try: + # Process each batch slice in final_video in-place + for info_idx, (write_start, write_end, batch_idx, ori_length) in enumerate(batch_info_list): + check_interrupt(ctx) + + debug.log(f"Post-processing batch {info_idx+1}/{num_valid_samples}", category="video", force=True) + debug.start_timer(f"postprocess_batch_{info_idx+1}") + + # Get slice from final_video - currently in [T, H, W, C] format, values in [-1, 1] + sample_thwc = ctx['final_video'][write_start:write_end] + + # For RGBA, we only process RGB channels for color correction + # Alpha was already written during alpha processing above + if ctx.get('is_rgba', False) and sample_thwc.shape[-1] == 4: + sample_thwc_rgb = sample_thwc[..., :3] # [T, H, W, 3] + sample = sample_thwc_rgb.permute(0, 3, 1, 2) # [T, H, W, 3] โ†’ [T, 3, H, W] + else: + sample = sample_thwc.permute(0, 3, 1, 2) # [T, H, W, C] โ†’ [T, C, H, W] + + # Move to VAE device for processing + sample = manage_tensor( + tensor=sample, + target_device=ctx['vae_device'], + tensor_name=f"sample_{info_idx+1}", + dtype=ctx['compute_dtype'], + debug=debug, + reason="post-processing", + indent_level=1 + ) + + # Reconstruct transformed video on-demand for color correction + input_video = None + if color_correction != "none" and ctx.get('batch_metadata') is not None: + if batch_idx < len(ctx['batch_metadata']) and ctx['batch_metadata'][batch_idx] is not None: + # Reconstruct transformation + transformed_video = _reconstruct_and_transform_batch(ctx, batch_idx, debug) + input_video = optimized_single_video_rearrange(transformed_video) + del transformed_video + + # For batches after the first with temporal overlap, the overlap frames + # were blended in Phase 3 and are not part of this slice. Skip them. + actual_overlap = ctx.get('actual_temporal_overlap', 0) + if info_idx > 0 and actual_overlap > 0: + input_video = input_video[actual_overlap:] + + # Trim input_video to match sample length (handles padding differences) + if input_video.shape[0] > sample.shape[0]: + input_video = input_video[:sample.shape[0]] + + # Trim spatial dimensions to true target size + if 'true_target_dims' in ctx: + true_h, true_w = ctx['true_target_dims'] + if input_video.shape[-2] != true_h or input_video.shape[-1] != true_w: + input_video = input_video[:, :, :true_h, :true_w] + + # Apply color correction if enabled (RGB only) + if color_correction != "none" and input_video is not None: + # Check if RGBA (samples are in T, C, H, W format at this point) + has_alpha = ctx.get('is_rgba', False) + alpha_channel = None + + if has_alpha: + # Check actual channel count + if sample.shape[1] == 4: + # Extract and temporarily store alpha for reattachment after color correction + alpha_channel = sample[:, 3:4, :, :] # (T, 1, H, W) + sample = sample[:, :3, :, :] # Keep only RGB (T, 3, H, W) + + # Ensure both tensors are on same device (GPU) for color correction + if input_video.device != sample.device: + input_video = manage_tensor( + tensor=input_video, + target_device=sample.device, + tensor_name=f"input_video_{info_idx+1}", + debug=debug, + reason="color correction", + indent_level=1 + ) + + # Apply selected color correction method + debug.start_timer(f"color_correction_{color_correction}") + + if color_correction == "lab": + debug.log("Applying LAB perceptual color transfer", category="video", force=True, indent_level=1) + sample = lab_color_transfer(sample, input_video, debug, luminance_weight=0.8) + elif color_correction == "wavelet_adaptive": + debug.log("Applying wavelet with adaptive saturation correction", category="video", force=True, indent_level=1) + sample = wavelet_adaptive_color_correction(sample, input_video, debug) + elif color_correction == "wavelet": + debug.log("Applying wavelet color reconstruction", category="video", force=True, indent_level=1) + sample = wavelet_reconstruction(sample, input_video, debug) + elif color_correction == "hsv": + debug.log("Applying HSV hue-conditional saturation matching", category="video", force=True, indent_level=1) + sample = hsv_saturation_histogram_match(sample, input_video, debug) + elif color_correction == "adain": + debug.log("Applying AdaIN color correction", category="video", force=True, indent_level=1) + sample = adaptive_instance_normalization(sample, input_video) + else: + debug.log(f"Unknown color correction method: {color_correction}", level="WARNING", category="video", force=True, indent_level=1) + + debug.end_timer(f"color_correction_{color_correction}", f"Color correction ({color_correction})") + + # Free the reconstructed transformed video + del input_video + + # Recombine with Alpha if it was present in input + if has_alpha and alpha_channel is not None: + # Concatenate in channels-first: (T, 3, H, W) + (T, 1, H, W) -> (T, 4, H, W) + sample = torch.cat([sample, alpha_channel], dim=1) + + else: + debug.log("Color correction disabled (set to none)", category="video", indent_level=1) + + # Convert to final format: [T, C, H, W] โ†’ [T, H, W, C] + sample = optimized_sample_to_image_format(sample) + + # Apply normalization only to RGB channels, preserve Alpha as-is + if ctx.get('is_rgba', False) and sample.shape[-1] == 4: + # Split RGBA: sample is (T, H, W, C) format after optimized_sample_to_image_format + rgb_channels = sample[..., :3] # (T, H, W, 3) + alpha_channel = sample[..., 3:4] # (T, H, W, 1) + + # Normalize only RGB from [-1, 1] to [0, 1] + rgb_channels.clamp_(-1, 1).mul_(0.5).add_(0.5) + + # Merge back with unchanged Alpha + sample = torch.cat([rgb_channels, alpha_channel], dim=-1) + else: + # RGB only: apply normalization as usual + sample.clamp_(-1, 1).mul_(0.5).add_(0.5) + + # Draw tile boundaries for debugging (if tile info available) + for phase, attr in [('encode', 'encode_tile_boundaries'), ('decode', 'decode_tile_boundaries')]: + tiles = getattr(debug, attr, None) + if tiles: + sample = _draw_tile_boundaries(sample, debug, tiles, phase) + break + + # Move to final_video device and write back in-place + sample = manage_tensor( + tensor=sample, + target_device=ctx['final_video'].device, + tensor_name=f"sample_{info_idx+1}_final", + dtype=ctx['compute_dtype'], + debug=debug, + reason="writing processed result to final_video", + indent_level=1 + ) + + # Write back to final_video in-place + # For RGBA, write only RGB channels (alpha already written during alpha processing) + if ctx.get('is_rgba', False) and ctx['final_video'].shape[-1] == 4: + ctx['final_video'][write_start:write_end, :, :, :3] = sample + else: + ctx['final_video'][write_start:write_end] = sample + + # Free sample memory + del sample, sample_thwc + + debug.end_timer(f"postprocess_batch_{info_idx+1}", f"Post-processed batch {info_idx+1}") + + # Update progress for main processing step + current_postprocessing_step += 1 + if progress_callback: + progress_callback(current_postprocessing_step, total_postprocessing_steps, + 1, "Phase 4: Post-processing") + + # Verify final assembly + if ctx['final_video'] is not None: + # Remove prepended frames if any were added at the start + frames_before_removal = ctx['final_video'].shape[0] + + if prepend_frames > 0: + if prepend_frames < ctx['final_video'].shape[0]: + debug.log(f"Removing {prepend_frames} prepended frames from output", category="video", force=True) + ctx['final_video'] = ctx['final_video'][prepend_frames:] + else: + debug.log(f"Warning: prepend_frames ({prepend_frames}) >= total frames ({ctx['final_video'].shape[0]}), skipping removal", + level="WARNING", category="video", force=True) + + final_shape = ctx['final_video'].shape + Tf, Hf, Wf, Cf = final_shape[0], final_shape[1], final_shape[2], final_shape[3] + channels_str = "RGBA" if Cf == 4 else "RGB" if Cf == 3 else f"{Cf}-channel" + + # Build message showing prepend and/or padding removal if applicable + frame_info = f"{Tf} frames" + adjustments = [] + + if prepend_frames > 0 and prepend_frames < frames_before_removal: + adjustments.append(f"{prepend_frames} prepend") + + if total_padding_removed > 0: + adjustments.append(f"{total_padding_removed} padding") + + # Use actual temporal overlap from encoding (may have been reset) + actual_overlap = ctx.get('actual_temporal_overlap', temporal_overlap) + + # Calculate and include temporal overlap blending info + if actual_overlap > 0: + frames_blended = (num_valid_samples - 1) * actual_overlap + adjustments.append(f"{frames_blended} overlap") + + if adjustments: + # Add back all removed/blended frames to get true computed count + total_computed = frames_before_removal + total_padding_removed + if actual_overlap > 0: + total_computed += (num_valid_samples - 1) * actual_overlap + frame_info += f" ({total_computed} computed with {' + '.join(adjustments)} removed)" + + debug.log(f"Output assembled: {frame_info}, Resolution: {Wf}x{Hf}px, Channels: {channels_str}", + category="generation", force=True) + else: + ctx['final_video'] = torch.empty((0, 0, 0, 0), dtype=ctx['compute_dtype']) + debug.log("No frames were processed", level="WARNING", category="generation", force=True) + + except Exception as e: + debug.log(f"Error in Phase 4 (Post-processing): {e}", level="ERROR", category="generation", force=True) + raise + finally: + # 1. Clean up decode_batch_info and padding stats + if 'decode_batch_info' in ctx: + del ctx['decode_batch_info'] + if 'total_padding_removed' in ctx: + del ctx['total_padding_removed'] + + # 2. Clean up video transform caches + if 'video_transform' in ctx and ctx['video_transform'] is not None: + if hasattr(ctx['video_transform'], 'transforms'): + for transform in ctx['video_transform'].transforms: + # Clear cache attributes + for cache_attr in ['cache', '_cache']: + if hasattr(transform, cache_attr): + setattr(transform, cache_attr, None) + # Clear remaining attributes + if hasattr(transform, '__dict__'): + transform.__dict__.clear() + del ctx['video_transform'] + + # 3. Clean up storage lists (all_latents, all_alpha_channels, etc.) + tensor_storage_keys = ['all_latents', 'all_alpha_channels', 'all_input_rgb'] + for key in tensor_storage_keys: + if key in ctx and ctx[key]: + release_tensor_collection(ctx[key]) + del ctx[key] + + # 4. Clean up non-tensor storage + if 'all_ori_lengths' in ctx: + del ctx['all_ori_lengths'] + if 'true_target_dims' in ctx: + del ctx['true_target_dims'] + if 'batch_metadata' in ctx: + del ctx['batch_metadata'] + if 'input_images' in ctx: + release_tensor_memory(ctx['input_images']) + del ctx['input_images'] + + debug.end_timer("phase4_postprocessing", "Phase 4: Post-processing complete", show_breakdown=True) + debug.log_memory_state("After phase 4 (Post-processing)", show_tensors=False) + + return ctx diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_utils.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9a4cb3c5ce75cac583db062175cb2812d632fa7e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/generation_utils.py @@ -0,0 +1,827 @@ +""" +Generation Utilities for SeedVR2 + +This module provides setup, configuration, and utility functions for the generation pipeline. +These are helper functions used to prepare, configure, and support the video upscaling process, +but do not contain the core 4-phase orchestration logic itself. + +Setup Functions: +- prepare_video_transforms: Create video transformation pipeline (resize, pad, normalize) +- setup_video_transform: Initialize transforms and compute target dimensions +- setup_generation_context: Initialize context with device configuration +- prepare_runner: Configure VideoDiffusionInfer with all settings + +Video Processing Utilities: +- pad_video_temporal: Unified temporal padding with reversed frames (prepend/append, any format) +- blend_overlapping_frames: Smooth blending for temporal overlap between batches + +Configuration Helpers: +- load_text_embeddings: Load positive/negative text embeddings for DiT +- calculate_optimal_batch_params: Compute batch processing parameters +- check_interrupt: Check for user interruption + +Debugging Utilities: +- _draw_tile_boundaries: Draw tile boundaries for debugging VAE tiling +- ensure_precision_initialized: Log model dtype information + +These utilities support the 4-phase pipeline implemented in generation_phases.py. +""" + +import os +import torch +from typing import Dict, List, Optional, Tuple, Any, Callable, Union +from torchvision.transforms import Compose, Lambda, Normalize + +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 + +# Get script directory for embeddings +script_directory = get_script_directory() + + +def prepare_video_transforms(resolution: int, max_resolution: int = 0, debug: Optional['Debug'] = None) -> Compose: + """ + Prepare optimized video transformation pipeline + + Args: + resolution (int): Target resolution for shortest edge + max_resolution (int): Maximum resolution for any edge (0 = no limit) + debug (Debug, optional): Debug instance for logging + + Returns: + Compose: Configured transformation pipeline + + Features: + - Resolution-aware upscaling (no downsampling) + - Optional max resolution constraint on longest edge + - Padding to divisible by 16 (no data loss) + - Proper normalization for model compatibility + - Memory-efficient tensor operations + """ + if debug: + msg = f"Initializing video transformation pipeline for {resolution}px (shortest edge)" + if max_resolution > 0: + msg += f", max {max_resolution}px (any edge)" + debug.log(msg, category="setup", indent_level=1) + + return Compose([ + NaResize( + resolution=resolution, + mode="side", + # Upsample image, model only trained for high res + downsample_only=False, + max_resolution=max_resolution, + ), + Lambda(lambda x: torch.clamp(x, 0.0, 1.0)), + DivisiblePad((16, 16)), + Normalize(0.5, 0.5), + Lambda(lambda x: x.permute(1, 0, 2, 3)), # t c h w -> c t h w (faster than Rearrange) + ]) + + +def setup_video_transform(ctx: Dict[str, Any], resolution: int, max_resolution: int = 0, + debug: Optional['Debug'] = None, + sample_frame: Optional[torch.Tensor] = None) -> Tuple[int, int, int, int]: + """ + Setup video transformation pipeline and compute target dimensions. + + Args: + ctx: Generation context dictionary + resolution: Target resolution for shortest edge + max_resolution: Maximum resolution for any edge (0 = no limit) + debug: Debug instance for logging + sample_frame: Optional sample frame tensor (C, H, W) to compute dimensions + + Returns: + (true_height, true_width, padded_height, padded_width) if dimensions computed, + (0, 0, 0, 0) otherwise + """ + # Check if transform exists AND is not None + existing_transform = ctx.get('video_transform') + + if existing_transform is not None: + # Transform exists - check if we need to compute dimensions + if 'true_target_dims' in ctx and sample_frame is not None: + # Return cached dimensions + recompute padded from sample + true_h, true_w = ctx['true_target_dims'] + transformed = existing_transform(sample_frame) + padded_h, padded_w = transformed.shape[-2:] + if debug: + debug.log("Reusing pre-initialized video transformation pipeline", category="reuse") + return true_h, true_w, padded_h, padded_w + elif debug: + debug.log("Reusing pre-initialized video transformation pipeline", category="reuse") + return 0, 0, 0, 0 + + # Create transformation pipeline (first time or after cleanup) + ctx['video_transform'] = prepare_video_transforms(resolution, max_resolution, debug) + + # Compute dimensions if sample frame provided + if sample_frame is not None: + # Get true target size (after resize, before padding) + temp_transform = Compose([ + NaResize(resolution=resolution, mode="side", downsample_only=False, max_resolution=max_resolution), + Lambda(lambda x: torch.clamp(x, 0.0, 1.0)) + ]) + resized_sample = temp_transform(sample_frame) + true_h, true_w = resized_sample.shape[-2:] + + # Round to even numbers for video codec compatibility (libx264 requirement) + true_h = (true_h // 2) * 2 + true_w = (true_w // 2) * 2 + + # Cache for later use in trimming + ctx['true_target_dims'] = (true_h, true_w) + + # Get padded dimensions + transformed_sample = ctx['video_transform'](sample_frame) + padded_h, padded_w = transformed_sample.shape[-2:] + + if debug: + if true_h == padded_h and true_w == padded_w: + debug.log(f"Target dimensions: {true_w}x{true_h} (no padding needed)", + category="setup", indent_level=1) + else: + debug.log(f"Target dimensions: {true_w}x{true_h} (padded to {padded_w}x{padded_h} for processing)", + category="setup", indent_level=1) + + del temp_transform, resized_sample, transformed_sample + return true_h, true_w, padded_h, padded_w + + return 0, 0, 0, 0 + + +def compute_generation_info( + ctx: Dict[str, Any], + images: torch.Tensor, + resolution: int = 1080, + max_resolution: int = 0, + batch_size: int = 5, + uniform_batch_size: bool = False, + seed: int = 42, + prepend_frames: int = 0, + temporal_overlap: int = 0, + debug: Optional['Debug'] = None +) -> Tuple[torch.Tensor, Dict[str, Any]]: + """ + Compute all generation parameters and dimensions for logging. + + Args: + ctx: Generation context dictionary + images: Input frames tensor [T, H, W, C] + resolution: Target resolution for shortest edge + max_resolution: Maximum resolution for any edge (0 = no limit) + batch_size: Frames per batch + uniform_batch_size: Whether to pad final batch to match batch_size + seed: Random seed + prepend_frames: Number of frames to prepend + temporal_overlap: Overlapping frames between batches + debug: Debug instance for logging + + Returns: + Tuple of (processed_images, info_dict) + - processed_images: Input images with prepending applied if needed + - info_dict: Information dictionary for logging + """ + # Track input frames before any modifications + input_frames = len(images) + input_h, input_w = images.shape[1], images.shape[2] + channels_info = "RGBA" if images.shape[-1] == 4 else "RGB" + + # Apply prepending if requested + if prepend_frames > 0: + images = pad_video_temporal(images, count=prepend_frames, temporal_dim=0, prepend=True, debug=debug) + + # Track total frames after prepending + total_frames = len(images) + ctx['total_frames'] = total_frames + + # Setup transform and compute dimensions on final frame count + sample_frame = images[0].permute(2, 0, 1).unsqueeze(0) + true_h, true_w, padded_h, padded_w = setup_video_transform( + ctx, resolution, max_resolution, debug, sample_frame + ) + del sample_frame + + info = { + 'input_frames': input_frames, + 'input_h': input_h, + 'input_w': input_w, + 'total_frames': total_frames, + 'true_h': true_h, + 'true_w': true_w, + 'padded_h': padded_h, + 'padded_w': padded_w, + 'channels_info': channels_info, + 'batch_size': batch_size, + 'uniform_batch_size': uniform_batch_size, + 'seed': seed, + 'prepend_frames': prepend_frames, + 'temporal_overlap': temporal_overlap, + 'resolution': resolution, + 'max_resolution': max_resolution + } + + return images, info + + +def log_generation_start(info: Dict[str, Any], debug: Optional['Debug'] = None) -> None: + """ + Log generation start information in a consistent format. + + Args: + info: Information dictionary from compute_generation_info() + debug: Debug instance for logging + """ + if debug is None: + return + + debug.log("", category="none", force=True) + debug.log("Starting upscaling generation...", category="generation", force=True) + + # Build concise parameter info + batch_text = f"Batch size: {info['batch_size']}" + if info.get('uniform_batch_size', False): + batch_text += " (uniform)" + params_info = batch_text + if info['prepend_frames'] > 0: + params_info += f", Prepend frames: {info['prepend_frames']}" + if info['temporal_overlap'] > 0: + params_info += f", Temporal overlap: {info['temporal_overlap']}" + params_info += f", Seed: {info['seed']}, Channels: {info['channels_info']}" + + # Build resolution constraint info + res_constraint = f"shortest edge: {info['resolution']}px" + if info['max_resolution'] > 0: + res_constraint += f", max edge: {info['max_resolution']}px" + + # Log dimension flow with full context + if info['true_h'] > 0: + frame_text = "frame" if info['input_frames'] <= 1 else "frames" + if info['true_h'] == info['padded_h'] and info['true_w'] == info['padded_w']: + debug.log( + f"Input: {info['input_frames']} {frame_text}, " + f"{info['input_w']}x{info['input_h']}px โ†’ Output: {info['true_w']}x{info['true_h']}px " + f"({res_constraint})", + category="generation", force=True, indent_level=1 + ) + else: + debug.log( + f"Input: {info['input_frames']} {frame_text}, " + f"{info['input_w']}x{info['input_h']}px โ†’ Padded: {info['padded_w']}x{info['padded_h']}px โ†’ " + f"Output: {info['true_w']}x{info['true_h']}px ({res_constraint})", + category="generation", force=True, indent_level=1 + ) + + debug.log(f"{params_info}", category="generation", force=True, indent_level=1) + + +def blend_overlapping_frames(prev_tail: torch.Tensor, cur_head: torch.Tensor, overlap: int) -> torch.Tensor: + """ + Blend two overlapping frame sequences in-place. + + Args: + prev_tail: Last `overlap` frames from previous batch [overlap, H, W, C] + cur_head: First `overlap` frames from current batch [overlap, H, W, C] + overlap: Number of overlapping frames + + Returns: + torch.Tensor: Blended frames [overlap, H, W, C] + """ + device = prev_tail.device + dtype = prev_tail.dtype + + # Smooth crossfade with Hann window for overlap >= 3, linear for smaller overlaps + if overlap >= 3: + t = torch.linspace(0.0, 1.0, steps=overlap, device=device, dtype=dtype) + blend_start = 1.0 / 3.0 + blend_end = 2.0 / 3.0 + u = ((t - blend_start) / (blend_end - blend_start)).clamp(0.0, 1.0) + w_prev_1d = 0.5 + 0.5 * torch.cos(torch.pi * u) # Hann window + else: + w_prev_1d = torch.linspace(1.0, 0.0, steps=overlap, device=device, dtype=dtype) + + w_prev = w_prev_1d.view(overlap, 1, 1, 1) + w_cur = 1.0 - w_prev + + return prev_tail * w_prev + cur_head * w_cur + + +def setup_generation_context( + dit_device: Optional[Union[str, torch.device]] = None, + vae_device: Optional[Union[str, torch.device]] = None, + dit_offload_device: Optional[Union[str, torch.device]] = None, + vae_offload_device: Optional[Union[str, torch.device]] = None, + tensor_offload_device: Optional[Union[str, torch.device]] = None, + debug: Optional['Debug'] = None +) -> Dict[str, Any]: + """ + Initialize generation context with device configuration. + + Processes device objects, configures environment variables, and creates the + generation context dictionary with all necessary state. + + Args: + dit_device: Device for DiT model (str or torch.device, defaults to 'cpu') + vae_device: Device for VAE model (str or torch.device, defaults to 'cpu') + dit_offload_device: Device to offload DiT to when not in use (optional) + vae_offload_device: Device to offload VAE to when not in use (optional) + tensor_offload_device: Device to offload intermediate tensors to (optional) + debug: Debug instance for logging + + Returns: + Dict[str, Any]: Generation context dictionary with torch.device objects + """ + # Normalize devices to torch.device objects (follows PyTorch convention) + def _normalize_device(device_spec: Optional[Union[str, torch.device]]) -> torch.device: + """Convert device specification to torch.device object.""" + if device_spec is None: + return torch.device("cpu") + if isinstance(device_spec, torch.device): + return device_spec + return torch.device(device_spec) + + dit_device = _normalize_device(dit_device) + 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 + # 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 + os.environ.setdefault("LOCAL_RANK", "0") + + # Detect ComfyUI integration for interrupt support + try: + import comfy.model_management + interrupt_fn = comfy.model_management.throw_exception_if_processing_interrupted + comfyui_available = True + except: + interrupt_fn = None + comfyui_available = False + + # Create generation context + ctx = { + 'dit_device': dit_device, + 'vae_device': vae_device, + 'dit_offload_device': dit_offload_device, + 'vae_offload_device': vae_offload_device, + 'tensor_offload_device': tensor_offload_device, + 'compute_dtype': COMPUTE_DTYPE, + 'interrupt_fn': interrupt_fn, + 'video_transform': None, + 'text_embeds': None, + 'all_latents': [], + 'all_upscaled_latents': [], + 'batch_samples': [], + 'final_video': None, + 'comfyui_available': comfyui_available, + } + + if debug: + # Build device configuration summary + offload_info = [] + if dit_offload_device: + offload_info.append(f"DiT offload={str(dit_offload_device)}") + if vae_offload_device: + offload_info.append(f"VAE offload={str(vae_offload_device)}") + if tensor_offload_device: + offload_info.append(f"Tensor offload={str(tensor_offload_device)}") + + offload_str = ", ".join(offload_info) if offload_info else "none" + + debug.log( + f"Generation context initialized: " + f"DiT={str(dit_device)}, VAE={str(vae_device)}, " + f"Offload=[{offload_str}], " + f"LOCAL_RANK={os.environ['LOCAL_RANK']}", + category="setup" + ) + 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 + + +def prepare_runner( + dit_model: str, + vae_model: str, + model_dir: str, + debug: 'Debug', + ctx: Dict[str, Any], + dit_cache: bool = False, + vae_cache: bool = False, + dit_id: Optional[int] = None, + vae_id: Optional[int] = None, + block_swap_config: Optional[Dict[str, Any]] = None, + encode_tiled: bool = False, + encode_tile_size: Optional[Tuple[int, int]] = None, + encode_tile_overlap: Optional[Tuple[int, int]] = None, + decode_tiled: bool = False, + decode_tile_size: Optional[Tuple[int, int]] = None, + decode_tile_overlap: Optional[Tuple[int, int]] = None, + tile_debug: str = "false", + attention_mode: str = 'sdpa', + torch_compile_args_dit: Optional[Dict[str, Any]] = None, + torch_compile_args_vae: Optional[Dict[str, Any]] = None +) -> Tuple['VideoDiffusionInfer', Dict[str, Any]]: + """ + Prepare runner with model state management and global cache integration. + Handles model changes and caching logic with independent DiT/VAE caching support. + + Args: + dit_model: DiT model filename (e.g., "seedvr2_ema_3b_fp16.safetensors") + vae_model: VAE model filename (e.g., "ema_vae_fp16.safetensors") + model_dir: Base directory containing model files + debug: Debug instance for logging (required) + ctx: Generation context from setup_generation_context + dit_cache: Whether to cache DiT model between runs + vae_cache: Whether to cache VAE model between runs + dit_id: Node instance ID for DiT model caching + vae_id: Node instance ID for VAE model caching + block_swap_config: Optional BlockSwap configuration for DiT memory optimization + encode_tiled: Enable tiled encoding to reduce VRAM during VAE encoding + encode_tile_size: Tile size for encoding (height, width) + encode_tile_overlap: Tile overlap for encoding (height, width) + decode_tiled: Enable tiled decoding to reduce VRAM during VAE decoding + 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', '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 + + Returns: + Tuple['VideoDiffusionInfer', Dict[str, Any]]: Tuple containing: + - VideoDiffusionInfer: Configured runner instance with models loaded and settings applied + - Dict[str, Any]: Cache context dictionary containing cache state and metadata with keys: + - 'global_cache': GlobalModelCache instance + - 'dit_cache', 'vae_cache': Caching enabled flags + - 'dit_id', 'vae_id': Node IDs for cache lookup + - 'cached_dit', 'cached_vae': Cached model instances (if found) + - 'reusing_runner': Flag indicating if runner template was reused + + Features: + - Independent DiT and VAE caching for flexible memory management + - Dynamic model reloading when models change + - Optional torch.compile optimization for inference speedup + - Separate encode/decode tiling configuration for optimal performance + - Memory optimization and BlockSwap integration + """ + dit_changed = False + vae_changed = False + + # Configure runner + debug.log("Configuring inference runner...", category="runner") + runner, cache_context = configure_runner( + dit_model=dit_model, + vae_model=vae_model, + base_cache_dir=model_dir, + debug=debug, + ctx=ctx, + dit_cache=dit_cache, + vae_cache=vae_cache, + dit_id=dit_id, + vae_id=vae_id, + block_swap_config=block_swap_config, + encode_tiled=encode_tiled, + encode_tile_size=encode_tile_size, + encode_tile_overlap=encode_tile_overlap, + decode_tiled=decode_tiled, + decode_tile_size=decode_tile_size, + decode_tile_overlap=decode_tile_overlap, + tile_debug=tile_debug, + attention_mode=attention_mode, + torch_compile_args_dit=torch_compile_args_dit, + torch_compile_args_vae=torch_compile_args_vae + ) + + return runner, cache_context + + +def load_text_embeddings(script_directory: str, device: torch.device, + dtype: torch.dtype, debug: Optional['Debug'] = None) -> Dict[str, List[torch.Tensor]]: + """ + Load and prepare text embeddings for generation + + Args: + script_directory (str): Script directory path + device (torch.device): Target device + dtype (torch.dtype): Target dtype + debug: Optional debug instance for logging + + Returns: + dict: Text embeddings dictionary + + Features: + - Adaptive dtype handling + - Device-optimized loading + - Memory-efficient embedding preparation + - Consistent movement logging + """ + 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'), weights_only=True) + + text_pos_embeds = manage_tensor( + tensor=text_pos_embeds, + target_device=device, + tensor_name="text_pos_embeds", + dtype=dtype, + debug=debug, + reason="DiT inference" + ) + text_neg_embeds = manage_tensor( + tensor=text_neg_embeds, + target_device=device, + tensor_name="text_neg_embeds", + dtype=dtype, + debug=debug, + reason="DiT inference" + ) + + return {"texts_pos": [text_pos_embeds], "texts_neg": [text_neg_embeds]} + + +def calculate_optimal_batch_params(total_frames: int, batch_size: int, + temporal_overlap: int) -> Dict[str, Any]: + """ + Calculate batch processing parameters. + + Args: + total_frames (int): Total number of frames to process + batch_size (int): Desired batch size + temporal_overlap (int): Number of overlapping frames between batches + + Returns: + dict: { + 'step': Effective step size between batches, + 'temporal_overlap': Adjusted temporal overlap, + 'best_batch': Optimal batch size matching video length (4n+1 format) + } + + The 4n+1 constraint (1, 5, 9, 13, 17, 21...) is required by the model. + """ + # Calculate step size + step = batch_size - temporal_overlap + if step <= 0: + step = batch_size + temporal_overlap = 0 + + # Find all valid 4n+1 batch sizes up to total_frames + valid_sizes = [i for i in range(1, total_frames + 1) if i % 4 == 1] + + # Best batch: largest valid size โ‰ค total_frames (maximizes temporal stability) + best_batch = max(valid_sizes) if valid_sizes else 1 + + return { + 'step': step, + 'temporal_overlap': temporal_overlap, + 'best_batch': best_batch + } + + +def pad_video_temporal(videos: torch.Tensor, count: int = 0, temporal_dim: int = 1, + prepend: bool = False, debug: Optional['Debug'] = None) -> torch.Tensor: + """ + Extend video with reversed frames for temporal continuity. + + Single source of truth for all temporal padding (prepend/append, 4n+1 constraint). + + Args: + videos: Video tensor (TCHW or CTHW format) + count: Frames to add (0 = auto-calculate for 4n+1 when prepend=False) + temporal_dim: Temporal axis (0=TCHW, 1=CTHW) + prepend: Add to start (True) or end (False) + debug: Optional debug logger + + Returns: + torch.Tensor: Extended video + """ + t = videos.size(temporal_dim) + + # Auto-calculate for 4n+1 constraint + if count == 0 and not prepend: + if t % 4 == 1: + return videos + count = ((t - 1) // 4 + 1) * 4 + 1 - t + + if count <= 0: + return videos + + if debug and prepend: + debug.log(f"Prepending {count} reversed frames to reduce start artifacts", + category="video", indent_level=1) + + # Helper to select frames along temporal dimension + def select(start, end): + return videos[start:end] if temporal_dim == 0 else videos[:, start:end] + + # Padding exceeds video length: repeat + reverse + if count >= t: + repeat_count = count - t + 1 + last = select(-1, None) + + # Repeat last frame + if temporal_dim == 0: + repeated = last.repeat(repeat_count, 1, 1, 1) + reversed_frames = select(1, None).flip(temporal_dim) if t > 1 else last[:0] + else: + repeated = last.expand(-1, repeat_count, -1, -1).contiguous() + reversed_frames = select(1, None).flip(temporal_dim) if t > 1 else last[:, :0] + + return torch.cat([repeated, reversed_frames, videos] if prepend else + [videos, reversed_frames, repeated], dim=temporal_dim) + + # Normal case: extract and reverse frames + if prepend: + reversed_frames = select(1, count+1).flip(temporal_dim) + else: + reversed_frames = select(-count-1, -1).flip(temporal_dim) + + return torch.cat([reversed_frames, videos] if prepend else + [videos, reversed_frames], dim=temporal_dim) + + +def check_interrupt(ctx: Dict[str, Any]) -> None: + """Single interrupt check to avoid redundant imports""" + if ctx.get('interrupt_fn') is not None: + ctx['interrupt_fn']() + + +def _draw_tile_boundaries(image: torch.Tensor, debug: 'Debug', tile_boundaries: list, phase: str) -> torch.Tensor: + """ + Draw tile boundary overlays on all frames for debugging (non-destructive). + + Args: + image: Image tensor [T, H, W, C] or [H, W, C] in range [0, 1] + debug: Debug instance for logging + tile_boundaries: List of tile boundary info dictionaries + phase: Phase name ('encode' or 'decode') for logging + + Returns: + Image with boundary overlays drawn inside tiles on all frames + """ + if not tile_boundaries: + return image + + # Try to import required libraries + try: + import cv2 + import numpy as np + import random + import colorsys + except ImportError as e: + debug.log(f"Tile debug ignored: missing imports ({e})", level="WARNING", category="video") + return image + + # Handle both [T, H, W, C] and [H, W, C] + squeeze_t = False + if image.ndim == 3: + image = image.unsqueeze(0) + squeeze_t = True + + T, H, W, C = image.shape + original_dtype = image.dtype + + log_frames = f"all {T} frames" if T > 1 else "1 frame" + debug.log(f"Drawing {phase} tile boundaries ({len(tile_boundaries)} tiles) on {log_frames}", category="video", indent_level=1, force=True) + + # Scale line thickness and font size based on video width + # Reference points: 512px (min) to 1920px (max) + min_width, max_width = 512, 1920 + min_line_thickness, max_line_thickness = 2, 6 + min_font_scale, max_font_scale = 0.8, 2.5 + min_text_thickness, max_text_thickness = 2, 4 + + # Calculate scale factor (clamped between 0 and 1) + scale_factor = max(0.0, min(1.0, (W - min_width) / (max_width - min_width))) + + # Apply scaling + line_thickness = int(min_line_thickness + scale_factor * (max_line_thickness - min_line_thickness)) + font_scale = min_font_scale + scale_factor * (max_font_scale - min_font_scale) + text_thickness = int(min_text_thickness + scale_factor * (max_text_thickness - min_text_thickness)) + + # Generate high-contrast colors using HSV color space + num_tiles = len(tile_boundaries) + colors = [] + for i in range(num_tiles): + hue = (i * 360 / num_tiles) % 360 + saturation = 0.9 + (i % 2) * 0.1 + brightness = 0.8 + ((i // 2) % 2) * 0.2 + r, g, b = colorsys.hsv_to_rgb(hue / 360, saturation, brightness) + colors.append((int(b * 255), int(g * 255), int(r * 255))) # BGR for OpenCV + + random.seed(42) + random.shuffle(colors) + + # Process all frames + annotated_frames = [] + for frame_idx in range(T): + # Convert frame to numpy (handle RGB and RGBA) + 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): + tile_id = tile_info['id'] + x, y = tile_info['x'], tile_info['y'] + w, h = tile_info['w'], tile_info['h'] + color = colors[idx] + + inset = line_thickness // 2 + + # Draw four edges + cv2.line(img, (x, y + inset), (x + w, y + inset), color, line_thickness) + cv2.line(img, (x, y + h - inset), (x + w, y + h - inset), color, line_thickness) + cv2.line(img, (x + inset, y), (x + inset, y + h), color, line_thickness) + cv2.line(img, (x + w - inset, y), (x + w - inset, y + h), color, line_thickness) + + # Draw tile number + text = str(tile_id) + font = cv2.FONT_HERSHEY_SIMPLEX + (text_w, text_h), baseline = cv2.getTextSize(text, font, font_scale, text_thickness) + margin = int(15 * scale_factor) if scale_factor > 0.5 else 8 # Scale margin too + text_x = x + margin + text_y = y + text_h + margin + cv2.putText(img, text, (text_x, text_y), font, font_scale, color, text_thickness, cv2.LINE_AA) + + # Convert back to tensor + frame_tensor = torch.from_numpy(img.astype(np.float32) / 255.0).to(device=image.device, dtype=original_dtype) + annotated_frames.append(frame_tensor) + + # Stack all frames + image = torch.stack(annotated_frames, dim=0) + + if squeeze_t: + image = image.squeeze(0) + + return image + + +def ensure_precision_initialized( + ctx: Dict[str, Any], + runner: 'VideoDiffusionInfer', + debug: Optional['Debug'] = None +) -> None: + """ + Log model dtypes for debugging. Compute dtype is hardcoded in context. + + Since compute_dtype is hardcoded to bfloat16 in setup_generation_context(), + this function only logs model dtypes for informational purposes. + + Args: + ctx: Generation context dictionary (compute_dtype already set) + runner: VideoDiffusionInfer instance with loaded models + debug: Optional Debug instance for logging + """ + if not debug: + return + + try: + # Get model dtypes for informational logging + dit_dtype = None + vae_dtype = None + + if runner.dit is not None: + try: + param_device = next(runner.dit.parameters()).device + if param_device.type != 'meta': + dit_dtype = next(runner.dit.parameters()).dtype + except StopIteration: + pass + + if runner.vae is not None: + try: + param_device = next(runner.vae.parameters()).device + if param_device.type != 'meta': + vae_dtype = next(runner.vae.parameters()).dtype + except StopIteration: + pass + + # Build precision info string + parts = [] + if dit_dtype is not None: + parts.append(f"DiT={dit_dtype}") + if vae_dtype is not None: + parts.append(f"VAE={vae_dtype}") + parts.append(f"compute={ctx['compute_dtype']}") + + if parts: + 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) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/infer.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/infer.py new file mode 100644 index 0000000000000000000000000000000000000000..a0869caec664c1a4232b3aa9988901ba51951542 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/infer.py @@ -0,0 +1,395 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import List, Optional, Tuple, Union +import torch +from einops import rearrange +from omegaconf import DictConfig, ListConfig +from torch import Tensor +from ..common.diffusion import ( + classifier_free_guidance_dispatcher, + create_sampler_from_config, + create_sampling_timesteps_from_config, + create_schedule_from_config, +) +from ..common.distributed import ( + get_device, +) +from ..optimization.performance import ( + optimized_channels_to_last, + optimized_channels_to_second +) +from ..models.dit_3b import na + + +class VideoDiffusionInfer(): + def __init__(self, config: DictConfig, debug: 'Debug', + encode_tiled: bool = False, encode_tile_size: Tuple[int, int] = (512, 512), + encode_tile_overlap: Tuple[int, int] = (64, 64), + decode_tiled: bool = False, decode_tile_size: Tuple[int, int] = (512, 512), + decode_tile_overlap: Tuple[int, int] = (64, 64), + tile_debug: str = "false"): + self.config = config + self.debug = debug + # Store separate encode and decode tiling parameters + self.encode_tiled = encode_tiled + self.encode_tile_size = encode_tile_size + self.encode_tile_overlap = encode_tile_overlap + self.decode_tiled = decode_tiled + self.decode_tile_size = decode_tile_size + self.decode_tile_overlap = decode_tile_overlap + self.tile_debug = tile_debug + + def get_condition(self, latent: Tensor, latent_blur: Tensor, task: str) -> Tensor: + t, h, w, c = latent.shape + cond = torch.zeros([t, h, w, c + 1], device=latent.device, dtype=latent.dtype) + if task == "t2v" or t == 1: + # t2i or t2v generation. + if task == "sr": + cond[:, ..., :-1] = latent_blur[:] + cond[:, ..., -1:] = 1.0 + return cond + if task == "i2v": + # i2v generation. + cond[:1, ..., :-1] = latent[:1] + cond[:1, ..., -1:] = 1.0 + return cond + if task == "v2v": + # v2v frame extension. + cond[:2, ..., :-1] = latent[:2] + cond[:2, ..., -1:] = 1.0 + return cond + if task == "sr": + # sr generation. + cond[:, ..., :-1] = latent_blur[:] + cond[:, ..., -1:] = 1.0 + return cond + raise NotImplementedError + + def configure_diffusion(self, device: Optional[torch.device] = None, dtype=torch.float32): + """ + Configure diffusion schedule and sampler. + + Args: + device: Device for schedule tensors. If None, uses get_device() + dtype: Data type for computations + """ + # Use provided device or fallback to standard detection + if device is None: + device = get_device() + elif not isinstance(device, torch.device): + device = torch.device(device) + + self.schedule = create_schedule_from_config( + config=self.config.diffusion.schedule, + device=device, + dtype=dtype, + ) + self.sampling_timesteps = create_sampling_timesteps_from_config( + config=self.config.diffusion.timesteps.sampling, + schedule=self.schedule, + device=device, + dtype=dtype, + ) + self.sampler = create_sampler_from_config( + config=self.config.diffusion.sampler, + schedule=self.schedule, + timesteps=self.sampling_timesteps, + ) + # Propagate debug to sampler + if hasattr(self, 'debug'): + self.sampler.debug = self.debug + + # -------------------------------- Helper ------------------------------- # + + @torch.no_grad() + def vae_encode(self, samples: List[Tensor]) -> List[Tensor]: + """VAE encode with configured dtype - converts samples to latents with optional tiling""" + use_sample = self.config.vae.get("use_sample", True) + latents = [] + if len(samples) > 0: + # Use VAE model's current device + # This ensures consistency with where the VAE model is loaded + try: + device = next(self.vae.parameters()).device + except StopIteration: + # Fallback if VAE has no parameters (shouldn't happen) + device = get_device() + + dtype = getattr(torch, self.config.vae.dtype) + scale = self.config.vae.scaling_factor + shift = self.config.vae.get("shifting_factor", 0.0) + + if isinstance(scale, ListConfig): + scale = torch.tensor(scale, device=device, dtype=dtype) + if isinstance(shift, ListConfig): + shift = torch.tensor(shift, device=device, dtype=dtype) + + # Group samples of the same shape to batches if enabled. + if self.config.vae.grouping: + batches, indices = na.pack(samples) + else: + batches = [sample.unsqueeze(0) for sample in samples] + + # VAE process by each group. + for sample in batches: + if hasattr(self.vae, "preprocess"): + sample = self.vae.preprocess(sample) + + # Detect VAE model dtype + try: + vae_dtype = next(self.vae.parameters()).dtype + except StopIteration: + 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: + 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: + 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, + 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) + + latent = latent.unsqueeze(2) if latent.ndim == 4 else latent + latent = optimized_channels_to_last(latent) + latent = (latent - shift) * scale + latents.append(latent) + + # Ungroup back to individual latent with the original order. + if self.config.vae.grouping: + latents = na.unpack(latents, indices) + else: + latents = [latent.squeeze(0) for latent in latents] + + self.debug.log(f"Latents shape: {latents[0].shape}", category="info", indent_level=1) + + return latents + + + @torch.no_grad() + def vae_decode(self, latents: List[Tensor]) -> List[Tensor]: + """VAE decode with configured dtype - converts latents to samples with optional tiling""" + samples = [] + if len(latents) > 0: + # Use VAE model's current device + # This ensures consistency with where the VAE model is loaded + try: + device = next(self.vae.parameters()).device + except StopIteration: + # Fallback if VAE has no parameters (shouldn't happen) + device = get_device() + + dtype = getattr(torch, self.config.vae.dtype) + scale = self.config.vae.scaling_factor + shift = self.config.vae.get("shifting_factor", 0.0) + + if isinstance(scale, ListConfig): + scale = torch.tensor(scale, device=device, dtype=dtype) + if isinstance(shift, ListConfig): + shift = torch.tensor(shift, device=device, dtype=dtype) + + # Group samples of the same shape to batches if enabled. + if self.config.vae.grouping: + latents, indices = na.pack(latents) + else: + latents = [latent.unsqueeze(0) for latent in latents] + + self.debug.log(f"Latents shape: {latents[0].shape}", category="info", indent_level=1) + + for i, latent in enumerate(latents): + latent = latent / scale + shift + latent = optimized_channels_to_second(latent) + latent = latent.squeeze(2) + + # Detect VAE model dtype + try: + vae_dtype = next(self.vae.parameters()).dtype + except StopIteration: + 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: + 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, + tiled=self.decode_tiled, tile_size=self.decode_tile_size, + tile_overlap=self.decode_tile_overlap + ).sample + + if hasattr(self.vae, "postprocess"): + sample = self.vae.postprocess(sample) + + samples.append(sample) + + if self.config.vae.grouping: + samples = na.unpack(samples, indices) + else: + samples = [sample.squeeze(0) for sample in samples] + + return samples + + + def timestep_transform(self, timesteps: Tensor, latents_shapes: Tensor): + # Skip if not needed. + if not self.config.diffusion.timesteps.get("transform", False): + return timesteps + + # Compute resolution. + vt = self.config.vae.model.get("temporal_downsample_factor", 4) + vs = self.config.vae.model.get("spatial_downsample_factor", 8) + frames = (latents_shapes[:, 0] - 1) * vt + 1 + heights = latents_shapes[:, 1] * vs + widths = latents_shapes[:, 2] * vs + + # Compute shift factor. + def get_lin_function(x1, y1, x2, y2): + m = (y2 - y1) / (x2 - x1) + b = y1 - m * x1 + return lambda x: m * x + b + + img_shift_fn = get_lin_function(x1=256 * 256, y1=1.0, x2=1024 * 1024, y2=3.2) + vid_shift_fn = get_lin_function(x1=256 * 256 * 37, y1=1.0, x2=1280 * 720 * 145, y2=5.0) + shift = torch.where( + frames > 1, + vid_shift_fn(heights * widths * frames), + img_shift_fn(heights * widths), + ) + + # Shift timesteps. + timesteps = timesteps / self.schedule.T + timesteps = shift * timesteps / (1 + (shift - 1) * timesteps) + timesteps = timesteps * self.schedule.T + return timesteps + + + @torch.no_grad() + def inference( + self, + noises: List[Tensor], + conditions: List[Tensor], + texts_pos: Union[List[str], List[Tensor], List[Tuple[Tensor]]], + texts_neg: Union[List[str], List[Tensor], List[Tuple[Tensor]]], + cfg_scale: Optional[float] = None, + ) -> List[Tensor]: + assert len(noises) == len(conditions) == len(texts_pos) == len(texts_neg) + batch_size = len(noises) + + # Return if empty. + if batch_size == 0: + return [] + + # Set cfg scale + if cfg_scale is None: + cfg_scale = self.config.diffusion.cfg.scale + + # Text embeddings. + assert type(texts_pos[0]) is type(texts_neg[0]) + if isinstance(texts_pos[0], str): + text_pos_embeds, text_pos_shapes = self.text_encode(texts_pos) + text_neg_embeds, text_neg_shapes = self.text_encode(texts_neg) + elif isinstance(texts_pos[0], tuple): + text_pos_embeds, text_pos_shapes = [], [] + text_neg_embeds, text_neg_shapes = [], [] + for pos in zip(*texts_pos): + emb, shape = na.flatten(pos) + text_pos_embeds.append(emb) + text_pos_shapes.append(shape) + for neg in zip(*texts_neg): + emb, shape = na.flatten(neg) + text_neg_embeds.append(emb) + text_neg_shapes.append(shape) + else: + text_pos_embeds, text_pos_shapes = na.flatten(texts_pos) + text_neg_embeds, text_neg_shapes = na.flatten(texts_neg) + + # Flatten. + latents, latents_shapes = na.flatten(noises) + latents_cond, _ = na.flatten(conditions) + + latents = self.sampler.sample( + x=latents, + f=lambda args: classifier_free_guidance_dispatcher( + pos=lambda: self.dit( + vid=torch.cat([args.x_t, latents_cond], dim=-1), + txt=text_pos_embeds, + vid_shape=latents_shapes, + txt_shape=text_pos_shapes, + timestep=args.t.repeat(batch_size), + ).vid_sample, + neg=lambda: self.dit( + vid=torch.cat([args.x_t, latents_cond], dim=-1), + txt=text_neg_embeds, + vid_shape=latents_shapes, + txt_shape=text_neg_shapes, + timestep=args.t.repeat(batch_size), + ).vid_sample, + scale=( + cfg_scale + if (args.i + 1) / len(self.sampler.timesteps) + <= self.config.diffusion.cfg.get("partial", 1) + else 1.0 + ), + rescale=self.config.diffusion.cfg.rescale, + ), + ) + + latents = na.unflatten(latents, latents_shapes) + + # Clean up temporary tensors + del latents_cond + del latents_shapes + del text_pos_embeds + del text_neg_embeds + del text_pos_shapes + del text_neg_shapes + + return latents \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_cache.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_cache.py new file mode 100644 index 0000000000000000000000000000000000000000..2c54ea514d108a7acea3faf4a869ddc06125d142 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_cache.py @@ -0,0 +1,239 @@ +""" +Global Model Cache for SeedVR2 +Enables independent DiT and VAE model sharing across multiple upscaler node instances +""" + +from typing import Dict, Any, Optional, Tuple +from ..optimization.memory_manager import release_model_memory + + +class GlobalModelCache: + """ + Global cache for sharing DiT and VAE models independently across upscaler instances. + Each upscaler gets its own runner but can share the underlying models. + Also caches runner templates for instant runner creation. + """ + + def __init__(self): + # Storage for cached DiT models: node_id -> (model, config) + self._dit_models: Dict[str, Tuple[Any, Dict]] = {} + # Storage for cached VAE models: node_id -> (model, config) + self._vae_models: Dict[str, Tuple[Any, Dict]] = {} + # Storage for runner templates: "dit_id+vae_id" -> runner + self._runner_templates: Dict[str, Any] = {} + + def get_dit(self, dit_config: Dict[str, Any], debug: Optional['Debug'] = None) -> Optional[Any]: + """ + Get cached DiT model if available. + + Args: + dit_config: Configuration dictionary with 'cache_model' and 'node_id' keys + debug: Optional debug instance for logging + + Returns: + Cached DiT model instance if found, None otherwise + """ + if not dit_config.get('cache_model', False): + return None + + node_id = dit_config.get('node_id') + if node_id in self._dit_models: + model, stored_config = self._dit_models[node_id] + return model + return None + + def get_vae(self, vae_config: Dict[str, Any], debug: Optional['Debug'] = None) -> Optional[Any]: + """ + Get cached VAE model if available. + + Args: + vae_config: Configuration dictionary with 'cache_model' and 'node_id' keys + debug: Optional debug instance for logging + + Returns: + Cached VAE model instance if found, None otherwise + """ + if not vae_config.get('cache_model', False): + return None + + node_id = vae_config.get('node_id') + if node_id in self._vae_models: + model, stored_config = self._vae_models[node_id] + return model + return None + + def get_runner(self, dit_id: Optional[int], vae_id: Optional[int], + debug: Optional['Debug'] = None) -> Optional[Any]: + """ + Get cached runner template if available. + + Args: + dit_id: DiT node ID for lookup + vae_id: VAE node ID for lookup + debug: Optional debug instance for logging + + Returns: + Cached runner template if found, None if either ID is None or not cached + """ + if dit_id is None or vae_id is None: + return None + + runner_key = f"{dit_id}+{vae_id}" + if runner_key in self._runner_templates: + return self._runner_templates[runner_key] + return None + + def set_dit(self, dit_config: Dict[str, Any], model: Any, model_name: str, debug: Optional['Debug'] = None) -> Optional[str]: + """ + Store DiT model in cache. + + Args: + dit_config: Configuration dictionary with 'cache_model' and 'node_id' keys + model: DiT model instance to cache + model_name: Name identifier for the model + debug: Optional debug instance for logging + + Returns: + Node ID string if cached successfully, None if caching disabled + """ + if not dit_config.get('cache_model', False): + return None + + node_id = dit_config.get('node_id') + self._dit_models[node_id] = (model, dit_config) + + if debug: + debug.log(f"DiT model cached in memory (node {node_id}): {model_name}", + category="cache", force=True) + + return node_id + + def set_vae(self, vae_config: Dict[str, Any], model: Any, model_name: str, debug: Optional['Debug'] = None) -> Optional[str]: + """ + Store VAE model in cache. + + Args: + vae_config: Configuration dictionary with 'cache_model' and 'node_id' keys + model: VAE model instance to cache + model_name: Name identifier for the model + debug: Optional debug instance for logging + + Returns: + Node ID string if cached successfully, None if caching disabled + """ + if not vae_config.get('cache_model', False): + return None + + node_id = vae_config.get('node_id') + self._vae_models[node_id] = (model, vae_config) + + if debug: + debug.log(f"VAE model cached in memory (node {node_id}): {model_name}", + category="cache", force=True) + + return node_id + + def set_runner(self, dit_id: Optional[int], vae_id: Optional[int], + runner: Any, debug: Optional['Debug'] = None) -> Optional[str]: + """ + Store runner template in cache. + + Args: + dit_id: DiT node ID for the runner template + vae_id: VAE node ID for the runner template + runner: Runner instance to cache as template + debug: Optional debug instance for logging + + Returns: + Runner key string (format: "dit_id+vae_id") if cached successfully, + None if either ID is None or runner already cached + """ + if dit_id is None or vae_id is None: + return None + + runner_key = f"{dit_id}+{vae_id}" + if runner_key not in self._runner_templates: + self._runner_templates[runner_key] = runner + if debug: + debug.log(f"Runner template cached in memory: nodes {runner_key}", category="cache", force=True) + return runner_key + + return None + + def remove_dit(self, dit_config: Dict[str, Any], debug: Optional['Debug'] = None) -> bool: + """ + Remove DiT model from cache if it exists. + + Args: + dit_config: Configuration dictionary with 'node_id' key + debug: Optional debug instance for logging + + Returns: + True if model was removed, False if not found in cache + + Note: + Also removes any runner templates that used this DiT model + """ + node_id = dit_config.get('node_id') + if node_id in self._dit_models: + if debug: + debug.log(f"Removing cached DiT: {node_id}", category="cache", force=True) + + model, stored_config = self._dit_models[node_id] + + # Release model memory + if model is not None: + release_model_memory(model=model, debug=debug) + + del self._dit_models[node_id] + + # Remove any runner templates that used this DiT + templates_to_remove = [k for k in self._runner_templates.keys() if k.startswith(str(node_id) + "+")] + for template_key in templates_to_remove: + del self._runner_templates[template_key] + + return True + return False + + def remove_vae(self, vae_config: Dict[str, Any], debug: Optional['Debug'] = None) -> bool: + """ + Remove VAE model from cache if it exists. + + Args: + vae_config: Configuration dictionary with 'node_id' key + debug: Optional debug instance for logging + + Returns: + True if model was removed, False if not found in cache + + Note: + Also removes any runner templates that used this VAE model + """ + node_id = vae_config.get('node_id') + if node_id in self._vae_models: + if debug: + debug.log(f"Removing cached VAE: {node_id}", category="cache", force=True) + + model, stored_config = self._vae_models[node_id] + + # Release model memory directly + if model is not None: + release_model_memory(model=model, debug=debug) + + del self._vae_models[node_id] + + # Remove any runner templates that used this VAE + templates_to_remove = [k for k in self._runner_templates.keys() if k.endswith("+" + str(node_id))] + for template_key in templates_to_remove: + del self._runner_templates[template_key] + + return True + return False + + +# Global singleton instance +_global_cache = GlobalModelCache() + +def get_global_cache() -> GlobalModelCache: + """Get the global model cache instance.""" + return _global_cache \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_configuration.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_configuration.py new file mode 100644 index 0000000000000000000000000000000000000000..61297627c20ce1a41e84c03d712d70438903717f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_configuration.py @@ -0,0 +1,1480 @@ +""" +Model Configuration and Setup for SeedVR2 + +This module orchestrates model configuration, caching, and runtime settings: +- Runner configuration and lifecycle management +- Model caching with validation and updates +- DiT and VAE model setup with meta device initialization +- Configuration change detection and updates +- torch.compile integration +- BlockSwap configuration +- VAE tiling configuration + +Key Features: +- Unified runner configuration via configure_runner() +- Global model caching with name validation +- Configuration change detection to avoid unnecessary reloads +- Meta device initialization for memory efficiency +- Dynamic configuration updates (BlockSwap, torch.compile, attention mode) +- Automatic cleanup of stale cache entries + +Main Functions: +- configure_runner: Main entry point for runner setup +- apply_model_specific_config: Apply BlockSwap and torch.compile to models +- _setup_models: Setup DiT and VAE models (cached or new) +- _update_dit_config: Update DiT configuration on cached models +- _update_vae_config: Update VAE configuration on cached models + +Cache Management: +- _initialize_cache_context: Initialize cache with validation +- _acquire_runner: Get or create runner instance +- _create_new_runner: Create new runner with config +- _configure_runner_settings: Configure runner settings for tiling/compile/BlockSwap + +Configuration Helpers: +- _configs_equal: Compare configurations for equality +- _describe_blockswap_config: Human-readable BlockSwap description +- _describe_compile_config: Human-readable torch.compile description +- _describe_attention_mode: Human-readable attention mode description +- _describe_tiling_config: Human-readable VAE tiling description +- _update_model_config: Generic config update handler with change detection + +Model Setup: +- _setup_dit_model: Setup DiT model from cache or create new structure +- _setup_vae_model: Setup VAE model from cache or create new structure + +torch.compile Support: +- _configure_torch_compile: Configure torch.compile settings +- _apply_torch_compile: Apply torch.compile to full model +- _apply_vae_submodule_compile: Apply torch.compile to VAE submodules +- _disable_compile_for_dynamic_modules: Mark dynamic modules no-compile +- _propagate_debug_to_modules: Propagate debug instance to submodules + +BlockSwap Management: +- _handle_blockswap_change: Handle BlockSwap configuration changes + +This module uses model_loader for actual weight loading operations. +""" + +import os +import torch +from omegaconf import OmegaConf +from typing import Dict, Any, Optional, Tuple, Union, Callable + +from .model_loader import ( + prepare_model_structure, + script_directory +) +from .infer import VideoDiffusionInfer +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 ( + CompatibleDiT, + TRITON_AVAILABLE, + validate_attention_mode +) +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 + + +def _configs_equal(config1: Optional[Dict[str, Any]], config2: Optional[Dict[str, Any]]) -> bool: + """ + Compare two configuration dictionaries for equality. + Handles None values properly. + + Args: + config1: First configuration dict (can be None) + config2: Second configuration dict (can be None) + + Returns: + True if configs are equivalent, False otherwise + """ + # Both None = equal + if config1 is None and config2 is None: + return True + + # One None, one not = different + if (config1 is None) != (config2 is None): + return False + + # Compare dictionary contents + return config1 == config2 + + +def _describe_blockswap_config(config: Optional[Dict[str, Any]]) -> str: + """ + Generate human-readable description of BlockSwap configuration. + + Args: + config: BlockSwap configuration dictionary + + Returns: + Human-readable description string + """ + if not is_blockswap_enabled(config): + return "disabled" + + blocks_to_swap = config.get("blocks_to_swap", 0) + swap_io_components = config.get("swap_io_components", False) + + block_text = "block" if blocks_to_swap <= 1 else "blocks" + parts = [f"{blocks_to_swap} {block_text}"] + if swap_io_components: + parts.append("I/O offload") + + return f"enabled ({', '.join(parts)})" + + +def _describe_compile_config(config: Optional[Dict[str, Any]]) -> str: + """ + Generate human-readable description of torch.compile configuration. + + Args: + config: torch.compile configuration dictionary + + Returns: + Human-readable description string + """ + if config is None or not config: + return "disabled" + + # Core parameters + mode = config.get("mode", "default") + backend = config.get("backend", "inductor") + + parts = [f"{mode} mode"] + + # Optional flags + if backend != "inductor": + parts.append(f"{backend} backend") + if config.get("fullgraph", False): + parts.append("fullgraph") + if config.get("dynamic", False): + parts.append("dynamic") + + # Dynamo tuning parameters (show if non-default) + cache_limit = config.get("dynamo_cache_size_limit", 64) + recompile_limit = config.get("dynamo_recompile_limit", 128) + + if cache_limit != 64: + parts.append(f"cache_limit={cache_limit}") + if recompile_limit != 128: + parts.append(f"recompile_limit={recompile_limit}") + + return f"enabled ({', '.join(parts)})" + + +def _describe_attention_mode(attention_mode: Optional[str]) -> str: + """ + Generate human-readable description of attention mode configuration. + + Args: + attention_mode: Attention mode string ('sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3') + + Returns: + Human-readable description string + """ + if attention_mode is None: + return "sdpa (default)" + + mode_descriptions = { + 'sdpa': 'PyTorch SDPA', + '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) + + +def _describe_tiling_config(encode_tiled: bool, encode_tile_size: Optional[Tuple[int, int]], + encode_tile_overlap: Optional[Tuple[int, int]], + decode_tiled: bool, decode_tile_size: Optional[Tuple[int, int]], + decode_tile_overlap: Optional[Tuple[int, int]]) -> str: + """ + Generate human-readable description of VAE tiling configuration. + + Args: + encode_tiled: Whether encode tiling is enabled + encode_tile_size: Tile size for encoding + encode_tile_overlap: Tile overlap for encoding + decode_tiled: Whether decode tiling is enabled + decode_tile_size: Tile size for decoding + decode_tile_overlap: Tile overlap for decoding + + Returns: + Human-readable description string + """ + if not encode_tiled and not decode_tiled: + return "disabled" + + parts = [] + if encode_tiled: + parts.append(f"encode Tile: {encode_tile_size}, Overlap: {encode_tile_overlap}") + if decode_tiled: + parts.append(f"decode Tile: {decode_tile_size}, Overlap: {decode_tile_overlap}") + + return "; ".join(parts) + + +def _update_model_config( + runner: 'VideoDiffusionInfer', + model_attr: str, + model_type: str, + new_configs: Dict[str, Any], + cached_config_attrs: Dict[str, str], + model_config_attrs: Dict[str, str], + config_describers: Dict[str, Callable], + special_handlers: Optional[Dict[str, Callable]] = None, + debug: Optional['Debug'] = None +) -> bool: + """ + Update model configuration uniformly for DiT or VAE. + + This generic function handles config comparison, logging, and attribute updates + for both DiT and VAE models, reducing code duplication. + + Args: + runner: VideoDiffusionInfer instance + model_attr: Attribute name for the model ('dit' or 'vae') + model_type: Model type string for logging ('DiT' or 'VAE') + new_configs: Dict mapping config names to new values + e.g., {'torch_compile': ..., 'block_swap': ...} + cached_config_attrs: Dict mapping config names to runner attribute names + e.g., {'torch_compile': '_dit_compile_args', ...} + model_config_attrs: Dict mapping config names to model storage attribute names. + e.g., {'torch_compile': '_config_compile', ...} + config_describers: Dict mapping config names to description functions + e.g., {'torch_compile': _describe_compile_config, ...} + special_handlers: Optional dict of config-specific handlers + e.g., {'block_swap': handle_blockswap_change} + debug: Debug instance + + Returns: + True if successful + """ + model = getattr(runner, model_attr) + + # Check if model is on meta device (not materialized yet) + try: + param_device = next(model.parameters()).device + if param_device.type == 'meta': + # Model not yet materialized - just update config attributes + for config_name, attr_name in cached_config_attrs.items(): + setattr(runner, attr_name, new_configs.get(config_name)) + # Store on model so config travels with the model when cached + for config_name, attr_name in cached_config_attrs.items(): + model_attr_name = f'_config_{config_name.split("_")[-1]}' + setattr(model, model_attr_name, new_configs.get(config_name)) + return True + except StopIteration: + pass + + # Get cached configurations and check what changed + config_changes = [] + changes_detected = {} + + for config_name, attr_name in cached_config_attrs.items(): + cached_value = getattr(runner, attr_name, None) + new_value = new_configs.get(config_name) + changed = not _configs_equal(cached_value, new_value) + changes_detected[config_name] = changed + + if changed: + # Get description function for this config type + desc_func = config_describers.get(config_name) + if desc_func: + if config_name == 'tiling': + # Special case for tiling which needs multiple params + old_cfg = cached_value or {} + new_cfg = new_value or {} + old_desc = desc_func( + old_cfg.get('encode_tiled', False), + old_cfg.get('encode_tile_size'), + old_cfg.get('encode_tile_overlap'), + old_cfg.get('decode_tiled', False), + old_cfg.get('decode_tile_size'), + old_cfg.get('decode_tile_overlap') + ) + new_desc = desc_func( + new_cfg.get('encode_tiled', False), + new_cfg.get('encode_tile_size'), + new_cfg.get('encode_tile_overlap'), + new_cfg.get('decode_tiled', False), + new_cfg.get('decode_tile_size'), + new_cfg.get('decode_tile_overlap') + ) + else: + old_desc = desc_func(cached_value) + new_desc = desc_func(new_value) + + # Format config name for display + display_name = config_name.replace('_', ' ').title() + if config_name == 'torch_compile': + display_name = 'torch.compile' + elif config_name == 'block_swap': + display_name = 'BlockSwap' + elif config_name == 'attention_mode': + display_name = 'Attention Mode' + + config_changes.append(f"{display_name}: {old_desc} โ†’ {new_desc}") + + # If nothing changed, reuse model as-is + if not any(changes_detected.values()): + debug.log(f"{model_type} configuration unchanged, reusing cached model", category=model_type.lower()) + # Still update attributes to ensure consistency + for config_name, attr_name in cached_config_attrs.items(): + setattr(runner, attr_name, new_configs.get(config_name)) + return True + + # Log configuration changes + debug.log(f"{model_type} configuration changed:", category=model_type.lower(), force=True) + for change in config_changes: + debug.log(f"{change}", category=model_type.lower(), force=True, indent_level=1) + + # Handle torch.compile unwrapping if needed + if changes_detected.get('torch_compile', False): + if model_attr == 'dit' and hasattr(model, '_orig_mod'): + debug.log(f"Removing torch.compile from {model_type}", category="setup") + model = model._orig_mod + setattr(runner, model_attr, model) + elif model_attr == 'vae': + # Unwrap compiled VAE submodules if present + if hasattr(model, 'encoder') and hasattr(model.encoder, '_orig_mod'): + debug.log(f"Removing torch.compile from {model_type} encoder", category="setup") + model.encoder = model.encoder._orig_mod + if hasattr(model, 'decoder') and hasattr(model.decoder, '_orig_mod'): + debug.log(f"Removing torch.compile from {model_type} decoder", category="setup") + model.decoder = model.decoder._orig_mod + + # Execute special handlers for config-specific logic + if special_handlers: + for config_name, handler in special_handlers.items(): + if changes_detected.get(config_name, False): + handler(runner, cached_config_attrs[config_name], + new_configs[config_name], debug) + + # Update config attributes + setattr(runner, model_attr, model) + for config_name, attr_name in cached_config_attrs.items(): + setattr(runner, attr_name, new_configs.get(config_name)) + + # Store on model so config travels with the model when cached + for config_name, attr_name in cached_config_attrs.items(): + # Use explicit mapping if provided, otherwise derive from config name + if model_config_attrs and config_name in model_config_attrs: + model_attr_name = model_config_attrs[config_name] + else: + model_attr_name = f'_config_{config_name.split("_")[-1]}' + setattr(model, model_attr_name, new_configs.get(config_name)) + + # Mark that configs need to be applied + config_needs_app_attr = f'_{model_attr}_config_needs_application' + setattr(runner, config_needs_app_attr, True) + + return True + + +def _handle_blockswap_change( + runner: 'VideoDiffusionInfer', + attr_name: str, + new_config: Optional[Dict[str, Any]], + debug: Optional['Debug'] = None +) -> None: + """ + Handle BlockSwap-specific configuration changes with proper cleanup. + + Called by _update_model_config when BlockSwap configuration changes are detected. + Manages transition between BlockSwap states (enabled/disabled/reconfigured) with + proper cleanup to avoid memory leaks and state corruption. + + Args: + runner: VideoDiffusionInfer instance with DiT model + attr_name: Runner attribute name storing cached config (e.g., '_dit_block_swap_config') + new_config: New BlockSwap configuration dict or None to disable + debug: Debug instance for logging + """ + cached_config = getattr(runner, attr_name, None) + + # Determine BlockSwap status from configs + had_blockswap = is_blockswap_enabled(cached_config) + has_blockswap = is_blockswap_enabled(new_config) + + # If old config had BlockSwap features, clean them up first + if had_blockswap and not has_blockswap: + # Disabling BlockSwap completely + debug.log("Disabling BlockSwap completely", category="blockswap") + cleanup_blockswap(runner=runner, keep_state_for_cache=False) + + # Mark as inactive so the new config can be applied + runner._blockswap_active = False + + +def _update_dit_config( + runner: 'VideoDiffusionInfer', + block_swap_config: Optional[Dict[str, Any]], + torch_compile_args: Optional[Dict[str, Any]], + attention_mode: Optional[str], + debug: Optional['Debug'] = None +) -> bool: + """ + Update DiT model configuration when reusing cached model. + + Compares new configuration settings against cached config to detect changes. + Handles BlockSwap, torch.compile, and attention_mode configuration updates with + proper cleanup and reapplication when settings change. + + Args: + runner: VideoDiffusionInfer instance with cached DiT model + block_swap_config: New BlockSwap configuration dict with keys: + - blocks_to_swap: int - Number of transformer blocks to offload + - swap_io_components: bool - Whether to offload I/O components + torch_compile_args: New torch.compile configuration dict with keys: + - backend: str - Compiler backend (default: "inductor") + - mode: str - Compilation mode (default: "default") + - fullgraph: bool - Require single graph compilation + - 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', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3') + debug: Debug instance for logging + + Returns: + bool: True if configuration was successfully updated + """ + return _update_model_config( + runner=runner, + model_attr='dit', + model_type='DiT', + new_configs={ + 'torch_compile': torch_compile_args, + 'block_swap': block_swap_config, + 'attention_mode': attention_mode + }, + cached_config_attrs={ + 'torch_compile': '_dit_compile_args', + 'block_swap': '_dit_block_swap_config', + 'attention_mode': '_dit_attention_mode' + }, + model_config_attrs={ + 'torch_compile': '_config_compile', + 'block_swap': '_config_swap', + 'attention_mode': '_config_attn' + }, + config_describers={ + 'torch_compile': _describe_compile_config, + 'block_swap': _describe_blockswap_config, + 'attention_mode': _describe_attention_mode + }, + special_handlers={ + 'block_swap': _handle_blockswap_change + }, + debug=debug + ) + + +def _update_vae_config( + runner: 'VideoDiffusionInfer', + torch_compile_args: Optional[Dict[str, Any]], + debug: Optional['Debug'] = None +) -> bool: + """ + Update VAE model configuration when reusing cached model. + + Compares new configuration settings against cached config to detect changes. + Handles torch.compile and tiling configuration updates with proper cleanup + and reapplication when settings change. + + Args: + runner: VideoDiffusionInfer instance with cached VAE model + torch_compile_args: New torch.compile configuration dict with keys: + - backend: str - Compiler backend (default: "inductor") + - mode: str - Compilation mode (default: "default") + - fullgraph: bool - Require single graph compilation + - dynamic: bool - Enable dynamic shapes + - dynamo_cache_size_limit: int - Cache size limit + - dynamo_recompile_limit: int - Recompilation limit + debug: Debug instance for logging + + Returns: + bool: True if configuration was successfully updated + """ + new_tiling_config = getattr(runner, '_new_vae_tiling_config', None) + + return _update_model_config( + runner=runner, + model_attr='vae', + model_type='VAE', + new_configs={ + 'torch_compile': torch_compile_args, + 'tiling': new_tiling_config + }, + cached_config_attrs={ + 'torch_compile': '_vae_compile_args', + 'tiling': '_vae_tiling_config' + }, + model_config_attrs={ + 'torch_compile': '_config_compile', + 'tiling': '_config_tiling' + }, + config_describers={ + 'torch_compile': _describe_compile_config, + 'tiling': _describe_tiling_config + }, + special_handlers=None, + debug=debug + ) + + +def _initialize_cache_context( + dit_cache: bool, + vae_cache: bool, + dit_id: Optional[int], + vae_id: Optional[int], + dit_model: str, + vae_model: str, + debug: Optional['Debug'] = None +) -> Dict[str, Any]: + """ + Initialize cache context with global cache lookups and model name validation. + + Checks the global cache for existing DiT/VAE models and validates that cached + models match the requested model names. Removes stale cache entries when model + names don't match. + + Args: + dit_cache: Whether DiT caching is enabled + vae_cache: Whether VAE caching is enabled + dit_id: Node ID for DiT model lookup + vae_id: Node ID for VAE model lookup + dit_model: Requested DiT model filename for validation + vae_model: Requested VAE model filename for validation + debug: Debug instance for logging + + Returns: + Dict[str, Any]: Cache context dictionary containing: + - global_cache: GlobalModelCache instance + - dit_cache: DiT caching enabled flag + - vae_cache: VAE caching enabled flag + - dit_id: DiT node ID + - vae_id: VAE node ID + - dit_model: DiT model name + - vae_model: VAE model name + - cached_dit: Cached DiT model instance (if found and valid) + - cached_vae: Cached VAE model instance (if found and valid) + - dit_newly_cached: Flag indicating if DiT was just cached + - vae_newly_cached: Flag indicating if VAE was just cached + - reusing_runner: Flag indicating if runner template is reused + + Note: + Automatically removes stale cache entries when model names don't match + the cached model's stored name (_model_name attribute). + """ + global_cache = get_global_cache() + context = { + 'global_cache': global_cache, + 'dit_cache': dit_cache, + 'vae_cache': vae_cache, + 'dit_id': dit_id, + 'vae_id': vae_id, + 'dit_model': dit_model, + 'vae_model': vae_model, + 'cached_dit': None, + 'cached_vae': None, + 'dit_newly_cached': False, + 'vae_newly_cached': False, + 'reusing_runner': False + } + + # Check for cached DiT model with model name validation + # Model name validation prevents stale cache when user switches models in UI + if dit_cache and dit_model and dit_id is not None: + cached_model = global_cache.get_dit({'node_id': dit_id, 'cache_model': True}, debug) + if cached_model: + # Verify cached model matches requested model by checking _model_name attribute + cached_model_name = getattr(cached_model, '_model_name', None) + if cached_model_name == dit_model: + # Cache hit with valid model - reuse it + context['cached_dit'] = cached_model + else: + # Model changed - remove stale cache and log the change + if cached_model_name: + debug.log(f"DiT model changed in cache ({cached_model_name} โ†’ {dit_model}), " + f"removing stale cached model", category="cache", force=True) + global_cache.remove_dit({'node_id': dit_id}, debug) + else: + # Caching disabled or no ID - clean up any existing cache for this node + if dit_id is not None: + global_cache.remove_dit({'node_id': dit_id}, debug) + + # Check for cached VAE model with model name validation + if vae_cache and vae_model and vae_id is not None: + cached_model = global_cache.get_vae({'node_id': vae_id, 'cache_model': True}, debug) + if cached_model: + # Verify cached model matches requested model by checking _model_name attribute + cached_model_name = getattr(cached_model, '_model_name', None) + if cached_model_name == vae_model: + context['cached_vae'] = cached_model + else: + # Model changed - remove stale cache and log the change + if cached_model_name: + debug.log(f"VAE model changed in cache ({cached_model_name} โ†’ {vae_model}), " + f"removing stale cached model", category="cache", force=True) + global_cache.remove_vae({'node_id': vae_id}, debug) + else: + if vae_id is not None: + global_cache.remove_vae({'node_id': vae_id}, debug) + + return context + + +def _acquire_runner( + cache_context: Dict[str, Any], + dit_model: str, + vae_model: str, + base_cache_dir: str, + debug: Optional['Debug'] = None +) -> VideoDiffusionInfer: + """ + Get or create VideoDiffusionInfer runner instance using unified caching approach. + + Checks global cache for existing runner template matching the DiT/VAE node ID pair. + If found and model names match, reuses the template. Otherwise creates new runner. + + Args: + cache_context: Cache context dict from _initialize_cache_context containing: + - global_cache: GlobalModelCache instance + - dit_id: DiT node ID for cache lookup + - vae_id: VAE node ID for cache lookup + - reusing_runner: Flag to be updated if template is reused + dit_model: DiT model filename for validation (e.g., "seedvr2_ema_3b_fp16.safetensors") + vae_model: VAE model filename for validation (e.g., "ema_vae_fp16.safetensors") + base_cache_dir: Base directory for model files + debug: Debug instance for logging + + Returns: + VideoDiffusionInfer: Runner instance (cached template or newly created) + """ + # Try to get runner template from global cache + template = cache_context['global_cache'].get_runner( + cache_context['dit_id'], cache_context['vae_id'], debug + ) + + if template: + # We have a template - check if we can use it + current_dit = getattr(template, '_dit_model_name', None) + current_vae = getattr(template, '_vae_model_name', None) + models_match = (current_dit == dit_model and current_vae == vae_model) + + if models_match: + # Perfect match - reuse template directly + runner_key = f"{cache_context['dit_id']}+{cache_context['vae_id']}" + debug.log(f"Reusing cached runner template: nodes {runner_key}", category="reuse", force=True) + cache_context['reusing_runner'] = True + return template + else: + # Template exists but models changed and no cached models - create new + return _create_new_runner(dit_model, vae_model, base_cache_dir, debug) + else: + # No template - create new runner + return _create_new_runner(dit_model, vae_model, base_cache_dir, debug) + + +def _create_new_runner( + dit_model: str, + vae_model: str, + base_cache_dir: str, + debug: Optional['Debug'] = None +) -> VideoDiffusionInfer: + """ + Create a new VideoDiffusionInfer runner instance from scratch. + + Loads appropriate configuration file based on model size (3B or 7B), creates + runner instance, and initializes with default settings. Called when no cached + runner template is available or when model selection changes. + + Args: + dit_model: DiT model filename (determines config selection) + - Contains "7b" โ†’ loads configs_7b/main.yaml + - Otherwise โ†’ loads configs_3b/main.yaml + vae_model: VAE model filename (stored for reference, not used in config selection) + base_cache_dir: Base directory for model files (not used directly but passed for context) + debug: Debug instance for logging and timing + + Returns: + VideoDiffusionInfer: Newly created runner with: + - Loaded OmegaConf configuration + - Initialized diffusion sampler and schedule + - Config set to mutable (readonly=False) + - No models loaded (structure only) + """ + debug.log(f"Creating new runner: DiT={dit_model}, VAE={vae_model}", + category="runner", force=True) + + debug.start_timer("config_load") + config_path = os.path.join(script_directory, + './configs_7b' if "7b" in dit_model else './configs_3b', + 'main.yaml') + config = load_config(config_path) + debug.end_timer("config_load", "Config loading") + + debug.start_timer("runner_video_infer") + runner = VideoDiffusionInfer(config, debug) + OmegaConf.set_readonly(runner.config, False) + debug.end_timer("runner_video_infer", "Video diffusion inference runner initialization") + + return runner + + +def configure_runner( + dit_model: str, + vae_model: str, + base_cache_dir: str, + debug: 'Debug', + ctx: Dict[str, Any], + dit_cache: bool = False, + vae_cache: bool = False, + dit_id: Optional[int] = None, + vae_id: Optional[int] = None, + block_swap_config: Optional[Dict[str, Any]] = None, + encode_tiled: bool = False, + encode_tile_size: Optional[Tuple[int, int]] = None, + encode_tile_overlap: Optional[Tuple[int, int]] = None, + decode_tiled: bool = False, + decode_tile_size: Optional[Tuple[int, int]] = None, + decode_tile_overlap: Optional[Tuple[int, int]] = None, + tile_debug: str = "false", + attention_mode: str = 'sdpa', + torch_compile_args_dit: Optional[Dict[str, Any]] = None, + torch_compile_args_vae: Optional[Dict[str, Any]] = None +) -> Tuple[VideoDiffusionInfer, Dict[str, Any]]: + """ + Configure VideoDiffusionInfer runner with model loading and settings. + + Handles model changes and caching logic with independent DiT/VAE caching support. + + Args: + dit_model: DiT model filename (e.g., "seedvr2_ema_3b_fp16.safetensors") + vae_model: VAE model filename (e.g., "ema_vae_fp16.safetensors") + base_cache_dir: Base directory containing model files + debug: Debug instance for logging (required) + ctx: Generation context from setup_generation_context + dit_cache: Whether to cache DiT model between runs + vae_cache: Whether to cache VAE model between runs + dit_id: Node instance ID for DiT model caching (required if dit_cache=True) + vae_id: Node instance ID for VAE model caching (required if vae_cache=True) + block_swap_config: Optional BlockSwap configuration for DiT memory optimization + encode_tiled: Enable tiled encoding to reduce VRAM during VAE encoding + encode_tile_size: Tile size for encoding (height, width) + encode_tile_overlap: Tile overlap for encoding (height, width) + decode_tiled: Enable tiled decoding to reduce VRAM during VAE decoding + 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', '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 + + Returns: + Tuple[VideoDiffusionInfer, Dict[str, Any]]: (configured runner, cache context dict) + + Features: + - Independent DiT and VAE caching for flexible memory management + - Dynamic model reloading when models change + - Optional torch.compile optimization for inference speedup + - Separate encode/decode tiling configuration for optimal performance + - Memory optimization and BlockSwap integration + + Raises: + ValueError: If debug instance is not provided + """ + + 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, + dit_model, vae_model, debug + ) + + # Phase 2: Get or create runner + runner = _acquire_runner( + cache_context, dit_model, vae_model, + base_cache_dir, debug + ) + + # Phase 3: Configure runner settings + _configure_runner_settings( + runner, ctx, + encode_tiled, encode_tile_size, encode_tile_overlap, + decode_tiled, decode_tile_size, decode_tile_overlap, + tile_debug, attention_mode, + torch_compile_args_dit, torch_compile_args_vae, + block_swap_config, debug + ) + + # Phase 4: Setup models (load from cache or create new) + _setup_models( + runner, cache_context, dit_model, vae_model, + base_cache_dir, block_swap_config, debug + ) + + return runner, cache_context + + +def _configure_runner_settings( + runner: VideoDiffusionInfer, + ctx: Dict[str, Any], + encode_tiled: bool, + encode_tile_size: Optional[Tuple[int, int]], + encode_tile_overlap: Optional[Tuple[int, int]], + decode_tiled: bool, + decode_tile_size: Optional[Tuple[int, int]], + decode_tile_overlap: Optional[Tuple[int, int]], + tile_debug: str, + attention_mode: str, + torch_compile_args_dit: Optional[Dict[str, Any]], + torch_compile_args_vae: Optional[Dict[str, Any]], + block_swap_config: Optional[Dict[str, Any]], + debug: Optional['Debug'] = None +) -> None: + """ + Configure runner settings for VAE tiling, torch.compile, and BlockSwap. + + Stores configuration settings on runner for later comparison and application. + Settings are stored in temporary "_new_*" attributes and later validated/applied + in _setup_models phase. This separation allows configuration change detection + when reusing cached models. + + Args: + runner: VideoDiffusionInfer instance to configure + ctx: Generation context from setup_generation_context + encode_tiled: Enable tiled VAE encoding to reduce VRAM during encoding + encode_tile_size: Tile dimensions (height, width) for encoding in pixels + encode_tile_overlap: Overlap dimensions (height, width) between encoding tiles + decode_tiled: Enable tiled VAE decoding to reduce VRAM during decoding + 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', '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 + debug: Debug instance (stored on runner for model access) + """ + # VAE tiling settings + runner.encode_tiled = encode_tiled + runner.encode_tile_size = encode_tile_size + runner.encode_tile_overlap = encode_tile_overlap + runner.decode_tiled = decode_tiled + runner.decode_tile_size = decode_tile_size + runner.decode_tile_overlap = decode_tile_overlap + runner.tile_debug = tile_debug + + # Store the new configs temporarily for later comparison + # Don't set them as attributes yet - let the update functions handle that + runner._new_dit_compile_args = torch_compile_args_dit + runner._new_vae_compile_args = torch_compile_args_vae + runner._new_dit_block_swap_config = block_swap_config + runner._new_dit_attention_mode = attention_mode + runner._new_vae_tiling_config = { + 'encode_tiled': encode_tiled, + 'encode_tile_size': encode_tile_size, + 'encode_tile_overlap': encode_tile_overlap, + 'decode_tiled': decode_tiled, + 'decode_tile_size': decode_tile_size, + 'decode_tile_overlap': decode_tile_overlap + } + + # Store device configuration on runner for submodule access (e.g., BlockSwap, Cleanup) + runner._dit_device = ctx['dit_device'] + runner._vae_device = ctx['vae_device'] + runner._dit_offload_device = ctx['dit_offload_device'] + runner._vae_offload_device = ctx['vae_offload_device'] + runner._tensor_offload_device = ctx['tensor_offload_device'] + runner._compute_dtype = ctx['compute_dtype'] + + runner.debug = debug + + +def _setup_models( + runner: VideoDiffusionInfer, + cache_context: Dict[str, Any], + dit_model: str, + vae_model: str, + base_cache_dir: str, + block_swap_config: Optional[Dict[str, Any]], + debug: Optional['Debug'] = None +) -> None: + """ + Setup DiT and VAE models from cache or create new structures. + + Central orchestration function that: + 1. Sets up DiT model (cached or new structure) + 2. Sets up VAE model (cached or new structure) + 3. Validates and updates configurations for cached models + 4. Stores initial configurations for new models + 5. Cleans up temporary configuration attributes + + This function coordinates the complex interaction between model caching, + configuration management, and the meta device initialization strategy. + + Args: + runner: VideoDiffusionInfer instance to setup + cache_context: Cache context from _initialize_cache_context with keys: + - cached_dit, cached_vae: Cached model instances (or None) + - dit_id, vae_id: Node IDs for caching + - Other cache state flags + dit_model: DiT model filename for loading/validation + vae_model: VAE model filename for loading/validation + base_cache_dir: Base directory containing model files + block_swap_config: BlockSwap configuration (passed to DiT setup) + debug: Debug instance for logging and timing + """ + debug.start_timer("model_structures") + + # Setup DiT + dit_created = _setup_dit_model(runner, cache_context, dit_model, base_cache_dir, + block_swap_config, debug) + + # Only update DiT config if model was cached/reused (not newly created) + if not dit_created and hasattr(runner, 'dit') and runner.dit is not None: + _update_dit_config(runner, runner._new_dit_block_swap_config, + runner._new_dit_compile_args, runner._new_dit_attention_mode, debug) + elif dit_created: + # For newly created models, just set initial config attributes (no comparison needed) + runner._dit_compile_args = runner._new_dit_compile_args + runner._dit_block_swap_config = runner._new_dit_block_swap_config + runner._dit_attention_mode = runner._new_dit_attention_mode + # Also store on model so config travels with the model when cached + if hasattr(runner, 'dit') and runner.dit: + runner.dit._config_compile = runner._new_dit_compile_args + runner.dit._config_swap = runner._new_dit_block_swap_config + runner.dit._config_attn = runner._new_dit_attention_mode + + # Setup VAE + vae_created = _setup_vae_model(runner, cache_context, vae_model, base_cache_dir, debug) + + # Only update VAE config if model was cached/reused (not newly created) + if not vae_created and hasattr(runner, 'vae') and runner.vae is not None: + _update_vae_config(runner, runner._new_vae_compile_args, debug) + elif vae_created: + # For newly created models, just set initial config attributes (no comparison needed) + runner._vae_compile_args = runner._new_vae_compile_args + runner._vae_tiling_config = runner._new_vae_tiling_config + # Also store on model so config travels with the model when cached + if hasattr(runner, 'vae') and runner.vae: + runner.vae._config_compile = runner._new_vae_compile_args + runner.vae._config_tiling = runner._new_vae_tiling_config + + # Clean up temporary attributes + for attr in ['_new_dit_compile_args', '_new_vae_compile_args', '_new_dit_block_swap_config', '_new_dit_attention_mode', '_new_vae_tiling_config']: + if hasattr(runner, attr): + delattr(runner, attr) + + debug.end_timer("model_structures", "Model structures prepared") + + +def _setup_dit_model( + runner: VideoDiffusionInfer, + cache_context: Dict[str, Any], + dit_model: str, + base_cache_dir: str, + block_swap_config: Optional[Dict[str, Any]], + debug: Optional['Debug'] = None +) -> bool: + """ + Setup DiT model from cache or create new meta device structure. + + Handles three scenarios: + 1. Model changed: Cleanup old model, create new structure + 2. Cached model available: Reuse cached model, restore config + 3. No model exists: Create new meta device structure + + Args: + runner: VideoDiffusionInfer instance to setup + cache_context: Cache context dict with keys: + - cached_dit: Cached DiT model instance (or None) + - dit_id: Node ID for cache operations + dit_model: DiT model filename (e.g., "seedvr2_ema_3b_fp16.safetensors") + base_cache_dir: Base directory containing model files + block_swap_config: BlockSwap configuration to store (not applied here) + debug: Debug instance for logging + + Returns: + bool: True if new model structure was created, False if cached model reused + """ + + # Check if model changed - clean up old model if different + current_dit_name = getattr(runner, '_dit_model_name', None) + if current_dit_name and current_dit_name != dit_model: + if hasattr(runner, 'dit') and runner.dit is not None: + debug.log(f"DiT model changed ({current_dit_name} โ†’ {dit_model}), cleaning old model", + category="cache", force=True) + cleanup_dit(runner=runner, debug=debug, cache_model=False) + + if cache_context['cached_dit'] is not None: + # Reuse cached DiT model + debug.log(f"Reusing cached DiT ({cache_context['dit_id']}): {dit_model}", + category="reuse", force=True) + runner.dit = cache_context['cached_dit'] + runner._dit_checkpoint = find_model_file(dit_model, base_cache_dir) + runner._dit_model_name = dit_model + + # Restore config attributes from model to runner (config travels with model) + runner._dit_compile_args = getattr(runner.dit, '_config_compile', None) + runner._dit_block_swap_config = getattr(runner.dit, '_config_swap', None) + runner._dit_attention_mode = getattr(runner.dit, '_config_attn', None) + + # blockswap_active will be set by apply_block_swap_to_dit + # when the model is materialized to the inference device + runner._blockswap_active = False + + return False + elif not hasattr(runner, 'dit') or runner.dit is None: + # Create new DiT model + # Set DiT dtype from runner's compute_dtype + # compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16) + # dit_dtype_str = str(compute_dtype).split('.')[-1] + # runner.config.dit.dtype = dit_dtype_str + # runner._dit_dtype_override = compute_dtype + + dit_checkpoint_path = find_model_file(dit_model, base_cache_dir) + runner = prepare_model_structure(runner, "dit", dit_checkpoint_path, + runner.config, debug, block_swap_config) + runner._dit_model_name = dit_model + return True + else: + # Model already exists from previous run with same runner + runner._dit_model_name = dit_model + return False + + +def _setup_vae_model( + runner: VideoDiffusionInfer, + cache_context: Dict[str, Any], + vae_model: str, + base_cache_dir: str, + debug: Optional['Debug'] = None +) -> bool: + """ + Setup VAE model from cache or create new meta device structure. + + Handles three scenarios: + 1. Model changed: Cleanup old model, configure and create new structure + 2. Cached model available: Reuse cached model, restore config + 3. No model exists: Configure VAE settings, create new meta device structure + + Args: + runner: VideoDiffusionInfer instance to setup + cache_context: Cache context dict with keys: + - cached_vae: Cached VAE model instance (or None) + - vae_id: Node ID for cache operations + vae_model: VAE model filename (e.g., "ema_vae_fp16.safetensors") + base_cache_dir: Base directory containing model files + debug: Debug instance for logging + + Returns: + bool: True if new model structure was created, False if cached model reused + """ + + # Check if model changed - clean up old model if different + current_vae_name = getattr(runner, '_vae_model_name', None) + if current_vae_name and current_vae_name != vae_model: + if hasattr(runner, 'vae') and runner.vae is not None: + debug.log(f"VAE model changed ({current_vae_name} โ†’ {vae_model}), cleaning old model", + category="cache", force=True) + cleanup_vae(runner=runner, debug=debug, cache_model=False) + + if cache_context['cached_vae'] is not None: + # Reuse cached VAE model + debug.log(f"Reusing cached VAE ({cache_context['vae_id']}): {vae_model}", + category="reuse", force=True) + runner.vae = cache_context['cached_vae'] + runner._vae_checkpoint = find_model_file(vae_model, base_cache_dir) + runner._vae_model_name = vae_model + + # Restore config attributes from model to runner (config travels with model) + runner._vae_compile_args = getattr(runner.vae, '_config_compile', None) + runner._vae_tiling_config = getattr(runner.vae, '_config_tiling', None) + + return False + elif not hasattr(runner, 'vae') or runner.vae is None: + # Create new VAE model + # Configure VAE + vae_config_path = os.path.join(script_directory, + 'src/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml') + vae_config = load_config(vae_config_path) + + spatial_downsample_factor = vae_config.get('spatial_downsample_factor', 8) + temporal_downsample_factor = vae_config.get('temporal_downsample_factor', 4) + vae_config.spatial_downsample_factor = spatial_downsample_factor + vae_config.temporal_downsample_factor = temporal_downsample_factor + + runner.config.vae.model = OmegaConf.merge(runner.config.vae.model, vae_config) + + # Set VAE dtype from runner's compute_dtype + compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16) + vae_dtype_str = str(compute_dtype).split('.')[-1] + runner.config.vae.dtype = vae_dtype_str + runner._vae_dtype_override = compute_dtype + + vae_checkpoint_path = find_model_file(vae_model, base_cache_dir) + runner = prepare_model_structure(runner, "vae", vae_checkpoint_path, + runner.config, debug, None) + + debug.log( + f"VAE downsample factors configured " + f"(spatial: {spatial_downsample_factor}x, " + f"temporal: {temporal_downsample_factor}x)", + category="vae" + ) + + runner._vae_model_name = vae_model + return True + else: + # Model already exists from previous run with same runner + runner._vae_model_name = vae_model + return False + + +def apply_model_specific_config(model: torch.nn.Module, runner: VideoDiffusionInfer, + config: OmegaConf, is_dit: bool, + debug: Optional['Debug'] = None) -> torch.nn.Module: + """ + Apply model-specific configurations (FP8, BlockSwap, torch.compile). + + This function is idempotent and can be safely called on both newly materialized + and already-configured models. It checks state flags to determine what needs + to be applied. + + Critical: For DiT, BlockSwap must be applied BEFORE torch.compile. + torch.compile captures the computational graph, so any wrapping done + after compilation (like BlockSwap's forward wrapping) won't work. + + Args: + model: Loaded model instance + runner: Runner to attach model to + config: Full configuration object + is_dit: Whether this is a DiT model (vs VAE) + debug: Debug instance + + Returns: + Configured model with BlockSwap and torch.compile applied if configured + """ + if is_dit: + # DiT-specific + # Apply compatibility wrapper with compute_dtype + if not isinstance(model, CompatibleDiT): + debug.log("Applying DiT compatibility wrapper", category="setup") + debug.start_timer("CompatibleDiT") + # Get compute_dtype from runner if available, fallback to bfloat16 + compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16) + model = CompatibleDiT(model, debug, compute_dtype=compute_dtype, skip_conversion=False) + debug.end_timer("CompatibleDiT", "Compatibility wrapper application") + else: + 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_attention_mode(requested_attention_mode, debug) + + # Get compute_dtype from runner + compute_dtype = getattr(runner, '_compute_dtype', torch.bfloat16) + debug.log(f"Applying {attention_mode} attention mode and {compute_dtype} compute dtype to model", category="setup") + + # Get the actual model (unwrap if needed) + actual_model = model.dit_model if hasattr(model, 'dit_model') else model + + # Update all FlashAttentionVarlen instances + updated_count = 0 + for module in actual_model.modules(): + if type(module).__name__ == 'FlashAttentionVarlen': + module.attention_mode = attention_mode + module.compute_dtype = compute_dtype + updated_count += 1 + + if updated_count > 0: + debug.log(f"Applied {attention_mode} and compute_dtype={compute_dtype} to {updated_count} modules", category="success") + + # Apply BlockSwap before torch.compile (only if not already active) + # BlockSwap wraps forward methods, and torch.compile needs to capture the wrapped version + if hasattr(runner, '_dit_block_swap_config') and runner._dit_block_swap_config: + # Check if BlockSwap needs to be applied + needs_blockswap = not (hasattr(runner, '_blockswap_active') and runner._blockswap_active) + if needs_blockswap: + runner.dit = model + apply_block_swap_to_dit(runner, runner._dit_block_swap_config, debug) + # Mark as active after successful application + runner._blockswap_active = True + + # Apply torch.compile after BlockSwap (only if not already compiled) + # Check: model has _orig_mod attribute means it's already torch.compiled + if hasattr(runner, '_dit_compile_args') and runner._dit_compile_args: + if not hasattr(model, '_orig_mod'): + model = _apply_torch_compile(model, runner._dit_compile_args, "DiT", debug) + + runner.dit = model + + # Clear the config application flag after successful application + if hasattr(runner, '_dit_config_needs_application'): + runner._dit_config_needs_application = False + + else: + # VAE-specific configurations + # Set to eval mode (no gradients needed for inference) + if model.training: + debug.log("VAE model set to eval mode (gradients disabled)", category="vae") + debug.start_timer("model_requires_grad") + model.requires_grad_(False).eval() + debug.end_timer("model_requires_grad", "VAE model set to eval mode") + + # Configure causal slicing if available - always apply as it's lightweight + if hasattr(model, "set_causal_slicing") and hasattr(config.vae, "slicing"): + debug.log("Configuring VAE causal slicing for temporal processing", category="vae") + debug.start_timer("vae_set_causal_slicing") + model.set_causal_slicing(**config.vae.slicing) + debug.end_timer("vae_set_causal_slicing", "VAE causal slicing configuration") + + # Set memory limits if available - always apply to ensure limits are set + if hasattr(model, "set_memory_limit") and hasattr(config.vae, "memory_limit"): + debug.log("Configuring VAE memory limits for causal convolutions", category="vae") + debug.start_timer("vae_set_memory_limit") + model.set_memory_limit(**config.vae.memory_limit) + debug.end_timer("vae_set_memory_limit", "VAE memory limits configured") + + # Apply torch.compile if configured (only if not already compiled) + if hasattr(runner, '_vae_compile_args') and runner._vae_compile_args: + # Check if encoder/decoder already compiled (have _orig_mod) + encoder_compiled = hasattr(model, 'encoder') and hasattr(model.encoder, '_orig_mod') + decoder_compiled = hasattr(model, 'decoder') and hasattr(model.decoder, '_orig_mod') + + if not (encoder_compiled and decoder_compiled): + model = _apply_vae_submodule_compile(model, runner._vae_compile_args, debug) + else: + debug.log("Reusing existing torch.compile for VAE submodules", category="reuse") + + # Propagate debug and tensor_offload_device to submodules + model.debug = debug + model.tensor_offload_device = runner._tensor_offload_device + _propagate_debug_to_modules(model, debug) + runner.vae = model + + # Clear the config application flag after successful application + if hasattr(runner, '_vae_config_needs_application'): + runner._vae_config_needs_application = False + + return model + + +def _configure_torch_compile(compile_args: Dict[str, Any], model_type: str, + debug: Optional['Debug'] = None) -> Tuple[Dict[str, Any], bool]: + """ + Extract and configure torch.compile settings with dependency validation. + + Centralizes common configuration logic for both full model and submodule compilation. + Validates that required dependencies (like Triton for inductor backend) are available + before attempting compilation to provide early, clear error messages. + + Args: + compile_args: Compilation configuration dictionary + model_type: Model type string for logging (e.g., "DiT", "VAE") + debug: Debug instance for logging + + Returns: + Tuple of (compile_settings dict, success boolean) + compile_settings contains: backend, mode, fullgraph, dynamic + + Raises: + RuntimeError: If inductor backend is requested but Triton is not available + """ + # Extract settings with defaults + settings = { + 'backend': compile_args.get('backend', 'inductor'), + 'mode': compile_args.get('mode', 'default'), + 'fullgraph': compile_args.get('fullgraph', False), + 'dynamic': compile_args.get('dynamic', False) + } + + dynamo_cache_size_limit = compile_args.get('dynamo_cache_size_limit', 64) + dynamo_recompile_limit = compile_args.get('dynamo_recompile_limit', 128) + + # Check Triton availability for inductor backend BEFORE attempting compilation + if settings['backend'] == 'inductor': + if not TRITON_AVAILABLE: + error_msg = ( + f"Cannot use torch.compile with 'inductor' backend: Triton is not installed.\n" + f"\n" + f"Triton is required for the inductor backend which performs kernel fusion and optimization.\n" + f"\n" + f"To fix this issue:\n" + f" 1. Install Triton: pip install triton\n" + f" 2. OR change backend to 'cudagraphs' (lightweight, no Triton needed)\n" + f" 3. OR disable torch.compile\n" + f"\n" + f"For more info: https://github.com/triton-lang/triton" + ) + debug.log(error_msg, level="ERROR", category="setup", force=True) + raise RuntimeError( + "torch.compile with inductor backend requires Triton. " + "Install with: pip install triton" + ) + + # Log compilation configuration + debug.log(f"Configuring torch.compile for {model_type}...", category="setup", force=True) + debug.log(f"Backend: {settings['backend']} | Mode: {settings['mode']} | " + f"Fullgraph: {settings['fullgraph']} | Dynamic: {settings['dynamic']}", + category="setup", indent_level=1) + debug.log(f"Dynamo cache_size_limit: {dynamo_cache_size_limit} | " + f"recompile_limit: {dynamo_recompile_limit}", category="setup", indent_level=1) + + # Configure torch._dynamo settings + try: + import torch._dynamo + torch._dynamo.config.cache_size_limit = dynamo_cache_size_limit + torch._dynamo.config.recompile_limit = dynamo_recompile_limit + debug.log(f"torch._dynamo configured successfully", category="success", indent_level=1) + return settings, True + except Exception as e: + debug.log(f"Could not configure torch._dynamo settings: {e}", + level="WARNING", category="setup", force=True, indent_level=1) + return settings, False + + +def _apply_torch_compile(model: torch.nn.Module, compile_args: Dict[str, Any], + model_type: str, debug: Optional['Debug'] = None) -> torch.nn.Module: + """ + Apply torch.compile to entire model with configured settings. + + Args: + model: Model to compile + compile_args: Compilation configuration + model_type: "DiT" or "VAE" for logging + debug: Debug instance + + Returns: + Compiled model, or original model if compilation fails + """ + try: + # Configure compilation settings + settings, _ = _configure_torch_compile(compile_args, model_type, debug) + + # Compile entire model + debug.start_timer(f"{model_type.lower()}_compile") + compiled_model = torch.compile(model, **settings) + debug.end_timer(f"{model_type.lower()}_compile", + f"{model_type} model wrapped for compilation", force=True) + debug.log(f"Actual compilation will happen on first batch (expect initial delay, then speedup)", category="info", indent_level=1) + + return compiled_model + + except Exception as e: + debug.log(f"torch.compile failed for {model_type}: {e}", + level="WARNING", category="setup", force=True) + debug.log(f"Falling back to uncompiled model", + level="WARNING", category="setup", force=True, indent_level=1) + return model + + +def _disable_compile_for_dynamic_modules(module: torch.nn.Module) -> None: + """ + Mark modules with dynamic shapes to be excluded from torch.compile. + This prevents recompilation issues with variable tensor sizes. + """ + for name, submodule in module.named_modules(): + if isinstance(submodule, InflatedCausalConv3d): + # Mark module to skip compilation + submodule._dynamo_disable = True + + +def _apply_vae_submodule_compile(model: torch.nn.Module, compile_args: Dict[str, Any], + debug: Optional['Debug'] = None) -> torch.nn.Module: + """ + Apply torch.compile to VAE core submodules instead of entire model. + + The VAE's high-level encode/decode methods contain complex control flow + (temporal slicing, tiling, stateful memory management) that prevents + torch.compile from optimizing effectively. Instead, we compile only the + core neural networks (encoder, decoder) which have straightforward forward + passes suitable for compilation. + + Note: quant_conv and post_quant_conv are disabled (None) in the current VAE + architecture/yaml, so they are not compiled. + + Args: + model: VAE model instance + compile_args: Compilation configuration + debug: Debug instance + + Returns: + VAE model with compiled submodules + """ + try: + # Configure compilation settings + settings, _ = _configure_torch_compile(compile_args, "VAE submodules", debug) + + # Compile submodules + compiled_modules = [] + debug.start_timer("vae_submodule_compile") + + if hasattr(model, 'encoder') and model.encoder is not None: + # Disable compilation for InflatedCausalConv3d modules due to dynamic shapes + _disable_compile_for_dynamic_modules(model.encoder) + model.encoder = torch.compile(model.encoder, **settings) + compiled_modules.append('encoder') + debug.log(f"VAE encoder found and added to compilation queue", category="success", indent_level=1) + + if hasattr(model, 'decoder') and model.decoder is not None: + # Disable compilation for InflatedCausalConv3d modules due to dynamic shapes + _disable_compile_for_dynamic_modules(model.decoder) + model.decoder = torch.compile(model.decoder, **settings) + compiled_modules.append('decoder') + debug.log(f"VAE decoder found and added to compilation queue", category="success", indent_level=1) + + debug.end_timer("vae_submodule_compile", + f"VAE submodules compiled: {', '.join(compiled_modules)}", force=True) + debug.log(f"Actual compilation will happen on first batch (expect initial delay, then speedup)", category="info", indent_level=1) + + return model + + except Exception as e: + debug.log(f"torch.compile failed for VAE submodules: {e}", + level="WARNING", category="setup", force=True) + debug.log(f"Falling back to uncompiled VAE", + level="WARNING", category="setup", force=True, indent_level=1) + return model + + +def _propagate_debug_to_modules(module: torch.nn.Module, debug: 'Debug') -> None: + """ + Propagate debug instance to specific submodules that need it. + Only targets modules that actually use debug to avoid unnecessary memory overhead. + + Args: + module: Parent module to propagate through + debug: Debug instance to attach + """ + if debug is None: + return # Early exit if no debug instance + + target_modules = {'ResnetBlock3D', 'Upsample3D', 'InflatedCausalConv3d', 'GroupNorm'} + + 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 \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_loader.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..3ce10788b6f4ed221c56310865eb8d45560b9645 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/core/model_loader.py @@ -0,0 +1,965 @@ +""" +Model Weight Loading for SeedVR2 + +This module handles all weight loading operations for DiT and VAE models: +- Loading state dictionaries from multiple formats (SafeTensors, PyTorch, GGUF) +- Materializing models from meta device to target device +- Applying weights with dtype conversion +- GGUF quantized model support with dequantization +- Meta buffer initialization for non-persistent buffers + +Key Features: +- Multi-format support: .safetensors, .pth, .gguf files +- Memory-efficient loading with meta device initialization +- Native FP8 weight handling with optimal performance +- GGUF quantization support (Q4_K_M, Q8_0, etc.) +- Automatic dtype conversion for compatibility +- Meta buffer initialization post-materialization + +Main Functions: +- load_quantized_state_dict: Load state dict from checkpoint file +- materialize_model: Move model from meta device and load weights +- prepare_model_structure: Create model structure on meta device + +GGUF Support: +- apply_gguf_parameters: Apply GGUF weights to model (handles meta and materialized) +- _load_gguf_state: Load GGUF quantized weights from file +- _load_gguf_weights: Apply GGUF weights to model with validation +- _validate_gguf_architecture: Validate GGUF model architecture +- _create_dequantize_method: Create dequantization callable +- _create_gguf_parameter: Create parameter preserving quantization info +- _set_parameter_on_meta_model: Set parameter on meta device model +- _set_parameter_on_materialized_model: Set parameter on materialized model +- _navigate_to_parameter: Navigate to module containing parameter +- _get_tensor_shape: Get logical shape of tensor (handling GGUF) +- _is_quantized_tensor: Check if tensor is GGUF quantized +- _report_parameter_mismatches: Report parameter mismatches + +Meta Buffer Initialization: +- initialize_meta_buffers: Initialize meta buffers with timing wrapper +- initialize_meta_buffers_impl: Initialize non-persistent buffers on target device + +Standard Loading: +- _load_model_weights: Orchestrate weight loading process +- _load_standard_weights: Apply SafeTensors/PyTorch weights +- _convert_state_dtype: Convert weight dtypes +- _log_weight_stats: Log weight statistics + +This module is used by model_configuration for weight loading during materialization. +""" + +import os +import torch +from omegaconf import OmegaConf +from typing import Dict, Any, Optional, Tuple, Union, Callable + +# Import SafeTensors with fallback +try: + from safetensors.torch import load_file as load_safetensors_file + SAFETENSORS_AVAILABLE = True +except ImportError: + SAFETENSORS_AVAILABLE = False + +from .infer import VideoDiffusionInfer +from ..common.config import create_object +from ..optimization.compatibility import ( + GGUF_AVAILABLE, + GGMLQuantizationType, + validate_gguf_availability +) + +# GGUF-specific imports (only when available) +if GGUF_AVAILABLE: + import gguf + import traceback + from ..optimization.gguf_dequant import dequantize_tensor + from ..optimization.gguf_ops import replace_linear_with_quantized + +from ..utils.constants import get_script_directory, suppress_tensor_warnings + +# Get script directory for config paths +script_directory = get_script_directory() + + +def load_quantized_state_dict(checkpoint_path: str, device: torch.device = torch.device("cpu"), + debug: Optional['Debug'] = None) -> Dict[str, torch.Tensor]: + """ + Load model state dict from checkpoint with support for multiple formats. + + Handles .safetensors, .gguf, and .pth files. GGUF models support quantization + for memory-efficient loading. Validates required libraries are installed. + + Args: + checkpoint_path: Path to checkpoint file + device: Target device for tensor placement (torch.device object, defaults to CPU) + debug: Optional Debug instance for logging + + Returns: + dict: State dictionary loaded with appropriate format handler + + Notes: + - SafeTensors files use optimized loading with direct device placement + - PyTorch files use memory-mapped loading to reduce RAM usage + """ + device_str = str(device) + + if checkpoint_path.endswith('.safetensors'): + if not SAFETENSORS_AVAILABLE: + error_msg = ( + f"Cannot load {os.path.basename(checkpoint_path)}\n" + f"SafeTensors library is required but not installed.\n" + f"Please install it with: pip install safetensors" + ) + if debug: + debug.log(error_msg, level="ERROR", category="dit", force=True) + 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) + + # 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( + checkpoint_path=checkpoint_path, + device=device, + debug=debug, + handle_prefix="model.diffusion_model." + ) + elif checkpoint_path.endswith('.pth'): + 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}") + + return state + + +def _load_gguf_state(checkpoint_path: str, device: torch.device, debug: Optional['Debug'] = None, + handle_prefix: str = "model.diffusion_model.") -> Dict[str, torch.Tensor]: + """ + Load GGUF state dict + + Args: + checkpoint_path: Path to GGUF file + device: Target device (torch.device object) + debug: Debug instance + handle_prefix: Prefix to strip from tensor names + + Returns: + State dictionary with loaded tensors + """ + reader = gguf.GGUFReader(checkpoint_path) + + # Filter and strip prefix + has_prefix = False + if handle_prefix is not None: + prefix_len = len(handle_prefix) + tensor_names = set(tensor.name for tensor in reader.tensors) + has_prefix = any(s.startswith(handle_prefix) for s in tensor_names) + + tensors = [] + for tensor in reader.tensors: + sd_key = tensor_name = tensor.name + if has_prefix: + if not tensor_name.startswith(handle_prefix): + continue + sd_key = tensor_name[prefix_len:] + tensors.append((sd_key, tensor)) + + state_dict = {} + total_tensors = len(reader.tensors) + + device_str = str(device) + debug.log(f"Loading {total_tensors} tensors to {str(device_str)}...", category="dit") + + # Suppress expected warnings: GGUF tensors are read-only numpy arrays that trigger warnings when converted + suppress_tensor_warnings() + + for i, (sd_key, tensor) in enumerate(tensors): + tensor_name = tensor.name + + # Create tensor directly on target device to avoid CPU->GPU copy overhead + # For meta-initialized models, this directly materializes to the target device + torch_tensor = torch.from_numpy(tensor.data).to(device, non_blocking=False) + + # Get original shape from metadata or infer from tensor shape + shape = _get_tensor_logical_shape(reader, tensor_name) + if shape is None: + shape = torch.Size(tuple(int(v) for v in reversed(tensor.shape))) + + # Handle tensors based on quantization type + if tensor.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + # For unquantized tensors, just reshape + torch_tensor = torch_tensor.view(*shape) + else: + # For quantized tensors, keep them quantized but track original shape + torch_tensor = GGUFTensor(torch_tensor, tensor_type=tensor.tensor_type, tensor_shape=shape, debug=debug) + + state_dict[sd_key] = torch_tensor + + # Progress reporting + if (i + 1) % 100 == 0: + debug.log(f"Loaded {i+1}/{total_tensors} tensors...", category="dit", indent_level=1) + + debug.log(f"Successfully loaded {len(state_dict)} tensors to {device_str}", category="success") + + return state_dict + + +def _get_tensor_logical_shape(reader: 'gguf.GGUFReader', tensor_name: str) -> Optional[torch.Size]: + """ + Extract the logical (unquantized) shape from GGUF metadata + """ + field_key = f"comfy.gguf.orig_shape.{tensor_name}" + field = reader.get_field(field_key) + if field is None: + return None + # Has original shape metadata, so we try to decode it. + if len(field.types) != 2 or field.types[0] != gguf.GGUFValueType.ARRAY or field.types[1] != gguf.GGUFValueType.INT32: + raise TypeError(f"Bad original shape metadata for {field_key}: Expected ARRAY of INT32, got {field.types}") + return torch.Size(tuple(int(field.parts[part_idx][0]) for part_idx in field.data)) + + +class GGUFTensor(torch.Tensor): + """ + Tensor wrapper for GGUF quantized tensors that preserves quantization info + """ + def __init__(self, *args, tensor_type, tensor_shape, **kwargs): + super().__init__() + self.tensor_type = tensor_type + self.tensor_shape = tensor_shape + + def __new__(cls, *args, tensor_type, tensor_shape, debug, **kwargs): + # Create tensor with requires_grad=False to avoid gradient issues + tensor = super().__new__(cls, *args, **kwargs) + tensor.requires_grad_(False) + tensor.tensor_type = tensor_type + tensor.tensor_shape = tensor_shape + tensor.debug = debug + return tensor + + def to(self, *args, **kwargs): + new = super().to(*args, **kwargs) + new.tensor_type = getattr(self, "tensor_type", None) + new.tensor_shape = getattr(self, "tensor_shape", self.tensor_shape if hasattr(self, "tensor_shape") else new.shape) + new.debug = getattr(self, "debug", None) + new.requires_grad_(False) # Ensure no gradients + return new + + @property + def shape(self): + # Always return the logical tensor shape, not the quantized data shape + if hasattr(self, "tensor_shape"): + return self.tensor_shape + else: + # Fallback to actual data shape if tensor_shape is not available + return self.size() + + def size(self, *args): + # Override size() to also return logical shape + if hasattr(self, "tensor_shape") and len(args) == 0: + return self.tensor_shape + elif hasattr(self, "tensor_shape") and len(args) == 1: + return self.tensor_shape[args[0]] + else: + return super().size(*args) + + def dequantize(self, device=None, dtype=torch.float16, dequant_dtype=None): + """Dequantize this tensor to its original shape""" + if device is None: + device = self.device + + # Suppress expected warning when converting from GGUFTensor subclass to regular tensor + suppress_tensor_warnings() + + # Check if already unquantized + if self.tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + # Return regular tensor, not GGUFTensor + result = self.to(device, dtype) + if isinstance(result, GGUFTensor): + # Convert to regular tensor to avoid __torch_function__ calls + result = torch.tensor(result, dtype=dtype, device=device, requires_grad=False) + return result + + # Try fast dequantization with crash protection + try: + result = dequantize_tensor(self, dtype, dequant_dtype) + final_result = result.to(device) + + # Ensure we return a regular tensor, not GGUFTensor + if isinstance(final_result, GGUFTensor): + final_result = torch.tensor(final_result.data, dtype=dtype, device=device, requires_grad=False) + + return final_result + except Exception as e: + self.debug.log(f"Fast dequantization failed: {e}", level="WARNING", category="dit", force=True) + self.debug.log(f"Falling back to numpy dequantization", level="WARNING", category="dit", force=True) + + # Fallback to numpy (slower but reliable) + try: + numpy_data = self.cpu().numpy() + dequantized = gguf.quants.dequantize(numpy_data, self.tensor_type) + result = torch.from_numpy(dequantized).to(device, dtype) + result.requires_grad_(False) + final_result = result.reshape(self.tensor_shape) + # from_numpy already returns a regular tensor, no conversion needed + return final_result + except Exception as e: + self.debug.log(f"Numpy fallback also failed: {e}", level="WARNING", category="dit", force=True) + self.debug.log(f"Tensor type: {self.tensor_type}", level="WARNING", category="dit", force=True, indent_level=1) + self.debug.log(f"Shape: {self.shape}", level="WARNING", category="dit", force=True, indent_level=1) + self.debug.log(f"Target shape: {self.tensor_shape}", level="WARNING", category="dit", force=True, indent_level=1) + traceback.print_exc() + + # Return regular tensor as last resort + result = self.to(device, dtype) + if isinstance(result, GGUFTensor): + result = torch.tensor(result.data, dtype=dtype, device=device, requires_grad=False) + return result + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + """Override torch function calls to automatically dequantize""" + if kwargs is None: + kwargs = {} + + # Find the GGUFTensor instance(s) in args + gguf_tensors = [arg for arg in args if isinstance(arg, cls)] + if not gguf_tensors: + return super().__torch_function__(func, types, args, kwargs) + + # Use the first GGUFTensor instance for attribute access + self = gguf_tensors[0] + + # Check if the tensor is fully constructed and still quantized + tensor_type = getattr(self, 'tensor_type', None) + if tensor_type is None: + # Tensor is either being constructed or already dequantized + return super().__torch_function__(func, types, args, kwargs) + + # Check if tensor is already unquantized (F32/F16) + if tensor_type in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + return super().__torch_function__(func, types, args, kwargs) + + # Check if debug exists before using it + debug = getattr(self, 'debug', None) + + # Handle linear operations specially + if func == torch.nn.functional.linear: + if len(args) >= 2 and isinstance(args[1], cls): # weight is the second argument + 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 linear dequantization: {e}", level="WARNING", category="dit", force=True) + debug.log(f"Function: {func}", level="WARNING", category="dit", force=True, indent_level=1) + debug.log(f"Args: {[arg.shape if hasattr(arg, 'shape') else type(arg) for arg in args]}", level="WARNING", category="dit", force=True, indent_level=1) + raise + + # Handle matrix multiplication operations that need dequantization + if func in {torch.matmul, torch.mm, torch.bmm, torch.addmm, torch.addmv, + torch.addr, torch.baddbmm, torch.chain_matmul}: + try: + new_args = [] + for arg in args: + if isinstance(arg, cls): + new_args.append(arg.dequantize()) + else: + new_args.append(arg) + return func(*tuple(new_args), **kwargs) + except Exception as e: + 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. + return super().__torch_function__(func, types, args, kwargs) + + +def prepare_model_structure( + runner: VideoDiffusionInfer, + model_type: str, + checkpoint_path: str, + config: OmegaConf, + debug: 'Debug', + block_swap_config: Optional[Dict[str, Any]] = None +) -> VideoDiffusionInfer: + """ + Prepare model structure on meta device without loading weights. + This uses zero memory as meta device doesn't allocate real memory. + + Args: + runner: VideoDiffusionInfer instance + model_type: "dit" or "vae" + checkpoint_path: Path to checkpoint (stored for later loading) + config: Model configuration + debug: Debug instance for logging (required) + block_swap_config: BlockSwap config (stored for DiT, optional) + + Returns: + runner: Updated runner with model structure on meta device + """ + if debug is None: + raise ValueError(f"Debug instance required for prepare_model_structure") + + is_dit = (model_type == "dit") + model_type_upper = "DiT" if is_dit else "VAE" + model_config = config.dit.model if is_dit else config.vae.model + + # Always create on meta device for zero memory usage + debug.log(f"Creating {model_type_upper} model structure on meta device", + category=model_type, force=True) + debug.start_timer(f"{model_type}_structure") + + with torch.device("meta"): + model = create_object(model_config) + + debug.end_timer(f"{model_type}_structure", f"{model_type_upper} structure created") + + # Store model and config for later materialization + if is_dit: + runner.dit = model + runner._dit_checkpoint = checkpoint_path + runner._dit_block_swap_config = block_swap_config + else: + runner.vae = model + runner._vae_checkpoint = checkpoint_path + + return runner + + +def materialize_model(runner: VideoDiffusionInfer, model_type: str, device: torch.device, + config: OmegaConf, debug: 'Debug') -> None: + """ + Materialize model weights from checkpoint to memory. + Call this right before the model is needed. + + Args: + runner: Runner with model structure on meta device + model_type: "dit" or "vae" + device: Target device for inference (torch.device object) + config: Full configuration + debug: Debug instance + """ + if debug is None: + raise ValueError(f"Debug instance required for materialize_model") + + is_dit = (model_type == "dit") + model_type_upper = "DiT" if is_dit else "VAE" + + # Get model and checkpoint path + if is_dit: + model = runner.dit + checkpoint_path = runner._dit_checkpoint + block_swap_config = runner._dit_block_swap_config + override_dtype = getattr(runner, '_dit_dtype_override', None) + else: + model = runner.vae + checkpoint_path = runner._vae_checkpoint + block_swap_config = None + override_dtype = getattr(runner, '_vae_dtype_override', None) + + # Check if already materialized + if model is None: + debug.log(f"No {model_type_upper} model structure found", level="WARNING", category=model_type, force=True) + return + param_device = next(model.parameters()).device + if param_device.type != 'meta': + debug.log(f"{model_type_upper} already materialized on {model.device}", category=model_type) + return + + # Determine target device for materialization + offload_device_str = None + if hasattr(runner, f'_{model_type}_offload_device'): + offload_device_str = getattr(runner, f'_{model_type}_offload_device') + + # If offload_device is set and not "none", materialize to offload device + if offload_device_str and offload_device_str != "none": + target_device = torch.device(offload_device_str) + offload_reason = " (offload device)" + else: + # Otherwise materialize to inference device + target_device = device + offload_reason = "" + + # Start materialization + debug.start_timer(f"{model_type}_materialize") + + # Load weights (this materializes from meta to target device) + model = _load_model_weights(model, checkpoint_path, target_device, True, + model_type_upper, offload_reason, debug, override_dtype) + + # Apply model-specific configurations (includes BlockSwap and torch.compile) + # Import here to avoid circular dependency + from .model_configuration import apply_model_specific_config + model = apply_model_specific_config(model, runner, config, is_dit, debug) + + debug.end_timer(f"{model_type}_materialize", f"{model_type_upper} materialized") + + # Clean up checkpoint paths (no longer needed after weights are loaded) + # Note: Config attributes (_dit_block_swap_config, _dit_compile_args) are preserved + # for configuration change detection on subsequent runs + if is_dit: + runner._dit_checkpoint = None + runner._dit_dtype_override = None + else: + runner._vae_checkpoint = None + runner._vae_dtype_override = None + + +def _load_model_weights(model: torch.nn.Module, checkpoint_path: str, target_device: torch.device, + used_meta: bool, model_type: str, cpu_reason: str, + debug: Optional['Debug'] = None, override_dtype: Optional[torch.dtype] = None) -> torch.nn.Module: + """ + Load model weights from checkpoint file with optimized GGUF support. + + For meta-initialized models, materializes to target device. + For standard models, loads weights and applies state dict. + + Args: + model: Model instance (may be on meta device) + checkpoint_path: Path to checkpoint file + target_device: Target device for weights (torch.device object) + used_meta: Whether model was created on meta device + model_type: Model type string for logging + cpu_reason: Reason string if using CPU + debug: Debug instance + override_dtype: Optional dtype override for weights + + Returns: + Model with loaded weights + """ + model_type_lower = model_type.lower() + + # Log loading action + action = "Materializing" if used_meta else "Loading" + target_device_str = str(target_device).upper() + debug.log(f"{action} {model_type} weights to {target_device_str}{cpu_reason}: {checkpoint_path}", + category=model_type_lower, force=True) + + # Load state dict from file + debug.start_timer(f"{model_type_lower}_weights_load") + state = load_quantized_state_dict(checkpoint_path, target_device, debug) + debug.end_timer(f"{model_type_lower}_weights_load", f"{model_type} weights loaded from file") + + # Apply dtype conversion if requested + if override_dtype is not None: + state = _convert_state_dtype(state, override_dtype, model_type, debug) + + # Log weight statistics + _log_weight_stats(state, used_meta, model_type, debug) + + # Handle GGUF or standard loading + if checkpoint_path.endswith('.gguf'): + model = _load_gguf_weights(model, state, used_meta, model_type_lower, debug) + else: + model = _load_standard_weights(model, state, used_meta, model_type, model_type_lower, debug) + + # Clean up state dict + del state + + # Initialize meta buffers if needed + if used_meta: + initialize_meta_buffers(model, target_device, debug) + + return model + + +def _convert_state_dtype(state: Dict[str, torch.Tensor], target_dtype: torch.dtype, + model_type: str, debug: Optional['Debug'] = None) -> Dict[str, torch.Tensor]: + """Convert floating point tensors in state dict to target dtype.""" + debug.log(f"Converting {model_type} weights to {target_dtype} during loading", category="precision") + debug.start_timer(f"{model_type.lower()}_dtype_convert") + + for key in state: + if torch.is_tensor(state[key]) and state[key].is_floating_point(): + state[key] = state[key].to(target_dtype) + + debug.end_timer(f"{model_type.lower()}_dtype_convert", f"{model_type} weights converted to {target_dtype}") + return state + + +def _log_weight_stats(state: Dict[str, torch.Tensor], used_meta: bool, model_type: str, debug: Optional['Debug'] = None) -> None: + """Log statistics about loaded weights.""" + num_params = len(state) + total_size_mb = sum(p.nelement() * p.element_size() for p in state.values()) / (1024 * 1024) + action = "Materializing" if used_meta else "Applying" + debug.log(f"{action} {model_type}: {num_params} parameters, {total_size_mb:.2f}MB total", + category=model_type.lower()) + + +def apply_gguf_parameters(model: torch.nn.Module, state: Dict[str, torch.Tensor], + model_state: Dict[str, torch.Tensor], debug: Optional['Debug'] = None) -> Dict[str, Any]: + """ + Apply GGUF parameters to model, handling both meta and materialized models. + + Returns: + Statistics dictionary with loaded count, quantized count, and parameter names + """ + loaded_names = set() + quantized_count = 0 + + for name, param in state.items(): + if name not in model_state: + continue + + model_param = model_state[name] + param_shape = _get_tensor_shape(param) + + if param_shape != model_param.shape: + debug.log(f"Unexpected shape mismatch for {name}: {param_shape} vs {model_param.shape}", + level="ERROR", category="dit", force=True) + raise ValueError(f"Shape mismatch for parameter {name}") + + # Apply parameter based on device type + with torch.no_grad(): + if model_param.device.type == 'meta': + _set_parameter_on_meta_model(model, name, param, debug) + else: + _set_parameter_on_materialized_model(model, name, param, debug) + + loaded_names.add(name) + if _is_quantized_tensor(param): + quantized_count += 1 + + return { + 'loaded': len(loaded_names), + 'quantized': quantized_count, + 'loaded_names': loaded_names + } + + +def _set_parameter_on_meta_model(model: torch.nn.Module, param_name: str, + param_value: torch.Tensor, debug: Optional['Debug'] = None) -> None: + """Set parameter on meta device model.""" + module, attr_name = _navigate_to_parameter(model, param_name) + new_param = _create_gguf_parameter(param_value, debug) + setattr(module, attr_name, new_param) + + +def _set_parameter_on_materialized_model(model: torch.nn.Module, param_name: str, + param_value: torch.Tensor, debug: Optional['Debug'] = None) -> None: + """Set parameter on already materialized model.""" + module, attr_name = _navigate_to_parameter(model, param_name) + + if _is_quantized_tensor(param_value): + # For quantized tensors, replace with wrapped parameter + new_param = _create_gguf_parameter(param_value, debug) + setattr(module, attr_name, new_param) + else: + # For regular tensors, just copy + existing_param = getattr(module, attr_name) + existing_param.copy_(param_value) + + +def _navigate_to_parameter(model: torch.nn.Module, param_path: str) -> Tuple[torch.nn.Module, str]: + """ + Navigate to the module containing a parameter. + + Args: + model: Root model + param_path: Dot-separated path to parameter + + Returns: + Tuple of (parent module, parameter name) + """ + path_parts = param_path.split('.') + module = model + + # Navigate to parent module + for part in path_parts[:-1]: + module = getattr(module, part) + + return module, path_parts[-1] + + +def _create_gguf_parameter(tensor: torch.Tensor, debug: Optional['Debug'] = None) -> torch.nn.Parameter: + """ + Create a parameter from a GGUF tensor, preserving quantization info. + + Args: + tensor: GGUF tensor (may be quantized) + debug: Debug instance for logging + + Returns: + Parameter with GGUF attributes and dequantize method if quantized + """ + param = torch.nn.Parameter(tensor, requires_grad=False) + + # Preserve GGUF attributes if present + if hasattr(tensor, 'tensor_type'): + param.tensor_type = tensor.tensor_type + param.tensor_shape = tensor.tensor_shape + + # Add dequantize method for runtime dequantization + param.gguf_dequantize = _create_dequantize_method(tensor, debug) + + return param + + +def _get_tensor_shape(tensor: torch.Tensor) -> torch.Size: + """Get the logical shape of a tensor (handling GGUF quantized tensors).""" + if hasattr(tensor, 'tensor_shape'): + return tensor.tensor_shape + return tensor.shape + + +def _is_quantized_tensor(tensor: torch.Tensor) -> bool: + """Check if a tensor is GGUF quantized.""" + return hasattr(tensor, 'tensor_type') and hasattr(tensor, 'tensor_shape') + + +def _report_parameter_mismatches(state: Dict[str, torch.Tensor], + model_state: Dict[str, torch.Tensor], + loaded_names: set, debug: Optional['Debug'] = None) -> None: + """Report any parameter mismatches between GGUF and model.""" + # Check for unmatched GGUF parameters + unmatched = [name for name in state if name not in model_state] + if unmatched: + debug.log(f"Warning: {len(unmatched)} parameters from GGUF not found in model", + level="WARNING", category="dit", force=True) + debug.log(f"First few unmatched: {unmatched[:5]}", level="WARNING", category="dit", force=True, indent_level=1) + + # Check for missing model parameters + missing = [name for name in model_state if name not in loaded_names] + if missing: + debug.log(f"Warning: {len(missing)} model parameters not loaded from GGUF", + level="WARNING", category="dit", force=True) + debug.log(f"First few missing: {missing[:5]}", level="WARNING", category="dit", force=True, indent_level=1) + + +def initialize_meta_buffers(model: torch.nn.Module, target_device: torch.device, debug: Optional['Debug'] = None) -> None: + """Initialize meta buffers with timing.""" + debug.start_timer("buffer_init") + initialized = initialize_meta_buffers_impl(model, target_device, debug) + if initialized > 0: + debug.log(f"Initialized {initialized} non-persistent buffers", category="success") + debug.end_timer("buffer_init", "Buffer initialization") + + +def initialize_meta_buffers_impl(model: torch.nn.Module, target_device: torch.device, debug: Optional['Debug'] = None) -> int: + """ + Initialize any buffers still on meta device after materialization. + + Non-persistent buffers aren't included in state_dict and remain on meta + device after load_state_dict. This function moves them to the target device. + + Args: + model: Model potentially containing meta device buffers + target_device: Target device for initialization (torch.device object) + debug: Debug instance for logging + + Returns: + Number of buffers initialized + """ + initialized_count = 0 + + # Simply initialize all meta device buffers to zeros on target device + for name, buffer in model.named_buffers(): + if buffer is not None and buffer.device.type == 'meta': + # Get the module that owns this buffer + module_path = name.rsplit('.', 1)[0] if '.' in name else '' + buffer_name = name.rsplit('.', 1)[1] if '.' in name else name + + # Get the actual module + if module_path: + module = model + for part in module_path.split('.'): + module = getattr(module, part) + else: + module = model + + # Create a zero tensor of the same shape on target device + # This is safe for all non-persistent buffers (caches, dummy tensors, etc.) + initialized_buffer = torch.zeros_like(buffer, device=target_device) + module.register_buffer(buffer_name, initialized_buffer, persistent=False) + initialized_count += 1 + + return initialized_count + + +def _load_standard_weights(model: torch.nn.Module, state: Dict[str, torch.Tensor], + used_meta: bool, model_type: str, model_type_lower: str, + debug: Optional['Debug'] = None) -> torch.nn.Module: + """Load standard (non-GGUF) weights into model.""" + debug.start_timer(f"{model_type_lower}_state_apply") + model.load_state_dict(state, strict=False, assign=True) + + action = "materialized" if used_meta else "applied" + debug.end_timer(f"{model_type_lower}_state_apply", f"{model_type} weights {action}") + + if used_meta: + debug.log(f"{model_type} materialized directly from meta with loaded weights", category=model_type_lower) + else: + debug.log(f"{model_type} weights applied", category=model_type_lower) + + return model + + +def _load_gguf_weights(model: torch.nn.Module, state: Dict[str, torch.Tensor], + used_meta: bool, model_type_lower: str, debug: Optional['Debug'] = None) -> torch.nn.Module: + """ + Load GGUF quantized weights into model with architecture validation. + + Args: + model: Target model + state: GGUF state dict with quantized tensors + used_meta: Whether model was initialized on meta device + model_type_lower: Lowercase model type for logging + debug: Debug instance + + Returns: + Model with GGUF weights loaded + """ + debug.log("Loading GGUF weights", category="dit") + + # Get model state dict for validation + model_state = model.state_dict() + + # Validate architecture compatibility + _validate_gguf_architecture(state, model_state, debug) + + # Load GGUF parameters + stats = apply_gguf_parameters(model, state, model_state, debug) + + # Log results + debug.log(f"GGUF loading complete: {stats['loaded']} parameters loaded", category="success") + debug.log(f"Quantized parameters: {stats['quantized']}", category="info") + + # Report any mismatches + _report_parameter_mismatches(state, model_state, stats['loaded_names'], debug) + + # Replace Linear/Conv2d layers with quantized versions for optimal precision handling + if stats['quantized'] > 0: + debug.log("Replacing layers with GGUF-optimized versions for precision handling", category="dit") + + replacements, quant_types = replace_linear_with_quantized(model, debug=debug) + + if replacements > 0: + debug.log(f"Replaced {replacements} layers with GGUF-optimized versions", category="success") + + # Show actual quantization types found and precision strategy + if quant_types: + qtypes_str = ', '.join([f"{qtype}:{count}" for qtype, count in quant_types.items()]) + debug.log( + f"GGUF precision path: {qtypes_str} โ†’ FP16 (preserve) โ†’ BF16/FP32 (compute)", + category="precision" + ) + else: + debug.log( + "GGUF precision: Dequantizing to FP16 first, then converting to compute dtype", + category="precision" + ) + else: + debug.log("Warning: No layers were replaced despite having quantized parameters", + level="WARNING", category="dit", force=True) + + return model + + +def _validate_gguf_architecture(state: Dict[str, torch.Tensor], + model_state: Dict[str, torch.Tensor], debug: Optional['Debug'] = None) -> None: + """ + Validate GGUF model architecture matches target model. + + Raises: + ValueError: If architecture mismatch is detected + """ + key_params = [ + "blocks.0.attn.proj_qkv.vid.weight", + "blocks.0.attn.proj_qkv.txt.weight", + "blocks.0.mlp.vid.proj_in.weight" + ] + + for key in key_params: + if key in state and key in model_state: + model_shape = model_state[key].shape + gguf_shape = _get_tensor_shape(state[key]) + + if model_shape != gguf_shape: + # Check if it's just a quantization difference + if hasattr(state[key], 'tensor_shape') and state[key].tensor_shape == model_shape: + continue + + raise ValueError( + f"GGUF model architecture mismatch: This GGUF model is incompatible with the current architecture.\n\n" + f"Detected mismatch:\n" + f" Parameter: {key}\n" + f" Expected shape: {model_shape}\n" + f" GGUF shape: {gguf_shape}\n\n" + f"Possible solutions:\n" + f"1. Use a GGUF model that matches the current architecture\n" + f"2. Try using a regular FP16 model instead\n" + f"3. Verify you're using the correct model variant (3B vs 7B)" + ) + + debug.log(f"Architecture check complete, no shape mismatch", category="success") + + +def _create_dequantize_method(tensor: torch.Tensor, debug: Optional['Debug'] = None) -> callable: + """ + Create a dequantization method for a GGUF tensor. + + Args: + tensor: GGUF quantized tensor with tensor_type and tensor_shape attributes + debug: Debug instance + + Returns: + Callable dequantization method + """ + def dequantize(device: Optional[torch.device] = None, + dtype: torch.dtype = torch.float16) -> torch.Tensor: + """Dequantize GGUF tensor on demand.""" + if hasattr(tensor, 'dequantize'): + return tensor.dequantize(device, dtype) + + try: + # Fallback to manual dequantization using gguf library + numpy_data = tensor.cpu().numpy() + dequantized = gguf.quants.dequantize(numpy_data, tensor.tensor_type) + result = torch.from_numpy(dequantized).to(device or tensor.device, dtype) + result.requires_grad_(False) + return result.reshape(tensor.tensor_shape) + except Exception as e: + if debug: + debug.log(f"Warning: Could not dequantize tensor: {e}", level="WARNING", category="dit", force=True) + return tensor.to(device or tensor.device, dtype) + + return dequantize \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/area_resize.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/area_resize.py new file mode 100644 index 0000000000000000000000000000000000000000..fc025da09de8bec6af951d4d8f5f253c778f1142 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/area_resize.py @@ -0,0 +1,141 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import math +import random +from typing import Union +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: + def __init__( + self, + max_area: float, + downsample_only: bool = False, + interpolation: InterpolationMode = InterpolationMode.BICUBIC, + ): + self.max_area = max_area + self.downsample_only = downsample_only + self.interpolation = interpolation + if is_mps_available(): + self.interpolation = InterpolationMode.BILINEAR + + def __call__(self, image: Union[torch.Tensor, Image.Image]): + + if isinstance(image, torch.Tensor): + height, width = image.shape[-2:] + elif isinstance(image, Image.Image): + width, height = image.size + else: + raise NotImplementedError + + scale = math.sqrt(self.max_area / (height * width)) + + # keep original height and width for small pictures. + scale = 1 if scale >= 1 and self.downsample_only else scale + + 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, + ) + + +class AreaRandomCrop: + def __init__( + self, + max_area: float, + ): + self.max_area = max_area + + def get_params(self, input_size, output_size): + """Get parameters for ``crop`` for a random crop. + + Args: + img (PIL Image): Image to be cropped. + output_size (tuple): Expected output size of the crop. + + Returns: + tuple: params (i, j, h, w) to be passed to ``crop`` for random crop. + """ + # w, h = _get_image_size(img) + h, w = input_size + th, tw = output_size + if w <= tw and h <= th: + return 0, 0, h, w + + i = random.randint(0, h - th) + j = random.randint(0, w - tw) + return i, j, th, tw + + def __call__(self, image: Union[torch.Tensor, Image.Image]): + if isinstance(image, torch.Tensor): + height, width = image.shape[-2:] + elif isinstance(image, Image.Image): + width, height = image.size + else: + raise NotImplementedError + + resized_height = math.sqrt(self.max_area / (width / height)) + resized_width = (width / height) * resized_height + + # print('>>>>>>>>>>>>>>>>>>>>>') + # print((height, width)) + # print( (resized_height, resized_width)) + + resized_height, resized_width = round(resized_height), round(resized_width) + i, j, h, w = self.get_params((height, width), (resized_height, resized_width)) + image = TVF.crop(image, i, j, h, w) + return image + +class ScaleResize: + def __init__( + self, + scale: float, + ): + self.scale = scale + + def __call__(self, image: Union[torch.Tensor, Image.Image]): + if isinstance(image, torch.Tensor): + height, width = image.shape[-2:] + interpolation_mode = InterpolationMode.BILINEAR + antialias = True if image.ndim == 4 else "warn" + elif isinstance(image, Image.Image): + width, height = image.size + interpolation_mode = InterpolationMode.LANCZOS + antialias = "warn" + else: + raise NotImplementedError + + scale = self.scale + + # keep original height and width for small pictures + + resized_height, resized_width = round(height * scale), round(width * scale) + image = TVF.resize( + image, + size=(resized_height, resized_width), + interpolation=interpolation_mode, + antialias=antialias, + ) + return image + \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/divisible_crop.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/divisible_crop.py new file mode 100644 index 0000000000000000000000000000000000000000..e09d5fb65d2a7732944b2984e244c09bad6a895b --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/divisible_crop.py @@ -0,0 +1,80 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Union +import torch +from PIL import Image +from torchvision.transforms import functional as TVF + + +class DivisibleCrop: + def __init__(self, factor): + if not isinstance(factor, tuple): + factor = (factor, factor) + + self.height_factor, self.width_factor = factor[0], factor[1] + + def __call__(self, image: Union[torch.Tensor, Image.Image]): + if isinstance(image, torch.Tensor): + height, width = image.shape[-2:] + elif isinstance(image, Image.Image): + width, height = image.size + else: + raise NotImplementedError + + cropped_height = height - (height % self.height_factor) + cropped_width = width - (width % self.width_factor) + + image = TVF.center_crop(img=image, output_size=(cropped_height, cropped_width)) + return image + + +class DivisiblePad: + """ + Pad image to make dimensions divisible by a factor. + Pads with black (0) to avoid data loss. + """ + def __init__(self, factor): + if not isinstance(factor, tuple): + factor = (factor, factor) + self.height_factor, self.width_factor = factor[0], factor[1] + + def __call__(self, image: Union[torch.Tensor, Image.Image]): + if isinstance(image, torch.Tensor): + height, width = image.shape[-2:] + elif isinstance(image, Image.Image): + width, height = image.size + else: + raise NotImplementedError + + # Calculate padding needed + pad_height = (self.height_factor - (height % self.height_factor)) % self.height_factor + pad_width = (self.width_factor - (width % self.width_factor)) % self.width_factor + + if pad_height == 0 and pad_width == 0: + return image + + # Pad symmetrically (or bottom/right) + if isinstance(image, torch.Tensor): + # Pad format: (left, right, top, bottom) + padding = (0, pad_width, 0, pad_height) + image = torch.nn.functional.pad(image, padding, mode='constant', value=0.0) + elif isinstance(image, Image.Image): + new_width = width + pad_width + new_height = height + pad_height + result = Image.new(image.mode, (new_width, new_height), (0, 0, 0)) + result.paste(image, (0, 0)) + image = result + + return image diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/na_resize.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/na_resize.py new file mode 100644 index 0000000000000000000000000000000000000000..e1111c75e1f13069af56f8b8d610473e276a97f2 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/na_resize.py @@ -0,0 +1,54 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import torch +from typing import Literal +from torchvision.transforms import CenterCrop, Compose, InterpolationMode, Resize + +from .area_resize import AreaResize +from .side_resize import SideResize +from ....optimization.memory_manager import is_mps_available + +def NaResize( + resolution: int, + mode: Literal["area", "side"], + downsample_only: bool, + max_resolution: int = 0, + interpolation: InterpolationMode = InterpolationMode.BICUBIC, +): + Interpolation = InterpolationMode.BILINEAR if is_mps_available() else interpolation + if mode == "area": + return AreaResize( + max_area=resolution**2, + downsample_only=downsample_only, + interpolation=Interpolation, + ) + if mode == "side": + return SideResize( + size=resolution, + max_size=max_resolution, + downsample_only=downsample_only, + interpolation=Interpolation, + ) + if mode == "square": + return Compose( + [ + Resize( + size=resolution, + interpolation=Interpolation, + ), + CenterCrop(resolution), + ] + ) + raise ValueError(f"Unknown resize mode: {mode}") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/side_resize.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/side_resize.py new file mode 100644 index 0000000000000000000000000000000000000000..01362aed00b20cb13aadae9a1a705139b799ce27 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/data/image/transforms/side_resize.py @@ -0,0 +1,76 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Union +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__( + self, + size: int, + max_size: int = 0, + downsample_only: bool = False, + interpolation: InterpolationMode = InterpolationMode.BICUBIC, + ): + self.size = size + self.max_size = max_size + self.downsample_only = downsample_only + self.interpolation = interpolation + if is_mps_available(): + self.interpolation = InterpolationMode.BILINEAR + + def __call__(self, image: Union[torch.Tensor, Image.Image]): + """ + Resize image with shortest edge set to size, optionally limiting longest edge. + + Args: + image (PIL Image or Tensor): Image to be scaled. + + Returns: + PIL Image or Tensor: Rescaled image with shortest edge = size, + and no edge exceeding max_size (if max_size > 0). + """ + if isinstance(image, torch.Tensor): + height, width = image.shape[-2:] + elif isinstance(image, Image.Image): + width, height = image.size + else: + raise NotImplementedError + + if self.downsample_only and min(width, height) < self.size: + size = min(width, height) + else: + size = self.size + + # Resize to shortest edge (disable antialias only for MPS tensors - not supported) + 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: + if isinstance(resized, torch.Tensor): + h, w = resized.shape[-2:] + else: + w, h = resized.size + + 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, antialias=antialias) + + return resized diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6414eec1256f05910011b0ce17bc59c3e96f520f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/__init__.py @@ -0,0 +1,39 @@ +""" +SeedVR2 ComfyUI Nodes +Central registry for all SeedVR2 nodes +""" + +from comfy_api.latest import ComfyExtension, io + +from .video_upscaler import SeedVR2VideoUpscaler +from .dit_model_loader import SeedVR2LoadDiTModel +from .vae_model_loader import SeedVR2LoadVAEModel +from .torch_compile_settings import SeedVR2TorchCompileSettings + + +class SeedVR2Extension(ComfyExtension): + """SeedVR2 ComfyUI Extension""" + + async def get_node_list(self) -> list[type[io.ComfyNode]]: + """Return list of all SeedVR2 nodes""" + return [ + SeedVR2VideoUpscaler, + SeedVR2LoadDiTModel, + SeedVR2LoadVAEModel, + SeedVR2TorchCompileSettings, + ] + + +async def comfy_entrypoint() -> ComfyExtension: + """ComfyUI V3 entry point""" + return SeedVR2Extension() + + +__all__ = [ + 'SeedVR2VideoUpscaler', + 'SeedVR2LoadDiTModel', + 'SeedVR2LoadVAEModel', + 'SeedVR2TorchCompileSettings', + 'SeedVR2Extension', + 'comfy_entrypoint', +] \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/dit_model_loader.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/dit_model_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..9d8204e8b8f66340a809548f525d483126970bbe --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/dit_model_loader.py @@ -0,0 +1,180 @@ +""" +SeedVR2 DiT Model Loader Node +Configure DiT (Diffusion Transformer) model with memory optimization +""" + +from comfy_api.latest import io +from comfy_execution.utils import get_executing_context +from typing import Dict, Any, Tuple +from ..utils.model_registry import get_available_dit_models, DEFAULT_DIT +from ..optimization.memory_manager import get_device_list + + +class SeedVR2LoadDiTModel(io.ComfyNode): + """ + Configure DiT (Diffusion Transformer) model loader with memory optimization + + Provides configuration for: + - Model selection and device placement + - BlockSwap memory optimization for limited VRAM + - Model caching between runs + - Optional torch.compile integration + + Returns: + SEEDVR2_DIT configuration dictionary for main upscaler node + """ + + @classmethod + def define_schema(cls) -> io.Schema: + devices = get_device_list() + dit_models = get_available_dit_models() + + return io.Schema( + node_id="SeedVR2LoadDiTModel", + display_name="SeedVR2 (Down)Load DiT Model", + category="SEEDVR2", + description=( + "Load and configure SeedVR2 DiT (Diffusion Transformer) model for video upscaling. " + "Supports BlockSwap memory optimization for low VRAM systems, model caching for batch processing, " + "multi-GPU offloading, and torch.compile acceleration. \n\n" + "Connect to Video Upscaler node." + ), + inputs=[ + io.Combo.Input("model", + options=dit_models, + default=DEFAULT_DIT, + tooltip=( + "DiT (Diffusion Transformer) model for video upscaling.\n" + "Models automatically download on first use.\n" + "Additional models can be added to the ComfyUI models folder." + ) + ), + io.Combo.Input("device", + options=devices, + default=devices[0], + tooltip="GPU device for DiT model inference (upscaling phase)" + ), + io.Int.Input("blocks_to_swap", + default=0, + min=0, + max=36, + step=1, + optional=True, + tooltip=( + "Number of transformer blocks to swap between devices for VRAM optimization.\n" + "โ€ข 0: Disabled (default)\n" + "โ€ข 3B model: 0-32 blocks\n" + "โ€ข 7B model: 0-36 blocks\n" + "\n" + "Requires offload_device to be set and different from device.\n" + "Not available on macOS (unified memory architecture)." + ) + ), + io.Boolean.Input("swap_io_components", + default=False, + optional=True, + tooltip=( + "Offload input/output embeddings and normalization layers to reduce VRAM.\n" + "Requires offload_device to be set and different from device.\n" + "Not available on macOS (unified memory architecture)." + ) + ), + io.Combo.Input("offload_device", + options=get_device_list(include_none=True, include_cpu=True), + default="none", + optional=True, + tooltip=( + "Device to offload DiT model when not actively processing.\n" + "โ€ข 'none': Keep model on inference device (default, fastest)\n" + "โ€ข 'cpu': Offload to system RAM (reduces VRAM usage)\n" + "โ€ข 'cuda:X': Offload to another GPU (good balance if available)\n" + "\n" + "Required for BlockSwap (blocks_to_swap or swap_io_components)." + ) + ), + io.Boolean.Input("cache_model", + default=False, + optional=True, + tooltip=( + "Keep DiT model loaded on offload_device between workflow runs.\n" + "Useful for batch processing to avoid repeated loading.\n" + "Requires offload_device to be set." + ) + ), + io.Combo.Input("attention_mode", + 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_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 and SageAttention provide speedup through optimized CUDA kernels on compatible GPUs." + ) + ), + io.Custom("TORCH_COMPILE_ARGS").Input("torch_compile_args", + optional=True, + tooltip=( + "Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node.\n" + "Provides 20-40% speedup with compatible PyTorch 2.0+ and Triton installation." + ) + ), + ], + outputs=[ + io.Custom("SEEDVR2_DIT").Output( + tooltip="DiT model configuration containing model path, device settings, BlockSwap parameters, and compilation options. Connect to Video Upscaler node." + ) + ] + ) + + @classmethod + def execute(cls, model: str, device: str, offload_device: str = "none", + cache_model: bool = False, blocks_to_swap: int = 0, + swap_io_components: bool = False, attention_mode: str = "sdpa", + torch_compile_args: Dict[str, Any] = None) -> io.NodeOutput: + """ + Create DiT model configuration for SeedVR2 main node + + Args: + model: Model filename to load + device: Target device for model execution + offload_device: Device to offload model to when not in use + 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', '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 cache_model is enabled but offload_device is not set + """ + # Validate cache_model configuration + if cache_model and offload_device == "none": + raise ValueError( + "Model caching (cache_model=True) requires offload_device to be set. " + f"Current: offload_device='{offload_device}'. " + "Please set offload_device to specify where the cached DiT model should be stored " + "(e.g., 'cpu' or another device). Set cache_model=False if you don't want to cache the model." + ) + + config = { + "model": model, + "device": device, + "offload_device": offload_device, + "cache_model": cache_model, + "blocks_to_swap": blocks_to_swap, + "swap_io_components": swap_io_components, + "attention_mode": attention_mode, + "torch_compile_args": torch_compile_args, + "node_id": get_executing_context().node_id, + } + + return io.NodeOutput(config) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/torch_compile_settings.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/torch_compile_settings.py new file mode 100644 index 0000000000000000000000000000000000000000..9958614f121b54bce5b365ed2c751b52d77b634f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/torch_compile_settings.py @@ -0,0 +1,121 @@ +""" +SeedVR2 Torch Compile Settings Node +Configure torch.compile optimization for DiT and VAE models +""" + +from comfy_api.latest import io +from typing import Dict, Any, Tuple + + +class SeedVR2TorchCompileSettings(io.ComfyNode): + """Configure torch.compile optimization for DiT and VAE models""" + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="SeedVR2TorchCompileSettings", + display_name="SeedVR2 Torch Compile Settings", + category="SEEDVR2", + description=( + "Configure SeedVR2 torch.compile optimization for 20-40% DiT speedup and 15-25% VAE speedup. " + "Trades longer first-run compilation time for faster inference.\n\n" + "Connect to DiT and/or VAE model loaders. Requires PyTorch 2.0+ and Triton for inductor backend." + ), + inputs=[ + io.Combo.Input("backend", + options=["inductor", "cudagraphs"], + default="inductor", + tooltip=( + "Compilation backend:\n" + "โ€ข inductor: Full optimization with Triton kernel generation and fusion (recommended)\n" + "โ€ข cudagraphs: Lightweight wrapper using CUDA graphs, no kernel optimization" + ) + ), + io.Combo.Input("mode", + options=["default", "reduce-overhead", "max-autotune", "max-autotune-no-cudagraphs"], + default="default", + tooltip=( + "Optimization level (compilation time vs runtime performance):\n" + "โ€ข default: Fast compilation with good speedup (recommended for development)\n" + "โ€ข reduce-overhead: Lower overhead, optimized for smaller models\n" + "โ€ข max-autotune: Slowest compilation, best runtime performance (recommended for production)\n" + "โ€ข max-autotune-no-cudagraphs: Like max-autotune but without CUDA graphs" + ) + ), + io.Boolean.Input("fullgraph", + default=False, + tooltip=( + "Compile entire model as single graph without breaks.\n" + "โ€ข False: Allow graph breaks for better compatibility (default)\n" + "โ€ข True: Enforce no breaks for maximum optimization (may fail with dynamic shapes)" + ) + ), + io.Boolean.Input("dynamic", + default=False, + tooltip=( + "Handle varying input shapes without recompilation.\n" + "โ€ข False: Specialize for exact input shapes (default)\n" + "โ€ข True: Create dynamic kernels that adapt to shape variations\n" + "\n" + "Enable when processing different resolutions or batch sizes." + ) + ), + io.Int.Input("dynamo_cache_size_limit", + default=64, + min=0, + max=1024, + step=1, + tooltip=( + "Maximum cached compiled versions per function (default: 64).\n" + "Controls how many shape variations to compile before stopping.\n" + "\n" + "โ€ข Increase: When processing many different input shapes (more memory usage)\n" + "โ€ข Decrease: When recompilation cost outweighs benefits (faster fallback to eager)" + ) + ), + io.Int.Input("dynamo_recompile_limit", + default=128, + min=0, + max=1024, + step=1, + tooltip=( + "Maximum recompilation attempts before fallback to eager mode (default: 128).\n" + "Safety limit to prevent infinite compilation loops.\n" + "\n" + "Only increase if you see 'hit config.recompile_limit' warnings and have bounded shape variations." + ) + ), + ], + outputs=[ + io.Custom("TORCH_COMPILE_ARGS").Output( + tooltip="torch.compile optimization settings including backend, mode, and Dynamo configuration. Connect to DiT and/or VAE model loader nodes." + ) + ] + ) + + @classmethod + def execute(cls, backend: str, mode: str, fullgraph: bool, dynamic: bool, + dynamo_cache_size_limit: int, dynamo_recompile_limit: int) -> io.NodeOutput: + """ + Create torch.compile configuration for model optimization + + Args: + backend: Compilation backend ("inductor" or "cudagraphs") + mode: Optimization mode ("default", "reduce-overhead", "max-autotune", etc.) + fullgraph: Whether to compile entire model as single graph + dynamic: Whether to handle varying input shapes without recompilation + dynamo_cache_size_limit: Maximum cached compiled versions per function + dynamo_recompile_limit: Maximum recompilation attempts before fallback + + Returns: + NodeOutput containing torch.compile configuration dictionary + """ + compile_args = { + "backend": backend, + "mode": mode, + "fullgraph": fullgraph, + "dynamic": dynamic, + "dynamo_cache_size_limit": dynamo_cache_size_limit, + "dynamo_recompile_limit": dynamo_recompile_limit, + } + return io.NodeOutput(compile_args) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/vae_model_loader.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/vae_model_loader.py new file mode 100644 index 0000000000000000000000000000000000000000..c6a3b0c3aa646eafa44fc3d9f73bdf34700f5f54 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/vae_model_loader.py @@ -0,0 +1,219 @@ +""" +SeedVR2 VAE Model Loader Node +Configure VAE (Variational Autoencoder) model with tiling support +""" + +from comfy_api.latest import io +from comfy_execution.utils import get_executing_context +from typing import Dict, Any, Tuple +from ..utils.model_registry import get_available_vae_models, DEFAULT_VAE +from ..optimization.memory_manager import get_device_list + + +class SeedVR2LoadVAEModel(io.ComfyNode): + """ + Configure VAE (Variational Autoencoder) model loader with tiling support + + Provides configuration for: + - Model selection and device placement + - Tiled encoding/decoding for VRAM reduction + - Tile size and overlap control + - Model caching between runs + - Optional torch.compile integration + + Returns: + SEEDVR2_VAE configuration dictionary for main upscaler node + """ + + @classmethod + def define_schema(cls) -> io.Schema: + devices = get_device_list() + vae_models = get_available_vae_models() + + return io.Schema( + node_id="SeedVR2LoadVAEModel", + display_name="SeedVR2 (Down)Load VAE Model", + category="SEEDVR2", + description=( + "Load and configure SeedVR2 VAE (Variational Autoencoder) for encoding/decoding video frames to/from latent space. " + "Supports tiled processing to handle high resolutions on limited VRAM, model caching, " + "multi-GPU offloading, and torch.compile acceleration. \n\n" + "Connect to Video Upscaler node." + ), + inputs=[ + io.Combo.Input("model", + options=vae_models, + default=DEFAULT_VAE, + tooltip=( + "VAE (Variational Autoencoder) model for encoding/decoding.\n" + "Models automatically download on first use.\n" + "Additional models can be added to the ComfyUI models folder." + ) + ), + io.Combo.Input("device", + options=devices, + default=devices[0], + tooltip="GPU device for VAE model inference (encoding/decoding phases)" + ), + io.Boolean.Input("encode_tiled", + default=False, + optional=True, + tooltip="Enable tiled encoding to reduce VRAM usage during the encoding phase" + ), + io.Int.Input("encode_tile_size", + default=1024, + min=64, + step=32, + optional=True, + tooltip=( + "Encoding tile size in pixels (default: 1024).\n" + "Applied to both height and width.\n" + "Lower values reduce VRAM usage but may increase processing time.\n" + "Only used when encode_tiled is enabled." + ) + ), + io.Int.Input("encode_tile_overlap", + default=128, + min=0, + step=32, + optional=True, + tooltip=( + "Pixel overlap between encoding tiles (default: 128).\n" + "Reduces visible seams between tiles through blending.\n" + "Higher values improve quality but slow processing.\n" + "Only used when encode_tiled is enabled." + ) + ), + io.Boolean.Input("decode_tiled", + default=False, + optional=True, + tooltip="Enable tiled decoding to reduce VRAM usage during the decoding phase" + ), + io.Int.Input("decode_tile_size", + default=1024, + min=64, + step=32, + optional=True, + tooltip=( + "Decoding tile size in pixels (default: 1024).\n" + "Applied to both height and width.\n" + "Lower values reduce VRAM usage but may increase processing time.\n" + "Only used when decode_tiled is enabled." + ) + ), + io.Int.Input("decode_tile_overlap", + default=128, + min=0, + step=32, + optional=True, + tooltip=( + "Pixel overlap between decoding tiles (default: 128).\n" + "Reduces visible seams between tiles through blending.\n" + "Higher values improve quality but slow processing.\n" + "Only used when decode_tiled is enabled." + ) + ), + io.Combo.Input("tile_debug", + options=["false", "encode", "decode"], + default="false", + optional=True, + tooltip=( + "Tile debug visualization mode:\n" + "โ€ข 'false': No visualization overlay (default)\n" + "โ€ข 'encode': Show encoding tile boundaries\n" + "โ€ข 'decode': Show decoding tile boundaries\n" + "\n" + "Only works when respective tiling is enabled." + ) + ), + io.Combo.Input("offload_device", + options=get_device_list(include_none=True, include_cpu=True), + default="none", + optional=True, + tooltip=( + "Device to offload VAE model when not actively processing.\n" + "โ€ข 'none': Keep model on inference device (default, fastest)\n" + "โ€ข 'cpu': Offload to system RAM (reduces VRAM usage)\n" + "โ€ข 'cuda:X': Offload to another GPU (good balance if available)" + ) + ), + io.Boolean.Input("cache_model", + default=False, + optional=True, + tooltip=( + "Keep VAE model loaded on offload_device between workflow runs.\n" + "Useful for batch processing to avoid repeated loading.\n" + "Requires offload_device to be set." + ) + ), + io.Custom("TORCH_COMPILE_ARGS").Input("torch_compile_args", + optional=True, + tooltip=( + "Optional torch.compile optimization settings from SeedVR2 Torch Compile Settings node.\n" + "Provides 15-25% speedup with compatible PyTorch 2.0+ and Triton installation." + ) + ), + ], + outputs=[ + io.Custom("SEEDVR2_VAE").Output( + tooltip="VAE model configuration containing model path, device settings, tiling parameters, and compilation options. Connect to Video Upscaler node." + ) + ] + ) + + @classmethod + def execute(cls, model: str, device: str, offload_device: str = "none", + cache_model: bool = False, encode_tiled: bool = False, + encode_tile_size: int = 512, encode_tile_overlap: int = 64, + decode_tiled: bool = False, decode_tile_size: int = 512, + decode_tile_overlap: int = 64, tile_debug: str = "false", + torch_compile_args: Dict[str, Any] = None + ) -> io.NodeOutput: + """ + Create VAE model configuration for SeedVR2 main node + + Args: + model: Model filename to load + device: Target device for model execution + offload_device: Device to offload model to when not in use + cache_model: Whether to keep model loaded between runs + encode_tiled: Enable tiled encoding + encode_tile_size: Tile size for encoding + encode_tile_overlap: Tile overlap for encoding + decode_tiled: Enable tiled decoding + decode_tile_size: Tile size for decoding + decode_tile_overlap: Tile overlap for decoding + tile_debug: Tile visualization mode (false/encode/decode) + torch_compile_args: Optional torch.compile configuration from settings node + + Returns: + NodeOutput containing configuration dictionary for SeedVR2 main node + + Raises: + ValueError: If cache_model is enabled but offload_device is invalid + """ + # Validate cache_model configuration + if cache_model and offload_device == "none": + raise ValueError( + "Model caching (cache_model=True) requires offload_device to be set. " + f"Current: offload_device='{offload_device}'. " + "Please set offload_device to specify where the cached VAE model should be stored " + "(e.g., 'cpu' or another device). Set cache_model=False if you don't want to cache the model." + ) + + config = { + "model": model, + "device": device, + "offload_device": offload_device, + "cache_model": cache_model, + "encode_tiled": encode_tiled, + "encode_tile_size": encode_tile_size, + "encode_tile_overlap": encode_tile_overlap, + "decode_tiled": decode_tiled, + "decode_tile_size": decode_tile_size, + "decode_tile_overlap": decode_tile_overlap, + "tile_debug": tile_debug, + "torch_compile_args": torch_compile_args, + "node_id": get_executing_context().node_id, + } + return io.NodeOutput(config) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/video_upscaler.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/video_upscaler.py new file mode 100644 index 0000000000000000000000000000000000000000..159ca2dcaf5ee77a80bd33928626d772ce4b8c75 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/interfaces/video_upscaler.py @@ -0,0 +1,580 @@ +""" +SeedVR2 Video Upscaler Node +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, __version__ +from ..utils.downloads import download_weight +from ..utils.debug import Debug +from ..core.generation_phases import ( + encode_all_batches, + upscale_all_batches, + decode_all_batches, + postprocess_all_batches +) +from ..core.generation_utils import ( + setup_generation_context, + prepare_runner, + compute_generation_info, + log_generation_start, + load_text_embeddings, + script_directory +) +from ..optimization.memory_manager import ( + cleanup_text_embeddings, + complete_cleanup, + get_device_list +) + +# Import ComfyUI progress reporting +try: + from comfy.utils import ProgressBar +except ImportError: + ProgressBar = None + + +class SeedVR2VideoUpscaler(io.ComfyNode): + """ + SeedVR2 Video Upscaler ComfyUI Node + + High-quality video upscaling using diffusion models with support for: + - Multiple model variants (3B/7B, FP16/FP8) + - Adaptive VRAM management + - Advanced dtype compatibility + - Optimized inference pipeline + - Real-time progress reporting + """ + + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id="SeedVR2VideoUpscaler", + 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. " + "Handles RGB/RGBA formats, maintains temporal consistency across frames, applies color correction, " + "and manages VRAM through intelligent tensor offloading. \n\n" + "Requires DiT and VAE model configurations." + ), + inputs=[ + io.Image.Input("image", + tooltip=( + "Input video frames as image batch.\n" + "Accepts both RGB (3-channel) and RGBA (4-channel) formats.\n" + "Output format will match input format." + ) + ), + io.Custom("SEEDVR2_DIT").Input("dit", + tooltip="DiT model configuration from SeedVR2 (Down)Load DiT Model node" + ), + io.Custom("SEEDVR2_VAE").Input("vae", + tooltip="VAE model configuration from SeedVR2 (Down)Load VAE Model node" + ), + io.Int.Input("seed", + default=42, + min=0, + max=2**32 - 1, + step=1, + tooltip=( + "Random seed for reproducible generation (default: 42).\n" + "Same seed with same inputs produces identical output." + ) + ), + io.Int.Input("resolution", + default=1080, + min=16, + max=16384, + step=2, + tooltip=( + "Target resolution for the shortest edge in pixels (default: 1080).\n" + "Automatically maintains aspect ratio of input.\n" + "Even values required for optimal processing." + ) + ), + io.Int.Input("max_resolution", + default=0, + min=0, + max=16384, + step=2, + tooltip=( + "Maximum resolution limit for any dimension (default: 0, no limit).\n" + "If any edge exceeds this after applying resolution,\n" + "both dimensions are scaled down proportionally.\n" + "Useful to prevent excessive VRAM usage on extreme aspect ratios." + ) + ), + io.Int.Input("batch_size", + default=5, + min=1, + max=16384, + step=4, + tooltip=( + "Number of frames processed together per batch (default: 5).\n" + "Must follow pattern 4n+1: 1, 5, 9, 13, 17, 21, ...\n" + "\n" + "โ€ข Higher values: Better temporal consistency and faster processing\n" + "โ€ข Lower values: Reduced VRAM usage\n" + "\n" + "Ideally match to shot length for best quality." + ) + ), + io.Boolean.Input("uniform_batch_size", + default=False, + tooltip=( + "Pad final batch to match batch_size (default: False).\n" + "Prevents temporal artifacts caused by small final batch.\n" + "Add extra compute but recommended for optimal quality." + ) + ), + io.Int.Input("temporal_overlap", + default=0, + min=0, + max=16, + step=1, + optional=True, + tooltip=( + "Overlapping frames between consecutive batches (default: 0, disabled).\n" + "Improves temporal consistency across batch boundaries through blending.\n" + "Values 1-4 work well for most content." + ) + ), + io.Int.Input("prepend_frames", + default=0, + min=0, + max=32, + step=1, + optional=True, + tooltip=( + "Number of frames to prepend (reversed from start) before processing (default: 0).\n" + "Helps reduce artifacts at video beginning.\n" + "Prepended frames are automatically removed from final output." + ) + ), + io.Combo.Input("color_correction", + options=["lab", "wavelet", "wavelet_adaptive", "hsv", "adain", "none"], + default="lab", + tooltip=( + "Corrects color shifts in upscaled output to match original input (default: lab).\n" + "The upscaling process may alter colors; this applies color grading to restore them.\n" + "\n" + "โ€ข lab: Perceptual color matching with detail preservation (recommended)\n" + "โ€ข wavelet: Frequency-based natural colors, preserves fine details\n" + "โ€ข wavelet_adaptive: Wavelet base with targeted saturation correction\n" + "โ€ข hsv: Hue-conditional saturation matching\n" + "โ€ข adain: Statistical style transfer approach\n" + "โ€ข none: No color correction applied" + ) + ), + io.Float.Input("input_noise_scale", + default=0.0, + min=0.0, + max=1.0, + step=0.001, + optional=True, + tooltip=( + "Input noise injection scale (default: 0.0, disabled).\n" + "Adds controlled variation to input images before encoding.\n" + "Range: 0.0 (no noise) to 1.0 (maximum noise).\n" + "Can help with certain types of artifacts." + ) + ), + io.Float.Input("latent_noise_scale", + default=0.0, + min=0.0, + max=1.0, + step=0.001, + optional=True, + tooltip=( + "Latent space noise injection scale (default: 0.0, disabled).\n" + "Adds controlled variation during the diffusion process.\n" + "Range: 0.0 (no noise) to 1.0 (maximum noise).\n" + "Can soften details if input_noise_scale doesn't help." + ) + ), + io.Combo.Input("offload_device", + options=get_device_list(include_none=True, include_cpu=True), + default="cpu", + optional=True, + tooltip=( + "Device for storing intermediate tensors between processing phases (default: cpu).\n" + "โ€ข 'none': Keep all tensors on inference device (fastest but highest VRAM usage)\n" + "โ€ข 'cpu': Offload to system RAM (recommended for long videos, slower transfers)\n" + "โ€ข 'cuda:X': Offload to another GPU (good balance if available, faster than CPU)" + ) + ), + io.Boolean.Input("enable_debug", + default=False, + optional=True, + tooltip=( + "Enable detailed debug logging (default: False).\n" + "Shows memory usage, timing information, and processing details.\n" + "Useful for troubleshooting errors and performance issues." + ) + ), + ], + outputs=[ + io.Image.Output( + tooltip="Upscaled video frames with color correction applied. Format (RGB/RGBA) matches input. Range [0, 1] normalized for ComfyUI compatibility." + ) + ] + ) + + @classmethod + def execute(cls, image: torch.Tensor, dit: Dict[str, Any], vae: Dict[str, Any], + seed: int, resolution: int = 1080, max_resolution: int = 0, batch_size: int = 5, + uniform_batch_size: bool = False, temporal_overlap: int = 0, prepend_frames: int = 0, + color_correction: str = "wavelet", input_noise_scale: float = 0.0, + latent_noise_scale: float = 0.0, offload_device: str = "none", + enable_debug: bool = False) -> io.NodeOutput: + """ + Execute SeedVR2 video upscaling with progress reporting + + Main entry point for ComfyUI node execution. + Automatically detects and preserves input format (RGB or RGBA). Handles model downloads, + configuration unpacking, and delegates to upscaling pipeline. + + Args: + image: Input video frames as tensor (N, H, W, C) in [0, 1] range + dit: DiT model configuration from SeedVR2LoadDiTModel node + vae: VAE model configuration from SeedVR2LoadVAEModel node + seed: Random seed for reproducible generation + resolution: Target resolution for shortest edge (maintains aspect ratio) + max_resolution: Maximum resolution for any edge (0 = no limit) + batch_size: Frames per batch (minimum 5 for temporal consistency) + uniform_batch_size: Whether to pad final batch to match batch_size + temporal_overlap: Overlapping frames between batches (0-16) + prepend_frames: Frames to prepend (0-32) to reduce initial artifacts. + color_correction: Color correction method + input_noise_scale: Input noise injection scale [0.0-1.0] + latent_noise_scale: Latent noise injection scale [0.0-1.0] + offload_device: Device to offload intermediate tensors + enable_debug: Enable detailed logging and memory tracking + + Returns: + NodeOutput containing upscaled video tensor (N, H', W', C) in [0, 1] range + + Raises: + ValueError: If model files cannot be downloaded or configuration is invalid + RuntimeError: If generation fails + """ + # Initialize debug (stateless - stored in local variable) + debug = Debug(enabled=enable_debug) + + # Track execution state in local variables (not instance) + runner = None + ctx = None + pbar = None + + # Define progress callback as local closure + def progress_callback(current_step: int, total_steps: int, + current_frames: int, phase_name: str) -> None: + """ + Update progress bar based on pipeline phase + + Args: + current_step: Current step within phase + total_steps: Total steps in phase + current_frames: Number of frames being processed + phase_name: Name of current phase + """ + if pbar is None: + return + + phase_weights = { + "Phase 1: Encoding": 0.2, + "Phase 2: Upscaling": 0.25, + "Phase 3: Decoding": 0.5, + "Phase 4: Post-processing": 0.05 + } + phase_offset = { + "Phase 1: Encoding": 0.0, + "Phase 2: Upscaling": 0.2, + "Phase 3: Decoding": 0.45, + "Phase 4: Post-processing": 0.95 + } + + # Extract phase key from phase_name + phase_key = phase_name.split(" (")[0] if " (" in phase_name else phase_name + + # Get weight and offset + weight = phase_weights.get(phase_key, 1.0) + offset = phase_offset.get(phase_key, 0.0) + + # Calculate weighted progress + phase_progress = (current_step / total_steps) if total_steps > 0 else 0 + overall_progress = offset + (phase_progress * weight) + + # Update the progress bar with the overall progress + progress_value = int(overall_progress * 100) + pbar.update_absolute(progress_value, 100) + + # Define cleanup as local function + def cleanup(dit_cache: bool = False, vae_cache: bool = False) -> None: + """Cleanup resources after upscaling""" + nonlocal runner, ctx + + # Use complete_cleanup for all cleanup operations + if runner: + complete_cleanup(runner=runner, debug=debug, + dit_cache=dit_cache, vae_cache=vae_cache) + + # Delete runner only if neither model is cached + if not (dit_cache or vae_cache): + runner = None + + # Clean up context text embeddings if they exist + if ctx: + cleanup_text_embeddings(ctx, debug) + ctx = None + + # Extract configuration from dict inputs + dit_model = dit["model"] + vae_model = vae["model"] + dit_device = torch.device(dit["device"]) + vae_device = torch.device(vae["device"]) + dit_id = dit["node_id"] + vae_id = vae["node_id"] + + # OPTIONAL inputs - use .get() with defaults + dit_cache = dit.get("cache_model", False) + attention_mode = dit.get("attention_mode", "sdpa") + vae_cache = vae.get("cache_model", False) + + # BlockSwap configuration - construct from individual values + blocks_to_swap = dit.get("blocks_to_swap", 0) + swap_io_components = dit.get("swap_io_components", False) + dit_offload_str = dit.get("offload_device", "none") + + block_swap_config = None + if blocks_to_swap > 0 or swap_io_components: + block_swap_config = { + "blocks_to_swap": blocks_to_swap, + "swap_io_components": swap_io_components, + } + if dit_offload_str != "none": + block_swap_config["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") + dit_offload_device = torch.device(dit_offload_str) if dit_offload_str != "none" else None + vae_offload_device = torch.device(vae_offload_str) if vae_offload_str != "none" else None + tensor_offload_device = torch.device(offload_device) if offload_device != "none" else None + + # VAE tiling configuration + encode_tiled = vae.get("encode_tiled", False) + encode_tile_size = vae.get("encode_tile_size", 512) + encode_tile_overlap = vae.get("encode_tile_overlap", 64) + decode_tiled = vae.get("decode_tiled", False) + decode_tile_size = vae.get("decode_tile_size", 512) + decode_tile_overlap = vae.get("decode_tile_overlap", 64) + tile_debug = vae.get("tile_debug", False) + + # TorchCompile args (optional connection, can be None) + dit_torch_compile_args = dit.get("torch_compile_args") + vae_torch_compile_args = vae.get("torch_compile_args") + + # Print header + debug.print_header() + + debug.start_timer("total_execution", force=True) + + debug.log("โ”โ”โ”โ”โ”โ”โ”โ”โ” Model Preparation โ”โ”โ”โ”โ”โ”โ”โ”โ”", category="none") + + # Initial memory state + debug.log_memory_state("Before model preparation", show_tensors=False, detailed_tensors=False) + debug.start_timer("model_preparation") + + # Check if download succeeded + debug.log("Checking and downloading models if needed...", category="download") + if not download_weight(dit_model=dit_model, vae_model=vae_model, debug=debug): + raise RuntimeError( + f"Failed to download required model files. " + f"DiT model: {dit_model}, VAE model: {vae_model}. " + "Please check the console output above for specific file failures and manual download instructions." + ) + + try: + # Initialize ComfyUI progress bar if available + if ProgressBar is not None: + pbar = ProgressBar(100) + + # Setup generation context with device configuration + ctx = setup_generation_context( + dit_device=dit_device, + vae_device=vae_device, + dit_offload_device=dit_offload_device, + vae_offload_device=vae_offload_device, + tensor_offload_device=tensor_offload_device, + debug=debug + ) + + # Prepare runner with model state management and global cache + runner, cache_context = prepare_runner( + dit_model=dit_model, + vae_model=vae_model, + model_dir=get_base_cache_dir(), + debug=debug, + ctx=ctx, + dit_cache=dit_cache, + vae_cache=vae_cache, + dit_id=dit_id, + vae_id=vae_id, + block_swap_config=block_swap_config, + encode_tiled=encode_tiled, + encode_tile_size=(encode_tile_size, encode_tile_size), + encode_tile_overlap=(encode_tile_overlap, encode_tile_overlap), + decode_tiled=decode_tiled, + decode_tile_size=(decode_tile_size, decode_tile_size), + decode_tile_overlap=(decode_tile_overlap, decode_tile_overlap), + tile_debug=tile_debug, + attention_mode=attention_mode, + torch_compile_args_dit=dit_torch_compile_args, + torch_compile_args_vae=vae_torch_compile_args + ) + + # 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) + + # Compute generation info and log start (handles prepending internally) + image, gen_info = compute_generation_info( + ctx=ctx, + images=image, + resolution=resolution, + max_resolution=max_resolution, + batch_size=batch_size, + uniform_batch_size=uniform_batch_size, + seed=seed, + prepend_frames=prepend_frames, + temporal_overlap=temporal_overlap, + debug=debug + ) + + # Log generation start in consistent format + log_generation_start(gen_info, debug) + + debug.start_timer("generation") + + # Phase 1: Encode + ctx = encode_all_batches( + runner, + ctx=ctx, + images=image, + debug=debug, + batch_size=batch_size, + uniform_batch_size=uniform_batch_size, + seed=seed, + progress_callback=progress_callback, + temporal_overlap=temporal_overlap, + resolution=resolution, + max_resolution=max_resolution, + input_noise_scale=input_noise_scale, + color_correction=color_correction + ) + + # Phase 2: Upscale + ctx = upscale_all_batches( + runner, + ctx=ctx, + debug=debug, + progress_callback=progress_callback, + seed=seed, + latent_noise_scale=latent_noise_scale, + cache_model=dit_cache + ) + + # Phase 3: Decode + ctx = decode_all_batches( + runner, + ctx=ctx, + debug=debug, + progress_callback=progress_callback, + cache_model=vae_cache + ) + + # Phase 4: Post-processing + ctx = postprocess_all_batches( + ctx=ctx, + debug=debug, + progress_callback=progress_callback, + color_correction=color_correction, + prepend_frames=prepend_frames, + temporal_overlap=temporal_overlap, + batch_size=batch_size + ) + + 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: + 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("Upscaling completed successfully!", category="success", force=True) + debug.end_timer("generation", "Video generation") + + # Final cleanup + debug.start_timer("final_cleanup") + cleanup(dit_cache=dit_cache, vae_cache=vae_cache) + debug.end_timer("final_cleanup", "Final cleanup") + + debug.log_memory_state("After all phases complete", show_tensors=False, detailed_tensors=False) + + # Final peak vram summary + debug.log_peak_memory_summary() + + # Final timing summary + debug.log("", category="none") + debug.log("โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€โ”€", category="none") + child_times = { + "Model preparation": debug.timer_durations.get("model_preparation", 0), + "Video generation": debug.timer_durations.get("generation", 0), + "Final cleanup": debug.timer_durations.get("final_cleanup", 0) + } + if "phase1_encoding" in debug.timer_durations: + child_times[" Phase 1: VAE encoding"] = debug.timer_durations.get("phase1_encoding", 0) + if "phase2_upscaling" in debug.timer_durations: + child_times[" Phase 2: DiT upscaling"] = debug.timer_durations.get("phase2_upscaling", 0) + if "phase3_decoding" in debug.timer_durations: + child_times[" Phase 3: VAE decoding"] = debug.timer_durations.get("phase3_decoding", 0) + if "phase4_postprocessing" in debug.timer_durations: + child_times[" Phase 4: Post-processing"] = debug.timer_durations.get("phase4_postprocessing", 0) + + total_execution_time = debug.end_timer("total_execution", "Total execution", show_breakdown=True, custom_children=child_times) + + if total_execution_time > 0: + fps = gen_info['total_frames'] / total_execution_time + debug.log(f"Average FPS: {fps:.2f} frames/sec", category="timing", force=True) + + # Print footer + debug.print_footer() + + debug.clear_history() + pbar = None + ctx = None + + # V3-compatible return with optional UI preview + return io.NodeOutput(sample) + + except Exception as e: + cleanup(dit_cache=dit_cache, vae_cache=vae_cache) + raise e \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/attention.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..03e3c634d7de24ee90243f9e1cd99ac92d786714 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/attention.py @@ -0,0 +1,148 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import torch +import torch.nn.functional as F + +# Import flash/sage attn with automatic fallback from compatibility layer +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 + + +def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False): + """ + A PyTorch-based implementation of variable-length attention to replace flash_attn_varlen_func. + It processes each sequence in the batch individually. + + NOTE: max_seqlen_q and max_seqlen_k are accepted for API compatibility but not used. + PyTorch's scaled_dot_product_attention automatically handles variable sequence lengths. + + COMPILE OPTIMIZATION: Uses torch.tensor_split to avoid .item() graph breaks + """ + # Split q, k, v using cumulative sequence lengths + # NOTE: torch.tensor_split requires int64 dtype and CPU device (PyTorch requirements) + q_splits = list(torch.tensor_split(q, cu_seqlens_q[1:-1].long().cpu(), dim=0)) + k_splits = list(torch.tensor_split(k, cu_seqlens_k[1:-1].long().cpu(), dim=0)) + v_splits = list(torch.tensor_split(v, cu_seqlens_k[1:-1].long().cpu(), dim=0)) + + # Process each sequence + output_splits = [] + for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits): + # Reshape for torch's scaled_dot_product_attention which expects (batch, heads, seq, dim). + # Here, we treat each sequence as a batch of 1. + q_i = q_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_q, head_dim) + k_i = k_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim) + v_i = v_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim) + + # Use PyTorch's built-in scaled dot-product attention. + output_i = F.scaled_dot_product_attention( + q_i, k_i, v_i, + dropout_p=dropout_p if not deterministic else 0.0, + is_causal=causal + ) + + # Reshape the output back to the original format (seq_len, heads, head_dim) + output_i = output_i.squeeze(0).permute(1, 0, 2) + output_splits.append(output_i) + + # Concatenate all outputs + return torch.cat(output_splits, dim=0) + + +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" + q = kwargs.get("query") or args[0] + k = kwargs.get("key") or args[1] + b, h, sq, d = q.shape + b, h, sk, d = k.shape + return b * h * (4 * d * (sq / 1e6) * (sk / 1e6)) + + def forward(self, *args, **kwargs): + return F.scaled_dot_product_attention(*args, **kwargs) + + +class FlashAttentionVarlen(nn.Module): + """ + Variable-length attention with configurable backend. + + Supported backends: + - sdpa: PyTorch SDPA (fully compilable, always available) + - flash_attn_2: Flash Attention 2 (Ampere+) + - 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): + """ + Initialize with specified attention backend. + + Args: + 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__() + self.attention_mode = attention_mode + self.compute_dtype = compute_dtype + + def tflops(self, args, kwargs, output) -> float: + cu_seqlens_q = kwargs["cu_seqlens_q"] + cu_seqlens_k = kwargs["cu_seqlens_k"] + _, h, d = output.shape + seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) / 1e6 + seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) / 1e6 + return h * (4 * d * (seqlens_q * seqlens_k).sum()) + + def forward(self, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs): + kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled() + + # Convert to pipeline compute_dtype if configured (handles FP8 โ†’ fp16/bf16) + if self.compute_dtype is not None and q.dtype != self.compute_dtype: + q = q.to(self.compute_dtype) + k = k.to(self.compute_dtype) + v = v.to(self.compute_dtype) + + if self.attention_mode == 'flash_attn_3': + 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( + q, k, v, cu_seqlens_q, cu_seqlens_k, + max_seqlen_q, max_seqlen_k, **kwargs + ) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/embedding.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..e972244f5767c9f34e5e77bb180ae720ce88b89c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/embedding.py @@ -0,0 +1,62 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Optional, Union +import torch +from diffusers.models.embeddings import get_timestep_embedding +from torch import nn + + +def emb_add(emb1: torch.Tensor, emb2: Optional[torch.Tensor]): + return emb1 if emb2 is None else emb1 + emb2 + + +class TimeEmbedding(nn.Module): + def __init__( + self, + sinusoidal_dim: int, + hidden_dim: int, + output_dim: int, + ): + super().__init__() + self.sinusoidal_dim = sinusoidal_dim + self.proj_in = nn.Linear(sinusoidal_dim, hidden_dim) + self.proj_hid = nn.Linear(hidden_dim, hidden_dim) + self.proj_out = nn.Linear(hidden_dim, output_dim) + self.act = nn.SiLU() + + def forward( + self, + timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], + device: torch.device, + dtype: torch.dtype, + ) -> torch.FloatTensor: + if not torch.is_tensor(timestep): + timestep = torch.tensor([timestep], device=device, dtype=dtype) + if timestep.ndim == 0: + timestep = timestep[None] + + emb = get_timestep_embedding( + timesteps=timestep, + embedding_dim=self.sinusoidal_dim, + flip_sin_to_cos=False, + downscale_freq_shift=0, + ) + emb = emb.to(dtype) + emb = self.proj_in(emb) + emb = self.act(emb) + emb = self.proj_hid(emb) + emb = self.act(emb) + emb = self.proj_out(emb) + return emb diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mlp.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..2d05cb021f3e3c6ac05c0e7ae1aa8a6d29475b87 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mlp.py @@ -0,0 +1,62 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Optional +import torch +import torch.nn.functional as F +from torch import nn + + +def get_mlp(mlp_type: Optional[str] = "normal"): + if mlp_type == "normal": + return MLP + elif mlp_type == "swiglu": + return SwiGLUMLP + + +class MLP(nn.Module): + def __init__( + self, + dim: int, + expand_ratio: int, + ): + super().__init__() + self.proj_in = nn.Linear(dim, dim * expand_ratio) + self.act = nn.GELU("tanh") + self.proj_out = nn.Linear(dim * expand_ratio, dim) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + x = self.proj_in(x) + x = self.act(x) + x = self.proj_out(x) + return x + + +class SwiGLUMLP(nn.Module): + def __init__( + self, + dim: int, + expand_ratio: int, + multiple_of: int = 256, + ): + super().__init__() + hidden_dim = int(2 * dim * expand_ratio / 3) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + self.proj_in_gate = nn.Linear(dim, hidden_dim, bias=False) + self.proj_out = nn.Linear(hidden_dim, dim, bias=False) + self.proj_in = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + x = self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x)) + return x diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mm.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mm.py new file mode 100644 index 0000000000000000000000000000000000000000..344f89a8fa22b9a5473b8d25f208085a630f0c85 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/mm.py @@ -0,0 +1,74 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Tuple +import torch +from torch import nn + + +@dataclass +class MMArg: + vid: Any + txt: Any + + +def get_args(key: str, args: List[Any]) -> List[Any]: + return [getattr(v, key) if isinstance(v, MMArg) else v for v in args] + + +def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]: + return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()} + + +class MMModule(nn.Module): + def __init__( + self, + module: Callable[..., nn.Module], + *args, + shared_weights: bool = False, + vid_only: bool = False, + **kwargs, + ): + super().__init__() + self.shared_weights = shared_weights + self.vid_only = vid_only + if self.shared_weights: + assert get_args("vid", args) == get_args("txt", args) + assert get_kwargs("vid", kwargs) == get_kwargs("txt", kwargs) + self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs)) + else: + self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs)) + self.txt = ( + module(*get_args("txt", args), **get_kwargs("txt", kwargs)) + if not vid_only + else None + ) + + def forward( + self, + vid: torch.FloatTensor, + txt: torch.FloatTensor, + *args, + **kwargs, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_module = self.vid if not self.shared_weights else self.all + vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs)) + if not self.vid_only: + txt_module = self.txt if not self.shared_weights else self.all + txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs)) + return vid, txt diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/modulation.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/modulation.py new file mode 100644 index 0000000000000000000000000000000000000000..854ae09a26dda7270945db80c37db8b35b133a97 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/modulation.py @@ -0,0 +1,121 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Callable, List, Optional +import torch +from einops import rearrange +from torch import nn + +from ...common.cache import Cache +from ...common.distributed.ops import slice_inputs + +# (dim: int, emb_dim: int) +ada_layer_type = Callable[[int, int], nn.Module] + + +def get_ada_layer(ada_layer: str) -> ada_layer_type: + if ada_layer == "single": + return AdaSingle + raise NotImplementedError(f"{ada_layer} is not supported") + + +def expand_dims(x: torch.Tensor, dim: int, ndim: int): + """ + Expand tensor "x" to "ndim" by adding empty dims at "dim". + Example: x is (b d), target ndim is 5, add dim at 1, return (b 1 1 1 d). + """ + shape = x.shape + shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:] + return x.reshape(shape) + + +class AdaSingle(nn.Module): + def __init__( + self, + dim: int, + emb_dim: int, + layers: List[str], + modes: List[str] = ["in", "out"], + ): + assert emb_dim == 6 * dim, "AdaSingle requires emb_dim == 6 * dim" + super().__init__() + self.dim = dim + self.emb_dim = emb_dim + self.layers = layers + for l in layers: + if "in" in modes: + self.register_parameter(f"{l}_shift", nn.Parameter(torch.randn(dim) / dim**0.5)) + self.register_parameter( + f"{l}_scale", nn.Parameter(torch.randn(dim) / dim**0.5 + 1) + ) + if "out" in modes: + self.register_parameter(f"{l}_gate", nn.Parameter(torch.randn(dim) / dim**0.5)) + + def forward( + self, + hid: torch.FloatTensor, # b ... c + emb: torch.FloatTensor, # b d + layer: str, + mode: str, + cache: Cache = Cache(disable=True), + branch_tag: str = "", + hid_len: Optional[torch.LongTensor] = None, # b + ) -> torch.FloatTensor: + idx = self.layers.index(layer) + emb = rearrange(emb, "b (d l g) -> b d l g", l=len(self.layers), g=3)[..., idx, :] + emb = expand_dims(emb, 1, hid.ndim + 1) + + if hid_len is not None: + emb = cache( + f"emb_repeat_{idx}_{branch_tag}", + lambda: slice_inputs( + torch.repeat_interleave(emb, hid_len, dim=0), + dim=0, + ), + ) + + shiftA, scaleA, gateA = emb.unbind(-1) + shiftB, scaleB, gateB = ( + getattr(self, f"{layer}_shift", None), + getattr(self, f"{layer}_scale", None), + getattr(self, f"{layer}_gate", None), + ) + + # Handle potential FP8 parameters - convert to input computation dtype + if hasattr(torch, 'float8_e4m3fn'): + fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2) + # Use input tensor's dtype as target (respects pipeline precision) + target_dtype = hid.dtype + + # Convert FP8 parameters to match input dtype for arithmetic operations + if shiftB is not None and shiftB.dtype in fp8_types: + shiftB = shiftB.to(target_dtype) + if scaleB is not None and scaleB.dtype in fp8_types: + scaleB = scaleB.to(target_dtype) + if gateB is not None and gateB.dtype in fp8_types: + gateB = gateB.to(target_dtype) + + if mode == "in": + return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB) + if mode == "out": + if gateB is not None: + return hid.mul_(gateA + gateB) + else: + # If no gate parameter, just use the embedding gate + return hid.mul_(gateA) + + raise NotImplementedError + + def extra_repr(self) -> str: + return f"dim={self.dim}, emb_dim={self.emb_dim}, layers={self.layers}" \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/na.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/na.py new file mode 100644 index 0000000000000000000000000000000000000000..26f17951a37015a6d33b39c3a51d37af2081b793 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/na.py @@ -0,0 +1,641 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Native Resolution Transformer (NaDiT) tensor manipulation utilities. + +TORCH.COMPILE OPTIMIZED VERSION +================================ +This module has been optimized for torch.compile compatibility by eliminating +data-dependent operations that cause graph breaks: + +Key Changes from Original: +- Replaced all .tolist() calls with pure tensor operations +- Minimized .item() calls (only used where required by einops API) +- Replaced list comprehensions with tensor-based splitting +- Added _tensor_split helper for compile-friendly splitting +- Proper device management to ensure tensors stay on correct devices +""" + +from itertools import chain +from typing import Callable, Dict, List, Tuple +import einops +import torch + + +def _tensor_split(tensor: torch.Tensor, lengths: torch.LongTensor, dim: int = 0) -> List[torch.Tensor]: + """ + Optimized compile-friendly split using torch.tensor_split. + + Uses PyTorch's native C++ implementation for fast eager mode execution + while remaining fully compatible with torch.compile symbolic tracing. + + Args: + tensor: Input tensor to split + lengths: Tensor of split lengths (1D) + dim: Dimension along which to split (default: 0) + + Returns: + List of split tensors + """ + if lengths.numel() == 0: + return [] + + if lengths.numel() == 1: + return [tensor] + + # Calculate split indices: torch.tensor_split splits BEFORE each index + # So we need cumsum[:-1] to get the split points + # NOTE: torch.tensor_split requires indices on CPU (PyTorch requirement) + split_indices = lengths.cumsum(0)[:-1].cpu() + + # Use torch.tensor_split - native C++ implementation + # Compilable: accepts tensor indices (symbolic shapes work) + # Fast: uses optimized CUDA/CPU kernels instead of Python loops + return list(torch.tensor_split(tensor, split_indices, dim=dim)) + + +def flatten( + hid: List[torch.FloatTensor], # List of (*** c) +) -> Tuple[ + torch.FloatTensor, # (L c) + torch.LongTensor, # (b n) +]: + """ + Flatten a list of tensors into a single tensor and track their shapes. + + Converts a list of tensors with potentially different spatial shapes but + same feature dimension into a flattened tensor and shape metadata. + + Args: + hid: List of tensors, each with shape (d1, d2, ..., dn, c) + + Returns: + Tuple of: + - Flattened tensor of shape (L, c) where L = sum of all spatial dimensions + - Shape tensor of shape (b, n) tracking original spatial dimensions + + COMPILE OPTIMIZATION: Uses tensor operations on correct device from input + """ + assert len(hid) > 0 + device = hid[0].device + + # Stack shape metadata - ensure tensors are created on correct device + shapes = [] + for x in hid: + shape_tensor = torch.tensor(x.shape[:-1], dtype=torch.long, device=device) + shapes.append(shape_tensor) + shape = torch.stack(shapes) + + # Flatten and concatenate + hid = torch.cat([x.flatten(0, -2) for x in hid]) + return hid, shape + + +def unflatten( + hid: torch.FloatTensor, # (L c) or (L ... c) + hid_shape: torch.LongTensor, # (b n) +) -> List[torch.Tensor]: # List of (*** c) or (*** ... c) + """ + Unflatten a tensor back to a list using shape metadata. + + Inverse operation of flatten(), reconstructing original tensor shapes. + + Args: + hid: Flattened tensor of shape (L, c) or (L, ..., c) + hid_shape: Shape metadata tensor of shape (b, n) + + Returns: + List of unflattened tensors with original spatial dimensions + + COMPILE OPTIMIZATION: + - Uses optimized _tensor_split with torch.tensor_split (fast) + - .cpu().numpy() conversion needed for torch.compile compatibility + (unflatten() requires concrete Python ints, not symbolic shapes) + """ + hid_len = hid_shape.prod(-1) + + # Use optimized tensor splitting (major performance improvement) + hid_list = _tensor_split(hid, hid_len, dim=0) + + # Unflatten each piece + # NOTE: .cpu().numpy() is required for torch.compile compatibility + # .tolist() would fail with symbolic shapes during compilation + result = [] + for i, x in enumerate(hid_list): + shape = hid_shape[i] + # Must use .cpu().numpy() for compilation compatibility + # Shape tensors are small, so CPU transfer overhead is minimal + target_shape = list(shape.cpu().numpy()) + result.append(x.unflatten(0, target_shape)) + + return result + + +def concat( + vid: torch.FloatTensor, # (VL ... c) + txt: torch.FloatTensor, # (TL ... c) + vid_len: torch.LongTensor, # (b) + txt_len: torch.LongTensor, # (b) +) -> torch.FloatTensor: # (L ... c) + """ + Interleave video and text tensors batch-wise. + + Splits video and text tensors by batch lengths, then interleaves them: + [vid_0, txt_0, vid_1, txt_1, ..., vid_b, txt_b] + + Args: + vid: Video features tensor (VL, c) + txt: Text features tensor (TL, c) + vid_len: Length of each video sequence (b,) + txt_len: Length of each text sequence (b,) + + Returns: + Interleaved tensor (L, c) where L = sum(vid_len) + sum(txt_len) + + COMPILE OPTIMIZATION: Uses _tensor_split for compile-friendly splitting + """ + # Use tensor-based splitting + vid_splits = _tensor_split(vid, vid_len, dim=0) + txt_splits = _tensor_split(txt, txt_len, dim=0) + + # Interleave + interleaved = [] + for v, t in zip(vid_splits, txt_splits): + interleaved.extend([v, t]) + + return torch.cat(interleaved) + + +def concat_idx( + vid_len: torch.LongTensor, # (b) + txt_len: torch.LongTensor, # (b) +) -> Tuple[ + Callable, + Callable, +]: + """ + Create index-based concatenation and un-concatenation functions. + + Pre-computes indices for efficient interleaving and de-interleaving operations. + Returns callable functions that can be reused multiple times. + + Args: + vid_len: Video sequence lengths (b,) + txt_len: Text sequence lengths (b,) + + Returns: + Tuple of (concat_fn, unconcat_fn): + - concat_fn: Lambda that interleaves vid and txt tensors + - unconcat_fn: Lambda that separates interleaved tensor back to vid and txt + + COMPILE OPTIMIZATION: Pre-computes all indices using tensor operations + """ + device = vid_len.device + vid_sum = vid_len.sum() + txt_sum = txt_len.sum() + + vid_idx = torch.arange(vid_sum, device=device) + txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device) + + # Build interleaving indices using compile-friendly _tensor_split + batch_size = len(vid_len) + vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0) + txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0) + + # Create interleaved target indices + tgt_idx_list = [] + for i in range(batch_size): + tgt_idx_list.append(vid_idx_splits[i]) + tgt_idx_list.append(txt_idx_splits[i]) + + tgt_idx = torch.cat(tgt_idx_list) + src_idx = torch.argsort(tgt_idx) + vid_idx_len = len(vid_idx) + + return ( + lambda vid, txt: torch.index_select(torch.cat([vid, txt]), 0, tgt_idx), + lambda all: torch.index_select(all, 0, src_idx).split([vid_idx_len, len(txt_idx)]), + ) + + +def unconcat( + all: torch.FloatTensor, # (L ... c) + vid_len: torch.LongTensor, # (b) + txt_len: torch.LongTensor, # (b) +) -> Tuple[ + torch.FloatTensor, # (VL ... c) + torch.FloatTensor, # (TL ... c) +]: + """ + De-interleave concatenated video and text tensors. + + Inverse of concat(). Separates an interleaved tensor back into video and text. + + Args: + all: Interleaved tensor (L, c) + vid_len: Video sequence lengths (b,) + txt_len: Text sequence lengths (b,) + + Returns: + Tuple of (vid, txt) tensors + + COMPILE OPTIMIZATION: Uses tensor operations to build interleave pattern + """ + batch_size = len(vid_len) + + # Create interleaved lengths: [vid_0, txt_0, vid_1, txt_1, ...] + interleave_len = torch.stack([vid_len, txt_len], dim=1).flatten() + + # Split using compile-friendly operation + all_splits = _tensor_split(all, interleave_len, dim=0) + + # Separate even (vid) and odd (txt) indices + vid_parts = [all_splits[i] for i in range(0, len(all_splits), 2)] + txt_parts = [all_splits[i] for i in range(1, len(all_splits), 2)] + + vid = torch.cat(vid_parts) + txt = torch.cat(txt_parts) + return vid, txt + + +def repeat_concat( + vid: torch.FloatTensor, # (VL ... c) + txt: torch.FloatTensor, # (TL ... c) + vid_len: torch.LongTensor, # (n*b) + txt_len: torch.LongTensor, # (b) + txt_repeat: torch.LongTensor, # (n) or (b) +) -> torch.FloatTensor: # (L ... c) + """ + Concatenate video and text with text repetition for window attention. + + For windowed attention, text features are repeated and interleaved with + multiple video windows: [vid_0, txt_0, vid_1, txt_0, vid_2, txt_0, ...] + + Args: + vid: Video features (VL, c) + txt: Text features (TL, c) + vid_len: Video window lengths (n*b,) where n=num_windows + txt_len: Text sequence lengths (b,) + txt_repeat: Number of times to repeat text (n,) or (b,) + + Returns: + Interleaved tensor with repeated text + + COMPILE OPTIMIZATION: Uses _tensor_split and tensor-based repetition + """ + # Split using compile-friendly operations + vid_splits = _tensor_split(vid, vid_len, dim=0) + txt_splits = _tensor_split(txt, txt_len, dim=0) + + # Handle txt_repeat shape flexibility + if txt_repeat.numel() == len(txt_splits): + repeat_counts = txt_repeat + else: + repeat_counts = txt_repeat.repeat(len(txt_splits)) + + # Interleave with repetition + result = [] + for i, v in enumerate(vid_splits): + result.append(v) + # Get corresponding text sample (cyclic) + batch_idx = i % len(txt_splits) if len(txt_splits) > 0 else 0 + if batch_idx < len(txt_splits): + result.append(txt_splits[batch_idx]) + + return torch.cat(result) + + +def repeat_concat_idx( + vid_len: torch.LongTensor, # (n*b) + txt_len: torch.LongTensor, # (b) + txt_repeat: torch.LongTensor, # (n) or scalar +) -> Tuple[ + Callable, + Callable, +]: + """ + Create index-based repeat-concatenation and un-concatenation with coalescing. + + Similar to concat_idx but handles text repetition for window attention. + The unconcat function coalesces (averages) repeated text features. + + Args: + vid_len: Video window lengths (n*b,) + txt_len: Text sequence lengths (b,) + txt_repeat: Repetition count (scalar or tensor) + + Returns: + Tuple of (concat_fn, unconcat_coalesce_fn): + - concat_fn: Interleaves with text repetition + - unconcat_coalesce_fn: Separates and averages repeated text + + Example: + Input: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10], repeat=3 + Concat: [0,1,2,9,10, 3,4,5,9,10, 6,7,8,9,10] + Unconcat: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10] (averaged) + + COMPILE OPTIMIZATION: + - Eliminates .tolist() calls that caused graph breaks + - Uses pure tensor operations for index building + - Minimizes data-dependent branching + """ + device = vid_len.device + vid_sum = vid_len.sum() + txt_sum = txt_len.sum() + + # Create base indices + vid_idx = torch.arange(vid_sum, device=device) + txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device) + + # Normalize txt_repeat to tensor + if isinstance(txt_repeat, (int, float)): + txt_repeat = torch.tensor([txt_repeat], dtype=torch.long, device=device) + elif txt_repeat.dim() == 0: + txt_repeat = txt_repeat.unsqueeze(0) + + # Calculate repeat pattern - keep as tensor to avoid graph breaks + batch_size = len(txt_len) + if txt_repeat.numel() == 1: + num_repeats_tensor = txt_repeat.reshape(-1) # Ensure 1D tensor + else: + # Use tensor operations for division + num_repeats_tensor = torch.tensor([len(vid_len) // batch_size], dtype=torch.long, device=device) + + # Build concatenated indices using compile-friendly _tensor_split + vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0) + txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0) + + tgt_idx_list = [] + for i in range(len(vid_len)): + # Add video window + tgt_idx_list.append(vid_idx_splits[i]) + + # Add corresponding text (with repeat) + batch_idx = i % batch_size + tgt_idx_list.append(txt_idx_splits[batch_idx]) + + tgt_idx = torch.cat(tgt_idx_list) + src_idx = torch.argsort(tgt_idx) + txt_idx_len = len(tgt_idx) - len(vid_idx) + + # Pre-compute split lengths for coalescing using tensor operations + repeat_txt_len = txt_len * num_repeats_tensor.squeeze() + + def unconcat_coalesce(all): + """ + Un-concat vid & txt, and coalesce the repeated txt by averaging. + + 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 unflatten with tensor dims (compile-friendly) + """ + vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len]) + + # Coalesce repeated text using unflatten and mean + txt_splits = _tensor_split(txt_out, repeat_txt_len, dim=0) + txt_out_coalesced = [] + + for txt in txt_splits: + # txt has shape (base_len * num_repeats, *other_dims) + # unflatten to (base_len, num_repeats, *other_dims) then average dim 1 + txt = txt.unflatten(0, (-1, num_repeats_tensor.squeeze())).mean(1) + txt_out_coalesced.append(txt) + + return vid_out, torch.cat(txt_out_coalesced) + + # Note: Using direct indexing instead of torch.index_select for backward compatibility + # Direct indexing is deterministic even with repeated indices + return ( + lambda vid, txt: torch.cat([vid, txt])[tgt_idx], + lambda all: unconcat_coalesce(all), + ) + + +def rearrange( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + pattern: str, + **kwargs: Dict[str, int], +) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, +]: + """ + Rearrange flattened tensor using einops pattern. + + Applies einops rearrange to each batch element independently. + + Args: + hid: Flattened tensor (L, c) + hid_shape: Shape metadata (b, n) + pattern: Einops rearrange pattern + **kwargs: Additional arguments for einops + + Returns: + Tuple of (rearranged tensor, new shape metadata) + """ + unflattened = unflatten(hid, hid_shape) + rearranged = [einops.rearrange(h, pattern, **kwargs) for h in unflattened] + return flatten(rearranged) + + +def rearrange_idx( + hid_shape: torch.LongTensor, # (b n) + pattern: str, + **kwargs: Dict[str, int], +) -> Tuple[Callable, Callable, torch.LongTensor]: + """ + Create index-based rearrange functions. + + Pre-computes indices for efficient rearrangement operations. + + Args: + hid_shape: Shape metadata (b, n) + pattern: Einops rearrange pattern + **kwargs: Additional arguments for einops + + Returns: + Tuple of (rearrange_fn, reverse_fn, new_shape) + """ + hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1) + tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs) + tgt_idx = tgt_idx.squeeze(-1) + src_idx = torch.argsort(tgt_idx) + return ( + lambda hid: torch.index_select(hid, 0, tgt_idx), + lambda hid: torch.index_select(hid, 0, src_idx), + tgt_shape, + ) + + +def repeat( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + pattern: str, + **kwargs: Dict[str, torch.LongTensor], # (b) +) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, +]: + """ + Repeat flattened tensor using einops pattern with per-batch parameters. + + Each batch element can have different repeat counts specified in kwargs. + + Args: + hid: Flattened tensor (L, c) + hid_shape: Shape metadata (b, n) + pattern: Einops repeat pattern (e.g., "l c -> t l c") + **kwargs: Repeat parameters as tensors (e.g., t=torch.tensor([2,3,4])) + + Returns: + Tuple of (repeated tensor, new shape metadata) + + COMPILE OPTIMIZATION: + - Minimizes .item() calls + - Only converts to int at the last moment for einops API requirement + """ + unflattened = unflatten(hid, hid_shape) + + # Build kwargs for each batch element + repeated = [] + for i in range(len(unflattened)): + # Extract values for einops (requires Python int) + batch_kwargs = {} + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + # Only convert to Python int where required by einops + batch_kwargs[k] = int(v[i].item()) + else: + batch_kwargs[k] = v + repeated.append(einops.repeat(unflattened[i], pattern, **batch_kwargs)) + + return flatten(repeated) + + +def pack( + samples: List[torch.Tensor], # List of (h w c). +) -> Tuple[ + List[torch.Tensor], # groups [(b1 h1 w1 c1), (b2 h2 w2 c2)] + List[List[int]], # reversal indices. +]: + """ + Group samples by shape and return grouped batches with reversal indices. + + Useful for batch processing samples with different spatial dimensions. + + Args: + samples: List of tensors with potentially different shapes + + Returns: + Tuple of (batched_groups, reversal_indices) for unpacking + """ + batches = {} + indices = {} + for i, sample in enumerate(samples): + shape = sample.shape + batches[shape] = batches.get(shape, []) + indices[shape] = indices.get(shape, []) + batches[shape].append(sample) + indices[shape].append(i) + + batches = list(map(torch.stack, batches.values())) + indices = list(indices.values()) + return batches, indices + + +def unpack( + batches: List[torch.Tensor], + indices: List[List[int]], +) -> List[torch.Tensor]: + """ + Unpack grouped batches back to original order. + + Inverse of pack(). + + Args: + batches: Grouped batches from pack() + indices: Reversal indices from pack() + + Returns: + List of tensors in original order + """ + samples = [None] * (max(chain(*indices)) + 1) + for batch, index in zip(batches, indices): + for sample, i in zip(batch.unbind(), index): + samples[i] = sample + return samples + + +def window( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + window_fn: Callable[[torch.Tensor], List[torch.Tensor]], +): + """ + Apply windowing function to create non-overlapping windows. + + Used for window attention mechanisms where sequences are split into windows. + + Args: + hid: Flattened tensor (L, c) + hid_shape: Shape metadata (b, n) + window_fn: Function that splits a tensor into windows + + Returns: + Tuple of (windowed_tensor, window_shapes, window_counts) + + COMPILE OPTIMIZATION: Uses tensor operation for window count tracking + """ + unflattened = unflatten(hid, hid_shape) + windowed = [window_fn(h) for h in unflattened] + + # Track window counts using tensor operations + device = hid_shape.device + hid_windows = torch.tensor([len(w) for w in windowed], dtype=torch.long, device=device) + + # Flatten all windows + all_windows = list(chain(*windowed)) + hid, hid_shape = flatten(all_windows) + return hid, hid_shape, hid_windows + + +def window_idx( + hid_shape: torch.LongTensor, # (b n) + window_fn: Callable[[torch.Tensor], List[torch.Tensor]], +): + """ + Create index-based windowing functions. + + Pre-computes indices for efficient windowing and reverse operations. + + Args: + hid_shape: Shape metadata (b, n) + window_fn: Function that splits a tensor into windows + + Returns: + Tuple of (window_fn, reverse_fn, window_shapes, window_counts) + """ + hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1) + tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn) + tgt_idx = tgt_idx.squeeze(-1) + src_idx = torch.argsort(tgt_idx) + return ( + lambda hid: torch.index_select(hid, 0, tgt_idx), + lambda hid: torch.index_select(hid, 0, src_idx), + tgt_shape, + tgt_windows, + ) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c1a9da26ef760575192042ea32b01bd9cd1a267d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/__init__.py @@ -0,0 +1,26 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from .mmsr_block import NaMMSRTransformerBlock + + +nadit_blocks = { + "mmdit_sr": NaMMSRTransformerBlock, +} + + +def get_nablock(block_type: str): + if block_type in nadit_blocks: + return nadit_blocks[block_type] + raise NotImplementedError(f"{block_type} is not supported") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a7561025245d888d26ade38f25668efb216cd907 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/__init__.py @@ -0,0 +1,25 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from .mmattn import NaMMAttention + +attns = { + "mm_full": NaMMAttention, +} + + +def get_attn(attn_type: str): + if attn_type in attns: + return attns[attn_type] + raise NotImplementedError(f"{attn_type} is not supported") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/mmattn.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/mmattn.py new file mode 100644 index 0000000000000000000000000000000000000000..a311449fb489f960e0a781ea1aa225dfceaf7f48 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/attention/mmattn.py @@ -0,0 +1,271 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Optional, Tuple, Union +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F +from torch.nn.modules.utils import _triple + +from .....common.cache import Cache +from .....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv +from .....common.half_precision_fixes import safe_pad_operation + +from ... import na +from ...attention import FlashAttentionVarlen +from ...mm import MMArg, MMModule +from ...normalization import norm_layer_type +from ...rope import get_na_rope +from ...window import get_window_op +from itertools import chain + + +class NaMMAttention(nn.Module): + def __init__( + self, + vid_dim: int, + txt_dim: int, + heads: int, + head_dim: int, + qk_bias: bool, + qk_norm: norm_layer_type, + qk_norm_eps: float, + rope_type: Optional[str], + rope_dim: int, + shared_weights: bool, + attention_mode: str = 'sdpa', + **kwargs, + ): + super().__init__() + dim = MMArg(vid_dim, txt_dim) + inner_dim = heads * head_dim + qkv_dim = inner_dim * 3 + self.head_dim = head_dim + self.proj_qkv = MMModule( + nn.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_weights + ) + self.proj_out = MMModule(nn.Linear, inner_dim, dim, shared_weights=shared_weights) + self.norm_q = MMModule( + qk_norm, + dim=head_dim, + eps=qk_norm_eps, + elementwise_affine=True, + shared_weights=shared_weights, + ) + self.norm_k = MMModule( + qk_norm, + dim=head_dim, + eps=qk_norm_eps, + elementwise_affine=True, + shared_weights=shared_weights, + ) + + self.rope = get_na_rope(rope_type=rope_type, dim=rope_dim) + self.attn = FlashAttentionVarlen(attention_mode=attention_mode) + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_qkv, txt_qkv = self.proj_qkv(vid, txt) + vid_qkv = gather_seq_scatter_heads_qkv( + vid_qkv, + seq_dim=0, + qkv_shape=vid_shape, + cache=cache.namespace("vid"), + ) + txt_qkv = gather_seq_scatter_heads_qkv( + txt_qkv, + seq_dim=0, + qkv_shape=txt_shape, + cache=cache.namespace("txt"), + ) + vid_qkv = rearrange(vid_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim) + txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim) + + vid_q, vid_k, vid_v = vid_qkv.unbind(1) + txt_q, txt_k, txt_v = txt_qkv.unbind(1) + + vid_q, txt_q = self.norm_q(vid_q, txt_q) + vid_k, txt_k = self.norm_k(vid_k, txt_k) + + if self.rope: + if self.rope.mm: + vid_q, vid_k, txt_q, txt_k = self.rope( + vid_q, vid_k, vid_shape, txt_q, txt_k, txt_shape, cache + ) + else: + vid_q, vid_k = self.rope(vid_q, vid_k, vid_shape, cache) + + vid_len = cache("vid_len", lambda: vid_shape.prod(-1)) + txt_len = cache("txt_len", lambda: txt_shape.prod(-1)) + all_len = cache("all_len", lambda: vid_len + txt_len) + + concat, unconcat = cache("mm_pnp", lambda: na.concat_idx(vid_len, txt_len)) + + # Attention handles dtype conversion internally using pipeline compute_dtype + attn = self.attn( + q=concat(vid_q, txt_q), + k=concat(vid_k, txt_k), + v=concat(vid_v, txt_v), + cu_seqlens_q=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()), + cu_seqlens_k=cache("mm_seqlens", lambda: safe_pad_operation(all_len.cumsum(0), (1, 0)).int()), + max_seqlen_q=cache("mm_maxlen", lambda: all_len.max()), + max_seqlen_k=cache("mm_maxlen", lambda: all_len.max()), + ).type_as(vid_q) + + attn = rearrange(attn, "l h d -> l (h d)") + vid_out, txt_out = unconcat(attn) + vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0) + txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0) + + vid_out, txt_out = self.proj_out(vid_out, txt_out) + return vid_out, txt_out + + +class NaSwinAttention(NaMMAttention): + def __init__( + self, + *args, + window: Union[int, Tuple[int, int, int]], + window_method: str, + attention_mode: str = 'sdpa', + **kwargs, + ): + super().__init__(*args, attention_mode=attention_mode, **kwargs) + self.window = _triple(window) + self.window_method = window_method + assert all(map(lambda v: isinstance(v, int) and v >= 0, self.window)) + + self.window_op = get_window_op(window_method) + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + + vid_qkv, txt_qkv = self.proj_qkv(vid, txt) + vid_qkv = gather_seq_scatter_heads_qkv( + vid_qkv, + seq_dim=0, + qkv_shape=vid_shape, + cache=cache.namespace("vid"), + ) + txt_qkv = gather_seq_scatter_heads_qkv( + txt_qkv, + seq_dim=0, + qkv_shape=txt_shape, + cache=cache.namespace("txt"), + ) + + # re-org the input seq for window attn + cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3") + + def make_window(x: torch.Tensor): + t, h, w, _ = x.shape + window_slices = self.window_op((t, h, w), self.window) + return [x[st, sh, sw] for (st, sh, sw) in window_slices] + + window_partition, window_reverse, window_shape, window_count = cache_win( + "win_transform", + lambda: na.window_idx(vid_shape, make_window), + ) + vid_qkv_win = window_partition(vid_qkv) + + vid_qkv_win = rearrange(vid_qkv_win, "l (o h d) -> l o h d", o=3, d=self.head_dim) + txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim) + + vid_q, vid_k, vid_v = vid_qkv_win.unbind(1) + txt_q, txt_k, txt_v = txt_qkv.unbind(1) + + vid_q, txt_q = self.norm_q(vid_q, txt_q) + vid_k, txt_k = self.norm_k(vid_k, txt_k) + + txt_len = cache("txt_len", lambda: txt_shape.prod(-1)) + + vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1)) + txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count)) + all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win) + concat_win, unconcat_win = cache_win( + "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count) + ) + + # window rope + if self.rope: + if self.rope.mm: + # repeat text q and k for window mmrope + _, num_h, _ = txt_q.shape + txt_q_repeat = rearrange(txt_q, "l h d -> l (h d)") + txt_q_repeat = na.unflatten(txt_q_repeat, txt_shape) + txt_q_repeat = [[x] * n for x, n in zip(txt_q_repeat, window_count)] + txt_q_repeat = list(chain(*txt_q_repeat)) + txt_q_repeat, txt_shape_repeat = na.flatten(txt_q_repeat) + txt_q_repeat = rearrange(txt_q_repeat, "l (h d) -> l h d", h=num_h) + + txt_k_repeat = rearrange(txt_k, "l h d -> l (h d)") + txt_k_repeat = na.unflatten(txt_k_repeat, txt_shape) + txt_k_repeat = [[x] * n for x, n in zip(txt_k_repeat, window_count)] + txt_k_repeat = list(chain(*txt_k_repeat)) + txt_k_repeat, _ = na.flatten(txt_k_repeat) + txt_k_repeat = rearrange(txt_k_repeat, "l (h d) -> l h d", h=num_h) + + vid_q, vid_k, txt_q, txt_k = self.rope( + vid_q, vid_k, window_shape, txt_q_repeat, txt_k_repeat, txt_shape_repeat, cache_win + ) + else: + vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win) + + # Attention handles dtype conversion internally using pipeline compute_dtype + out = self.attn( + q=concat_win(vid_q, txt_q), + k=concat_win(vid_k, txt_k), + v=concat_win(vid_v, txt_v), + cu_seqlens_q=cache_win( + "vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int() + ), + cu_seqlens_k=cache_win( + "vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int() + ), + max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max()), + max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max()), + ).type_as(vid_q) + + # text pooling + vid_out, txt_out = unconcat_win(out) + + vid_out = rearrange(vid_out, "l h d -> l (h d)") + txt_out = rearrange(txt_out, "l h d -> l (h d)") + vid_out = window_reverse(vid_out) + + vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0) + txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0) + + vid_out, txt_out = self.proj_out(vid_out, txt_out) + + return vid_out, txt_out \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/mmsr_block.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/mmsr_block.py new file mode 100644 index 0000000000000000000000000000000000000000..2f9a562c6024595a748c28e7af32d5a6c3bf61c9 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nablocks/mmsr_block.py @@ -0,0 +1,128 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Tuple +import torch +import torch.nn as nn + +# from ..cache import Cache +from ....common.cache import Cache + +from .attention.mmattn import NaSwinAttention +from ..mm import MMArg +from ..modulation import ada_layer_type +from ..normalization import norm_layer_type +from ..mm import MMArg, MMModule +from ..mlp import get_mlp + + +class NaMMSRTransformerBlock(nn.Module): + def __init__( + self, + *, + vid_dim: int, + txt_dim: int, + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm: norm_layer_type, + norm_eps: float, + ada: ada_layer_type, + qk_bias: bool, + qk_norm: norm_layer_type, + mlp_type: str, + shared_weights: bool, + rope_type: str, + rope_dim: int, + is_last_layer: bool, + attention_mode: str = 'sdpa', + **kwargs, + ): + super().__init__() + dim = MMArg(vid_dim, txt_dim) + self.attn_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights,) + + self.attn = NaSwinAttention( + vid_dim=vid_dim, + txt_dim=txt_dim, + heads=heads, + head_dim=head_dim, + qk_bias=qk_bias, + qk_norm=qk_norm, + qk_norm_eps=norm_eps, + rope_type=rope_type, + rope_dim=rope_dim, + shared_weights=shared_weights, + attention_mode=attention_mode, + window=kwargs.pop("window", None), + window_method=kwargs.pop("window_method", None), + ) + + self.mlp_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False, shared_weights=shared_weights, vid_only=is_last_layer) + self.mlp = MMModule( + get_mlp(mlp_type), + dim=dim, + expand_ratio=expand_ratio, + shared_weights=shared_weights, + vid_only=is_last_layer + ) + self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"], shared_weights=shared_weights, vid_only=is_last_layer) + self.is_last_layer = is_last_layer + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + emb: torch.FloatTensor, + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + torch.LongTensor, + torch.LongTensor, + ]: + hid_len = MMArg( + cache("vid_len", lambda: vid_shape.prod(-1)), + cache("txt_len", lambda: txt_shape.prod(-1)), + ) + ada_kwargs = { + "emb": emb, + "hid_len": hid_len, + "cache": cache, + "branch_tag": MMArg("vid", "txt"), + } + + vid_attn, txt_attn = self.attn_norm(vid, txt) + + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs) + vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs) + vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt) + + vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn) + # ADD BY NUMZ + if vid_mlp.dtype != vid_attn.dtype: + vid_mlp = vid_mlp.to(vid_attn.dtype) + if txt_mlp.dtype != txt_attn.dtype: + txt_mlp = txt_mlp.to(txt_attn.dtype) + # END BY NUMZ + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs) + vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs) + vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn) + + return vid_mlp, txt_mlp, vid_shape, txt_shape diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nadit.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nadit.py new file mode 100644 index 0000000000000000000000000000000000000000..6f3b88deae03f23cbab2febf6034362ad806bd3d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/nadit.py @@ -0,0 +1,248 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from dataclasses import dataclass +from typing import List, Optional, Tuple, Union, Callable +import torch +from torch import nn + +from ...common.cache import Cache +from ...common.distributed.ops import slice_inputs + +from . import na +from .embedding import TimeEmbedding +from .modulation import get_ada_layer +from .nablocks import get_nablock +from .normalization import get_norm_layer +from .patch import get_na_patch_layers + +# Fake func, no checkpointing is required for inference +def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs): + return module(*args, **kwargs) + +@dataclass +class NaDiTOutput: + vid_sample: torch.Tensor + + +class NaDiT(nn.Module): + """ + Native Resolution Diffusion Transformer (NaDiT) + """ + + gradient_checkpointing = False + + def __init__( + self, + vid_in_channels: int, + vid_out_channels: int, + vid_dim: int, + txt_in_dim: Union[int, List[int]], + txt_dim: Optional[int], + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm: Optional[str], + norm_eps: float, + ada: str, + qk_bias: bool, + qk_norm: Optional[str], + patch_size: Union[int, Tuple[int, int, int]], + num_layers: int, + block_type: Union[str, Tuple[str]], + mm_layers: Union[int, Tuple[bool]], + mlp_type: str = "normal", + patch_type: str = "v1", + rope_type: Optional[str] = "rope3d", + rope_dim: Optional[int] = None, + window: Optional[Tuple] = None, + window_method: Optional[Tuple[str]] = None, + msa_type: Optional[Tuple[str]] = None, + mca_type: Optional[Tuple[str]] = None, + txt_in_norm: Optional[str] = None, + txt_in_norm_scale_factor: int = 0.01, + txt_proj_type: Optional[str] = "linear", + vid_out_norm: Optional[str] = None, + attention_mode: str = 'sdpa', + **kwargs, + ): + ada = get_ada_layer(ada) + norm = get_norm_layer(norm) + qk_norm = get_norm_layer(qk_norm) + rope_dim = rope_dim if rope_dim is not None else head_dim // 2 + if isinstance(block_type, str): + block_type = [block_type] * num_layers + elif len(block_type) != num_layers: + raise ValueError("The ``block_type`` list should equal to ``num_layers``.") + super().__init__() + NaPatchIn, NaPatchOut = get_na_patch_layers(patch_type) + self.vid_in = NaPatchIn( + in_channels=vid_in_channels, + patch_size=patch_size, + dim=vid_dim, + ) + if not isinstance(txt_in_dim, int): + self.txt_in = nn.ModuleList([]) + for in_dim in txt_in_dim: + txt_norm_layer = get_norm_layer(txt_in_norm)(txt_dim, norm_eps, True) + if txt_proj_type == "linear": + txt_proj_layer = nn.Linear(in_dim, txt_dim) + else: + txt_proj_layer = nn.Sequential( + nn.Linear(in_dim, in_dim), nn.GELU("tanh"), nn.Linear(in_dim, txt_dim) + ) + torch.nn.init.constant_(txt_norm_layer.weight, txt_in_norm_scale_factor) + self.txt_in.append( + nn.Sequential( + txt_proj_layer, + txt_norm_layer, + ) + ) + else: + self.txt_in = ( + nn.Linear(txt_in_dim, txt_dim) + if txt_in_dim and txt_in_dim != txt_dim + else nn.Identity() + ) + self.emb_in = TimeEmbedding( + sinusoidal_dim=256, + hidden_dim=max(vid_dim, txt_dim), + output_dim=emb_dim, + ) + + if window is None or isinstance(window[0], int): + window = [window] * num_layers + if window_method is None or isinstance(window_method, str): + window_method = [window_method] * num_layers + + if msa_type is None or isinstance(msa_type, str): + msa_type = [msa_type] * num_layers + if mca_type is None or isinstance(mca_type, str): + mca_type = [mca_type] * num_layers + + self.blocks = nn.ModuleList( + [ + get_nablock(block_type[i])( + vid_dim=vid_dim, + txt_dim=txt_dim, + emb_dim=emb_dim, + heads=heads, + head_dim=head_dim, + expand_ratio=expand_ratio, + norm=norm, + norm_eps=norm_eps, + ada=ada, + qk_bias=qk_bias, + qk_norm=qk_norm, + shared_weights=not ( + (i < mm_layers) if isinstance(mm_layers, int) else mm_layers[i] + ), + mlp_type=mlp_type, + window=window[i], + window_method=window_method[i], + msa_type=msa_type[i], + mca_type=mca_type[i], + rope_type=rope_type, + rope_dim=rope_dim, + is_last_layer=(i == num_layers - 1), + attention_mode=attention_mode, + **kwargs, + ) + for i in range(num_layers) + ] + ) + + self.vid_out_norm = None + if vid_out_norm is not None: + self.vid_out_norm = get_norm_layer(vid_out_norm)( + dim=vid_dim, + eps=norm_eps, + elementwise_affine=True, + ) + self.vid_out_ada = ada( + dim=vid_dim, + emb_dim=emb_dim, + layers=["out"], + modes=["in"], + ) + + self.vid_out = NaPatchOut( + out_channels=vid_out_channels, + patch_size=patch_size, + dim=vid_dim, + ) + + def set_gradient_checkpointing(self, enable: bool): + self.gradient_checkpointing = enable + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: Union[torch.FloatTensor, List[torch.FloatTensor]], # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: Union[torch.LongTensor, List[torch.LongTensor]], # b 1 + timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], # b + disable_cache: bool = False, # for test + ): + cache = Cache(disable=disable_cache) + + # slice vid after patching in when using sequence parallelism + if isinstance(txt, list): + assert isinstance(self.txt_in, nn.ModuleList) + txt = [ + na.unflatten(fc(i), s) for fc, i, s in zip(self.txt_in, txt, txt_shape) + ] # B L D + txt, txt_shape = na.flatten([torch.cat(t, dim=0) for t in zip(*txt)]) + txt = slice_inputs(txt, dim=0) + else: + txt = slice_inputs(txt, dim=0) + txt = self.txt_in(txt) + + # Video input. + # Sequence parallel slicing is done inside patching class. + vid, vid_shape = self.vid_in(vid, vid_shape, cache) + + # Embedding input. + emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype) + + # Body + for i, block in enumerate(self.blocks): + vid, txt, vid_shape, txt_shape = gradient_checkpointing( + enabled=(self.gradient_checkpointing and self.training), + module=block, + vid=vid, + txt=txt, + vid_shape=vid_shape, + txt_shape=txt_shape, + emb=emb, + cache=cache, + ) + + # Video output norm. + if self.vid_out_norm: + vid = self.vid_out_norm(vid) + vid = self.vid_out_ada( + vid, + emb=emb, + layer="out", + mode="in", + hid_len=cache("vid_len", lambda: vid_shape.prod(-1)), + cache=cache, + branch_tag="vid", + ) + + # Video output. + vid, vid_shape = self.vid_out(vid, vid_shape, cache) + return NaDiTOutput(vid_sample=vid) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/normalization.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..155c9981b818980d28c2a445855d7d1c808db23d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/normalization.py @@ -0,0 +1,150 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Callable, Optional +from diffusers.models.normalization import RMSNorm +from torch import nn +import torch +import torch.nn.functional as F +import numbers +from torch.nn.parameter import Parameter +from torch.nn import init + +# (dim: int, eps: float, elementwise_affine: bool) +norm_layer_type = Callable[[int, float, bool], nn.Module] + + +class CustomLayerNorm(nn.Module): + """ + Custom LayerNorm implementation to replace Apex FusedLayerNorm + """ + def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True): + super(CustomLayerNorm, self).__init__() + + if isinstance(normalized_shape, numbers.Integral): + normalized_shape = (normalized_shape,) + self.normalized_shape = torch.Size(normalized_shape) + self.eps = eps + self.elementwise_affine = elementwise_affine + + if self.elementwise_affine: + self.weight = Parameter(torch.Tensor(*normalized_shape)) + self.bias = Parameter(torch.Tensor(*normalized_shape)) + else: + self.register_parameter('weight', None) + self.register_parameter('bias', None) + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + init.ones_(self.weight) + init.zeros_(self.bias) + + def forward(self, input): + # ๐Ÿš€ FP8 COMPATIBILITY: Convert parameters to match input dtype + # This prevents "Promotion for Float8 Types is not supported" errors + weight = self.weight + bias = self.bias + + if self.elementwise_affine and weight is not None: + if weight.dtype != input.dtype: + weight = weight.to(input.dtype) + if bias is not None and bias.dtype != input.dtype: + bias = bias.to(input.dtype) + + return F.layer_norm( + input, self.normalized_shape, weight, bias, self.eps) + + +class CustomRMSNorm(nn.Module): + """ + Custom RMSNorm implementation to replace Apex FusedRMSNorm + """ + def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True): + super(CustomRMSNorm, self).__init__() + + if isinstance(normalized_shape, numbers.Integral): + normalized_shape = (normalized_shape,) + self.normalized_shape = torch.Size(normalized_shape) + self.eps = eps + self.elementwise_affine = elementwise_affine + + if self.elementwise_affine: + self.weight = Parameter(torch.ones(*normalized_shape)) + else: + self.register_parameter('weight', None) + + def forward(self, input): + # RMS normalization: x / sqrt(mean(x^2) + eps) * weight + dims = tuple(range(-len(self.normalized_shape), 0)) + + # Calculate RMS: sqrt(mean(x^2)) + variance = input.pow(2).mean(dim=dims, keepdim=True) + rms = torch.sqrt(variance + self.eps) + + # Normalize + normalized = input / rms + + if self.elementwise_affine: + # Convert FP8 weight to match input dtype for arithmetic operations + if hasattr(torch, 'float8_e4m3fn'): + fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2) + if self.weight.dtype in fp8_types: + # Use input dtype as target (respects pipeline precision) + weight = self.weight.to(input.dtype) + return normalized * weight + + return normalized * self.weight + return normalized + + +def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type: + + def _norm_layer(dim: int, eps: float, elementwise_affine: bool): + if norm_type is None: + return nn.Identity() + + if norm_type == "layer": + return nn.LayerNorm( + normalized_shape=dim, + eps=eps, + elementwise_affine=elementwise_affine, + ) + + if norm_type == "rms": + return RMSNorm( + dim=dim, + eps=eps, + elementwise_affine=elementwise_affine, + ) + + if norm_type == "fusedln": + # Use custom LayerNorm instead of Apex FusedLayerNorm + return CustomLayerNorm( + normalized_shape=dim, + elementwise_affine=elementwise_affine, + eps=eps, + ) + + if norm_type == "fusedrms": + # Use custom RMSNorm instead of Apex FusedRMSNorm + return CustomRMSNorm( + normalized_shape=dim, + elementwise_affine=elementwise_affine, + eps=eps, + ) + + raise NotImplementedError(f"{norm_type} is not supported") + + return _norm_layer \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4e3c9783163f1e671f2d946dfad39ca33b12843d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/__init__.py @@ -0,0 +1,19 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +def get_na_patch_layers(patch_type="v1"): + assert patch_type in ["v1"] + if patch_type == "v1": + from .patch_v1 import NaPatchIn, NaPatchOut + return NaPatchIn, NaPatchOut diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/patch_v1.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/patch_v1.py new file mode 100644 index 0000000000000000000000000000000000000000..14b5a62d5fd81ed07e9ae7873900d428dcae73f0 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/patch/patch_v1.py @@ -0,0 +1,127 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Tuple, Union +import torch +from einops import rearrange +from torch import nn +from torch.nn.modules.utils import _triple + +from ....common.cache import Cache +from ....common.distributed.ops import gather_outputs, slice_inputs + +from .. import na + + +class PatchIn(nn.Module): + def __init__( + self, + in_channels: int, + patch_size: Union[int, Tuple[int, int, int]], + dim: int, + ): + super().__init__() + t, h, w = _triple(patch_size) + self.patch_size = t, h, w + self.proj = nn.Linear(in_channels * t * h * w, dim) + + def forward( + self, + vid: torch.Tensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + if t > 1: + assert vid.size(2) % t == 1 + vid = torch.cat([vid[:, :, :1]] * (t - 1) + [vid], dim=2) + vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w) + vid = self.proj(vid) + return vid + + +class PatchOut(nn.Module): + def __init__( + self, + out_channels: int, + patch_size: Union[int, Tuple[int, int, int]], + dim: int, + ): + super().__init__() + t, h, w = _triple(patch_size) + self.patch_size = t, h, w + self.proj = nn.Linear(dim, out_channels * t * h * w) + + def forward( + self, + vid: torch.Tensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + vid = self.proj(vid) + vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w) + if t > 1: + vid = vid[:, :, (t - 1) :] + return vid + + +class NaPatchIn(PatchIn): + def forward( + self, + vid: torch.Tensor, # l c + vid_shape: torch.LongTensor, + cache: Cache = Cache(disable=True), # for test + ) -> torch.Tensor: + cache = cache.namespace("patch") + vid_shape_before_patchify = cache("vid_shape_before_patchify", lambda: vid_shape) + t, h, w = self.patch_size + if not (t == h == w == 1): + vid = na.unflatten(vid, vid_shape) + for i in range(len(vid)): + if t > 1 and vid_shape_before_patchify[i, 0] % t != 0: + vid[i] = torch.cat([vid[i][:1]] * (t - vid[i].size(0) % t) + [vid[i]], dim=0) + vid[i] = rearrange(vid[i], "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w) + vid, vid_shape = na.flatten(vid) + + # slice vid after patching in when using sequence parallelism + vid = slice_inputs(vid, dim=0) + vid = self.proj(vid) + return vid, vid_shape + + +class NaPatchOut(PatchOut): + def forward( + self, + vid: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, + cache: Cache = Cache(disable=True), # for test + ) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, + ]: + cache = cache.namespace("patch") + vid_shape_before_patchify = cache.get("vid_shape_before_patchify") + + t, h, w = self.patch_size + vid = self.proj(vid) + # gather vid before patching out when enabling sequence parallelism + vid = gather_outputs( + vid, gather_dim=0, padding_dim=0, unpad_shape=vid_shape, cache=cache.namespace("vid") + ) + if not (t == h == w == 1): + vid = na.unflatten(vid, vid_shape) + for i in range(len(vid)): + vid[i] = rearrange(vid[i], "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w) + if t > 1 and vid_shape_before_patchify[i, 0] % t != 0: + vid[i] = vid[i][(t - vid_shape_before_patchify[i, 0] % t) :] + vid, vid_shape = na.flatten(vid) + + return vid, vid_shape diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/rope.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/rope.py new file mode 100644 index 0000000000000000000000000000000000000000..b74c3bd40b3e03cd037c04fa84860c1a6a3e4250 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/rope.py @@ -0,0 +1,184 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from functools import lru_cache +from typing import Optional, Tuple +import torch +from einops import rearrange +from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb +from torch import nn + +from ...common.cache import Cache + + +class RotaryEmbeddingBase(nn.Module): + def __init__(self, dim: int, rope_dim: int): + super().__init__() + self.rope = RotaryEmbedding( + dim=dim // rope_dim, + freqs_for="pixel", + max_freq=256, + ) + # 1. Set model.requires_grad_(True) after model creation will make + # the `requires_grad=False` for rope freqs no longer hold. + # 2. Even if we don't set requires_grad_(True) explicitly, + # FSDP is not memory efficient when handling fsdp_wrap + # with mixed requires_grad=True/False. + # With above consideration, it is easier just remove the freqs + # out of nn.Parameters when `learned_freq=False` + freqs = self.rope.freqs + del self.rope.freqs + self.rope.register_buffer("freqs", freqs.data) + + @lru_cache(maxsize=128) + def get_axial_freqs(self, *dims): + return self.rope.get_axial_freqs(*dims) + + +class RotaryEmbedding3d(RotaryEmbeddingBase): + def __init__(self, dim: int): + super().__init__(dim, rope_dim=3) + self.mm = False + + def forward( + self, + q: torch.FloatTensor, # b h l d + k: torch.FloatTensor, # b h l d + size: Tuple[int, int, int], + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + T, H, W = size + freqs = self.get_axial_freqs(T, H, W) + q = rearrange(q, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W) + k = rearrange(k, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W) + q = apply_rotary_emb(freqs, q.float()).to(q.dtype) + k = apply_rotary_emb(freqs, k.float()).to(k.dtype) + q = rearrange(q, "b h T H W d -> b h (T H W) d") + k = rearrange(k, "b h T H W d -> b h (T H W) d") + return q, k + + +class MMRotaryEmbeddingBase(RotaryEmbeddingBase): + def __init__(self, dim: int, rope_dim: int): + super().__init__(dim, rope_dim) + self.rope = RotaryEmbedding( + dim=dim // rope_dim, + freqs_for="lang", + theta=10000, + ) + freqs = self.rope.freqs + del self.rope.freqs + self.rope.register_buffer("freqs", freqs.data) + self.mm = True + + +class NaMMRotaryEmbedding3d(MMRotaryEmbeddingBase): + def __init__(self, dim: int): + super().__init__(dim, rope_dim=3) + + def forward( + self, + vid_q: torch.FloatTensor, # L h d + vid_k: torch.FloatTensor, # L h d + vid_shape: torch.LongTensor, # B 3 + txt_q: torch.FloatTensor, # L h d + txt_k: torch.FloatTensor, # L h d + txt_shape: torch.LongTensor, # B 1 + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_freqs, txt_freqs = cache( + "mmrope_freqs_3d", + lambda: self.get_freqs(vid_shape, txt_shape), + ) + target_device = vid_q.device + if vid_freqs.device != target_device: + vid_freqs = vid_freqs.to(target_device) + if txt_freqs.device != target_device: + txt_freqs = txt_freqs.to(target_device) + vid_q = rearrange(vid_q, "L h d -> h L d") + vid_k = rearrange(vid_k, "L h d -> h L d") + vid_q = apply_rotary_emb(vid_freqs, vid_q.float()).to(vid_q.dtype) + vid_k = apply_rotary_emb(vid_freqs, vid_k.float()).to(vid_k.dtype) + vid_q = rearrange(vid_q, "h L d -> L h d") + vid_k = rearrange(vid_k, "h L d -> L h d") + + txt_q = rearrange(txt_q, "L h d -> h L d") + txt_k = rearrange(txt_k, "L h d -> h L d") + txt_q = apply_rotary_emb(txt_freqs, txt_q.float()).to(txt_q.dtype) + txt_k = apply_rotary_emb(txt_freqs, txt_k.float()).to(txt_k.dtype) + txt_q = rearrange(txt_q, "h L d -> L h d") + txt_k = rearrange(txt_k, "h L d -> L h d") + return vid_q, vid_k, txt_q, txt_k + + @torch._dynamo.disable # Disable compilation: .tolist() is data-dependent and causes graph breaks + def get_freqs( + self, + vid_shape: torch.LongTensor, + txt_shape: torch.LongTensor, + ) -> Tuple[ + torch.Tensor, + torch.Tensor, + ]: + """ + Generate RoPE frequencies for variable batch shapes. + + Note: This method uses @torch._dynamo.disable because it requires + data-dependent control flow (shape.tolist()) that cannot be symbolically + traced by torch.compile. The cache() wrapper in the forward pass memoizes + results to reduce recomputation overhead. + """ + # Calculate actual max dimensions needed for this batch + max_temporal = 0 + max_height = 0 + max_width = 0 + max_txt_len = 0 + + for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()): + max_temporal = max(max_temporal, l + f) # Need up to l+f for temporal + max_height = max(max_height, h) + max_width = max(max_width, w) + max_txt_len = max(max_txt_len, l) + + # Compute frequencies for actual max dimensions needed + # Add small buffer to improve cache hits across similar batches + vid_freqs = self.get_axial_freqs( + min(max_temporal + 16, 1024), # Cap at 1024, add small buffer + min(max_height + 4, 128), # Cap at 128, add small buffer + min(max_width + 4, 128) # Cap at 128, add small buffer + ) + txt_freqs = self.get_axial_freqs(min(max_txt_len + 16, 1024)) + + # Now slice as before + vid_freq_list, txt_freq_list = [], [] + for (f, h, w), l in zip(vid_shape.tolist(), txt_shape[:, 0].tolist()): + vid_freq = vid_freqs[l : l + f, :h, :w].reshape(-1, vid_freqs.size(-1)) + txt_freq = txt_freqs[:l].repeat(1, 3).reshape(-1, vid_freqs.size(-1)) + vid_freq_list.append(vid_freq) + txt_freq_list.append(txt_freq) + return torch.cat(vid_freq_list, dim=0), torch.cat(txt_freq_list, dim=0) + + +def get_na_rope(rope_type: Optional[str], dim: int): + if rope_type is None: + return None + if rope_type == "mmrope3d": + return NaMMRotaryEmbedding3d(dim=dim) + raise NotImplementedError(f"{rope_type} is not supported.") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/window.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/window.py new file mode 100644 index 0000000000000000000000000000000000000000..b7475921ae283cf76d82bff7521233c133f54bfd --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_3b/window.py @@ -0,0 +1,83 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from math import ceil +from typing import Tuple +import math + +def get_window_op(name: str): + if name == "720pwin_by_size_bysize": + return make_720Pwindows_bysize + if name == "720pswin_by_size_bysize": + return make_shifted_720Pwindows_bysize + raise ValueError(f"Unknown windowing method: {name}") + + +# -------------------------------- Windowing -------------------------------- # +def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]): + t, h, w = size + resized_nt, resized_nh, resized_nw = num_windows + #cal windows under 720p + scale = math.sqrt((45 * 80) / (h * w)) + resized_h, resized_w = round(h * scale), round(w * scale) + wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size. + wt = ceil(min(t, 30) / resized_nt) # window size. + nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww) # window size. + return [ + ( + slice(it * wt, min((it + 1) * wt, t)), + slice(ih * wh, min((ih + 1) * wh, h)), + slice(iw * ww, min((iw + 1) * ww, w)), + ) + for iw in range(nw) + if min((iw + 1) * ww, w) > iw * ww + for ih in range(nh) + if min((ih + 1) * wh, h) > ih * wh + for it in range(nt) + if min((it + 1) * wt, t) > it * wt + ] + +def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]): + t, h, w = size + resized_nt, resized_nh, resized_nw = num_windows + #cal windows under 720p + scale = math.sqrt((45 * 80) / (h * w)) + resized_h, resized_w = round(h * scale), round(w * scale) + wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size. + wt = ceil(min(t, 30) / resized_nt) # window size. + + st, sh, sw = ( # shift size. + 0.5 if wt < t else 0, + 0.5 if wh < h else 0, + 0.5 if ww < w else 0, + ) + nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww) # window size. + nt, nh, nw = ( # number of window. + nt + 1 if st > 0 else 1, + nh + 1 if sh > 0 else 1, + nw + 1 if sw > 0 else 1, + ) + return [ + ( + slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)), + slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)), + slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)), + ) + for iw in range(nw) + if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0) + for ih in range(nh) + if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0) + for it in range(nt) + if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0) + ] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/attention.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..03e3c634d7de24ee90243f9e1cd99ac92d786714 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/attention.py @@ -0,0 +1,148 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import torch +import torch.nn.functional as F + +# Import flash/sage attn with automatic fallback from compatibility layer +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 + + +def pytorch_varlen_attention(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q=None, max_seqlen_k=None, dropout_p=0.0, softmax_scale=None, causal=False, deterministic=False): + """ + A PyTorch-based implementation of variable-length attention to replace flash_attn_varlen_func. + It processes each sequence in the batch individually. + + NOTE: max_seqlen_q and max_seqlen_k are accepted for API compatibility but not used. + PyTorch's scaled_dot_product_attention automatically handles variable sequence lengths. + + COMPILE OPTIMIZATION: Uses torch.tensor_split to avoid .item() graph breaks + """ + # Split q, k, v using cumulative sequence lengths + # NOTE: torch.tensor_split requires int64 dtype and CPU device (PyTorch requirements) + q_splits = list(torch.tensor_split(q, cu_seqlens_q[1:-1].long().cpu(), dim=0)) + k_splits = list(torch.tensor_split(k, cu_seqlens_k[1:-1].long().cpu(), dim=0)) + v_splits = list(torch.tensor_split(v, cu_seqlens_k[1:-1].long().cpu(), dim=0)) + + # Process each sequence + output_splits = [] + for q_i, k_i, v_i in zip(q_splits, k_splits, v_splits): + # Reshape for torch's scaled_dot_product_attention which expects (batch, heads, seq, dim). + # Here, we treat each sequence as a batch of 1. + q_i = q_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_q, head_dim) + k_i = k_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim) + v_i = v_i.permute(1, 0, 2).unsqueeze(0) # (1, heads, seq_len_k, head_dim) + + # Use PyTorch's built-in scaled dot-product attention. + output_i = F.scaled_dot_product_attention( + q_i, k_i, v_i, + dropout_p=dropout_p if not deterministic else 0.0, + is_causal=causal + ) + + # Reshape the output back to the original format (seq_len, heads, head_dim) + output_i = output_i.squeeze(0).permute(1, 0, 2) + output_splits.append(output_i) + + # Concatenate all outputs + return torch.cat(output_splits, dim=0) + + +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" + q = kwargs.get("query") or args[0] + k = kwargs.get("key") or args[1] + b, h, sq, d = q.shape + b, h, sk, d = k.shape + return b * h * (4 * d * (sq / 1e6) * (sk / 1e6)) + + def forward(self, *args, **kwargs): + return F.scaled_dot_product_attention(*args, **kwargs) + + +class FlashAttentionVarlen(nn.Module): + """ + Variable-length attention with configurable backend. + + Supported backends: + - sdpa: PyTorch SDPA (fully compilable, always available) + - flash_attn_2: Flash Attention 2 (Ampere+) + - 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): + """ + Initialize with specified attention backend. + + Args: + 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__() + self.attention_mode = attention_mode + self.compute_dtype = compute_dtype + + def tflops(self, args, kwargs, output) -> float: + cu_seqlens_q = kwargs["cu_seqlens_q"] + cu_seqlens_k = kwargs["cu_seqlens_k"] + _, h, d = output.shape + seqlens_q = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]) / 1e6 + seqlens_k = (cu_seqlens_k[1:] - cu_seqlens_k[:-1]) / 1e6 + return h * (4 * d * (seqlens_q * seqlens_k).sum()) + + def forward(self, q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, **kwargs): + kwargs["deterministic"] = torch.are_deterministic_algorithms_enabled() + + # Convert to pipeline compute_dtype if configured (handles FP8 โ†’ fp16/bf16) + if self.compute_dtype is not None and q.dtype != self.compute_dtype: + q = q.to(self.compute_dtype) + k = k.to(self.compute_dtype) + v = v.to(self.compute_dtype) + + if self.attention_mode == 'flash_attn_3': + 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( + q, k, v, cu_seqlens_q, cu_seqlens_k, + max_seqlen_q, max_seqlen_k, **kwargs + ) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..3195b400a407b871a6c19b67cf25239c5c3f196d --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/__init__.py @@ -0,0 +1,25 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from .mmdit_window_block import MMWindowTransformerBlock + +dit_blocks = { + "mmdit_window": MMWindowTransformerBlock, +} + + +def get_block(block_type: str): + if block_type in dit_blocks: + return dit_blocks[block_type] + raise NotImplementedError(f"{block_type} is not supported") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/mmdit_window_block.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/mmdit_window_block.py new file mode 100644 index 0000000000000000000000000000000000000000..ee9ae9ae8d34b94c56bc6dadc34ebdd99e00c5e3 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/blocks/mmdit_window_block.py @@ -0,0 +1,233 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Tuple, Union +import torch +from einops import rearrange +from torch import nn +from torch.nn import functional as F +from torch.nn.modules.utils import _triple +from ....common.half_precision_fixes import safe_pad_operation +from ....common.distributed.ops import ( + gather_heads, + gather_heads_scatter_seq, + gather_seq_scatter_heads_qkv, + scatter_heads, +) + +from ..attention import TorchAttention +from ..mlp import get_mlp +from ..mm import MMArg, MMModule +from ..modulation import ada_layer_type +from ..normalization import norm_layer_type +from ..rope import RotaryEmbedding3d + + +class MMWindowAttention(nn.Module): + def __init__( + self, + vid_dim: int, + txt_dim: int, + heads: int, + head_dim: int, + qk_bias: bool, + qk_rope: bool, + qk_norm: norm_layer_type, + qk_norm_eps: float, + window: Union[int, Tuple[int, int, int]], + window_method: str, + shared_qkv: bool, + ): + super().__init__() + dim = MMArg(vid_dim, txt_dim) + inner_dim = heads * head_dim + qkv_dim = inner_dim * 3 + + self.window = _triple(window) + self.window_method = window_method + assert all(map(lambda v: isinstance(v, int) and v >= 0, self.window)) + + self.head_dim = head_dim + self.proj_qkv = MMModule(nn.Linear, dim, qkv_dim, bias=qk_bias, shared_weights=shared_qkv) + self.proj_out = MMModule(nn.Linear, inner_dim, dim, shared_weights=shared_qkv) + self.norm_q = MMModule(qk_norm, dim=head_dim, eps=qk_norm_eps, elementwise_affine=True) + self.norm_k = MMModule(qk_norm, dim=head_dim, eps=qk_norm_eps, elementwise_affine=True) + self.rope = RotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None + self.attn = TorchAttention() + + def forward( + self, + vid: torch.FloatTensor, # b T H W c + txt: torch.FloatTensor, # b L c + txt_mask: torch.BoolTensor, # b L + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + # Project q, k, v. + vid_qkv, txt_qkv = self.proj_qkv(vid, txt) + vid_qkv = gather_seq_scatter_heads_qkv(vid_qkv, seq_dim=2) + _, T, H, W, _ = vid_qkv.shape + _, L, _ = txt.shape + + if self.window_method == "win": + nt, nh, nw = self.window + tt, hh, ww = T // nt, H // nh, W // nw + elif self.window_method == "win_by_size": + tt, hh, ww = self.window + tt, hh, ww = ( + tt if tt > 0 else T, + hh if hh > 0 else H, + ww if ww > 0 else W, + ) + nt, nh, nw = T // tt, H // hh, W // ww + else: + raise NotImplementedError + + vid_qkv = rearrange(vid_qkv, "b T H W (o h d) -> o b h (T H W) d", o=3, d=self.head_dim) + txt_qkv = rearrange(txt_qkv, "b L (o h d) -> o b h L d", o=3, d=self.head_dim) + txt_qkv = scatter_heads(txt_qkv, dim=2) + + vid_q, vid_k, vid_v = vid_qkv.unbind() + txt_q, txt_k, txt_v = txt_qkv.unbind() + + vid_q, txt_q = self.norm_q(vid_q, txt_q) + vid_k, txt_k = self.norm_k(vid_k, txt_k) + + if self.rope: + vid_q, vid_k = self.rope(vid_q, vid_k, (T, H, W)) + + def vid_window(v): + return rearrange( + v, + "b h (nt tt nh hh nw ww) d -> b h (nt nh nw) (tt hh ww) d", + hh=hh, + ww=ww, + tt=tt, + nh=nh, + nw=nw, + nt=nt, + ) + + def txt_window(t): + return rearrange(t, "b h L d -> b h 1 L d").expand(-1, -1, nt * nh * nw, -1, -1) + + # Process video attention. + vid_msk = safe_pad_operation(txt_mask, (tt * hh * ww, 0), value=True) + vid_msk = rearrange(vid_msk, "b l -> b 1 1 1 l").expand(-1, 1, 1, tt * hh * ww, -1) + vid_out = self.attn( + vid_window(vid_q), + torch.cat([vid_window(vid_k), txt_window(txt_k)], dim=-2), + torch.cat([vid_window(vid_v), txt_window(txt_v)], dim=-2), + vid_msk, + ) + vid_out = rearrange( + vid_out, + "b h (nt nh nw) (tt hh ww) d -> b (nt tt) (nh hh) (nw ww) (h d)", + hh=hh, + ww=ww, + tt=tt, + nh=nh, + nw=nw, + ) + vid_out = gather_heads_scatter_seq(vid_out, head_dim=4, seq_dim=2) + + # Process text attention. + txt_msk = safe_pad_operation(txt_mask, (T * H * W, 0), value=True) + txt_msk = rearrange(txt_msk, "b l -> b 1 1 l").expand(-1, 1, L, -1) + txt_out = self.attn( + txt_q, + torch.cat([vid_k, txt_k], dim=-2), + torch.cat([vid_v, txt_v], dim=-2), + txt_msk, + ) + txt_out = rearrange(txt_out, "b h L d -> b L (h d)") + txt_out = gather_heads(txt_out, dim=2) + + # Project output. + vid_out, txt_out = self.proj_out(vid_out, txt_out) + return vid_out, txt_out + + +class MMWindowTransformerBlock(nn.Module): + def __init__( + self, + *, + vid_dim: int, + txt_dim: int, + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm: norm_layer_type, + norm_eps: float, + ada: ada_layer_type, + qk_bias: bool, + qk_rope: bool, + qk_norm: norm_layer_type, + window: Union[int, Tuple[int, int, int]], + window_method: str, + shared_qkv: bool, + shared_mlp: bool, + mlp_type: str, + **kwargs, + ): + super().__init__() + dim = MMArg(vid_dim, txt_dim) + self.attn_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False) + self.attn = MMWindowAttention( + vid_dim=vid_dim, + txt_dim=txt_dim, + heads=heads, + head_dim=head_dim, + qk_bias=qk_bias, + qk_rope=qk_rope, + qk_norm=qk_norm, + qk_norm_eps=norm_eps, + window=window, + window_method=window_method, + shared_qkv=shared_qkv, + ) + self.mlp_norm = MMModule(norm, dim=dim, eps=norm_eps, elementwise_affine=False) + self.mlp = MMModule( + get_mlp(mlp_type), + dim=dim, + expand_ratio=expand_ratio, + shared_weights=shared_mlp, + ) + self.ada = MMModule(ada, dim=dim, emb_dim=emb_dim, layers=["attn", "mlp"]) + + def forward( + self, + vid: torch.FloatTensor, + txt: torch.FloatTensor, + txt_mask: torch.BoolTensor, + emb: torch.FloatTensor, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_attn, txt_attn = self.attn_norm(vid, txt) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, emb=emb, layer="attn", mode="in") + vid_attn, txt_attn = self.attn(vid_attn, txt_attn, txt_mask=txt_mask) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, emb=emb, layer="attn", mode="out") + vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt) + + vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, emb=emb, layer="mlp", mode="in") + vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, emb=emb, layer="mlp", mode="out") + vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn) + + return vid_mlp, txt_mlp diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/embedding.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..e972244f5767c9f34e5e77bb180ae720ce88b89c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/embedding.py @@ -0,0 +1,62 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Optional, Union +import torch +from diffusers.models.embeddings import get_timestep_embedding +from torch import nn + + +def emb_add(emb1: torch.Tensor, emb2: Optional[torch.Tensor]): + return emb1 if emb2 is None else emb1 + emb2 + + +class TimeEmbedding(nn.Module): + def __init__( + self, + sinusoidal_dim: int, + hidden_dim: int, + output_dim: int, + ): + super().__init__() + self.sinusoidal_dim = sinusoidal_dim + self.proj_in = nn.Linear(sinusoidal_dim, hidden_dim) + self.proj_hid = nn.Linear(hidden_dim, hidden_dim) + self.proj_out = nn.Linear(hidden_dim, output_dim) + self.act = nn.SiLU() + + def forward( + self, + timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], + device: torch.device, + dtype: torch.dtype, + ) -> torch.FloatTensor: + if not torch.is_tensor(timestep): + timestep = torch.tensor([timestep], device=device, dtype=dtype) + if timestep.ndim == 0: + timestep = timestep[None] + + emb = get_timestep_embedding( + timesteps=timestep, + embedding_dim=self.sinusoidal_dim, + flip_sin_to_cos=False, + downscale_freq_shift=0, + ) + emb = emb.to(dtype) + emb = self.proj_in(emb) + emb = self.act(emb) + emb = self.proj_hid(emb) + emb = self.act(emb) + emb = self.proj_out(emb) + return emb diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mlp.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mlp.py new file mode 100644 index 0000000000000000000000000000000000000000..2d05cb021f3e3c6ac05c0e7ae1aa8a6d29475b87 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mlp.py @@ -0,0 +1,62 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Optional +import torch +import torch.nn.functional as F +from torch import nn + + +def get_mlp(mlp_type: Optional[str] = "normal"): + if mlp_type == "normal": + return MLP + elif mlp_type == "swiglu": + return SwiGLUMLP + + +class MLP(nn.Module): + def __init__( + self, + dim: int, + expand_ratio: int, + ): + super().__init__() + self.proj_in = nn.Linear(dim, dim * expand_ratio) + self.act = nn.GELU("tanh") + self.proj_out = nn.Linear(dim * expand_ratio, dim) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + x = self.proj_in(x) + x = self.act(x) + x = self.proj_out(x) + return x + + +class SwiGLUMLP(nn.Module): + def __init__( + self, + dim: int, + expand_ratio: int, + multiple_of: int = 256, + ): + super().__init__() + hidden_dim = int(2 * dim * expand_ratio / 3) + hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + self.proj_in_gate = nn.Linear(dim, hidden_dim, bias=False) + self.proj_out = nn.Linear(hidden_dim, dim, bias=False) + self.proj_in = nn.Linear(dim, hidden_dim, bias=False) + + def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: + x = self.proj_out(F.silu(self.proj_in_gate(x)) * self.proj_in(x)) + return x diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mm.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mm.py new file mode 100644 index 0000000000000000000000000000000000000000..49be1f5915a61d8ea27f3e3718f35e5c9af662e7 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/mm.py @@ -0,0 +1,67 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Tuple +import torch +from torch import nn + + +@dataclass +class MMArg: + vid: Any + txt: Any + + +def get_args(key: str, args: List[Any]) -> List[Any]: + return [getattr(v, key) if isinstance(v, MMArg) else v for v in args] + + +def get_kwargs(key: str, kwargs: Dict[str, Any]) -> Dict[str, Any]: + return {k: getattr(v, key) if isinstance(v, MMArg) else v for k, v in kwargs.items()} + + +class MMModule(nn.Module): + def __init__( + self, + module: Callable[..., nn.Module], + *args, + shared_weights: bool = False, + **kwargs, + ): + super().__init__() + self.shared_weights = shared_weights + if self.shared_weights: + assert get_args("vid", args) == get_args("txt", args) + assert get_kwargs("vid", kwargs) == get_kwargs("txt", kwargs) + self.all = module(*get_args("vid", args), **get_kwargs("vid", kwargs)) + else: + self.vid = module(*get_args("vid", args), **get_kwargs("vid", kwargs)) + self.txt = module(*get_args("txt", args), **get_kwargs("txt", kwargs)) + + def forward( + self, + vid: torch.FloatTensor, + txt: torch.FloatTensor, + *args, + **kwargs, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + vid_module = self.vid if not self.shared_weights else self.all + txt_module = self.txt if not self.shared_weights else self.all + vid = vid_module(vid, *get_args("vid", args), **get_kwargs("vid", kwargs)) + txt = txt_module(txt, *get_args("txt", args), **get_kwargs("txt", kwargs)) + return vid, txt diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/modulation.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/modulation.py new file mode 100644 index 0000000000000000000000000000000000000000..38fff07ebdf08f5e12432176f1569225d7459900 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/modulation.py @@ -0,0 +1,112 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Callable, List, Optional +import torch +from einops import rearrange +from torch import nn + +from ...common.cache import Cache +from ...common.distributed.ops import slice_inputs + +# (dim: int, emb_dim: int) +ada_layer_type = Callable[[int, int], nn.Module] + + +def get_ada_layer(ada_layer: str) -> ada_layer_type: + if ada_layer == "single": + return AdaSingle + raise NotImplementedError(f"{ada_layer} is not supported") + + +def expand_dims(x: torch.Tensor, dim: int, ndim: int): + """ + Expand tensor "x" to "ndim" by adding empty dims at "dim". + Example: x is (b d), target ndim is 5, add dim at 1, return (b 1 1 1 d). + """ + shape = x.shape + shape = shape[:dim] + (1,) * (ndim - len(shape)) + shape[dim:] + return x.reshape(shape) + + +class AdaSingle(nn.Module): + def __init__( + self, + dim: int, + emb_dim: int, + layers: List[str], + ): + assert emb_dim == 6 * dim, "AdaSingle requires emb_dim == 6 * dim" + super().__init__() + self.dim = dim + self.emb_dim = emb_dim + self.layers = layers + for l in layers: + self.register_parameter(f"{l}_shift", nn.Parameter(torch.randn(dim) / dim**0.5)) + self.register_parameter(f"{l}_scale", nn.Parameter(torch.randn(dim) / dim**0.5 + 1)) + self.register_parameter(f"{l}_gate", nn.Parameter(torch.randn(dim) / dim**0.5)) + + def forward( + self, + hid: torch.FloatTensor, # b ... c + emb: torch.FloatTensor, # b d + layer: str, + mode: str, + cache: Cache = Cache(disable=True), + branch_tag: str = "", + hid_len: Optional[torch.LongTensor] = None, # b + ) -> torch.FloatTensor: + idx = self.layers.index(layer) + emb = rearrange(emb, "b (d l g) -> b d l g", l=len(self.layers), g=3)[..., idx, :] + emb = expand_dims(emb, 1, hid.ndim + 1) + + if hid_len is not None: + emb = cache( + f"emb_repeat_{idx}_{branch_tag}", + lambda: slice_inputs( + torch.repeat_interleave(emb, hid_len, dim=0), + dim=0, + ), + ) + + shiftA, scaleA, gateA = emb.unbind(-1) + shiftB, scaleB, gateB = ( + getattr(self, f"{layer}_shift"), + getattr(self, f"{layer}_scale"), + getattr(self, f"{layer}_gate"), + ) + + # Handle potential FP8 parameters - convert to input computation dtype + if hasattr(torch, 'float8_e4m3fn'): + fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2) + # Use input tensor's dtype as target (respects pipeline precision) + target_dtype = hid.dtype + + # Convert FP8 parameters to match input dtype for arithmetic operations + if shiftB is not None and shiftB.dtype in fp8_types: + shiftB = shiftB.to(target_dtype) + if scaleB is not None and scaleB.dtype in fp8_types: + scaleB = scaleB.to(target_dtype) + if gateB is not None and gateB.dtype in fp8_types: + gateB = gateB.to(target_dtype) + + if mode == "in": + return hid.mul_(scaleA + scaleB).add_(shiftA + shiftB) + if mode == "out": + return hid.mul_(gateA + gateB) + + raise NotImplementedError + + def extra_repr(self) -> str: + return f"dim={self.dim}, emb_dim={self.emb_dim}, layers={self.layers}" \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/na.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/na.py new file mode 100644 index 0000000000000000000000000000000000000000..26f17951a37015a6d33b39c3a51d37af2081b793 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/na.py @@ -0,0 +1,641 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +""" +Native Resolution Transformer (NaDiT) tensor manipulation utilities. + +TORCH.COMPILE OPTIMIZED VERSION +================================ +This module has been optimized for torch.compile compatibility by eliminating +data-dependent operations that cause graph breaks: + +Key Changes from Original: +- Replaced all .tolist() calls with pure tensor operations +- Minimized .item() calls (only used where required by einops API) +- Replaced list comprehensions with tensor-based splitting +- Added _tensor_split helper for compile-friendly splitting +- Proper device management to ensure tensors stay on correct devices +""" + +from itertools import chain +from typing import Callable, Dict, List, Tuple +import einops +import torch + + +def _tensor_split(tensor: torch.Tensor, lengths: torch.LongTensor, dim: int = 0) -> List[torch.Tensor]: + """ + Optimized compile-friendly split using torch.tensor_split. + + Uses PyTorch's native C++ implementation for fast eager mode execution + while remaining fully compatible with torch.compile symbolic tracing. + + Args: + tensor: Input tensor to split + lengths: Tensor of split lengths (1D) + dim: Dimension along which to split (default: 0) + + Returns: + List of split tensors + """ + if lengths.numel() == 0: + return [] + + if lengths.numel() == 1: + return [tensor] + + # Calculate split indices: torch.tensor_split splits BEFORE each index + # So we need cumsum[:-1] to get the split points + # NOTE: torch.tensor_split requires indices on CPU (PyTorch requirement) + split_indices = lengths.cumsum(0)[:-1].cpu() + + # Use torch.tensor_split - native C++ implementation + # Compilable: accepts tensor indices (symbolic shapes work) + # Fast: uses optimized CUDA/CPU kernels instead of Python loops + return list(torch.tensor_split(tensor, split_indices, dim=dim)) + + +def flatten( + hid: List[torch.FloatTensor], # List of (*** c) +) -> Tuple[ + torch.FloatTensor, # (L c) + torch.LongTensor, # (b n) +]: + """ + Flatten a list of tensors into a single tensor and track their shapes. + + Converts a list of tensors with potentially different spatial shapes but + same feature dimension into a flattened tensor and shape metadata. + + Args: + hid: List of tensors, each with shape (d1, d2, ..., dn, c) + + Returns: + Tuple of: + - Flattened tensor of shape (L, c) where L = sum of all spatial dimensions + - Shape tensor of shape (b, n) tracking original spatial dimensions + + COMPILE OPTIMIZATION: Uses tensor operations on correct device from input + """ + assert len(hid) > 0 + device = hid[0].device + + # Stack shape metadata - ensure tensors are created on correct device + shapes = [] + for x in hid: + shape_tensor = torch.tensor(x.shape[:-1], dtype=torch.long, device=device) + shapes.append(shape_tensor) + shape = torch.stack(shapes) + + # Flatten and concatenate + hid = torch.cat([x.flatten(0, -2) for x in hid]) + return hid, shape + + +def unflatten( + hid: torch.FloatTensor, # (L c) or (L ... c) + hid_shape: torch.LongTensor, # (b n) +) -> List[torch.Tensor]: # List of (*** c) or (*** ... c) + """ + Unflatten a tensor back to a list using shape metadata. + + Inverse operation of flatten(), reconstructing original tensor shapes. + + Args: + hid: Flattened tensor of shape (L, c) or (L, ..., c) + hid_shape: Shape metadata tensor of shape (b, n) + + Returns: + List of unflattened tensors with original spatial dimensions + + COMPILE OPTIMIZATION: + - Uses optimized _tensor_split with torch.tensor_split (fast) + - .cpu().numpy() conversion needed for torch.compile compatibility + (unflatten() requires concrete Python ints, not symbolic shapes) + """ + hid_len = hid_shape.prod(-1) + + # Use optimized tensor splitting (major performance improvement) + hid_list = _tensor_split(hid, hid_len, dim=0) + + # Unflatten each piece + # NOTE: .cpu().numpy() is required for torch.compile compatibility + # .tolist() would fail with symbolic shapes during compilation + result = [] + for i, x in enumerate(hid_list): + shape = hid_shape[i] + # Must use .cpu().numpy() for compilation compatibility + # Shape tensors are small, so CPU transfer overhead is minimal + target_shape = list(shape.cpu().numpy()) + result.append(x.unflatten(0, target_shape)) + + return result + + +def concat( + vid: torch.FloatTensor, # (VL ... c) + txt: torch.FloatTensor, # (TL ... c) + vid_len: torch.LongTensor, # (b) + txt_len: torch.LongTensor, # (b) +) -> torch.FloatTensor: # (L ... c) + """ + Interleave video and text tensors batch-wise. + + Splits video and text tensors by batch lengths, then interleaves them: + [vid_0, txt_0, vid_1, txt_1, ..., vid_b, txt_b] + + Args: + vid: Video features tensor (VL, c) + txt: Text features tensor (TL, c) + vid_len: Length of each video sequence (b,) + txt_len: Length of each text sequence (b,) + + Returns: + Interleaved tensor (L, c) where L = sum(vid_len) + sum(txt_len) + + COMPILE OPTIMIZATION: Uses _tensor_split for compile-friendly splitting + """ + # Use tensor-based splitting + vid_splits = _tensor_split(vid, vid_len, dim=0) + txt_splits = _tensor_split(txt, txt_len, dim=0) + + # Interleave + interleaved = [] + for v, t in zip(vid_splits, txt_splits): + interleaved.extend([v, t]) + + return torch.cat(interleaved) + + +def concat_idx( + vid_len: torch.LongTensor, # (b) + txt_len: torch.LongTensor, # (b) +) -> Tuple[ + Callable, + Callable, +]: + """ + Create index-based concatenation and un-concatenation functions. + + Pre-computes indices for efficient interleaving and de-interleaving operations. + Returns callable functions that can be reused multiple times. + + Args: + vid_len: Video sequence lengths (b,) + txt_len: Text sequence lengths (b,) + + Returns: + Tuple of (concat_fn, unconcat_fn): + - concat_fn: Lambda that interleaves vid and txt tensors + - unconcat_fn: Lambda that separates interleaved tensor back to vid and txt + + COMPILE OPTIMIZATION: Pre-computes all indices using tensor operations + """ + device = vid_len.device + vid_sum = vid_len.sum() + txt_sum = txt_len.sum() + + vid_idx = torch.arange(vid_sum, device=device) + txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device) + + # Build interleaving indices using compile-friendly _tensor_split + batch_size = len(vid_len) + vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0) + txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0) + + # Create interleaved target indices + tgt_idx_list = [] + for i in range(batch_size): + tgt_idx_list.append(vid_idx_splits[i]) + tgt_idx_list.append(txt_idx_splits[i]) + + tgt_idx = torch.cat(tgt_idx_list) + src_idx = torch.argsort(tgt_idx) + vid_idx_len = len(vid_idx) + + return ( + lambda vid, txt: torch.index_select(torch.cat([vid, txt]), 0, tgt_idx), + lambda all: torch.index_select(all, 0, src_idx).split([vid_idx_len, len(txt_idx)]), + ) + + +def unconcat( + all: torch.FloatTensor, # (L ... c) + vid_len: torch.LongTensor, # (b) + txt_len: torch.LongTensor, # (b) +) -> Tuple[ + torch.FloatTensor, # (VL ... c) + torch.FloatTensor, # (TL ... c) +]: + """ + De-interleave concatenated video and text tensors. + + Inverse of concat(). Separates an interleaved tensor back into video and text. + + Args: + all: Interleaved tensor (L, c) + vid_len: Video sequence lengths (b,) + txt_len: Text sequence lengths (b,) + + Returns: + Tuple of (vid, txt) tensors + + COMPILE OPTIMIZATION: Uses tensor operations to build interleave pattern + """ + batch_size = len(vid_len) + + # Create interleaved lengths: [vid_0, txt_0, vid_1, txt_1, ...] + interleave_len = torch.stack([vid_len, txt_len], dim=1).flatten() + + # Split using compile-friendly operation + all_splits = _tensor_split(all, interleave_len, dim=0) + + # Separate even (vid) and odd (txt) indices + vid_parts = [all_splits[i] for i in range(0, len(all_splits), 2)] + txt_parts = [all_splits[i] for i in range(1, len(all_splits), 2)] + + vid = torch.cat(vid_parts) + txt = torch.cat(txt_parts) + return vid, txt + + +def repeat_concat( + vid: torch.FloatTensor, # (VL ... c) + txt: torch.FloatTensor, # (TL ... c) + vid_len: torch.LongTensor, # (n*b) + txt_len: torch.LongTensor, # (b) + txt_repeat: torch.LongTensor, # (n) or (b) +) -> torch.FloatTensor: # (L ... c) + """ + Concatenate video and text with text repetition for window attention. + + For windowed attention, text features are repeated and interleaved with + multiple video windows: [vid_0, txt_0, vid_1, txt_0, vid_2, txt_0, ...] + + Args: + vid: Video features (VL, c) + txt: Text features (TL, c) + vid_len: Video window lengths (n*b,) where n=num_windows + txt_len: Text sequence lengths (b,) + txt_repeat: Number of times to repeat text (n,) or (b,) + + Returns: + Interleaved tensor with repeated text + + COMPILE OPTIMIZATION: Uses _tensor_split and tensor-based repetition + """ + # Split using compile-friendly operations + vid_splits = _tensor_split(vid, vid_len, dim=0) + txt_splits = _tensor_split(txt, txt_len, dim=0) + + # Handle txt_repeat shape flexibility + if txt_repeat.numel() == len(txt_splits): + repeat_counts = txt_repeat + else: + repeat_counts = txt_repeat.repeat(len(txt_splits)) + + # Interleave with repetition + result = [] + for i, v in enumerate(vid_splits): + result.append(v) + # Get corresponding text sample (cyclic) + batch_idx = i % len(txt_splits) if len(txt_splits) > 0 else 0 + if batch_idx < len(txt_splits): + result.append(txt_splits[batch_idx]) + + return torch.cat(result) + + +def repeat_concat_idx( + vid_len: torch.LongTensor, # (n*b) + txt_len: torch.LongTensor, # (b) + txt_repeat: torch.LongTensor, # (n) or scalar +) -> Tuple[ + Callable, + Callable, +]: + """ + Create index-based repeat-concatenation and un-concatenation with coalescing. + + Similar to concat_idx but handles text repetition for window attention. + The unconcat function coalesces (averages) repeated text features. + + Args: + vid_len: Video window lengths (n*b,) + txt_len: Text sequence lengths (b,) + txt_repeat: Repetition count (scalar or tensor) + + Returns: + Tuple of (concat_fn, unconcat_coalesce_fn): + - concat_fn: Interleaves with text repetition + - unconcat_coalesce_fn: Separates and averages repeated text + + Example: + Input: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10], repeat=3 + Concat: [0,1,2,9,10, 3,4,5,9,10, 6,7,8,9,10] + Unconcat: vid=[0,1,2,3,4,5,6,7,8], txt=[9,10] (averaged) + + COMPILE OPTIMIZATION: + - Eliminates .tolist() calls that caused graph breaks + - Uses pure tensor operations for index building + - Minimizes data-dependent branching + """ + device = vid_len.device + vid_sum = vid_len.sum() + txt_sum = txt_len.sum() + + # Create base indices + vid_idx = torch.arange(vid_sum, device=device) + txt_idx = torch.arange(vid_sum, vid_sum + txt_sum, device=device) + + # Normalize txt_repeat to tensor + if isinstance(txt_repeat, (int, float)): + txt_repeat = torch.tensor([txt_repeat], dtype=torch.long, device=device) + elif txt_repeat.dim() == 0: + txt_repeat = txt_repeat.unsqueeze(0) + + # Calculate repeat pattern - keep as tensor to avoid graph breaks + batch_size = len(txt_len) + if txt_repeat.numel() == 1: + num_repeats_tensor = txt_repeat.reshape(-1) # Ensure 1D tensor + else: + # Use tensor operations for division + num_repeats_tensor = torch.tensor([len(vid_len) // batch_size], dtype=torch.long, device=device) + + # Build concatenated indices using compile-friendly _tensor_split + vid_idx_splits = _tensor_split(vid_idx, vid_len, dim=0) + txt_idx_splits = _tensor_split(txt_idx, txt_len, dim=0) + + tgt_idx_list = [] + for i in range(len(vid_len)): + # Add video window + tgt_idx_list.append(vid_idx_splits[i]) + + # Add corresponding text (with repeat) + batch_idx = i % batch_size + tgt_idx_list.append(txt_idx_splits[batch_idx]) + + tgt_idx = torch.cat(tgt_idx_list) + src_idx = torch.argsort(tgt_idx) + txt_idx_len = len(tgt_idx) - len(vid_idx) + + # Pre-compute split lengths for coalescing using tensor operations + repeat_txt_len = txt_len * num_repeats_tensor.squeeze() + + def unconcat_coalesce(all): + """ + Un-concat vid & txt, and coalesce the repeated txt by averaging. + + 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 unflatten with tensor dims (compile-friendly) + """ + vid_out, txt_out = all[src_idx].split([len(vid_idx), txt_idx_len]) + + # Coalesce repeated text using unflatten and mean + txt_splits = _tensor_split(txt_out, repeat_txt_len, dim=0) + txt_out_coalesced = [] + + for txt in txt_splits: + # txt has shape (base_len * num_repeats, *other_dims) + # unflatten to (base_len, num_repeats, *other_dims) then average dim 1 + txt = txt.unflatten(0, (-1, num_repeats_tensor.squeeze())).mean(1) + txt_out_coalesced.append(txt) + + return vid_out, torch.cat(txt_out_coalesced) + + # Note: Using direct indexing instead of torch.index_select for backward compatibility + # Direct indexing is deterministic even with repeated indices + return ( + lambda vid, txt: torch.cat([vid, txt])[tgt_idx], + lambda all: unconcat_coalesce(all), + ) + + +def rearrange( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + pattern: str, + **kwargs: Dict[str, int], +) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, +]: + """ + Rearrange flattened tensor using einops pattern. + + Applies einops rearrange to each batch element independently. + + Args: + hid: Flattened tensor (L, c) + hid_shape: Shape metadata (b, n) + pattern: Einops rearrange pattern + **kwargs: Additional arguments for einops + + Returns: + Tuple of (rearranged tensor, new shape metadata) + """ + unflattened = unflatten(hid, hid_shape) + rearranged = [einops.rearrange(h, pattern, **kwargs) for h in unflattened] + return flatten(rearranged) + + +def rearrange_idx( + hid_shape: torch.LongTensor, # (b n) + pattern: str, + **kwargs: Dict[str, int], +) -> Tuple[Callable, Callable, torch.LongTensor]: + """ + Create index-based rearrange functions. + + Pre-computes indices for efficient rearrangement operations. + + Args: + hid_shape: Shape metadata (b, n) + pattern: Einops rearrange pattern + **kwargs: Additional arguments for einops + + Returns: + Tuple of (rearrange_fn, reverse_fn, new_shape) + """ + hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1) + tgt_idx, tgt_shape = rearrange(hid_idx, hid_shape, pattern, **kwargs) + tgt_idx = tgt_idx.squeeze(-1) + src_idx = torch.argsort(tgt_idx) + return ( + lambda hid: torch.index_select(hid, 0, tgt_idx), + lambda hid: torch.index_select(hid, 0, src_idx), + tgt_shape, + ) + + +def repeat( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + pattern: str, + **kwargs: Dict[str, torch.LongTensor], # (b) +) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, +]: + """ + Repeat flattened tensor using einops pattern with per-batch parameters. + + Each batch element can have different repeat counts specified in kwargs. + + Args: + hid: Flattened tensor (L, c) + hid_shape: Shape metadata (b, n) + pattern: Einops repeat pattern (e.g., "l c -> t l c") + **kwargs: Repeat parameters as tensors (e.g., t=torch.tensor([2,3,4])) + + Returns: + Tuple of (repeated tensor, new shape metadata) + + COMPILE OPTIMIZATION: + - Minimizes .item() calls + - Only converts to int at the last moment for einops API requirement + """ + unflattened = unflatten(hid, hid_shape) + + # Build kwargs for each batch element + repeated = [] + for i in range(len(unflattened)): + # Extract values for einops (requires Python int) + batch_kwargs = {} + for k, v in kwargs.items(): + if isinstance(v, torch.Tensor): + # Only convert to Python int where required by einops + batch_kwargs[k] = int(v[i].item()) + else: + batch_kwargs[k] = v + repeated.append(einops.repeat(unflattened[i], pattern, **batch_kwargs)) + + return flatten(repeated) + + +def pack( + samples: List[torch.Tensor], # List of (h w c). +) -> Tuple[ + List[torch.Tensor], # groups [(b1 h1 w1 c1), (b2 h2 w2 c2)] + List[List[int]], # reversal indices. +]: + """ + Group samples by shape and return grouped batches with reversal indices. + + Useful for batch processing samples with different spatial dimensions. + + Args: + samples: List of tensors with potentially different shapes + + Returns: + Tuple of (batched_groups, reversal_indices) for unpacking + """ + batches = {} + indices = {} + for i, sample in enumerate(samples): + shape = sample.shape + batches[shape] = batches.get(shape, []) + indices[shape] = indices.get(shape, []) + batches[shape].append(sample) + indices[shape].append(i) + + batches = list(map(torch.stack, batches.values())) + indices = list(indices.values()) + return batches, indices + + +def unpack( + batches: List[torch.Tensor], + indices: List[List[int]], +) -> List[torch.Tensor]: + """ + Unpack grouped batches back to original order. + + Inverse of pack(). + + Args: + batches: Grouped batches from pack() + indices: Reversal indices from pack() + + Returns: + List of tensors in original order + """ + samples = [None] * (max(chain(*indices)) + 1) + for batch, index in zip(batches, indices): + for sample, i in zip(batch.unbind(), index): + samples[i] = sample + return samples + + +def window( + hid: torch.FloatTensor, # (L c) + hid_shape: torch.LongTensor, # (b n) + window_fn: Callable[[torch.Tensor], List[torch.Tensor]], +): + """ + Apply windowing function to create non-overlapping windows. + + Used for window attention mechanisms where sequences are split into windows. + + Args: + hid: Flattened tensor (L, c) + hid_shape: Shape metadata (b, n) + window_fn: Function that splits a tensor into windows + + Returns: + Tuple of (windowed_tensor, window_shapes, window_counts) + + COMPILE OPTIMIZATION: Uses tensor operation for window count tracking + """ + unflattened = unflatten(hid, hid_shape) + windowed = [window_fn(h) for h in unflattened] + + # Track window counts using tensor operations + device = hid_shape.device + hid_windows = torch.tensor([len(w) for w in windowed], dtype=torch.long, device=device) + + # Flatten all windows + all_windows = list(chain(*windowed)) + hid, hid_shape = flatten(all_windows) + return hid, hid_shape, hid_windows + + +def window_idx( + hid_shape: torch.LongTensor, # (b n) + window_fn: Callable[[torch.Tensor], List[torch.Tensor]], +): + """ + Create index-based windowing functions. + + Pre-computes indices for efficient windowing and reverse operations. + + Args: + hid_shape: Shape metadata (b, n) + window_fn: Function that splits a tensor into windows + + Returns: + Tuple of (window_fn, reverse_fn, window_shapes, window_counts) + """ + hid_idx = torch.arange(hid_shape.prod(-1).sum(), device=hid_shape.device).unsqueeze(-1) + tgt_idx, tgt_shape, tgt_windows = window(hid_idx, hid_shape, window_fn) + tgt_idx = tgt_idx.squeeze(-1) + src_idx = torch.argsort(tgt_idx) + return ( + lambda hid: torch.index_select(hid, 0, tgt_idx), + lambda hid: torch.index_select(hid, 0, src_idx), + tgt_shape, + tgt_windows, + ) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..afa206db157786d9e4cf830bec09bd3a390bd9a8 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/__init__.py @@ -0,0 +1,25 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from .mmsr_block import NaMMSRTransformerBlock + +nadit_blocks = { + "mmdit_sr": NaMMSRTransformerBlock, +} + + +def get_nablock(block_type: str): + if block_type in nadit_blocks: + return nadit_blocks[block_type] + raise NotImplementedError(f"{block_type} is not supported") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/mmsr_block.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/mmsr_block.py new file mode 100644 index 0000000000000000000000000000000000000000..fe9010ac688287a6f072413a9cd1820b27ecf193 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nablocks/mmsr_block.py @@ -0,0 +1,250 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Tuple, Union +import torch +from einops import rearrange +from torch.nn import functional as F + +# from ..cache import Cache +from ....common.cache import Cache +from ....common.distributed.ops import gather_heads_scatter_seq, gather_seq_scatter_heads_qkv + +from .. import na +from ..attention import FlashAttentionVarlen +from ..blocks.mmdit_window_block import MMWindowAttention, MMWindowTransformerBlock +from ..mm import MMArg +from ..modulation import ada_layer_type +from ..normalization import norm_layer_type +from ..rope import NaRotaryEmbedding3d +from ..window import get_window_op +from ....common.half_precision_fixes import safe_pad_operation + +class NaSwinAttention(MMWindowAttention): + def __init__( + self, + vid_dim: int, + txt_dim: int, + heads: int, + head_dim: int, + qk_bias: bool, + qk_rope: bool, + qk_norm: norm_layer_type, + qk_norm_eps: float, + window: Union[int, Tuple[int, int, int]], + window_method: str, + shared_qkv: bool, + attention_mode: str = 'sdpa', + **kwargs, + ): + super().__init__( + vid_dim=vid_dim, + txt_dim=txt_dim, + heads=heads, + head_dim=head_dim, + qk_bias=qk_bias, + qk_rope=qk_rope, + qk_norm=qk_norm, + qk_norm_eps=qk_norm_eps, + window=window, + window_method=window_method, + shared_qkv=shared_qkv, + ) + self.rope = NaRotaryEmbedding3d(dim=head_dim // 2) if qk_rope else None + self.attn = FlashAttentionVarlen(attention_mode=attention_mode) + self.window_op = get_window_op(window_method) + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + + vid_qkv, txt_qkv = self.proj_qkv(vid, txt) + vid_qkv = gather_seq_scatter_heads_qkv( + vid_qkv, + seq_dim=0, + qkv_shape=vid_shape, + cache=cache.namespace("vid"), + ) + txt_qkv = gather_seq_scatter_heads_qkv( + txt_qkv, + seq_dim=0, + qkv_shape=txt_shape, + cache=cache.namespace("txt"), + ) + + # re-org the input seq for window attn + cache_win = cache.namespace(f"{self.window_method}_{self.window}_sd3") + + def make_window(x: torch.Tensor): + t, h, w, _ = x.shape + window_slices = self.window_op((t, h, w), self.window) + return [x[st, sh, sw] for (st, sh, sw) in window_slices] + + window_partition, window_reverse, window_shape, window_count = cache_win( + "win_transform", + lambda: na.window_idx(vid_shape, make_window), + ) + vid_qkv_win = window_partition(vid_qkv) + + vid_qkv_win = rearrange(vid_qkv_win, "l (o h d) -> l o h d", o=3, d=self.head_dim) + txt_qkv = rearrange(txt_qkv, "l (o h d) -> l o h d", o=3, d=self.head_dim) + + vid_q, vid_k, vid_v = vid_qkv_win.unbind(1) + txt_q, txt_k, txt_v = txt_qkv.unbind(1) + + vid_q, txt_q = self.norm_q(vid_q, txt_q) + vid_k, txt_k = self.norm_k(vid_k, txt_k) + + txt_len = cache("txt_len", lambda: txt_shape.prod(-1)) + + vid_len_win = cache_win("vid_len", lambda: window_shape.prod(-1)) + txt_len_win = cache_win("txt_len", lambda: txt_len.repeat_interleave(window_count)) + all_len_win = cache_win("all_len", lambda: vid_len_win + txt_len_win) + concat_win, unconcat_win = cache_win( + "mm_pnp", lambda: na.repeat_concat_idx(vid_len_win, txt_len, window_count) + ) + + # window rope + if self.rope: + vid_q, vid_k = self.rope(vid_q, vid_k, window_shape, cache_win) + + # Attention handles dtype conversion internally using pipeline compute_dtype + out = self.attn( + q=concat_win(vid_q, txt_q), + k=concat_win(vid_k, txt_k), + v=concat_win(vid_v, txt_v), + cu_seqlens_q=cache_win( + "vid_seqlens_q", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int() + ), + cu_seqlens_k=cache_win( + "vid_seqlens_k", lambda: safe_pad_operation(all_len_win.cumsum(0), (1, 0)).int() + ), + max_seqlen_q=cache_win("vid_max_seqlen_q", lambda: all_len_win.max()), + max_seqlen_k=cache_win("vid_max_seqlen_k", lambda: all_len_win.max()), + ).type_as(vid_q) + + # text pooling + vid_out, txt_out = unconcat_win(out) + + vid_out = rearrange(vid_out, "l h d -> l (h d)") + txt_out = rearrange(txt_out, "l h d -> l (h d)") + vid_out = window_reverse(vid_out) + + vid_out = gather_heads_scatter_seq(vid_out, head_dim=1, seq_dim=0) + txt_out = gather_heads_scatter_seq(txt_out, head_dim=1, seq_dim=0) + + vid_out, txt_out = self.proj_out(vid_out, txt_out) + + return vid_out, txt_out + + +class NaMMSRTransformerBlock(MMWindowTransformerBlock): + def __init__( + self, + *, + vid_dim: int, + txt_dim: int, + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm: norm_layer_type, + norm_eps: float, + ada: ada_layer_type, + qk_bias: bool, + qk_rope: bool, + qk_norm: norm_layer_type, + shared_qkv: bool, + shared_mlp: bool, + mlp_type: str, + **kwargs, + ): + super().__init__( + vid_dim=vid_dim, + txt_dim=txt_dim, + emb_dim=emb_dim, + heads=heads, + head_dim=head_dim, + expand_ratio=expand_ratio, + norm=norm, + norm_eps=norm_eps, + ada=ada, + qk_bias=qk_bias, + qk_rope=qk_rope, + qk_norm=qk_norm, + shared_qkv=shared_qkv, + shared_mlp=shared_mlp, + mlp_type=mlp_type, + **kwargs, + ) + + self.attn = NaSwinAttention( + vid_dim=vid_dim, + txt_dim=txt_dim, + heads=heads, + head_dim=head_dim, + qk_bias=qk_bias, + qk_rope=qk_rope, + qk_norm=qk_norm, + qk_norm_eps=norm_eps, + shared_qkv=shared_qkv, + **kwargs, + ) + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + emb: torch.FloatTensor, + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + torch.LongTensor, + torch.LongTensor, + ]: + hid_len = MMArg( + cache("vid_len", lambda: vid_shape.prod(-1)), + cache("txt_len", lambda: txt_shape.prod(-1)), + ) + ada_kwargs = { + "emb": emb, + "hid_len": hid_len, + "cache": cache, + "branch_tag": MMArg("vid", "txt"), + } + + vid_attn, txt_attn = self.attn_norm(vid, txt) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="in", **ada_kwargs) + vid_attn, txt_attn = self.attn(vid_attn, txt_attn, vid_shape, txt_shape, cache) + vid_attn, txt_attn = self.ada(vid_attn, txt_attn, layer="attn", mode="out", **ada_kwargs) + vid_attn, txt_attn = (vid_attn + vid), (txt_attn + txt) + + vid_mlp, txt_mlp = self.mlp_norm(vid_attn, txt_attn) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="in", **ada_kwargs) + vid_mlp, txt_mlp = self.mlp(vid_mlp, txt_mlp) + vid_mlp, txt_mlp = self.ada(vid_mlp, txt_mlp, layer="mlp", mode="out", **ada_kwargs) + vid_mlp, txt_mlp = (vid_mlp + vid_attn), (txt_mlp + txt_attn) + + return vid_mlp, txt_mlp, vid_shape, txt_shape diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nadit.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nadit.py new file mode 100644 index 0000000000000000000000000000000000000000..37eb1a46813ec12c09e5a0cdbeff69a9857f714c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/nadit.py @@ -0,0 +1,353 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from dataclasses import dataclass +from typing import Optional, Tuple, Union, Callable +import torch +from torch import nn + +from ...common.cache import Cache +from ...common.distributed.ops import slice_inputs + +from . import na +from .embedding import TimeEmbedding +from .modulation import get_ada_layer +from .nablocks import get_nablock +from .normalization import get_norm_layer +from .patch import NaPatchIn, NaPatchOut + +# Fake func, no checkpointing is required for inference +def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs): + return module(*args, **kwargs) + +@dataclass +class NaDiTOutput: + vid_sample: torch.Tensor + + +class NaDiT(nn.Module): + """ + Native Resolution Diffusion Transformer (NaDiT) + """ + + gradient_checkpointing = False + + def __init__( + self, + vid_in_channels: int, + vid_out_channels: int, + vid_dim: int, + txt_in_dim: Optional[int], + txt_dim: Optional[int], + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm: Optional[str], + norm_eps: float, + ada: str, + qk_bias: bool, + qk_rope: bool, + qk_norm: Optional[str], + patch_size: Union[int, Tuple[int, int, int]], + num_layers: int, + block_type: Union[str, Tuple[str]], + shared_qkv: bool = False, + shared_mlp: bool = False, + mlp_type: str = "normal", + window: Optional[Tuple] = None, + window_method: Optional[Tuple[str]] = None, + temporal_window_size: int = None, + temporal_shifted: bool = False, + attention_mode: str = 'sdpa', + **kwargs, + ): + ada = get_ada_layer(ada) + norm = get_norm_layer(norm) + qk_norm = get_norm_layer(qk_norm) + if isinstance(block_type, str): + block_type = [block_type] * num_layers + elif len(block_type) != num_layers: + raise ValueError("The ``block_type`` list should equal to ``num_layers``.") + super().__init__() + self.vid_in = NaPatchIn( + in_channels=vid_in_channels, + patch_size=patch_size, + dim=vid_dim, + ) + self.txt_in = ( + nn.Linear(txt_in_dim, txt_dim) + if txt_in_dim and txt_in_dim != txt_dim + else nn.Identity() + ) + self.emb_in = TimeEmbedding( + sinusoidal_dim=256, + hidden_dim=max(vid_dim, txt_dim), + output_dim=emb_dim, + ) + + if window is None or isinstance(window[0], int): + window = [window] * num_layers + if window_method is None or isinstance(window_method, str): + window_method = [window_method] * num_layers + if temporal_window_size is None or isinstance(temporal_window_size, int): + temporal_window_size = [temporal_window_size] * num_layers + if temporal_shifted is None or isinstance(temporal_shifted, bool): + temporal_shifted = [temporal_shifted] * num_layers + + self.blocks = nn.ModuleList( + [ + get_nablock(block_type[i])( + vid_dim=vid_dim, + txt_dim=txt_dim, + emb_dim=emb_dim, + heads=heads, + head_dim=head_dim, + expand_ratio=expand_ratio, + norm=norm, + norm_eps=norm_eps, + ada=ada, + qk_bias=qk_bias, + qk_rope=qk_rope, + qk_norm=qk_norm, + shared_qkv=shared_qkv, + shared_mlp=shared_mlp, + mlp_type=mlp_type, + window=window[i], + window_method=window_method[i], + temporal_window_size=temporal_window_size[i], + temporal_shifted=temporal_shifted[i], + attention_mode=attention_mode, + **kwargs, + ) + for i in range(num_layers) + ] + ) + self.vid_out = NaPatchOut( + out_channels=vid_out_channels, + patch_size=patch_size, + dim=vid_dim, + ) + + self.need_txt_repeat = block_type[0] in [ + "mmdit_stwin", + "mmdit_stwin_spatial", + "mmdit_stwin_3d_spatial", + ] + + def set_gradient_checkpointing(self, enable: bool): + self.gradient_checkpointing = enable + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], # b + disable_cache: bool = True, # for test + ): + # Text input. + if txt_shape.size(-1) == 1 and self.need_txt_repeat: + txt, txt_shape = na.repeat(txt, txt_shape, "l c -> t l c", t=vid_shape[:, 0]) + # slice vid after patching in when using sequence parallelism + txt = slice_inputs(txt, dim=0) + txt = self.txt_in(txt) + + # Video input. + # Sequence parallel slicing is done inside patching class. + vid, vid_shape = self.vid_in(vid, vid_shape) + + # Embedding input. + emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype) + + # Body + cache = Cache(disable=disable_cache) + for i, block in enumerate(self.blocks): + vid, txt, vid_shape, txt_shape = gradient_checkpointing( + enabled=(self.gradient_checkpointing and self.training), + module=block, + vid=vid, + txt=txt, + vid_shape=vid_shape, + txt_shape=txt_shape, + emb=emb, + cache=cache, + ) + + vid, vid_shape = self.vid_out(vid, vid_shape, cache) + return NaDiTOutput(vid_sample=vid) + + +class NaDiTUpscaler(nn.Module): + """ + Native Resolution Diffusion Transformer (NaDiT) + """ + + gradient_checkpointing = False + + def __init__( + self, + vid_in_channels: int, + vid_out_channels: int, + vid_dim: int, + txt_in_dim: Optional[int], + txt_dim: Optional[int], + emb_dim: int, + heads: int, + head_dim: int, + expand_ratio: int, + norm: Optional[str], + norm_eps: float, + ada: str, + qk_bias: bool, + qk_rope: bool, + qk_norm: Optional[str], + patch_size: Union[int, Tuple[int, int, int]], + num_layers: int, + block_type: Union[str, Tuple[str]], + shared_qkv: bool = False, + shared_mlp: bool = False, + mlp_type: str = "normal", + window: Optional[Tuple] = None, + window_method: Optional[Tuple[str]] = None, + temporal_window_size: int = None, + temporal_shifted: bool = False, + **kwargs, + ): + ada = get_ada_layer(ada) + norm = get_norm_layer(norm) + qk_norm = get_norm_layer(qk_norm) + if isinstance(block_type, str): + block_type = [block_type] * num_layers + elif len(block_type) != num_layers: + raise ValueError("The ``block_type`` list should equal to ``num_layers``.") + super().__init__() + self.vid_in = NaPatchIn( + in_channels=vid_in_channels, + patch_size=patch_size, + dim=vid_dim, + ) + self.txt_in = ( + nn.Linear(txt_in_dim, txt_dim) + if txt_in_dim and txt_in_dim != txt_dim + else nn.Identity() + ) + self.emb_in = TimeEmbedding( + sinusoidal_dim=256, + hidden_dim=max(vid_dim, txt_dim), + output_dim=emb_dim, + ) + + self.emb_scale = TimeEmbedding( + sinusoidal_dim=256, + hidden_dim=max(vid_dim, txt_dim), + output_dim=emb_dim, + ) + + if window is None or isinstance(window[0], int): + window = [window] * num_layers + if window_method is None or isinstance(window_method, str): + window_method = [window_method] * num_layers + if temporal_window_size is None or isinstance(temporal_window_size, int): + temporal_window_size = [temporal_window_size] * num_layers + if temporal_shifted is None or isinstance(temporal_shifted, bool): + temporal_shifted = [temporal_shifted] * num_layers + + self.blocks = nn.ModuleList( + [ + get_nablock(block_type[i])( + vid_dim=vid_dim, + txt_dim=txt_dim, + emb_dim=emb_dim, + heads=heads, + head_dim=head_dim, + expand_ratio=expand_ratio, + norm=norm, + norm_eps=norm_eps, + ada=ada, + qk_bias=qk_bias, + qk_rope=qk_rope, + qk_norm=qk_norm, + shared_qkv=shared_qkv, + shared_mlp=shared_mlp, + mlp_type=mlp_type, + window=window[i], + window_method=window_method[i], + temporal_window_size=temporal_window_size[i], + temporal_shifted=temporal_shifted[i], + attention_mode=attention_mode, + **kwargs, + ) + for i in range(num_layers) + ] + ) + self.vid_out = NaPatchOut( + out_channels=vid_out_channels, + patch_size=patch_size, + dim=vid_dim, + ) + + self.need_txt_repeat = block_type[0] in [ + "mmdit_stwin", + "mmdit_stwin_spatial", + "mmdit_stwin_3d_spatial", + ] + + def set_gradient_checkpointing(self, enable: bool): + self.gradient_checkpointing = enable + + def forward( + self, + vid: torch.FloatTensor, # l c + txt: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, # b 3 + txt_shape: torch.LongTensor, # b 1 + timestep: Union[int, float, torch.IntTensor, torch.FloatTensor], # b + downscale: Union[int, float, torch.IntTensor, torch.FloatTensor], # b + disable_cache: bool = False, # for test + ): + + # Text input. + if txt_shape.size(-1) == 1 and self.need_txt_repeat: + txt, txt_shape = na.repeat(txt, txt_shape, "l c -> t l c", t=vid_shape[:, 0]) + # slice vid after patching in when using sequence parallelism + txt = slice_inputs(txt, dim=0) + txt = self.txt_in(txt) + + # Video input. + # Sequence parallel slicing is done inside patching class. + vid, vid_shape = self.vid_in(vid, vid_shape) + + # Embedding input. + emb = self.emb_in(timestep, device=vid.device, dtype=vid.dtype) + emb_scale = self.emb_scale(downscale, device=vid.device, dtype=vid.dtype) + emb = emb + emb_scale + + # Body + cache = Cache(disable=disable_cache) + for i, block in enumerate(self.blocks): + vid, txt, vid_shape, txt_shape = gradient_checkpointing( + enabled=(self.gradient_checkpointing and self.training), + module=block, + vid=vid, + txt=txt, + vid_shape=vid_shape, + txt_shape=txt_shape, + emb=emb, + cache=cache, + ) + + vid, vid_shape = self.vid_out(vid, vid_shape, cache) + return NaDiTOutput(vid_sample=vid) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/normalization.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..c34513ad3b5f481845bb29b4f64b6b4029b225da --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/normalization.py @@ -0,0 +1,139 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Callable, Optional +from diffusers.models.normalization import RMSNorm +from torch import nn +import torch +import torch.nn.functional as F +import numbers +from torch.nn.parameter import Parameter +from torch.nn import init + +# (dim: int, eps: float, elementwise_affine: bool) +norm_layer_type = Callable[[int, float, bool], nn.Module] + + +class CustomLayerNorm(nn.Module): + """ + Custom LayerNorm implementation to replace Apex FusedLayerNorm + """ + def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True): + super(CustomLayerNorm, self).__init__() + + if isinstance(normalized_shape, numbers.Integral): + normalized_shape = (normalized_shape,) + self.normalized_shape = torch.Size(normalized_shape) + self.eps = eps + self.elementwise_affine = elementwise_affine + + if self.elementwise_affine: + self.weight = Parameter(torch.Tensor(*normalized_shape)) + self.bias = Parameter(torch.Tensor(*normalized_shape)) + else: + self.register_parameter('weight', None) + self.register_parameter('bias', None) + self.reset_parameters() + + def reset_parameters(self): + if self.elementwise_affine: + init.ones_(self.weight) + init.zeros_(self.bias) + + def forward(self, input): + return F.layer_norm( + input, self.normalized_shape, self.weight, self.bias, self.eps) + + +class CustomRMSNorm(nn.Module): + """ + Custom RMSNorm implementation to replace Apex FusedRMSNorm + """ + def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True): + super(CustomRMSNorm, self).__init__() + + if isinstance(normalized_shape, numbers.Integral): + normalized_shape = (normalized_shape,) + self.normalized_shape = torch.Size(normalized_shape) + self.eps = eps + self.elementwise_affine = elementwise_affine + + if self.elementwise_affine: + self.weight = Parameter(torch.ones(*normalized_shape)) + else: + self.register_parameter('weight', None) + + def forward(self, input): + # RMS normalization: x / sqrt(mean(x^2) + eps) * weight + dims = tuple(range(-len(self.normalized_shape), 0)) + + # Calculate RMS: sqrt(mean(x^2)) + variance = input.pow(2).mean(dim=dims, keepdim=True) + rms = torch.sqrt(variance + self.eps) + + # Normalize + normalized = input / rms + + if self.elementwise_affine: + # Convert FP8 weight to match input dtype for arithmetic operations + if hasattr(torch, 'float8_e4m3fn'): + fp8_types = (torch.float8_e4m3fn, torch.float8_e5m2) + if self.weight.dtype in fp8_types: + # Use input dtype as target (respects pipeline precision) + weight = self.weight.to(input.dtype) + return normalized * weight + + return normalized * self.weight + return normalized + + +def get_norm_layer(norm_type: Optional[str]) -> norm_layer_type: + + def _norm_layer(dim: int, eps: float, elementwise_affine: bool): + if norm_type is None: + return nn.Identity() + + if norm_type == "layer": + return nn.LayerNorm( + normalized_shape=dim, + eps=eps, + elementwise_affine=elementwise_affine, + ) + + if norm_type == "rms": + return RMSNorm( + dim=dim, + eps=eps, + elementwise_affine=elementwise_affine, + ) + + if norm_type == "fusedln": + # Use custom LayerNorm instead of Apex FusedLayerNorm + return CustomLayerNorm( + normalized_shape=dim, + elementwise_affine=elementwise_affine, + eps=eps, + ) + + if norm_type == "fusedrms": + # Use custom RMSNorm instead of Apex FusedRMSNorm + return CustomRMSNorm( + normalized_shape=dim, + elementwise_affine=elementwise_affine, + eps=eps, + ) + + raise NotImplementedError(f"{norm_type} is not supported") + + return _norm_layer \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/patch.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/patch.py new file mode 100644 index 0000000000000000000000000000000000000000..570e3ae40a743fd434d1dd34e14dc6cc4192783e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/patch.py @@ -0,0 +1,112 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Tuple, Union +import torch +from einops import rearrange +from torch import nn +from torch.nn.modules.utils import _triple + +from ...common.cache import Cache +from ...common.distributed.ops import gather_outputs, slice_inputs + +from . import na + + +class PatchIn(nn.Module): + def __init__( + self, + in_channels: int, + patch_size: Union[int, Tuple[int, int, int]], + dim: int, + ): + super().__init__() + t, h, w = _triple(patch_size) + self.patch_size = t, h, w + self.proj = nn.Linear(in_channels * t * h * w, dim) + + def forward( + self, + vid: torch.Tensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + vid = rearrange(vid, "b c (T t) (H h) (W w) -> b T H W (t h w c)", t=t, h=h, w=w) + vid = self.proj(vid) + return vid + + +class PatchOut(nn.Module): + def __init__( + self, + out_channels: int, + patch_size: Union[int, Tuple[int, int, int]], + dim: int, + ): + super().__init__() + t, h, w = _triple(patch_size) + self.patch_size = t, h, w + self.proj = nn.Linear(dim, out_channels * t * h * w) + + def forward( + self, + vid: torch.Tensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + vid = self.proj(vid) + vid = rearrange(vid, "b T H W (t h w c) -> b c (T t) (H h) (W w)", t=t, h=h, w=w) + return vid + + +class NaPatchIn(PatchIn): + def forward( + self, + vid: torch.Tensor, # l c + vid_shape: torch.LongTensor, + ) -> torch.Tensor: + t, h, w = self.patch_size + if not (t == h == w == 1): + vid, vid_shape = na.rearrange( + vid, vid_shape, "(T t) (H h) (W w) c -> T H W (t h w c)", t=t, h=h, w=w + ) + # slice vid after patching in when using sequence parallelism + vid = slice_inputs(vid, dim=0) + vid = self.proj(vid) + return vid, vid_shape + + +class NaPatchOut(PatchOut): + def forward( + self, + vid: torch.FloatTensor, # l c + vid_shape: torch.LongTensor, + cache: Cache = Cache(disable=True), + ) -> Tuple[ + torch.FloatTensor, + torch.LongTensor, + ]: + t, h, w = self.patch_size + vid = self.proj(vid) + # gather vid before patching out when enabling sequence parallelism + vid = gather_outputs( + vid, + gather_dim=0, + padding_dim=0, + unpad_shape=vid_shape, + cache=cache.namespace("vid"), + ) + if not (t == h == w == 1): + vid, vid_shape = na.rearrange( + vid, vid_shape, "T H W (t h w c) -> (T t) (H h) (W w) c", t=t, h=h, w=w + ) + return vid, vid_shape diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/rope.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/rope.py new file mode 100644 index 0000000000000000000000000000000000000000..68a91355c6cff163aa578b845b7eddd8efe9a295 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/rope.py @@ -0,0 +1,111 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from functools import lru_cache +from typing import Tuple +import torch +from einops import rearrange +from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb +from torch import nn + +from ...common.cache import Cache + + +class RotaryEmbeddingBase(nn.Module): + def __init__(self, dim: int, rope_dim: int): + super().__init__() + self.rope = RotaryEmbedding( + dim=dim // rope_dim, + freqs_for="pixel", + max_freq=256, + ) + # 1. Set model.requires_grad_(True) after model creation will make + # the `requires_grad=False` for rope freqs no longer hold. + # 2. Even if we don't set requires_grad_(True) explicitly, + # FSDP is not memory efficient when handling fsdp_wrap + # with mixed requires_grad=True/False. + # With above consideration, it is easier just remove the freqs + # out of nn.Parameters when `learned_freq=False` + freqs = self.rope.freqs + del self.rope.freqs + self.rope.register_buffer("freqs", freqs.data) + + @lru_cache(maxsize=128) + def get_axial_freqs(self, *dims): + return self.rope.get_axial_freqs(*dims) + + +class RotaryEmbedding3d(RotaryEmbeddingBase): + def __init__(self, dim: int): + super().__init__(dim, rope_dim=3) + + def forward( + self, + q: torch.FloatTensor, # b h l d + k: torch.FloatTensor, # b h l d + size: Tuple[int, int, int], + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + T, H, W = size + freqs = self.get_axial_freqs(T, H, W) + q = rearrange(q, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W) + k = rearrange(k, "b h (T H W) d -> b h T H W d", T=T, H=H, W=W) + q = apply_rotary_emb(freqs, q) + k = apply_rotary_emb(freqs, k) + q = rearrange(q, "b h T H W d -> b h (T H W) d") + k = rearrange(k, "b h T H W d -> b h (T H W) d") + return q, k + + +class NaRotaryEmbedding3d(RotaryEmbedding3d): + def forward( + self, + q: torch.FloatTensor, # L h d + k: torch.FloatTensor, # L h d + shape: torch.LongTensor, + cache: Cache, + ) -> Tuple[ + torch.FloatTensor, + torch.FloatTensor, + ]: + freqs = cache("rope_freqs_3d", lambda: self.get_freqs(shape)) + freqs = freqs.to(device=q.device, dtype=q.dtype) + q = rearrange(q, "L h d -> h L d") + k = rearrange(k, "L h d -> h L d") + q = apply_rotary_emb(freqs, q.float()).to(q.dtype) + k = apply_rotary_emb(freqs, k.float()).to(k.dtype) + q = rearrange(q, "h L d -> L h d") + k = rearrange(k, "h L d -> L h d") + return q, k + + @torch._dynamo.disable # Disable compilation: shape.tolist() is data-dependent and causes graph breaks + def get_freqs( + self, + shape: torch.LongTensor, + ) -> torch.Tensor: + """ + Generate RoPE frequencies for video and text with adaptive dimensions. + + Note: This method uses @torch._dynamo.disable because it requires + data-dependent control flow (.tolist()) that cannot be symbolically + traced by torch.compile. The cache() wrapper in the forward pass memoizes + results to reduce recomputation overhead. + """ + freq_list = [] + for f, h, w in shape.tolist(): + freqs = self.get_axial_freqs(f, h, w) + freq_list.append(freqs.view(-1, freqs.size(-1))) + return torch.cat(freq_list, dim=0) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/window.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/window.py new file mode 100644 index 0000000000000000000000000000000000000000..b7475921ae283cf76d82bff7521233c133f54bfd --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/dit_7b/window.py @@ -0,0 +1,83 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from math import ceil +from typing import Tuple +import math + +def get_window_op(name: str): + if name == "720pwin_by_size_bysize": + return make_720Pwindows_bysize + if name == "720pswin_by_size_bysize": + return make_shifted_720Pwindows_bysize + raise ValueError(f"Unknown windowing method: {name}") + + +# -------------------------------- Windowing -------------------------------- # +def make_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]): + t, h, w = size + resized_nt, resized_nh, resized_nw = num_windows + #cal windows under 720p + scale = math.sqrt((45 * 80) / (h * w)) + resized_h, resized_w = round(h * scale), round(w * scale) + wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size. + wt = ceil(min(t, 30) / resized_nt) # window size. + nt, nh, nw = ceil(t / wt), ceil(h / wh), ceil(w / ww) # window size. + return [ + ( + slice(it * wt, min((it + 1) * wt, t)), + slice(ih * wh, min((ih + 1) * wh, h)), + slice(iw * ww, min((iw + 1) * ww, w)), + ) + for iw in range(nw) + if min((iw + 1) * ww, w) > iw * ww + for ih in range(nh) + if min((ih + 1) * wh, h) > ih * wh + for it in range(nt) + if min((it + 1) * wt, t) > it * wt + ] + +def make_shifted_720Pwindows_bysize(size: Tuple[int, int, int], num_windows: Tuple[int, int, int]): + t, h, w = size + resized_nt, resized_nh, resized_nw = num_windows + #cal windows under 720p + scale = math.sqrt((45 * 80) / (h * w)) + resized_h, resized_w = round(h * scale), round(w * scale) + wh, ww = ceil(resized_h / resized_nh), ceil(resized_w / resized_nw) # window size. + wt = ceil(min(t, 30) / resized_nt) # window size. + + st, sh, sw = ( # shift size. + 0.5 if wt < t else 0, + 0.5 if wh < h else 0, + 0.5 if ww < w else 0, + ) + nt, nh, nw = ceil((t - st) / wt), ceil((h - sh) / wh), ceil((w - sw) / ww) # window size. + nt, nh, nw = ( # number of window. + nt + 1 if st > 0 else 1, + nh + 1 if sh > 0 else 1, + nw + 1 if sw > 0 else 1, + ) + return [ + ( + slice(max(int((it - st) * wt), 0), min(int((it - st + 1) * wt), t)), + slice(max(int((ih - sh) * wh), 0), min(int((ih - sh + 1) * wh), h)), + slice(max(int((iw - sw) * ww), 0), min(int((iw - sw + 1) * ww), w)), + ) + for iw in range(nw) + if min(int((iw - sw + 1) * ww), w) > max(int((iw - sw) * ww), 0) + for ih in range(nh) + if min(int((ih - sh + 1) * wh), h) > max(int((ih - sh) * wh), 0) + for it in range(nt) + if min(int((it - st + 1) * wt), t) > max(int((it - st) * wt), 0) + ] diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/attn_video_vae.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/attn_video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..2a7d70bd03e2389421ced10fea26b18032137269 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/attn_video_vae.py @@ -0,0 +1,1733 @@ +# Copyright (c) 2023 HuggingFace Team +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache License, Version 2.0 (the "License") +# +# This file has been modified by ByteDance Ltd. and/or its affiliates. on 1st June 2025 +# +# Original file was released under Apache License, Version 2.0 (the "License"), with the full license text +# available at http://www.apache.org/licenses/LICENSE-2.0. +# +# This modified file is released under the same license. + + +from contextlib import nullcontext +from typing import Literal, Optional, Tuple, Union +import diffusers +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.attention_processor import Attention, SpatialNorm +from diffusers.models.autoencoders.vae import DecoderOutput, DiagonalGaussianDistribution +from diffusers.models.downsampling import Downsample2D +from diffusers.models.lora import LoRACompatibleConv +from diffusers.models.modeling_outputs import AutoencoderKLOutput +from diffusers.models.resnet import ResnetBlock2D +from diffusers.models.unets.unet_2d_blocks import DownEncoderBlock2D, UpDecoderBlock2D +from diffusers.models.upsampling import Upsample2D +from diffusers.utils import is_torch_version +from diffusers.utils.accelerate_utils import apply_forward_hook +from einops import rearrange +from ....common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation + +from ....common.distributed.advanced import get_sequence_parallel_world_size +from ....common.logger import get_logger +from .causal_inflation_lib import ( + InflatedCausalConv3d, + causal_norm_wrapper, + init_causal_conv3d, + remove_head, +) +from .context_parallel_lib import ( + causal_conv_gather_outputs, + causal_conv_slice_inputs, +) +from .global_config import set_norm_limit +from .types import ( + CausalAutoencoderOutput, + CausalDecoderOutput, + CausalEncoderOutput, + MemoryState, + _inflation_mode_t, + _memory_device_t, + _receptive_field_t, +) +from ....optimization.memory_manager import retry_on_oom + +logger = get_logger(__name__) # pylint: disable=invalid-name + +class Upsample3D(Upsample2D): + """A 3D upsampling layer with an optional convolution.""" + + def __init__( + self, + *args, + inflation_mode: _inflation_mode_t = "tail", + temporal_up: bool = False, + spatial_up: bool = True, + slicing: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + conv = self.conv if self.name == "conv" else self.Conv2d_0 + + assert type(conv) is not nn.ConvTranspose2d + # Note: lora_layer is not passed into constructor in the original implementation. + # So we make a simplification. + conv = init_causal_conv3d( + self.channels, + self.out_channels, + 3, + padding=1, + inflation_mode=inflation_mode, + ) + + self.temporal_up = temporal_up + self.spatial_up = spatial_up + self.temporal_ratio = 2 if temporal_up else 1 + self.spatial_ratio = 2 if spatial_up else 1 + self.slicing = slicing + + assert not self.interpolate + # [Override] MAGViT v2 implementation + if not self.interpolate: + upscale_ratio = (self.spatial_ratio**2) * self.temporal_ratio + self.upscale_conv = nn.Conv3d( + self.channels, self.channels * upscale_ratio, kernel_size=1, padding=0 + ) + identity = ( + torch.eye(self.channels) + .repeat(upscale_ratio, 1) + .reshape_as(self.upscale_conv.weight) + ) + self.upscale_conv.weight.data.copy_(identity) + nn.init.zeros_(self.upscale_conv.bias) + + if self.name == "conv": + self.conv = conv + else: + self.Conv2d_0 = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + output_size: Optional[int] = None, + memory_state: MemoryState = MemoryState.DISABLED, + **kwargs, + ) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if hasattr(self, "norm") and self.norm is not None: + # [Overridden] change to causal norm. + hidden_states = causal_norm_wrapper(self.norm, hidden_states) + + if self.use_conv_transpose: + return self.conv(hidden_states) + + if self.slicing: + split_size = hidden_states.size(2) // 2 + hidden_states = list( + hidden_states.split([split_size, hidden_states.size(2) - split_size], dim=2) + ) + else: + hidden_states = [hidden_states] + + for i in range(len(hidden_states)): + def upscale_and_rearrange(): + temp = self.upscale_conv(hidden_states[i]) + return rearrange( + temp, + "b (x y z c) f h w -> b c (f z) (h x) (w y)", + x=self.spatial_ratio, + y=self.spatial_ratio, + z=self.temporal_ratio, + ) + + hidden_states[i] = retry_on_oom( + upscale_and_rearrange, + debug=getattr(self, 'debug', None), + operation_name="Upsample3D.upscale_conv" + ) + + # [Overridden] For causal temporal conv + if self.temporal_up and memory_state != MemoryState.ACTIVE: + hidden_states[0] = remove_head(hidden_states[0]) + + if not self.slicing: + hidden_states = hidden_states[0] + + if self.use_conv: + def apply_conv(): + if self.name == "conv": + return self.conv(hidden_states, memory_state=memory_state) + else: + return self.Conv2d_0(hidden_states, memory_state=memory_state) + + hidden_states = retry_on_oom( + apply_conv, + debug=getattr(self, 'debug', None), + operation_name="Upsample3D.conv" + ) + + if not self.slicing: + return hidden_states + else: + return torch.cat(hidden_states, dim=2) + + +class Downsample3D(Downsample2D): + """A 3D downsampling layer with an optional convolution.""" + + def __init__( + self, + *args, + inflation_mode: _inflation_mode_t = "tail", + spatial_down: bool = False, + temporal_down: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + conv = self.conv + self.temporal_down = temporal_down + self.spatial_down = spatial_down + + self.temporal_ratio = 2 if temporal_down else 1 + self.spatial_ratio = 2 if spatial_down else 1 + + self.temporal_kernel = 3 if temporal_down else 1 + self.spatial_kernel = 3 if spatial_down else 1 + + if type(conv) in [nn.Conv2d, LoRACompatibleConv]: + # Note: lora_layer is not passed into constructor in the original implementation. + # So we make a simplification. + conv = init_causal_conv3d( + self.channels, + self.out_channels, + kernel_size=(self.temporal_kernel, self.spatial_kernel, self.spatial_kernel), + stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio), + padding=( + 1 if self.temporal_down else 0, + self.padding if self.spatial_down else 0, + self.padding if self.spatial_down else 0, + ), + inflation_mode=inflation_mode, + ) + elif type(conv) is nn.AvgPool2d: + assert self.channels == self.out_channels + conv = nn.AvgPool3d( + kernel_size=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio), + stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio), + ) + else: + raise NotImplementedError + + if self.name == "conv": + self.Conv2d_0 = conv + self.conv = conv + else: + self.conv = conv + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state: MemoryState = MemoryState.DISABLED, + **kwargs, + ) -> torch.FloatTensor: + + assert hidden_states.shape[1] == self.channels + + if hasattr(self, "norm") and self.norm is not None: + # [Overridden] change to causal norm. + hidden_states = causal_norm_wrapper(self.norm, hidden_states) + + if self.use_conv and self.padding == 0 and self.spatial_down: + pad = (0, 1, 0, 1) + hidden_states = safe_pad_operation(hidden_states, pad, mode="constant", value=0) + + assert hidden_states.shape[1] == self.channels + + hidden_states = self.conv(hidden_states, memory_state=memory_state) + + return hidden_states + + +class ResnetBlock3D(ResnetBlock2D): + def __init__( + self, + *args, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + slicing: bool = False, + **kwargs, + ): + super().__init__(*args, **kwargs) + self.conv1 = init_causal_conv3d( + self.in_channels, + self.out_channels, + kernel_size=(1, 3, 3) if time_receptive_field == "half" else (3, 3, 3), + stride=1, + padding=(0, 1, 1) if time_receptive_field == "half" else (1, 1, 1), + inflation_mode=inflation_mode, + ) + + self.conv2 = init_causal_conv3d( + self.out_channels, + self.conv2.out_channels, + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + if self.up: + assert type(self.upsample) is Upsample2D + self.upsample = Upsample3D( + self.in_channels, + use_conv=False, + inflation_mode=inflation_mode, + slicing=slicing, + ) + elif self.down: + assert type(self.downsample) is Downsample2D + self.downsample = Downsample3D( + self.in_channels, + use_conv=False, + padding=1, + name="op", + inflation_mode=inflation_mode, + ) + + if self.use_in_shortcut: + self.conv_shortcut = init_causal_conv3d( + self.in_channels, + self.conv_shortcut.out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=(self.conv_shortcut.bias is not None), + inflation_mode=inflation_mode, + ) + + + def forward( + self, input_tensor, temb, memory_state: MemoryState = MemoryState.DISABLED, **kwargs + ): + hidden_states = input_tensor + + hidden_states = causal_norm_wrapper(self.norm1, hidden_states) + hidden_states = retry_on_oom( + self.nonlinearity, + hidden_states, + debug=getattr(self, 'debug', None), + operation_name="ResnetBlock3D.nonlinearity" + ) + + if self.upsample is not None: + # upsample_nearest_nhwc fails with large batch sizes. + # see https://github.com/huggingface/diffusers/issues/984 + if hidden_states.shape[0] >= 64: + input_tensor = input_tensor.contiguous() + hidden_states = hidden_states.contiguous() + input_tensor = self.upsample(input_tensor, memory_state=memory_state) + hidden_states = self.upsample(hidden_states, memory_state=memory_state) + elif self.downsample is not None: + input_tensor = self.downsample(input_tensor, memory_state=memory_state) + hidden_states = self.downsample(hidden_states, memory_state=memory_state) + + hidden_states = self.conv1(hidden_states, memory_state=memory_state) + + if self.time_emb_proj is not None: + if not self.skip_time_act: + temb = self.nonlinearity(temb) + temb = self.time_emb_proj(temb)[:, :, None, None] + + if temb is not None and self.time_embedding_norm == "default": + hidden_states = hidden_states + temb + + hidden_states = causal_norm_wrapper(self.norm2, hidden_states) + + if temb is not None and self.time_embedding_norm == "scale_shift": + scale, shift = torch.chunk(temb, 2, dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.nonlinearity(hidden_states) + + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states, memory_state=memory_state) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor, memory_state=memory_state) + + output_tensor = (input_tensor + hidden_states) / self.output_scale_factor + + return output_tensor + + +class DownEncoderBlock3D(DownEncoderBlock2D): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_downsample: bool = True, + downsample_padding: int = 1, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_down: bool = True, + spatial_down: bool = True, + ): + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + num_layers=num_layers, + resnet_eps=resnet_eps, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_pre_norm=resnet_pre_norm, + output_scale_factor=output_scale_factor, + add_downsample=add_downsample, + downsample_padding=downsample_padding, + ) + resnets = [] + temporal_modules = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + # [Override] Replace module. + ResnetBlock3D( + in_channels=in_channels, + out_channels=out_channels, + temb_channels=None, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + temporal_modules.append(nn.Identity()) + + self.resnets = nn.ModuleList(resnets) + self.temporal_modules = nn.ModuleList(temporal_modules) + + if add_downsample: + self.downsamplers = nn.ModuleList( + [ + # [Override] Replace module. + Downsample3D( + out_channels, + use_conv=True, + out_channels=out_channels, + padding=downsample_padding, + name="op", + temporal_down=temporal_down, + spatial_down=spatial_down, + inflation_mode=inflation_mode, + ) + ] + ) + else: + self.downsamplers = None + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state: MemoryState = MemoryState.DISABLED, + **kwargs, + ) -> torch.FloatTensor: + for resnet, temporal in zip(self.resnets, self.temporal_modules): + hidden_states = resnet(hidden_states, temb=None, memory_state=memory_state) + hidden_states = temporal(hidden_states) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, memory_state=memory_state) + + return hidden_states + + +class UpDecoderBlock3D(UpDecoderBlock2D): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + output_scale_factor: float = 1.0, + add_upsample: bool = True, + temb_channels: Optional[int] = None, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_up: bool = True, + spatial_up: bool = True, + slicing: bool = False, + ): + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + num_layers=num_layers, + resnet_eps=resnet_eps, + resnet_time_scale_shift=resnet_time_scale_shift, + resnet_act_fn=resnet_act_fn, + resnet_groups=resnet_groups, + resnet_pre_norm=resnet_pre_norm, + output_scale_factor=output_scale_factor, + add_upsample=add_upsample, + temb_channels=temb_channels, + ) + resnets = [] + temporal_modules = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + # [Override] Replace module. + ResnetBlock3D( + in_channels=input_channels, + out_channels=out_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + slicing=slicing, + ) + ) + + temporal_modules.append(nn.Identity()) + + self.resnets = nn.ModuleList(resnets) + self.temporal_modules = nn.ModuleList(temporal_modules) + + if add_upsample: + # [Override] Replace module & use learnable upsample + self.upsamplers = nn.ModuleList( + [ + Upsample3D( + out_channels, + use_conv=True, + out_channels=out_channels, + temporal_up=temporal_up, + spatial_up=spatial_up, + interpolate=False, + inflation_mode=inflation_mode, + slicing=slicing, + ) + ] + ) + else: + self.upsamplers = None + + def forward( + self, + hidden_states: torch.FloatTensor, + temb: Optional[torch.FloatTensor] = None, + memory_state: MemoryState = MemoryState.DISABLED, + ) -> torch.FloatTensor: + for resnet, temporal in zip(self.resnets, self.temporal_modules): + hidden_states = resnet(hidden_states, temb=None, memory_state=memory_state) + hidden_states = temporal(hidden_states) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, memory_state=memory_state) + + return hidden_states + + +class UNetMidBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + temb_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + resnet_eps: float = 1e-6, + resnet_time_scale_shift: str = "default", # default, spatial + resnet_act_fn: str = "swish", + resnet_groups: int = 32, + resnet_pre_norm: bool = True, + add_attention: bool = True, + attention_head_dim: int = 1, + output_scale_factor: float = 1.0, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + ): + super().__init__() + resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) + self.add_attention = add_attention + + # there is always at least one resnet + resnets = [ + # [Override] Replace module. + ResnetBlock3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ] + attentions = [] + + if attention_head_dim is None: + logger.warn( + f"It is not recommend to pass `attention_head_dim=None`. " + f"Defaulting `attention_head_dim` to `in_channels`: {in_channels}." + ) + attention_head_dim = in_channels + + for _ in range(num_layers): + if self.add_attention: + attentions.append( + Attention( + in_channels, + heads=in_channels // attention_head_dim, + dim_head=attention_head_dim, + rescale_output_factor=output_scale_factor, + eps=resnet_eps, + norm_num_groups=( + resnet_groups if resnet_time_scale_shift == "default" else None + ), + spatial_norm_dim=( + temb_channels if resnet_time_scale_shift == "spatial" else None + ), + residual_connection=True, + bias=True, + upcast_softmax=True, + _from_deprecated_attn_block=True, + ) + ) + else: + attentions.append(None) + + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=in_channels, + temb_channels=temb_channels, + eps=resnet_eps, + groups=resnet_groups, + dropout=dropout, + time_embedding_norm=resnet_time_scale_shift, + non_linearity=resnet_act_fn, + output_scale_factor=output_scale_factor, + pre_norm=resnet_pre_norm, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + + self.attentions = nn.ModuleList(attentions) + self.resnets = nn.ModuleList(resnets) + + def forward(self, hidden_states, temb=None, memory_state: MemoryState = MemoryState.DISABLED): + video_length, frame_height, frame_width = hidden_states.size()[-3:] + hidden_states = self.resnets[0](hidden_states, temb, memory_state=memory_state) + for attn, resnet in zip(self.attentions, self.resnets[1:]): + if attn is not None: + hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w") + hidden_states = attn(hidden_states, temb=temb) + hidden_states = rearrange( + hidden_states, "(b f) c h w -> b c f h w", f=video_length + ) + hidden_states = resnet(hidden_states, temb, memory_state=memory_state) + + return hidden_states + + +class Encoder3D(nn.Module): + r""" + [Override] override most logics to support extra condition input and causal conv + + The `Encoder` layer of a variational autoencoder that encodes + its input into a latent representation. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + down_block_types (`Tuple[str, ...]`, *optional*, defaults to `("DownEncoderBlock2D",)`): + The types of down blocks to use. + See `~diffusers.models.unet_2d_blocks.get_down_block` + for available options. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. + See `~diffusers.models.activations.get_activation` for available options. + double_z (`bool`, *optional*, defaults to `True`): + Whether to double the number of output channels for the last block. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str, ...] = ("DownEncoderBlock3D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + double_z: bool = True, + mid_block_add_attention=True, + # [Override] add extra_cond_dim, temporal down num + temporal_down_num: int = 2, + extra_cond_dim: int = None, + gradient_checkpoint: bool = False, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + ): + super().__init__() + self.layers_per_block = layers_per_block + self.temporal_down_num = temporal_down_num + + self.conv_in = init_causal_conv3d( + in_channels, + block_out_channels[0], + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.mid_block = None + self.down_blocks = nn.ModuleList([]) + self.extra_cond_dim = extra_cond_dim + + self.conv_extra_cond = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i, down_block_type in enumerate(down_block_types): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + # [Override] to support temporal down block design + is_temporal_down_block = i >= len(block_out_channels) - self.temporal_down_num - 1 + # Note: take the last ones + + assert down_block_type == "DownEncoderBlock3D" + + down_block = DownEncoderBlock3D( + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=not is_final_block, + resnet_eps=1e-6, + downsample_padding=0, + # Note: Don't know why set it as 0 + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + temporal_down=is_temporal_down_block, + spatial_down=True, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + self.down_blocks.append(down_block) + + def zero_module(module): + # Zero out the parameters of a module and return it. + for p in module.parameters(): + p.detach().zero_() + return module + + self.conv_extra_cond.append( + zero_module( + nn.Conv3d(extra_cond_dim, output_channel, kernel_size=1, stride=1, padding=0) + ) + if self.extra_cond_dim is not None and self.extra_cond_dim > 0 + else None + ) + + # mid + self.mid_block = UNetMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default", + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=None, + add_attention=mid_block_add_attention, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + # out + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[-1], num_groups=norm_num_groups, eps=1e-6 + ) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = init_causal_conv3d( + block_out_channels[-1], conv_out_channels, 3, padding=1, inflation_mode=inflation_mode + ) + + self.gradient_checkpointing = gradient_checkpoint + + def forward( + self, + sample: torch.FloatTensor, + extra_cond=None, + memory_state: MemoryState = MemoryState.DISABLED, + ) -> torch.FloatTensor: + r"""The forward method of the `Encoder` class.""" + sample = self.conv_in(sample, memory_state=memory_state) + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + # down + # [Override] add extra block and extra cond + for down_block, extra_block in zip(self.down_blocks, self.conv_extra_cond): + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(down_block), sample, memory_state, use_reentrant=False + ) + if extra_block is not None: + sample = sample + safe_interpolate_operation(extra_block(extra_cond), size=sample.shape[2:]) + + # middle + sample = self.mid_block(sample, memory_state=memory_state) + + # sample = torch.utils.checkpoint.checkpoint( + # create_custom_forward(self.mid_block), sample, use_reentrant=False + # ) + + else: + # down + # [Override] add extra block and extra cond + for down_block, extra_block in zip(self.down_blocks, self.conv_extra_cond): + sample = down_block(sample, memory_state=memory_state) + if extra_block is not None: + sample = sample + safe_interpolate_operation(extra_block(extra_cond), size=sample.shape[2:]) + + # middle + sample = self.mid_block(sample, memory_state=memory_state) + + # post-process + sample = causal_norm_wrapper(self.conv_norm_out, sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample, memory_state=memory_state) + + return sample + + +class Decoder3D(nn.Module): + r""" + The `Decoder` layer of a variational autoencoder that + decodes its latent representation into an output sample. + + Args: + in_channels (`int`, *optional*, defaults to 3): + The number of input channels. + out_channels (`int`, *optional*, defaults to 3): + The number of output channels. + up_block_types (`Tuple[str, ...]`, *optional*, defaults to `("UpDecoderBlock2D",)`): + The types of up blocks to use. + See `~diffusers.models.unet_2d_blocks.get_up_block` for available options. + block_out_channels (`Tuple[int, ...]`, *optional*, defaults to `(64,)`): + The number of output channels for each block. + layers_per_block (`int`, *optional*, defaults to 2): + The number of layers per block. + norm_num_groups (`int`, *optional*, defaults to 32): + The number of groups for normalization. + act_fn (`str`, *optional*, defaults to `"silu"`): + The activation function to use. + See `~diffusers.models.activations.get_activation` for available options. + norm_type (`str`, *optional*, defaults to `"group"`): + The normalization type to use. Can be either `"group"` or `"spatial"`. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + up_block_types: Tuple[str, ...] = ("UpDecoderBlock3D",), + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + norm_num_groups: int = 32, + act_fn: str = "silu", + norm_type: str = "group", # group, spatial + mid_block_add_attention=True, + # [Override] add temporal up block + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_up_num: int = 2, + slicing_up_num: int = 0, + gradient_checkpoint: bool = False, + ): + super().__init__() + self.layers_per_block = layers_per_block + self.temporal_up_num = temporal_up_num + + self.conv_in = init_causal_conv3d( + in_channels, + block_out_channels[-1], + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.mid_block = None + self.up_blocks = nn.ModuleList([]) + + temb_channels = in_channels if norm_type == "spatial" else None + + # mid + self.mid_block = UNetMidBlock3D( + in_channels=block_out_channels[-1], + resnet_eps=1e-6, + resnet_act_fn=act_fn, + output_scale_factor=1, + resnet_time_scale_shift="default" if norm_type == "group" else norm_type, + attention_head_dim=block_out_channels[-1], + resnet_groups=norm_num_groups, + temb_channels=temb_channels, + add_attention=mid_block_add_attention, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + #print(f"slicing_up_num: {slicing_up_num}") + for i, up_block_type in enumerate(up_block_types): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + + is_final_block = i == len(block_out_channels) - 1 + is_temporal_up_block = i < self.temporal_up_num + is_slicing_up_block = i >= len(block_out_channels) - slicing_up_num + # Note: Keep symmetric + + assert up_block_type == "UpDecoderBlock3D" + up_block = UpDecoderBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=not is_final_block, + resnet_eps=1e-6, + resnet_act_fn=act_fn, + resnet_groups=norm_num_groups, + resnet_time_scale_shift=norm_type, + temb_channels=temb_channels, + temporal_up=is_temporal_up_block, + slicing=is_slicing_up_block, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + self.up_blocks.append(up_block) + prev_output_channel = output_channel + + # out + if norm_type == "spatial": + self.conv_norm_out = SpatialNorm(block_out_channels[0], temb_channels) + else: + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=1e-6 + ) + self.conv_act = nn.SiLU() + self.conv_out = init_causal_conv3d( + block_out_channels[0], out_channels, 3, padding=1, inflation_mode=inflation_mode + ) + + self.gradient_checkpointing = gradient_checkpoint + + # Note: Just copy from Decoder. + def forward( + self, + sample: torch.FloatTensor, + latent_embeds: Optional[torch.FloatTensor] = None, + memory_state: MemoryState = MemoryState.DISABLED, + ) -> torch.FloatTensor: + r"""The forward method of the `Decoder` class.""" + + sample = self.conv_in(sample, memory_state=memory_state) + + if self.training and self.gradient_checkpointing: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs) + + return custom_forward + + if is_torch_version(">=", "1.11.0"): + sample = self.mid_block(sample, latent_embeds, memory_state=memory_state) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), + sample, + latent_embeds, + memory_state, + use_reentrant=False, + ) + else: + # middle + sample = self.mid_block(sample, latent_embeds, memory_state=memory_state) + + # up + for up_block in self.up_blocks: + sample = torch.utils.checkpoint.checkpoint( + create_custom_forward(up_block), sample, latent_embeds, memory_state + ) + else: + # middle + sample = self.mid_block(sample, latent_embeds, memory_state=memory_state) + + # up + for up_block in self.up_blocks: + sample = up_block(sample, latent_embeds, memory_state=memory_state) + + # post-process + sample = causal_norm_wrapper(self.conv_norm_out, sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample, memory_state=memory_state) + + return sample + + +class AutoencoderKL(diffusers.AutoencoderKL): + """ + We simply inherit the model code from diffusers + """ + + def __init__(self, attention: bool = True, *args, **kwargs): + super().__init__(*args, **kwargs) + + # A hacky way to remove attention. + if not attention: + self.encoder.mid_block.attentions = torch.nn.ModuleList([None]) + self.decoder.mid_block.attentions = torch.nn.ModuleList([None]) + + def load_state_dict(self, state_dict, strict=True, assign=False): + # Newer version of diffusers changed the model keys, + # causing incompatibility with old checkpoints. + # They provided a method for conversion. We call conversion before loading state_dict. + convert_deprecated_attention_blocks = getattr( + self, "_convert_deprecated_attention_blocks", None + ) + if callable(convert_deprecated_attention_blocks): + convert_deprecated_attention_blocks(state_dict) + return super().load_state_dict(state_dict, strict, assign) + + +class VideoAutoencoderKL(diffusers.AutoencoderKL): + """ + We simply inherit the model code from diffusers + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + down_block_types: Tuple[str] = ("DownEncoderBlock3D",), + up_block_types: Tuple[str] = ("UpDecoderBlock3D",), + block_out_channels: Tuple[int] = (64,), + layers_per_block: int = 1, + act_fn: str = "silu", + latent_channels: int = 4, + norm_num_groups: int = 32, + sample_size: int = 32, + scaling_factor: float = 0.18215, + force_upcast: float = True, + attention: bool = True, + temporal_scale_num: int = 2, + slicing_up_num: int = 0, + gradient_checkpoint: bool = False, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "full", + slicing_sample_min_size: int = 32, + use_quant_conv: bool = True, + use_post_quant_conv: bool = True, + *args, + **kwargs, + ): + 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 = max(1, slicing_sample_min_size // (2**temporal_scale_num)) + + super().__init__( + in_channels=in_channels, + out_channels=out_channels, + # [Override] make sure it can be normally initialized + down_block_types=tuple( + [down_block_type.replace("3D", "2D") for down_block_type in down_block_types] + ), + up_block_types=tuple( + [up_block_type.replace("3D", "2D") for up_block_type in up_block_types] + ), + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + latent_channels=latent_channels, + norm_num_groups=norm_num_groups, + sample_size=sample_size, + scaling_factor=scaling_factor, + force_upcast=force_upcast, + *args, + **kwargs, + ) + + # pass init params to Encoder + self.encoder = Encoder3D( + in_channels=in_channels, + out_channels=latent_channels, + down_block_types=down_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + act_fn=act_fn, + norm_num_groups=norm_num_groups, + double_z=True, + extra_cond_dim=extra_cond_dim, + # [Override] add temporal_down_num parameter + temporal_down_num=temporal_scale_num, + gradient_checkpoint=gradient_checkpoint, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + # pass init params to Decoder + self.decoder = Decoder3D( + in_channels=latent_channels, + out_channels=out_channels, + up_block_types=up_block_types, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + norm_num_groups=norm_num_groups, + act_fn=act_fn, + # [Override] add temporal_up_num parameter + temporal_up_num=temporal_scale_num, + slicing_up_num=slicing_up_num, + gradient_checkpoint=gradient_checkpoint, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + self.quant_conv = ( + init_causal_conv3d( + in_channels=2 * latent_channels, + out_channels=2 * latent_channels, + kernel_size=1, + inflation_mode=inflation_mode, + ) + if use_quant_conv + else None + ) + self.post_quant_conv = ( + init_causal_conv3d( + in_channels=latent_channels, + out_channels=latent_channels, + kernel_size=1, + inflation_mode=inflation_mode, + ) + if use_post_quant_conv + else None + ) + + # A hacky way to remove attention. + if not attention: + self.encoder.mid_block.attentions = torch.nn.ModuleList([None]) + self.decoder.mid_block.attentions = torch.nn.ModuleList([None]) + + @apply_forward_hook + def encode(self, x: torch.FloatTensor, return_dict: bool = True, + tiled: bool = False, tile_size: Tuple[int, int] = (512, 512), + tile_overlap: Tuple[int, int] = (64, 64)) -> AutoencoderKLOutput: + if tiled: + h = self.tiled_encode(x, tile_size=tile_size, tile_overlap=tile_overlap) + else: + h = self.slicing_encode(x) + + posterior = DiagonalGaussianDistribution(h) + + if not return_dict: + return (posterior,) + + return AutoencoderKLOutput(latent_dist=posterior) + + @apply_forward_hook + def decode(self, z: torch.Tensor, return_dict: bool = True, + tiled: bool = False, tile_size: Tuple[int, int] = (512, 512), + tile_overlap: Tuple[int, int] = (64, 64)) -> Union[DecoderOutput, torch.Tensor]: + + if tiled: + decoded = self.tiled_decode(z, tile_size=tile_size, tile_overlap=tile_overlap) + else: + decoded = self.slicing_decode(z) + + if not return_dict: + return (decoded,) + + return DecoderOutput(sample=decoded) + + def _encode( + self, x: torch.Tensor, memory_state: MemoryState = MemoryState.DISABLED) -> torch.Tensor: + # Only transfer if not already on correct device + _x = x if x.device == self.device else x.to(self.device) + + _x = causal_conv_slice_inputs(_x, self.slicing_sample_min_size, memory_state=memory_state) + h = self.encoder(_x, memory_state=memory_state) + + if self.quant_conv is not None: + output = self.quant_conv(h, memory_state=memory_state) + else: + output = h + + 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) + + def _decode( + self, z: torch.Tensor, memory_state: MemoryState = MemoryState.DISABLED) -> torch.Tensor: + # Only transfer if not already on correct device + _z = z if z.device == self.device else z.to(self.device) + + _z = causal_conv_slice_inputs(_z, self.slicing_latent_min_size, memory_state=memory_state) + + if self.post_quant_conv is not None: + _z = self.post_quant_conv(_z, memory_state=memory_state) + + 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) + + def slicing_encode(self, x: torch.Tensor) -> torch.Tensor: + sp_size = get_sequence_parallel_world_size() + if self.use_slicing and (x.shape[2] - 1) > self.slicing_sample_min_size * sp_size: + x_slices = x[:, :, 1:].split(split_size=self.slicing_sample_min_size * sp_size, dim=2) + encoded_slices = [ + self._encode( + torch.cat((x[:, :, :1], x_slices[0]), dim=2), + memory_state=MemoryState.INITIALIZING, + ) + ] + for x_idx in range(1, len(x_slices)): + encoded_slices.append( + self._encode(x_slices[x_idx], memory_state=MemoryState.ACTIVE) + ) + out = torch.cat(encoded_slices, dim=2) + # Clear memory efficiently + modules_with_memory = [m for m in self.modules() + if isinstance(m, InflatedCausalConv3d) and m.memory is not None] + for m in modules_with_memory: + m.memory = None + return out + else: + return self._encode(x) + + def slicing_decode(self, z: torch.Tensor) -> torch.Tensor: + sp_size = get_sequence_parallel_world_size() + if self.use_slicing and (z.shape[2] - 1) > self.slicing_latent_min_size * sp_size: + z_slices = z[:, :, 1:].split(split_size=self.slicing_latent_min_size * sp_size, dim=2) + decoded_slices = [ + self._decode( + torch.cat((z[:, :, :1], z_slices[0]), dim=2), + memory_state=MemoryState.INITIALIZING + ) + ] + for z_idx in range(1, len(z_slices)): + decoded_slices.append( + self._decode(z_slices[z_idx], memory_state=MemoryState.ACTIVE) + ) + out = torch.cat(decoded_slices, dim=2) + # Clear memory efficiently + modules_with_memory = [m for m in self.modules() + if isinstance(m, InflatedCausalConv3d) and m.memory is not None] + for m in modules_with_memory: + m.memory = None + return out + else: + return self._decode(z) + + def tiled_encode(self, x: torch.Tensor, tile_size: Tuple[int, int] = (512, 512), + tile_overlap: Tuple[int, int] = (64, 64)) -> torch.Tensor: + r""" + Encodes an input tensor `x` by splitting it into spatial tiles in latent space. Temporal is handled by `slicing_encode`. + `tile_size` and `tile_overlap` are interpreted in output-space pixels and converted to latent-space. + """ + # Ensure 5D [B, C, F, H, W] + if x.ndim != 5: + x = x.unsqueeze(2) + + b, c, f, H, W = x.shape + tile_h, tile_w = tile_size + + # Only tile if input resolution requires multiple tiles + if H <= tile_h and W <= tile_w: + return self.slicing_encode(x) + else: + 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 + + # Convert output-space tiling params to latent-space + tile_h, tile_w = tile_size + overlap_h, overlap_w = tile_overlap + + latent_tile_h = max(1, tile_h // scale_factor) + latent_tile_w = max(1, tile_w // scale_factor) + 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)) + + stride_h = max(1, latent_tile_h - latent_overlap_h) + stride_w = max(1, latent_tile_w - latent_overlap_w) + + H_lat_total = (H + scale_factor - 1) // scale_factor + W_lat_total = (W + scale_factor - 1) // scale_factor + + result = None + count = None + + num_tiles = ((max(H_lat_total - latent_overlap_h, 1) + stride_h - 1) // stride_h) \ + * ((max(W_lat_total - latent_overlap_w, 1) + stride_w - 1) // stride_w) + + # Log once at start instead of per-tile + if self.debug: + self.debug.log( + f"Encoding {num_tiles} tiles (Tile: {tile_size}, Overlap: {tile_overlap})", + category="vae", + ) + + # Pre-compute common ramp values + ramp_cache = {} + if latent_overlap_h > 0: + t_h = torch.linspace(0, 1, steps=latent_overlap_h, device=x.device, dtype=x.dtype) + ramp_cache['h'] = 0.5 - 0.5 * torch.cos(t_h * torch.pi) + if latent_overlap_w > 0: + t_w = torch.linspace(0, 1, steps=latent_overlap_w, device=x.device, dtype=x.dtype) + ramp_cache['w'] = 0.5 - 0.5 * torch.cos(t_w * torch.pi) + + tile_id = 0 + for y_lat in range(0, H_lat_total, stride_h): + y_lat_end = min(y_lat + latent_tile_h, H_lat_total) + for x_lat in range(0, W_lat_total, stride_w): + x_lat_end = min(x_lat + latent_tile_w, W_lat_total) + + # Skip if fully within overlap of previous tiles + if (y_lat > 0 and (y_lat_end - y_lat) <= latent_overlap_h) or \ + (x_lat > 0 and (x_lat_end - x_lat) <= latent_overlap_w): + continue + + # Map latent tile to output-space crop + y_out = y_lat * scale_factor + x_out = x_lat * scale_factor + y_out_end = min(y_lat_end * scale_factor, H) + x_out_end = min(x_lat_end * scale_factor, W) + + tile_id += 1 + + # Store tile boundary info for debug visualization + if self.debug and hasattr(self.debug, 'encode_tile_boundaries'): + self.debug.encode_tile_boundaries.append({ + 'id': tile_id, + 'y': y_out, + 'x': x_out, + 'h': y_out_end - y_out, + 'w': x_out_end - x_out + }) + + tile_sample = x[:, :, :, y_out:y_out_end, x_out:x_out_end] + + # Log progress periodically instead of every tile (at 1, 6, 11, 16, ...) + if self.debug and (tile_id % 5 == 1 or tile_id == num_tiles): + if tile_id == num_tiles: + # Only log final tile if not covered by previous range + if (tile_id - 1) % 5 == 0: + self.debug.log(f"Encoding tile {tile_id} / {num_tiles}", category="vae", indent_level=1) + else: + end_tile = min(tile_id + 4, num_tiles) + self.debug.log(f"Encoding tiles {tile_id}-{end_tile} / {num_tiles}", category="vae", indent_level=1) + + encoded_tile = self.slicing_encode(tile_sample) + + # Initialize output size using first encoded tile + if result is None: + b_out, c_out, f_lat, _, _ = encoded_tile.shape + + # Accumulate on offload device if specified and different, else on inference device + device = getattr(self, 'tensor_offload_device', None) + if device is None or device == encoded_tile.device: + device = encoded_tile.device + + result = torch.zeros( + (b_out, c_out, f_lat, H_lat_total, W_lat_total), + device=device, + dtype=encoded_tile.dtype, + ) + count = torch.zeros((1, 1, 1, H_lat_total, W_lat_total), device=device, dtype=encoded_tile.dtype) + + eff_h_lat = min(y_lat_end - y_lat, encoded_tile.shape[3], result.shape[3] - y_lat) + eff_w_lat = min(x_lat_end - x_lat, encoded_tile.shape[4], result.shape[4] - x_lat) + + encoded_tile = encoded_tile[:, :, : result.shape[2], :eff_h_lat, :eff_w_lat] + + # Build faded masks + ov_h = max(0, min(latent_overlap_h, eff_h_lat - 1)) + ov_w = max(0, min(latent_overlap_w, eff_w_lat - 1)) + + weight_h = torch.ones((eff_h_lat,), device=encoded_tile.device, dtype=encoded_tile.dtype) + weight_w = torch.ones((eff_w_lat,), device=encoded_tile.device, dtype=encoded_tile.dtype) + + # Apply fades only on interior edges using cached ramps (avoid fading on outer image borders) + if ov_h > 0: + if y_lat > 0: # Not top edge + weight_h[:ov_h] = ramp_cache['h'][:ov_h] + if y_lat_end < H_lat_total: # Not bottom edge + weight_h[-ov_h:] = 1 - ramp_cache['h'][:ov_h] + if ov_w > 0: + if x_lat > 0: # Not left edge + weight_w[:ov_w] = ramp_cache['w'][:ov_w] + if x_lat_end < W_lat_total: # Not right edge + weight_w[-ov_w:] = 1 - ramp_cache['w'][:ov_w] + + # Separable application (no 2D mask to save memory) + weight_h_5d = weight_h.view(1, 1, 1, eff_h_lat, 1) + weight_w_5d = weight_w.view(1, 1, 1, 1, eff_w_lat) + encoded_tile.mul_(weight_h_5d).mul_(weight_w_5d) + + # Accumulate (move to result device if different) + if result.device != encoded_tile.device: + encoded_tile = encoded_tile.to(result.device) + weight_h_5d = weight_h_5d.to(result.device) + weight_w_5d = weight_w_5d.to(result.device) + + result[:, :, : encoded_tile.shape[2], y_lat : y_lat + eff_h_lat, x_lat : x_lat + eff_w_lat] += encoded_tile + count[:, :, :, y_lat : y_lat + eff_h_lat, x_lat : x_lat + eff_w_lat].addcmul_(weight_h_5d, weight_w_5d) + + # Move result back to inference device if needed and normalize + if result.device != x.device: + result = result.to(x.device) + count = count.to(x.device) + result.div_(count.clamp(min=1e-6)) + + if x.shape[2] == 1: # single frame + result = result.squeeze(2) + + return result + + def tiled_decode(self, z: torch.Tensor, tile_size: Tuple[int, int] = (512, 512), tile_overlap: Tuple[int, int] = (64, 64)) -> torch.Tensor: + r""" + Decodes a latent tensor `z` by splitting it into spatial tiles only. Temporal is handled by `slicing_decode`. + """ + if z.ndim != 5: + z = z.unsqueeze(2) + + b, c, f, H, W = z.shape + + # Spatial scale factor (output/latent) + scale_factor = self.spatial_downsample_factor + + # Convert output-space tiling params to latent-space for spatial tiling + tile_h, tile_w = tile_size + overlap_h, overlap_w = tile_overlap + + latent_tile_h = max(1, tile_h // scale_factor) + latent_tile_w = max(1, tile_w // scale_factor) + + # Only tile if latent resolution requires multiple tiles + if H <= latent_tile_h and W <= latent_tile_w: + return self.slicing_decode(z) + else: + 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)) + + stride_h = max(1, latent_tile_h - latent_overlap_h) + stride_w = max(1, latent_tile_w - latent_overlap_w) + + # Allocate later using first decoded results + result = None + count = None + + num_tiles = ((max(H - latent_overlap_h, 1) + stride_h - 1) // stride_h) \ + * ((max(W - latent_overlap_w, 1) + stride_w - 1) // stride_w) + + # Log once at start instead of per-tile + if self.debug: + self.debug.log( + f"Decoding {num_tiles} tiles (Tile: {tile_size}, Overlap: {tile_overlap})", + category="vae", + ) + + # Pre-compute common ramp values (small memory, big time save) + ramp_cache = {} + if overlap_h > 0: + t_h = torch.linspace(0, 1, steps=overlap_h, device=z.device, dtype=z.dtype) + ramp_cache['h'] = 0.5 - 0.5 * torch.cos(t_h * torch.pi) + if overlap_w > 0: + t_w = torch.linspace(0, 1, steps=overlap_w, device=z.device, dtype=z.dtype) + ramp_cache['w'] = 0.5 - 0.5 * torch.cos(t_w * torch.pi) + + tile_id = 0 + for y_lat in range(0, H, stride_h): + y_lat_end = min(y_lat + latent_tile_h, H) + for x_lat in range(0, W, stride_w): + x_lat_end = min(x_lat + latent_tile_w, W) + + # Skip if fully within overlap of previous tiles + if (y_lat > 0 and (y_lat_end - y_lat) <= latent_overlap_h) or \ + (x_lat > 0 and (x_lat_end - x_lat) <= latent_overlap_w): + continue + + tile_id += 1 + + # Store tile boundary info for debug visualization + if self.debug and hasattr(self.debug, 'decode_tile_boundaries'): + # Map to output space + y_out = y_lat * scale_factor + x_out = x_lat * scale_factor + y_out_end = y_lat_end * scale_factor + x_out_end = x_lat_end * scale_factor + self.debug.decode_tile_boundaries.append({ + 'id': tile_id, + 'y': y_out, + 'x': x_out, + 'h': y_out_end - y_out, + 'w': x_out_end - x_out + }) + + tile_latent = z[:, :, :, y_lat:y_lat_end, x_lat:x_lat_end] + + # Log progress periodically instead of every tile (at 1, 6, 11, 16, ...) + if self.debug and (tile_id % 5 == 1 or tile_id == num_tiles): + if tile_id == num_tiles: + # Only log final tile if not covered by previous range + if (tile_id - 1) % 5 == 0: + self.debug.log(f"Decoding tile {tile_id} / {num_tiles}", category="vae", indent_level=1) + else: + end_tile = min(tile_id + 4, num_tiles) + self.debug.log(f"Decoding tiles {tile_id}-{end_tile} / {num_tiles}", category="vae", indent_level=1) + + decoded_tile = self.slicing_decode(tile_latent) + + # Initialize result tensors using actual decoded shapes on first tile + if result is None: + b_out, c_out, out_f_tile, _, _ = decoded_tile.shape + output_h = H * scale_factor + output_w = W * scale_factor + + # Accumulate on offload device if specified and different, else on inference device + device = getattr(self, 'tensor_offload_device', None) + if device is None or device == decoded_tile.device: + device = decoded_tile.device + + result = torch.zeros((b_out, c_out, out_f_tile, output_h, output_w), device=device, dtype=decoded_tile.dtype) + count = torch.zeros((1, 1, 1, output_h, output_w), device=device, dtype=decoded_tile.dtype) + + # Corresponding output-space placement + y_out, y_out_end = y_lat * scale_factor, y_lat_end * scale_factor + x_out, x_out_end = x_lat * scale_factor, x_lat_end * scale_factor + + h_out = y_out_end - y_out + w_out = x_out_end - x_out + + # Build faded masks + ov_h_out = max(0, min(overlap_h, h_out - 1)) + ov_w_out = max(0, min(overlap_w, w_out - 1)) + + weight_h = torch.ones((h_out,), device=decoded_tile.device, dtype=decoded_tile.dtype) + weight_w = torch.ones((w_out,), device=decoded_tile.device, dtype=decoded_tile.dtype) + + # Apply fades only on interior edges using cached ramps (avoid fading on outer image borders) + if ov_h_out > 0: + if y_lat > 0: # Not top edge + weight_h[:ov_h_out] = ramp_cache['h'][:ov_h_out] + if y_lat_end < H: # Not bottom edge + weight_h[-ov_h_out:] = 1 - ramp_cache['h'][:ov_h_out] + if ov_w_out > 0: + if x_lat > 0: # Not left edge + weight_w[:ov_w_out] = ramp_cache['w'][:ov_w_out] + if x_lat_end < W: # Not right edge + weight_w[-ov_w_out:] = 1 - ramp_cache['w'][:ov_w_out] + + # Separable application (no 2D mask to save memory) + weight_h_5d = weight_h.view(1, 1, 1, h_out, 1) + weight_w_5d = weight_w.view(1, 1, 1, 1, w_out) + decoded_tile.mul_(weight_h_5d).mul_(weight_w_5d) + + # Accumulate (move to result device if different) + if result.device != decoded_tile.device: + decoded_tile = decoded_tile.to(result.device) + weight_h_5d = weight_h_5d.to(result.device) + weight_w_5d = weight_w_5d.to(result.device) + + result[:, :, : decoded_tile.shape[2], y_out:y_out_end, x_out:x_out_end] += decoded_tile + count[:, :, :, y_out:y_out_end, x_out:x_out_end].addcmul_(weight_h_5d, weight_w_5d) + + # Move result back to inference device if needed and normalize + if result.device != z.device: + result = result.to(z.device) + count = count.to(z.device) + result.div_(count.clamp(min=1e-6)) # In-place normalize + + if z.shape[2] == 1: # single frame + result = result.squeeze(2) + + return result + + def forward( + self, x: torch.FloatTensor, mode: Literal["encode", "decode", "all"] = "all", **kwargs + ): + # x: [b c t h w] + if mode == "encode": + h = self.encode(x) + return h.latent_dist + elif mode == "decode": + h = self.decode(x) + return h.sample + else: + h = self.encode(x) + h = self.decode(h.latent_dist.mode()) + return h.sample + + def load_state_dict(self, state_dict, strict=False, assign=False): + # Newer version of diffusers changed the model keys, + # causing incompatibility with old checkpoints. + # They provided a method for conversion. + # We call conversion before loading state_dict. + convert_deprecated_attention_blocks = getattr( + self, "_convert_deprecated_attention_blocks", None + ) + if callable(convert_deprecated_attention_blocks): + convert_deprecated_attention_blocks(state_dict) + return super().load_state_dict(state_dict, strict, assign) + + +class VideoAutoencoderKLWrapper(VideoAutoencoderKL): + def __init__( + self, + *args, + spatial_downsample_factor: int, + temporal_downsample_factor: int, + freeze_encoder: bool, + **kwargs, + ): + self.spatial_downsample_factor = spatial_downsample_factor + self.temporal_downsample_factor = temporal_downsample_factor + self.freeze_encoder = freeze_encoder + super().__init__(*args, **kwargs) + + def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput: + with torch.no_grad() if self.freeze_encoder else nullcontext(): + z, p = self.encode(x) + x = self.decode(z).sample + return CausalAutoencoderOutput(x, z, p) + + def encode(self, x: torch.FloatTensor, return_dict: bool = True, + tiled: bool = False, tile_size: Tuple[int, int] = (512, 512), + tile_overlap: Tuple[int, int] = (64, 64)) -> CausalEncoderOutput: + if x.ndim == 4: + x = x.unsqueeze(2) + p = super().encode(x, return_dict=return_dict, tiled=tiled, tile_size=tile_size, + tile_overlap=tile_overlap).latent_dist + # Use deterministic mode for tiled encoding to avoid artifacts + z = p.mode().squeeze(2) + return CausalEncoderOutput(z, p) + + def decode(self, z: torch.Tensor, return_dict: bool = True, + tiled: bool = False, tile_size: Tuple[int, int] = (512, 512), + tile_overlap: Tuple[int, int] = (64, 64)) -> CausalDecoderOutput: + if z.ndim == 4: + z = z.unsqueeze(2) + x = super().decode(z, return_dict=return_dict, tiled=tiled, tile_size=tile_size, + tile_overlap=tile_overlap).sample.squeeze(2) + return CausalDecoderOutput(x) + + def preprocess(self, x: torch.Tensor): + # x should in [B, C, T, H, W], [B, C, H, W] + assert x.ndim == 4 or x.size(2) % 4 == 1 + return x + + def postprocess(self, x: torch.Tensor): + # x should in [B, C, T, H, W], [B, C, H, W] + return x + + def set_causal_slicing( + self, + *, + split_size: Optional[int], + memory_device: _memory_device_t, + ): + assert ( + split_size is None or memory_device is not None + ), "if split_size is set, memory_device must not be None." + if split_size is not None: + self.enable_slicing() + self.slicing_sample_min_size = split_size + self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor) + else: + self.disable_slicing() + for module in self.modules(): + if isinstance(module, InflatedCausalConv3d): + module.set_memory_device(memory_device) + + def set_memory_limit(self, conv_max_mem: Optional[float], norm_max_mem: Optional[float]): + set_norm_limit(norm_max_mem) + for m in self.modules(): + if isinstance(m, InflatedCausalConv3d): + m.set_memory_limit(conv_max_mem if conv_max_mem is not None else float("inf")) + \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/causal_inflation_lib.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/causal_inflation_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..664f8b93b44bca89806c37107c669a70e2bc3c79 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/causal_inflation_lib.py @@ -0,0 +1,503 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +import math +from contextlib import contextmanager +from typing import List, Optional, Union +import torch +import torch.nn.functional as F +from diffusers.models.normalization import RMSNorm +from einops import rearrange +from torch import Tensor, nn +from torch.nn import Conv3d + +from .context_parallel_lib import cache_send_recv, get_cache_size +from .global_config import get_norm_limit +from .types import MemoryState, _inflation_mode_t, _memory_device_t +from ....common.half_precision_fixes import safe_pad_operation +from ....optimization.memory_manager import retry_on_oom +from ....optimization.compatibility import NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND + +# Single GPU inference - no distributed processing needed +#print("Warning: Using single GPU inference mode - distributed features disabled in causal_inflation_lib") + +# Mock distributed functions for single GPU inference +def get_sequence_parallel_group(): + return None + +def get_sequence_parallel_rank(): + return 0 + +def get_sequence_parallel_world_size(): + return 1 + +def get_next_sequence_parallel_rank(): + return 0 + +def get_prev_sequence_parallel_rank(): + return 0 + + +@contextmanager +def ignore_padding(model): + orig_padding = model.padding + model.padding = (0, 0, 0) + try: + yield + finally: + model.padding = orig_padding + + +class InflatedCausalConv3d(Conv3d): + def __init__( + self, + *args, + inflation_mode: _inflation_mode_t, + memory_device: _memory_device_t = "same", + **kwargs, + ): + self.inflation_mode = inflation_mode + self.memory = None + super().__init__(*args, **kwargs) + self.temporal_padding = self.padding[0] + self.memory_device = memory_device + self.padding = (0, *self.padding[1:]) # Remove temporal pad to keep causal. + self.memory_limit = float("inf") + + def set_memory_limit(self, value: float): + self.memory_limit = value + + def set_memory_device(self, memory_device: _memory_device_t): + self.memory_device = memory_device + + def _conv_forward(self, input, weight, bias, *args, **kwargs): + """ + Override _conv_forward to work around NVIDIA Conv3d memory bug. + + Bug: PyTorch 2.9-2.10 with cuDNN >= 91002 uses 3x memory for Conv3d + with fp16/bfloat16 weights due to buggy dispatch layer. + + 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) and + hasattr(torch.backends.cudnn, 'is_available') and + torch.backends.cudnn.is_available() and + getattr(torch.backends.cudnn, 'enabled', True)): + try: + # Direct cuDNN call bypasses buggy PyTorch dispatch layer (NVIDIA only) + out = torch.cudnn_convolution( + input, weight, self.padding, self.stride, self.dilation, self.groups, + benchmark=False, deterministic=False, allow_tf32=True + ) + 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, + x, + *, + split_dim=3, + padding=(0, 0, 0, 0, 0, 0), + prev_cache=None, + ): + # Compatible with no limit. + if math.isinf(self.memory_limit): + if prev_cache is not None: + x = torch.cat([prev_cache, x], dim=split_dim - 1) + return super().forward(x) + + # Compute tensor shape after concat & padding. + shape = torch.tensor(x.size()) + if prev_cache is not None: + shape[split_dim - 1] += prev_cache.size(split_dim - 1) + shape[-3:] += torch.tensor(padding).view(3, 2).sum(-1).flip(0) + memory_occupy = shape.prod() * x.element_size() / 1024**3 # GiB + if memory_occupy < self.memory_limit or split_dim == x.ndim: + x_concat = x + if prev_cache is not None: + x_concat = torch.cat([prev_cache, x], dim=split_dim - 1) + + def pad_and_forward(): + padded = safe_pad_operation(x_concat, padding, mode='constant', value=0.0) + with ignore_padding(self): + return Conv3d.forward(self, padded) + + return retry_on_oom( + pad_and_forward, + debug=getattr(self, 'debug', None), + operation_name="InflatedCausalConv3d.pad_and_forward" + ) + + # Exceed memory limit, splitting tensor + + # Split input (& prev_cache). + num_splits = math.ceil(memory_occupy / self.memory_limit) + size_per_split = x.size(split_dim) // num_splits + split_sizes = [size_per_split] * (num_splits - 1) + split_sizes += [x.size(split_dim) - sum(split_sizes)] + + x = list(x.split(split_sizes, dim=split_dim)) + if prev_cache is not None: + prev_cache = list(prev_cache.split(split_sizes, dim=split_dim)) + # Loop Fwd. + cache = None + for idx in range(len(x)): + # Concat prev cache from last dim + if prev_cache is not None: + x[idx] = torch.cat([prev_cache[idx], x[idx]], dim=split_dim - 1) + + # Get padding pattern. + lpad_dim = (x[idx].ndim - split_dim - 1) * 2 + rpad_dim = lpad_dim + 1 + padding = list(padding) + padding[lpad_dim] = self.padding[split_dim - 2] if idx == 0 else 0 + padding[rpad_dim] = self.padding[split_dim - 2] if idx == len(x) - 1 else 0 + pad_len = padding[lpad_dim] + padding[rpad_dim] + padding = tuple(padding) + + # Prepare cache for next slice (this dim). + next_cache = None + cache_len = cache.size(split_dim) if cache is not None else 0 + next_catch_size = get_cache_size( + conv_module=self, + input_len=x[idx].size(split_dim) + cache_len, + pad_len=pad_len, + dim=split_dim - 2, + ) + if next_catch_size != 0: + assert next_catch_size <= x[idx].size(split_dim) + next_cache = ( + x[idx].transpose(0, split_dim)[-next_catch_size:].transpose(0, split_dim) + ) + + # Recursive. + x[idx] = self.memory_limit_conv( + x[idx], + split_dim=split_dim + 1, + padding=padding, + prev_cache=cache + ) + + # Update cache. + cache = next_cache + + output = retry_on_oom( + torch.cat, + x, + split_dim, + debug=getattr(self, 'debug', None), + operation_name="InflatedCausalConv3d.concat_splits" + ) + return output + + def forward( + self, + input: Union[Tensor, List[Tensor]], + memory_state: MemoryState = MemoryState.UNSET + ) -> Tensor: + assert memory_state != MemoryState.UNSET + if memory_state != MemoryState.ACTIVE: + self.memory = None + if ( + math.isinf(self.memory_limit) + and torch.is_tensor(input) + and get_sequence_parallel_group() is None + ): + return self.basic_forward(input, memory_state) + return self.slicing_forward(input, memory_state) + + def basic_forward(self, input: Tensor, memory_state: MemoryState = MemoryState.UNSET): + mem_size = self.stride[0] - self.kernel_size[0] + if (self.memory is not None) and (memory_state == MemoryState.ACTIVE): + input = extend_head(input, memory=self.memory, times=-1) + else: + input = extend_head(input, times=self.temporal_padding * 2) + memory = ( + input[:, :, mem_size:].detach() + if (mem_size != 0 and memory_state != MemoryState.DISABLED) + else None + ) + if ( + memory_state != MemoryState.DISABLED + and not self.training + and (self.memory_device is not None) + ): + self.memory = memory + if self.memory_device == "cpu" and self.memory is not None: + self.memory = self.memory.to("cpu") + return super().forward(input) + + def slicing_forward( + self, + input: Union[Tensor, List[Tensor]], + memory_state: MemoryState = MemoryState.UNSET, + ) -> Tensor: + squeeze_out = False + if torch.is_tensor(input): + input = [input] + squeeze_out = True + + cache_size = self.kernel_size[0] - self.stride[0] + cache = cache_send_recv( + input, cache_size=cache_size, memory=self.memory, times=self.temporal_padding * 2 + ) + + # Single GPU inference - simplified memory management + if ( + memory_state in [MemoryState.INITIALIZING, MemoryState.ACTIVE] # use_slicing + and not self.training + and (self.memory_device is not None) + and cache_size != 0 + ): + if cache_size > input[-1].size(2) and cache is not None and len(input) == 1: + input[0] = torch.cat([cache, input[0]], dim=2) + cache = None + if cache_size <= input[-1].size(2): + self.memory = input[-1][:, :, -cache_size:].detach().contiguous() + if self.memory_device == "cpu" and self.memory is not None: + self.memory = self.memory.to("cpu") + + padding = tuple(x for x in reversed(self.padding) for _ in range(2)) + for i in range(len(input)): + # Prepare cache for next input slice. + next_cache = None + cache_size = 0 + if i < len(input) - 1: + cache_len = cache.size(2) if cache is not None else 0 + cache_size = get_cache_size(self, input[i].size(2) + cache_len, pad_len=0) + if cache_size != 0: + if cache_size > input[i].size(2) and cache is not None: + input[i] = torch.cat([cache, input[i]], dim=2) + cache = None + assert cache_size <= input[i].size(2), f"{cache_size} > {input[i].size(2)}" + next_cache = input[i][:, :, -cache_size:] + + # Conv forward for this input slice. + input[i] = self.memory_limit_conv( + input[i], + padding=padding, + prev_cache=cache + ) + + # Update cache. + cache = next_cache + + return input[0] if squeeze_out else input + + def tflops(self, args, kwargs, output) -> float: + if torch.is_tensor(output): + output_numel = output.numel() + elif isinstance(output, list): + output_numel = sum(o.numel() for o in output) + else: + raise NotImplementedError + return (2 * math.prod(self.kernel_size) * self.in_channels * (output_numel / 1e6)) / 1e6 + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + if self.inflation_mode != "none": + state_dict = modify_state_dict( + self, + state_dict, + prefix, + inflate_weight_fn=inflate_weight, + inflate_bias_fn=inflate_bias, + ) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + (strict and self.inflation_mode == "none"), + missing_keys, + unexpected_keys, + error_msgs, + ) + + +def init_causal_conv3d( + *args, + inflation_mode: _inflation_mode_t, + **kwargs, +): + """ + Initialize a Causal-3D convolution layer. + Parameters: + inflation_mode: Listed as below. It's compatible with all the 3D-VAE checkpoints we have. + - none: No inflation will be conducted. + The loading logic of state dict will fall back to default. + - tail / replicate: Refer to the definition of `InflatedCausalConv3d`. + """ + return InflatedCausalConv3d(*args, inflation_mode=inflation_mode, **kwargs) + + +def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor: + if isinstance(norm_layer, (nn.LayerNorm, RMSNorm)): + if x.ndim == 4: + x = rearrange(x, "b c h w -> b h w c") + x = norm_layer(x) + x = rearrange(x, "b h w c -> b c h w") + return x + if x.ndim == 5: + x = rearrange(x, "b c t h w -> b t h w c") + x = norm_layer(x) + x = rearrange(x, "b t h w c -> b c t h w") + return x + if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)): + if x.ndim <= 4: + return norm_layer(x) + if x.ndim == 5: + t = x.size(2) + x = rearrange(x, "b c t h w -> (b t) c h w") + memory_occupy = x.numel() * x.element_size() / 1024**3 + if isinstance(norm_layer, nn.GroupNorm) and memory_occupy > get_norm_limit(): + num_chunks = min(4 if x.element_size() == 2 else 2, norm_layer.num_groups) + assert norm_layer.num_groups % num_chunks == 0 + num_groups_per_chunk = norm_layer.num_groups // num_chunks + + x = list(x.chunk(num_chunks, dim=1)) + weights = norm_layer.weight.chunk(num_chunks, dim=0) + biases = norm_layer.bias.chunk(num_chunks, dim=0) + + for i, (w, b) in enumerate(zip(weights, biases)): + def apply_group_norm(): + return F.group_norm(x[i], num_groups_per_chunk, w, b, norm_layer.eps) + + x[i] = retry_on_oom( + apply_group_norm, + debug=getattr(norm_layer, 'debug', None), + operation_name=f"GroupNorm.chunk_{i}" + ) + x[i] = x[i] + + x = retry_on_oom( + torch.cat, + x, + dim=1, + debug=getattr(norm_layer, 'debug', None), + operation_name="GroupNorm.concat_chunks" + ) + else: + x = retry_on_oom( + norm_layer, + x, + debug=getattr(norm_layer, 'debug', None), + operation_name="GroupNorm.direct" + ) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + return x + raise NotImplementedError + + +def remove_head(tensor: Tensor, times: int = 1) -> Tensor: + """ + Remove duplicated first frame features in the up-sampling process. + """ + # Single GPU inference - always process + if times == 0: + return tensor + return torch.cat(tensors=(tensor[:, :, :1], tensor[:, :, times + 1 :]), dim=2) + + +def extend_head(tensor: Tensor, times: int = 2, memory: Optional[Tensor] = None) -> Tensor: + """ + When memory is None: + - Duplicate first frame features in the down-sampling process. + When memory is not None: + - Concatenate memory features with the input features to keep temporal consistency. + """ + if memory is not None: + return torch.cat((memory.to(tensor), tensor), dim=2) + assert times >= 0, "Invalid input for function 'extend_head'!" + if times == 0: + return tensor + else: + tile_repeat = [1] * tensor.ndim + tile_repeat[2] = times + return torch.cat(tensors=(torch.tile(tensor[:, :, :1], tile_repeat), tensor), dim=2) + + +def inflate_weight(weight_2d: torch.Tensor, weight_3d: torch.Tensor, inflation_mode: str): + """ + Inflate a 2D convolution weight matrix to a 3D one. + Parameters: + weight_2d: The weight matrix of 2D conv to be inflated. + weight_3d: The weight matrix of 3D conv to be initialized. + inflation_mode: the mode of inflation + """ + assert inflation_mode in ["tail", "replicate"] + assert weight_3d.shape[:2] == weight_2d.shape[:2] + with torch.no_grad(): + if inflation_mode == "replicate": + depth = weight_3d.size(2) + weight_3d.copy_(weight_2d.unsqueeze(2).repeat(1, 1, depth, 1, 1) / depth) + else: + weight_3d.fill_(0.0) + weight_3d[:, :, -1].copy_(weight_2d) + return weight_3d + + +def inflate_bias(bias_2d: torch.Tensor, bias_3d: torch.Tensor, inflation_mode: str): + """ + Inflate a 2D convolution bias tensor to a 3D one + Parameters: + bias_2d: The bias tensor of 2D conv to be inflated. + bias_3d: The bias tensor of 3D conv to be initialized. + inflation_mode: Placeholder to align `inflate_weight`. + """ + assert bias_3d.shape == bias_2d.shape + with torch.no_grad(): + bias_3d.copy_(bias_2d) + return bias_3d + + +def modify_state_dict(layer, state_dict, prefix, inflate_weight_fn, inflate_bias_fn): + """ + the main function to inflated 2D parameters to 3D. + """ + weight_name = prefix + "weight" + bias_name = prefix + "bias" + if weight_name in state_dict: + weight_2d = state_dict[weight_name] + if weight_2d.dim() == 4: + # Assuming the 2D weights are 4D tensors (out_channels, in_channels, h, w) + weight_3d = inflate_weight_fn( + weight_2d=weight_2d, + weight_3d=layer.weight, + inflation_mode=layer.inflation_mode, + ) + state_dict[weight_name] = weight_3d + else: + return state_dict + # It's a 3d state dict, should not do inflation on both bias and weight. + if bias_name in state_dict: + bias_2d = state_dict[bias_name] + if bias_2d.dim() == 1: + # Assuming the 2D biases are 1D tensors (out_channels,) + bias_3d = inflate_bias_fn( + bias_2d=bias_2d, + bias_3d=layer.bias, + inflation_mode=layer.inflation_mode, + ) + state_dict[bias_name] = bias_3d + return state_dict diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/context_parallel_lib.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/context_parallel_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..663525e37c3d6fb7a75b48fb97a5b2b91e798698 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/context_parallel_lib.py @@ -0,0 +1,67 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import List +import torch +import torch.nn.functional as F +from torch import Tensor + +from .types import MemoryState + +# Single GPU inference - no distributed processing needed +# print("Warning: Using single GPU inference mode - distributed features disabled") + + +def causal_conv_slice_inputs(x, split_size, memory_state): + # Single GPU inference - no slicing needed, return full tensor + return x + + +def causal_conv_gather_outputs(x): + # Single GPU inference - no gathering needed, return tensor as is + return x + + +def get_output_len(conv_module, input_len, pad_len, dim=0): + dilated_kernerl_size = conv_module.dilation[dim] * (conv_module.kernel_size[dim] - 1) + 1 + output_len = (input_len + pad_len - dilated_kernerl_size) // conv_module.stride[dim] + 1 + return output_len + + +def get_cache_size(conv_module, input_len, pad_len, dim=0): + dilated_kernerl_size = conv_module.dilation[dim] * (conv_module.kernel_size[dim] - 1) + 1 + output_len = (input_len + pad_len - dilated_kernerl_size) // conv_module.stride[dim] + 1 + remain_len = ( + input_len + pad_len - ((output_len - 1) * conv_module.stride[dim] + dilated_kernerl_size) + ) + overlap_len = dilated_kernerl_size - conv_module.stride[dim] + cache_len = overlap_len + remain_len # >= 0 + + assert output_len > 0 + return cache_len + + +def cache_send_recv(tensor: List[Tensor], cache_size, times, memory=None): + # Single GPU inference - simplified cache handling + recv_buffer = None + + # Handle memory buffer for single GPU case + if memory is not None: + recv_buffer = memory.to(tensor[0]) + elif times > 0: + tile_repeat = [1] * tensor[0].ndim + tile_repeat[2] = times + recv_buffer = torch.tile(tensor[0][:, :, :1], tile_repeat) + + return recv_buffer diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/global_config.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/global_config.py new file mode 100644 index 0000000000000000000000000000000000000000..863117570a8aadde38b8eae8f1aa16480cd9f7ca --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/global_config.py @@ -0,0 +1,28 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from typing import Optional + +_NORM_LIMIT = float("inf") + + +def get_norm_limit(): + return _NORM_LIMIT + + +def set_norm_limit(value: Optional[float] = None): + global _NORM_LIMIT + if value is None: + value = float("inf") + _NORM_LIMIT = value diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/inflated_layers.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/inflated_layers.py new file mode 100644 index 0000000000000000000000000000000000000000..4b8e6dfb698e477f9af268f3dcca99a7b6837e63 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/inflated_layers.py @@ -0,0 +1,106 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from functools import partial +from typing import Literal, Optional +from torch import Tensor +from torch.nn import Conv3d + +from .inflated_lib import ( + MemoryState, + extend_head, + inflate_bias, + inflate_weight, + modify_state_dict, +) + +_inflation_mode_t = Literal["none", "tail", "replicate"] +_memory_device_t = Optional[Literal["cpu", "same"]] + + +class InflatedCausalConv3d(Conv3d): + def __init__( + self, + *args, + inflation_mode: _inflation_mode_t, + memory_device: _memory_device_t = "same", + **kwargs, + ): + self.inflation_mode = inflation_mode + self.memory = None + super().__init__(*args, **kwargs) + self.temporal_padding = self.padding[0] + self.memory_device = memory_device + self.padding = (0, *self.padding[1:]) # Remove temporal pad to keep causal. + + def set_memory_device(self, memory_device: _memory_device_t): + self.memory_device = memory_device + + def forward(self, input: Tensor, memory_state: MemoryState = MemoryState.DISABLED) -> Tensor: + mem_size = self.stride[0] - self.kernel_size[0] + if (self.memory is not None) and (memory_state == MemoryState.ACTIVE): + input = extend_head(input, memory=self.memory) + else: + input = extend_head(input, times=self.temporal_padding * 2) + memory = ( + input[:, :, mem_size:].detach() + if (mem_size != 0 and memory_state != MemoryState.DISABLED) + else None + ) + if ( + memory_state != MemoryState.DISABLED + and not self.training + and (self.memory_device is not None) + ): + self.memory = memory + if self.memory_device == "cpu" and self.memory is not None: + self.memory = self.memory.to("cpu") + return super().forward(input) + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + if self.inflation_mode != "none": + state_dict = modify_state_dict( + self, + state_dict, + prefix, + inflate_weight_fn=partial(inflate_weight, position="tail"), + inflate_bias_fn=partial(inflate_bias, position="tail"), + ) + super()._load_from_state_dict( + state_dict, + prefix, + local_metadata, + (strict and self.inflation_mode == "none"), + missing_keys, + unexpected_keys, + error_msgs, + ) + + +def init_causal_conv3d( + *args, + inflation_mode: _inflation_mode_t, + **kwargs, +): + """ + Initialize a Causal-3D convolution layer. + Parameters: + inflation_mode: Listed as below. It's compatible with all the 3D-VAE checkpoints we have. + - none: No inflation will be conducted. + The loading logic of state dict will fall back to default. + - tail / replicate: Refer to the definition of `InflatedCausalConv3d`. + """ + return InflatedCausalConv3d(*args, inflation_mode=inflation_mode, **kwargs) diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/inflated_lib.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/inflated_lib.py new file mode 100644 index 0000000000000000000000000000000000000000..486c63ca121f066585aff710e1b99d33e386f5e4 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/inflated_lib.py @@ -0,0 +1,156 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from enum import Enum +from typing import Optional +import numpy as np +import torch +from diffusers.models.normalization import RMSNorm +from einops import rearrange +from torch import Tensor, nn + +from ....common.logger import get_logger + +logger = get_logger(__name__) + + +class MemoryState(Enum): + """ + State[Disabled]: No memory bank will be enabled. + State[Initializing]: The model is handling the first clip, + need to reset / initialize the memory bank. + State[Active]: There has been some data in the memory bank. + """ + + DISABLED = 0 + INITIALIZING = 1 + ACTIVE = 2 + + +def causal_norm_wrapper(norm_layer: nn.Module, x: torch.Tensor) -> torch.Tensor: + if isinstance(norm_layer, (nn.LayerNorm, RMSNorm)): + if x.ndim == 4: + x = rearrange(x, "b c h w -> b h w c") + x = norm_layer(x) + x = rearrange(x, "b h w c -> b c h w") + return x + if x.ndim == 5: + x = rearrange(x, "b c t h w -> b t h w c") + x = norm_layer(x) + x = rearrange(x, "b t h w c -> b c t h w") + return x + if isinstance(norm_layer, (nn.GroupNorm, nn.BatchNorm2d, nn.SyncBatchNorm)): + if x.ndim <= 4: + return norm_layer(x) + if x.ndim == 5: + t = x.size(2) + x = rearrange(x, "b c t h w -> (b t) c h w") + x = norm_layer(x) + x = rearrange(x, "(b t) c h w -> b c t h w", t=t) + return x + raise NotImplementedError + + +def remove_head(tensor: Tensor, times: int = 1) -> Tensor: + """ + Remove duplicated first frame features in the up-sampling process. + """ + if times == 0: + return tensor + return torch.cat(tensors=(tensor[:, :, :1], tensor[:, :, times + 1 :]), dim=2) + + +def extend_head( + tensor: Tensor, times: Optional[int] = 2, memory: Optional[Tensor] = None +) -> Tensor: + """ + When memory is None: + - Duplicate first frame features in the down-sampling process. + When memory is not None: + - Concatenate memory features with the input features to keep temporal consistency. + """ + if times == 0: + return tensor + if memory is not None: + return torch.cat((memory.to(tensor), tensor), dim=2) + else: + tile_repeat = np.ones(tensor.ndim).astype(int) + tile_repeat[2] = times + return torch.cat(tensors=(torch.tile(tensor[:, :, :1], list(tile_repeat)), tensor), dim=2) + + +def inflate_weight(weight_2d: torch.Tensor, weight_3d: torch.Tensor, inflation_mode: str): + """ + Inflate a 2D convolution weight matrix to a 3D one. + Parameters: + weight_2d: The weight matrix of 2D conv to be inflated. + weight_3d: The weight matrix of 3D conv to be initialized. + inflation_mode: the mode of inflation + """ + assert inflation_mode in ["constant", "replicate"] + assert weight_3d.shape[:2] == weight_2d.shape[:2] + with torch.no_grad(): + if inflation_mode == "replicate": + depth = weight_3d.size(2) + weight_3d.copy_(weight_2d.unsqueeze(2).repeat(1, 1, depth, 1, 1) / depth) + else: + weight_3d.fill_(0.0) + weight_3d[:, :, -1].copy_(weight_2d) + return weight_3d + + +def inflate_bias(bias_2d: torch.Tensor, bias_3d: torch.Tensor, inflation_mode: str): + """ + Inflate a 2D convolution bias tensor to a 3D one + Parameters: + bias_2d: The bias tensor of 2D conv to be inflated. + bias_3d: The bias tensor of 3D conv to be initialized. + inflation_mode: Placeholder to align `inflate_weight`. + """ + assert bias_3d.shape == bias_2d.shape + with torch.no_grad(): + bias_3d.copy_(bias_2d) + return bias_3d + + +def modify_state_dict(layer, state_dict, prefix, inflate_weight_fn, inflate_bias_fn): + """ + the main function to inflated 2D parameters to 3D. + """ + weight_name = prefix + "weight" + bias_name = prefix + "bias" + if weight_name in state_dict: + weight_2d = state_dict[weight_name] + if weight_2d.dim() == 4: + # Assuming the 2D weights are 4D tensors (out_channels, in_channels, h, w) + weight_3d = inflate_weight_fn( + weight_2d=weight_2d, + weight_3d=layer.weight, + inflation_mode=layer.inflation_mode, + ) + state_dict[weight_name] = weight_3d + else: + return state_dict + # It's a 3d state dict, should not do inflation on both bias and weight. + if bias_name in state_dict: + bias_2d = state_dict[bias_name] + if bias_2d.dim() == 1: + # Assuming the 2D biases are 1D tensors (out_channels,) + bias_3d = inflate_bias_fn( + bias_2d=bias_2d, + bias_3d=layer.bias, + inflation_mode=layer.inflation_mode, + ) + state_dict[bias_name] = bias_3d + return state_dict diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/types.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/types.py new file mode 100644 index 0000000000000000000000000000000000000000..5a030d2d284f9535f2a84c1f9befcd3f82d8d9ff --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/types.py @@ -0,0 +1,76 @@ +# // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +# // +# // 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. + +from enum import Enum +from typing import Dict, Literal, NamedTuple, Optional +import torch + +_receptive_field_t = Literal["half", "full"] +_inflation_mode_t = Literal["none", "tail", "replicate"] +_memory_device_t = Optional[Literal["cpu", "same"]] +_gradient_checkpointing_t = Optional[Literal["half", "full"]] +_selective_checkpointing_t = Optional[Literal["coarse", "fine"]] + +class DiagonalGaussianDistribution: + def __init__(self, mean: torch.Tensor, logvar: torch.Tensor): + self.mean = mean + self.logvar = torch.clamp(logvar, -30.0, 20.0) + self.std = torch.exp(0.5 * self.logvar) + self.var = torch.exp(self.logvar) + + def mode(self) -> torch.Tensor: + return self.mean + + def sample(self) -> torch.FloatTensor: + return self.mean + self.std * torch.randn_like(self.mean) + + def kl(self) -> torch.Tensor: + return 0.5 * torch.sum( + self.mean**2 + self.var - 1.0 - self.logvar, + dim=list(range(1, self.mean.ndim)), + ) + +class MemoryState(Enum): + """ + State[Disabled]: No memory bank will be enabled. + State[Initializing]: The model is handling the first clip, need to reset the memory bank. + State[Active]: There has been some data in the memory bank. + State[Unset]: Error state, indicating users didn't pass correct memory state in. + """ + + DISABLED = 0 + INITIALIZING = 1 + ACTIVE = 2 + UNSET = 3 + + +class QuantizerOutput(NamedTuple): + latent: torch.Tensor + extra_loss: torch.Tensor + statistics: Dict[str, torch.Tensor] + + +class CausalAutoencoderOutput(NamedTuple): + sample: torch.Tensor + latent: torch.Tensor + posterior: Optional[DiagonalGaussianDistribution] + + +class CausalEncoderOutput(NamedTuple): + latent: torch.Tensor + posterior: Optional[DiagonalGaussianDistribution] + + +class CausalDecoderOutput(NamedTuple): + sample: torch.Tensor diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/video_vae.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/video_vae.py new file mode 100644 index 0000000000000000000000000000000000000000..08239016b2254b952298cb78f319cc8d82206709 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/modules/video_vae.py @@ -0,0 +1,956 @@ +# Copyright (c) 2023 HuggingFace Team +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates. +# SPDX-License-Identifier: Apache License, Version 2.0 (the "License") +# +# This file has been modified by ByteDance Ltd. and/or its affiliates. on 1st June 2025 +# +# Original file was released under Apache License, Version 2.0 (the "License"), with the full license text +# available at http://www.apache.org/licenses/LICENSE-2.0. +# +# This modified file is released under the same license. + +from contextlib import nullcontext +from typing import Optional, Tuple, Literal, Callable, Union + +import torch +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution +from einops import rearrange +from ....common.half_precision_fixes import safe_pad_operation + +from ....common.distributed.advanced import get_sequence_parallel_world_size +from ....common.logger import get_logger +from .causal_inflation_lib import ( + InflatedCausalConv3d, + causal_norm_wrapper, + init_causal_conv3d, + remove_head, +) +from .context_parallel_lib import ( + causal_conv_gather_outputs, + causal_conv_slice_inputs, +) +from .global_config import set_norm_limit +from .types import ( + CausalAutoencoderOutput, + CausalDecoderOutput, + CausalEncoderOutput, + MemoryState, + _inflation_mode_t, + _memory_device_t, + _receptive_field_t, + _selective_checkpointing_t, +) + +logger = get_logger(__name__) # pylint: disable=invalid-name + +# Fake func, no checkpointing is required for inference +def gradient_checkpointing(module: Union[Callable, nn.Module], *args, enabled: bool, **kwargs): + return module(*args, **kwargs) + +class ResnetBlock2D(nn.Module): + r""" + A Resnet block. + + Parameters: + in_channels (`int`): The number of channels in the input. + out_channels (`int`, *optional*, default to be `None`): + The number of output channels for the first conv2d layer. + If None, same as `in_channels`. + dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. + """ + + def __init__( + self, *, in_channels: int, out_channels: Optional[int] = None, dropout: float = 0.0 + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + + self.nonlinearity = nn.SiLU() + + self.norm1 = torch.nn.GroupNorm( + num_groups=32, num_channels=in_channels, eps=1e-6, affine=True + ) + + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + + self.norm2 = torch.nn.GroupNorm( + num_groups=32, num_channels=out_channels, eps=1e-6, affine=True + ) + + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + + self.use_in_shortcut = self.in_channels != out_channels + + self.conv_shortcut = None + if self.use_in_shortcut: + self.conv_shortcut = nn.Conv2d( + in_channels, out_channels, kernel_size=1, stride=1, padding=0 + ) + + def forward(self, input_tensor: torch.Tensor) -> torch.Tensor: + hidden = input_tensor + + hidden = self.norm1(hidden) + hidden = self.nonlinearity(hidden) + hidden = self.conv1(hidden) + + hidden = self.norm2(hidden) + hidden = self.nonlinearity(hidden) + hidden = self.dropout(hidden) + hidden = self.conv2(hidden) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor) + + output_tensor = input_tensor + hidden + + return output_tensor + +class Upsample3D(nn.Module): + """A 3D upsampling layer.""" + + def __init__( + self, + channels: int, + inflation_mode: _inflation_mode_t = "tail", + temporal_up: bool = False, + spatial_up: bool = True, + slicing: bool = False, + ): + super().__init__() + self.channels = channels + self.conv = init_causal_conv3d( + self.channels, self.channels, kernel_size=3, padding=1, inflation_mode=inflation_mode + ) + + self.temporal_up = temporal_up + self.spatial_up = spatial_up + self.temporal_ratio = 2 if temporal_up else 1 + self.spatial_ratio = 2 if spatial_up else 1 + self.slicing = slicing + + upscale_ratio = (self.spatial_ratio**2) * self.temporal_ratio + self.upscale_conv = nn.Conv3d( + self.channels, self.channels * upscale_ratio, kernel_size=1, padding=0 + ) + identity = ( + torch.eye(self.channels).repeat(upscale_ratio, 1).reshape_as(self.upscale_conv.weight) + ) + + self.upscale_conv.weight.data.copy_(identity) + nn.init.zeros_(self.upscale_conv.bias) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state: MemoryState, + ) -> torch.FloatTensor: + return gradient_checkpointing( + self.custom_forward, + hidden_states, + memory_state, + enabled=self.training and self.gradient_checkpointing, + ) + + def custom_forward( + self, + hidden_states: torch.FloatTensor, + memory_state: MemoryState, + ) -> torch.FloatTensor: + assert hidden_states.shape[1] == self.channels + + if self.slicing: + split_size = hidden_states.size(2) // 2 + hidden_states = list( + hidden_states.split([split_size, hidden_states.size(2) - split_size], dim=2) + ) + else: + hidden_states = [hidden_states] + + for i in range(len(hidden_states)): + hidden_states[i] = self.upscale_conv(hidden_states[i]) + hidden_states[i] = rearrange( + hidden_states[i], + "b (x y z c) f h w -> b c (f z) (h x) (w y)", + x=self.spatial_ratio, + y=self.spatial_ratio, + z=self.temporal_ratio, + ) + + # [Overridden] For causal temporal conv + if self.temporal_up and memory_state != MemoryState.ACTIVE: + hidden_states[0] = remove_head(hidden_states[0]) + + if self.slicing: + hidden_states = self.conv(hidden_states, memory_state=memory_state) + return torch.cat(hidden_states, dim=2) + else: + return self.conv(hidden_states[0], memory_state=memory_state) + + +class Downsample3D(nn.Module): + """A 3D downsampling layer.""" + + def __init__( + self, + channels: int, + inflation_mode: _inflation_mode_t = "tail", + temporal_down: bool = False, + spatial_down: bool = True, + ): + super().__init__() + self.channels = channels + self.temporal_down = temporal_down + self.spatial_down = spatial_down + + self.temporal_ratio = 2 if temporal_down else 1 + self.spatial_ratio = 2 if spatial_down else 1 + + self.temporal_kernel = 3 if temporal_down else 1 + self.spatial_kernel = 3 if spatial_down else 1 + + self.conv = init_causal_conv3d( + self.channels, + self.channels, + kernel_size=(self.temporal_kernel, self.spatial_kernel, self.spatial_kernel), + stride=(self.temporal_ratio, self.spatial_ratio, self.spatial_ratio), + padding=((1 if self.temporal_down else 0), 0, 0), + inflation_mode=inflation_mode, + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.FloatTensor, + memory_state: MemoryState, + ) -> torch.FloatTensor: + return gradient_checkpointing( + self.custom_forward, + hidden_states, + memory_state, + enabled=self.training and self.gradient_checkpointing, + ) + + def custom_forward( + self, + hidden_states: torch.FloatTensor, + memory_state: MemoryState, + ) -> torch.FloatTensor: + + assert hidden_states.shape[1] == self.channels + + if self.spatial_down: + hidden_states = safe_pad_operation(hidden_states, (0, 1, 0, 1), mode="constant", value=0) + + hidden_states = self.conv(hidden_states, memory_state=memory_state) + return hidden_states + + +class ResnetBlock3D(ResnetBlock2D): + def __init__( + self, + *args, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + **kwargs, + ): + super().__init__(*args, **kwargs) + self.conv1 = init_causal_conv3d( + self.in_channels, + self.out_channels, + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.conv2 = init_causal_conv3d( + self.out_channels, + self.out_channels, + kernel_size=(1, 3, 3) if time_receptive_field == "half" else (3, 3, 3), + stride=1, + padding=(0, 1, 1) if time_receptive_field == "half" else (1, 1, 1), + inflation_mode=inflation_mode, + ) + + if self.use_in_shortcut: + self.conv_shortcut = init_causal_conv3d( + self.in_channels, + self.out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=(self.conv_shortcut.bias is not None), + inflation_mode=inflation_mode, + ) + self.gradient_checkpointing = False + + def forward(self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET): + return gradient_checkpointing( + self.custom_forward, + input_tensor, + memory_state, + enabled=self.training and self.gradient_checkpointing, + ) + + def custom_forward( + self, input_tensor: torch.Tensor, memory_state: MemoryState = MemoryState.UNSET + ): + assert memory_state != MemoryState.UNSET + hidden_states = input_tensor + + hidden_states = causal_norm_wrapper(self.norm1, hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.conv1(hidden_states, memory_state=memory_state) + + hidden_states = causal_norm_wrapper(self.norm2, hidden_states) + hidden_states = self.nonlinearity(hidden_states) + hidden_states = self.dropout(hidden_states) + hidden_states = self.conv2(hidden_states, memory_state=memory_state) + + if self.conv_shortcut is not None: + input_tensor = self.conv_shortcut(input_tensor, memory_state=memory_state) + + output_tensor = input_tensor + hidden_states + + return output_tensor + + +class DownEncoderBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + add_downsample: bool = True, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_down: bool = True, + spatial_down: bool = True, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + in_channels = in_channels if i == 0 else out_channels + resnets.append( + ResnetBlock3D( + in_channels=in_channels, + out_channels=out_channels, + dropout=dropout, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + self.downsamplers = None + if add_downsample: + # Todo: Refactor this line before V5 Image VAE Training. + self.downsamplers = nn.ModuleList( + [ + Downsample3D( + channels=out_channels, + inflation_mode=inflation_mode, + temporal_down=temporal_down, + spatial_down=spatial_down, + ) + ] + ) + + def forward( + self, hidden_states: torch.FloatTensor, memory_state: MemoryState + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, memory_state=memory_state) + + if self.downsamplers is not None: + for downsampler in self.downsamplers: + hidden_states = downsampler(hidden_states, memory_state=memory_state) + + return hidden_states + + +class UpDecoderBlock3D(nn.Module): + def __init__( + self, + in_channels: int, + out_channels: int, + dropout: float = 0.0, + num_layers: int = 1, + add_upsample: bool = True, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_up: bool = True, + spatial_up: bool = True, + slicing: bool = False, + ): + super().__init__() + resnets = [] + + for i in range(num_layers): + input_channels = in_channels if i == 0 else out_channels + + resnets.append( + ResnetBlock3D( + in_channels=input_channels, + out_channels=out_channels, + dropout=dropout, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + ) + + self.resnets = nn.ModuleList(resnets) + + self.upsamplers = None + # Todo: Refactor this line before V5 Image VAE Training. + if add_upsample: + self.upsamplers = nn.ModuleList( + [ + Upsample3D( + channels=out_channels, + inflation_mode=inflation_mode, + temporal_up=temporal_up, + spatial_up=spatial_up, + slicing=slicing, + ) + ] + ) + + def forward( + self, hidden_states: torch.FloatTensor, memory_state: MemoryState + ) -> torch.FloatTensor: + for resnet in self.resnets: + hidden_states = resnet(hidden_states, memory_state=memory_state) + + if self.upsamplers is not None: + for upsampler in self.upsamplers: + hidden_states = upsampler(hidden_states, memory_state=memory_state) + + return hidden_states + + +class UNetMidBlock3D(nn.Module): + def __init__( + self, + channels: int, + dropout: float = 0.0, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + ): + super().__init__() + self.resnets = nn.ModuleList( + [ + ResnetBlock3D( + in_channels=channels, + out_channels=channels, + dropout=dropout, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ), + ResnetBlock3D( + in_channels=channels, + out_channels=channels, + dropout=dropout, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ), + ] + ) + + def forward(self, hidden_states: torch.Tensor, memory_state: MemoryState): + for resnet in self.resnets: + hidden_states = resnet(hidden_states, memory_state) + return hidden_states + + +class Encoder3D(nn.Module): + r""" + The `Encoder` layer of a variational autoencoder that encodes + its input into a latent representation. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + double_z: bool = True, + temporal_down_num: int = 2, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",), + ): + super().__init__() + self.layers_per_block = layers_per_block + + self.temporal_down_num = temporal_down_num + + self.conv_in = init_causal_conv3d( + in_channels, + block_out_channels[0], + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.down_blocks = nn.ModuleList([]) + + # down + output_channel = block_out_channels[0] + for i in range(len(block_out_channels)): + input_channel = output_channel + output_channel = block_out_channels[i] + is_final_block = i == len(block_out_channels) - 1 + is_temporal_down_block = i >= len(block_out_channels) - self.temporal_down_num - 1 + # Note: take the last one + + down_block = DownEncoderBlock3D( + num_layers=self.layers_per_block, + in_channels=input_channel, + out_channels=output_channel, + add_downsample=not is_final_block, + temporal_down=is_temporal_down_block, + spatial_down=True, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + self.down_blocks.append(down_block) + + # mid + self.mid_block = UNetMidBlock3D( + channels=block_out_channels[-1], + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + # out + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[-1], num_groups=32, eps=1e-6 + ) + self.conv_act = nn.SiLU() + + conv_out_channels = 2 * out_channels if double_z else out_channels + self.conv_out = init_causal_conv3d( + block_out_channels[-1], conv_out_channels, 3, padding=1, inflation_mode=inflation_mode + ) + + assert len(selective_checkpointing) == len(self.down_blocks) + self.set_gradient_checkpointing(selective_checkpointing) + + def set_gradient_checkpointing(self, checkpointing_types): + gradient_checkpointing = [] + for down_block, sac_type in zip(self.down_blocks, checkpointing_types): + if sac_type == "coarse": + gradient_checkpointing.append(True) + elif sac_type == "fine": + for n, m in down_block.named_modules(): + if hasattr(m, "gradient_checkpointing"): + m.gradient_checkpointing = True + logger.debug(f"set gradient_checkpointing: {n}") + gradient_checkpointing.append(False) + else: + gradient_checkpointing.append(False) + self.gradient_checkpointing = gradient_checkpointing + logger.info(f"[Encoder3D] gradient_checkpointing: {checkpointing_types}") + + def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor: + r"""The forward method of the `Encoder` class.""" + sample = self.conv_in(sample, memory_state=memory_state) + # down + for down_block, sac in zip(self.down_blocks, self.gradient_checkpointing): + sample = gradient_checkpointing( + down_block, + sample, + memory_state=memory_state, + enabled=self.training and sac, + ) + + # middle + sample = self.mid_block(sample, memory_state=memory_state) + + # post-process + sample = causal_norm_wrapper(self.conv_norm_out, sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample, memory_state=memory_state) + + return sample + + +class Decoder3D(nn.Module): + r""" + The `Decoder` layer of a variational autoencoder that + decodes its latent representation into an output sample. + """ + + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + block_out_channels: Tuple[int, ...] = (64,), + layers_per_block: int = 2, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + temporal_up_num: int = 2, + slicing_up_num: int = 0, + selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",), + ): + super().__init__() + self.layers_per_block = layers_per_block + self.temporal_up_num = temporal_up_num + + self.conv_in = init_causal_conv3d( + in_channels, + block_out_channels[-1], + kernel_size=3, + stride=1, + padding=1, + inflation_mode=inflation_mode, + ) + + self.up_blocks = nn.ModuleList([]) + + # mid + self.mid_block = UNetMidBlock3D( + channels=block_out_channels[-1], + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + # up + reversed_block_out_channels = list(reversed(block_out_channels)) + output_channel = reversed_block_out_channels[0] + for i in range(len(reversed_block_out_channels)): + prev_output_channel = output_channel + output_channel = reversed_block_out_channels[i] + + is_final_block = i == len(block_out_channels) - 1 + is_temporal_up_block = i < self.temporal_up_num + is_slicing_up_block = i >= len(block_out_channels) - slicing_up_num + # Note: Keep symmetric + + up_block = UpDecoderBlock3D( + num_layers=self.layers_per_block + 1, + in_channels=prev_output_channel, + out_channels=output_channel, + add_upsample=not is_final_block, + temporal_up=is_temporal_up_block, + slicing=is_slicing_up_block, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + self.up_blocks.append(up_block) + + # out + self.conv_norm_out = nn.GroupNorm( + num_channels=block_out_channels[0], num_groups=32, eps=1e-6 + ) + self.conv_act = nn.SiLU() + self.conv_out = init_causal_conv3d( + block_out_channels[0], out_channels, 3, padding=1, inflation_mode=inflation_mode + ) + + assert len(selective_checkpointing) == len(self.up_blocks) + self.set_gradient_checkpointing(selective_checkpointing) + + def set_gradient_checkpointing(self, checkpointing_types): + gradient_checkpointing = [] + for up_block, sac_type in zip(self.up_blocks, checkpointing_types): + if sac_type == "coarse": + gradient_checkpointing.append(True) + elif sac_type == "fine": + for n, m in up_block.named_modules(): + if hasattr(m, "gradient_checkpointing"): + m.gradient_checkpointing = True + logger.debug(f"set gradient_checkpointing: {n}") + gradient_checkpointing.append(False) + else: + gradient_checkpointing.append(False) + self.gradient_checkpointing = gradient_checkpointing + logger.info(f"[Decoder3D] gradient_checkpointing: {checkpointing_types}") + + def forward(self, sample: torch.FloatTensor, memory_state: MemoryState) -> torch.FloatTensor: + r"""The forward method of the `Decoder` class.""" + + sample = self.conv_in(sample, memory_state=memory_state) + + # middle + sample = self.mid_block(sample, memory_state=memory_state) + + # up + for up_block, sac in zip(self.up_blocks, self.gradient_checkpointing): + sample = gradient_checkpointing( + up_block, + sample, + memory_state=memory_state, + enabled=self.training and sac, + ) + + # post-process + sample = causal_norm_wrapper(self.conv_norm_out, sample) + sample = self.conv_act(sample) + sample = self.conv_out(sample, memory_state=memory_state) + + return sample + + +class VideoAutoencoderKL(nn.Module): + def __init__( + self, + in_channels: int = 3, + out_channels: int = 3, + block_out_channels: Tuple[int] = (64,), + layers_per_block: int = 1, + latent_channels: int = 4, + use_quant_conv: bool = True, + use_post_quant_conv: bool = True, + enc_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",), + dec_selective_checkpointing: Tuple[_selective_checkpointing_t] = ("none",), + temporal_scale_num: int = 3, + slicing_up_num: int = 0, + inflation_mode: _inflation_mode_t = "tail", + time_receptive_field: _receptive_field_t = "half", + slicing_sample_min_size: int = None, + spatial_downsample_factor: int = 16, + temporal_downsample_factor: int = 8, + freeze_encoder: bool = False, + ): + super().__init__() + self.spatial_downsample_factor = spatial_downsample_factor + self.temporal_downsample_factor = temporal_downsample_factor + self.freeze_encoder = freeze_encoder + 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 = max(1, slicing_sample_min_size // (2**temporal_scale_num)) + + # pass init params to Encoder + self.encoder = Encoder3D( + in_channels=in_channels, + out_channels=latent_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + double_z=True, + temporal_down_num=temporal_scale_num, + selective_checkpointing=enc_selective_checkpointing, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + # pass init params to Decoder + self.decoder = Decoder3D( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=block_out_channels, + layers_per_block=layers_per_block, + # [Override] add temporal_up_num parameter + temporal_up_num=temporal_scale_num, + slicing_up_num=slicing_up_num, + selective_checkpointing=dec_selective_checkpointing, + inflation_mode=inflation_mode, + time_receptive_field=time_receptive_field, + ) + + self.quant_conv = ( + init_causal_conv3d( + in_channels=2 * latent_channels, + out_channels=2 * latent_channels, + kernel_size=1, + inflation_mode=inflation_mode, + ) + if use_quant_conv + else None + ) + self.post_quant_conv = ( + init_causal_conv3d( + in_channels=latent_channels, + out_channels=latent_channels, + kernel_size=1, + inflation_mode=inflation_mode, + ) + if use_post_quant_conv + else None + ) + + self.use_slicing = False + + def enable_slicing(self): + self.use_slicing = True + + def disable_slicing(self): + self.use_slicing = False + + def encode(self, x: torch.FloatTensor) -> CausalEncoderOutput: + if x.ndim == 4: + x = x.unsqueeze(2) + h = self.slicing_encode(x) + p = DiagonalGaussianDistribution(h) + z = p.sample() + return CausalEncoderOutput(z, p) + + def decode(self, z: torch.FloatTensor) -> CausalDecoderOutput: + if z.ndim == 4: + z = z.unsqueeze(2) + x = self.slicing_decode(z) + return CausalDecoderOutput(x) + + def _encode(self, x: torch.Tensor, memory_state: MemoryState) -> torch.Tensor: + x = causal_conv_slice_inputs(x, self.slicing_sample_min_size, memory_state=memory_state) + h = self.encoder(x, memory_state=memory_state) + h = self.quant_conv(h, memory_state=memory_state) if self.quant_conv is not None else h + h = causal_conv_gather_outputs(h) + return h + + def _decode(self, z: torch.Tensor, memory_state: MemoryState) -> torch.Tensor: + z = causal_conv_slice_inputs(z, self.slicing_latent_min_size, memory_state=memory_state) + z = ( + self.post_quant_conv(z, memory_state=memory_state) + if self.post_quant_conv is not None + else z + ) + x = self.decoder(z, memory_state=memory_state) + x = causal_conv_gather_outputs(x) + return x + + def slicing_encode(self, x: torch.Tensor) -> torch.Tensor: + sp_size = get_sequence_parallel_world_size() + if self.use_slicing and (x.shape[2] - 1) > self.slicing_sample_min_size * sp_size: + x_slices = x[:, :, 1:].split(split_size=self.slicing_sample_min_size * sp_size, dim=2) + encoded_slices = [ + self._encode( + torch.cat((x[:, :, :1], x_slices[0]), dim=2), + memory_state=MemoryState.INITIALIZING, + ) + ] + for x_idx in range(1, len(x_slices)): + encoded_slices.append( + self._encode(x_slices[x_idx], memory_state=MemoryState.ACTIVE) + ) + return torch.cat(encoded_slices, dim=2) + else: + return self._encode(x, memory_state=MemoryState.DISABLED) + + def slicing_decode(self, z: torch.Tensor) -> torch.Tensor: + sp_size = get_sequence_parallel_world_size() + if self.use_slicing and (z.shape[2] - 1) > self.slicing_latent_min_size * sp_size: + z_slices = z[:, :, 1:].split(split_size=self.slicing_latent_min_size * sp_size, dim=2) + decoded_slices = [ + self._decode( + torch.cat((z[:, :, :1], z_slices[0]), dim=2), + memory_state=MemoryState.INITIALIZING, + ) + ] + for z_idx in range(1, len(z_slices)): + decoded_slices.append( + self._decode(z_slices[z_idx], memory_state=MemoryState.ACTIVE) + ) + return torch.cat(decoded_slices, dim=2) + else: + return self._decode(z, memory_state=MemoryState.DISABLED) + + def forward(self, x: torch.FloatTensor) -> CausalAutoencoderOutput: + with torch.no_grad() if self.freeze_encoder else nullcontext(): + z, p = self.encode(x) + x = self.decode(z).sample + return CausalAutoencoderOutput(x, z, p) + + def preprocess(self, x: torch.Tensor): + # x should in [B, C, T, H, W], [B, C, H, W] + assert x.ndim == 4 or x.size(2) % self.temporal_downsample_factor == 1 + return x + + def postprocess(self, x: torch.Tensor): + # x should in [B, C, T, H, W], [B, C, H, W] + return x + + def set_causal_slicing( + self, + *, + split_size: Optional[int], + memory_device: _memory_device_t, + ): + assert ( + split_size is None or memory_device is not None + ), "if split_size is set, memory_device must not be None." + if split_size is not None: + self.enable_slicing() + self.slicing_sample_min_size = split_size + self.slicing_latent_min_size = max(1, split_size // self.temporal_downsample_factor) + else: + self.disable_slicing() + for module in self.modules(): + if isinstance(module, InflatedCausalConv3d): + module.set_memory_device(memory_device) + + def set_memory_limit(self, conv_max_mem: Optional[float], norm_max_mem: Optional[float]): + set_norm_limit(norm_max_mem) + for m in self.modules(): + if isinstance(m, InflatedCausalConv3d): + m.set_memory_limit(conv_max_mem if conv_max_mem is not None else float("inf")) + + +class VideoAutoencoderKLWrapper(VideoAutoencoderKL): + def __init__( + self, *args, spatial_downsample_factor: int, temporal_downsample_factor: int, **kwargs + ): + self.spatial_downsample_factor = spatial_downsample_factor + self.temporal_downsample_factor = temporal_downsample_factor + super().__init__(*args, **kwargs) + + def forward(self, x) -> CausalAutoencoderOutput: + z, _, p = self.encode(x) + x, _ = self.decode(z) + return CausalAutoencoderOutput(x, z, None, p) + + def encode(self, x) -> CausalEncoderOutput: + if x.ndim == 4: + x = x.unsqueeze(2) + p = super().encode(x).latent_dist + z = p.sample().squeeze(2) + return CausalEncoderOutput(z, None, p) + + def decode(self, z) -> CausalDecoderOutput: + if z.ndim == 4: + z = z.unsqueeze(2) + x = super().decode(z).sample.squeeze(2) + return CausalDecoderOutput(x, None) + + def preprocess(self, x): + # x should in [B, C, T, H, W], [B, C, H, W] + assert x.ndim == 4 or x.size(2) % 4 == 1 + return x + + def postprocess(self, x): + # x should in [B, C, T, H, W], [B, C, H, W] + return x + + def set_causal_slicing( + self, + *, + split_size: Optional[int], + memory_device: Optional[Literal["cpu", "same"]], + ): + assert ( + split_size is None or memory_device is not None + ), "if split_size is set, memory_device must not be None." + if split_size is not None: + self.enable_slicing() + else: + self.disable_slicing() + self.slicing_sample_min_size = split_size + if split_size is not None: + 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) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml new file mode 100644 index 0000000000000000000000000000000000000000..949122667aacd5937b82bcbf9a2abc54a2e5ce1c --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/models/video_vae_v3/s8_c16_t4_inflation_sd3.yaml @@ -0,0 +1,28 @@ +act_fn: silu +block_out_channels: + - 128 + - 256 + - 512 + - 512 +down_block_types: + - DownEncoderBlock3D + - DownEncoderBlock3D + - DownEncoderBlock3D + - DownEncoderBlock3D +in_channels: 3 +latent_channels: 16 +layers_per_block: 2 +norm_num_groups: 32 +out_channels: 3 +slicing_sample_min_size: 4 +temporal_scale_num: 2 +inflation_mode: pad +up_block_types: + - UpDecoderBlock3D + - UpDecoderBlock3D + - UpDecoderBlock3D + - UpDecoderBlock3D +spatial_downsample_factor: 8 +temporal_downsample_factor: 4 +use_quant_conv: False +use_post_quant_conv: False diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/blockswap.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/blockswap.py new file mode 100644 index 0000000000000000000000000000000000000000..0ecbdb627efa403709e60177b34b93ed2c6d04c7 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/blockswap.py @@ -0,0 +1,938 @@ +""" +BlockSwap Module for SeedVR2 + +This module implements dynamic block swapping between GPU and CPU memory +to enable running large models on limited VRAM systems. + +Key Features: +- Dynamic transformer block offloading during inference +- Non-blocking GPU transfers for optimal performance +- RoPE computation fallback to CPU on OOM +- Minimal performance overhead with intelligent caching +- I/O component offloading for maximum memory savings +""" + +import time +import types +import torch +import weakref + +from typing import Dict, Any, List, Optional +from .memory_manager import clear_memory +from .compatibility import call_rope_with_stability +from ..common.distributed import get_device + + +def is_blockswap_enabled(config: Optional[Dict[str, Any]]) -> bool: + """ + Check if BlockSwap configuration indicates BlockSwap should be enabled. + + BlockSwap is enabled if either blocks_to_swap > 0 OR swap_io_components is True. + This is the authoritative function for determining BlockSwap status from configuration. + + Args: + config: BlockSwap configuration dictionary with optional keys: + - blocks_to_swap: Number of blocks to offload (0 = disabled) + - swap_io_components: Whether to offload I/O components + + Returns: + True if BlockSwap should be active, False otherwise + """ + if not config: + return False + + blocks_to_swap = config.get("blocks_to_swap", 0) + swap_io_components = config.get("swap_io_components", False) + + 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 +@torch._dynamo.disable +def _get_swap_start_time(debug, enabled: bool) -> Optional[float]: + """Get start time for swap operation if debug is enabled.""" + return time.time() if debug and enabled else None + + +@torch._dynamo.disable +def _log_swap_timing(debug, t_start: Optional[float], component_id, component_type: str) -> None: + """Log swap timing if start time was captured.""" + if debug and t_start is not None: + debug.log_swap_time( + component_id=component_id, + duration=time.time() - t_start, + component_type=component_type + ) + + +def get_module_memory_mb(module: torch.nn.Module) -> float: + """ + Calculate memory usage of a module in MB. + + Args: + module: PyTorch module to measure + + Returns: + Memory usage in megabytes + """ + total_bytes = sum( + param.nelement() * param.element_size() + for param in module.parameters() + if param.data is not None + ) + return total_bytes / (1024 * 1024) + + +def apply_block_swap_to_dit( + runner: 'VideoDiffusionInfer', + block_swap_config: Dict[str, Any], + debug: 'Debug' +) -> None: + """ + Apply block swapping configuration to a DiT model with OOM protection. + + This is the main entry point for configuring block swapping on a model. + Handles block selection, I/O component offloading, device placement, and + forward method wrapping for dynamic memory management. + + Args: + runner: VideoDiffusionInfer instance containing the model + block_swap_config: Configuration dictionary with keys: + - blocks_to_swap: Number of blocks to swap (from the start) + - swap_io_components: Whether to offload I/O components + - enable_debug: Whether to enable debug logging + - offload_device: Device to offload to (default: 'cpu') + debug: Debug instance for logging (required) + """ + # Early return if BlockSwap not enabled + if not is_blockswap_enabled(block_swap_config): + return + + blocks_to_swap = block_swap_config.get("blocks_to_swap", 0) + swap_io_components = block_swap_config.get("swap_io_components", False) + + # Early return only if both block swap and I/O swap are disabled + if blocks_to_swap <= 0 and not swap_io_components: + return + + if debug is None: + if hasattr(runner, 'debug') and runner.debug is not None: + debug = runner.debug + else: + raise ValueError("Debug instance must be provided to apply_block_swap_to_dit") + + debug.start_timer("apply_blockswap") + + # Get the actual model (handle CompatibleDiT wrapper) + model = runner.dit + if hasattr(model, "dit_model"): + model = model.dit_model + + # Determine devices + if hasattr(runner, '_dit_device'): + device = runner._dit_device + else: + device = get_device() + offload_device = block_swap_config.get("offload_device", torch.device('cpu')) + + # Validate model structure + if not hasattr(model, "blocks"): + debug.log("Model doesn't have 'blocks' attribute for BlockSwap", level="ERROR", category="blockswap", force=True) + return + + total_blocks = len(model.blocks) + + # Clamp blocks_to_swap to available blocks BEFORE logging + effective_blocks = min(blocks_to_swap, total_blocks) if blocks_to_swap > 0 else 0 + + # 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: + model.blocks_to_swap = effective_blocks - 1 # Convert to 0-indexed + else: + # No block swapping, set to -1 so no blocks match the swap condition + model.blocks_to_swap = -1 + + model.main_device = device + model.offload_device = offload_device + + # Configure I/O components + io_config = _configure_io_components(model, device, offload_device, + swap_io_components, debug) + memory_stats = _configure_blocks(model, device, offload_device, debug) + memory_stats['io_components'] = io_config['components'] + memory_stats['io_memory_mb'] = io_config['memory_mb'] + memory_stats['gpu_components'] = io_config['gpu_components'] + memory_stats['io_gpu_memory_mb'] = io_config['gpu_memory_mb'] + + # Log memory summary + _log_memory_summary(memory_stats, offload_device, device, swap_io_components, + debug) + + # Wrap block forward methods for dynamic swapping (only if blocks_to_swap > 0) + if blocks_to_swap > 0: + for b, block in enumerate(model.blocks): + if b <= model.blocks_to_swap: + _wrap_block_forward(block, b, model, debug) + + # Patch RoPE modules for robust error handling + _patch_rope_for_blockswap(model, debug) + + # Mark BlockSwap as active + runner._blockswap_active = True + + # Store configuration for debugging and cleanup + model._block_swap_config = { + "blocks_swapped": blocks_to_swap, + "swap_io_components": swap_io_components, + "total_blocks": total_blocks, + "offload_device": offload_device, + "main_device": device, + "offload_memory": memory_stats['offload_memory'], + "main_memory": memory_stats['main_memory'] + } + + # Protect model from being moved entirely + _protect_model_from_move(model, runner, debug) + + debug.log("BlockSwap configuration complete", category="success") + debug.end_timer("apply_blockswap", "BlockSwap configuration application") + + +def _configure_io_components( + model: torch.nn.Module, + device: torch.device, + offload_device: torch.device, + swap_io_components: bool, + debug: 'Debug' +) -> Dict[str, Any]: + """ + Configure I/O component placement and wrapping with memory tracking. + + Handles all non-block modules (embeddings, normalization layers, etc.) by + either keeping them on GPU or offloading them with dynamic swapping wrappers. + + Args: + model: DiT model containing named children to configure + device: Main computation device (typically GPU) + offload_device: Device for offloaded components (typically CPU) + swap_io_components: If True, offload I/O components with dynamic swapping + debug: Debug instance for logging (required) + + Returns: + Dictionary containing: + - components: List of offloaded component names + - memory_mb: Total memory of offloaded components in MB + - gpu_components: List of components remaining on GPU + - gpu_memory_mb: Total memory of GPU components in MB + """ + io_components_offloaded = [] + io_components_on_gpu = [] + io_memory_mb = 0.0 + io_gpu_memory_mb = 0.0 + + # Handle I/O modules with dynamic swapping + for name, module in model.named_children(): + if name != "blocks": + module_memory = get_module_memory_mb(module) + + if swap_io_components: + module.to(offload_device) + _wrap_io_forward(module, name, model, debug) + io_components_offloaded.append(name) + io_memory_mb += module_memory + debug.log(f"{name} โ†’ {str(offload_device).upper()} ({module_memory:.2f}MB, dynamic swapping)", category="blockswap", indent_level=1) + else: + module.to(device) + io_components_on_gpu.append(name) + io_gpu_memory_mb += module_memory + debug.log(f"{name} โ†’ {str(device).upper()} ({module_memory:.2f}MB)", category="blockswap", indent_level=1) + + return { + 'components': io_components_offloaded, + 'memory_mb': io_memory_mb, + 'gpu_components': io_components_on_gpu, + 'gpu_memory_mb': io_gpu_memory_mb + } + + +def _configure_blocks( + model: torch.nn.Module, + device: torch.device, + offload_device: torch.device, + debug: 'Debug' +) -> Dict[str, float]: + """ + Configure transformer block placement and calculate memory statistics. + + Moves blocks to their designated devices based on model.blocks_to_swap + attribute. Blocks with index <= blocks_to_swap go to offload device, + others stay on main device. + + Args: + model: DiT model with blocks attribute and blocks_to_swap configured + device: Main computation device for non-swapped blocks + offload_device: Device for swapped blocks + debug: Debug instance for logging (required) + + Returns: + Dictionary containing: + - offload_memory: Total memory of offloaded blocks in MB + - main_memory: Total memory of blocks on main device in MB + - io_components: Empty list (populated by caller) + """ + total_offload_memory = 0.0 + total_main_memory = 0.0 + + # Move blocks based on swap configuration + for b, block in enumerate(model.blocks): + block_memory = get_module_memory_mb(block) + + if b > model.blocks_to_swap: + block.to(device) + total_main_memory += block_memory + else: + block.to(offload_device, non_blocking=False) + total_offload_memory += block_memory + + # Ensure all buffers match their containing module's device + for b, block in enumerate(model.blocks): + target_device = device if b > model.blocks_to_swap else offload_device + for name, buffer in block.named_buffers(): + if buffer.device != torch.device(target_device): + buffer.data = buffer.data.to(target_device, non_blocking=False) + + return { + "offload_memory": total_offload_memory, + "main_memory": total_main_memory, + "io_components": [] # Will be populated by caller + } + + +def _log_memory_summary( + memory_stats: Dict[str, float], + offload_device: torch.device, + device: torch.device, + swap_io_components: bool, + debug: 'Debug' +) -> None: + """ + Log comprehensive memory usage summary for BlockSwap configuration. + + Displays detailed breakdown of memory distribution across devices, + including transformer blocks and I/O components. + + Args: + memory_stats: Dictionary containing: + - offload_memory: Memory offloaded from blocks (MB) + - main_memory: Memory remaining on main device (MB) + - io_memory_mb: Memory from offloaded I/O components (MB) + - io_gpu_memory_mb: Memory from I/O components on GPU (MB) + offload_device: Device used for offloading + device: Main computation device + swap_io_components: Whether I/O components are being swapped + debug: Debug instance for logging (required) + """ + debug.log("BlockSwap memory configuration:", category="blockswap") + + # Log transformer blocks memory + blocks_offloaded = memory_stats['offload_memory'] + blocks_on_gpu = memory_stats['main_memory'] + + offload_str = str(offload_device) + device_str = str(device) + + if blocks_on_gpu == 0: + debug.log(f"Transformer blocks: {blocks_offloaded:.2f}MB on {offload_str} (dynamic swapping)", category="blockswap", indent_level=1) + else: + debug.log(f"Transformer blocks: {blocks_on_gpu:.2f}MB on {device_str}, {blocks_offloaded:.2f}MB on {offload_str}", category="blockswap", indent_level=1) + + # Always log I/O components (whether swapping or not) + io_memory = memory_stats.get('io_memory_mb', 0.0) + io_gpu_memory = memory_stats.get('io_gpu_memory_mb', 0.0) + + if swap_io_components and io_memory > 0: + io_components = memory_stats.get('io_components', []) + debug.log(f"I/O components: {io_memory:.2f}MB on {offload_str} (dynamic swapping)", category="blockswap", indent_level=1) + debug.log(f"{', '.join(io_components)}", category="blockswap", indent_level=2) + elif io_gpu_memory > 0: + io_gpu_components = memory_stats.get('gpu_components', []) + debug.log(f"I/O components: {io_gpu_memory:.2f}MB on {device_str}", category="blockswap", indent_level=1) + debug.log(f"{', '.join(io_gpu_components)}", category="blockswap", indent_level=2) + + # Log total VRAM savings + total_offloaded = blocks_offloaded + (io_memory if swap_io_components else 0) + if total_offloaded > 0: + debug.log(f"Total VRAM saved: {total_offloaded:.2f}MB (~{total_offloaded/1024:.2f}GB)", category="blockswap", indent_level=1) + + +def _wrap_block_forward( + block: torch.nn.Module, + block_idx: int, + model: torch.nn.Module, + debug: 'Debug' +) -> None: + """ + Wrap individual transformer block forward for dynamic device swapping. + + Creates a wrapped forward method that automatically: + 1. Moves block to GPU before computation + 2. Executes original forward pass + 3. Moves block back to offload device after computation + 4. Logs timing and manages memory pressure + + Uses weak references to prevent memory leaks from closure retention. + + Args: + block: Individual transformer block to wrap + block_idx: Index of this block in model.blocks + model: Parent DiT model (used for device references) + debug: Debug instance for logging (required) + """ + if hasattr(block, '_original_forward'): + return # Already wrapped + + # Store original forward method + original_forward = block.forward + + # Create weak references + model_ref = weakref.ref(model) + debug_ref = weakref.ref(debug) + + # Store block_idx on the block itself to avoid closure issues + block._block_idx = block_idx + + def wrapped_forward(self, *args, **kwargs): + # Retrieve weak references + model = model_ref() + debug = debug_ref() + + if not model: + # Model has been garbage collected, fall back to original + return original_forward(*args, **kwargs) + + # Check if block swap is active for this block + if hasattr(model, 'blocks_to_swap') and self._block_idx <= model.blocks_to_swap: + # Use dynamo-disabled helper to get start time (avoids compilation warnings) + t_start = _get_swap_start_time(debug, debug.enabled if debug else False) + + # Only move to GPU if necessary + current_device = next(self.parameters()).device + target_device = torch.device(model.main_device) + + if current_device != target_device: + self.to(model.main_device, non_blocking=False) + + # Execute forward pass with OOM protection + output = original_forward(*args, **kwargs) + + # Move back to offload device + self.to(model.offload_device, non_blocking=False) + + # Use dynamo-disabled helper to log timing (avoids compilation warnings) + _log_swap_timing(debug, t_start, self._block_idx, "block") + + # Only clear cache under memory pressure + clear_memory(debug=debug, deep=False, force=False, timer_name="wrap_block_forward") + else: + output = original_forward(*args, **kwargs) + + return output + + # Bind the wrapped function as a method to the block + block.forward = types.MethodType(wrapped_forward, block) + + # Store reference to original forward for cleanup + block._original_forward = original_forward + + +def _wrap_io_forward( + module: torch.nn.Module, + module_name: str, + model: torch.nn.Module, + debug: 'Debug' +) -> None: + """ + Wrap I/O component forward for dynamic device swapping. + + Similar to _wrap_block_forward but for I/O components (embeddings, + normalization layers, etc.). Handles swapping between GPU and CPU + during forward passes. + + Uses weak references to prevent circular dependencies and memory leaks. + + Args: + module: I/O component module to wrap + module_name: Name identifier for logging (e.g., 'x_embedder') + model: Parent DiT model (used for device references) + debug: Debug instance for logging (required) + """ + if hasattr(module, '_is_io_wrapped') and module._is_io_wrapped: + debug.log(f"Reusing existing I/O wrapper for {module_name}", category="reuse") + return # Already wrapped + + # Store original forward method + original_forward = module.forward + + # Create weak references + model_ref = weakref.ref(model) + debug_ref = weakref.ref(debug) if debug else lambda: None + + # Store module name on the module itself + module._module_name = module_name + module._original_forward = original_forward + + def wrapped_io_forward(self, *args, **kwargs): + # Retrieve weak references + model = model_ref() + debug = debug_ref() + + if not model: + # Model has been garbage collected, fall back to original + return self._original_forward(*args, **kwargs) + + # Use dynamo-disabled helper to get start time (avoids compilation warnings) + t_start = _get_swap_start_time(debug, debug.enabled if debug else False) + + # Check current device to avoid unnecessary moves + current_device = next(self.parameters()).device + target_device = torch.device(model.main_device) + + # Move to GPU for computation if needed + if current_device != target_device: + self.to(model.main_device, non_blocking=False) + + # Execute forward pass + output = self._original_forward(*args, **kwargs) + + # Move back to offload device + self.to(model.offload_device, non_blocking=False) + + # Use dynamo-disabled helper to log timing (avoids compilation warnings) + _log_swap_timing(debug, t_start, self._module_name, "I/O") + + # Only clear cache under memory pressure + clear_memory(debug=debug, deep=False, force=False, timer_name="wrap_block_forward") + + return output + + # Bind as a method + module.forward = types.MethodType(wrapped_io_forward, module) + module._is_io_wrapped = True + + # Store module reference for restoration + if not hasattr(model, '_io_swappers'): + model._io_swappers = [] + model._io_swappers.append((module, module_name)) + + +def _patch_rope_for_blockswap( + model: torch.nn.Module, + debug: 'Debug' +) -> None: + """ + Patch RoPE (Rotary Position Embedding) modules for device-aware fallback. + + Adds CPU fallback logic to RoPE modules to handle device mismatch errors + that can occur during BlockSwap operations. Complements the stability + wrapper from compatibility.py with device-specific error handling. + + Args: + model: DiT model containing RoPE modules to patch + debug: Debug instance for logging (required) + """ + rope_patches = [] + + for name, module in model.named_modules(): + if "rope" in name.lower() and hasattr(module, "get_axial_freqs"): + # Skip if already wrapped by blockswap + if hasattr(module, '_blockswap_wrapped') and module._blockswap_wrapped: + continue + + # Get current method (might be stability-wrapped) + current_method = module.get_axial_freqs + + # Create device-aware wrapper with proper closure handling + def make_device_aware_wrapper(module_name, current_fn): + def device_aware_rope_wrapper(self, *args, **kwargs): + try: + # Try current method (original or stability-wrapped) + return current_fn(*args, **kwargs) + except (RuntimeError, torch.cuda.OutOfMemoryError) as e: + error_msg = str(e).lower() + # Only handle device/memory specific errors + if any(x in error_msg for x in ["device", "memory", "allocation"]): + debug.log(f"RoPE OOM for {module_name}", level="WARNING", category="rope", force=True) + debug.log(f"Clearing RoPE cache and retrying", category="info", force=True) + + # Get current device from parameters + try: + current_device = next(self.parameters()).device + except StopIteration: + # Fallback: use model's main_device if BlockSwap has set it, else use offload_device + if hasattr(model, 'main_device'): + current_device = torch.device(model.main_device) + elif hasattr(model, 'offload_device'): + current_device = torch.device(model.offload_device) + + # Try clearing cache first (non-invasive fix) + if hasattr(current_fn, 'cache_clear'): + current_fn.cache_clear() + try: + # Retry on same device after clearing cache + return current_fn(*args, **kwargs) + except Exception as retry_error: + # Cache clear wasn't enough, need more drastic measures + debug.log(f"Cache clear insufficient for {module_name}, falling back to CPU", level="WARNING", category="rope", force=True) + + # Fallback to CPU computation with stability + self.cpu() + + try: + # Use call_rope_with_stability for CPU computation + # This ensures cache is cleared and autocast disabled + original_fn = getattr(self, '_original_get_axial_freqs', current_fn) + result = call_rope_with_stability(original_fn, *args, **kwargs) + + # Move module back to original device + self.to(current_device) + + # Move result to appropriate device if it's a tensor + if hasattr(result, 'to'): + target_device = args[0].device if len(args) > 0 and hasattr(args[0], 'device') else current_device + return result.to(target_device) + return result + + except Exception as cpu_error: + # Always restore device even on error + self.to(current_device) + raise cpu_error + else: + # Not a device error, let it bubble up + raise + + return device_aware_rope_wrapper + + # Apply wrapper + module.get_axial_freqs = types.MethodType( + make_device_aware_wrapper(name, current_method), + module + ) + module._blockswap_wrapped = True + + # Store for cleanup (use original or previously stored) + original_method = getattr(module, '_original_get_axial_freqs', current_method) + rope_patches.append((module, original_method)) + + if rope_patches: + model._rope_patches = rope_patches + debug.log(f"Patched {len(rope_patches)} RoPE modules with device handling", category="success") + + +def _protect_model_from_move( + model: torch.nn.Module, + runner: 'VideoDiffusionInfer', + debug: 'Debug' +) -> None: + """ + Protect model from unintended full device movement during BlockSwap. + + 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 model flag. + + Args: + model: DiT model to protect + runner: VideoDiffusionInfer instance (for active status check) + debug: Debug instance for logging (required) + """ + if not hasattr(model, '_original_to'): + # Store runner reference as weak reference to avoid circular refs + model._blockswap_runner_ref = weakref.ref(runner) + model._original_to = model.to + + # 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 hasattr(runner_obj, "_blockswap_active"): + blockswap_is_active = runner_obj._blockswap_active + + # Block attempts to move model away from configured offload device when active + if blockswap_is_active and str(device) != str(blockswap_offload_device): + # Get debug instance from runner if available + debug_instance = None + if runner_ref: + runner_obj = runner_ref() + if runner_obj and hasattr(runner_obj, 'debug'): + debug_instance = runner_obj.debug + + if debug_instance: + debug_instance.log( + f"Blocked attempt to move BlockSwap model from {blockswap_offload_device} to {device}", + level="WARNING", category="blockswap", force=True + ) + return self + + # Allow movement (either bypass is enabled or target is offload device) + if hasattr(self, '_original_to'): + return self._original_to(device, *args, **kwargs) + else: + # Fallback - shouldn't happen + return super(type(self), self).to(device, *args, **kwargs) + + # Bind as a method to the model instance + model.to = types.MethodType(protected_model_to, model) + + +def set_blockswap_bypass(runner, bypass: bool, debug): + """ + Set or unset bypass flag for BlockSwap protection. + Used for offloading to temporarily allow model movement. + + Args: + runner: Runner instance with BlockSwap + bypass: True to bypass protection, False to enforce it + debug: Debug instance for logging + """ + if not hasattr(runner, "_blockswap_active") or not runner._blockswap_active: + return + + # 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") + else: + debug.log("BlockSwap protection renabled to avoid accidentally offloading the entire DiT model", category="success") + + +def cleanup_blockswap(runner, keep_state_for_cache=False): + """ + Clean up BlockSwap configuration based on caching mode. + + When caching (keep_state_for_cache=True): + - Keep all BlockSwap configuration intact + - Only mark as inactive for safety during non-inference operations + + When not caching (keep_state_for_cache=False): + - Full cleanup of all BlockSwap state + + Args: + runner: VideoDiffusionInfer instance to clean up + keep_state_for_cache: If True, preserve BlockSwap state for reuse + """ + # Get debug instance from runner + if not hasattr(runner, 'debug') or runner.debug is None: + raise ValueError("Debug instance must be available on runner for cleanup_blockswap") + + debug = runner.debug + + # 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(model, "_block_swap_config") or + hasattr(model, "_blockswap_bypass_protection") + ) + + if not has_blockswap_state: + return + + debug.log("Starting BlockSwap cleanup", category="cleanup") + + if keep_state_for_cache: + # 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(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 CompatibleDiT wrapper) + model = runner.dit + if hasattr(model, "dit_model"): + model = model.dit_model + + # 1. Restore block forward methods + if hasattr(model, 'blocks'): + restored_count = 0 + for block in model.blocks: + if hasattr(block, '_original_forward'): + block.forward = block._original_forward + delattr(block, '_original_forward') + restored_count += 1 + + # Clean up wrapper attributes + for attr in ['_block_idx', '_model_ref', '_debug_ref', '_blockswap_wrapped']: + if hasattr(block, attr): + delattr(block, attr) + + if restored_count > 0: + debug.log(f"Restored {restored_count} block forward methods", category="success") + + # 2. Restore RoPE patches + if hasattr(model, '_rope_patches'): + for module, original_method in model._rope_patches: + module.get_axial_freqs = original_method + # Clean up wrapper attributes + for attr in ['_rope_wrapped', '_original_get_axial_freqs']: + if hasattr(module, attr): + delattr(module, attr) + debug.log(f"Restored {len(model._rope_patches)} RoPE methods", category="success") + delattr(model, '_rope_patches') + + # 3. Restore I/O component forward methods and move to offload device + if hasattr(model, '_io_swappers'): + for module, module_name in model._io_swappers: + if hasattr(module, '_original_forward'): + module.forward = module._original_forward + # Clean up wrapper attributes + for attr in ['_original_forward', '_model_ref', '_debug_ref', + '_module_name', '_is_io_wrapped']: + if hasattr(module, attr): + delattr(module, attr) + debug.log(f"Restored {len(model._io_swappers)} I/O components", category="success") + delattr(model, '_io_swappers') + + # Move all IO components to offload device during full cleanup + if hasattr(model, 'offload_device'): + offload_device = model.offload_device + moved_count = 0 + for name, module in model.named_children(): + if name != "blocks": + module.to(offload_device) + moved_count += 1 + if moved_count > 0: + debug.log(f"Moved {moved_count} IO components to offload device", category="success") + + # 4. Restore original .to() method + if hasattr(model, '_original_to'): + model.to = model._original_to + delattr(model, '_original_to') + debug.log("Restored original .to() method", category="success") + + # 5. Clean up BlockSwap-specific attributes + for attr in ['_blockswap_runner_ref', 'blocks_to_swap', 'main_device', + 'offload_device']: + if hasattr(model, attr): + delattr(model, attr) + + # 6. Clean up runner attributes + runner._blockswap_active = False + + # Remove all config attributes + for attr in ['_cached_blockswap_config', '_block_swap_config', '_blockswap_debug']: + if hasattr(runner, attr): + delattr(runner, attr) + + debug.log("BlockSwap cleanup complete", category="success") diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/compatibility.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/compatibility.py new file mode 100644 index 0000000000000000000000000000000000000000..c462022bed8ce602116fd1d6d8d88c4ff49adccf --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/compatibility.py @@ -0,0 +1,955 @@ +""" +Compatibility module for SeedVR2 +Contains FP8/FP16 compatibility layers and wrappers for different model architectures + +Extracted from: seedvr2.py (lines 1045-1630) +""" + +# 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/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 attention mode availability with automatic fallback. + + Args: + requested_mode: 'sdpa', 'flash_attn_2', 'flash_attn_3', 'sageattn_2', or 'sageattn_3' + debug: Optional debug instance for logging + + Returns: + Validated mode that is 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 = ( + "Cannot use 'flash_attn_3' attention mode: Flash Attention is not installed.\n" + "\n" + "Flash Attention 3 provides maximum speedup on Hopper+ 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' + + # 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 + TRITON_AVAILABLE = True +except ImportError: + TRITON_AVAILABLE = False + + +# 3. GGUF - Required for quantized model loading +try: + import gguf + from gguf import GGMLQuantizationType + GGUF_AVAILABLE = True +except ImportError: + GGUF_AVAILABLE = False + gguf = None + GGMLQuantizationType = None + + +def validate_gguf_availability(operation: str = "load GGUF model", debug=None) -> None: + """ + Validate GGUF availability and raise error if not installed. + + Args: + operation: Description of the operation requiring GGUF + debug: Optional debug instance for logging + + Raises: + RuntimeError: If GGUF is not available + """ + if not GGUF_AVAILABLE: + error_msg = ( + f"Cannot {operation}: GGUF library is not installed.\n" + f"\n" + f"GGUF provides quantized model support for memory-efficient loading.\n" + f"\n" + f"To fix this issue:\n" + f" 1. Install GGUF: pip install gguf\n" + f" 2. OR use a non-quantized model format (.safetensors)\n" + f"\n" + f"For more info: https://github.com/ggerganov/ggml" + ) + if debug: + debug.log(error_msg, level="ERROR", category="setup", force=True) + raise RuntimeError(f"GGUF library required to {operation}") + + +# 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+ with cuDNN >= 91002 uses 3x memory for Conv3d + with fp16/bfloat16 due to buggy dispatch layer. + """ + try: + # 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 + version_str = torch.__version__.split('+')[0] + parts = version_str.split('.') + torch_version = tuple(int(p) for p in parts[:2]) + + # Bug affects PyTorch 2.9 and later versions + if torch_version < (2, 9): + return False + + if not hasattr(torch.backends.cudnn, 'version'): + return False + + cudnn_version = torch.backends.cudnn.version() + if cudnn_version is None or cudnn_version < 91002: + return False + + return True + except: + return False + +NVIDIA_CONV3D_MEMORY_BUG_WORKAROUND = _check_conv3d_memory_bug() + + +# Log all optimization status once globally (cross-process) using environment variable +if not os.environ.get("SEEDVR2_OPTIMIZATIONS_LOGGED"): + os.environ["SEEDVR2_OPTIMIZATIONS_LOGGED"] = "1" + + # Build status strings + sage_status = "โœ…" if SAGE_ATTN_AVAILABLE else "โŒ" + flash_status = "โœ…" if FLASH_ATTN_AVAILABLE else "โŒ" + triton_status = "โœ…" if TRITON_AVAILABLE else "โŒ" + + # Count available optimizations + available = [SAGE_ATTN_AVAILABLE, FLASH_ATTN_AVAILABLE, TRITON_AVAILABLE] + num_available = sum(available) + + if num_available == 3: + print(f"โšก SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}") + elif num_available == 0: + print(f"โš ๏ธ SeedVR2 optimizations check: SageAttention {sage_status} | Flash Attention {flash_status} | Triton {triton_status}") + print("๐Ÿ’ก For best performance: pip install sageattention flash-attn triton") + else: + 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: + torch_ver = torch.__version__.split('+')[0] + cudnn_ver = torch.backends.cudnn.version() + 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 (CUDA only) + This prevents artifacts in FP8/mixed precision models. + """ + if hasattr(method, 'cache_clear'): + method.cache_clear() + + # 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 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/BFloat16/Float32: Uses native precision throughout + - GGUF: On-the-fly dequantization to compute_dtype + - MPS: Forces all parameters to compute_dtype (unified memory requires dtype consistency) + - RoPE: Converted from FP8 to compute_dtype for numerical consistency + + Optimizations: + - RoPE Stabilization: Error handling for numerical stability in mixed precision + - MPS Compatibility: Unified dtype conversion for Apple Silicon backends + """ + + def __init__(self, dit_model, debug: 'Debug', compute_dtype: torch.dtype = torch.bfloat16, skip_conversion: bool = False): + super().__init__() + self.dit_model = dit_model + self.debug = debug + self.compute_dtype = compute_dtype + self.model_dtype = self._detect_model_dtype() + self.is_fp8_model = self.model_dtype in (torch.float8_e4m3fn, torch.float8_e5m2) + self.is_fp16_model = self.model_dtype == torch.float16 + + # Only convert if not already done (e.g., when reusing cached weights) + if not skip_conversion and self.is_fp8_model: + # FP8 models need RoPE frequency conversion to compute dtype + model_variant = self._get_model_variant() + self.debug.log(f"Detected NaDiT {model_variant} FP8 - Converting RoPE freqs for FP8 compatibility", + category="precision") + 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") + + # MPS requires unified dtype for all parameters/buffers (no autocast fallback) + # 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") + + # Apply RoPE stabilization for numerical stability + self.debug.log(f"Stabilizing RoPE computations for numerical stability", category="setup") + self.debug.start_timer("_stabilize_rope_computations") + self._stabilize_rope_computations() + self.debug.end_timer("_stabilize_rope_computations", "RoPE stabilization") + + def _detect_model_dtype(self) -> torch.dtype: + """Detect main model dtype""" + try: + return next(self.dit_model.parameters()).dtype + except: + return torch.bfloat16 + + def _get_model_variant(self) -> str: + """Detect model variant from module path""" + model_module = str(self.dit_model.__class__.__module__).lower() + if 'dit_7b' in model_module: + return "7B" + elif 'dit_3b' in model_module: + return "3B" + else: + return "Unknown" + + def _convert_rope_freqs(self, target_dtype: torch.dtype = torch.bfloat16) -> None: + """ + Convert RoPE frequency buffers from FP8 to target dtype for compatibility. + + Args: + target_dtype: Target dtype for RoPE freqs (default: bfloat16 for stability) + """ + converted = 0 + for module in self.dit_model.modules(): + if 'RotaryEmbedding' in type(module).__name__: + if hasattr(module, 'rope') and hasattr(module.rope, 'freqs'): + if module.rope.freqs.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + if module.rope.freqs.device.type == "mps": + module.rope.freqs.data = module.rope.freqs.to("cpu").to(target_dtype).to("mps") + else: + module.rope.freqs.data = module.rope.freqs.to(target_dtype) + converted += 1 + self.debug.log(f"Converted {converted} RoPE frequency buffers from FP8 to {target_dtype} for compatibility", category="success") + + def _force_nadit_precision(self, target_dtype: torch.dtype = torch.bfloat16) -> None: + """ + Force ALL NaDiT parameters to target dtype to avoid promotion errors (MPS requirement). + + Args: + target_dtype: Target dtype for all parameters (default: bfloat16 for MPS compatibility) + """ + converted_count = 0 + original_dtype = None + + # Convert ALL parameters to target dtype + for name, param in self.dit_model.named_parameters(): + if original_dtype is None: + original_dtype = param.dtype + if param.dtype != target_dtype: + if param.device.type == "mps": + temp_cpu = param.data.to("cpu") + temp_converted = temp_cpu.to(target_dtype) + param.data = temp_converted.to("mps") + del temp_cpu, temp_converted + else: + param.data = param.data.to(target_dtype) + converted_count += 1 + + # 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") + temp_converted = temp_cpu.to(target_dtype) + buffer.data = temp_converted.to("mps") + del temp_cpu, temp_converted + else: + buffer.data = buffer.data.to(target_dtype) + converted_count += 1 + + self.debug.log(f"Converted {converted_count} NaDiT parameters/buffers to {target_dtype} for MPS", category="success") + + # Update detected dtype + self.model_dtype = target_dtype + self.is_fp8_model = (target_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)) + + def _stabilize_rope_computations(self): + """ + Add error handling to RoPE computations to prevent artifacts. + + Wraps the get_axial_freqs method of RoPE modules with a try-except handler. + During normal operation, uses the original cached method for performance. + Only on exceptions (e.g., numerical instability, NaN propagation) does it + intervene by clearing the cache and retrying the computation through + call_rope_with_stability. + + This prevents artifacts in FP8, mixed precision, and edge cases while + maintaining optimal performance for normal operations. + """ + if not hasattr(self.dit_model, 'blocks'): + return + + rope_count = 0 + + # Wrap RoPE modules to handle numerical instability + for name, module in self.dit_model.named_modules(): + if "rope" in name.lower() and hasattr(module, "get_axial_freqs"): + # Check if already wrapped + if hasattr(module, '_rope_wrapped'): + continue + + original_method = module.get_axial_freqs + + # Mark as wrapped and store original + module._rope_wrapped = 'stability' + module._original_get_axial_freqs = original_method + + # Error handler that prevents NaN propagation + def stable_rope_computation(self, *args, **kwargs): + try: + return original_method(*args, **kwargs) + except Exception: + return call_rope_with_stability(original_method, *args, **kwargs) + + module.get_axial_freqs = types.MethodType(stable_rope_computation, module) + rope_count += 1 + + if rope_count > 0: + self.debug.log(f"Stabilized {rope_count} RoPE modules", category="success") + + def forward(self, *args, **kwargs): + """ + Forward pass with minimal dtype conversion overhead + + Conversion strategy: + - FP16/BFloat16/Float32 models: Use native precision (no conversion needed) + - FP8 models: Convert FP8 tensors to compute_dtype for arithmetic operations + (FP8 parameters stay in FP8 for memory efficiency, only converted for computation) + """ + + # Only convert if we have an FP8 model for arithmetic operations + if self.is_fp8_model: + fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2) + target_dtype = self.compute_dtype + + # Convert args + converted_args = [] + for arg in args: + if isinstance(arg, torch.Tensor) and arg.dtype in fp8_dtypes: + converted_args.append(arg.to(target_dtype)) + else: + converted_args.append(arg) + + # Convert kwargs + converted_kwargs = {} + for key, value in kwargs.items(): + if isinstance(value, torch.Tensor) and value.dtype in fp8_dtypes: + converted_kwargs[key] = value.to(target_dtype) + else: + converted_kwargs[key] = value + + args = tuple(converted_args) + kwargs = converted_kwargs + + # Execute forward pass + try: + return self.dit_model(*args, **kwargs) + except Exception as e: + self.debug.log(f"Forward pass error: {e}", level="ERROR", category="generation", force=True) + if self.is_fp8_model: + self.debug.log(f"FP8 model - converted FP8 tensors to {self.compute_dtype}", category="info", force=True) + else: + self.debug.log(f"{self.model_dtype} model - no conversion applied", category="info", force=True) + raise + + def __getattr__(self, name): + """Redirect all other attributes to original model""" + if name in ['dit_model', 'model_dtype', 'is_fp8_model', 'is_fp16_model']: + return super().__getattr__(name) + return getattr(self.dit_model, name) + + def __setattr__(self, name, value): + """Redirect assignments to original model except for our attributes""" + if name in ['dit_model', 'model_dtype', 'is_fp8_model', 'is_fp16_model']: + super().__setattr__(name, value) + else: + if hasattr(self, 'dit_model'): + setattr(self.dit_model, name, value) + else: + super().__setattr__(name, value) + \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/gguf_dequant.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/gguf_dequant.py new file mode 100644 index 0000000000000000000000000000000000000000..e98a9a6fb1c78ca0f1b3e857a226930bda1c5f37 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/gguf_dequant.py @@ -0,0 +1,344 @@ +""" +Fast GGUF dequantization functions +Adapted from ComfyUI-GGUF +Optimized for SeedVR2 with proper debug logging and error handling +""" +import torch +import traceback +from typing import Optional, Tuple, List +from ..utils.constants import QK_K, K_SCALE_SIZE, suppress_tensor_warnings +from ..optimization.compatibility import GGUF_AVAILABLE, validate_gguf_availability + +# Import GGUF library +if GGUF_AVAILABLE: + import gguf + TORCH_COMPATIBLE_QTYPES = (None, gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16) +else: + gguf = None + TORCH_COMPATIBLE_QTYPES = (None,) + + +def is_torch_compatible(tensor: torch.Tensor) -> bool: + return tensor is None or getattr(tensor, "tensor_type", None) in TORCH_COMPATIBLE_QTYPES + + +def is_quantized(tensor: torch.Tensor) -> bool: + return not is_torch_compatible(tensor) + + +@torch._dynamo.disable +def dequantize_tensor(tensor: torch.Tensor, dtype: Optional[torch.dtype] = None, + dequant_dtype: Optional[torch.dtype] = None, + debug: Optional['Debug'] = None) -> torch.Tensor: + """ + Fast dequantization using optimized PyTorch operations + Returns regular PyTorch tensors to avoid infinite loops + + Args: + tensor: GGUF tensor to dequantize + dtype: Target dtype for final result + dequant_dtype: Intermediate dtype for dequantization + debug: Optional Debug instance for logging + """ + qtype = getattr(tensor, "tensor_type", None) + oshape = getattr(tensor, "tensor_shape", tensor.shape) + + # Suppress tensor copy warning - we intentionally convert GGUFTensor to regular tensor + suppress_tensor_warnings() + + if qtype in TORCH_COMPATIBLE_QTYPES: + result = tensor.to(dtype) + # Ensure we return a regular tensor, not GGUFTensor + if hasattr(result, 'tensor_type'): + result = torch.tensor(result.data, dtype=dtype, device=result.device, requires_grad=False) + return result + elif qtype in dequantize_functions: + dequant_dtype = dtype if dequant_dtype == "target" else dequant_dtype + # Use tensor.data like ComfyUI-GGUF does + result = dequantize(tensor.data, qtype, oshape, dtype=dequant_dtype, debug=debug) + final_result = result.to(dtype) + + # Ensure we return a regular tensor, not GGUFTensor + if hasattr(final_result, 'tensor_type'): + final_result = torch.tensor(final_result.data, dtype=dtype, device=final_result.device, requires_grad=False) + + return final_result + else: + raise NotImplementedError(f"No dequantization for {qtype}") + + +@torch._dynamo.disable +def dequantize(data: torch.Tensor, qtype: 'gguf.GGMLQuantizationType', + oshape: Tuple[int, ...], dtype: Optional[torch.dtype] = None, + debug: Optional['Debug'] = None) -> torch.Tensor: + """ + Dequantize tensor back to usable shape/dtype using fast operations + + Args: + data: Quantized data to dequantize + qtype: GGUF quantization type + oshape: Original shape to restore + dtype: Target dtype (default: torch.float16) + debug: Optional Debug instance for logging + + Returns: + Dequantized tensor in original shape + + Raises: + ValueError: If quantization type is not supported + RuntimeError: If dequantization fails + """ + if not GGUF_AVAILABLE: + validate_gguf_availability("dequantize GGUF tensor", debug) + + if dtype is None: + dtype = torch.float16 + + if qtype not in dequantize_functions: + raise ValueError(f"Unsupported quantization type: {qtype}") + try: + if debug: + debug.start_timer(f"dequant_{qtype.name if hasattr(qtype, 'name') else qtype}") + + block_size, type_size = gguf.GGML_QUANT_SIZES[qtype] + dequantize_blocks = dequantize_functions[qtype] + + # Ensure data is contiguous and properly formatted + data = data.contiguous() + + # More robust reshaping to handle edge cases + rows = data.reshape((-1, data.shape[-1])).view(torch.uint8) + + # Calculate number of blocks + n_blocks = rows.numel() // type_size + if rows.numel() % type_size != 0: + error_msg = f"Data size {rows.numel()} not divisible by type_size {type_size}. This usually indicates corrupted GGUF data." + if debug: + debug.log(error_msg, level="ERROR", category="precision", force=True) + raise ValueError(error_msg) + blocks = rows.reshape((n_blocks, type_size)) + + # Call the dequantization function + blocks = dequantize_blocks(blocks, block_size, type_size, dtype) + + result = blocks.reshape(oshape) + + # Ensure the result tensor is properly formatted + result = result.contiguous() + result.requires_grad_(False) + + if debug: + debug.end_timer(f"dequant_{qtype.name if hasattr(qtype, 'name') else qtype}", + f"Dequantized {qtype} tensor") + + return result + except Exception as e: + if debug: + debug.log(f"Error in dequantize: {e}", level="ERROR", category="precision", force=True) + debug.log(f"Data shape: {data.shape if 'data' in locals() else 'unknown'}", level="ERROR", category="precision", force=True) + debug.log(f"qtype: {qtype}, oshape: {oshape}", level="ERROR", category="precision", force=True) + debug.log(f"Traceback: {traceback.format_exc()}", level="ERROR", category="precision", force=True) + else: + traceback.print_exc() + raise + + +def to_uint32(x: torch.Tensor) -> torch.Tensor: + # no uint32 :( + x = x.view(torch.uint8).to(torch.int32) + return (x[:, 0] | x[:, 1] << 8 | x[:, 2] << 16 | x[:, 3] << 24).unsqueeze(1) + + +def split_block_dims(blocks: torch.Tensor, *args: int) -> List[torch.Tensor]: + n_max = blocks.shape[1] + dims = list(args) + [n_max - sum(args)] + return torch.split(blocks, dims, dim=1) + + +def get_scale_min(scales: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + n_blocks = scales.shape[0] + scales = scales.view(torch.uint8) + scales = scales.reshape((n_blocks, 3, 4)) + + d, m, m_d = torch.split(scales, scales.shape[-2] // 3, dim=-2) + + sc = torch.cat([d & 0x3F, (m_d & 0x0F) | ((d >> 2) & 0x30)], dim=-1) + min = torch.cat([m & 0x3F, (m_d >> 4) | ((m >> 2) & 0x30)], dim=-1) + + return (sc.reshape((n_blocks, 8)), min.reshape((n_blocks, 8))) + + +def dequantize_blocks_Q4_K(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q4_K dequantization""" + n_blocks = blocks.shape[0] + d, dmin, scales, qs = split_block_dims(blocks, 2, 2, K_SCALE_SIZE) + d = d.view(torch.float16).to(dtype) + dmin = dmin.view(torch.float16).to(dtype) + sc, m = get_scale_min(scales) + d = (d * sc).reshape((n_blocks, -1, 1)) + dm = (dmin * m).reshape((n_blocks, -1, 1)) + qs = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) + qs = (qs & 0x0F).reshape((n_blocks, -1, 32)) + return (d * qs - dm).reshape((n_blocks, QK_K)) + + +def dequantize_blocks_Q8_0(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Fast Q8_0 dequantization""" + d, x = split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + x = x.view(torch.int8) + return (d * x) + + +def dequantize_blocks_BF16(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """BF16 dequantization""" + return (blocks.view(torch.int16).to(torch.int32) << 16).view(torch.float32) + + +def dequantize_blocks_Q5_1(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q5_1 dequantization""" + n_blocks = blocks.shape[0] + d, m, qh, qs = split_block_dims(blocks, 2, 2, 4) + d = d.view(torch.float16).to(dtype) + m = m.view(torch.float16).to(dtype) + qh = to_uint32(qh) + qh = qh.reshape((n_blocks, 1)) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) + ql = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1) + qh = (qh & 1).to(torch.uint8) + ql = (ql & 0x0F).reshape((n_blocks, -1)) + qs = (ql | (qh << 4)) + return (d * qs) + m + + +def dequantize_blocks_Q5_0(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q5_0 dequantization""" + n_blocks = blocks.shape[0] + d, qh, qs = split_block_dims(blocks, 2, 4) + d = d.view(torch.float16).to(dtype) + qh = to_uint32(qh) + qh = qh.reshape(n_blocks, 1) >> torch.arange(32, device=d.device, dtype=torch.int32).reshape(1, 32) + ql = qs.reshape(n_blocks, -1, 1, block_size // 2) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1) + qh = (qh & 1).to(torch.uint8) + ql = (ql & 0x0F).reshape(n_blocks, -1) + qs = (ql | (qh << 4)).to(torch.int8) - 16 + return (d * qs) + + +def dequantize_blocks_Q4_1(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q4_1 dequantization""" + n_blocks = blocks.shape[0] + d, m, qs = split_block_dims(blocks, 2, 2) + d = d.view(torch.float16).to(dtype) + m = m.view(torch.float16).to(dtype) + qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape(1, 1, 2, 1) + qs = (qs & 0x0F).reshape(n_blocks, -1) + return (d * qs) + m + + +def dequantize_blocks_Q4_0(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q4_0 dequantization""" + n_blocks = blocks.shape[0] + d, qs = split_block_dims(blocks, 2) + d = d.view(torch.float16).to(dtype) + qs = qs.reshape((n_blocks, -1, 1, block_size // 2)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) + qs = (qs & 0x0F).reshape((n_blocks, -1)).to(torch.int8) - 8 + return (d * qs) + + +def dequantize_blocks_Q6_K(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q6_K dequantization""" + n_blocks = blocks.shape[0] + ql, qh, scales, d, = split_block_dims(blocks, QK_K // 2, QK_K // 4, QK_K // 16) + scales = scales.view(torch.int8).to(dtype) + d = d.view(torch.float16).to(dtype) + d = (d * scales).reshape((n_blocks, QK_K // 16, 1)) + ql = ql.reshape((n_blocks, -1, 1, 64)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) + ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) + qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) + qh = (qh & 0x03).reshape((n_blocks, -1, 32)) + q = (ql | (qh << 4)).to(torch.int8) - 32 + q = q.reshape((n_blocks, QK_K // 16, -1)) + return (d * q).reshape((n_blocks, QK_K)) + + +def dequantize_blocks_Q5_K(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q5_K dequantization""" + n_blocks = blocks.shape[0] + d, dmin, scales, qh, qs = split_block_dims(blocks, 2, 2, K_SCALE_SIZE, QK_K // 8) + d = d.view(torch.float16).to(dtype) + dmin = dmin.view(torch.float16).to(dtype) + sc, m = get_scale_min(scales) + d = (d * sc).reshape((n_blocks, -1, 1)) + dm = (dmin * m).reshape((n_blocks, -1, 1)) + ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 1, 2, 1)) + qh = qh.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([i for i in range(8)], device=d.device, dtype=torch.uint8).reshape((1, 1, 8, 1)) + ql = (ql & 0x0F).reshape((n_blocks, -1, 32)) + qh = (qh & 0x01).reshape((n_blocks, -1, 32)) + q = (ql | (qh << 4)) + return (d * q - dm).reshape((n_blocks, QK_K)) + + +def dequantize_blocks_Q3_K(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q3_K dequantization""" + n_blocks = blocks.shape[0] + hmask, qs, scales, d = split_block_dims(blocks, QK_K // 8, QK_K // 4, 12) + d = d.view(torch.float16).to(dtype) + lscales, hscales = scales[:, :8], scales[:, 8:] + lscales = lscales.reshape((n_blocks, 1, 8)) >> torch.tensor([0, 4], device=d.device, dtype=torch.uint8).reshape((1, 2, 1)) + lscales = lscales.reshape((n_blocks, 16)) + hscales = hscales.reshape((n_blocks, 1, 4)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 4, 1)) + hscales = hscales.reshape((n_blocks, 16)) + scales = (lscales & 0x0F) | ((hscales & 0x03) << 4) + scales = (scales.to(torch.int8) - 32) + dl = (d * scales).reshape((n_blocks, 16, 1)) + ql = qs.reshape((n_blocks, -1, 1, 32)) >> torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) + qh = hmask.reshape(n_blocks, -1, 1, 32) >> torch.tensor([i for i in range(8)], device=d.device, dtype=torch.uint8).reshape((1, 1, 8, 1)) + ql = ql.reshape((n_blocks, 16, QK_K // 16)) & 3 + qh = (qh.reshape((n_blocks, 16, QK_K // 16)) & 1) ^ 1 + q = (ql.to(torch.int8) - (qh << 2).to(torch.int8)) + return (dl * q).reshape((n_blocks, QK_K)) + + +def dequantize_blocks_Q2_K(blocks: torch.Tensor, block_size: int, type_size: int, + dtype: Optional[torch.dtype] = None, debug: Optional['Debug'] = None) -> torch.Tensor: + """Q2_K dequantization""" + n_blocks = blocks.shape[0] + scales, qs, d, dmin = split_block_dims(blocks, QK_K // 16, QK_K // 4, 2) + d = d.view(torch.float16).to(dtype) + dmin = dmin.view(torch.float16).to(dtype) + dl = (d * (scales & 0xF)).reshape((n_blocks, QK_K // 16, 1)) + ml = (dmin * (scales >> 4)).reshape((n_blocks, QK_K // 16, 1)) + shift = torch.tensor([0, 2, 4, 6], device=d.device, dtype=torch.uint8).reshape((1, 1, 4, 1)) + qs = (qs.reshape((n_blocks, -1, 1, 32)) >> shift) & 3 + qs = qs.reshape((n_blocks, QK_K // 16, 16)) + qs = dl * qs - ml + return qs.reshape((n_blocks, -1)) + + +# Main dequantization function lookup +if GGUF_AVAILABLE: + dequantize_functions = { + gguf.GGMLQuantizationType.BF16: dequantize_blocks_BF16, + gguf.GGMLQuantizationType.Q8_0: dequantize_blocks_Q8_0, + gguf.GGMLQuantizationType.Q5_1: dequantize_blocks_Q5_1, + gguf.GGMLQuantizationType.Q5_0: dequantize_blocks_Q5_0, + gguf.GGMLQuantizationType.Q4_1: dequantize_blocks_Q4_1, + gguf.GGMLQuantizationType.Q4_0: dequantize_blocks_Q4_0, + gguf.GGMLQuantizationType.Q6_K: dequantize_blocks_Q6_K, + gguf.GGMLQuantizationType.Q5_K: dequantize_blocks_Q5_K, + gguf.GGMLQuantizationType.Q4_K: dequantize_blocks_Q4_K, + gguf.GGMLQuantizationType.Q3_K: dequantize_blocks_Q3_K, + gguf.GGMLQuantizationType.Q2_K: dequantize_blocks_Q2_K, + } +else: + dequantize_functions = {} \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/gguf_ops.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/gguf_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..50dc1e927eb68ff495fac779be0d8511d99659c6 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/gguf_ops.py @@ -0,0 +1,284 @@ +""" +GGUF Quantization Operations Module for SeedVR2 + +Provides runtime quantization support with proper debug logging and memory management +""" + +import torch +import torch.nn as nn +from typing import Optional, Dict, Tuple +from ..optimization.compatibility import GGUF_AVAILABLE, validate_gguf_availability + +# Import GGUF-specific modules +if GGUF_AVAILABLE: + import gguf + from .gguf_dequant import dequantize_tensor, is_torch_compatible, is_quantized as is_gguf_quantized +else: + gguf = None + + # Fallback functions when GGUF not available + def is_torch_compatible(tensor: torch.Tensor) -> bool: + return True + + def is_gguf_quantized(tensor: torch.Tensor) -> bool: + return False + + def dequantize_tensor(tensor: torch.Tensor, dtype: Optional[torch.dtype] = None, + dequant_dtype: Optional[torch.dtype] = None, + debug: Optional['Debug'] = None) -> torch.Tensor: + if dtype is not None: + return tensor.to(dtype) + return tensor + + +class _GGUFQuantizedBase(nn.Module): + """Base class for GGUF quantized layers with shared dequantization logic""" + + def __init__(self, debug: Optional['Debug'] = None): + super().__init__() + self.weight = None + self.bias = None + self.weight_qtype = None + self.weight_shape = None + self.debug = debug + + def load_quantized_weight(self, weight_tensor: torch.Tensor, bias_tensor: Optional[torch.Tensor] = None) -> None: + """Load quantized weight tensor with debug logging""" + if hasattr(weight_tensor, 'tensor_type') and hasattr(weight_tensor, 'tensor_shape'): + # This is a quantized tensor - register as non-persistent buffer + # Non-persistent buffers are moved with model.to() but not saved in state_dict + # Check if buffer already exists (for model reloading scenarios) + if hasattr(self, 'quantized_weight'): + # Replace existing buffer + self._buffers['quantized_weight'] = weight_tensor + else: + # Register new buffer + self.register_buffer('quantized_weight', weight_tensor, persistent=False) + + self.weight_qtype = weight_tensor.tensor_type + self.weight_shape = weight_tensor.tensor_shape + else: + # Regular tensor + self.weight = nn.Parameter(weight_tensor, requires_grad=False) + + if bias_tensor is not None: + self.bias = nn.Parameter(bias_tensor, requires_grad=False) + + def _apply(self, fn): + """ + Override _apply to handle GGUF quantized weight device movement properly. + This ensures model.to(device) works correctly with quantized weights. + """ + # Handle quantized weight if it exists + if 'quantized_weight' in self._buffers: + # Apply the function (device/dtype conversion) to the quantized weight + self._buffers['quantized_weight'] = fn(self._buffers['quantized_weight']) + + # Let parent class handle parameters and other buffers + return super()._apply(fn) + + @torch._dynamo.disable + def dequantize_weight(self, device: Optional[torch.device] = None, dtype: torch.dtype = torch.float16) -> torch.Tensor: + """Dequantize weight tensor on-the-fly""" + if self.quantized_weight is None: + return self.weight + + if self.debug: + self.debug.start_timer("gguf_dequant") + + try: + # Check if already unquantized + if GGUF_AVAILABLE and self.weight_qtype in {gguf.GGMLQuantizationType.F32, gguf.GGMLQuantizationType.F16}: + return self.quantized_weight.to(device, dtype) + + # Use our optimized dequantize_tensor function + result = dequantize_tensor(self.quantized_weight, dtype=dtype, + dequant_dtype=dtype, debug=self.debug) + if device: + result = result.to(device) + return result + + except Exception as e: + if self.debug: + self.debug.log(f"Dequantization failed: {e}", + level="ERROR", category="precision", force=True) + raise RuntimeError(f"Failed to dequantize weight: {e}") + finally: + if self.debug: + self.debug.end_timer("gguf_dequant", "Weight dequantization completed") + + @torch._dynamo.disable + def get_dequantized_weight_for_compute(self, input: torch.Tensor) -> torch.Tensor: + """ + Get dequantized weight optimized for computation. + + For GGUF models: Dequantizes to FP16 first (original precision), + then converts to compute dtype to preserve maximum precision. + + Args: + input: Input tensor (used to determine device and compute dtype) + + Returns: + Dequantized weight tensor ready for computation + """ + device = input.device + compute_dtype = input.dtype + + # Dequantize GGUF to FP16 first (original precision), then convert to compute dtype + # This preserves maximum precision during dequantization + dequant_dtype = torch.float16 + weight = self.dequantize_weight(device, dequant_dtype) + + # Convert to compute dtype if different + if weight.dtype != compute_dtype: + weight = weight.to(compute_dtype) + + return weight + +class GGUFQuantizedLinear(_GGUFQuantizedBase): + """Quantized Linear layer with on-the-fly dequantization""" + + def __init__(self, in_features: int, out_features: int, bias: bool = True, + device=None, dtype=None, debug: Optional['Debug'] = None): + super().__init__(debug) + self.in_features = in_features + self.out_features = out_features + + def forward(self, input: torch.Tensor) -> torch.Tensor: + # Get precision-optimized weight + weight = self.get_dequantized_weight_for_compute(input) + + # Validate weight shape after dequantization + expected_shape = (self.out_features, self.in_features) + if weight.shape != expected_shape: + if self.debug: + self.debug.log(f"Linear weight shape mismatch: got {weight.shape}, expected {expected_shape}", + level="ERROR", category="precision", force=True) + raise RuntimeError(f"Dequantized linear weight has incorrect shape: {weight.shape} vs expected {expected_shape}") + + return torch.nn.functional.linear(input, weight, self.bias) + +class GGUFQuantizedConv2d(_GGUFQuantizedBase): + """Quantized Conv2d layer with on-the-fly dequantization""" + + def __init__(self, in_channels: int, out_channels: int, kernel_size, stride=1, + padding=0, dilation=1, groups=1, bias=True, device=None, dtype=None, + debug: Optional['Debug'] = None): + super().__init__(debug) + self.in_channels = in_channels + self.out_channels = out_channels + self.kernel_size = kernel_size + self.stride = stride + self.padding = padding + self.dilation = dilation + self.groups = groups + + def forward(self, input: torch.Tensor) -> torch.Tensor: + # Get precision-optimized weight + weight = self.get_dequantized_weight_for_compute(input) + + # Validate weight shape after dequantization + if isinstance(self.kernel_size, int): + kernel_h = kernel_w = self.kernel_size + else: + kernel_h, kernel_w = self.kernel_size + + expected_shape = (self.out_channels, self.in_channels // self.groups, kernel_h, kernel_w) + if weight.shape != expected_shape: + if self.debug: + self.debug.log(f"Conv2d weight shape mismatch: got {weight.shape}, expected {expected_shape}", + level="ERROR", category="precision", force=True) + raise RuntimeError(f"Dequantized conv weight has incorrect shape: {weight.shape} vs expected {expected_shape}") + + # Standard conv2d operation + return torch.nn.functional.conv2d(input, weight, self.bias, self.stride, + self.padding, self.dilation, self.groups) + + +def is_quantized_tensor(tensor: torch.Tensor) -> bool: + """Check if a tensor is GGUF quantized""" + if not GGUF_AVAILABLE: + return False + + # Check multiple indicators + if hasattr(tensor, 'tensor_type') and hasattr(tensor, 'tensor_shape'): + return True + if hasattr(tensor, '_gguf_quantized'): + return True + if 'GGUFTensor' in str(type(tensor)): + return True + return is_gguf_quantized(tensor) if callable(is_gguf_quantized) else False + + +def replace_linear_with_quantized(module, debug: Optional['Debug'] = None, prefix="") -> Tuple[int, Dict[str, int]]: + """ + Replace Linear and Conv2d layers with quantized versions if they have GGUF quantized weights + + Args: + module: The module to process + debug: Optional Debug instance for logging + prefix: Prefix for recursive calls (used internally) + + Returns: + Tuple of (replacements_made, quant_types) where: + - replacements_made: Number of layers replaced + - quant_types: Dict mapping quantization type names to counts + """ + if not GGUF_AVAILABLE: + return 0, {} + + replacements_made = 0 + quant_types = {} # Track quantization types found + + for name, child in module.named_children(): + if isinstance(child, nn.Linear): + if is_quantized_tensor(child.weight): + # Track quantization type + if hasattr(child.weight, 'tensor_type'): + qtype_name = child.weight.tensor_type.name if hasattr(child.weight.tensor_type, 'name') else str(child.weight.tensor_type) + quant_types[qtype_name] = quant_types.get(qtype_name, 0) + 1 + + # Create quantized linear layer + quantized_linear = GGUFQuantizedLinear( + child.in_features, + child.out_features, + bias=child.bias is not None, + debug=debug + ) + quantized_linear.load_quantized_weight(child.weight, child.bias) + setattr(module, name, quantized_linear) + replacements_made += 1 + + elif isinstance(child, nn.Conv2d): + if is_quantized_tensor(child.weight): + # Track quantization type + if hasattr(child.weight, 'tensor_type'): + qtype_name = child.weight.tensor_type.name if hasattr(child.weight.tensor_type, 'name') else str(child.weight.tensor_type) + quant_types[qtype_name] = quant_types.get(qtype_name, 0) + 1 + + # Create quantized conv2d layer + quantized_conv = GGUFQuantizedConv2d( + child.in_channels, + child.out_channels, + child.kernel_size, + stride=child.stride, + padding=child.padding, + dilation=child.dilation, + groups=child.groups, + bias=child.bias is not None, + debug=debug + ) + quantized_conv.load_quantized_weight(child.weight, child.bias) + setattr(module, name, quantized_conv) + replacements_made += 1 + else: + # Recursively replace in child modules + full_name = f"{prefix}.{name}" if prefix else name + child_replacements, child_qtypes = replace_linear_with_quantized(child, debug, full_name) + replacements_made += child_replacements + # Merge quantization types from child + for qtype, count in child_qtypes.items(): + quant_types[qtype] = quant_types.get(qtype, 0) + count + + # Always return the tuple + return replacements_made, quant_types \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/memory_manager.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/memory_manager.py new file mode 100644 index 0000000000000000000000000000000000000000..780c6909d02a45fb901c0c5077bfef98708efe1e --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/memory_manager.py @@ -0,0 +1,1231 @@ +""" +Memory management module for SeedVR2 +Handles VRAM usage, cache management, and memory optimization + +Extracted from: seedvr2.py (lines 373-405, 607-626, 1016-1044) +""" + +import torch +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]: + """ + Get list of available compute devices for SeedVR2 + + Args: + include_none: If True, prepend "none" to the device list (for offload options) + include_cpu: If True, include "cpu" in the device list (for offload options only) + Note: On MPS-only systems, "cpu" is automatically excluded since + unified memory architecture makes CPU offloading meaningless + + Returns: + List of device strings (e.g., ["cuda:0", "cuda:1"] or ["none", "cpu", "cuda:0", "cuda:1"]) + """ + devs = [] + has_cuda = False + has_mps = False + + try: + if is_cuda_available(): + devs += [f"cuda:{i}" for i in range(torch.cuda.device_count())] + has_cuda = True + except Exception: + pass + + try: + if is_mps_available(): + devs.append("mps") # MPS doesn't use device indices + has_mps = True + except Exception: + pass + + # Build result list with optional prefixes + result = [] + if include_none: + result.append("none") + + # Only include "cpu" option if: + # 1. It was requested (include_cpu=True), AND + # 2. Either CUDA is available OR MPS is not the only option + # Rationale: On MPS-only systems with unified memory architecture, + # CPU offloading is semantically meaningless as CPU and GPU share the same memory pool + if include_cpu and (has_cuda or not has_mps): + result.append("cpu") + + result.extend(devs) + + return result if result else [] + + +def get_basic_vram_info(device: Optional[torch.device] = None) -> Dict[str, Any]: + """ + Get basic VRAM availability info (free and total memory). + Used for capacity planning and initial checks. + + Args: + device: Optional device to query. If None, uses cuda:0 + + Returns: + dict: {"free_gb": float, "total_gb": float} or {"error": str} + """ + try: + 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 is_mps_available(): + # MPS doesn't support per-device queries or mem_get_info + # Use system memory as proxy + mem = psutil.virtual_memory() + free_memory = mem.total - mem.used + total_memory = mem.total + else: + return {"error": "No GPU backend available (CUDA/MPS)"} + + return { + "free_gb": free_memory / (1024**3), + "total_gb": total_memory / (1024**3) + } + except Exception as e: + return {"error": f"Failed to get memory info: {str(e)}"} + + +# Initial VRAM check at module load +vram_info = get_basic_vram_info(device=None) +if "error" not in vram_info: + 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, float]: + """ + Get current VRAM usage metrics for monitoring. + Used for tracking memory consumption during processing. + + Args: + device: Optional device to query. If None, uses cuda:0 + debug: Optional debug instance for logging + + Returns: + tuple: (allocated_gb, reserved_gb, peak_allocated_gb, peak_reserved_gb) + Returns (0, 0, 0, 0) if no GPU available + """ + try: + 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) + peak_allocated = torch.cuda.max_memory_allocated(device) / (1024**3) + peak_reserved = torch.cuda.max_memory_reserved(device) / (1024**3) + 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) + # MPS doesn't track peak separately + 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, 0.0 + + +def get_ram_usage(debug: Optional['Debug'] = None) -> Tuple[float, float, float, float]: + """ + Get current RAM usage metrics for the current process. + Provides accurate tracking of process-specific memory consumption. + + Args: + debug: Optional debug instance for logging + + Returns: + tuple: (process_gb, available_gb, total_gb, used_by_others_gb) + Returns (0, 0, 0, 0) if psutil not available or on error + """ + try: + if not psutil: + return 0.0, 0.0, 0.0, 0.0 + + # Get current process memory + process = psutil.Process() + process_memory = process.memory_info() + process_gb = process_memory.rss / (1024**3) + + # Get system memory + sys_memory = psutil.virtual_memory() + total_gb = sys_memory.total / (1024**3) + available_gb = sys_memory.available / (1024**3) + + # Calculate memory used by other processes + # This is the CORRECT calculation: + total_used_gb = total_gb - available_gb # Total memory used by ALL processes + used_by_others_gb = max(0, total_used_gb - process_gb) # Subtract current process + + return process_gb, available_gb, total_gb, used_by_others_gb + + except Exception as e: + if debug: + debug.log(f"Failed to get RAM usage: {e}", level="WARNING", category="memory", force=True) + return 0.0, 0.0, 0.0, 0.0 + + +# Global cache for OS libraries (initialized once) +_os_memory_lib = None + + +def clear_memory(debug: Optional['Debug'] = None, deep: bool = False, force: bool = True, + timer_name: Optional[str] = None) -> None: + """ + Clear memory caches with two-tier approach for optimal performance. + + Args: + debug: Debug instance for logging (optional) + force: If True, always clear. If False, only clear when <5% free + deep: If True, perform deep cleanup including GC and OS operations. + If False (default), only perform minimal GPU cache clearing. + timer_name: Optional suffix for timer names to make them unique per invocation + + Two-tier approach: + - Minimal mode (deep=False): GPU cache operations (~1-5ms) + Used for frequent calls during batch processing + - Deep mode (deep=True): Complete cleanup with GC and OS operations (~10-50ms) + Used at key points like model switches or final cleanup + """ + global _os_memory_lib + + # Create unique timer names if suffix provided + if timer_name: + main_timer = f"memory_clear_{timer_name}" + gpu_timer = f"gpu_cache_clear_{timer_name}" + gc_timer = f"garbage_collection_{timer_name}" + os_timer = f"os_memory_release_{timer_name}" + completion_msg = f"clear_memory() completion ({timer_name})" + else: + main_timer = "memory_clear" + gpu_timer = "gpu_cache_clear" + gc_timer = "garbage_collection" + os_timer = "os_memory_release" + completion_msg = "clear_memory() completion" + + # Start timer for entire operation + if debug: + debug.start_timer(main_timer) + + # Check if we should clear based on memory pressure + if not force: + should_clear = False + + # Use existing function for memory info + mem_info = get_basic_vram_info(device=None) + + 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 = "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 is_mps_available(): + mem = psutil.virtual_memory() + if mem.available < mem.total * 0.05: + should_clear = True + if debug: + debug.log(f"RAM pressure: {mem.available/(1024**3):.2f}GB free of {mem.total/(1024**3):.2f}GB", category="memory") + + if not should_clear: + # End timer before early return to keep stack clean + if debug: + debug.end_timer(main_timer) + return + + # Determine cleanup level + cleanup_mode = "deep" if deep else "minimal" + if debug: + debug.log(f"Clearing memory caches ({cleanup_mode})...", category="cleanup") + + # ===== MINIMAL OPERATIONS (Always performed) ===== + # Step 1: Clear GPU caches - Fast operations (~1-5ms) + if debug: + debug.start_timer(gpu_timer) + + if is_cuda_available(): + torch.cuda.empty_cache() + torch.cuda.ipc_collect() + elif is_mps_available(): + torch.mps.empty_cache() + + if debug: + debug.end_timer(gpu_timer, "GPU cache clearing") + + # ===== DEEP OPERATIONS (Only when deep=True) ===== + if deep: + # Step 2: Deep garbage collection (expensive ~5-20ms) + if debug: + debug.start_timer(gc_timer) + + gc.collect(2) + + if debug: + debug.end_timer(gc_timer, "Garbage collection") + + # Step 3: Return memory to OS (platform-specific, ~5-30ms) + if debug: + debug.start_timer(os_timer) + + try: + if sys.platform == 'linux': + # Linux: malloc_trim + import ctypes # Import only when needed + if _os_memory_lib is None: + _os_memory_lib = ctypes.CDLL("libc.so.6") + _os_memory_lib.malloc_trim(0) + + elif sys.platform == 'win32': + # Windows: Trim working set + import ctypes # Import only when needed + if _os_memory_lib is None: + _os_memory_lib = ctypes.windll.kernel32 + handle = _os_memory_lib.GetCurrentProcess() + _os_memory_lib.SetProcessWorkingSetSize(handle, -1, -1) + + elif is_mps_available(): + # macOS with MPS + import ctypes # Import only when needed + import ctypes.util + if _os_memory_lib is None: + libc_path = ctypes.util.find_library('c') + if libc_path: + _os_memory_lib = ctypes.CDLL(libc_path) + + if _os_memory_lib: + _os_memory_lib.sync() + except Exception as e: + if debug: + debug.log(f"Failed to perform OS memory operations: {e}", level="WARNING", category="memory", force=True) + + if debug: + debug.end_timer(os_timer, "OS memory release") + + # End overall timer + if debug: + debug.end_timer(main_timer, completion_msg) + + +def retry_on_oom(func, *args, debug=None, operation_name="operation", **kwargs): + """ + Execute function with single OOM retry after memory cleanup. + + Args: + func: Callable to execute + *args: Positional arguments for func + debug: Debug instance for logging (optional) + operation_name: Name for logging + **kwargs: Keyword arguments for func + + Returns: + Result of func(*args, **kwargs) + """ + try: + return func(*args, **kwargs) + except (torch.cuda.OutOfMemoryError, RuntimeError) as e: + # Only handle OOM errors + if not any(x in str(e).lower() for x in ["out of memory", "allocation on device"]): + raise + + if debug: + debug.log(f"OOM during {operation_name}: {e}", level="WARNING", category="memory", force=True) + debug.log(f"Clearing memory and retrying", category="info", force=True) + + # Clear memory + clear_memory(debug=debug, deep=True, force=True, timer_name=operation_name) + # Let memory settle + time.sleep(0.5) + debug.log_memory_state("After memory clearing", show_tensors=False, detailed_tensors=False) + + # Single retry + try: + result = func(*args, **kwargs) + if debug: + debug.log(f"Retry successful for {operation_name}", category="success", force=True) + return result + except Exception as retry_e: + if debug: + debug.log(f"Retry failed for {operation_name}: {retry_e}", level="ERROR", category="memory", force=True) + raise + + +def reset_vram_peak(device: Optional[torch.device] = None, debug: Optional['Debug'] = None) -> None: + """ + Reset VRAM peak memory statistics for fresh tracking. + + Args: + device: Optional device to reset stats for. If None, uses cuda:0 + debug: Optional debug instance for logging + """ + if debug and debug.enabled: + debug.log("Resetting VRAM peak memory statistics", category="memory") + try: + if is_cuda_available(): + if device is None: + device = torch.device("cuda:0") + elif not isinstance(device, torch.device): + device = torch.device(device) + torch.cuda.reset_peak_memory_stats(device) + # Note: MPS doesn't support peak memory reset - no action needed + except Exception as e: + if debug and debug.enabled: + debug.log(f"Failed to reset peak memory stats: {e}", level="WARNING", category="memory", force=True) + + +def clear_rope_lru_caches(model: Optional[torch.nn.Module], debug: Optional['Debug'] = None) -> int: + """ + Clear ALL LRU caches from RoPE modules. + + Args: + model: PyTorch model to clear caches from + debug: Optional debug instance for logging + + Returns: + Number of caches cleared + """ + if model is None: + return 0 + + cleared_count = 0 + try: + for name, module in model.named_modules(): + if hasattr(module, 'get_axial_freqs') and hasattr(module.get_axial_freqs, 'cache_clear'): + try: + module.get_axial_freqs.cache_clear() + cleared_count += 1 + except Exception as e: + if debug: + debug.log(f"Failed to clear RoPE LRU cache for module {name}: {e}", level="WARNING", category="memory", force=True) + except (AttributeError, RuntimeError) as e: + if debug: + debug.log(f"Failed to iterate model modules for RoPE LRU cache clearing: {e}", level="WARNING", category="memory", force=True) + + return cleared_count + + +def release_tensor_memory(tensor: Optional[torch.Tensor]) -> None: + """Release tensor memory from any device (CPU/CUDA/MPS)""" + if tensor is not None and torch.is_tensor(tensor): + # Release storage for all devices (CPU, CUDA, MPS) + if tensor.numel() > 0: + tensor.data.set_() + tensor.grad = None + + +def release_tensor_collection(collection: Any, recursive: bool = True) -> None: + """ + Release GPU memory from tensors in any collection (list, tuple, dict, or single tensor). + + Args: + collection: Tensor, list, tuple, dict, or nested structure to release + recursive: If True, handle nested structures recursively + + Examples: + release_tensor_collection(tensor) # Single tensor + release_tensor_collection([tensor1, tensor2]) # List of tensors + release_tensor_collection([[t1, t2], [t3, t4]]) # Nested lists + release_tensor_collection({'a': tensor}) # Dict values + """ + if collection is None: + return + + if torch.is_tensor(collection): + release_tensor_memory(collection) + elif isinstance(collection, dict): + for value in collection.values(): + if recursive: + release_tensor_collection(value, recursive=True) + elif torch.is_tensor(value): + release_tensor_memory(value) + elif isinstance(collection, (list, tuple)): + for item in collection: + if recursive: + release_tensor_collection(item, recursive=True) + elif torch.is_tensor(item): + release_tensor_memory(item) + + +def release_text_embeddings(*embeddings: torch.Tensor, debug: Optional['Debug'] = None, names: Optional[List[str]] = None) -> None: + """ + Release memory for text embeddings + + Args: + *embeddings: Variable number of embedding tensors to release + debug: Optional debug instance for logging + names: Optional list of names for logging + """ + for i, embedding in enumerate(embeddings): + if embedding is not None: + release_tensor_memory(embedding) + if debug and names and i < len(names): + debug.log(f"Cleaned up {names[i]}", category="cleanup") + + +def cleanup_text_embeddings(ctx: Dict[str, Any], debug: Optional['Debug'] = None) -> None: + """ + Clean up text embeddings from a context dictionary. + Extracts embeddings, releases memory, and clears the context entry. + + Args: + ctx: Context dictionary potentially containing 'text_embeds' + debug: Optional debug instance for logging + """ + if not ctx or not ctx.get('text_embeds'): + return + + embeddings = [] + names = [] + for key, embeds_list in ctx['text_embeds'].items(): + if embeds_list: + embeddings.extend(embeds_list) + names.append(key) + + if embeddings: + release_text_embeddings(embeddings, names, debug) + + if debug: + debug.log(f"Cleaned up text embeddings: {', '.join(names)}", category="cleanup") + + ctx['text_embeds'] = None + + +def release_model_memory(model: Optional[torch.nn.Module], debug: Optional['Debug'] = None) -> None: + """ + Release all GPU/MPS memory from model in-place without CPU transfer. + + Args: + model: PyTorch model to release memory from + debug: Optional debug instance for logging + """ + if model is None: + return + + try: + # Clear gradients first + model.zero_grad(set_to_none=True) + + # Release GPU memory directly without CPU transfer + released_params = 0 + released_buffers = 0 + + for param in model.parameters(): + if param.is_cuda or param.is_mps: + if param.numel() > 0: + param.data.set_() + released_params += 1 + param.grad = None + + for buffer in model.buffers(): + if buffer.is_cuda or buffer.is_mps: + if buffer.numel() > 0: + buffer.data.set_() + released_buffers += 1 + + if debug and (released_params > 0 or released_buffers > 0): + debug.log(f"Released memory from {released_params} params and {released_buffers} buffers", category="success") + + except (AttributeError, RuntimeError) as e: + if debug: + debug.log(f"Failed to release model memory: {e}", level="WARNING", category="memory", force=True) + + +def manage_tensor( + tensor: torch.Tensor, + target_device: torch.device, + tensor_name: str = "tensor", + dtype: Optional[torch.dtype] = None, + non_blocking: bool = False, + debug: Optional['Debug'] = None, + reason: Optional[str] = None, + indent_level: int = 0 +) -> torch.Tensor: + """ + Unified tensor management for device movement and dtype conversion. + + Handles both device transfers (CPU โ†” GPU) and dtype conversions (e.g., float16 โ†’ bfloat16) + with intelligent early-exit optimization and comprehensive logging. + + Args: + tensor: Tensor to manage + target_device: Target device (torch.device object) + tensor_name: Descriptive name for logging (e.g., "latent", "sample", "alpha_channel") + dtype: Optional target dtype to cast to (if None, keeps original dtype) + non_blocking: Whether to use non-blocking transfer + debug: Debug instance for logging + reason: Optional reason for the operation (e.g., "inference", "offload", "dtype alignment") + indent_level: Indentation level for debug logging (0=no indent, 1=2 spaces, etc.) + + Returns: + Tensor on target device with optional dtype conversion + + Note: + - Skips operation if tensor already has target device and dtype (zero-copy) + - Uses PyTorch's optimized .to() for efficient device/dtype handling + - Logs all operations consistently for tracking and debugging + """ + if tensor is None: + return tensor + + # Get current state + current_device = tensor.device + current_dtype = tensor.dtype + target_dtype = dtype if dtype is not None else current_dtype + + # Check if movement is actually needed + 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: + # Already on target device and dtype - skip + return tensor + + # Determine reason for movement + if reason is None: + if needs_device_move and needs_dtype_change: + reason = "device and dtype conversion" + elif needs_device_move: + reason = "device movement" + else: + reason = "dtype conversion" + + # Log the movement + if debug: + current_device_str = _device_str(current_device) + target_device_str = _device_str(target_device) + + dtype_info = "" + if needs_dtype_change: + dtype_info = f", {current_dtype} โ†’ {target_dtype}" + + debug.log( + f"Moving {tensor_name} from {current_device_str} to {target_device_str}{dtype_info} ({reason})", + category="general", + indent_level=indent_level + ) + + # Perform the operation based on what needs to change + if needs_device_move and needs_dtype_change: + # Both device and dtype need to change + return tensor.to(target_device, dtype=target_dtype, non_blocking=non_blocking) + elif needs_device_move: + # Only device needs to change + return tensor.to(target_device, non_blocking=non_blocking) + else: + # Only dtype needs to change + return tensor.to(dtype=target_dtype) + + +def manage_model_device(model: torch.nn.Module, target_device: torch.device, model_name: str, + debug: Optional['Debug'] = None, reason: Optional[str] = None, + runner: Optional[Any] = None) -> bool: + """ + Move model to target device with optimizations. + Handles BlockSwap-enabled models transparently. + + Args: + model: The model to move + target_device: Target device (torch.device object, e.g., torch.device('cuda:0')) + model_name: Name for logging (e.g., "VAE", "DiT") + debug: Debug instance for logging + reason: Optional custom reason for the movement + runner: Optional runner instance for BlockSwap detection + + Returns: + bool: True if model was moved, False if already on target device + """ + if model is None: + return False + + # Check if this is a BlockSwap-enabled DiT model + is_blockswap_model = False + actual_model = model + if runner and model_name == "DiT": + # Import here to avoid circular dependency + from .blockswap import is_blockswap_enabled + # Check if BlockSwap config exists and is enabled + has_blockswap_config = ( + hasattr(runner, '_dit_block_swap_config') and + is_blockswap_enabled(runner._dit_block_swap_config) + ) + + if has_blockswap_config: + is_blockswap_model = True + # Get the actual model (handle CompatibleDiT wrapper) + if hasattr(model, "dit_model"): + actual_model = model.dit_model + + # Get current device + try: + current_device = next(model.parameters()).device + except StopIteration: + return False + + # Extract device type for comparison (both are torch.device objects) + target_type = target_device.type + current_device_upper = _device_str(current_device) + target_device_upper = _device_str(target_device) + + # Compare normalized device types + if current_device_upper == target_device_upper and not is_blockswap_model: + # Already on target device type, no movement needed + if debug: + debug.log(f"{model_name} already on {current_device_upper}, skipping movement", category="general") + return False + + # Handle BlockSwap models specially + if is_blockswap_model: + return _handle_blockswap_model_movement( + runner, actual_model, current_device, target_device, target_type, + model_name, debug, reason + ) + + # Standard model movement (non-BlockSwap) + return _standard_model_movement( + model, current_device, target_device, target_type, model_name, + debug, reason + ) + + +def _handle_blockswap_model_movement(runner: Any, model: torch.nn.Module, + current_device: torch.device, target_device: torch.device, + target_type: str, model_name: str, + debug: Optional['Debug'] = None, reason: Optional[str] = None) -> bool: + """ + Handle device movement for BlockSwap-enabled models. + + Args: + runner: Runner instance with BlockSwap configuration + model: Model to move (actual unwrapped model) + current_device: Current device of the model + target_device: Target device (torch.device object) + target_type: Target device type (cpu/cuda/mps) + model_name: Model name for logging + debug: Debug instance + reason: Movement reason + + Returns: + bool: True if model was moved + """ + # Import here to avoid circular dependency + from .blockswap import set_blockswap_bypass + + if target_type == "cpu": + # Moving to offload device (typically CPU) + # Check if any parameter is on GPU (for accurate logging) + actual_source_device = None + for param in model.parameters(): + if param.device.type in ['cuda', 'mps']: + actual_source_device = param.device + break + + 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 {_device_str(target_device)} ({reason or 'model caching'})", category="general") + + # Enable bypass to allow movement + set_blockswap_bypass(runner=runner, bypass=True, debug=debug) + + # Start timer + timer_name = f"{model_name.lower()}_to_{target_type}" + if debug: + debug.start_timer(timer_name) + + # Move entire model to target offload device + model.to(target_device) + model.zero_grad(set_to_none=True) + + if debug: + 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(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") + return False + + # Get actual current device for accurate logging + actual_current_device = None + for param in model.parameters(): + if param.device.type != 'meta': + actual_current_device = param.device + break + + 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 {_device_str(target_device)} ({reason or 'inference requirement'})", category="general") + + timer_name = f"{model_name.lower()}_to_gpu" + if debug: + debug.start_timer(timer_name) + + # 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 = model._block_swap_config.get("offload_device") + if not offload_device: + raise ValueError("BlockSwap config missing offload_device") + + # Move blocks according to BlockSwap configuration + for b, block in enumerate(model.blocks): + if b > model.blocks_to_swap: + # This block should be on GPU + block.to(target_device) + else: + # This block stays on offload device (will be swapped during forward) + block.to(offload_device) + + # Handle I/O components + 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": + module.to(target_device) + else: + # I/O components stay on offload device + for name, module in model.named_children(): + if name != "blocks": + module.to(offload_device) + + if debug: + # Get actual configuration from runner + if hasattr(model, '_block_swap_config'): + blocks_on_gpu = model._block_swap_config.get('total_blocks', 32) - model._block_swap_config.get('blocks_swapped', 16) + total_blocks = model._block_swap_config.get('total_blocks', 32) + 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 {_device_str(main_device)})", category="success") + else: + debug.log("BlockSwap blocks restored to configured devices", category="success") + + + # Reactivate BlockSwap now that blocks are restored to their configured devices + runner._blockswap_active = True + + # Disable bypass, re-enable protection + set_blockswap_bypass(runner=runner, bypass=False, debug=debug) + + if debug: + debug.end_timer(timer_name, "BlockSwap model restored") + + return True + + +def _standard_model_movement(model: torch.nn.Module, current_device: torch.device, + target_device: torch.device, target_type: str, model_name: str, + debug: Optional['Debug'] = None, reason: Optional[str] = None) -> bool: + """ + Handle standard (non-BlockSwap) model movement. + + Args: + model: Model to move + current_device: Current device of the model + target_device: Target device (torch.device object) + target_type: Target device type + model_name: Model name for logging + debug: Debug instance + reason: Movement reason + + Returns: + bool: True if model was moved + """ + # Check if model is on meta device - can't move meta tensors + if current_device.type == 'meta': + if debug: + debug.log(f"{model_name} is on meta device - skipping movement (will materialize when needed)", + category=model_name.lower()) + return False + + # Determine reason for movement + reason = reason or "inference requirement" + + # Log the movement with full device strings + if debug: + current_device_str = _device_str(current_device) + 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 + timer_name = f"{model_name.lower()}_to_{'gpu' if target_type != 'cpu' else 'cpu'}" + if debug: + debug.start_timer(timer_name) + + # Move model and clear gradients + model.to(target_device) + model.zero_grad(set_to_none=True) + + # Clear VAE memory buffers when moving to CPU + if target_type == 'cpu' and model_name == "VAE": + cleared_count = 0 + for module in model.modules(): + if hasattr(module, 'memory') and module.memory is not None: + if torch.is_tensor(module.memory) and (module.memory.is_cuda or module.memory.is_mps): + module.memory = None + cleared_count += 1 + if cleared_count > 0 and debug: + debug.log(f"Cleared {cleared_count} VAE memory buffers", category="success") + + # End timer + if debug: + debug.end_timer(timer_name, f"{model_name} moved to {_device_str(target_device)}") + + return True + + +def clear_runtime_caches(runner: Any, debug: Optional['Debug'] = None) -> int: + """ + Clear all runtime caches and temporary attributes. + """ + if not runner: + return 0 + + if debug: + debug.start_timer("runtime_cache_clear") + + cleaned_items = 0 + + # 1. Clear main runner cache + if hasattr(runner, 'cache') and hasattr(runner.cache, 'cache'): + if debug: + debug.start_timer("runner_cache_clear") + + cache_entries = len(runner.cache.cache) + + # Properly release tensor memory and delete as we go + for key in list(runner.cache.cache.keys()): + value = runner.cache.cache[key] + if torch.is_tensor(value): + release_tensor_memory(value) + elif isinstance(value, (list, tuple)): + for item in value: + if torch.is_tensor(item): + release_tensor_memory(item) + # Delete immediately to release reference + del runner.cache.cache[key] + + # Final clear for safety + runner.cache.cache.clear() + cleaned_items += cache_entries + + if debug: + debug.end_timer("runner_cache_clear", f"Clearing main runner cache entries") + + if cache_entries > 0: + debug.log(f"Cleared {cache_entries} runtime cache entries", category="success") + + # 2. Clear RoPE caches + if hasattr(runner, 'dit'): + if debug: + debug.start_timer("rope_cache_clear") + + model = runner.dit + if hasattr(model, 'dit_model'): # Handle wrapper + model = model.dit_model + + rope_cleared = clear_rope_lru_caches(model=model, debug=debug) + cleaned_items += rope_cleared + if debug: + debug.end_timer("rope_cache_clear", "Clearing RoPE LRU caches") + + if rope_cleared > 0: + debug.log(f"Cleared {rope_cleared} RoPE LRU caches", category="success") + + # 3. Clear temporary attributes + temp_attrs = ['_temp_cache', '_block_cache', '_swap_cache', '_generation_cache', + '_rope_cache', '_intermediate_cache', '_backward_cache'] + + for obj in [runner, getattr(runner, 'dit', None), getattr(runner, 'vae', None)]: + if obj is None: + continue + + actual_obj = obj.dit_model if hasattr(obj, 'dit_model') else obj + + for attr in temp_attrs: + if hasattr(actual_obj, attr): + delattr(actual_obj, attr) + cleaned_items += 1 + + if debug: + debug.end_timer("runtime_cache_clear", f"clear_runtime_caches() completion") + + return cleaned_items + + +def cleanup_dit(runner: Any, debug: Optional['Debug'] = None, cache_model: bool = False) -> None: + """ + Cleanup DiT model and BlockSwap state after upscaling phase. + Called at the end of upscale_all_batches when DiT is no longer needed. + + Args: + runner: Runner instance containing DiT model + debug: Debug instance for logging + cache_model: If True, move DiT to offload_device; if False, delete completely + """ + if not runner or not hasattr(runner, 'dit'): + return + + if debug: + debug.log("Cleaning up DiT components", category="cleanup") + + # 1. Clear DiT-specific runtime caches first + if hasattr(runner, 'dit'): + model = runner.dit + if hasattr(model, 'dit_model'): # Handle wrapper + model = model.dit_model + + # Clear RoPE caches + rope_cleared = clear_rope_lru_caches(model=model, debug=debug) + if rope_cleared > 0 and debug: + debug.log(f"Cleared {rope_cleared} RoPE LRU caches", category="success") + + # Clear DiT temporary attributes + temp_attrs = ['_temp_cache', '_block_cache', '_swap_cache', '_generation_cache', + '_rope_cache', '_intermediate_cache', '_backward_cache'] + + actual_obj = model.dit_model if hasattr(model, 'dit_model') else model + for attr in temp_attrs: + if hasattr(actual_obj, attr): + delattr(actual_obj, attr) + + # 2. Handle model offloading (for caching or before deletion) + try: + param_device = next(runner.dit.parameters()).device + + # Move model off GPU if needed + if param_device.type not in ['meta', 'cpu']: + # MPS: skip CPU movement before deletion (unified memory, just causes sync) + if param_device.type == 'mps' and not cache_model: + if debug: + debug.log("DiT on MPS - skipping CPU movement before deletion", category="cleanup") + else: + offload_target = getattr(runner, '_dit_offload_device', None) + if offload_target is None or offload_target == 'none': + offload_target = torch.device('cpu') + 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: + pass + + # 3. Clean BlockSwap after model movement + if hasattr(runner, "_blockswap_active") and runner._blockswap_active: + # Import here to avoid circular dependency + from .blockswap import cleanup_blockswap + cleanup_blockswap(runner=runner, keep_state_for_cache=cache_model) + + # 4. Complete cleanup if not caching + if not cache_model: + release_model_memory(model=runner.dit, debug=debug) + runner.dit = None + if debug: + debug.log("DiT model deleted", category="cleanup") + + # Clear DiT config attributes - not needed when model is not cached (will be recreated) + if hasattr(runner, '_dit_compile_args'): + delattr(runner, '_dit_compile_args') + if hasattr(runner, '_dit_block_swap_config'): + delattr(runner, '_dit_block_swap_config') + if hasattr(runner, '_dit_attention_mode'): + delattr(runner, '_dit_attention_mode') + + # 5. Clear DiT temporary attributes (should be already cleared in materialize_model) + runner._dit_checkpoint = None + runner._dit_dtype_override = None + + # 6. Clear DiT-related components and temporary attributes + runner.sampler = None + runner.sampling_timesteps = None + runner.schedule = None + + +def cleanup_vae(runner: Any, debug: Optional['Debug'] = None, cache_model: bool = False) -> None: + """ + Cleanup VAE model after decoding phase. + Called at the end of decode_all_batches when VAE is no longer needed. + + Args: + runner: Runner instance containing VAE model + debug: Debug instance for logging + cache_model: If True, move VAE to offload_device; if False, delete completely + """ + if not runner or not hasattr(runner, 'vae'): + return + + if debug: + debug.log("Cleaning up VAE components", category="cleanup") + + # 1. Clear VAE-specific temporary attributes + if hasattr(runner, 'vae'): + temp_attrs = ['_temp_cache', '_block_cache', '_swap_cache', '_generation_cache', + '_rope_cache', '_intermediate_cache', '_backward_cache'] + + for attr in temp_attrs: + if hasattr(runner.vae, attr): + delattr(runner.vae, attr) + + # 2. Handle model offloading (for caching or before deletion) + try: + param_device = next(runner.vae.parameters()).device + + # Move model off GPU if needed + if param_device.type not in ['meta', 'cpu']: + # MPS: skip CPU movement before deletion (unified memory, just causes sync) + if param_device.type == 'mps' and not cache_model: + if debug: + debug.log("VAE on MPS - skipping CPU movement before deletion", category="cleanup") + else: + offload_target = getattr(runner, '_vae_offload_device', None) + if offload_target is None or offload_target == 'none': + offload_target = torch.device('cpu') + 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: + pass + + # 3. Complete cleanup if not caching + if not cache_model: + release_model_memory(model=runner.vae, debug=debug) + runner.vae = None + if debug: + debug.log("VAE model deleted", category="cleanup") + + # Clear VAE config attributes - not needed when model is not cached (will be recreated) + if hasattr(runner, '_vae_compile_args'): + delattr(runner, '_vae_compile_args') + if hasattr(runner, '_vae_tiling_config'): + delattr(runner, '_vae_tiling_config') + + # 3. Clear VAE temporary attributes (should be already cleared in materialize_model) + runner._vae_checkpoint = None + runner._vae_dtype_override = None + + +def complete_cleanup(runner: Any, debug: Optional['Debug'] = None, dit_cache: bool = False, vae_cache: bool = False) -> None: + """ + Complete cleanup of runner and remaining components with independent model caching support. + This is a lightweight cleanup for final stage, as model-specific cleanup + happens in their respective phases (cleanup_dit, cleanup_vae). + + Args: + runner: Runner instance to clean up + debug: Debug instance for logging + dit_cache: If True, preserve DiT model on offload_device for future runs + vae_cache: If True, preserve VAE model on offload_device for future runs + + Behavior: + - Can cache DiT and VAE independently for flexible memory management + - Preserves _dit_model_name and _vae_model_name when either model is cached for change detection + - Clears all temporary attributes and runtime caches + - Performs deep memory cleanup only when both models are fully released + + Note: + Model name tracking (_dit_model_name, _vae_model_name) is only cleared if neither + model is cached, enabling proper model change detection on subsequent runs. + """ + if not runner: + return + + if debug: + cleanup_type = "partial cleanup" if (dit_cache or vae_cache) else "full cleanup" + debug.log(f"Starting {cleanup_type}", category="cleanup") + + # 1. Cleanup any remaining models if they still exist + # (This handles cases where phases were skipped or errored) + if hasattr(runner, 'dit') and runner.dit is not None: + cleanup_dit(runner=runner, debug=debug, cache_model=dit_cache) + + if hasattr(runner, 'vae') and runner.vae is not None: + cleanup_vae(runner=runner, debug=debug, cache_model=vae_cache) + + # 2. Clear remaining runtime caches + clear_runtime_caches(runner=runner, debug=debug) + + # 3. Clear config and other non-model components when fully releasing runner + if not (dit_cache or vae_cache): + # Full cleanup - clear config and model tracking + runner.config = None + runner._dit_model_name = None + runner._vae_model_name = None + + # 4. Final memory cleanup + clear_memory(debug=debug, deep=True, force=True, timer_name="complete_cleanup") + + # 5. Clear cuBLAS workspaces + torch._C._cuda_clearCublasWorkspaces() if hasattr(torch._C, '_cuda_clearCublasWorkspaces') else None + + # Log what models are cached for next run + if dit_cache or vae_cache: + cached_models = [] + if dit_cache and hasattr(runner, '_dit_model_name'): + cached_models.append(f"DiT ({runner._dit_model_name})") + if vae_cache and hasattr(runner, '_vae_model_name'): + cached_models.append(f"VAE ({runner._vae_model_name})") + + if cached_models: + models_str = " and ".join(cached_models) + debug.log(f"Models cached for next run: {models_str}", category="cache", force=True) + + if debug: + debug.log(f"Completed {cleanup_type}", category="success") \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/performance.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/performance.py new file mode 100644 index 0000000000000000000000000000000000000000..0e9dd8619322def3372be6489269e5e2498d42c7 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/optimization/performance.py @@ -0,0 +1,169 @@ +""" +Performance optimization module for SeedVR2 +Contains optimized tensor operations and video processing functions + +Extracted from: seedvr2.py (lines 1633-1730) +""" + +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 + Replaces slow loops with vectorized operations + + Transforms: + - 3D: c h w -> t c h w (with t=1) + - 4D: c t h w -> t c h w + + Expected gains: 5-10x faster than naive loops + + Args: + video_tensors: List of video tensors to rearrange + + Returns: + List of rearranged tensors in t c h w format + + Raises: + ValueError: If video tensor has invalid dimensions (not 3D or 4D) + """ + if not video_tensors: + return [] + + # ๐Ÿ” Analyze dimensions to optimize processing + videos_3d = [] + videos_4d = [] + indices_3d = [] + indices_4d = [] + + for i, video in enumerate(video_tensors): + if video.ndim == 3: + videos_3d.append(video) + indices_3d.append(i) + elif video.ndim == 4: + videos_4d.append(video) + indices_4d.append(i) + else: + raise ValueError(f"Video tensor at index {i} has invalid dimensions: {video.ndim}. Expected 3D or 4D.") + + # ๐ŸŽฏ Prepare final result + samples = [None] * len(video_tensors) + + # ๐Ÿš€ BATCH PROCESSING for 3D videos (c h w -> 1 c h w) + if videos_3d: + # Stack + permute (faster than rearrange) + # c h w -> c 1 h w -> 1 c h w + batch_3d = torch.stack([v.unsqueeze(1) for v in videos_3d]) # [batch, c, 1, h, w] + batch_3d = batch_3d.permute(0, 2, 1, 3, 4) # [batch, 1, c, h, w] + + for i, idx in enumerate(indices_3d): + samples[idx] = batch_3d[i] # [1, c, h, w] + + # ๐Ÿš€ BATCH PROCESSING for 4D videos (c t h w -> t c h w) + if videos_4d: + # Check if all 4D videos have the same shape for maximum optimization + shapes = [v.shape for v in videos_4d] + if len(set(shapes)) == 1: + # ๐ŸŽฏ MAXIMUM OPTIMIZATION: All shapes identical + # Stack + permute in single operation + batch_4d = torch.stack(videos_4d) # [batch, c, t, h, w] + batch_4d = batch_4d.permute(0, 2, 1, 3, 4) # [batch, t, c, h, w] + + for i, idx in enumerate(indices_4d): + samples[idx] = batch_4d[i] # [t, c, h, w] + else: + # ๐Ÿ”„ FALLBACK: Different shapes, optimized individual processing + for i, idx in enumerate(indices_4d): + # Use permute instead of rearrange (faster) + samples[idx] = videos_4d[i].permute(1, 0, 2, 3) # c t h w -> t c h w + + return samples + + +def optimized_single_video_rearrange(video: torch.Tensor) -> torch.Tensor: + """ + ๐Ÿš€ OPTIMIZED version for single video tensor + Replaces rearrange() with native PyTorch operations + + Transforms: + - 3D: c h w -> 1 c h w (add temporal dimension) + - 4D: c t h w -> t c h w (permute dimensions) + + Expected gains: 2-5x faster than rearrange() + + Args: + video: Input video tensor + + Returns: + Rearranged tensor with temporal dimension first + """ + if video.ndim == 3: + # c h w -> 1 c h w (add temporal dimension t=1) + return video.unsqueeze(0) + else: # ndim == 4 + # c t h w -> t c h w (permute channels and temporal) + return video.permute(1, 0, 2, 3) + + +def optimized_sample_to_image_format(sample: torch.Tensor) -> torch.Tensor: + """ + ๐Ÿš€ OPTIMIZED version to convert sample to image format + Replaces rearrange() with native PyTorch operations + + Transforms: + - 3D: c h w -> 1 h w c (add temporal dimension + permute to image format) + - 4D: t c h w -> t h w c (permute to image format) + + Expected gains: 2-5x faster than rearrange() + + Args: + sample: Input sample tensor + + Returns: + Tensor in image format (channels last) + """ + if sample.ndim == 3: + # c h w -> 1 h w c (add temporal dimension then permute) + return sample.unsqueeze(0).permute(0, 2, 3, 1) + else: # ndim == 4 + # t c h w -> t h w c (permute channels to last) + return sample.permute(0, 2, 3, 1) + + + diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/__init__.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/color_fix.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/color_fix.py new file mode 100644 index 0000000000000000000000000000000000000000..32ea9435bd3d596a9e5d774fe73acddc28c29b02 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/color_fix.py @@ -0,0 +1,873 @@ +""" +Color Correction Module for SeedVR2 + +Provides perceptually-accurate color correction methods to match upscaled video +frames to their original color characteristics. All methods preserve spatial details +while transferring color distributions. + +Available Methods: +- lab: Full perceptual color matching with optional detail preservation (recommended) +- wavelet: Frequency-based color transfer preserving high-frequency details +- wavelet_adaptive: Wavelet with targeted saturation correction +- hsv: Hue-conditional saturation histogram matching +- adain: Adaptive instance normalization style transfer +""" + +import torch +from PIL import Image +from torch import Tensor +from torch.nn import functional as F +from typing import Optional +from torchvision.transforms import ToTensor, ToPILImage +from ..common.half_precision_fixes import safe_pad_operation, safe_interpolate_operation, ensure_float32_precision + + +def adain_color_fix(target: Image.Image, source: Image.Image) -> Image.Image: + """ + Apply AdaIN color correction to PIL images. + + Args: + target: PIL Image with desired details + source: PIL Image with desired colors + + Returns: + PIL Image with corrected colors + """ + to_tensor = ToTensor() + target_tensor = to_tensor(target).unsqueeze(0) + source_tensor = to_tensor(source).unsqueeze(0) + + result_tensor = adaptive_instance_normalization(target_tensor, source_tensor) + + to_image = ToPILImage() + result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0)) + + return result_image + + +def wavelet_color_fix(target: Image.Image, source: Image.Image, debug: Optional['Debug'] = None) -> Image.Image: + """ + Apply wavelet-based color correction to PIL images. + + Args: + target: PIL Image with desired details + source: PIL Image with desired colors + debug: Debug instance for logging + + Returns: + PIL Image with corrected colors + """ + to_tensor = ToTensor() + target_tensor = to_tensor(target).unsqueeze(0) + source_tensor = to_tensor(source).unsqueeze(0) + + result_tensor = wavelet_reconstruction(target_tensor, source_tensor, debug) + + to_image = ToPILImage() + result_image = to_image(result_tensor.squeeze(0).clamp_(0.0, 1.0)) + + return result_image + + +def calc_mean_std(feat: Tensor, eps: float = 1e-5) -> tuple[Tensor, Tensor]: + """ + Calculate channel-wise mean and standard deviation. + + Args: + feat: 4D tensor [B, C, H, W] + eps: Small value to prevent division by zero + + Returns: + Tuple of (mean, std) tensors [B, C, 1, 1] + """ + size = feat.size() + assert len(size) == 4, 'The input feature should be 4D tensor.' + b, c = size[:2] + + feat_var = feat.reshape(b, c, -1).var(dim=2) + eps + feat_std = feat_var.sqrt().reshape(b, c, 1, 1) + feat_mean = feat.reshape(b, c, -1).mean(dim=2).reshape(b, c, 1, 1) + + return feat_mean, feat_std + + +def adaptive_instance_normalization(content_feat: Tensor, style_feat: Tensor) -> Tensor: + """ + Adaptive Instance Normalization (AdaIN) for style transfer. + + Transfers the color distribution (mean and variance) from style to content + while preserving content structure. + + Args: + content_feat: Target tensor [B, C, H, W] with desired structure + style_feat: Source tensor [B, C, H, W] with desired color statistics + + Returns: + Normalized tensor with style statistics and content structure + """ + style_mean, style_std = calc_mean_std(style_feat) + content_mean, content_std = calc_mean_std(content_feat) + + # Normalize content to zero mean, unit variance (avoid expand by using broadcasting) + normalized_feat = (content_feat - content_mean) / content_std + del content_mean, content_std + + # Apply style statistics (broadcasting avoids expand) + result = normalized_feat * style_std + style_mean + del normalized_feat, style_mean, style_std + + return result + + +def wavelet_blur(image: Tensor, radius: int) -> Tensor: + """ + Apply Gaussian-like blur using dilated convolution for wavelet decomposition. + + Automatically limits radius to prevent numerical instability at high resolutions. + Supports arbitrary number of channels. + + Args: + image: Input tensor [B, C, H, W] + radius: Dilation radius for blur kernel + + Returns: + Blurred tensor [B, C, H, W] + """ + # Prevent excessive dilation that causes OOM/numerical issues + # Conservative limit: 1/8 of smallest spatial dimension + max_safe_radius = max(1, min(image.shape[-2:]) // 8) + if radius > max_safe_radius: + radius = max_safe_radius + + num_channels = image.shape[1] + + # 3x3 Gaussian-approximation kernel + kernel_vals = [ + [0.0625, 0.125, 0.0625], + [0.125, 0.25, 0.125], + [0.0625, 0.125, 0.0625], + ] + kernel = torch.tensor(kernel_vals, dtype=image.dtype, device=image.device) + kernel = kernel[None, None].repeat(num_channels, 1, 1, 1) + + # Apply padding and grouped convolution + image = safe_pad_operation(image, (radius, radius, radius, radius), mode='replicate') + output = F.conv2d(image, kernel, groups=num_channels, dilation=radius) + + return output + + +def wavelet_decomposition(image: Tensor, levels: int = 5) -> tuple[Tensor, Tensor]: + """ + Multi-scale wavelet decomposition to separate frequency components. + + Decomposes image into high-frequency (details/edges) and low-frequency + (color/illumination) components using iterative Gaussian pyramid. + + Args: + image: Input tensor [B, C, H, W] + levels: Number of decomposition levels (default: 5) + + Returns: + Tuple of (high_freq, low_freq) tensors + - high_freq: Detail information [B, C, H, W] + - low_freq: Color/illumination information [B, C, H, W] + """ + high_freq = torch.zeros_like(image) + + for i in range(levels): + radius = 2 ** i + low_freq = wavelet_blur(image, radius) + high_freq.add_(image).sub_(low_freq) + image = low_freq + + return high_freq, low_freq + + +def wavelet_reconstruction(content_feat: Tensor, style_feat: Tensor, debug: Optional['Debug'] = None) -> Tensor: + """ + Apply wavelet-based color transfer from style to content. + + Preserves high-frequency details from content while adopting low-frequency + color information from style using multi-resolution wavelet decomposition. + + Algorithm: + 1. Decompose both images into high (detail) and low (color) frequency components + 2. Combine content's high frequencies with style's low frequencies + 3. Reconstruct the image preserving details with transferred colors + + Args: + content_feat: Target tensor with desired details [B, C, H, W] in [-1,1] + style_feat: Source tensor with desired colors [B, C, H, W] in [-1,1] + debug: Debug instance for logging + + Returns: + Tensor: Reconstructed tensor with content details and style colors in [-1,1] + """ + # Handle dimension mismatch if needed + if content_feat.shape != style_feat.shape: + debug.log(f"Dimension mismatch: content {content_feat.shape} vs style {style_feat.shape}", + level="WARNING", category="precision", force=True) + + # Resize style to match content spatial dimensions + if len(content_feat.shape) >= 3: + # safe_interpolate_operation handles FP16 conversion automatically + style_feat = safe_interpolate_operation( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False + ) + debug.log(f"Style resized to: {style_feat.shape}", category="precision", force=True) + + # Decompose both features into frequency components + content_high_freq, content_low_freq = wavelet_decomposition(content_feat) + del content_low_freq # Free memory immediately + + style_high_freq, style_low_freq = wavelet_decomposition(style_feat) + del style_high_freq # Free memory immediately + + # Safety check (should not happen after resize) + if content_high_freq.shape != style_low_freq.shape: + debug.log(f"Final dimension adjustment needed", level="WARNING", category="precision", force=True) + style_low_freq = safe_interpolate_operation( + style_low_freq, + size=content_high_freq.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # Reconstruct: content details + style color (in-place on content_high_freq) + content_high_freq.add_(style_low_freq) + + # 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 content_high_freq.clamp_(-1.0, 1.0) + + +def lab_color_transfer( + content_feat: Tensor, + style_feat: Tensor, + debug, + luminance_weight: float = 0.8 +) -> Tensor: + """ + Perceptually-accurate color transfer using CIELAB color space. + + Combines wavelet reconstruction (for spatial continuity) with LAB histogram + matching (for precise color matching). This eliminates tile artifacts while + providing superior color accuracy. + + Algorithm: + 1. Apply wavelet reconstruction to get artifact-free base + 2. Convert both wavelet result and style to LAB color space + 3. Apply histogram matching to LAB channels + 4. Convert back to RGB + + Args: + content_feat: Target tensor [B, C, H, W] in [-1, 1] with upscaled details + style_feat: Source tensor [B, C, H, W] in [-1, 1] with original colors + debug: Debug instance for logging + luminance_weight: How much content luminance to preserve (0.0-1.0) + 0.0 = full color match, 1.0 = preserve all detail + Default: 0.8 (slight color adjustment, strong detail preservation) + + Returns: + Color-corrected tensor [B, C, H, W] in [-1, 1] + """ + # 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}", + level="WARNING", category="precision", force=True + ) + style_feat = safe_interpolate_operation( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # 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 + rgb_to_xyz_matrix = torch.tensor([ + [0.4124564, 0.3575761, 0.1804375], + [0.2126729, 0.7151522, 0.0721750], + [0.0193339, 0.1191920, 0.9503041] + ], dtype=torch.float32, device=device) + + xyz_to_rgb_matrix = torch.tensor([ + [ 3.2404542, -1.5371385, -0.4985314], + [-0.9692660, 1.8760108, 0.0415560], + [ 0.0556434, -0.2040259, 1.0572252] + ], dtype=torch.float32, device=device) + + # LAB conversion constants + epsilon = 6.0 / 29.0 + kappa = (29.0 / 3.0) ** 3 + + # Convert from [-1, 1] to [0, 1] range (in-place) + content_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) + + # Convert to LAB color space + content_lab = _rgb_to_lab_batch(content_feat, device, rgb_to_xyz_matrix, epsilon, kappa) + del content_feat + + style_lab = _rgb_to_lab_batch(style_feat, device, rgb_to_xyz_matrix, epsilon, kappa) + 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) + matched_b = _histogram_matching_channel(content_lab[:, 2], style_lab[:, 2], device) + + # Handle luminance with weighted blending + if luminance_weight < 1.0: + # Partially match luminance for better overall color accuracy + matched_L = _histogram_matching_channel(content_lab[:, 0], style_lab[:, 0], device) + # Blend: preserve some content L* for detail, adopt some style L* for color + result_L = content_lab[:, 0].mul(luminance_weight).add_(matched_L.mul(1.0 - luminance_weight)) + del matched_L + else: + # Fully preserve content luminance + result_L = content_lab[:, 0] + + del content_lab, style_lab + + # Reconstruct LAB with corrected channels + result_lab = torch.stack([result_L, matched_a, matched_b], dim=1) + del result_L, matched_a, matched_b + + # Convert back to RGB + result_rgb = _lab_to_rgb_batch(result_lab, device, xyz_to_rgb_matrix, epsilon, kappa) + del result_lab, xyz_to_rgb_matrix + + # Convert back to [-1, 1] range (in-place) + result = result_rgb.mul_(2.0).sub_(1.0) + del result_rgb + + # Restore original dtype + if result.dtype != original_dtype: + result = result.to(original_dtype) + + debug.log(f"LAB color transfer completed (luminance_weight={luminance_weight})", category="video", indent_level=1) + + return result + + +def _rgb_to_lab_batch(rgb: Tensor, device: torch.device, matrix: Tensor, epsilon: float, kappa: float) -> Tensor: + """Convert batch of RGB images to CIELAB color space using D65 illuminant.""" + # Apply sRGB gamma correction (linearize) + mask = rgb > 0.04045 + rgb_linear = torch.where( + mask, + torch.pow((rgb + 0.055) / 1.055, 2.4), + rgb / 12.92 + ) + del mask + + # Matrix multiplication: RGB -> XYZ + B, C, H, W = rgb_linear.shape + 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 + + xyz = xyz_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2) + del xyz_flat + + # Normalize by D65 white point (in-place) + xyz[:, 0].div_(0.95047) # X + # xyz[:, 1] /= 1.00000 # Y (no-op, skip) + xyz[:, 2].div_(1.08883) # Z + + # XYZ to LAB transformation + epsilon_cubed = epsilon ** 3 + mask = xyz > epsilon_cubed + f_xyz = torch.where( + mask, + torch.pow(xyz, 1.0 / 3.0), + xyz.mul(kappa).add_(16.0).div_(116.0) + ) + del xyz, mask + + # Extract channels and compute LAB + L = f_xyz[:, 1].mul(116.0).sub_(16.0) # Lightness [0, 100] + a = (f_xyz[:, 0] - f_xyz[:, 1]).mul_(500.0) # Green-Red [-128, 127] + b = (f_xyz[:, 1] - f_xyz[:, 2]).mul_(200.0) # Blue-Yellow [-128, 127] + del f_xyz + + return torch.stack([L, a, b], dim=1) + + +def _lab_to_rgb_batch(lab: Tensor, device: torch.device, matrix_inv: Tensor, epsilon: float, kappa: float) -> Tensor: + """Convert batch of CIELAB images to RGB color space.""" + L, a, b = lab[:, 0], lab[:, 1], lab[:, 2] + + # LAB to XYZ + fy = (L + 16.0) / 116.0 + fx = a.div(500.0).add_(fy) + fz = fy - b / 200.0 + del L, a, b + + # XYZ transformation + x = torch.where( + fx > epsilon, + torch.pow(fx, 3.0), + fx.mul(116.0).sub_(16.0).div_(kappa) + ) + y = torch.where( + fy > epsilon, + torch.pow(fy, 3.0), + fy.mul(116.0).sub_(16.0).div_(kappa) + ) + z = torch.where( + fz > epsilon, + torch.pow(fz, 3.0), + fz.mul(116.0).sub_(16.0).div_(kappa) + ) + del fx, fy, fz + + # Apply D65 white point (in-place) + x.mul_(0.95047) + # y *= 1.00000 # (no-op, skip) + z.mul_(1.08883) + + xyz = torch.stack([x, y, z], dim=1) + del x, y, z + + # Matrix multiplication: XYZ -> RGB + B, C, H, W = xyz.shape + 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 + + rgb_linear = rgb_linear_flat.reshape(B, H, W, 3).permute(0, 3, 1, 2) + del rgb_linear_flat + + # Apply inverse gamma correction (delinearize) + mask = rgb_linear > 0.0031308 + rgb = torch.where( + mask, + torch.pow(torch.clamp(rgb_linear, min=0.0), 1.0 / 2.4).mul_(1.055).sub_(0.055), + rgb_linear * 12.92 + ) + del mask, rgb_linear + + return torch.clamp(rgb, 0.0, 1.0) + + +def _histogram_matching_channel(source: Tensor, reference: Tensor, device: torch.device) -> Tensor: + """ + Match histogram of source channel to reference using CDF mapping. + + Args: + source: Source channel tensor [B, H, W] + reference: Reference channel tensor [B, H, W] + + Returns: + Matched channel tensor [B, H, W] + """ + original_shape = source.shape + + # Flatten + source_flat = source.flatten() + reference_flat = reference.flatten() + + # 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) + n_reference = len(reference_sorted) + + if n_source == n_reference: + matched_sorted = reference_sorted + else: + # Interpolate reference to match source quantiles + source_quantiles = torch.linspace(0, 1, n_source, device=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, reference_sorted + + del source_sorted, source_flat + + # Reconstruct using argsort (portable across CUDA/ROCm/MPS) + 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) + + +def hsv_saturation_histogram_match(content_feat: Tensor, style_feat: Tensor, debug: Optional['Debug'] = None) -> Tensor: + """ + Hue-conditional saturation histogram matching in HSV color space. + + Matches saturation distribution from style to content separately for each hue bin + to handle color-specific oversaturation (e.g., overly saturated reds). + + Based on: Neumann et al. (2005) "Color Style Transfer using Hue, Lightness + and Saturation Histogram Matching" + + Algorithm: + 1. Convert both images to HSV color space + 2. Divide hue circle into 12 bins (30ยฐ each) + 3. For each hue bin, match saturation histograms independently + 4. Reconstruct HSV with matched saturation + 5. Convert back to RGB + + Args: + content_feat: Target tensor [B, C, H, W] in [-1, 1] with upscaled details + style_feat: Source tensor [B, C, H, W] in [-1, 1] with original saturation + debug: Debug instance for logging + + Returns: + Saturation-corrected tensor [B, C, H, W] in [-1, 1] + """ + # Handle spatial dimension mismatch + if content_feat.shape != style_feat.shape: + debug.log( + f"HSV: Resizing style {style_feat.shape} to match content {content_feat.shape}", + level="WARNING", category="precision", force=True + ) + style_feat = safe_interpolate_operation( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # Convert to float32 for processing + content_feat, original_dtype = ensure_float32_precision(content_feat) + style_feat, _ = ensure_float32_precision(style_feat) + + # 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) + style_rgb = style_feat.add(1.0).mul_(0.5).clamp_(0.0, 1.0) + del content_feat, style_feat + + # Convert to HSV + content_hsv = _rgb_to_hsv_batch(content_rgb) + del content_rgb + + style_hsv = _rgb_to_hsv_batch(style_rgb) + del style_rgb + + # Extract channels + content_h = content_hsv[:, 0] + content_s = content_hsv[:, 1] + content_v = content_hsv[:, 2] + del content_hsv + + style_h = style_hsv[:, 0] + style_s = style_hsv[:, 1] + del style_hsv + + # Match saturation per hue bin + matched_s = _hue_conditional_saturation_match(content_h, content_s, style_h, style_s) + del style_h, style_s + + # Reconstruct HSV: preserve H and V from content, use matched S + result_hsv = torch.stack([content_h, matched_s, content_v], dim=1) + del content_h, matched_s, content_v + + # Convert back to RGB + result_rgb = _hsv_to_rgb_batch(result_hsv) + del result_hsv + result_rgb.clamp_(0.0, 1.0) + + # Convert back to [-1, 1] range (in-place) + result = result_rgb.mul_(2.0).sub_(1.0) + del result_rgb + + # Restore original dtype + if result.dtype != original_dtype: + result = result.to(original_dtype) + + debug.log("HSV hue-conditional saturation matching completed", category="video", indent_level=1) + + return result + + +def _rgb_to_hsv_batch(rgb: Tensor) -> Tensor: + """Convert batch of RGB images to HSV color space. All channels in [0, 1].""" + r, g, b = rgb[:, 0], rgb[:, 1], rgb[:, 2] + + maxc = torch.max(rgb, dim=1)[0] + minc = torch.min(rgb, dim=1)[0] + rangec = maxc - minc + + # Avoid division by zero + rangec_nz = torch.where(rangec > 1e-10, rangec, torch.ones_like(rangec)) + + # Hue calculation (in [0, 1]) + h = torch.zeros_like(maxc) + + mask_r = (maxc == r) & (rangec > 1e-10) + h[mask_r] = ((g[mask_r] - b[mask_r]) / rangec_nz[mask_r]) % 6.0 + del mask_r + + mask_g = (maxc == g) & (rangec > 1e-10) + h[mask_g] = ((b[mask_g] - r[mask_g]) / rangec_nz[mask_g]) + 2.0 + del mask_g + + mask_b = (maxc == b) & (rangec > 1e-10) + h[mask_b] = ((r[mask_b] - g[mask_b]) / rangec_nz[mask_b]) + 4.0 + del mask_b + + h.div_(6.0) # Normalize to [0, 1] (in-place) + + # Saturation calculation + s = torch.where(maxc > 1e-10, rangec / torch.clamp(maxc, min=1e-10), torch.zeros_like(maxc)) + del rangec, rangec_nz + + # Value calculation + v = maxc + + return torch.stack([h, s, v], dim=1) + + +def _hsv_to_rgb_batch(hsv: Tensor) -> Tensor: + """Convert batch of HSV images to RGB color space.""" + h = hsv[:, 0].mul(6.0) # Convert to [0, 6] + s = hsv[:, 1] + v = hsv[:, 2] + + i = torch.floor(h).long() % 6 + f = h - torch.floor(h) + + p = v * (1.0 - s) + q = v * (1.0 - s * f) + t = v * (1.0 - s * (1.0 - f)) + + # Initialize RGB channels + r = torch.zeros_like(v) + g = torch.zeros_like(v) + b = torch.zeros_like(v) + + # Apply hue sector transformations + mask0 = (i == 0) + r[mask0], g[mask0], b[mask0] = v[mask0], t[mask0], p[mask0] + del mask0 + + mask1 = (i == 1) + r[mask1], g[mask1], b[mask1] = q[mask1], v[mask1], p[mask1] + del mask1 + + mask2 = (i == 2) + r[mask2], g[mask2], b[mask2] = p[mask2], v[mask2], t[mask2] + del mask2 + + mask3 = (i == 3) + r[mask3], g[mask3], b[mask3] = p[mask3], q[mask3], v[mask3] + del mask3 + + mask4 = (i == 4) + r[mask4], g[mask4], b[mask4] = t[mask4], p[mask4], v[mask4] + del mask4 + + mask5 = (i == 5) + r[mask5], g[mask5], b[mask5] = v[mask5], p[mask5], q[mask5] + del mask5, i, f, p, q, t, h, s, v + + return torch.stack([r, g, b], dim=1) + + +def _hue_conditional_saturation_match( + content_h: Tensor, + content_s: Tensor, + style_h: Tensor, + style_s: Tensor +) -> Tensor: + """ + Match saturation histogram conditionally per hue bin. + + Divides hue circle into 12 bins (30ยฐ each) and matches saturation + separately for each bin to handle color-specific oversaturation. + """ + num_bins = 12 + bin_width = 1.0 / num_bins + min_pixels = 100 # Minimum pixels for reliable histogram matching + + matched_s = content_s.clone() + device = content_s.device + + for bin_idx in range(num_bins): + bin_start = bin_idx * bin_width + bin_end = (bin_idx + 1) * bin_width + + # Handle hue wrap-around for red (0ยฐ/360ยฐ) + if bin_idx == 0: + content_mask = ((content_h >= 0) & (content_h < bin_end)) | (content_h >= (1.0 - bin_width)) + style_mask = ((style_h >= 0) & (style_h < bin_end)) | (style_h >= (1.0 - bin_width)) + else: + content_mask = (content_h >= bin_start) & (content_h < bin_end) + style_mask = (style_h >= bin_start) & (style_h < bin_end) + + # Extract saturation values for this hue bin + content_s_bin = content_s[content_mask] + style_s_bin = style_s[style_mask] + + # Only match if both bins have sufficient pixels + if len(content_s_bin) > min_pixels and len(style_s_bin) > min_pixels: + matched_s_bin = _histogram_match_1d(content_s_bin, style_s_bin, device) + matched_s[content_mask] = matched_s_bin + del matched_s_bin + + del content_mask, style_mask, content_s_bin, style_s_bin + + return matched_s + + +def _histogram_match_1d(source: Tensor, reference: Tensor, device: torch.device) -> Tensor: + """Match 1D histogram using CDF mapping.""" + source_sorted, source_indices = torch.sort(source) + reference_sorted, _ = torch.sort(reference) + + n_source = len(source_sorted) + n_reference = len(reference_sorted) + + if n_source == n_reference: + matched_sorted = reference_sorted + else: + source_quantiles = torch.linspace(0, 1, n_source, device=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, reference_sorted + + del source_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 + + +def wavelet_adaptive_color_correction(content_feat: Tensor, style_feat: Tensor, debug: Optional['Debug'] = None) -> Tensor: + """ + Adaptive hybrid color correction combining wavelet and HSV methods. + + Uses wavelet as the base correction for natural colors, then selectively + applies HSV saturation correction only to oversaturated regions. + + Algorithm: + 1. Apply wavelet reconstruction (natural color base) + 2. Apply HSV saturation matching (targeted correction) + 3. Detect oversaturated pixels by comparing saturation levels + 4. Blend HSV correction only into oversaturated areas via sigmoid + 5. Keep wavelet colors everywhere else + + Args: + content_feat: Target tensor [B, C, H, W] in [-1, 1] with upscaled details + style_feat: Source tensor [B, C, H, W] in [-1, 1] with original colors + debug: Debug instance for logging + + Returns: + Adaptively corrected tensor [B, C, H, W] in [-1, 1] + """ + # Handle spatial dimension mismatch + if content_feat.shape != style_feat.shape: + debug.log( + f"Wavelet Adaptive: Resizing style {style_feat.shape} to match content {content_feat.shape}", + level="WARNING", category="precision", force=True + ) + style_feat = safe_interpolate_operation( + style_feat, + size=content_feat.shape[-2:], + mode='bilinear', + align_corners=False + ) + + # Convert to float32 for processing + content_feat, original_dtype = ensure_float32_precision(content_feat) + style_feat, _ = ensure_float32_precision(style_feat) + + # Step 1: Apply wavelet (base correction) + wavelet_result = wavelet_reconstruction(content_feat, style_feat, debug) + + # Step 2: Apply HSV saturation matching (targeted correction) + hsv_result = hsv_saturation_histogram_match(content_feat, style_feat, debug) + + # Step 3: Compute saturation maps to detect oversaturation + content_sat = _get_saturation_map(content_feat) + style_sat = _get_saturation_map(style_feat) + wavelet_sat = _get_saturation_map(wavelet_result) + + # Step 4: Create adaptive blend mask based on saturation difference + sat_difference = content_sat - style_sat + del content_sat + + # Parameters for blending + oversaturation_threshold = 0.15 # Saturation difference threshold + blend_sharpness = 5.0 # Sigmoid sharpness for smooth transitions + + # Sigmoid blend: 0 = use wavelet, 1 = use HSV correction + blend_weight = torch.sigmoid(blend_sharpness * (sat_difference - oversaturation_threshold)) + del sat_difference + + # Only correct if wavelet itself is still oversaturated + wavelet_oversaturated = (wavelet_sat - style_sat) > (oversaturation_threshold * 0.5) + del wavelet_sat, style_sat + + blend_weight.mul_(wavelet_oversaturated.float()) + del wavelet_oversaturated + blend_weight.clamp_(0.0, 1.0) + + # Step 5: Adaptive blending (in-place where possible) + result = wavelet_result.mul(1.0 - blend_weight).add_(hsv_result.mul(blend_weight)) + + # Log statistics + correction_pct = (blend_weight > 0.01).float().mean().item() * 100 + del blend_weight, wavelet_result, hsv_result + + # Restore original dtype + if result.dtype != original_dtype: + result = result.to(original_dtype) + + debug.log(f"Wavelet Adaptive: {correction_pct:.1f}% pixels use HSV correction", category="video", indent_level=1) + + return result + + +def _get_saturation_map(tensor: Tensor) -> Tensor: + """Extract saturation channel from RGB tensor in [-1, 1].""" + rgb = tensor.add(1.0).mul_(0.5).clamp_(0.0, 1.0) + + maxc = torch.max(rgb, dim=1, keepdim=True)[0] + minc = torch.min(rgb, dim=1, keepdim=True)[0] + del rgb + + saturation = torch.where( + maxc > 1e-10, + (maxc - minc) / torch.clamp(maxc, min=1e-10), + torch.zeros_like(maxc) + ) + del maxc, minc + + return saturation \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/constants.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..088d538253df6ea7011b4920be0036634441923f --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/constants.py @@ -0,0 +1,161 @@ +""" +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 +from typing import Optional + +# Model folder names +SEEDVR2_FOLDER_NAME = "SEEDVR2" # Physical folder name on disk +SEEDVR2_MODEL_TYPE = "seedvr2" # Model type identifier for ComfyUI + +# Supported model file formats +SUPPORTED_MODEL_EXTENSIONS = {'.safetensors', '.gguf'} + +# GGUF Quantization Constants +QK_K = 256 +K_SCALE_SIZE = 12 +GGUF_BLOCK_SIZE = 32 +GGUF_TYPE_SIZE = 64 + +# Download configuration +HUGGINGFACE_BASE_URL = "https://huggingface.co/{repo}/resolve/main/{filename}" +DOWNLOAD_CHUNK_SIZE = 8192 * 1024 # 8MB chunks for hash calculation +DOWNLOAD_MAX_RETRIES = 3 +DOWNLOAD_RETRY_DELAY = 2 # seconds + +def get_script_directory() -> str: + """Get the root script directory path (3 levels up from this file)""" + return os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + + +def get_base_cache_dir() -> str: + """ + Get the default model cache directory path. + + Returns the path without creating the directory. + + Returns: + str: Path to default cache directory + """ + try: + import folder_paths # Only works if ComfyUI is available + cache_dir = os.path.join(folder_paths.models_dir, SEEDVR2_FOLDER_NAME) + folder_paths.add_model_folder_path(SEEDVR2_MODEL_TYPE, cache_dir) + except: + cache_dir = f"./models/{SEEDVR2_FOLDER_NAME}" + + return cache_dir + + +def get_all_model_paths() -> list: + """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() + + # Case-insensitive lookup: search through all registered folder types + # 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()] + + +def get_all_model_files() -> dict: + """ + Get a mapping of all model files to their full paths across all registered directories. + + Returns: + dict: Mapping of filename -> full path for all discovered model files + """ + model_files = {} + all_paths = get_all_model_paths() + + for path in all_paths: + if os.path.exists(path): + for file in os.listdir(path): + if is_supported_model_file(file): + # Only keep first occurrence of each file (priority order) + if file not in model_files: + model_files[file] = os.path.join(path, file) + + return model_files + + +def find_model_file(filename: str, fallback_dir: Optional[str] = None) -> str: + """ + Find a model file in any registered path. + + Args: + filename: Name of the model file to find + fallback_dir: Directory to use if file not found in any registered path + + Returns: + str: Full path to the model file + """ + # Get all model files + model_files = get_all_model_files() + + # Return path if found + if filename in model_files: + return model_files[filename] + + # Fallback to specified directory or base cache dir + if fallback_dir: + return os.path.join(fallback_dir, filename) + else: + return os.path.join(get_base_cache_dir(), filename) + + +def get_validation_cache_path(cache_dir: Optional[str] = None) -> str: + """ + Get path to model validation cache file. + + Args: + cache_dir: Optional directory for cache file. If None, uses default base cache dir. + + Returns: + Full path to validation cache JSON file + """ + if cache_dir is None: + cache_dir = get_base_cache_dir() + return os.path.join(cache_dir, ".validation_cache.json") + + +def is_supported_model_file(filename: str) -> bool: + """Check if a file has a supported model extension""" + return any(filename.endswith(ext) for ext in SUPPORTED_MODEL_EXTENSIONS) + + +def suppress_tensor_warnings() -> None: + """ + Suppress common tensor conversion and numpy array warnings that are expected behavior + when working with GGUF tensors and numpy arrays. + """ + warnings.filterwarnings("ignore", message="To copy construct from a tensor", category=UserWarning) + warnings.filterwarnings("ignore", message="The given NumPy array is not writable", category=UserWarning) \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/debug.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/debug.py new file mode 100644 index 0000000000000000000000000000000000000000..f5059de3cb6db80f441f5422aafe09eea46e3ca0 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/debug.py @@ -0,0 +1,800 @@ +""" +Unified debugging system for SeedVR2 generation pipeline + +Provides structured logging, memory tracking, and performance monitoring +for all pipeline stages, including BlockSwap operations. +""" + +import time +import torch +import gc +from typing import Optional, List, Dict, Any, Union +from datetime import datetime +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: + """ + Unified debug logging for generation pipeline and BlockSwap monitoring + + Features: + - Structured logging with categories + - Memory tracking (VRAM/RAM) + - Timing utilities + - BlockSwap operation tracking + - Minimal overhead when disabled + - Timestamped logs for better troubleshooting + - Force parameters for critical logs + """ + + # Icon mapping for different categories + CATEGORY_ICONS = { + "general": "๐Ÿ”„", # General operations/processing + "timing": "โšก", # Performance timing + "memory": "๐Ÿ“Š", # Memory usage tracking + "cache": "๐Ÿ’พ", # Cache operations + "cleanup": "๐Ÿงน", # Cleanup operations + "setup": "๐Ÿ”ง", # Configuration/setup + "generation": "๐ŸŽฌ", # Generation process + "dit": "๐Ÿš€", # Model loading/operations + "blockswap": "๐Ÿ”€", # BlockSwap operations + "download": "๐Ÿ“ฅ", # Download operations + "success": "โœ…", # Successful completion + "warning": "โš ๏ธ", # Warnings + "error": "โŒ", # Errors + "info": "โ„น๏ธ", # Statistics/info + "tip" :"๐Ÿ’ก", # Tip/suggestion + "video": "๐Ÿ“น", # Video/sequence info + "reuse": "โ™ป๏ธ", # Reusing/recycling + "runner": "๐Ÿƒ", # Runner operations + "vae": "๐ŸŽจ", # VAE operations\ + "precision": "๐ŸŽฏ", # Precision + "device": "๐Ÿ–ฅ๏ธ", # Device info + "file": "๐Ÿ“‚", # File operations + "alpha": "๐Ÿ‘ป", # Alpha operations + "starlove": "โญ๐Ÿ’", # Star + love + "dialogue": "๐Ÿ’ฌ", # Dialogue + "none" : "", + } + + def __init__(self, enabled: bool = False, show_timestamps: bool = True): + self.enabled = enabled + self.show_timestamps = show_timestamps + self.timers: Dict[str, float] = {} + self.memory_checkpoints: List[Dict[str, Any]] = [] + self.max_checkpoints = 100 + self.timer_hierarchy: Dict[str, List[str]] = {} + self.timer_durations: Dict[str, float] = {} + self.timer_messages: Dict[str, str] = {} + self.swap_times: List[Dict[str, Any]] = [] + self.current_phase: Optional[str] = None + self.vram_history: List[float] = [] + self.active_timer_stack: List[str] = [] + self.timer_namespace: str = "" + 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: + """ + Log a categorized message with optional timestamp and indentation + + Args: + message: Message to log + level: Log level (INFO, WARN, ERROR) + category: Category for the message + force: If True, always log regardless of enabled state (for critical messages) + indent_level: Indentation level (0=no indent, 1=2 spaces, 2=4 spaces, etc.) + """ + # Always log forced messages or if debugging is enabled - early return if not + if not (self.enabled or force): + return + + # Get icon for category, fallback to general icon + icon = self.CATEGORY_ICONS.get(category, self.CATEGORY_ICONS["general"]) + + # Format prefix based on level + if level == "WARNING": + icon = self.CATEGORY_ICONS["warning"] + elif level == "ERROR": + icon = self.CATEGORY_ICONS["error"] + + # Build the log message with optional timestamp + if self.show_timestamps: + timestamp = datetime.now().strftime("%H:%M:%S.%f")[:-3] + prefix = f"[{timestamp}] {icon}" + else: + prefix = f"{icon}" + + if level != "INFO": + prefix += f" [{level}]" + + # Add indentation + indent = " " * (indent_level * 2) + + print(f"{prefix} {indent}{message}", flush=True) + + def print_header(self, cli: bool = False) -> None: + """Print the header with banner - always displayed""" + # 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, indent_level=1) + self.log("โ–ˆโ–ˆโ•”โ•โ•โ•โ•โ•โ–ˆโ–ˆโ•”โ•โ•โ•โ•โ•โ–ˆโ–ˆโ•”โ•โ•โ•โ•โ•โ–ˆโ–ˆโ•”โ•โ•โ–ˆโ–ˆโ•—โ–ˆโ–ˆโ•‘ โ–ˆโ–ˆโ•‘โ–ˆโ–ˆโ•”โ•โ•โ–ˆโ–ˆโ•— โ•šโ•โ•โ•โ•โ–ˆโ–ˆโ•— โ–ˆโ–ˆโ•”โ•โ•โ•โ•โ•", category="none", force=True, indent_level=1) + self.log("โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•—โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•— โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•— โ–ˆโ–ˆโ•‘ โ–ˆโ–ˆโ•‘โ–ˆโ–ˆโ•‘ โ–ˆโ–ˆโ•‘โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•”โ• โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•”โ• โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•—", category="none", force=True, indent_level=1) + self.log("โ•šโ•โ•โ•โ•โ–ˆโ–ˆโ•‘โ–ˆโ–ˆโ•”โ•โ•โ• โ–ˆโ–ˆโ•”โ•โ•โ• โ–ˆโ–ˆโ•‘ โ–ˆโ–ˆโ•‘โ•šโ–ˆโ–ˆโ•— โ–ˆโ–ˆโ•”โ•โ–ˆโ–ˆโ•”โ•โ•โ–ˆโ–ˆโ•— โ–ˆโ–ˆโ•”โ•โ•โ•โ• โ•šโ•โ•โ•โ•โ–ˆโ–ˆโ•‘", category="none", force=True, indent_level=1) + self.log("โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•‘โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•—โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•—โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•”โ• โ•šโ–ˆโ–ˆโ–ˆโ–ˆโ•”โ• โ–ˆโ–ˆโ•‘ โ–ˆโ–ˆโ•‘ โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•— โ–ˆโ–ˆโ•— โ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ–ˆโ•‘", category="none", force=True, indent_level=1) + self.log("โ•šโ•โ•โ•โ•โ•โ•โ•โ•šโ•โ•โ•โ•โ•โ•โ•โ•šโ•โ•โ•โ•โ•โ•โ•โ•šโ•โ•โ•โ•โ•โ• โ•šโ•โ•โ•โ• โ•šโ•โ• โ•šโ•โ• โ•šโ•โ•โ•โ•โ•โ•โ• โ•šโ•โ• โ•šโ•โ•โ•โ•โ•โ•โ•", category="none", force=True, indent_level=1) + # 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: + try: + distro = platform.freedesktop_os_release() + 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, 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="starlove", force=True) + + @torch._dynamo.disable # Skip tracing to avoid time.time() warnings + def start_timer(self, name: str, force: bool = False) -> None: + """ + Start a named timer + + Args: + name: Timer name + force: If True, start timer even when debug is disabled + """ + if self.enabled or force: + # Apply namespace if set + if self.timer_namespace: + name = f"{self.timer_namespace}_{name}" + + self.timers[name] = time.time() + + # Track phase for memory peak monitoring + if name.startswith("phase") and name.endswith(("_encoding", "_upscaling", "_decoding", "_postprocessing")): + # Extract phase number (e.g., "phase3_decoding" -> "3") + phase_num = name.split("_")[0].replace("phase", "") + self.current_phase = f"phase{phase_num}" + + # Auto-hierarchy: if there's an active timer, this is a child + if self.active_timer_stack: + parent = self.active_timer_stack[-1] + if parent not in self.timer_hierarchy: + self.timer_hierarchy[parent] = [] + # Only add if not already a child (prevents duplicates) + if name not in self.timer_hierarchy[parent]: + self.timer_hierarchy[parent].append(name) + + # Push to stack + self.active_timer_stack.append(name) + + @torch._dynamo.disable # Skip tracing to avoid time.time() warnings + def end_timer(self, name: str, message: Optional[str] = None, + force: bool = False, show_breakdown: bool = False, + custom_children: Optional[Dict[str, float]] = None) -> float: + """ + End a timer and optionally log its duration + + Args: + name: Timer name + message: Optional message to log with the duration + force: If True, log even when debug is disabled (for critical timings) + show_breakdown: If True, show breakdown of child timers + custom_children: Optional dict of child timer names and durations to override automatic hierarchy + + Returns: + Duration in seconds (0.0 if timer not found) + """ + # Apply namespace if set + if self.timer_namespace: + name = f"{self.timer_namespace}_{name}" + + # Check if timer exists + if name not in self.timers: + return 0.0 + + duration = time.time() - self.timers[name] + self.timer_durations[name] = duration + # Store the message for later use in summary + if message: + self.timer_messages[name] = message + del self.timers[name] + + # Pop from stack if this is the current active timer + if self.active_timer_stack and self.active_timer_stack[-1] == name: + self.active_timer_stack.pop() + + # If debug is disabled and not forcing, return early + if not self.enabled and not force: + return duration + + # ONLY log if show_breakdown is True - this means it's a major summary timer + if message and show_breakdown: + # Use custom children if provided, otherwise use automatic hierarchy + if custom_children: + children = custom_children + child_total = sum(children.values()) + unaccounted = duration - child_total + + self.log(f"{message}: {duration:.2f}s", category="timing", force=force) + + # Sort custom children by duration for better readability + sorted_children = sorted(children.items(), key=lambda x: x[1], reverse=True) + + for child_name, child_duration in sorted_children: + if child_duration >= 0.01: # Only show if >= 10ms + self.log(f"โ””โ”€ {child_name}: {child_duration:.2f}s", category="timing", force=force, indent_level=1) + else: + # Use automatic hierarchy tracking + children = self.timer_hierarchy.get(name, []) + child_total = sum(self.timer_durations.get(child, 0) for child in children) + unaccounted = duration - child_total + + self.log(f"{message}: {duration:.2f}s", category="timing", force=force) + + # Sort children by duration for better readability + sorted_children = sorted(children, key=lambda c: self.timer_durations.get(c, 0), reverse=True) + + for child in sorted_children: + child_duration = self.timer_durations.get(child, 0) + if child_duration >= 0.01: # Only show if >= 10ms + child_message = self.timer_messages.get(child, child) + self.log(f"โ””โ”€ {child_message}: {child_duration:.2f}s", category="timing", force=force, indent_level=1) + + # Recursively show grandchildren + if child in self.timer_hierarchy: + grandchildren = self.timer_hierarchy[child] + sorted_grandchildren = sorted(grandchildren, key=lambda c: self.timer_durations.get(c, 0), reverse=True) + + for grandchild in sorted_grandchildren: + grandchild_duration = self.timer_durations.get(grandchild, 0) + if grandchild_duration >= 0.01: # Only show if >= 10ms + grandchild_message = self.timer_messages.get(grandchild, grandchild) + self.log(f"โ””โ”€ {grandchild_message}: {grandchild_duration:.2f}s", category="timing", force=force, indent_level=2) + + if unaccounted > 0.01: # Show if more than 10ms unaccounted + self.log(f"โ””โ”€ (other operations): {unaccounted:.2f}s", category="timing", force=force, indent_level=1) + + return duration + + def log_memory_state(self, label: str, show_diff: bool = True, show_tensors: bool = False, + detailed_tensors: bool = False, force: bool = False) -> None: + """ + Log current memory state with minimal overhead. + + Args: + label: Description for this checkpoint + show_diff: Show change from last checkpoint + show_tensors: Include tensor counts + detailed_tensors: Show detailed tensor analysis (use sparingly) + force: If True, always log regardless of enabled state + """ + if not (self.enabled or force): + return + + # Collect memory metrics efficiently + memory_info = self._collect_memory_metrics() + + # Show category + self.log(f"{label}:", category="memory", force=force) + + # Show VRAM + if memory_info['summary_vram']: + self.log(f"{memory_info['summary_vram']}", category="memory", force=force) + + # Show RAM + if memory_info['summary_ram']: + self.log(f"{memory_info['summary_ram']}", category="memory", force=force) + + # Show tensors + if show_tensors: + tensor_stats = self._collect_tensor_stats(detailed=detailed_tensors) + self.log(f"{tensor_stats['summary']}", category="memory", force=force) + + # Show diff from last checkpoint + if show_diff and self.memory_checkpoints: + self._log_memory_diff(current_metrics=memory_info, force=force) + + # 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 peaks if we're in an active phase + if self.current_phase: + if memory_info['vram_peak_alloc'] > 0: + self.phase_vram_peaks_alloc[self.current_phase] = max( + 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.""" + 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_alloc': 0.0, + 'vram_peak_rsv': 0.0, + 'vram_overflow': 0.0, + 'ram_process': 0.0, + 'ram_available': 0.0, + 'ram_total': 0.0, + 'ram_others': 0.0, + 'summary_vram': "", + 'summary_ram': "" + } + + if has_gpu: + metrics['vram_allocated'], metrics['vram_reserved'], metrics['vram_peak_alloc'], metrics['vram_peak_rsv'] = get_vram_usage(device=None, debug=self) + vram_info = get_basic_vram_info(device=None) + + 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 = "Unified Memory" if is_mps else "VRAM" + metrics['summary_vram'] = ( + f" [{backend}] {metrics['vram_allocated']:.2f}GB allocated / " + f"{metrics['vram_reserved']:.2f}GB reserved / " + f"Peak: {metrics['vram_peak_alloc']:.2f}GB / " + f"{metrics['vram_free']:.2f}GB free / " + f"{metrics['vram_total']:.2f}GB total" + ) + + self.vram_history.append(metrics['vram_reserved']) + + # 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 / " + f"{metrics['ram_others']:.2f}GB others / " + f"{metrics['ram_available']:.2f}GB free / " + f"{metrics['ram_total']:.2f}GB total" + ) + + return metrics + + def _collect_tensor_stats(self, detailed: bool = False) -> Dict[str, Any]: + """Collect tensor statistics with minimal overhead.""" + stats = { + 'gpu_count': 0, + 'cpu_count': 0, + 'total_count': 0, + 'summary': "", + 'details': None + } + + if detailed: + stats['details'] = { + 'gpu_tensors': [], + 'large_cpu_tensors': [], + 'shape_patterns': {}, + 'module_types': {} + } + + # Single pass through gc objects + for obj in gc.get_objects(): + try: + if torch.is_tensor(obj): + stats['total_count'] += 1 + is_gpu = obj.is_cuda or (hasattr(obj, 'is_mps') and obj.is_mps) + + if is_gpu: + stats['gpu_count'] += 1 + else: + stats['cpu_count'] += 1 + + # Collect detailed info if requested + if detailed and obj.numel() > 0: + size_mb = obj.element_size() * obj.nelement() / (1024**2) + + if is_gpu or size_mb > 10: # Only track GPU tensors or large CPU tensors + tensor_info = { + 'shape': tuple(obj.shape), + 'dtype': str(obj.dtype), + 'size_mb': size_mb, + 'requires_grad': obj.requires_grad + } + + if is_gpu: + stats['details']['gpu_tensors'].append(tensor_info) + elif size_mb > 10: # Large CPU tensors (>10MB) + stats['details']['large_cpu_tensors'].append(tensor_info) + + # Track shape patterns + shape_key = str(tuple(obj.shape)) + stats['details']['shape_patterns'][shape_key] = stats['details']['shape_patterns'].get(shape_key, 0) + 1 + + elif detailed and isinstance(obj, torch.nn.Module): + module_type = type(obj).__name__ + stats['details']['module_types'][module_type] = stats['details']['module_types'].get(module_type, 0) + 1 + + except (ReferenceError, AttributeError): + # Object was deleted or doesn't have expected attributes + pass + + stats['summary'] = f" [Tensors] {stats['gpu_count']} GPU / {stats['cpu_count']} CPU / {stats['total_count']} total" + + return stats + + def _log_detailed_tensor_analysis(self, details: Dict[str, Any], force: bool = False) -> None: + """Log detailed tensor analysis when requested.""" + + # GPU tensors + if details['gpu_tensors']: + gpu_total_gb = sum(t['size_mb'] for t in details['gpu_tensors']) / 1024 + self.log(f"GPU tensors: {len(details['gpu_tensors'])} using {gpu_total_gb:.2f}GB", category="memory", force=force, indent_level=1) + + # Show top 5 largest + largest = sorted(details['gpu_tensors'], key=lambda x: x['size_mb'], reverse=True)[:5] + for t in largest: + self.log(f"{t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force, indent_level=1) + + # Large CPU tensors + if details['large_cpu_tensors']: + cpu_large_gb = sum(t['size_mb'] for t in details['large_cpu_tensors']) / 1024 + self.log(f"Large CPU tensors (>10MB):", category="memory", force=force, indent_level=1) + self.log(f"{len(details['large_cpu_tensors'])} using {cpu_large_gb:.2f}GB", category="memory", force=force, indent_level=1) + + # Show top 3 largest + largest = sorted(details['large_cpu_tensors'], key=lambda x: x['size_mb'], reverse=True)[:3] + for t in largest: + self.log(f"{t['shape']}: {t['size_mb']:.2f}MB, {t['dtype']}", category="memory", force=force, indent_level=1) + + # Common shape patterns + if details['shape_patterns']: + common_shapes = sorted(details['shape_patterns'].items(), + key=lambda x: x[1], reverse=True)[:5] + if len(common_shapes) > 0: + self.log("Common tensor shapes:", category="memory", force=force, indent_level=1) + for shape, count in common_shapes: + if count > 1: + self.log(f"{shape}: {count} instances", category="memory", force=force, indent_level=1) + + # Module instances + if details['module_types']: + multi_instance = [(k, v) for k, v in details['module_types'].items() if v > 1] + if multi_instance: + self.log("Multiple module instances:", category="memory", force=force, indent_level=1) + for mtype, count in sorted(multi_instance, key=lambda x: x[1], reverse=True)[:5]: + self.log(f"{mtype}: {count} instances", category="memory", force=force, indent_level=1) + + def _log_memory_diff(self, current_metrics: Dict[str, Any], force: bool = False) -> None: + """Log memory changes from last checkpoint.""" + last = self.memory_checkpoints[-1] + + vram_diff = current_metrics['vram_allocated'] - last.get('vram_allocated', 0) + ram_diff = current_metrics['ram_process'] - last.get('ram_process', 0) + + diffs = [] + if abs(vram_diff) > 0.01: + sign = "+" if vram_diff > 0 else "" + diffs.append(f"VRAM {sign}{vram_diff:.2f}GB") + if abs(ram_diff) > 0.01: + sign = "+" if ram_diff > 0 else "" + diffs.append(f"RAM {sign}{ram_diff:.2f}GB") + + if diffs: + self.log(f"Memory changes: {', '.join(diffs)}", category="memory", force=force, indent_level=1) + + def log_peak_memory_summary(self, force: bool = True) -> None: + """Display peak memory usage across all phases.""" + if not self.phase_vram_peaks_alloc and not self.phase_ram_peaks: + return + + phase_names = { + 'phase1': 'VAE encoding', + 'phase2': 'DiT upscaling', + 'phase3': 'VAE decoding', + '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 memory by phase:", category="memory", force=force) + + 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] + phase_name = phase_names.get(phase_key, phase_key) + 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_alloc = max(self.phase_vram_peaks_alloc.values()) if self.phase_vram_peaks_alloc else 0 + 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: + """Store checkpoint with memory limit to prevent leaks.""" + checkpoint = { + 'label': label, + 'timestamp': time.time(), + 'vram_allocated': metrics['vram_allocated'], + 'vram_reserved': metrics['vram_reserved'], + 'vram_free': metrics['vram_free'], + 'ram_process': metrics['ram_process'], + 'ram_available': metrics['ram_available'], + 'ram_others': metrics['ram_others'] + } + + self.memory_checkpoints.append(checkpoint) + + # Prevent memory leak by limiting checkpoint history + if len(self.memory_checkpoints) > self.max_checkpoints: + # Keep first and last N/2 checkpoints for better history coverage + mid = self.max_checkpoints // 2 + self.memory_checkpoints = (self.memory_checkpoints[:mid] + + self.memory_checkpoints[-mid:]) + + def log_swap_time(self, component_id: Union[int, str], duration: float, + component_type: str = "block", force: bool = False) -> None: + """ + Log swap timing information for BlockSwap operations + + Args: + component_id: Identifier for the component being swapped + duration: Duration of the swap in seconds + component_type: Type of component ('block' or other) + force: If True, always log regardless of enabled state + """ + if self.enabled or force: + # Store timing data + self.swap_times.append({ + 'component_id': component_id, + 'component_type': component_type, + 'duration': duration, + }) + + # Format message based on component type + if component_type == "block": + message = f"Block {component_id} swap: {duration*1000:.2f}ms" + else: + message = f"{component_type} {component_id} swap: {duration*1000:.2f}ms" + + self.log(message, category="blockswap", force=force) + + def get_swap_summary(self) -> Dict[str, Any]: + """Get summary of swap operations for analysis""" + if not self.swap_times: + return {} + + # Group by component type + block_swaps = [s for s in self.swap_times if s['component_type'] == 'block'] + io_swaps = [s for s in self.swap_times if s['component_type'] != 'block'] + + # Calculate statistics + summary = { + 'total_swaps': len(self.swap_times), + 'block_swaps': len(block_swaps), + 'io_swaps': len(io_swaps), + } + + if block_swaps: + block_times = [s['duration'] for s in block_swaps] + summary['block_avg_ms'] = sum(block_times) * 1000 / len(block_times) + summary['block_total_ms'] = sum(block_times) * 1000 + summary['block_min_ms'] = min(block_times) * 1000 + summary['block_max_ms'] = max(block_times) * 1000 + + # Track which blocks are swapped most frequently + block_frequency = {} + for swap in block_swaps: + block_id = swap['component_id'] + block_frequency[block_id] = block_frequency.get(block_id, 0) + 1 + summary['most_swapped_block'] = max(block_frequency, key=block_frequency.get) + summary['most_swapped_count'] = block_frequency[summary['most_swapped_block']] + + if io_swaps: + io_times = [s['duration'] for s in io_swaps] + summary['io_avg_ms'] = sum(io_times) * 1000 / len(io_times) + summary['io_total_ms'] = sum(io_times) * 1000 + + # Track which I/O components are swapped + io_components = list(set(s['component_id'] for s in io_swaps)) + summary['io_components_swapped'] = io_components + + # VRAM efficiency metrics + if self.vram_history: + summary['peak_vram_gb'] = max(self.vram_history) + summary['avg_vram_gb'] = sum(self.vram_history) / len(self.vram_history) + summary['vram_variation_gb'] = max(self.vram_history) - min(self.vram_history) + + return summary + + def clear_history(self) -> None: + """Clear all history tracking""" + self.timers.clear() + self.memory_checkpoints.clear() + self.swap_times.clear() + self.vram_history.clear() + self.timer_hierarchy.clear() + self.timer_durations.clear() + self.timer_messages.clear() + self.active_timer_stack.clear() + self.phase_vram_peaks_alloc.clear() + self.phase_vram_peaks_rsv.clear() + self.phase_ram_peaks.clear() + self.current_phase = None \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/downloads.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/downloads.py new file mode 100644 index 0000000000000000000000000000000000000000..2569514f95e5399935432a594a8320604eb4a690 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/downloads.py @@ -0,0 +1,284 @@ +""" +Downloads utility module for SeedVR2 +Handles model and VAE downloads from HuggingFace repositories with integrity validation +""" + +import os +import hashlib +import json +import urllib.request +from typing import Optional +from tqdm import tqdm +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, + DOWNLOAD_CHUNK_SIZE, + DOWNLOAD_MAX_RETRIES, + DOWNLOAD_RETRY_DELAY +) + +def load_validation_cache(cache_dir: Optional[str] = None): + """ + Load validation cache. + + Args: + cache_dir: Optional directory containing cache file. If None, uses default. + """ + cache_path = get_validation_cache_path(cache_dir) + if os.path.exists(cache_path): + try: + with open(cache_path, 'r') as f: + return json.load(f) + except: + pass + return {} + + +def save_validation_cache(cache, cache_dir: Optional[str] = None): + """ + Save validation cache. + + Args: + cache: Cache dictionary to save + cache_dir: Optional directory for cache file. If None, uses default. + """ + cache_path = get_validation_cache_path(cache_dir) + try: + # Ensure directory exists + os.makedirs(os.path.dirname(cache_path), exist_ok=True) + with open(cache_path, 'w') as f: + json.dump(cache, f, indent=2) + except: + pass + + +def is_file_validated_cached(filepath: str, cache_dir: Optional[str] = None) -> bool: + """ + Check if file is validated using cache (fast). + + Args: + filepath: Path to file to check + cache_dir: Optional directory containing cache file. If None, uses default. + """ + if not os.path.exists(filepath): + return False + + cache = load_validation_cache(cache_dir) + filename = os.path.basename(filepath) + + if filename in cache: + cached = cache[filename] + # Check if file hasn't changed since validation + if (cached.get('size') == os.path.getsize(filepath) and + abs(cached.get('mtime', 0) - os.path.getmtime(filepath)) < 2): + return True + + return False + + +def validate_file(filepath: str, expected_hash: Optional[str] = None, cache_dir: Optional[str] = None) -> bool: + """ + File validation with hash check and cache update. + + Args: + filepath: Path to file to validate + expected_hash: Optional SHA256 hash to verify against + cache_dir: Optional directory for cache file. If None, uses default. + """ + if not os.path.exists(filepath): + return False + + file_size = os.path.getsize(filepath) + if file_size == 0: + return False + + # Quick safetensors header check + if filepath.endswith('.safetensors'): + try: + with open(filepath, 'rb') as f: + header_size = int.from_bytes(f.read(8), 'little') + if header_size <= 0 or header_size > file_size: + return False + except: + return False + + # Only calculate hash if we have an expected hash to compare + if expected_hash: + sha256 = hashlib.sha256() + with open(filepath, "rb") as f: + while chunk := f.read(DOWNLOAD_CHUNK_SIZE): + sha256.update(chunk) + actual_hash = sha256.hexdigest() + + if actual_hash != expected_hash: + return False + + # Update cache with validated info + cache = load_validation_cache(cache_dir) + cache[os.path.basename(filepath)] = { + "size": os.path.getsize(filepath), + "mtime": os.path.getmtime(filepath), + "hash": expected_hash + } + save_validation_cache(cache, cache_dir) + + return True + + +def download_with_resume(url: str, filepath: str, debug=None) -> bool: + """Download with resume support and progress bar""" + temp_file = f"{filepath}.download" + existing_size = os.path.getsize(temp_file) if os.path.exists(temp_file) else 0 + + headers = {'Range': f'bytes={existing_size}-'} if existing_size > 0 else {} + + try: + req = urllib.request.Request(url, headers=headers) + with urllib.request.urlopen(req, timeout=30) as response: + content_length = int(response.headers.get('Content-Length', 0)) + total_size = existing_size + content_length + + pbar = tqdm(total=total_size, initial=existing_size, unit='B', + unit_scale=True, unit_divisor=1024, + desc=os.path.basename(filepath)) + + with open(temp_file, 'ab' if existing_size else 'wb') as f: + while chunk := response.read(DOWNLOAD_CHUNK_SIZE): + f.write(chunk) + pbar.update(len(chunk)) + pbar.close() + + os.replace(temp_file, filepath) + return True + + except Exception as e: + if debug: + debug.log(f"Download error: {e}", level="ERROR", category="download", force=True) + return False + + +def download_weight(dit_model: str, vae_model: str, model_dir: Optional[str] = None, debug=None) -> bool: + """Download SeedVR2 DiT and VAE models with integrity checking""" + cache_dir = model_dir or get_base_cache_dir() + os.makedirs(cache_dir, exist_ok=True) + + files_to_download = [ + (dit_model, MODEL_REGISTRY.get(dit_model)), + (vae_model, MODEL_REGISTRY.get(vae_model)) + ] + + for filename, model_info in files_to_download: + if not model_info: + if debug: + debug.log(f"{filename} not in registry, skipping validation", + 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 + if os.path.exists(filepath): + if debug: + 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 + if debug: + debug.log(f"File corrupted: {filename}, re-downloading...", + level="WARNING", category="download", force=True) + os.remove(filepath) + # Clear from cache + cache = load_validation_cache(cache_dir) + if filename in cache: + del cache[filename] + save_validation_cache(cache, cache_dir) + + # Download file + url = HUGGINGFACE_BASE_URL.format(repo=repo, filename=filename) + temp_file = f"{filepath}.download" + + if os.path.exists(temp_file) and debug: + size_gb = os.path.getsize(temp_file) / (1024**3) + debug.log(f"Resuming {filepath} from {size_gb:.2f}GB (source: {url})", + category="download", force=True) + elif debug: + debug.log(f"Downloading {filepath} from {url}...", + category="download", force=True) + + # Download with retries + success = False + for attempt in range(DOWNLOAD_MAX_RETRIES): + if attempt > 0: + time.sleep(DOWNLOAD_RETRY_DELAY * attempt) + if debug: + debug.log(f"Retry {attempt}/{DOWNLOAD_MAX_RETRIES}", + category="download", force=True) + + if download_with_resume(url, filepath, debug): + # Validate downloaded file + if validate_file(filepath, expected_hash): + if debug: + debug.log(f"Downloaded and validated: {filename}", + category="success", force=True) + success = True + break + else: + # Remove corrupted download + for f in [filepath, temp_file]: + if os.path.exists(f): + os.remove(f) + if debug: + debug.log(f"Downloaded file failed validation, retrying...", + level="WARNING", category="download", force=True) + + if not success: + if debug: + debug.log(f"Failed to download {filename} after {DOWNLOAD_MAX_RETRIES} attempts", + level="ERROR", category="download", force=True) + debug.log(f"Manual download: https://huggingface.co/{repo}/blob/main/{filename}", + category="info", force=True) + debug.log(f"Save to: {filepath}", category="info", force=True) + return False + + return True \ No newline at end of file diff --git a/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/model_registry.py b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/model_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..c3628e20bd86af19f5dd0d74ae7daba8789f74b0 --- /dev/null +++ b/v3-nodes/ComfyUI-SeedVR2_VideoUpscaler/src/utils/model_registry.py @@ -0,0 +1,91 @@ +""" +Model Registry for SeedVR2 +Central registry for model definitions, repositories, and metadata +""" + +import os +from typing import List, Optional +from dataclasses import dataclass +from .constants import get_all_model_files + +# Model class imports using relative imports +from ..models.dit_3b.nadit import NaDiT as NaDiT3B +from ..models.dit_7b.nadit import NaDiT as NaDiT7B +from ..models.video_vae_v3.modules.attn_video_vae import VideoAutoencoderKLWrapper + +# Model classes - simple registry with clear keys +MODEL_CLASSES = { + "dit_3b.nadit": NaDiT3B, + "dit_7b.nadit": NaDiT7B, + "video_vae_v3.modules.attn_video_vae": VideoAutoencoderKLWrapper, +} + +@dataclass +class ModelInfo: + """Model metadata""" + repo: str = "numz/SeedVR2_comfyUI" + category: str = "dit" # 'model' or 'vae' + precision: str = "fp16" # 'fp16', 'fp8_e4m3fn', 'Q4_K_M', etc. + size: str = "3B" # '3B', '7B', etc. + variant: Optional[str] = None # 'sharp', etc. + sha256: Optional[str] = None # Cached hash + +# Model registry with metadata +MODEL_REGISTRY = { + # 3B models + "seedvr2_ema_3b-Q4_K_M.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="3B", precision="Q4_K_M", sha256="e665e3909de1a8c88a69c609bca9d43ff5a134647face2ce4497640cc3597f0e"), + "seedvr2_ema_3b-Q8_0.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="3B", precision="Q8_0", sha256="be0d60083a2051a265eb4b77f28edf494e6db67ffc250216f32b72292e5cbd96"), + "seedvr2_ema_3b_fp8_e4m3fn.safetensors": ModelInfo(size="3B", precision="fp8_e4m3fn", sha256="3bf1e43ebedd570e7e7a0b1b60d6a02e105978f505c8128a241cde99a8240cff"), + "seedvr2_ema_3b_fp16.safetensors": ModelInfo(size="3B", precision="fp16", sha256="2fd0e03a3dad24e07086750360727ca437de4ecd456f769856e960ae93e2b304"), + + # 7B models + "seedvr2_ema_7b-Q4_K_M.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="Q4_K_M", sha256="db9cb2ad90ebd40d2e8c29da2b3fc6fd03ba87cd58cbadceccca13ad27162789"), + "seedvr2_ema_7b_fp8_e4m3fn_mixed_block35_fp16.safetensors": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="fp8_e4m3fn_mixed_block35_fp16", sha256="3d68b5ec0b295ae28092e355c8cad870edd00b817b26587d0cb8f9dd2df19bb2"), + "seedvr2_ema_7b_fp16.safetensors": ModelInfo(size="7B", precision="fp16", sha256="7b8241aa957606ab6cfb66edabc96d43234f9819c5392b44d2492d9f0b0bbe4a"), + + # 7B sharp variants + "seedvr2_ema_7b_sharp-Q4_K_M.gguf": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="Q4_K_M", variant="sharp", sha256="7aed800ac4eb8e0d18569a954c0ff35f5a1caa3ed5d920e66cc31405f75b6e69"), + "seedvr2_ema_7b_sharp_fp8_e4m3fn_mixed_block35_fp16.safetensors": ModelInfo(repo="AInVFX/SeedVR2_comfyUI", size="7B", precision="fp8_e4m3fn_mixed_block35_fp16", variant="sharp", sha256="0d2c5b8be0fda94351149c5115da26aef4f4932a7a2a928c6f184dda9186e0be"), + "seedvr2_ema_7b_sharp_fp16.safetensors": ModelInfo(size="7B", precision="fp16", variant="sharp", sha256="20a93e01ff24beaeebc5de4e4e5be924359606c356c9c51509fba245bd2d77dd"), + + # VAE models + "ema_vae_fp16.safetensors": ModelInfo(category="vae", precision="fp16", sha256="20678548f420d98d26f11442d3528f8b8c94e57ee046ef93dbb7633da8612ca1"), +} + +# Configuration constants +DEFAULT_DIT = "seedvr2_ema_3b_fp8_e4m3fn.safetensors" +DEFAULT_VAE = "ema_vae_fp16.safetensors" + +def get_default_models(category: str) -> List[str]: + """Get list of default models""" + return [name for name, info in MODEL_REGISTRY.items() if info.category == category] + +def get_model_repo(model_name: str) -> str: + """Get repository for a specific model""" + return MODEL_REGISTRY.get(model_name, ModelInfo()).repo + +def get_available_dit_models() -> List[str]: + """Get all available DiT models including those discovered on disk""" + model_list = get_default_models("dit") + + try: + # Get all model files from all paths + model_files = get_all_model_files() + + # Add files not in registry + discovered_models = [ + filename for filename in model_files + if filename not in MODEL_REGISTRY + ] + + # Add discovered models to the list + model_list.extend(sorted(discovered_models)) + except: + pass + + return model_list + +def get_available_vae_models() -> List[str]: + """Get all available VAE models from the registry""" + model_list = get_default_models("vae") + return model_list \ No newline at end of file