From 024e627b471ac1ee04755039d5cb5f23b019762d Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Sat, 15 Aug 2026 21:16:14 +0200 Subject: [PATCH] feat: add SPEED and LDF-VFI interpolation Integrate checksum-pinned runtimes, harden interpolation and cleanup paths, and add a SPEED/LDF model-lab workflow. --- NOTICE | 11 + README.md | 126 ++- __init__.py | 12 + assets/model-comparison.svg | 40 +- .../tween_speed_ldf_model_lab.json | 757 ++++++++++++++++++ external_sources.py | 173 ++++ ldf_backend.py | 437 ++++++++++ nodes.py | 712 +++++++++++++--- pyproject.toml | 7 +- requirements.txt | 3 + speed_backend.py | 158 ++++ 11 files changed, 2286 insertions(+), 150 deletions(-) create mode 100644 NOTICE create mode 100644 example_workflows/tween_speed_ldf_model_lab.json create mode 100644 external_sources.py create mode 100644 ldf_backend.py create mode 100644 speed_backend.py diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..0c9db9c --- /dev/null +++ b/NOTICE @@ -0,0 +1,11 @@ +ComfyUI-Tween + +The LDF-VFI sequence sampling adapter in ldf_backend.py is based on the +official LDF-VFI implementation: +https://github.com/xypeng9903/LDF-VFI + +LDF-VFI is distributed under the Apache License, Version 2.0. Its official +runtime is downloaded on demand and retains its upstream LICENSE file. + +SPEED runtime source is not redistributed by this project. It is downloaded +on demand from the official repository at a checksum-pinned commit. diff --git a/README.md b/README.md index f340146..6d24cc8 100644 --- a/README.md +++ b/README.md @@ -3,9 +3,9 @@ [![ComfyUI](https://img.shields.io/badge/ComfyUI-Custom_Node-0a7ef0)](https://registry.comfy.org/) [![Python 3.10+](https://img.shields.io/badge/Python-3.10+-3776AB?logo=python&logoColor=white)](https://www.python.org/) [![License](https://img.shields.io/badge/License-Apache_2.0-green.svg)](https://www.apache.org/licenses/LICENSE-2.0) -[![Models](https://img.shields.io/badge/VFI_Models-4-8B5CF6)](#which-model-should-i-use) +[![Models](https://img.shields.io/badge/VFI_Models-6-8B5CF6)](#which-model-should-i-use) -Four video frame interpolation models in one package — **BIM-VFI**, **EMA-VFI**, **SGM-VFI**, and **GIMM-VFI**. Designed for long videos with thousands of frames without running out of VRAM. +Six video frame interpolation models in one package — **BIM-VFI**, **EMA-VFI**, **SGM-VFI**, **GIMM-VFI**, **SPEED**, and **LDF-VFI**. Pairwise models include chunked/segmented processing; LDF-VFI adds holistic long-sequence diffusion interpolation.

Model Comparison @@ -21,11 +21,19 @@ git clone https://github.com/Ethanfel/ComfyUI-Tween.git pip install -r requirements.txt ``` -All dependencies (`gdown`, `timm`, `omegaconf`, `easydict`, `yacs`, `einops`, `huggingface_hub`) are declared in `pyproject.toml` and `requirements.txt`, installed automatically by ComfyUI Manager or pip. +Dependencies are declared in `pyproject.toml` and `requirements.txt` and are installed automatically by ComfyUI Manager or pip. LDF-VFI requires PyTorch 2.5+ plus a current `diffusers`/`accelerate` stack. -### cupy (required for BIM-VFI, SGM-VFI, GIMM-VFI) +### Demo workflow -[cupy](https://cupy.dev/) provides GPU-accelerated optical flow warping. **EMA-VFI works without it.** +Import [`example_workflows/tween_speed_ldf_model_lab.json`](example_workflows/tween_speed_ldf_model_lab.json) for the recommended starter graph. It requires [ComfyUI-VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite) for video loading and encoding. + +- The enabled SPEED branch loads a 25-frame, 24 FPS sample, automatically tunes memory settings, interpolates to 48 FPS, preserves audio, and writes `Tween/demo_speed_24_to_48`. +- The LDF-VFI branch is visibly grouped and muted by default so the workflow does not unexpectedly download ~6.4 GB or reserve ~20 GB VRAM. Enable its three coral nodes when you want to compare the sequence-native model. +- Keep the loader's `force_rate`, Tween's `source_fps`/`target_fps`, and Video Combine's `frame_rate` synchronized when changing cadence. + +### cupy (accelerates BIM-VFI, SGM-VFI, and GIMM-VFI) + +[cupy](https://cupy.dev/) provides GPU-accelerated optical flow warping. **EMA-VFI, SPEED, and LDF-VFI do not use it.** 1. Find your CUDA version: ```bash @@ -55,19 +63,16 @@ All dependencies (`gdown`, `timm`, `omegaconf`, `easydict`, `yacs`, `einops`, `h ## Which model should I use? -| | BIM-VFI | EMA-VFI | SGM-VFI | GIMM-VFI | -|---|---------|---------|---------|----------| -| **Best for** | General-purpose | Fast, low VRAM | Large motion | High multipliers (4x/8x) | -| **Quality** | Highest | Good | Best on large motion | Good | -| **Speed** | Moderate | Fastest | Slowest | Fast for 4x/8x | -| **VRAM** | ~2 GB/pair | ~1.5 GB/pair | ~3 GB/pair | ~2.5 GB/pair | -| **Params** | ~17 M | ~14–65 M | ~15 M + GMFlow | ~80 M (RAFT) / ~123 M (FlowFormer) | -| **Arbitrary timestep** | Yes | Yes (`_t` checkpoint) | No (fixed 0.5) | Yes (native) | -| **4x/8x** | Recursive passes | Recursive passes | Recursive passes | Single forward pass | -| **Requires cupy** | Yes | No | Yes | Yes | -| **Paper** | CVPR 2025 | CVPR 2023 | CVPR 2024 | NeurIPS 2024 | +| Model | Best for | Multiplier path | Typical VRAM | Trade-off | +|-------|----------|-----------------|--------------|-----------| +| **BIM-VFI** | Strong general pairwise quality | Recursive 2x/4x/8x | ~2 GB/pair | Research/education license | +| **EMA-VFI** | Speed and lower VRAM | Recursive 2x/4x/8x | ~1.5 GB/pair | Less robust on extreme motion | +| **SGM-VFI** | Large motion | Recursive 2x/4x/8x | ~3 GB/pair | Slowest pairwise option | +| **GIMM-VFI** | Arbitrary timesteps, efficient 4x/8x | Native multi-frame per pair | ~2.5 GB/pair | Still frame-pair-centric | +| **SPEED** | New high-quality midpoint generation | One diffusion step at 2x; recursive 4x/8x | ~2.3–2.6 GB at benchmark resolutions | Stochastic, ~447 MB checkpoint | +| **LDF-VFI** | Long-range temporal coherence and 2x–16x | Native sequence diffusion | ~20 GB | ~6.4 GB weights; much slower | -**TL;DR:** Start with **BIM-VFI** for best quality. Use **EMA-VFI** for speed or if you can't install cupy. Use **SGM-VFI** for large camera motion. Use **GIMM-VFI** for 4x/8x without recursive passes. +**TL;DR:** Try **SPEED** as the modern pairwise default. Use **EMA-VFI** when latency matters, **SGM-VFI** for difficult large motion, **GIMM-VFI** for lightweight arbitrary timesteps, and **LDF-VFI** when sequence consistency matters more than speed or memory. ## VRAM Guide @@ -78,9 +83,11 @@ All dependencies (`gdown`, `timm`, `omegaconf`, `easydict`, `yacs`, `einops`, `h | 48 GB+ | `batch_size=4–16, all_on_gpu=true` | | 96 GB+ | `batch_size=8–16, all_on_gpu=true, chunk_size=0` | +SPEED generally fits the 24 GB tier at HD resolutions. LDF-VFI is a separate workload: its official 8x quick start requires about 20 GB, and higher resolutions may require smaller VAE tiles or more VRAM. + ## Nodes -All Interpolate nodes share a common set of controls: +The pairwise Interpolate nodes (BIM/EMA/SGM/GIMM/SPEED) share these controls: | Input | Description | |-------|-------------| @@ -93,7 +100,7 @@ All Interpolate nodes share a common set of controls: | **all_on_gpu** | Keep all intermediate frames on GPU (fast, needs large VRAM) | | **clear_cache_after_n_frames** | Clear CUDA cache every N pairs to prevent VRAM buildup | | **source_fps** | Input frame rate. Required when target_fps > 0 | -| **target_fps** | Target output FPS. When > 0, overrides multiplier — auto-computes the optimal power-of-2 oversample then selects frames at exact target timestamps. 0 = use multiplier | +| **target_fps** | Target output FPS. When > 0, overrides multiplier — auto-computes a power-of-2 oversample up to 8x, then selects the nearest generated frame for each target timestamp. 0 = use multiplier | | Output | Description | |--------|-------------| @@ -202,14 +209,65 @@ Same pattern as other Segment nodes. +

+SPEED + +#### Load SPEED Model + +Downloads the official `speed.pt` checkpoint from [zhZ524/SPEED](https://huggingface.co/zhZ524/SPEED) to `ComfyUI/models/speed-vfi/`. The loader also fetches a checksum-pinned snapshot of the official runtime on first use; Tween does not bundle that source. + +| Input | Description | +|-------|-------------| +| **model_path** | Checkpoint from `models/speed-vfi/` (official default is ~447 MB) | +| **precision** | `auto` prefers BF16, then FP16; FP32 is available for comparison | + +#### SPEED Interpolate / Segment Interpolate + +Uses the same batching, chunking, segment, and exact-target-FPS controls as BIM-VFI, plus a `seed` input for repeatable starting pixel noise. Keeping the seed on the interpolation node lets it change without reloading the model. SPEED is repeatable for the same seed and execution settings; changing batch, chunk, or segment boundaries can change how its stochastic noise is assigned. The released model predicts only the midpoint, so 4x and 8x are recursive passes. Inputs are padded to the model's 64-pixel divisor and cropped back automatically. + +
+ +
+LDF-VFI + +#### Load LDF-VFI Model + +Downloads the official transformer and conditional VAE from [onecat-ai/LDF-VFI](https://huggingface.co/onecat-ai/LDF-VFI) to `ComfyUI/models/ldf-vfi/` (~6.4 GB total). A checksum-pinned Apache-2.0 runtime snapshot is fetched on first use. Loading stays on CPU until the interpolation node executes. + +| Input | Description | +|-------|-------------| +| **tile_size / tile_overlap** | Spatial VAE tiling and seam blending; default 256/64 | +| **vae_batch_size** | Lower first if VAE encode/decode runs out of VRAM | +| **attention_type** | Official `slide_chunk_all_block_2x1x1` sparse attention is recommended | + +#### LDF-VFI Sequence Interpolate + +LDF-VFI is not a pairwise node. It processes the ordered source batch with the paper's skip-concat autoregressive sampler and internally chunks long sequences without breaking temporal context. + +| Input | Description | +|-------|-------------| +| **temporal_factor** | Any integer from 2x through 16x | +| **sampling_steps** | Diffusion steps per temporal block; official quick start uses 16 | +| **t_shift / t_cond** | Official defaults are 8.0 / 0.1 | +| **seed** | Repeatable VAE and diffusion sampling | +| **offload_after** | Return transformer and VAE to CPU after generation | +| **source_fps / target_fps** | Optional exact-FPS selection using the smallest sufficient native factor | + +The second output, `generated_sequence`, is the full native-factor sequence before exact-FPS selection. LDF has no Segment node because externally splitting the sequence would discard the long-range context it is designed to preserve. + +
+ ### Tween Concat Videos -Concatenates segment video files into a single video using ffmpeg. Connect from any Segment Interpolate's model output to ensure it runs after all segments are saved. Works with all four models. +Concatenates segment video files into a single video using ffmpeg. Connect from any pairwise Segment Interpolate's model output to ensure it runs after all segments are saved. ### Output frame count -- **Multiplier mode:** 2x = 2N-1, 4x = 4N-3, 8x = 8N-7 -- **Target FPS mode:** `floor((N-1) / source_fps * target_fps) + 1` frames. Automatically oversamples to the nearest power-of-2 above the ratio, then selects frames at exact target timestamps. Downsampling (target < source) also works — frames are selected from the input with no model calls. +- **Pairwise multiplier mode:** 2x = 2N-1, 4x = 4N-3, 8x = 8N-7 +- **LDF-VFI native factor:** factor `F` = `F(N-1)+1`, for any integer `F` from 2 through 16 +- **Target FPS mode:** `floor((N-1) / source_fps * target_fps) + 1` frames. Pairwise nodes oversample to the nearest power-of-2 above the ratio (up to 8x), then select the nearest generated frame for each target timestamp. Downsampling (target < source) also works — frames are selected from the input with no model calls. LDF-VFI supports native factors up to 16x. + +In target-FPS Segment mode, a very small `segment_size` can cover less than one output-frame interval while downsampling. Increase `segment_size` if the node reports that the segment contains no target timestamps; returning a placeholder frame would make concatenated timing incorrect. ## Acknowledgments @@ -219,8 +277,10 @@ Concatenates segment video files into a single video using ffmpeg. Connect from | **EMA-VFI** | Zhang et al. (MCG-NJU) | CVPR 2023 | [Paper](https://arxiv.org/abs/2303.00440) · [Code](https://github.com/MCG-NJU/EMA-VFI) | | **SGM-VFI** | Zhang et al. (MCG-NJU) | CVPR 2024 | [Paper](https://arxiv.org/abs/2404.06913) · [Code](https://github.com/MCG-NJU/SGM-VFI) | | **GIMM-VFI** | Guo, Li, Loy (S-Lab NTU) | NeurIPS 2024 | [Paper](https://arxiv.org/abs/2407.08680) · [Code](https://github.com/GSeanCDAT/GIMM-VFI) | +| **SPEED** | Zhang et al. | ACM MM 2026 | [Paper](https://arxiv.org/abs/2607.15585) · [Code](https://github.com/bbldCVer/SPEED) · [Model](https://huggingface.co/zhZ524/SPEED) | +| **LDF-VFI** | Peng et al. | CVPR 2026 | [Paper](https://arxiv.org/abs/2601.14959) · [Code](https://github.com/xypeng9903/LDF-VFI) · [Model](https://huggingface.co/onecat-ai/LDF-VFI) | -GIMM-VFI adaptation from [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/ComfyUI-GIMM-VFI) with checkpoints from [Kijai/GIMM-VFI_safetensors](https://huggingface.co/Kijai/GIMM-VFI_safetensors). Architecture files in `bim_vfi_arch/`, `ema_vfi_arch/`, `sgm_vfi_arch/`, and `gimm_vfi_arch/` are vendored from their respective repositories with minimal modifications. +GIMM-VFI adaptation from [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/ComfyUI-GIMM-VFI) with checkpoints from [Kijai/GIMM-VFI_safetensors](https://huggingface.co/Kijai/GIMM-VFI_safetensors). Architecture files in `bim_vfi_arch/`, `ema_vfi_arch/`, `sgm_vfi_arch/`, and `gimm_vfi_arch/` are vendored from their respective repositories with minimal modifications. SPEED and LDF-VFI use checksum-pinned official source snapshots downloaded into their model directories on demand.
BibTeX citations @@ -253,6 +313,22 @@ GIMM-VFI adaptation from [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/Comfy booktitle={Advances in Neural Information Processing Systems (NeurIPS)}, year={2024} } + +@misc{zhang2026speed, + title={SPEED: One-Step Pixel Diffusion for High-quality Video Frame Interpolation}, + author={Zhang, Zihao and Zhao, Haoyu and Yang, Siqian and Wu, Yidi and Jiang, Yudong and Wu, Zuxuan}, + year={2026}, + eprint={2607.15585}, + archivePrefix={arXiv} +} + +@misc{peng2026holistic, + title={Towards Holistic Modeling for Video Frame Interpolation with Auto-regressive Diffusion Transformers}, + author={Peng, Xinyu and Li, Han and Huang, Yuyang and Zheng, Ziyang and Wang, Yaoming and Chen, Xin and Dai, Wenrui and Li, Chenglin and Zou, Junni and Xiong, Hongkai}, + year={2026}, + eprint={2601.14959}, + archivePrefix={arXiv} +} ```
@@ -261,6 +337,8 @@ GIMM-VFI adaptation from [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/Comfy **BIM-VFI:** Research and education only. Commercial use requires permission from Prof. Munchurl Kim (mkimee@kaist.ac.kr). See the [original repository](https://github.com/KAIST-VICLab/BiM-VFI). -**EMA-VFI, SGM-VFI, GIMM-VFI:** [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0). GIMM-VFI ComfyUI adaptation based on [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/ComfyUI-GIMM-VFI). +**EMA-VFI, SGM-VFI, GIMM-VFI, LDF-VFI:** [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0). GIMM-VFI ComfyUI adaptation based on [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/ComfyUI-GIMM-VFI). + +**SPEED:** The official source repository did not include a license file when this integration was pinned. Tween does not redistribute that source; the loader downloads it directly from the official repository. Review the upstream terms before redistribution or commercial use. The checkpoint is likewise downloaded from its official Hugging Face repository. **This wrapper code:** [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0) diff --git a/__init__.py b/__init__.py index d150021..7667d86 100644 --- a/__init__.py +++ b/__init__.py @@ -3,6 +3,8 @@ from .nodes import ( LoadEMAVFIModel, EMAVFIInterpolate, EMAVFISegmentInterpolate, LoadSGMVFIModel, SGMVFIInterpolate, SGMVFISegmentInterpolate, LoadGIMMVFIModel, GIMMVFIInterpolate, GIMMVFISegmentInterpolate, + LoadSPEEDVFIModel, SPEEDVFIInterpolate, SPEEDVFISegmentInterpolate, + LoadLDFVFIModel, LDFVFIInterpolate, VFIOptimizer, ) @@ -20,6 +22,11 @@ NODE_CLASS_MAPPINGS = { "LoadGIMMVFIModel": LoadGIMMVFIModel, "GIMMVFIInterpolate": GIMMVFIInterpolate, "GIMMVFISegmentInterpolate": GIMMVFISegmentInterpolate, + "LoadSPEEDVFIModel": LoadSPEEDVFIModel, + "SPEEDVFIInterpolate": SPEEDVFIInterpolate, + "SPEEDVFISegmentInterpolate": SPEEDVFISegmentInterpolate, + "LoadLDFVFIModel": LoadLDFVFIModel, + "LDFVFIInterpolate": LDFVFIInterpolate, "VFIOptimizer": VFIOptimizer, } @@ -37,5 +44,10 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LoadGIMMVFIModel": "Load GIMM-VFI Model", "GIMMVFIInterpolate": "GIMM-VFI Interpolate", "GIMMVFISegmentInterpolate": "GIMM-VFI Segment Interpolate", + "LoadSPEEDVFIModel": "Load SPEED Model", + "SPEEDVFIInterpolate": "SPEED Interpolate", + "SPEEDVFISegmentInterpolate": "SPEED Segment Interpolate", + "LoadLDFVFIModel": "Load LDF-VFI Model", + "LDFVFIInterpolate": "LDF-VFI Sequence Interpolate", "VFIOptimizer": "VFI Optimizer", } diff --git a/assets/model-comparison.svg b/assets/model-comparison.svg index 9998f88..3a1dd08 100644 --- a/assets/model-comparison.svg +++ b/assets/model-comparison.svg @@ -1,4 +1,4 @@ - + @@ -12,13 +12,13 @@ - + BIM-VFI - ★ Recommended · Best quality · CVPR 2025 + Strong pairwise quality · CVPR 2025 Quality @@ -81,4 +81,38 @@ Apache 2.0 + + + + + SPEED + Modern pairwise default · ACM MM 2026 + + Quality + + + Speed + + + VRAM + + + Check upstream terms + + + + + LDF-VFI + Long-sequence coherence · CVPR 2026 + + Quality + + + Speed + + + VRAM + + + Apache 2.0 diff --git a/example_workflows/tween_speed_ldf_model_lab.json b/example_workflows/tween_speed_ldf_model_lab.json new file mode 100644 index 0000000..2b501f4 --- /dev/null +++ b/example_workflows/tween_speed_ldf_model_lab.json @@ -0,0 +1,757 @@ +{ + "id": "d1b1b840-08a9-4f93-9a1a-5bb9bdb05c39", + "revision": 0, + "last_node_id": 13, + "last_link_id": 11, + "nodes": [ + { + "id": 1, + "type": "Note", + "pos": [ + 0, + 0 + ], + "size": [ + 1420, + 112 + ], + "flags": {}, + "order": 0, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "Node name for S&R": "Note" + }, + "widgets_values": [ + "TWEEN MODEL LAB — SPEED + LDF-VFI\n1) Choose a short input video below. 2) Queue Prompt to run the enabled SPEED path. 3) The LDF-VFI path is muted by default; enable its three coral nodes only when you want the heavier sequence model." + ], + "color": "#1f3a4d", + "bgcolor": "#294d66" + }, + { + "id": 2, + "type": "VHS_LoadVideoPath", + "pos": [ + 20, + 190 + ], + "size": [ + 300, + 310 + ], + "flags": {}, + "order": 1, + "mode": 0, + "inputs": [ + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 1, + 8 + ] + }, + { + "name": "frame_count", + "type": "INT", + "links": null + }, + { + "name": "audio", + "type": "AUDIO", + "links": [ + 7, + 11 + ] + }, + { + "name": "video_info", + "type": "VHS_VIDEOINFO", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "Node name for S&R": "VHS_LoadVideoPath" + }, + "widgets_values": { + "video": "", + "force_rate": 24, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 25, + "skip_first_frames": 0, + "select_every_nth": 1, + "format": "AnimateDiff", + "videopreview": { + "hidden": false, + "paused": false, + "params": { + "filename": "", + "type": "path", + "format": "video/", + "force_rate": 24, + "custom_width": 0, + "custom_height": 0, + "frame_load_cap": 25, + "skip_first_frames": 0, + "select_every_nth": 1 + } + } + } + }, + { + "id": 3, + "type": "Note", + "pos": [ + 20, + 520 + ], + "size": [ + 300, + 166 + ], + "flags": {}, + "order": 2, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "Node name for S&R": "Note" + }, + "widgets_values": [ + "INPUT SAFETY DEFAULTS\n• force_rate = 24 FPS\n• frame_load_cap = 25 (~1 second)\n• dimensions = source size\n\nFor another source rate, change force_rate here and source_fps in the chosen Tween node. The Video Combine frame rate must match target_fps." + ], + "color": "#3f3520", + "bgcolor": "#594b2d" + }, + { + "id": 4, + "type": "LoadSPEEDVFIModel", + "pos": [ + 390, + 170 + ], + "size": [ + 300, + 105 + ], + "flags": {}, + "order": 3, + "mode": 0, + "inputs": [], + "outputs": [ + { + "name": "model", + "type": "SPEED_VFI_MODEL", + "links": [ + 2, + 4 + ] + } + ], + "properties": { + "aux_id": "ComfyUI-Tween.git", + "Node name for S&R": "LoadSPEEDVFIModel" + }, + "widgets_values": [ + "speed.pt", + "auto" + ], + "color": "#243b32", + "bgcolor": "#315244" + }, + { + "id": 5, + "type": "VFIOptimizer", + "pos": [ + 390, + 330 + ], + "size": [ + 310, + 150 + ], + "flags": {}, + "order": 4, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 1 + }, + { + "name": "model", + "type": "*", + "link": 2 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 3 + ] + }, + { + "name": "settings", + "type": "VFI_SETTINGS", + "links": [ + 5 + ] + } + ], + "properties": { + "aux_id": "ComfyUI-Tween.git", + "Node name for S&R": "VFIOptimizer" + }, + "widgets_values": [ + 2, + 0 + ], + "color": "#243b32", + "bgcolor": "#315244" + }, + { + "id": 6, + "type": "SPEEDVFIInterpolate", + "pos": [ + 760, + 240 + ], + "size": [ + 355, + 340 + ], + "flags": {}, + "order": 5, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 3 + }, + { + "name": "model", + "type": "SPEED_VFI_MODEL", + "link": 4 + }, + { + "name": "settings", + "shape": 7, + "type": "VFI_SETTINGS", + "link": 5 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 6 + ] + }, + { + "name": "oversampled", + "type": "IMAGE", + "links": null + } + ], + "properties": { + "aux_id": "ComfyUI-Tween.git", + "Node name for S&R": "SPEEDVFIInterpolate" + }, + "widgets_values": [ + 2, + 10, + true, + false, + 1, + 0, + 24, + 48, + 0 + ], + "color": "#243b32", + "bgcolor": "#315244" + }, + { + "id": 7, + "type": "VHS_VideoCombine", + "pos": [ + 1170, + 240 + ], + "size": [ + 360, + 340 + ], + "flags": {}, + "order": 6, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 6 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": 7 + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 48, + "loop_count": 0, + "filename_prefix": "Tween/demo_speed_24_to_48", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": {} + } + }, + "color": "#243b32", + "bgcolor": "#315244" + }, + { + "id": 8, + "type": "Note", + "pos": [ + 390, + 20 + ], + "size": [ + 1140, + 110 + ], + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "Node name for S&R": "Note" + }, + "widgets_values": [ + "SPEED — RECOMMENDED, ENABLED\nOne-step diffusion midpoint interpolation. This branch uses target-FPS mode (24→48), a fixed seed, and VFI Optimizer settings. First use downloads the official ~447 MB checkpoint and checksum-pinned runtime." + ], + "color": "#234936", + "bgcolor": "#2d6046" + }, + { + "id": 9, + "type": "Note", + "pos": [ + 390, + 760 + ], + "size": [ + 1140, + 125 + ], + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "Node name for S&R": "Note" + }, + "widgets_values": [ + "LDF-VFI — ADVANCED, MUTED BY DEFAULT\nRequires an NVIDIA BF16 GPU and roughly 20 GB VRAM; first use downloads ~6.4 GB. To run it, select the three coral nodes below and set Mode → Always. Disable/mute the SPEED output node if you only want LDF. Use 16 sampling steps for the official quality-oriented quick-start setting." + ], + "color": "#5a2d2d", + "bgcolor": "#743b3b" + }, + { + "id": 10, + "type": "LoadLDFVFIModel", + "pos": [ + 390, + 930 + ], + "size": [ + 340, + 190 + ], + "flags": {}, + "order": 9, + "mode": 2, + "inputs": [], + "outputs": [ + { + "name": "model", + "type": "LDF_VFI_MODEL", + "links": [ + 9 + ] + } + ], + "properties": { + "aux_id": "ComfyUI-Tween.git", + "Node name for S&R": "LoadLDFVFIModel" + }, + "widgets_values": [ + "onecat-ai/LDF-VFI", + 256, + 64, + 4, + "slide_chunk_all_block_2x1x1" + ], + "color": "#5a2d2d", + "bgcolor": "#743b3b" + }, + { + "id": 11, + "type": "LDFVFIInterpolate", + "pos": [ + 790, + 930 + ], + "size": [ + 370, + 320 + ], + "flags": {}, + "order": 10, + "mode": 2, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 8 + }, + { + "name": "model", + "type": "LDF_VFI_MODEL", + "link": 9 + } + ], + "outputs": [ + { + "name": "images", + "type": "IMAGE", + "links": [ + 10 + ] + }, + { + "name": "generated_sequence", + "type": "IMAGE", + "links": null + } + ], + "properties": { + "aux_id": "ComfyUI-Tween.git", + "Node name for S&R": "LDFVFIInterpolate" + }, + "widgets_values": [ + 2, + 8, + 8, + 0.1, + 42, + true, + 24, + 48 + ], + "color": "#5a2d2d", + "bgcolor": "#743b3b" + }, + { + "id": 12, + "type": "VHS_VideoCombine", + "pos": [ + 1210, + 930 + ], + "size": [ + 360, + 340 + ], + "flags": {}, + "order": 11, + "mode": 2, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 10 + }, + { + "name": "audio", + "shape": 7, + "type": "AUDIO", + "link": 11 + }, + { + "name": "meta_batch", + "shape": 7, + "type": "VHS_BatchManager", + "link": null + }, + { + "name": "vae", + "shape": 7, + "type": "VAE", + "link": null + } + ], + "outputs": [ + { + "name": "Filenames", + "type": "VHS_FILENAMES", + "links": null + } + ], + "properties": { + "cnr_id": "comfyui-videohelpersuite", + "Node name for S&R": "VHS_VideoCombine" + }, + "widgets_values": { + "frame_rate": 48, + "loop_count": 0, + "filename_prefix": "Tween/demo_ldf_24_to_48", + "format": "video/h264-mp4", + "pix_fmt": "yuv420p", + "crf": 19, + "save_metadata": true, + "trim_to_audio": false, + "pingpong": false, + "save_output": true, + "videopreview": { + "hidden": false, + "paused": false, + "params": {} + } + }, + "color": "#5a2d2d", + "bgcolor": "#743b3b" + }, + { + "id": 13, + "type": "Note", + "pos": [ + 20, + 930 + ], + "size": [ + 300, + 236 + ], + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [], + "outputs": [], + "properties": { + "Node name for S&R": "Note" + }, + "widgets_values": [ + "WHY TWO PATHS?\n\nSPEED\n• Best default for normal pairwise VFI\n• Much smaller and faster\n• 2x midpoint; 4x/8x recursive\n\nLDF-VFI\n• Sequence-native temporal modeling\n• Better suited to long-range coherence\n• 2x–16x, but far heavier\n\nThe LDF demo uses 8 steps for a quicker test. Raise to 16 after the graph is working." + ], + "color": "#40303b", + "bgcolor": "#594353" + } + ], + "links": [ + [ + 1, + 2, + 0, + 5, + 0, + "IMAGE" + ], + [ + 2, + 4, + 0, + 5, + 1, + "*" + ], + [ + 3, + 5, + 0, + 6, + 0, + "IMAGE" + ], + [ + 4, + 4, + 0, + 6, + 1, + "SPEED_VFI_MODEL" + ], + [ + 5, + 5, + 1, + 6, + 2, + "VFI_SETTINGS" + ], + [ + 6, + 6, + 0, + 7, + 0, + "IMAGE" + ], + [ + 7, + 2, + 2, + 7, + 1, + "AUDIO" + ], + [ + 8, + 2, + 0, + 11, + 0, + "IMAGE" + ], + [ + 9, + 10, + 0, + 11, + 1, + "LDF_VFI_MODEL" + ], + [ + 10, + 11, + 0, + 12, + 0, + "IMAGE" + ], + [ + 11, + 2, + 2, + 12, + 1, + "AUDIO" + ] + ], + "groups": [ + { + "id": 1, + "title": "1 · INPUT CLIP", + "bounding": [ + -20, + 135, + 360, + 575 + ], + "color": "#8a6d2f", + "font_size": 24, + "flags": {} + }, + { + "id": 2, + "title": "2 · SPEED — RECOMMENDED / ENABLED", + "bounding": [ + 350, + -20, + 1240, + 650 + ], + "color": "#347a56", + "font_size": 24, + "flags": {} + }, + { + "id": 3, + "title": "3 · LDF-VFI — ADVANCED / MUTED", + "bounding": [ + 350, + 730, + 1240, + 590 + ], + "color": "#8a4646", + "font_size": 24, + "flags": {} + } + ], + "config": {}, + "extra": { + "workflowRendererVersion": "LG", + "ue_links": [], + "links_added_by_ue": [], + "ds": { + "scale": 0.76, + "offset": [ + 80, + 80 + ] + }, + "frontendVersion": "1.45.19", + "VHS_latentpreview": true, + "VHS_latentpreviewrate": 0, + "VHS_MetadataImage": true, + "VHS_KeepIntermediate": true + }, + "version": 0.4 +} diff --git a/external_sources.py b/external_sources.py new file mode 100644 index 0000000..06d7b72 --- /dev/null +++ b/external_sources.py @@ -0,0 +1,173 @@ +"""Pinned, lazy installers for optional upstream model runtimes. + +Tween does not redistribute these projects. Their official source archives are +downloaded only when a corresponding loader node is executed, verified against +a pinned SHA-256 digest, and kept next to that model's checkpoints. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +import os +from pathlib import Path +import shutil +import tempfile +import threading +import urllib.error +import urllib.request +import zipfile + + +logger = logging.getLogger("Tween") + + +UPSTREAM_SOURCES = { + "speed": { + "project": "SPEED", + "commit": "40fadbe85c88cc6e4015062389da464fd7e85ab9", + "url": ( + "https://codeload.github.com/bbldCVer/SPEED/zip/" + "40fadbe85c88cc6e4015062389da464fd7e85ab9" + ), + "sha256": "9e9cc71bfeaf7a62008950b8f234f5f035df27b65a5fc0464caee2542f47f68c", + "required": "src/models/model.py", + }, + "ldf": { + "project": "LDF-VFI", + "commit": "61b34d2379df8a313e8e4cb467cc2f74c52b45d7", + "url": ( + "https://codeload.github.com/xypeng9903/LDF-VFI/zip/" + "61b34d2379df8a313e8e4cb467cc2f74c52b45d7" + ), + "sha256": "3a903aeb5353c7e5eb932f129d975d1750246502937d8f7b283b61a269c23668", + "required": "training/models/precond.py", + }, +} +_SOURCE_LOCKS = {name: threading.Lock() for name in UPSTREAM_SOURCES} + + +def _download(url: str, destination: Path) -> None: + request = urllib.request.Request(url, headers={"User-Agent": "ComfyUI-Tween"}) + try: + with urllib.request.urlopen(request, timeout=60) as response, destination.open("wb") as output: + shutil.copyfileobj(response, output, length=1024 * 1024) + except (OSError, urllib.error.URLError) as exc: + raise RuntimeError( + f"Could not download optional upstream runtime from {url}. " + "Check network access and retry the loader node." + ) from exc + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _safe_extract(archive: Path, destination: Path) -> Path: + with zipfile.ZipFile(archive) as source_zip: + members = source_zip.infolist() + if not members: + raise RuntimeError(f"Downloaded source archive is empty: {archive}") + + destination_resolved = destination.resolve() + for member in members: + member_path = (destination / member.filename).resolve() + if os.path.commonpath((destination_resolved, member_path)) != str(destination_resolved): + raise RuntimeError(f"Unsafe path in source archive: {member.filename}") + source_zip.extractall(destination) + + top_level = {Path(member.filename).parts[0] for member in members if member.filename} + if len(top_level) != 1: + raise RuntimeError("Expected one top-level directory in the upstream source archive") + return destination / top_level.pop() + + +def _ensure_upstream_source_unlocked(name: str, model_dir: str | os.PathLike[str]) -> str: + try: + spec = UPSTREAM_SOURCES[name] + except KeyError as exc: + raise ValueError(f"Unknown Tween upstream source: {name}") from exc + + model_root = Path(model_dir) + source_dir = model_root / "_upstream" + required_file = source_dir / spec["required"] + if required_file.is_file(): + marker_path = source_dir / ".tween-source.json" + if marker_path.is_file(): + try: + marker = json.loads(marker_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as exc: + raise RuntimeError(f"Invalid source marker: {marker_path}") from exc + if ( + marker.get("commit") != spec["commit"] + or marker.get("archive_sha256") != spec["sha256"] + ): + raise RuntimeError( + f"{spec['project']} runtime at {source_dir} is pinned to a different commit. " + "Remove _upstream and run the loader again." + ) + else: + logger.warning( + "Using manually installed %s runtime at %s (no Tween verification marker)", + spec["project"], source_dir, + ) + return str(source_dir) + + if source_dir.exists(): + raise RuntimeError( + f"Incomplete {spec['project']} runtime at {source_dir}. " + "Remove that _upstream directory and run the loader again." + ) + + model_root.mkdir(parents=True, exist_ok=True) + logger.info( + "Downloading pinned %s runtime (%s) to %s", + spec["project"], spec["commit"][:12], source_dir, + ) + + with tempfile.TemporaryDirectory(prefix="tween-source-", dir=model_root) as temp_name: + temp_dir = Path(temp_name) + archive = temp_dir / "source.zip" + _download(spec["url"], archive) + + actual_digest = _sha256(archive) + if actual_digest != spec["sha256"]: + raise RuntimeError( + f"Checksum mismatch for {spec['project']} source archive: " + f"expected {spec['sha256']}, got {actual_digest}" + ) + + extracted = _safe_extract(archive, temp_dir / "extract") + if not (extracted / spec["required"]).is_file(): + raise RuntimeError( + f"The {spec['project']} archive does not contain {spec['required']}" + ) + + marker = { + "project": spec["project"], + "commit": spec["commit"], + "archive_sha256": spec["sha256"], + "source_url": spec["url"], + } + (extracted / ".tween-source.json").write_text( + json.dumps(marker, indent=2) + "\n", encoding="utf-8" + ) + shutil.move(str(extracted), str(source_dir)) + + logger.info("Installed %s runtime at %s", spec["project"], source_dir) + return str(source_dir) + + +def ensure_upstream_source(name: str, model_dir: str | os.PathLike[str]) -> str: + """Return a verified upstream checkout, downloading it once per process.""" + try: + source_lock = _SOURCE_LOCKS[name] + except KeyError as exc: + raise ValueError(f"Unknown Tween upstream source: {name}") from exc + with source_lock: + return _ensure_upstream_source_unlocked(name, model_dir) diff --git a/ldf_backend.py b/ldf_backend.py new file mode 100644 index 0000000..51c2fe7 --- /dev/null +++ b/ldf_backend.py @@ -0,0 +1,437 @@ +"""Sequence-native ComfyUI adapter for the official Apache-2.0 LDF-VFI runtime.""" + +from __future__ import annotations + +import importlib +import logging +from pathlib import Path +import sys +import threading +import types + +from einops import rearrange, repeat +import torch +import torch.nn.functional as F + + +logger = logging.getLogger("Tween") +_LDF_NAMESPACE = "_tween_ldf_upstream" +_LDF_IMPORT_LOCK = threading.RLock() + + +def _cuda_bf16_supported(device: torch.device) -> bool: + if device.type != "cuda": + return False + with torch.cuda.device(device): + return torch.cuda.is_bf16_supported() + + +def _namespace_package(name: str, path: Path): + package = sys.modules.get(name) + if package is not None: + return package + package = types.ModuleType(name) + package.__path__ = [str(path)] + package.__package__ = name + sys.modules[name] = package + return package + + +def load_ldf_runtime(source_root: str): + """Load LDF under an isolated namespace without polluting ``training``.""" + root = Path(source_root).resolve() + required = root / "training" / "models" / "precond.py" + if not required.is_file(): + raise RuntimeError(f"Invalid LDF-VFI source directory: missing {required}") + + with _LDF_IMPORT_LOCK: + cached = sys.modules.get(f"{_LDF_NAMESPACE}.models.precond") + if cached is not None: + transformer = importlib.import_module(f"{_LDF_NAMESPACE}.models.transformer_wan") + return { + "Precond": cached.Precond, + "ConditionalVAE": cached.Wan2_1SpatialTiledConditionEncoder3Dv2, + "MaskEncoder": cached.MaskSpatialTiledEncoder3D, + "Transformer": transformer.WanTransformer3DModel, + } + + training_root = _namespace_package(_LDF_NAMESPACE, root / "training") + saved_training_modules = { + name: module for name, module in tuple(sys.modules.items()) + if name == "training" or name.startswith("training.") + } + for name in saved_training_modules: + sys.modules.pop(name, None) + # One upstream transformer import is absolute (training.distributed.util). + # Temporarily alias only while importing, then restore the host process. + sys.modules["training"] = training_root + try: + precond = importlib.import_module(f"{_LDF_NAMESPACE}.models.precond") + transformer = importlib.import_module(f"{_LDF_NAMESPACE}.models.transformer_wan") + except (ImportError, AttributeError) as exc: + for name in tuple(sys.modules): + if name == _LDF_NAMESPACE or name.startswith(f"{_LDF_NAMESPACE}."): + sys.modules.pop(name, None) + raise RuntimeError( + "LDF-VFI requires PyTorch 2.5+ and diffusers 0.33+. " + "Install Tween's current requirements and restart ComfyUI." + ) from exc + finally: + for name in tuple(sys.modules): + if name == "training" or name.startswith("training."): + sys.modules.pop(name, None) + sys.modules.update(saved_training_modules) + + return { + "Precond": precond.Precond, + "ConditionalVAE": precond.Wan2_1SpatialTiledConditionEncoder3Dv2, + "MaskEncoder": precond.MaskSpatialTiledEncoder3D, + "Transformer": transformer.WanTransformer3DModel, + } + + +def _upsample_nearest(frames: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Expand sparse source frames to every temporal position in a window.""" + kept_indices = torch.where(mask)[0] + if kept_indices.numel() == 0: + raise RuntimeError("LDF-VFI received a temporal window with no source frames") + positions = torch.arange(mask.shape[0], device=mask.device) + nearest = (positions[:, None] - kept_indices[None, :]).abs().argmin(dim=1) + return frames[nearest] + + +class LDFVFIModel: + """Long-sequence diffusion interpolation using LDF's skip-concat sampler.""" + + TRAIN_FRAMES = 60 + TILE_TIME = 20 + CONDITION_TILES = 1 + + def __init__( + self, + model_root: str, + vae_path: str, + source_root: str, + tile_size: int = 256, + tile_overlap: int = 64, + vae_batch_size: int = 8, + attention_type: str = "slide_chunk_all_block_2x1x1", + ): + if tile_size % 8 or tile_overlap % 8: + raise ValueError("LDF-VFI tile size and overlap must be divisible by 8") + if tile_overlap >= tile_size: + raise ValueError("LDF-VFI tile overlap must be smaller than tile size") + + runtime = load_ldf_runtime(source_root) + self.device = "cpu" + self.dtype = torch.bfloat16 + self.vae_path = vae_path + self._ConditionalVAE = runtime["ConditionalVAE"] + self._MaskEncoder = runtime["MaskEncoder"] + self._Precond = runtime["Precond"] + + logger.info("Loading LDF-VFI transformer from %s", model_root) + try: + transformer = runtime["Transformer"].from_pretrained( + model_root, + subfolder="transformer", + torch_dtype=self.dtype, + low_cpu_mem_usage=True, + ) + except TypeError: + transformer = runtime["Transformer"].from_pretrained( + model_root, subfolder="transformer", torch_dtype=self.dtype + ) + transformer.set_attention_type(attention_type) + transformer.requires_grad_(False).eval() + + stride = tile_size - tile_overlap + tiled_kwargs = { + "tile_sample_min_height": tile_size, + "tile_sample_min_width": tile_size, + "tile_sample_min_time": self.TILE_TIME, + "tile_sample_stride_height": stride, + "tile_sample_stride_width": stride, + "spatial_compression_ratio": 8, + "temporal_compression_ratio": 4, + } + # The same conditional VAE can encode conditions and decode predictions; + # sharing it avoids loading a second ~800 MB copy as the reference CLI does. + self.vae = self._ConditionalVAE(vae_path, vae_batch_size, **tiled_kwargs) + self.mask_encoder = self._MaskEncoder(**tiled_kwargs) + self.model = self._Precond( + transformer=transformer, + vae=self.vae, + lq_encoder=self.vae, + msk_encoder=self.mask_encoder, + ) + self.model.requires_grad_(False).eval() + + @property + def transformer(self): + return self.model.transformer + + def _move_auxiliary_models(self, device: torch.device) -> None: + self.mask_encoder.mask_encoder.to(device=device, dtype=self.dtype) + if self.vae.vae is None: + if device.type != "cpu": + self.vae.init(device) + return + self.vae.vae.to(device=device, dtype=self.dtype) + if hasattr(self.vae, "mean") and hasattr(self.vae, "std"): + self.vae.mean = self.vae.mean.to(device=device, dtype=self.dtype) + self.vae.std = self.vae.std.to(device=device, dtype=self.dtype) + self.vae.scale = [self.vae.mean, 1.0 / self.vae.std] + + def to(self, device): + target = torch.device(device) + if target.type == "cuda" and not _cuda_bf16_supported(target): + raise RuntimeError("LDF-VFI requires a CUDA GPU with BF16 support (Ampere or newer)") + self.transformer.to(device=target, dtype=self.dtype) + self._move_auxiliary_models(target) + self.device = str(target) + return self + + def clear_cache(self) -> None: + for name in ( + "_swin_attention_mask", + "_sliding_chunk_attention_mask", + "_sliding_window_attention_mask", + ): + method = getattr(self.transformer, name, None) + cache_clear = getattr(method, "cache_clear", None) + if cache_clear is not None: + cache_clear() + + def _prepare_condition(self, frames, mask, device): + if mask.shape[0] > self.TRAIN_FRAMES: + raise ValueError("Internal LDF temporal window exceeds the training window") + if mask.shape[0] < self.TRAIN_FRAMES: + mask = F.pad(mask, (0, self.TRAIN_FRAMES - mask.shape[0])) + dense = _upsample_nearest(frames, mask) + dense = rearrange(dense, "t c h w -> 1 c t h w") + dense = dense.to(device=device, dtype=self.dtype, non_blocking=True).mul(2).sub(1) + dense_mask = repeat( + mask, "t -> 1 1 t h w", h=dense.shape[-2], w=dense.shape[-1] + ).to(device=device, dtype=self.dtype) + condition = self.vae.encode(dense, for_train=True) + encoded_mask = self.mask_encoder.encode(dense_mask, for_train=True) + return dense, dense_mask, condition, encoded_mask + + def _time_schedule(self, num_steps: int, t_shift: float) -> torch.Tensor: + schedule = torch.linspace(1.0, 0.0, steps=num_steps + 1) + return t_shift * schedule / (1 + (t_shift - 1) * schedule) + + def _predict_step(self, latent, timestep, condition, encoded_mask): + return self.model.predict_v(latent, timestep, condition, encoded_mask) + + def _sample_free(self, condition, encoded_mask, schedule, device, progress): + latent = torch.randn_like(condition) + for index in range(schedule.shape[0] - 1): + progress() + timestep = torch.full( + condition.shape[:-4], float(schedule[index]), device=device, dtype=self.dtype + ) + velocity = self._predict_step(latent, timestep, condition, encoded_mask) + step_size = float(schedule[index + 1] - schedule[index]) + latent = latent + velocity * step_size + return latent + + def _sample_between(self, previous, following, condition, encoded_mask, + schedule, t_cond, device, progress): + previous_noisy = previous * (1 - t_cond) + torch.randn_like(previous) * t_cond + following_noisy = following * (1 - t_cond) + torch.randn_like(following) * t_cond + middle = torch.randn_like(condition[:, self.CONDITION_TILES:-self.CONDITION_TILES]) + previous_t = torch.full( + previous.shape[:-4], t_cond, device=device, dtype=self.dtype + ) + following_t = torch.full( + following.shape[:-4], t_cond, device=device, dtype=self.dtype + ) + for index in range(schedule.shape[0] - 1): + progress() + latent = torch.cat((previous_noisy, middle, following_noisy), dim=1) + middle_t = torch.full( + middle.shape[:-4], float(schedule[index]), device=device, dtype=self.dtype + ) + timestep = torch.cat((previous_t, middle_t, following_t), dim=1) + velocity = self._predict_step( + latent, timestep, condition, encoded_mask + )[:, self.CONDITION_TILES:-self.CONDITION_TILES] + step_size = float(schedule[index + 1] - schedule[index]) + middle = middle + velocity * step_size + return middle + + def _sample_tail(self, previous, condition, encoded_mask, schedule, + t_cond, device, progress): + previous_noisy = previous * (1 - t_cond) + torch.randn_like(previous) * t_cond + tail = torch.randn_like(condition[:, self.CONDITION_TILES:]) + previous_t = torch.full( + previous.shape[:-4], t_cond, device=device, dtype=self.dtype + ) + for index in range(schedule.shape[0] - 1): + progress() + latent = torch.cat((previous_noisy, tail), dim=1) + tail_t = torch.full( + tail.shape[:-4], float(schedule[index]), device=device, dtype=self.dtype + ) + timestep = torch.cat((previous_t, tail_t), dim=1) + velocity = self._predict_step( + latent, timestep, condition, encoded_mask + )[:, self.CONDITION_TILES:] + step_size = float(schedule[index + 1] - schedule[index]) + tail = tail + velocity * step_size + return tail + + def _decode(self, latent, dense, dense_mask, height, width): + latent = rearrange( + latent, "1 nt nh nw c t h w -> 1 nt c t (nh h) (nw w)" + ) + prediction = self.vae.decode(latent, dense, dense_mask)[..., :height, :width] + return rearrange(prediction, "1 c t h w -> t c h w").add(1).mul(0.5).clamp_(0, 1).float().cpu() + + @staticmethod + def sampling_block_count(num_input_frames: int, temporal_factor: int) -> int: + total_length = num_input_frames * temporal_factor + t0 = 40 + stride = 20 + blocks = 1 + while t0 + stride <= total_length - 1: + blocks += 2 + t0 += stride * 2 + if t0 < total_length: + blocks += 1 + return blocks + + @torch.no_grad() + def interpolate_sequence( + self, + frames: torch.Tensor, + temporal_factor: int, + num_steps: int = 16, + t_shift: float = 8.0, + t_cond: float = 0.1, + seed: int = 42, + progress_callback=None, + ) -> torch.Tensor: + if not 2 <= temporal_factor <= 16: + raise ValueError("LDF-VFI temporal factor must be between 2 and 16") + if num_steps < 1: + raise ValueError("LDF-VFI num_steps must be at least 1") + if t_shift <= 0: + raise ValueError("LDF-VFI t_shift must be greater than 0") + if not 0 <= t_cond <= 1: + raise ValueError("LDF-VFI t_cond must be between 0 and 1") + if frames.shape[0] < 2: + return frames + + device = next(self.transformer.parameters()).device + if device.type != "cuda": + raise RuntimeError("Move LDF-VFI to a CUDA device before interpolation") + source = frames.detach().float().cpu() + height, width = source.shape[-2:] + schedule = self._time_schedule(num_steps, t_shift) + progress_callback = progress_callback or (lambda: None) + + cuda_index = device.index + if cuda_index is None: + cuda_index = torch.cuda.current_device() + rng_context = torch.random.fork_rng(devices=[cuda_index]) + with rng_context: + torch.manual_seed(int(seed)) + with torch.cuda.device(device): + torch.cuda.manual_seed(int(seed)) + chunks = self._interpolate_sequence_impl( + source, temporal_factor, schedule, t_cond, device, progress_callback + ) + + result = torch.cat(chunks, dim=0) + expected = (frames.shape[0] - 1) * temporal_factor + 1 + return result[:expected] + + def _interpolate_sequence_impl(self, source, factor, schedule, t_cond, + device, progress): + total_mask = torch.zeros(source.shape[0] * factor, dtype=torch.bool) + total_mask[::factor] = True + n_tiles = self.TRAIN_FRAMES // self.TILE_TIME + output_tiles = n_tiles - self.CONDITION_TILES + outputs = [] + + # First chunk. + mask = total_mask[:self.TRAIN_FRAMES] + input_end = int(mask.sum()) + dense, dense_mask, condition, encoded_mask = self._prepare_condition( + source[:input_end], mask, device + ) + latent = self._sample_free(condition, encoded_mask, schedule, device, progress) + latent = latent[:, :output_tiles] + previous = latent[:, -self.CONDITION_TILES:] + decode_time = output_tiles * self.TILE_TIME + outputs.append(self._decode( + latent, dense[:, :, :decode_time], dense_mask[:, :, :decode_time], + source.shape[-2], source.shape[-1], + )) + + t0 = output_tiles * self.TILE_TIME + stride = (n_tiles - self.CONDITION_TILES * 2) * self.TILE_TIME + while t0 + stride <= total_mask.shape[0] - 1: + # A future/skip chunk establishes the far-side condition. + input_start = t0 + stride - self.CONDITION_TILES * self.TILE_TIME + input_end_t = t0 + stride * 2 + self.CONDITION_TILES * self.TILE_TIME + mask = total_mask[input_start:input_end_t] + source_start = int(total_mask[:input_start].sum()) + source_end = int(total_mask[:input_end_t].sum()) + dense, dense_mask, condition, encoded_mask = self._prepare_condition( + source[source_start:source_end], mask, device + ) + skip = self._sample_free(condition, encoded_mask, schedule, device, progress) + skip = skip[:, self.CONDITION_TILES:-self.CONDITION_TILES] + following = skip[:, :self.CONDITION_TILES] + previous_next = skip[:, -self.CONDITION_TILES:] + start_time = self.CONDITION_TILES * self.TILE_TIME + end_time = output_tiles * self.TILE_TIME + decoded_skip = self._decode( + skip, dense[:, :, start_time:end_time], dense_mask[:, :, start_time:end_time], + source.shape[-2], source.shape[-1], + ) + + # Fill the gap between the preceding and skip chunks. + input_start = t0 - self.CONDITION_TILES * self.TILE_TIME + input_end_t = t0 + stride + self.CONDITION_TILES * self.TILE_TIME + mask = total_mask[input_start:input_end_t] + source_start = int(total_mask[:input_start].sum()) + source_end = int(total_mask[:input_end_t].sum()) + dense, dense_mask, condition, encoded_mask = self._prepare_condition( + source[source_start:source_end], mask, device + ) + middle = self._sample_between( + previous, following, condition, encoded_mask, + schedule, t_cond, device, progress, + ) + decoded_middle = self._decode( + middle, dense[:, :, start_time:end_time], dense_mask[:, :, start_time:end_time], + source.shape[-2], source.shape[-1], + ) + outputs.extend((decoded_middle, decoded_skip)) + previous = previous_next + t0 += stride * 2 + + # Remaining tail. + if t0 < total_mask.shape[0]: + input_start = t0 - self.CONDITION_TILES * self.TILE_TIME + mask = total_mask[input_start:] + source_start = int(total_mask[:input_start].sum()) + dense, dense_mask, condition, encoded_mask = self._prepare_condition( + source[source_start:], mask, device + ) + tail = self._sample_tail( + previous, condition, encoded_mask, schedule, + t_cond, device, progress, + ) + start_time = self.CONDITION_TILES * self.TILE_TIME + outputs.append(self._decode( + tail, dense[:, :, start_time:], dense_mask[:, :, start_time:], + source.shape[-2], source.shape[-1], + )) + + return outputs diff --git a/nodes.py b/nodes.py index 51ca52c..3730c23 100644 --- a/nodes.py +++ b/nodes.py @@ -2,15 +2,24 @@ import math import os import glob import logging +import re import shutil import subprocess +import sys import tempfile import time import torch import folder_paths from comfy.utils import ProgressBar +try: + import comfy.model_management as model_management +except ImportError: # Allows lightweight tests outside a full ComfyUI install. + model_management = None from .inference import BiMVFIModel, EMAVFIModel, SGMVFIModel, GIMMVFIModel +from .speed_backend import SpeedVFIModel +from .ldf_backend import LDFVFIModel +from .external_sources import ensure_upstream_source from .bim_vfi_arch import clear_backwarp_cache from .ema_vfi_arch import clear_warp_cache as clear_ema_warp_cache from .sgm_vfi_arch import clear_warp_cache as clear_sgm_warp_cache @@ -19,6 +28,68 @@ from .gimm_vfi_arch import clear_gimm_caches logger = logging.getLogger("Tween") +def _get_torch_device(): + """Honor ComfyUI's selected device instead of assuming the default CUDA GPU.""" + if model_management is not None: + return torch.device(model_management.get_torch_device()) + return torch.device("cuda" if torch.cuda.is_available() else "cpu") + + +def _soft_empty_cache(): + if model_management is not None: + model_management.soft_empty_cache() + elif torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _throw_if_interrupted(): + if model_management is not None: + model_management.throw_exception_if_processing_interrupted() + + +def _target_fps_enabled(source_fps, target_fps): + if target_fps > 0 and source_fps <= 0: + raise ValueError("source_fps must be greater than 0 when target_fps is enabled") + return target_fps > 0 and source_fps > 0 + + +def _interpolate_batch_with_offload(model, frames0, frames1, device, keep_device): + """Run one pair batch and reliably return an optionally offloaded model to CPU.""" + if keep_device: + return model.interpolate_batch(frames0, frames1, time_step=0.5) + + try: + model.to(device) + return model.interpolate_batch(frames0, frames1, time_step=0.5) + finally: + inference_failed = sys.exc_info()[0] is not None + try: + model.to("cpu") + except Exception: + # Preserve an inference exception if one is already active. + if not inference_failed: + raise + logger.exception("Failed to offload VFI model after an inference error") + + +def _interpolate_multi_with_offload(model, frame0, frame1, num_intermediates, + device, keep_device): + if keep_device: + return model.interpolate_multi(frame0, frame1, num_intermediates) + + try: + model.to(device) + return model.interpolate_multi(frame0, frame1, num_intermediates) + finally: + inference_failed = sys.exc_info()[0] is not None + try: + model.to("cpu") + except Exception: + if not inference_failed: + raise + logger.exception("Failed to offload VFI model after an inference error") + + def _get_system_ram_gb(): @@ -62,8 +133,11 @@ def _clear_model_cache(model): clear_sgm_warp_cache() elif isinstance(model, GIMMVFIModel): clear_gimm_caches() - if torch.cuda.is_available(): - torch.cuda.empty_cache() + elif isinstance(model, SpeedVFIModel): + model.clear_cache() + elif isinstance(model, LDFVFIModel): + model.clear_cache() + _soft_empty_cache() def _compute_target_fps_params(source_fps, target_fps): @@ -77,6 +151,12 @@ def _compute_target_fps_params(source_fps, target_fps): return 0, 1 # no interpolation needed (downsampling or same fps) num_passes = math.ceil(math.log2(ratio)) mult = 2 ** num_passes + if mult > 8: + raise ValueError( + f"Pairwise Tween nodes support target-FPS interpolation up to 8x; " + f"{source_fps:g} -> {target_fps:g} FPS requires {mult}x. " + "Use LDF-VFI for ratios up to 16x or interpolate in multiple stages." + ) return num_passes, mult @@ -84,7 +164,7 @@ def _select_target_fps_frames(frames, source_fps, target_fps, mult, num_input): """Pick frames from oversampled [M,C,H,W] tensor to hit target FPS timing. For downsampling (mult=1, ratio<=1), selects from original input frames. - For upsampling, selects from the oversampled sequence at target timestamps. + For upsampling, selects the nearest oversampled frame for each target timestamp. """ duration = (num_input - 1) / source_fps num_output = int(math.floor(duration * target_fps)) + 1 @@ -130,6 +210,24 @@ GIMM_MODEL_DIR = os.path.join(folder_paths.models_dir, "gimm-vfi") if not os.path.exists(GIMM_MODEL_DIR): os.makedirs(GIMM_MODEL_DIR, exist_ok=True) +# SPEED +SPEED_HF_REPO = "zhZ524/SPEED" +SPEED_DEFAULT_MODEL = "speed.pt" +SPEED_MODEL_DIR = os.path.join(folder_paths.models_dir, "speed-vfi") +if not os.path.exists(SPEED_MODEL_DIR): + os.makedirs(SPEED_MODEL_DIR, exist_ok=True) + +# LDF-VFI +LDF_HF_REPO = "onecat-ai/LDF-VFI" +LDF_MODEL_FILES = ( + "transformer/config.json", + "transformer/diffusion_pytorch_model.safetensors", + "Wan2.1_VAE_cond_v2.pth", +) +LDF_MODEL_DIR = os.path.join(folder_paths.models_dir, "ldf-vfi") +if not os.path.exists(LDF_MODEL_DIR): + os.makedirs(LDF_MODEL_DIR, exist_ok=True) + def get_available_models(): """List available checkpoint files in the bim-vfi model directory.""" @@ -205,6 +303,8 @@ class LoadBIMVFIModel: class BIMVFIInterpolate: + MODEL_LABEL = "BIM-VFI" + @classmethod def INPUT_TYPES(cls): return { @@ -245,7 +345,7 @@ class BIMVFIInterpolate: }), "target_fps": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01, - "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes the optimal power-of-2 oversample then selects frames. 0 = use multiplier.", + "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes a power-of-2 oversample up to 8x, then selects frames. 0 = use multiplier.", }), }, "optional": { @@ -273,27 +373,24 @@ class BIMVFIInterpolate: Interpolated frames as [M, C, H, W] tensor on storage_device """ for pass_idx in range(num_passes): - logger.info(f"BIM-VFI: pass {pass_idx + 1}/{num_passes}, {frames.shape[0]} -> {2 * frames.shape[0] - 1} frames") + logger.info(f"{self.MODEL_LABEL}: pass {pass_idx + 1}/{num_passes}, {frames.shape[0]} -> {2 * frames.shape[0] - 1} frames") new_frames = [] num_pairs = frames.shape[0] - 1 pairs_since_clear = 0 for i in range(0, num_pairs, batch_size): + _throw_if_interrupted() batch_end = min(i + batch_size, num_pairs) actual_batch = batch_end - i frames0 = frames[i:batch_end] frames1 = frames[i + 1:batch_end + 1] - if not keep_device: - model.to(device) - - mids = model.interpolate_batch(frames0, frames1, time_step=0.5) + mids = _interpolate_batch_with_offload( + model, frames0, frames1, device, keep_device + ) mids = mids.to(storage_device) - if not keep_device: - model.to("cpu") - for j in range(actual_batch): new_frames.append(frames[i + j:i + j + 1]) new_frames.append(mids[j:j+1]) @@ -302,17 +399,14 @@ class BIMVFIInterpolate: pbar.update_absolute(step_ref[0]) pairs_since_clear += actual_batch - if pairs_since_clear >= clear_cache_after_n_frames and torch.cuda.is_available(): - clear_backwarp_cache() - torch.cuda.empty_cache() + if pairs_since_clear >= clear_cache_after_n_frames: + _clear_model_cache(model) pairs_since_clear = 0 new_frames.append(frames[-1:]) frames = torch.cat(new_frames, dim=0) - if torch.cuda.is_available(): - clear_backwarp_cache() - torch.cuda.empty_cache() + _clear_model_cache(model) return frames @@ -328,16 +422,21 @@ class BIMVFIInterpolate: def interpolate(self, images, model, multiplier, clear_cache_after_n_frames, keep_device, all_on_gpu, batch_size, chunk_size, - source_fps=0.0, target_fps=0.0, settings=None): + source_fps=0.0, target_fps=0.0, settings=None, seed=None): batch_size, chunk_size, keep_device, all_on_gpu, clear_cache_after_n_frames = \ _apply_vfi_settings(settings, batch_size, chunk_size, keep_device, all_on_gpu, clear_cache_after_n_frames) + if seed is not None and hasattr(model, "set_seed"): + model.set_seed(seed) + if hasattr(model, "reset_seed"): + model.reset_seed() + if images.shape[0] < 2: return (images, images) # Target FPS mode: auto-compute multiplier from fps ratio - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) if use_target_fps: num_passes, mult = _compute_target_fps_params(source_fps, target_fps) if num_passes == 0: @@ -354,14 +453,14 @@ class BIMVFIInterpolate: if use_target_fps: if num_passes == 0: expected = int(math.floor((N - 1) / source_fps * target_fps)) + 1 - logger.info(f"BIM-VFI: {N} frames, {source_fps}fps -> {target_fps}fps (downsampling), expected output: {expected} frames") + logger.info(f"{self.MODEL_LABEL}: {N} frames, {source_fps}fps -> {target_fps}fps (downsampling), expected output: {expected} frames") else: expected_target = int(math.floor((N - 1) / source_fps * target_fps)) + 1 - logger.info(f"BIM-VFI: interpolating {N} frames, {source_fps}fps -> {target_fps}fps (oversample {mult}x, {num_passes} pass(es)), expected output: {expected_target} frames") + logger.info(f"{self.MODEL_LABEL}: interpolating {N} frames, {source_fps}fps -> {target_fps}fps (oversample {mult}x, {num_passes} pass(es)), expected output: {expected_target} frames") else: - logger.info(f"BIM-VFI: interpolating {N} frames, {mult}x ({num_passes} pass(es)), expected output: {expected} frames") + logger.info(f"{self.MODEL_LABEL}: interpolating {N} frames, {mult}x ({num_passes} pass(es)), expected output: {expected} frames") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True @@ -386,7 +485,7 @@ class BIMVFIInterpolate: break if len(chunks) > 1: - logger.info(f"BIM-VFI: processing in {len(chunks)} chunk(s)") + logger.info(f"{self.MODEL_LABEL}: processing in {len(chunks)} chunk(s)") # Calculate total progress steps across all chunks total_steps = sum(self._count_steps(ce - cs, num_passes) for cs, ce in chunks) @@ -427,7 +526,7 @@ class BIMVFIInterpolate: # Convert back to ComfyUI [B, H, W, C], on CPU result = result.cpu().permute(0, 2, 3, 1) - logger.info(f"BIM-VFI: done, {result.shape[0]} output frames") + logger.info(f"{self.MODEL_LABEL}: done, {result.shape[0]} output frames") return (result, oversampled) @@ -463,13 +562,18 @@ class BIMVFISegmentInterpolate(BIMVFIInterpolate): def interpolate(self, images, model, multiplier, clear_cache_after_n_frames, keep_device, all_on_gpu, batch_size, chunk_size, segment_index, segment_size, - source_fps=0.0, target_fps=0.0, settings=None): + source_fps=0.0, target_fps=0.0, settings=None, seed=None): batch_size, chunk_size, keep_device, all_on_gpu, clear_cache_after_n_frames = \ _apply_vfi_settings(settings, batch_size, chunk_size, keep_device, all_on_gpu, clear_cache_after_n_frames) + if seed is not None and hasattr(model, "set_seed"): + model.set_seed(seed) + if hasattr(model, "reset_seed"): + model.reset_seed() + total_input = images.shape[0] - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) # Compute segment boundaries (1-frame overlap) start = segment_index * (segment_size - 1) @@ -480,7 +584,7 @@ class BIMVFISegmentInterpolate(BIMVFIInterpolate): return (images[:1], model) segment_images = images[start:end] - logger.info(f"BIM-VFI segment {segment_index}: input frames [{start}:{end}] of {total_input}") + logger.info(f"{self.MODEL_LABEL} segment {segment_index}: input frames [{start}:{end}] of {total_input}") if use_target_fps: num_passes, mult = _compute_target_fps_params(source_fps, target_fps) @@ -498,9 +602,12 @@ class BIMVFISegmentInterpolate(BIMVFIInterpolate): j_end = min(int(math.floor(seg_end_time * target_fps)), total_output - 1) if j_start > j_end: - return (images[:1], model) + raise ValueError( + "This segment contains no frames at the requested target FPS. " + "Increase segment_size or use the non-segment Interpolate node." + ) - logger.info(f"BIM-VFI segment {segment_index}: target fps output j=[{j_start}..{j_end}]") + logger.info(f"{self.MODEL_LABEL} segment {segment_index}: target fps output j=[{j_start}..{j_end}]") if num_passes == 0: # Downsampling — select from segment input directly @@ -516,7 +623,7 @@ class BIMVFISegmentInterpolate(BIMVFIInterpolate): return (result, model) # Oversample segment using computed num_passes - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True storage_device = device if all_on_gpu else torch.device("cpu") @@ -549,6 +656,7 @@ class BIMVFISegmentInterpolate(BIMVFIInterpolate): (result, _) = super().interpolate( segment_images, model, multiplier, clear_cache_after_n_frames, keep_device, all_on_gpu, batch_size, chunk_size, + seed=seed, ) if is_continuation: @@ -627,6 +735,13 @@ class TweenConcatVideos: if not os.path.isdir(out_dir): raise ValueError(f"Output directory does not exist: {out_dir}") + if ( + not output_filename + or os.path.basename(output_filename) != output_filename + or any(char in output_filename for char in "\r\n") + ): + raise ValueError("output_filename must be a plain filename without newlines") + output_path = os.path.abspath(os.path.join(out_dir, output_filename)) # Find segment files matching the prefix safe_prefix = glob.escape(filename_prefix) @@ -635,7 +750,16 @@ class TweenConcatVideos: segments.extend( glob.glob(os.path.join(out_dir, f"{safe_prefix}_*.{ext}")) ) - segments.sort() + segments = [ + segment for segment in segments + if os.path.abspath(segment) != output_path + ] + if any("\n" in segment or "\r" in segment for segment in segments): + raise ValueError("Segment filenames cannot contain newline characters") + segments.sort(key=lambda path: [ + int(part) if part.isdigit() else part.lower() + for part in re.split(r"(\d+)", os.path.basename(path)) + ]) if not segments: raise FileNotFoundError( @@ -655,7 +779,6 @@ class TweenConcatVideos: escaped = os.path.abspath(seg).replace("\\", "\\\\").replace("'", "\\'") f.write(f"file '{escaped}'\n") - output_path = os.path.join(out_dir, output_filename) ffmpeg = self._find_ffmpeg() cmd = [ @@ -710,7 +833,7 @@ class VFIOptimizer: "tooltip": "Input images — only the first 2 frames are used for calibration.", }), "model": ("*", { - "tooltip": "Any VFI model (BIM, EMA, SGM, GIMM). Used for benchmark inference.", + "tooltip": "Any pairwise VFI model (BIM, EMA, SGM, GIMM, SPEED). LDF-VFI uses a separate sequence pipeline.", }), "min_free_vram_gb": ("FLOAT", { "default": 2.0, "min": 0.0, "max": 48.0, "step": 0.5, @@ -744,11 +867,13 @@ class VFIOptimizer: }) def optimize(self, images, model, min_free_vram_gb, force_batch_size=0): - if images.shape[0] < 2 or not torch.cuda.is_available(): - logger.info("VFI Optimizer: <2 frames or no CUDA, returning conservative defaults") + if isinstance(model, LDFVFIModel): + logger.info("VFI Optimizer: LDF-VFI is sequence-native; returning conservative defaults") + return self._conservative_defaults(images) + device = _get_torch_device() + if images.shape[0] < 2 or device.type != "cuda": + logger.info("VFI Optimizer: <2 frames or selected device is not CUDA; returning conservative defaults") return self._conservative_defaults(images) - - device = torch.device("cuda") # --- Static analysis: model VRAM --- model_params = getattr(model, "model", model) @@ -765,6 +890,7 @@ class VFIOptimizer: frame1 = images[1:2].permute(0, 3, 1, 2) try: + _throw_if_interrupted() model.to(device) torch.cuda.reset_peak_memory_stats(device) mem_before = torch.cuda.memory_allocated(device) @@ -778,18 +904,20 @@ class VFIOptimizer: peak_mem = torch.cuda.max_memory_allocated(device) per_pair_vram_bytes = peak_mem - mem_before except Exception as e: + interrupt_type = getattr(model_management, "InterruptProcessingException", ()) + if interrupt_type and isinstance(e, interrupt_type): + raise logger.warning(f"VFI Optimizer: calibration failed ({e}), returning conservative defaults") - try: - _clear_model_cache(model) - model.to("cpu") - except Exception: - pass return self._conservative_defaults(images) finally: - _clear_model_cache(model) - model.to("cpu") - if torch.cuda.is_available(): - torch.cuda.empty_cache() + try: + model.to("cpu") + except Exception: + logger.exception("VFI Optimizer: model offload failed") + try: + _clear_model_cache(model) + except Exception: + logger.exception("VFI Optimizer: cache cleanup failed") per_pair_vram_mb = max(per_pair_vram_bytes / (1024 ** 2), 1.0) @@ -984,7 +1112,7 @@ class EMAVFIInterpolate: }), "target_fps": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01, - "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes the optimal power-of-2 oversample then selects frames. 0 = use multiplier.", + "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes a power-of-2 oversample up to 8x, then selects frames. 0 = use multiplier.", }), }, "optional": { @@ -1011,21 +1139,18 @@ class EMAVFIInterpolate: pairs_since_clear = 0 for i in range(0, num_pairs, batch_size): + _throw_if_interrupted() batch_end = min(i + batch_size, num_pairs) actual_batch = batch_end - i frames0 = frames[i:batch_end] frames1 = frames[i + 1:batch_end + 1] - if not keep_device: - model.to(device) - - mids = model.interpolate_batch(frames0, frames1, time_step=0.5) + mids = _interpolate_batch_with_offload( + model, frames0, frames1, device, keep_device + ) mids = mids.to(storage_device) - if not keep_device: - model.to("cpu") - for j in range(actual_batch): new_frames.append(frames[i + j:i + j + 1]) new_frames.append(mids[j:j+1]) @@ -1034,17 +1159,16 @@ class EMAVFIInterpolate: pbar.update_absolute(step_ref[0]) pairs_since_clear += actual_batch - if pairs_since_clear >= clear_cache_after_n_frames and torch.cuda.is_available(): + if pairs_since_clear >= clear_cache_after_n_frames: clear_ema_warp_cache() - torch.cuda.empty_cache() + _soft_empty_cache() pairs_since_clear = 0 new_frames.append(frames[-1:]) frames = torch.cat(new_frames, dim=0) - if torch.cuda.is_available(): - clear_ema_warp_cache() - torch.cuda.empty_cache() + clear_ema_warp_cache() + _soft_empty_cache() return frames @@ -1069,7 +1193,7 @@ class EMAVFIInterpolate: return (images, images) # Target FPS mode: auto-compute multiplier from fps ratio - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) if use_target_fps: num_passes, mult = _compute_target_fps_params(source_fps, target_fps) if num_passes == 0: @@ -1092,7 +1216,7 @@ class EMAVFIInterpolate: else: logger.info(f"EMA-VFI: interpolating {N} frames, {mult}x ({num_passes} pass(es)), expected output: {expected} frames") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True @@ -1200,7 +1324,7 @@ class EMAVFISegmentInterpolate(EMAVFIInterpolate): all_on_gpu, clear_cache_after_n_frames) total_input = images.shape[0] - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) # Compute segment boundaries (1-frame overlap) start = segment_index * (segment_size - 1) @@ -1227,7 +1351,10 @@ class EMAVFISegmentInterpolate(EMAVFIInterpolate): j_end = min(int(math.floor(seg_end_time * target_fps)), total_output - 1) if j_start > j_end: - return (images[:1], model) + raise ValueError( + "This segment contains no frames at the requested target FPS. " + "Increase segment_size or use the non-segment Interpolate node." + ) logger.info(f"EMA-VFI segment {segment_index}: target fps output j=[{j_start}..{j_end}]") @@ -1244,7 +1371,7 @@ class EMAVFISegmentInterpolate(EMAVFIInterpolate): return (result, model) # Oversample segment using computed num_passes - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True storage_device = device if all_on_gpu else torch.device("cpu") @@ -1412,7 +1539,7 @@ class SGMVFIInterpolate: }), "target_fps": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01, - "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes the optimal power-of-2 oversample then selects frames. 0 = use multiplier.", + "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes a power-of-2 oversample up to 8x, then selects frames. 0 = use multiplier.", }), }, "optional": { @@ -1439,21 +1566,18 @@ class SGMVFIInterpolate: pairs_since_clear = 0 for i in range(0, num_pairs, batch_size): + _throw_if_interrupted() batch_end = min(i + batch_size, num_pairs) actual_batch = batch_end - i frames0 = frames[i:batch_end] frames1 = frames[i + 1:batch_end + 1] - if not keep_device: - model.to(device) - - mids = model.interpolate_batch(frames0, frames1, time_step=0.5) + mids = _interpolate_batch_with_offload( + model, frames0, frames1, device, keep_device + ) mids = mids.to(storage_device) - if not keep_device: - model.to("cpu") - for j in range(actual_batch): new_frames.append(frames[i + j:i + j + 1]) new_frames.append(mids[j:j+1]) @@ -1462,17 +1586,16 @@ class SGMVFIInterpolate: pbar.update_absolute(step_ref[0]) pairs_since_clear += actual_batch - if pairs_since_clear >= clear_cache_after_n_frames and torch.cuda.is_available(): + if pairs_since_clear >= clear_cache_after_n_frames: clear_sgm_warp_cache() - torch.cuda.empty_cache() + _soft_empty_cache() pairs_since_clear = 0 new_frames.append(frames[-1:]) frames = torch.cat(new_frames, dim=0) - if torch.cuda.is_available(): - clear_sgm_warp_cache() - torch.cuda.empty_cache() + clear_sgm_warp_cache() + _soft_empty_cache() return frames @@ -1497,7 +1620,7 @@ class SGMVFIInterpolate: return (images, images) # Target FPS mode: auto-compute multiplier from fps ratio - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) if use_target_fps: num_passes, mult = _compute_target_fps_params(source_fps, target_fps) if num_passes == 0: @@ -1520,7 +1643,7 @@ class SGMVFIInterpolate: else: logger.info(f"SGM-VFI: interpolating {N} frames, {mult}x ({num_passes} pass(es)), expected output: {expected} frames") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True @@ -1628,7 +1751,7 @@ class SGMVFISegmentInterpolate(SGMVFIInterpolate): all_on_gpu, clear_cache_after_n_frames) total_input = images.shape[0] - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) # Compute segment boundaries (1-frame overlap) start = segment_index * (segment_size - 1) @@ -1655,7 +1778,10 @@ class SGMVFISegmentInterpolate(SGMVFIInterpolate): j_end = min(int(math.floor(seg_end_time * target_fps)), total_output - 1) if j_start > j_end: - return (images[:1], model) + raise ValueError( + "This segment contains no frames at the requested target FPS. " + "Increase segment_size or use the non-segment Interpolate node." + ) logger.info(f"SGM-VFI segment {segment_index}: target fps output j=[{j_start}..{j_end}]") @@ -1672,7 +1798,7 @@ class SGMVFISegmentInterpolate(SGMVFIInterpolate): return (result, model) # Oversample segment using computed num_passes - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True storage_device = device if all_on_gpu else torch.device("cpu") @@ -1713,6 +1839,354 @@ class SGMVFISegmentInterpolate(SGMVFIInterpolate): return (result, model) +# --------------------------------------------------------------------------- +# LDF-VFI nodes +# --------------------------------------------------------------------------- + +def ensure_ldf_model_files(): + missing = [ + filename for filename in LDF_MODEL_FILES + if not os.path.isfile(os.path.join(LDF_MODEL_DIR, filename)) + ] + if not missing: + return + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise RuntimeError( + "huggingface_hub is required to download LDF-VFI. " + "Install it with: pip install huggingface_hub" + ) from exc + + logger.warning( + "Downloading LDF-VFI from %s (~6.4 GB total). This only happens once.", + LDF_HF_REPO, + ) + for filename in missing: + logger.info("Downloading LDF-VFI file: %s", filename) + downloaded = hf_hub_download( + repo_id=LDF_HF_REPO, + filename=filename, + local_dir=LDF_MODEL_DIR, + ) + if not os.path.isfile(downloaded): + raise RuntimeError(f"Failed to download LDF-VFI file: {filename}") + + +class LoadLDFVFIModel: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": (["onecat-ai/LDF-VFI"], { + "default": "onecat-ai/LDF-VFI", + "tooltip": "Official LDF-VFI transformer + conditional VAE. Downloads ~6.4 GB on first use and needs about 20 GB VRAM.", + }), + "tile_size": ("INT", { + "default": 256, "min": 128, "max": 1024, "step": 8, + "tooltip": "Spatial VAE tile size. Larger tiles can improve throughput but use more VRAM.", + }), + "tile_overlap": ("INT", { + "default": 64, "min": 8, "max": 512, "step": 8, + "tooltip": "Spatial overlap blended between VAE tiles. Must be smaller than tile_size.", + }), + "vae_batch_size": ("INT", { + "default": 8, "min": 1, "max": 32, "step": 1, + "tooltip": "VAE temporal-tile batch size. Lower this first if VAE encoding or decoding runs out of VRAM.", + }), + "attention_type": ([ + "slide_chunk_all_block_2x1x1", + "slide_chunk_all_block", + "slide_chunk_all", + "full", + ], { + "default": "slide_chunk_all_block_2x1x1", + "tooltip": "Official quick-start sparse attention is recommended. Full attention is extremely memory-intensive.", + }), + } + } + + RETURN_TYPES = ("LDF_VFI_MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "load_model" + CATEGORY = "video/LDF-VFI" + + def load_model(self, model, tile_size, tile_overlap, vae_batch_size, attention_type): + del model # The combo documents the fixed official checkpoint. + if tile_overlap >= tile_size: + raise ValueError("LDF-VFI tile_overlap must be smaller than tile_size") + source_root = ensure_upstream_source("ldf", LDF_MODEL_DIR) + ensure_ldf_model_files() + wrapper = LDFVFIModel( + model_root=LDF_MODEL_DIR, + vae_path=os.path.join(LDF_MODEL_DIR, "Wan2.1_VAE_cond_v2.pth"), + source_root=source_root, + tile_size=tile_size, + tile_overlap=tile_overlap, + vae_batch_size=vae_batch_size, + attention_type=attention_type, + ) + logger.info( + "LDF-VFI loaded on CPU (tile=%s, overlap=%s, VAE batch=%s)", + tile_size, tile_overlap, vae_batch_size, + ) + return (wrapper,) + + +class LDFVFIInterpolate: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "images": ("IMAGE", { + "tooltip": "Ordered source sequence. LDF models it holistically with internal skip-concat chunks.", + }), + "model": ("LDF_VFI_MODEL", { + "tooltip": "LDF-VFI model from the Load LDF-VFI Model node.", + }), + "temporal_factor": ("INT", { + "default": 8, "min": 2, "max": 16, "step": 1, + "tooltip": "Native temporal upsampling factor. LDF supports every integer from 2x through 16x.", + }), + "sampling_steps": ("INT", { + "default": 16, "min": 1, "max": 100, "step": 1, + "tooltip": "Diffusion steps per temporal chunk. Official quick start uses 16; fewer is faster but may reduce quality.", + }), + "t_shift": ("FLOAT", { + "default": 8.0, "min": 0.1, "max": 20.0, "step": 0.1, + "tooltip": "Diffusion timestep shift. Official quick start uses 8.", + }), + "t_cond": ("FLOAT", { + "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, + "tooltip": "Noise applied to autoregressive boundary latents during skip-concat sampling.", + }), + "seed": ("INT", { + "default": 42, "min": 0, "max": 0x7FFFFFFF, "step": 1, + "tooltip": "Seed for VAE sampling and diffusion noise.", + }), + "offload_after": ("BOOLEAN", { + "default": True, + "tooltip": "Move the ~6.4 GB model stack back to CPU after generation to release VRAM.", + }), + "source_fps": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01, + "tooltip": "Input FPS. Set with target_fps to produce the requested output cadence.", + }), + "target_fps": ("FLOAT", { + "default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01, + "tooltip": "Optional exact output FPS. Chooses the smallest native integer factor up to 16x, then selects the nearest generated frames.", + }), + } + } + + RETURN_TYPES = ("IMAGE", "IMAGE") + RETURN_NAMES = ("images", "generated_sequence") + FUNCTION = "interpolate" + CATEGORY = "video/LDF-VFI" + + def interpolate(self, images, model, temporal_factor, sampling_steps, + t_shift, t_cond, seed, offload_after, + source_fps=0.0, target_fps=0.0): + if images.shape[0] < 2: + return (images, images) + device = _get_torch_device() + if device.type != "cuda": + raise RuntimeError( + f"LDF-VFI requires an NVIDIA CUDA GPU with BF16 support; " + f"ComfyUI selected {device}" + ) + + use_target_fps = _target_fps_enabled(source_fps, target_fps) + if use_target_fps: + ratio = target_fps / source_fps + if ratio <= 1: + source = images.permute(0, 3, 1, 2) + selected = _select_target_fps_frames( + source, source_fps, target_fps, 1, source.shape[0] + ).permute(0, 2, 3, 1).cpu() + return (selected, images) + temporal_factor = math.ceil(ratio) + if temporal_factor > 16: + raise ValueError( + f"LDF-VFI supports at most 16x, but {source_fps} -> {target_fps} FPS needs {temporal_factor}x" + ) + + source = images.permute(0, 3, 1, 2).contiguous() + blocks = model.sampling_block_count(source.shape[0], temporal_factor) + pbar = ProgressBar(blocks * sampling_steps) + progress_step = [0] + + def update_progress(): + _throw_if_interrupted() + progress_step[0] += 1 + pbar.update_absolute(progress_step[0]) + + logger.info( + "LDF-VFI: %s source frames, %sx, %s sampling steps, %s diffusion blocks", + source.shape[0], temporal_factor, sampling_steps, blocks, + ) + try: + model.to(device) + generated = model.interpolate_sequence( + source, + temporal_factor=temporal_factor, + num_steps=sampling_steps, + t_shift=t_shift, + t_cond=t_cond, + seed=seed, + progress_callback=update_progress, + ) + finally: + generation_failed = sys.exc_info()[0] is not None + cleanup_error = None + if offload_after: + try: + model.to("cpu") + except Exception as exc: + cleanup_error = exc + logger.exception("Failed to offload LDF-VFI after generation") + try: + _clear_model_cache(model) + except Exception as exc: + if cleanup_error is None: + cleanup_error = exc + logger.exception("Failed to clear LDF-VFI caches") + if cleanup_error is not None and not generation_failed: + raise RuntimeError("LDF-VFI cleanup failed") from cleanup_error + + generated_sequence = generated.permute(0, 2, 3, 1).cpu() + if use_target_fps: + generated = _select_target_fps_frames( + generated, source_fps, target_fps, + temporal_factor, source.shape[0], + ) + result = generated.permute(0, 2, 3, 1).cpu() + logger.info("LDF-VFI: done, %s output frames", result.shape[0]) + return (result, generated_sequence) + + +# --------------------------------------------------------------------------- +# SPEED nodes +# --------------------------------------------------------------------------- + +def get_available_speed_models(): + models = [] + if os.path.isdir(SPEED_MODEL_DIR): + for filename in os.listdir(SPEED_MODEL_DIR): + if filename.endswith((".pt", ".pth", ".ckpt")): + models.append(filename) + if not models: + models.append(SPEED_DEFAULT_MODEL) + return sorted(models) + + +def download_speed_model(filename, dest_dir): + try: + from huggingface_hub import hf_hub_download + except ImportError as exc: + raise RuntimeError( + "huggingface_hub is required to auto-download SPEED. " + "Install it with: pip install huggingface_hub" + ) from exc + + logger.info("Downloading %s from Hugging Face (%s)...", filename, SPEED_HF_REPO) + downloaded = hf_hub_download( + repo_id=SPEED_HF_REPO, + filename=filename, + local_dir=dest_dir, + ) + if not os.path.isfile(downloaded): + raise RuntimeError(f"Failed to download SPEED checkpoint to {downloaded}") + return downloaded + + +class LoadSPEEDVFIModel: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model_path": (get_available_speed_models(), { + "default": SPEED_DEFAULT_MODEL, + "tooltip": "Checkpoint in models/speed-vfi/. The official ~447 MB checkpoint and pinned runtime source download on first use.", + }), + "precision": (["auto", "bf16", "fp16", "fp32"], { + "default": "auto", + "tooltip": "Inference precision. Auto uses BF16 on supported CUDA GPUs, otherwise FP16. FP32 is slower and uses more VRAM.", + }), + } + } + + RETURN_TYPES = ("SPEED_VFI_MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "load_model" + CATEGORY = "video/SPEED" + + def load_model(self, model_path, precision): + source_root = ensure_upstream_source("speed", SPEED_MODEL_DIR) + full_path = os.path.join(SPEED_MODEL_DIR, model_path) + if not os.path.isfile(full_path): + full_path = download_speed_model(model_path, SPEED_MODEL_DIR) + + wrapper = SpeedVFIModel( + checkpoint_path=full_path, + source_root=source_root, + precision=precision, + device="cpu", + ) + logger.info("SPEED loaded (precision=%s)", precision) + return (wrapper,) + + +class SPEEDVFIInterpolate(BIMVFIInterpolate): + MODEL_LABEL = "SPEED" + CATEGORY = "video/SPEED" + + @classmethod + def INPUT_TYPES(cls): + inputs = BIMVFIInterpolate.INPUT_TYPES() + inputs["required"]["model"] = ("SPEED_VFI_MODEL", { + "tooltip": "SPEED model from the Load SPEED Model node.", + }) + inputs["required"]["multiplier"] = ([2, 4, 8], { + "default": 2, + "tooltip": "SPEED is midpoint-only: 4x and 8x use recursive midpoint passes.", + }) + inputs["required"]["seed"] = ("INT", { + "default": 0, "min": 0, "max": 0x7FFFFFFF, "step": 1, + "tooltip": "SPEED starts from random pixel noise. The seed makes the same execution settings repeatable without reloading the model.", + }) + inputs["required"]["batch_size"][1]["tooltip"] += ( + " SPEED is stochastic, so changing batch size can change how seeded noise is assigned to pairs." + ) + inputs["required"]["chunk_size"][1]["tooltip"] = ( + "Process the source in overlapping chunks to bound VRAM. SPEED remains repeatable for the same " + "seed and settings, but changing chunk or batch boundaries can change its stochastic result." + ) + inputs["required"]["keep_device"][1]["tooltip"] = ( + "Keep the ~235M-parameter model on GPU between batches. Faster, but uses substantial VRAM." + ) + return inputs + + +class SPEEDVFISegmentInterpolate(BIMVFISegmentInterpolate): + MODEL_LABEL = "SPEED" + RETURN_TYPES = ("IMAGE", "SPEED_VFI_MODEL") + CATEGORY = "video/SPEED" + + @classmethod + def INPUT_TYPES(cls): + inputs = SPEEDVFIInterpolate.INPUT_TYPES() + inputs["required"]["segment_index"] = ("INT", { + "default": 0, "min": 0, "max": 10000, "step": 1, + "tooltip": "Zero-based segment to process. Adjacent segments overlap by one source frame.", + }) + inputs["required"]["segment_size"] = ("INT", { + "default": 500, "min": 2, "max": 10000, "step": 1, + "tooltip": "Input frames per segment. Save each segment before processing the next to bound RAM.", + }) + return inputs + + # --------------------------------------------------------------------------- # GIMM-VFI nodes # --------------------------------------------------------------------------- @@ -1857,7 +2331,7 @@ class GIMMVFIInterpolate: }), "target_fps": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1000.0, "step": 0.01, - "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes the optimal power-of-2 oversample then selects frames. 0 = use multiplier.", + "tooltip": "Target output FPS. When > 0, overrides multiplier and auto-computes a power-of-2 oversample up to 8x, then selects frames. 0 = use multiplier.", }), }, "optional": { @@ -1884,18 +2358,15 @@ class GIMMVFIInterpolate: pairs_since_clear = 0 for i in range(num_pairs): + _throw_if_interrupted() frame0 = frames[i:i+1] frame1 = frames[i+1:i+2] - if not keep_device: - model.to(device) - - mids = model.interpolate_multi(frame0, frame1, num_intermediates) + mids = _interpolate_multi_with_offload( + model, frame0, frame1, num_intermediates, device, keep_device + ) mids = [m.to(storage_device) for m in mids] - if not keep_device: - model.to("cpu") - new_frames.append(frames[i:i+1]) for m in mids: new_frames.append(m) @@ -1904,17 +2375,16 @@ class GIMMVFIInterpolate: pbar.update_absolute(step_ref[0]) pairs_since_clear += 1 - if pairs_since_clear >= clear_cache_after_n_frames and torch.cuda.is_available(): + if pairs_since_clear >= clear_cache_after_n_frames: clear_gimm_caches() - torch.cuda.empty_cache() + _soft_empty_cache() pairs_since_clear = 0 new_frames.append(frames[-1:]) result = torch.cat(new_frames, dim=0) - if torch.cuda.is_available(): - clear_gimm_caches() - torch.cuda.empty_cache() + clear_gimm_caches() + _soft_empty_cache() return result @@ -1929,21 +2399,18 @@ class GIMMVFIInterpolate: pairs_since_clear = 0 for i in range(0, num_pairs, batch_size): + _throw_if_interrupted() batch_end = min(i + batch_size, num_pairs) actual_batch = batch_end - i frames0 = frames[i:batch_end] frames1 = frames[i + 1:batch_end + 1] - if not keep_device: - model.to(device) - - mids = model.interpolate_batch(frames0, frames1, time_step=0.5) + mids = _interpolate_batch_with_offload( + model, frames0, frames1, device, keep_device + ) mids = mids.to(storage_device) - if not keep_device: - model.to("cpu") - for j in range(actual_batch): new_frames.append(frames[i + j:i + j + 1]) new_frames.append(mids[j:j+1]) @@ -1952,17 +2419,16 @@ class GIMMVFIInterpolate: pbar.update_absolute(step_ref[0]) pairs_since_clear += actual_batch - if pairs_since_clear >= clear_cache_after_n_frames and torch.cuda.is_available(): + if pairs_since_clear >= clear_cache_after_n_frames: clear_gimm_caches() - torch.cuda.empty_cache() + _soft_empty_cache() pairs_since_clear = 0 new_frames.append(frames[-1:]) frames = torch.cat(new_frames, dim=0) - if torch.cuda.is_available(): - clear_gimm_caches() - torch.cuda.empty_cache() + clear_gimm_caches() + _soft_empty_cache() return frames @@ -1988,7 +2454,7 @@ class GIMMVFIInterpolate: return (images, images) # Target FPS mode: auto-compute multiplier from fps ratio - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) if use_target_fps: num_passes, mult = _compute_target_fps_params(source_fps, target_fps) if num_passes == 0: @@ -2000,6 +2466,12 @@ class GIMMVFIInterpolate: else: mult = multiplier + if not single_pass or use_target_fps: + num_passes_recursive = ( + num_passes if use_target_fps + else {2: 1, 4: 2, 8: 3}[multiplier] + ) + N = images.shape[0] expected = mult * (N - 1) + 1 if use_target_fps: @@ -2010,15 +2482,10 @@ class GIMMVFIInterpolate: expected_target = int(math.floor((N - 1) / source_fps * target_fps)) + 1 logger.info(f"GIMM-VFI: interpolating {N} frames, {source_fps}fps -> {target_fps}fps (oversample {mult}x, {num_passes} pass(es)), expected output: {expected_target} frames") else: - logger.info(f"GIMM-VFI: interpolating {N} frames, {mult}x ({num_passes if not single_pass else 'single-pass'}), expected output: {expected} frames") + mode = f"{num_passes_recursive} recursive pass(es)" if not single_pass else "single-pass" + logger.info(f"GIMM-VFI: interpolating {N} frames, {mult}x ({mode}), expected output: {expected} frames") - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - - if not single_pass or use_target_fps: - if use_target_fps: - num_passes_recursive = num_passes - else: - num_passes_recursive = {2: 1, 4: 2, 8: 3}[multiplier] + device = _get_torch_device() if all_on_gpu: keep_device = True @@ -2136,7 +2603,7 @@ class GIMMVFISegmentInterpolate(GIMMVFIInterpolate): all_on_gpu, clear_cache_after_n_frames) total_input = images.shape[0] - use_target_fps = target_fps > 0 and source_fps > 0 + use_target_fps = _target_fps_enabled(source_fps, target_fps) # Compute segment boundaries (1-frame overlap) start = segment_index * (segment_size - 1) @@ -2163,7 +2630,10 @@ class GIMMVFISegmentInterpolate(GIMMVFIInterpolate): j_end = min(int(math.floor(seg_end_time * target_fps)), total_output - 1) if j_start > j_end: - return (images[:1], model) + raise ValueError( + "This segment contains no frames at the requested target FPS. " + "Increase segment_size or use the non-segment Interpolate node." + ) logger.info(f"GIMM-VFI segment {segment_index}: target fps output j=[{j_start}..{j_end}]") @@ -2180,7 +2650,7 @@ class GIMMVFISegmentInterpolate(GIMMVFIInterpolate): return (result, model) # Oversample segment directly - device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + device = _get_torch_device() if all_on_gpu: keep_device = True storage_device = device if all_on_gpu else torch.device("cpu") diff --git a/pyproject.toml b/pyproject.toml index d249ce9..17bb1e9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "comfyui-tween" -description = "Video frame interpolation nodes for ComfyUI using BIM-VFI, EMA-VFI, SGM-VFI, and GIMM-VFI. Designed for long videos with thousands of frames." -version = "1.1.0" +description = "Video frame interpolation nodes for ComfyUI using BIM-VFI, EMA-VFI, SGM-VFI, GIMM-VFI, SPEED, and LDF-VFI." +version = "1.2.0" license = "Apache-2.0" requires-python = ">=3.10" dependencies = [ @@ -12,6 +12,9 @@ dependencies = [ "easydict", "einops", "huggingface_hub", + "diffusers>=0.33.1,<0.40", + "accelerate>=1.5,<2", + "safetensors", ] [project.urls] diff --git a/requirements.txt b/requirements.txt index 8e45d1e..5ca465f 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,3 +5,6 @@ yacs easydict einops huggingface_hub +diffusers>=0.33.1,<0.40 +accelerate>=1.5,<2 +safetensors diff --git a/speed_backend.py b/speed_backend.py new file mode 100644 index 0000000..997a3d0 --- /dev/null +++ b/speed_backend.py @@ -0,0 +1,158 @@ +"""ComfyUI inference adapter for the official SPEED runtime.""" + +from __future__ import annotations + +from contextlib import nullcontext +import importlib +import logging +from pathlib import Path +import sys +import types + +import torch + + +logger = logging.getLogger("Tween") +_SPEED_NAMESPACE = "_tween_speed_upstream" + + +def _cuda_bf16_supported(device: torch.device) -> bool: + if device.type != "cuda": + return False + with torch.cuda.device(device): + return torch.cuda.is_bf16_supported() + + +def _namespace_package(name: str, path: Path) -> None: + if name in sys.modules: + return + package = types.ModuleType(name) + package.__path__ = [str(path)] + package.__package__ = name + sys.modules[name] = package + + +def load_speed_model_class(source_root: str): + """Import SpeedDiT without adding the upstream repository to sys.path.""" + root = Path(source_root).resolve() + model_file = root / "src" / "models" / "model.py" + if not model_file.is_file(): + raise RuntimeError(f"Invalid SPEED source directory: missing {model_file}") + + _namespace_package(_SPEED_NAMESPACE, root) + _namespace_package(f"{_SPEED_NAMESPACE}.src", root / "src") + _namespace_package(f"{_SPEED_NAMESPACE}.src.models", root / "src" / "models") + module = importlib.import_module(f"{_SPEED_NAMESPACE}.src.models.model") + return module.SpeedDiT + + +class SpeedVFIModel: + """Midpoint interpolation wrapper around the official SPEED SpeedDiT.""" + + def __init__(self, checkpoint_path: str, source_root: str, + precision: str = "auto", device: str = "cpu"): + SpeedDiT = load_speed_model_class(source_root) + self.model = SpeedDiT( + hidden_dim=768, + head_dim=64, + depths=(2, 6, 4), + patch_sizes=(64, 32, 16), + ) + self.precision = precision + self.device = str(device) + self._seed = 0 + self._generator = None + self._generator_device = None + self._load_checkpoint(checkpoint_path) + self.model.requires_grad_(False).eval() + self.to(device) + + def _load_checkpoint(self, checkpoint_path: str) -> None: + checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) + state_dict = checkpoint.get("model", checkpoint) if isinstance(checkpoint, dict) else checkpoint + if not isinstance(state_dict, dict): + raise TypeError(f"SPEED checkpoint does not contain a state dict: {checkpoint_path}") + if state_dict and all(key.startswith("module.") for key in state_dict): + state_dict = {key[len("module."):]: value for key, value in state_dict.items()} + self.model.load_state_dict(state_dict, strict=True) + + def _autocast_dtype(self, device: torch.device): + if self.precision == "fp32" or device.type != "cuda": + return None + if self.precision == "fp16": + return torch.float16 + if self.precision == "bf16": + if not _cuda_bf16_supported(device): + raise RuntimeError( + "SPEED BF16 precision requires a CUDA GPU with BF16 support; " + "select auto or fp16 on this GPU" + ) + return torch.bfloat16 + if _cuda_bf16_supported(device): + return torch.bfloat16 + return torch.float16 + + def to(self, device): + target = torch.device(device) + self.device = str(target) + # Match the official runtime: retain FP32 weights and use autocast for + # CUDA inference. This also avoids mixed-dtype timestep embedding bugs. + self.model.to(device=target, dtype=torch.float32) + # Keep the generator alive across CPU offloading. Recreating it on every + # pair batch would restart the noise stream whenever keep_device=False. + # _get_generator replaces it automatically if inference changes device. + return self + + def set_seed(self, seed: int) -> None: + self._seed = int(seed) + self._generator = None + self._generator_device = None + + def reset_seed(self) -> None: + self.set_seed(self._seed) + + def clear_cache(self) -> None: + rope = getattr(self.model, "rope_embedder", None) + cache = getattr(rope, "rope_cache", None) + if cache is not None: + cache.clear() + + def _get_generator(self, device: torch.device) -> torch.Generator: + device_name = str(device) + if self._generator is None or self._generator_device != device_name: + self._generator = torch.Generator(device=device) + self._generator.manual_seed(self._seed) + self._generator_device = device_name + return self._generator + + @torch.no_grad() + def interpolate_batch(self, frames0, frames1, time_step=0.5): + if abs(float(time_step) - 0.5) > 1e-6: + raise ValueError("SPEED's released checkpoint supports midpoint interpolation only") + + device = next(self.model.parameters()).device + frame0 = frames0.to(device=device, dtype=torch.float32, non_blocking=True).mul(2).sub(1) + frame1 = frames1.to(device=device, dtype=torch.float32, non_blocking=True).mul(2).sub(1) + cond_frames = torch.cat((frame0, frame1), dim=0) + noisy_frames = torch.randn( + frame0.shape, + generator=self._get_generator(device), + device=device, + dtype=torch.float32, + ) + timestep = torch.full( + (frame0.shape[0],), 1000.0, device=device, dtype=torch.float32 + ) + autocast_dtype = self._autocast_dtype(device) + autocast = ( + torch.autocast(device_type="cuda", dtype=autocast_dtype) + if autocast_dtype is not None + else nullcontext() + ) + with autocast: + prediction = self.model( + noisy_frames=noisy_frames, + cond_frames=cond_frames, + timestep=timestep, + ) + return prediction.div(2).add(0.5).clamp_(0, 1).float()