Compare commits
9
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
15b88d1689 | ||
|
|
211481558c | ||
|
|
333a281eef | ||
|
|
5240e27038 | ||
|
|
a4831d7b71 | ||
|
|
fa6d8a7d88 | ||
|
|
cb60fe2542 | ||
|
|
b2350e7f08 | ||
|
|
5b10a1a594 |
@@ -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.
|
||||||
@@ -3,9 +3,9 @@
|
|||||||
[](https://registry.comfy.org/)
|
[](https://registry.comfy.org/)
|
||||||
[](https://www.python.org/)
|
[](https://www.python.org/)
|
||||||
[](https://www.apache.org/licenses/LICENSE-2.0)
|
[](https://www.apache.org/licenses/LICENSE-2.0)
|
||||||
[](#which-model-should-i-use)
|
[](#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.
|
||||||
|
|
||||||
<p align="center">
|
<p align="center">
|
||||||
<img src="assets/model-comparison.svg" alt="Model Comparison" width="720"/>
|
<img src="assets/model-comparison.svg" alt="Model Comparison" width="720"/>
|
||||||
@@ -21,11 +21,18 @@ git clone https://github.com/Ethanfel/ComfyUI-Tween.git
|
|||||||
pip install -r requirements.txt
|
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. There is intentionally no custom `install.py`, avoiding a second redundant dependency-install pass after Manager processes `requirements.txt`. 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_bim_model_lab.json`](example_workflows/tween_speed_bim_model_lab.json) for the recommended starter graph. It requires [ComfyUI-VideoHelperSuite](https://github.com/Kosinkadink/ComfyUI-VideoHelperSuite) for video loading and encoding.
|
||||||
|
|
||||||
|
- The SPEED and BIM-VFI branches load the same 25-frame, 24 FPS sample, tune memory settings independently, interpolate to 48 FPS, preserve audio, and save separate comparison videos.
|
||||||
|
- 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. It is deliberately **not installed automatically**, because replacing or mixing CUDA-specific cupy wheels can disrupt other ComfyUI nodes. BIM-VFI, SGM-VFI, and GIMM-VFI work without it through their PyTorch fallback. **EMA-VFI, SPEED, and LDF-VFI do not use it.**
|
||||||
|
|
||||||
1. Find your CUDA version:
|
1. Find your CUDA version:
|
||||||
```bash
|
```bash
|
||||||
@@ -36,10 +43,11 @@ All dependencies (`gdown`, `timm`, `omegaconf`, `easydict`, `yacs`, `einops`, `h
|
|||||||
|
|
||||||
| CUDA | Command |
|
| CUDA | Command |
|
||||||
|------|---------|
|
|------|---------|
|
||||||
|
| 13.x | `pip install cupy-cuda13x` |
|
||||||
| 12.x | `pip install cupy-cuda12x` |
|
| 12.x | `pip install cupy-cuda12x` |
|
||||||
| 11.x | `pip install cupy-cuda11x` |
|
| 11.x | `pip install cupy-cuda11x` |
|
||||||
|
|
||||||
> Make sure to run pip in the same Python environment as ComfyUI. If cupy is missing, the Load node shows an error with your CUDA version and the exact install command.
|
> Make sure to run pip in the same Python environment as ComfyUI, and uninstall any different cupy wheel variant first. If cupy is absent or incompatible, Tween safely uses its PyTorch fallback.
|
||||||
|
|
||||||
<details>
|
<details>
|
||||||
<summary>cupy troubleshooting</summary>
|
<summary>cupy troubleshooting</summary>
|
||||||
@@ -48,26 +56,23 @@ All dependencies (`gdown`, `timm`, `omegaconf`, `easydict`, `yacs`, `einops`, `h
|
|||||||
|---------|----------|
|
|---------|----------|
|
||||||
| `ModuleNotFoundError: No module named 'cupy'` | Install cupy using the steps above |
|
| `ModuleNotFoundError: No module named 'cupy'` | Install cupy using the steps above |
|
||||||
| `cupy` installed but `ImportError` at runtime | CUDA version mismatch — uninstall and reinstall the correct version |
|
| `cupy` installed but `ImportError` at runtime | CUDA version mismatch — uninstall and reinstall the correct version |
|
||||||
| Install hangs or takes very long | cupy wheels are ~800 MB, be patient |
|
| Install hangs or takes very long | Confirm pip selected a prebuilt wheel for your Python and CUDA versions |
|
||||||
| Docker / no build tools | Use the prebuilt wheel: `pip install cupy-cuda12x` (not bare `cupy` which compiles from source) |
|
| Docker / no build tools | Use the matching prebuilt `cupy-cudaXXx` wheel, not bare `cupy` which compiles from source |
|
||||||
|
|
||||||
</details>
|
</details>
|
||||||
|
|
||||||
## Which model should I use?
|
## Which model should I use?
|
||||||
|
|
||||||
| | BIM-VFI | EMA-VFI | SGM-VFI | GIMM-VFI |
|
| Model | Best for | Multiplier path | Typical VRAM | Trade-off |
|
||||||
|---|---------|---------|---------|----------|
|
|-------|----------|-----------------|--------------|-----------|
|
||||||
| **Best for** | General-purpose | Fast, low VRAM | Large motion | High multipliers (4x/8x) |
|
| **BIM-VFI** | Strong general pairwise quality | Recursive 2x/4x/8x | ~2 GB/pair | Research/education license |
|
||||||
| **Quality** | Highest | Good | Best on large motion | Good |
|
| **EMA-VFI** | Speed and lower VRAM | Recursive 2x/4x/8x | ~1.5 GB/pair | Less robust on extreme motion |
|
||||||
| **Speed** | Moderate | Fastest | Slowest | Fast for 4x/8x |
|
| **SGM-VFI** | Large motion | Recursive 2x/4x/8x | ~3 GB/pair | Slowest pairwise option |
|
||||||
| **VRAM** | ~2 GB/pair | ~1.5 GB/pair | ~3 GB/pair | ~2.5 GB/pair |
|
| **GIMM-VFI** | Arbitrary timesteps, efficient 4x/8x | Native multi-frame per pair | ~2.5 GB/pair | Still frame-pair-centric |
|
||||||
| **Params** | ~17 M | ~14–65 M | ~15 M + GMFlow | ~80 M (RAFT) / ~123 M (FlowFormer) |
|
| **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 |
|
||||||
| **Arbitrary timestep** | Yes | Yes (`_t` checkpoint) | No (fixed 0.5) | Yes (native) |
|
| **LDF-VFI** | Long-range temporal coherence and 2x–16x | Native sequence diffusion | ~20 GB | ~6.4 GB weights; much slower |
|
||||||
| **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 |
|
|
||||||
|
|
||||||
**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
|
## 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` |
|
| 48 GB+ | `batch_size=4–16, all_on_gpu=true` |
|
||||||
| 96 GB+ | `batch_size=8–16, all_on_gpu=true, chunk_size=0` |
|
| 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
|
## Nodes
|
||||||
|
|
||||||
All Interpolate nodes share a common set of controls:
|
The pairwise Interpolate nodes (BIM/EMA/SGM/GIMM/SPEED) share these controls:
|
||||||
|
|
||||||
| Input | Description |
|
| 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) |
|
| **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 |
|
| **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 |
|
| **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 |
|
| Output | Description |
|
||||||
|--------|-------------|
|
|--------|-------------|
|
||||||
@@ -110,8 +117,11 @@ Loads the BiM-VFI checkpoint. Auto-downloads from Google Drive on first use to `
|
|||||||
| Input | Description |
|
| Input | Description |
|
||||||
|-------|-------------|
|
|-------|-------------|
|
||||||
| **model_path** | Checkpoint from `models/bim-vfi/` |
|
| **model_path** | Checkpoint from `models/bim-vfi/` |
|
||||||
| **auto_pyr_level** | Auto pyramid level by resolution (<540p=3, 540p=5, 1080p=6, 4K=7) |
|
| **auto_pyr_level** | Official automatic pyramid policy (below 1080p=5, 1080p=6, 4K=7) |
|
||||||
| **pyr_level** | Manual pyramid level (3–7), used when auto is off |
|
| **pyr_level** | Manual pyramid level (3–7), used when auto is off |
|
||||||
|
| **artifact_safe_mode** | Disables the RGB refinement residual to suppress wrong-edge/halo artifacts caused by flow misalignment in blurry or large-motion shots. Off preserves official behavior and can retain more detail on easy shots |
|
||||||
|
|
||||||
|
`artifact_safe_mode` implements the workaround recommended by the official BIM-VFI maintainer for [wrong-edge artifacts caused by severely misaligned warped inputs](https://github.com/KAIST-VICLab/BiM-VFI/issues/1). Enable it selectively for affected footage.
|
||||||
|
|
||||||
#### BIM-VFI Interpolate
|
#### BIM-VFI Interpolate
|
||||||
|
|
||||||
@@ -202,14 +212,65 @@ Same pattern as other Segment nodes.
|
|||||||
|
|
||||||
</details>
|
</details>
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary><strong>SPEED</strong></summary>
|
||||||
|
|
||||||
|
#### 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.
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
|
<details>
|
||||||
|
<summary><strong>LDF-VFI</strong></summary>
|
||||||
|
|
||||||
|
#### 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.
|
||||||
|
|
||||||
|
</details>
|
||||||
|
|
||||||
### Tween Concat Videos
|
### 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
|
### Output frame count
|
||||||
|
|
||||||
- **Multiplier mode:** 2x = 2N-1, 4x = 4N-3, 8x = 8N-7
|
- **Pairwise 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.
|
- **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
|
## Acknowledgments
|
||||||
|
|
||||||
@@ -219,8 +280,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) |
|
| **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) |
|
| **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) |
|
| **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.
|
||||||
|
|
||||||
<details>
|
<details>
|
||||||
<summary>BibTeX citations</summary>
|
<summary>BibTeX citations</summary>
|
||||||
@@ -253,6 +316,22 @@ GIMM-VFI adaptation from [kijai/ComfyUI-GIMM-VFI](https://github.com/kijai/Comfy
|
|||||||
booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
|
booktitle={Advances in Neural Information Processing Systems (NeurIPS)},
|
||||||
year={2024}
|
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}
|
||||||
|
}
|
||||||
```
|
```
|
||||||
|
|
||||||
</details>
|
</details>
|
||||||
@@ -261,6 +340,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).
|
**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)
|
**This wrapper code:** [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0)
|
||||||
|
|||||||
+12
@@ -3,6 +3,8 @@ from .nodes import (
|
|||||||
LoadEMAVFIModel, EMAVFIInterpolate, EMAVFISegmentInterpolate,
|
LoadEMAVFIModel, EMAVFIInterpolate, EMAVFISegmentInterpolate,
|
||||||
LoadSGMVFIModel, SGMVFIInterpolate, SGMVFISegmentInterpolate,
|
LoadSGMVFIModel, SGMVFIInterpolate, SGMVFISegmentInterpolate,
|
||||||
LoadGIMMVFIModel, GIMMVFIInterpolate, GIMMVFISegmentInterpolate,
|
LoadGIMMVFIModel, GIMMVFIInterpolate, GIMMVFISegmentInterpolate,
|
||||||
|
LoadSPEEDVFIModel, SPEEDVFIInterpolate, SPEEDVFISegmentInterpolate,
|
||||||
|
LoadLDFVFIModel, LDFVFIInterpolate,
|
||||||
VFIOptimizer,
|
VFIOptimizer,
|
||||||
)
|
)
|
||||||
|
|
||||||
@@ -20,6 +22,11 @@ NODE_CLASS_MAPPINGS = {
|
|||||||
"LoadGIMMVFIModel": LoadGIMMVFIModel,
|
"LoadGIMMVFIModel": LoadGIMMVFIModel,
|
||||||
"GIMMVFIInterpolate": GIMMVFIInterpolate,
|
"GIMMVFIInterpolate": GIMMVFIInterpolate,
|
||||||
"GIMMVFISegmentInterpolate": GIMMVFISegmentInterpolate,
|
"GIMMVFISegmentInterpolate": GIMMVFISegmentInterpolate,
|
||||||
|
"LoadSPEEDVFIModel": LoadSPEEDVFIModel,
|
||||||
|
"SPEEDVFIInterpolate": SPEEDVFIInterpolate,
|
||||||
|
"SPEEDVFISegmentInterpolate": SPEEDVFISegmentInterpolate,
|
||||||
|
"LoadLDFVFIModel": LoadLDFVFIModel,
|
||||||
|
"LDFVFIInterpolate": LDFVFIInterpolate,
|
||||||
"VFIOptimizer": VFIOptimizer,
|
"VFIOptimizer": VFIOptimizer,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -37,5 +44,10 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
|||||||
"LoadGIMMVFIModel": "Load GIMM-VFI Model",
|
"LoadGIMMVFIModel": "Load GIMM-VFI Model",
|
||||||
"GIMMVFIInterpolate": "GIMM-VFI Interpolate",
|
"GIMMVFIInterpolate": "GIMM-VFI Interpolate",
|
||||||
"GIMMVFISegmentInterpolate": "GIMM-VFI Segment 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",
|
"VFIOptimizer": "VFI Optimizer",
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 720 320" width="720" height="320">
|
<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 720 475" width="720" height="475">
|
||||||
<defs>
|
<defs>
|
||||||
<linearGradient id="gQ" x1="0" y1="0" x2="1" y2="0">
|
<linearGradient id="gQ" x1="0" y1="0" x2="1" y2="0">
|
||||||
<stop offset="0%" stop-color="#7aa2f7"/><stop offset="100%" stop-color="#7dcfff"/>
|
<stop offset="0%" stop-color="#7aa2f7"/><stop offset="100%" stop-color="#7dcfff"/>
|
||||||
@@ -12,13 +12,13 @@
|
|||||||
</defs>
|
</defs>
|
||||||
|
|
||||||
<!-- Background -->
|
<!-- Background -->
|
||||||
<rect width="720" height="320" rx="16" fill="#0d1117"/>
|
<rect width="720" height="475" rx="16" fill="#0d1117"/>
|
||||||
|
|
||||||
<!-- ═══ BIM-VFI (top-left) ═══ -->
|
<!-- ═══ BIM-VFI (top-left) ═══ -->
|
||||||
<rect x="10" y="10" width="340" height="145" rx="10" fill="#161b22" stroke="#30363d" stroke-width="1"/>
|
<rect x="10" y="10" width="340" height="145" rx="10" fill="#161b22" stroke="#30363d" stroke-width="1"/>
|
||||||
<rect x="11" y="22" width="3" height="121" fill="#3fb950"/>
|
<rect x="11" y="22" width="3" height="121" fill="#3fb950"/>
|
||||||
<text x="30" y="38" fill="#e6edf3" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="15" font-weight="600">BIM-VFI</text>
|
<text x="30" y="38" fill="#e6edf3" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="15" font-weight="600">BIM-VFI</text>
|
||||||
<text x="30" y="56" fill="#3fb950" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">★ Recommended · Best quality · CVPR 2025</text>
|
<text x="30" y="56" fill="#3fb950" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Strong pairwise quality · CVPR 2025</text>
|
||||||
<line x1="30" y1="64" x2="330" y2="64" stroke="#30363d" stroke-width="0.5"/>
|
<line x1="30" y1="64" x2="330" y2="64" stroke="#30363d" stroke-width="0.5"/>
|
||||||
<text x="30" y="82" fill="#7aa2f7" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Quality</text>
|
<text x="30" y="82" fill="#7aa2f7" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Quality</text>
|
||||||
<rect x="88" y="72" width="244" height="11" rx="3" fill="#21262d"/>
|
<rect x="88" y="72" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
@@ -81,4 +81,38 @@
|
|||||||
<rect x="448" y="263" width="244" height="11" rx="3" fill="#21262d"/>
|
<rect x="448" y="263" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
<rect x="448" y="263" width="146" height="11" rx="3" fill="url(#gV)" opacity="0.85"/>
|
<rect x="448" y="263" width="146" height="11" rx="3" fill="url(#gV)" opacity="0.85"/>
|
||||||
<text x="632" y="298" fill="#8b949e" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="10">Apache 2.0</text>
|
<text x="632" y="298" fill="#8b949e" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="10">Apache 2.0</text>
|
||||||
|
|
||||||
|
<!-- ═══ SPEED (third-row left) ═══ -->
|
||||||
|
<rect x="10" y="320" width="340" height="145" rx="10" fill="#161b22" stroke="#30363d" stroke-width="1"/>
|
||||||
|
<rect x="11" y="332" width="3" height="121" fill="#f2cc60"/>
|
||||||
|
<text x="30" y="348" fill="#e6edf3" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="15" font-weight="600">SPEED</text>
|
||||||
|
<text x="30" y="366" fill="#f2cc60" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Modern pairwise default · ACM MM 2026</text>
|
||||||
|
<line x1="30" y1="374" x2="330" y2="374" stroke="#30363d" stroke-width="0.5"/>
|
||||||
|
<text x="30" y="392" fill="#7aa2f7" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Quality</text>
|
||||||
|
<rect x="88" y="382" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
|
<rect x="88" y="382" width="220" height="11" rx="3" fill="url(#gQ)" opacity="0.85"/>
|
||||||
|
<text x="30" y="410" fill="#9ece6a" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Speed</text>
|
||||||
|
<rect x="88" y="400" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
|
<rect x="88" y="400" width="190" height="11" rx="3" fill="url(#gS)" opacity="0.85"/>
|
||||||
|
<text x="30" y="428" fill="#bb9af7" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">VRAM</text>
|
||||||
|
<rect x="88" y="418" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
|
<rect x="88" y="418" width="175" height="11" rx="3" fill="url(#gV)" opacity="0.85"/>
|
||||||
|
<text x="253" y="453" fill="#f0883e" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="10">Check upstream terms</text>
|
||||||
|
|
||||||
|
<!-- ═══ LDF-VFI (third-row right) ═══ -->
|
||||||
|
<rect x="370" y="320" width="340" height="145" rx="10" fill="#161b22" stroke="#30363d" stroke-width="1"/>
|
||||||
|
<rect x="371" y="332" width="3" height="121" fill="#ff7b72"/>
|
||||||
|
<text x="390" y="348" fill="#e6edf3" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="15" font-weight="600">LDF-VFI</text>
|
||||||
|
<text x="390" y="366" fill="#ff7b72" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Long-sequence coherence · CVPR 2026</text>
|
||||||
|
<line x1="390" y1="374" x2="690" y2="374" stroke="#30363d" stroke-width="0.5"/>
|
||||||
|
<text x="390" y="392" fill="#7aa2f7" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Quality</text>
|
||||||
|
<rect x="448" y="382" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
|
<rect x="448" y="382" width="230" height="11" rx="3" fill="url(#gQ)" opacity="0.85"/>
|
||||||
|
<text x="390" y="410" fill="#9ece6a" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">Speed</text>
|
||||||
|
<rect x="448" y="400" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
|
<rect x="448" y="400" width="48" height="11" rx="3" fill="url(#gS)" opacity="0.85"/>
|
||||||
|
<text x="390" y="428" fill="#bb9af7" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="11">VRAM</text>
|
||||||
|
<rect x="448" y="418" width="244" height="11" rx="3" fill="#21262d"/>
|
||||||
|
<rect x="448" y="418" width="40" height="11" rx="3" fill="url(#gV)" opacity="0.85"/>
|
||||||
|
<text x="632" y="453" fill="#8b949e" font-family="-apple-system,BlinkMacSystemFont,'Segoe UI','Noto Sans',Helvetica,Arial,sans-serif" font-size="10">Apache 2.0</text>
|
||||||
</svg>
|
</svg>
|
||||||
|
|||||||
|
Before Width: | Height: | Size: 7.8 KiB After Width: | Height: | Size: 11 KiB |
@@ -12,13 +12,17 @@ from ..utils.padder import InputPadder
|
|||||||
|
|
||||||
|
|
||||||
class BiMVFI(nn.Module):
|
class BiMVFI(nn.Module):
|
||||||
def __init__(self, pyr_level=3, feat_channels=32, **kwargs):
|
def __init__(self, pyr_level=3, feat_channels=32,
|
||||||
|
artifact_safe_mode=False, **kwargs):
|
||||||
super(BiMVFI, self).__init__()
|
super(BiMVFI, self).__init__()
|
||||||
self.pyr_level = pyr_level
|
self.pyr_level = pyr_level
|
||||||
self.mfe = ResNetPyramid(feat_channels)
|
self.mfe = ResNetPyramid(feat_channels)
|
||||||
self.cfe = ResNetPyramid(feat_channels)
|
self.cfe = ResNetPyramid(feat_channels)
|
||||||
self.bimfn = BiMFN(feat_channels)
|
self.bimfn = BiMFN(feat_channels)
|
||||||
self.sn = SynthesisNetwork(feat_channels)
|
self.sn = SynthesisNetwork(
|
||||||
|
feat_channels,
|
||||||
|
use_rgb_refine_residual=not artifact_safe_mode,
|
||||||
|
)
|
||||||
self.feat_channels = feat_channels
|
self.feat_channels = feat_channels
|
||||||
self.caun = CAUN(feat_channels)
|
self.caun = CAUN(feat_channels)
|
||||||
|
|
||||||
|
|||||||
+12
-2
@@ -5,8 +5,9 @@ from .backwarp import backwarp
|
|||||||
|
|
||||||
|
|
||||||
class SynthesisNetwork(nn.Module):
|
class SynthesisNetwork(nn.Module):
|
||||||
def __init__(self, feat_channels):
|
def __init__(self, feat_channels, use_rgb_refine_residual=True):
|
||||||
super(SynthesisNetwork, self).__init__()
|
super(SynthesisNetwork, self).__init__()
|
||||||
|
self.use_rgb_refine_residual = use_rgb_refine_residual
|
||||||
input_channels = 6 + 1
|
input_channels = 6 + 1
|
||||||
self.conv_down1 = nn.Sequential(
|
self.conv_down1 = nn.Sequential(
|
||||||
nn.Conv2d(input_channels, feat_channels, 7, padding=3),
|
nn.Conv2d(input_channels, feat_channels, 7, padding=3),
|
||||||
@@ -59,6 +60,13 @@ class SynthesisNetwork(nn.Module):
|
|||||||
warped_img1 = backwarp(i1, flow_t1)
|
warped_img1 = backwarp(i1, flow_t1)
|
||||||
return warped_img0, warped_img1, warped_c0, warped_c1
|
return warped_img0, warped_img1, warped_c0, warped_c1
|
||||||
|
|
||||||
|
def merge_warped_images(self, warped_img0, warped_img1,
|
||||||
|
blending_mask, refine_res):
|
||||||
|
merged_img = warped_img0 * blending_mask + warped_img1 * (1 - blending_mask)
|
||||||
|
if self.use_rgb_refine_residual:
|
||||||
|
merged_img = merged_img + refine_res
|
||||||
|
return merged_img
|
||||||
|
|
||||||
def forward(self, i0, i1, c0_pyr, c1_pyr, bi_flow_pyr, occ):
|
def forward(self, i0, i1, c0_pyr, c1_pyr, bi_flow_pyr, occ):
|
||||||
warped_img0, warped_img1, warped_c0, warped_c1 = \
|
warped_img0, warped_img1, warped_c0, warped_c1 = \
|
||||||
self.get_warped_representations(
|
self.get_warped_representations(
|
||||||
@@ -82,7 +90,9 @@ class SynthesisNetwork(nn.Module):
|
|||||||
occ_res = refine[:, 3:]
|
occ_res = refine[:, 3:]
|
||||||
occ_out = occ + occ_res
|
occ_out = occ + occ_res
|
||||||
blending_mask = torch.sigmoid(occ_out)
|
blending_mask = torch.sigmoid(occ_out)
|
||||||
merged_img = (warped_img0 * blending_mask + warped_img1 * (1 - blending_mask)) + refine_res
|
merged_img = self.merge_warped_images(
|
||||||
|
warped_img0, warped_img1, blending_mask, refine_res
|
||||||
|
)
|
||||||
interp_img = merged_img
|
interp_img = merged_img
|
||||||
|
|
||||||
extra_dict = {}
|
extra_dict = {}
|
||||||
|
|||||||
@@ -0,0 +1,323 @@
|
|||||||
|
{
|
||||||
|
"last_node_id": 9,
|
||||||
|
"last_link_id": 14,
|
||||||
|
"nodes": [
|
||||||
|
{
|
||||||
|
"id": 1,
|
||||||
|
"type": "VHS_LoadVideoPath",
|
||||||
|
"pos": [20, 170],
|
||||||
|
"size": [300, 310],
|
||||||
|
"flags": {},
|
||||||
|
"order": 0,
|
||||||
|
"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, 14]},
|
||||||
|
{"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": 2,
|
||||||
|
"type": "LoadSPEEDVFIModel",
|
||||||
|
"pos": [390, 150],
|
||||||
|
"size": [300, 105],
|
||||||
|
"flags": {},
|
||||||
|
"order": 1,
|
||||||
|
"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": 3,
|
||||||
|
"type": "VFIOptimizer",
|
||||||
|
"pos": [390, 320],
|
||||||
|
"size": [310, 150],
|
||||||
|
"flags": {},
|
||||||
|
"order": 2,
|
||||||
|
"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": 4,
|
||||||
|
"type": "SPEEDVFIInterpolate",
|
||||||
|
"pos": [760, 220],
|
||||||
|
"size": [355, 360],
|
||||||
|
"flags": {},
|
||||||
|
"order": 3,
|
||||||
|
"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},
|
||||||
|
{"name": "elapsed_seconds", "type": "FLOAT", "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": 5,
|
||||||
|
"type": "VHS_VideoCombine",
|
||||||
|
"pos": [1170, 220],
|
||||||
|
"size": [360, 340],
|
||||||
|
"flags": {},
|
||||||
|
"order": 4,
|
||||||
|
"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": 6,
|
||||||
|
"type": "LoadBIMVFIModel",
|
||||||
|
"pos": [390, 790],
|
||||||
|
"size": [300, 150],
|
||||||
|
"flags": {},
|
||||||
|
"order": 5,
|
||||||
|
"mode": 0,
|
||||||
|
"inputs": [],
|
||||||
|
"outputs": [
|
||||||
|
{"name": "model", "type": "BIM_VFI_MODEL", "links": [9, 11]}
|
||||||
|
],
|
||||||
|
"properties": {
|
||||||
|
"aux_id": "ComfyUI-Tween.git",
|
||||||
|
"Node name for S&R": "LoadBIMVFIModel"
|
||||||
|
},
|
||||||
|
"widgets_values": ["bim_vfi.pth", true, 3, false],
|
||||||
|
"color": "#28384a",
|
||||||
|
"bgcolor": "#36506b"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": 7,
|
||||||
|
"type": "VFIOptimizer",
|
||||||
|
"pos": [390, 970],
|
||||||
|
"size": [310, 150],
|
||||||
|
"flags": {},
|
||||||
|
"order": 6,
|
||||||
|
"mode": 0,
|
||||||
|
"inputs": [
|
||||||
|
{"name": "images", "type": "IMAGE", "link": 8},
|
||||||
|
{"name": "model", "type": "*", "link": 9}
|
||||||
|
],
|
||||||
|
"outputs": [
|
||||||
|
{"name": "images", "type": "IMAGE", "links": [10]},
|
||||||
|
{"name": "settings", "type": "VFI_SETTINGS", "links": [12]}
|
||||||
|
],
|
||||||
|
"properties": {
|
||||||
|
"aux_id": "ComfyUI-Tween.git",
|
||||||
|
"Node name for S&R": "VFIOptimizer"
|
||||||
|
},
|
||||||
|
"widgets_values": [2, 0],
|
||||||
|
"color": "#28384a",
|
||||||
|
"bgcolor": "#36506b"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": 8,
|
||||||
|
"type": "BIMVFIInterpolate",
|
||||||
|
"pos": [760, 860],
|
||||||
|
"size": [355, 340],
|
||||||
|
"flags": {},
|
||||||
|
"order": 7,
|
||||||
|
"mode": 0,
|
||||||
|
"inputs": [
|
||||||
|
{"name": "images", "type": "IMAGE", "link": 10},
|
||||||
|
{"name": "model", "type": "BIM_VFI_MODEL", "link": 11},
|
||||||
|
{"name": "settings", "shape": 7, "type": "VFI_SETTINGS", "link": 12}
|
||||||
|
],
|
||||||
|
"outputs": [
|
||||||
|
{"name": "images", "type": "IMAGE", "links": [13]},
|
||||||
|
{"name": "oversampled", "type": "IMAGE", "links": null},
|
||||||
|
{"name": "elapsed_seconds", "type": "FLOAT", "links": null}
|
||||||
|
],
|
||||||
|
"properties": {
|
||||||
|
"aux_id": "ComfyUI-Tween.git",
|
||||||
|
"Node name for S&R": "BIMVFIInterpolate"
|
||||||
|
},
|
||||||
|
"widgets_values": [2, 10, true, false, 1, 0, 24, 48],
|
||||||
|
"color": "#28384a",
|
||||||
|
"bgcolor": "#36506b"
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": 9,
|
||||||
|
"type": "VHS_VideoCombine",
|
||||||
|
"pos": [1170, 860],
|
||||||
|
"size": [360, 340],
|
||||||
|
"flags": {},
|
||||||
|
"order": 8,
|
||||||
|
"mode": 0,
|
||||||
|
"inputs": [
|
||||||
|
{"name": "images", "type": "IMAGE", "link": 13},
|
||||||
|
{"name": "audio", "shape": 7, "type": "AUDIO", "link": 14},
|
||||||
|
{"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_bim_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": "#28384a",
|
||||||
|
"bgcolor": "#36506b"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"links": [
|
||||||
|
[1, 1, 0, 3, 0, "IMAGE"],
|
||||||
|
[2, 2, 0, 3, 1, "*"],
|
||||||
|
[3, 3, 0, 4, 0, "IMAGE"],
|
||||||
|
[4, 2, 0, 4, 1, "SPEED_VFI_MODEL"],
|
||||||
|
[5, 3, 1, 4, 2, "VFI_SETTINGS"],
|
||||||
|
[6, 4, 0, 5, 0, "IMAGE"],
|
||||||
|
[7, 1, 2, 5, 1, "AUDIO"],
|
||||||
|
[8, 1, 0, 7, 0, "IMAGE"],
|
||||||
|
[9, 6, 0, 7, 1, "*"],
|
||||||
|
[10, 7, 0, 8, 0, "IMAGE"],
|
||||||
|
[11, 6, 0, 8, 1, "BIM_VFI_MODEL"],
|
||||||
|
[12, 7, 1, 8, 2, "VFI_SETTINGS"],
|
||||||
|
[13, 8, 0, 9, 0, "IMAGE"],
|
||||||
|
[14, 1, 2, 9, 1, "AUDIO"]
|
||||||
|
],
|
||||||
|
"groups": [
|
||||||
|
{
|
||||||
|
"id": 1,
|
||||||
|
"title": "INPUT · 24 FPS / 25 FRAMES",
|
||||||
|
"bounding": [-20, 100, 360, 470],
|
||||||
|
"color": "#625d52",
|
||||||
|
"font_size": 22,
|
||||||
|
"flags": {}
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": 2,
|
||||||
|
"title": "SPEED",
|
||||||
|
"bounding": [350, 100, 1240, 520],
|
||||||
|
"color": "#347a56",
|
||||||
|
"font_size": 22,
|
||||||
|
"flags": {}
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"id": 3,
|
||||||
|
"title": "BIM-VFI",
|
||||||
|
"bounding": [350, 740, 1240, 510],
|
||||||
|
"color": "#3f6488",
|
||||||
|
"font_size": 22,
|
||||||
|
"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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
+12
-5
@@ -17,12 +17,18 @@ logger = logging.getLogger("Tween")
|
|||||||
class BiMVFIModel:
|
class BiMVFIModel:
|
||||||
"""Clean inference wrapper around BiMVFI for ComfyUI integration."""
|
"""Clean inference wrapper around BiMVFI for ComfyUI integration."""
|
||||||
|
|
||||||
def __init__(self, checkpoint_path, pyr_level=3, auto_pyr_level=True, device="cpu"):
|
def __init__(self, checkpoint_path, pyr_level=3, auto_pyr_level=True,
|
||||||
|
artifact_safe_mode=False, device="cpu"):
|
||||||
self.pyr_level = pyr_level
|
self.pyr_level = pyr_level
|
||||||
self.auto_pyr_level = auto_pyr_level
|
self.auto_pyr_level = auto_pyr_level
|
||||||
|
self.artifact_safe_mode = artifact_safe_mode
|
||||||
self.device = device
|
self.device = device
|
||||||
|
|
||||||
self.model = BiMVFI(pyr_level=pyr_level, feat_channels=32)
|
self.model = BiMVFI(
|
||||||
|
pyr_level=pyr_level,
|
||||||
|
feat_channels=32,
|
||||||
|
artifact_safe_mode=artifact_safe_mode,
|
||||||
|
)
|
||||||
self._load_checkpoint(checkpoint_path)
|
self._load_checkpoint(checkpoint_path)
|
||||||
self.model.eval()
|
self.model.eval()
|
||||||
self.model.to(device)
|
self.model.to(device)
|
||||||
@@ -61,10 +67,11 @@ class BiMVFIModel:
|
|||||||
return 7
|
return 7
|
||||||
elif h >= 1080:
|
elif h >= 1080:
|
||||||
return 6
|
return 6
|
||||||
elif h >= 540:
|
|
||||||
return 5
|
|
||||||
else:
|
else:
|
||||||
return 3
|
# Match the official video inference path. Level 3 can miss
|
||||||
|
# large motion even at low resolutions; it remains available
|
||||||
|
# as a manual speed/quality tradeoff.
|
||||||
|
return 5
|
||||||
return self.pyr_level
|
return self.pyr_level
|
||||||
|
|
||||||
@torch.no_grad()
|
@torch.no_grad()
|
||||||
|
|||||||
-46
@@ -1,46 +0,0 @@
|
|||||||
import subprocess
|
|
||||||
import sys
|
|
||||||
import os
|
|
||||||
|
|
||||||
|
|
||||||
def get_cupy_package():
|
|
||||||
"""Detect PyTorch's CUDA version and return the matching cupy package name."""
|
|
||||||
try:
|
|
||||||
import torch
|
|
||||||
if not torch.cuda.is_available():
|
|
||||||
return None
|
|
||||||
cuda_version = torch.version.cuda
|
|
||||||
if cuda_version is None:
|
|
||||||
return None
|
|
||||||
major = int(cuda_version.split(".")[0])
|
|
||||||
cupy_pkg = f"cupy-cuda{major}x"
|
|
||||||
return cupy_pkg
|
|
||||||
except Exception:
|
|
||||||
return None
|
|
||||||
|
|
||||||
|
|
||||||
def install():
|
|
||||||
# Install core requirements first
|
|
||||||
requirements_path = os.path.join(os.path.dirname(__file__), "requirements.txt")
|
|
||||||
subprocess.check_call([
|
|
||||||
sys.executable, "-m", "pip", "install", "-r", requirements_path
|
|
||||||
])
|
|
||||||
|
|
||||||
# Try to install cupy for NVIDIA users (optional, improves performance)
|
|
||||||
cupy_pkg = get_cupy_package()
|
|
||||||
if cupy_pkg:
|
|
||||||
try:
|
|
||||||
subprocess.check_call([
|
|
||||||
sys.executable, "-m", "pip", "install", cupy_pkg
|
|
||||||
])
|
|
||||||
print(f"[Tween] cupy installed ({cupy_pkg}) — fast CUDA kernels enabled")
|
|
||||||
except subprocess.CalledProcessError:
|
|
||||||
print(f"[Tween] WARNING: Could not install {cupy_pkg}. "
|
|
||||||
f"BIM-VFI, SGM-VFI, and GIMM-VFI will use slower PyTorch fallback.")
|
|
||||||
else:
|
|
||||||
print("[Tween] cupy skipped (no NVIDIA CUDA). "
|
|
||||||
"BIM-VFI, SGM-VFI, and GIMM-VFI will use PyTorch fallback.")
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__":
|
|
||||||
install()
|
|
||||||
+481
@@ -0,0 +1,481 @@
|
|||||||
|
"""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)")
|
||||||
|
# LDF's custom Wan fork expects its complete condition embedder to use
|
||||||
|
# one dtype: time_embedder feeds time_proj directly without an explicit
|
||||||
|
# cast. Match the official generator's model.to(..., dtype=BF16) call;
|
||||||
|
# preserving diffusers' generic FP32-module policy leaves that pair as
|
||||||
|
# Float/BFloat16 and fails in the first sampling step.
|
||||||
|
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)"
|
||||||
|
)
|
||||||
|
if dense.ndim != 5 or dense_mask.ndim != 5:
|
||||||
|
raise RuntimeError(
|
||||||
|
"LDF-VFI decode conditions must use [batch, channels, time, height, width]"
|
||||||
|
)
|
||||||
|
|
||||||
|
temporal_tiles = latent.shape[1]
|
||||||
|
if dense.shape[2] % temporal_tiles:
|
||||||
|
raise RuntimeError(
|
||||||
|
f"LDF-VFI condition length {dense.shape[2]} is not divisible by "
|
||||||
|
f"the {temporal_tiles} decode tiles"
|
||||||
|
)
|
||||||
|
|
||||||
|
# Mirror the official generate.vae_decode adapter. The conditional VAE
|
||||||
|
# consumes one condition and mask tile per latent temporal tile, not a
|
||||||
|
# single continuous 5-D condition tensor.
|
||||||
|
decode_height = latent.shape[-2] * self.vae.spatial_compression_ratio
|
||||||
|
decode_width = latent.shape[-1] * self.vae.spatial_compression_ratio
|
||||||
|
pad_height = decode_height - dense.shape[-2]
|
||||||
|
pad_width = decode_width - dense.shape[-1]
|
||||||
|
if pad_height < 0 or pad_width < 0:
|
||||||
|
raise RuntimeError(
|
||||||
|
"LDF-VFI decoded latent is smaller than its conditioning frames; "
|
||||||
|
"check the VAE tile and overlap settings"
|
||||||
|
)
|
||||||
|
dense = F.pad(dense, (0, pad_width, 0, pad_height))
|
||||||
|
dense = rearrange(
|
||||||
|
dense, "b c (nt t) h w -> b nt c t h w", nt=temporal_tiles
|
||||||
|
)
|
||||||
|
|
||||||
|
dense_mask = dense_mask[..., 0, 0]
|
||||||
|
dense_mask = repeat(
|
||||||
|
dense_mask, "b c t -> b c t h w", h=decode_height, w=decode_width
|
||||||
|
)
|
||||||
|
dense_mask = rearrange(
|
||||||
|
dense_mask, "b c (nt t) h w -> b nt c t h w", nt=temporal_tiles
|
||||||
|
)
|
||||||
|
|
||||||
|
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
|
||||||
+9
-2
@@ -1,7 +1,7 @@
|
|||||||
[project]
|
[project]
|
||||||
name = "comfyui-tween"
|
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."
|
description = "Video frame interpolation nodes for ComfyUI using BIM-VFI, EMA-VFI, SGM-VFI, GIMM-VFI, SPEED, and LDF-VFI."
|
||||||
version = "1.1.0"
|
version = "1.2.2"
|
||||||
license = "Apache-2.0"
|
license = "Apache-2.0"
|
||||||
requires-python = ">=3.10"
|
requires-python = ">=3.10"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
@@ -12,6 +12,9 @@ dependencies = [
|
|||||||
"easydict",
|
"easydict",
|
||||||
"einops",
|
"einops",
|
||||||
"huggingface_hub",
|
"huggingface_hub",
|
||||||
|
"diffusers>=0.33.1,<0.40",
|
||||||
|
"accelerate>=1.5,<2",
|
||||||
|
"safetensors",
|
||||||
]
|
]
|
||||||
|
|
||||||
[project.urls]
|
[project.urls]
|
||||||
@@ -20,3 +23,7 @@ Repository = "https://github.com/Ethanfel/ComfyUI-Tween"
|
|||||||
[tool.comfy]
|
[tool.comfy]
|
||||||
PublisherId = "ethanfel"
|
PublisherId = "ethanfel"
|
||||||
DisplayName = "Tween - Video Frame Interpolation"
|
DisplayName = "Tween - Video Frame Interpolation"
|
||||||
|
|
||||||
|
[tool.pytest.ini_options]
|
||||||
|
testpaths = ["tests"]
|
||||||
|
addopts = "--confcutdir=tests"
|
||||||
|
|||||||
@@ -5,3 +5,6 @@ yacs
|
|||||||
easydict
|
easydict
|
||||||
einops
|
einops
|
||||||
huggingface_hub
|
huggingface_hub
|
||||||
|
diffusers>=0.33.1,<0.40
|
||||||
|
accelerate>=1.5,<2
|
||||||
|
safetensors
|
||||||
|
|||||||
@@ -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()
|
||||||
@@ -0,0 +1,82 @@
|
|||||||
|
import importlib
|
||||||
|
from pathlib import Path
|
||||||
|
import sys
|
||||||
|
import types
|
||||||
|
|
||||||
|
import torch
|
||||||
|
|
||||||
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
PACKAGE_NAME = "_tween_bim_tests"
|
||||||
|
|
||||||
|
|
||||||
|
def _load_inference_module():
|
||||||
|
package = types.ModuleType(PACKAGE_NAME)
|
||||||
|
package.__path__ = [str(REPO_ROOT)]
|
||||||
|
package.__package__ = PACKAGE_NAME
|
||||||
|
sys.modules.setdefault(PACKAGE_NAME, package)
|
||||||
|
return importlib.import_module(f"{PACKAGE_NAME}.inference")
|
||||||
|
|
||||||
|
|
||||||
|
def test_auto_pyramid_levels_match_official_video_inference():
|
||||||
|
model_class = _load_inference_module().BiMVFIModel
|
||||||
|
model = model_class.__new__(model_class)
|
||||||
|
model.auto_pyr_level = True
|
||||||
|
|
||||||
|
assert model._get_pyr_level(240) == 5
|
||||||
|
assert model._get_pyr_level(539) == 5
|
||||||
|
assert model._get_pyr_level(720) == 5
|
||||||
|
assert model._get_pyr_level(1079) == 5
|
||||||
|
assert model._get_pyr_level(1080) == 6
|
||||||
|
assert model._get_pyr_level(2159) == 6
|
||||||
|
assert model._get_pyr_level(2160) == 7
|
||||||
|
|
||||||
|
|
||||||
|
def test_manual_pyramid_level_remains_available():
|
||||||
|
model_class = _load_inference_module().BiMVFIModel
|
||||||
|
model = model_class.__new__(model_class)
|
||||||
|
model.auto_pyr_level = False
|
||||||
|
model.pyr_level = 3
|
||||||
|
|
||||||
|
assert model._get_pyr_level(720) == 3
|
||||||
|
|
||||||
|
|
||||||
|
def test_artifact_safe_mode_removes_only_rgb_refinement_residual():
|
||||||
|
module = _load_inference_module()
|
||||||
|
official = module.BiMVFI(pyr_level=3, feat_channels=1)
|
||||||
|
artifact_safe = module.BiMVFI(
|
||||||
|
pyr_level=3,
|
||||||
|
feat_channels=1,
|
||||||
|
artifact_safe_mode=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert official.sn.use_rgb_refine_residual is True
|
||||||
|
assert artifact_safe.sn.use_rgb_refine_residual is False
|
||||||
|
|
||||||
|
official_keys = official.state_dict().keys()
|
||||||
|
artifact_safe_keys = artifact_safe.state_dict().keys()
|
||||||
|
assert official_keys == artifact_safe_keys
|
||||||
|
|
||||||
|
warped0 = torch.full((1, 3, 2, 2), 0.2)
|
||||||
|
warped1 = torch.full((1, 3, 2, 2), 0.8)
|
||||||
|
mask = torch.full((1, 1, 2, 2), 0.25)
|
||||||
|
residual = torch.full((1, 3, 2, 2), 0.1)
|
||||||
|
blend = warped0 * mask + warped1 * (1 - mask)
|
||||||
|
|
||||||
|
assert torch.equal(
|
||||||
|
official.sn.merge_warped_images(warped0, warped1, mask, residual),
|
||||||
|
blend + residual,
|
||||||
|
)
|
||||||
|
assert torch.equal(
|
||||||
|
artifact_safe.sn.merge_warped_images(warped0, warped1, mask, residual),
|
||||||
|
blend,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def test_demo_exposes_artifact_safe_mode_without_enabling_it():
|
||||||
|
import json
|
||||||
|
|
||||||
|
workflow_path = REPO_ROOT / "example_workflows" / "tween_speed_bim_model_lab.json"
|
||||||
|
workflow = json.loads(workflow_path.read_text(encoding="utf-8"))
|
||||||
|
loader = next(node for node in workflow["nodes"] if node["type"] == "LoadBIMVFIModel")
|
||||||
|
|
||||||
|
assert loader["widgets_values"] == ["bim_vfi.pth", True, 3, False]
|
||||||
@@ -0,0 +1,53 @@
|
|||||||
|
import json
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
WORKFLOW_PATH = (
|
||||||
|
Path(__file__).resolve().parents[1]
|
||||||
|
/ "example_workflows"
|
||||||
|
/ "tween_speed_bim_model_lab.json"
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def _workflow():
|
||||||
|
return json.loads(WORKFLOW_PATH.read_text(encoding="utf-8"))
|
||||||
|
|
||||||
|
|
||||||
|
def test_speed_bim_demo_is_clean_and_current():
|
||||||
|
workflow = _workflow()
|
||||||
|
node_types = [node["type"] for node in workflow["nodes"]]
|
||||||
|
|
||||||
|
assert "Note" not in node_types
|
||||||
|
assert not any("LDF" in node_type for node_type in node_types)
|
||||||
|
assert node_types.count("SPEEDVFIInterpolate") == 1
|
||||||
|
assert node_types.count("BIMVFIInterpolate") == 1
|
||||||
|
assert [group["title"] for group in workflow["groups"]] == [
|
||||||
|
"INPUT · 24 FPS / 25 FRAMES",
|
||||||
|
"SPEED",
|
||||||
|
"BIM-VFI",
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_speed_bim_demo_links_are_internally_consistent():
|
||||||
|
workflow = _workflow()
|
||||||
|
nodes = {node["id"]: node for node in workflow["nodes"]}
|
||||||
|
node_ids = set(nodes)
|
||||||
|
links = {link[0]: link for link in workflow["links"]}
|
||||||
|
|
||||||
|
assert workflow["last_node_id"] == max(node_ids)
|
||||||
|
assert workflow["last_link_id"] == max(links)
|
||||||
|
|
||||||
|
for link_id, source_id, source_slot, target_id, target_slot, _ in workflow["links"]:
|
||||||
|
assert link_id in links
|
||||||
|
assert source_id in node_ids
|
||||||
|
assert target_id in node_ids
|
||||||
|
assert link_id in nodes[source_id]["outputs"][source_slot]["links"]
|
||||||
|
assert nodes[target_id]["inputs"][target_slot]["link"] == link_id
|
||||||
|
|
||||||
|
for node in workflow["nodes"]:
|
||||||
|
for input_slot in node.get("inputs", []):
|
||||||
|
if input_slot.get("link") is not None:
|
||||||
|
assert input_slot["link"] in links
|
||||||
|
for output_slot in node.get("outputs", []):
|
||||||
|
for link_id in output_slot.get("links") or []:
|
||||||
|
assert link_id in links
|
||||||
@@ -0,0 +1,26 @@
|
|||||||
|
from pathlib import Path
|
||||||
|
import tomllib
|
||||||
|
|
||||||
|
|
||||||
|
REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||||
|
|
||||||
|
|
||||||
|
def _requirements():
|
||||||
|
return [
|
||||||
|
line.strip()
|
||||||
|
for line in (REPO_ROOT / "requirements.txt").read_text().splitlines()
|
||||||
|
if line.strip() and not line.lstrip().startswith("#")
|
||||||
|
]
|
||||||
|
|
||||||
|
|
||||||
|
def test_comfy_manager_has_no_redundant_install_script():
|
||||||
|
assert not (REPO_ROOT / "install.py").exists()
|
||||||
|
|
||||||
|
|
||||||
|
def test_declared_dependencies_stay_aligned_and_exclude_optional_cupy():
|
||||||
|
with (REPO_ROOT / "pyproject.toml").open("rb") as file:
|
||||||
|
project_dependencies = tomllib.load(file)["project"]["dependencies"]
|
||||||
|
|
||||||
|
requirements = _requirements()
|
||||||
|
assert requirements == project_dependencies
|
||||||
|
assert not any("cupy" in dependency.lower() for dependency in requirements)
|
||||||
@@ -0,0 +1,140 @@
|
|||||||
|
import torch
|
||||||
|
import pytest
|
||||||
|
from types import SimpleNamespace
|
||||||
|
|
||||||
|
from ldf_backend import LDFVFIModel
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingConditionalVAE:
|
||||||
|
spatial_compression_ratio = 8
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.decode_shapes = None
|
||||||
|
|
||||||
|
def decode(self, latent, condition, mask):
|
||||||
|
self.decode_shapes = (
|
||||||
|
tuple(latent.shape),
|
||||||
|
tuple(condition.shape),
|
||||||
|
tuple(mask.shape),
|
||||||
|
)
|
||||||
|
# A small real result lets _decode finish without materializing the
|
||||||
|
# full 720p tensors used for the shape-only inputs above.
|
||||||
|
return torch.zeros(1, 3, 40, 1, 1)
|
||||||
|
|
||||||
|
|
||||||
|
class _RecordingTransformer:
|
||||||
|
def __init__(self):
|
||||||
|
self.to_kwargs = None
|
||||||
|
|
||||||
|
def to(self, **kwargs):
|
||||||
|
self.to_kwargs = kwargs
|
||||||
|
return self
|
||||||
|
|
||||||
|
|
||||||
|
class _ShapeOnlyLDF(LDFVFIModel):
|
||||||
|
"""Exercise sequence chunking without loading the multi-GB checkpoint."""
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
self.decode_tile_counts = []
|
||||||
|
|
||||||
|
def _prepare_condition(self, frames, mask, device):
|
||||||
|
assert frames.shape[0] == int(mask.sum())
|
||||||
|
dense = torch.empty(1, 3, self.TRAIN_FRAMES, 1, 1)
|
||||||
|
dense_mask = torch.empty(1, 1, self.TRAIN_FRAMES, 1, 1)
|
||||||
|
latent = torch.empty(
|
||||||
|
1, self.TRAIN_FRAMES // self.TILE_TIME, 1, 1, 1, 1, 1, 1
|
||||||
|
)
|
||||||
|
return dense, dense_mask, latent, latent
|
||||||
|
|
||||||
|
def _sample_free(self, condition, encoded_mask, schedule, device, progress):
|
||||||
|
return condition
|
||||||
|
|
||||||
|
def _sample_between(
|
||||||
|
self, previous, following, condition, encoded_mask,
|
||||||
|
schedule, t_cond, device, progress,
|
||||||
|
):
|
||||||
|
return condition[:, self.CONDITION_TILES:-self.CONDITION_TILES]
|
||||||
|
|
||||||
|
def _sample_tail(
|
||||||
|
self, previous, condition, encoded_mask, schedule,
|
||||||
|
t_cond, device, progress,
|
||||||
|
):
|
||||||
|
return condition[:, self.CONDITION_TILES:]
|
||||||
|
|
||||||
|
def _decode(self, latent, dense, dense_mask, height, width):
|
||||||
|
temporal_tiles = latent.shape[1]
|
||||||
|
assert dense.shape[2] == temporal_tiles * self.TILE_TIME
|
||||||
|
assert dense_mask.shape[2] == dense.shape[2]
|
||||||
|
self.decode_tile_counts.append(temporal_tiles)
|
||||||
|
return torch.empty(dense.shape[2], 3, height, width)
|
||||||
|
|
||||||
|
|
||||||
|
def test_decode_tiles_40_frame_720p_condition_for_conditional_vae():
|
||||||
|
"""Regression for the 5-D condition passed to upstream VAE.decode."""
|
||||||
|
model = LDFVFIModel.__new__(LDFVFIModel)
|
||||||
|
model.vae = _RecordingConditionalVAE()
|
||||||
|
|
||||||
|
# Representative first LDF output block at 720p. Meta tensors exercise
|
||||||
|
# exact shape transforms without allocating hundreds of MB in the test.
|
||||||
|
latent = torch.empty((1, 2, 4, 7, 4, 5, 24, 24), device="meta")
|
||||||
|
condition = torch.empty((1, 3, 40, 720, 1280), device="meta")
|
||||||
|
mask = torch.empty((1, 1, 40, 720, 1280), device="meta")
|
||||||
|
|
||||||
|
result = model._decode(latent, condition, mask, height=720, width=1280)
|
||||||
|
|
||||||
|
assert model.vae.decode_shapes == (
|
||||||
|
(1, 2, 4, 5, 96, 168),
|
||||||
|
(1, 2, 3, 20, 768, 1344),
|
||||||
|
(1, 2, 1, 20, 768, 1344),
|
||||||
|
)
|
||||||
|
assert result.shape == (40, 3, 1, 1)
|
||||||
|
|
||||||
|
|
||||||
|
def test_device_move_casts_complete_ldf_transformer_to_bfloat16():
|
||||||
|
model = LDFVFIModel.__new__(LDFVFIModel)
|
||||||
|
transformer = _RecordingTransformer()
|
||||||
|
model.model = SimpleNamespace(transformer=transformer)
|
||||||
|
model.dtype = torch.bfloat16
|
||||||
|
model._move_auxiliary_models = lambda device: None
|
||||||
|
|
||||||
|
model.to("cpu")
|
||||||
|
|
||||||
|
assert transformer.to_kwargs == {
|
||||||
|
"device": torch.device("cpu"),
|
||||||
|
"dtype": torch.bfloat16,
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def test_decode_rejects_condition_length_that_cannot_tile():
|
||||||
|
model = LDFVFIModel.__new__(LDFVFIModel)
|
||||||
|
model.vae = _RecordingConditionalVAE()
|
||||||
|
latent = torch.empty((1, 2, 1, 1, 4, 5, 1, 1), device="meta")
|
||||||
|
condition = torch.empty((1, 3, 39, 8, 8), device="meta")
|
||||||
|
mask = torch.empty((1, 1, 39, 8, 8), device="meta")
|
||||||
|
|
||||||
|
try:
|
||||||
|
model._decode(latent, condition, mask, height=8, width=8)
|
||||||
|
except RuntimeError as error:
|
||||||
|
assert "not divisible" in str(error)
|
||||||
|
else:
|
||||||
|
raise AssertionError("Expected an invalid temporal tile error")
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.parametrize("temporal_factor", [2, 3, 8, 16])
|
||||||
|
def test_sequence_chunk_conditions_align_for_reported_391_frames(temporal_factor):
|
||||||
|
model = _ShapeOnlyLDF()
|
||||||
|
source = torch.empty(391, 3, 1, 1)
|
||||||
|
chunks = model._interpolate_sequence_impl(
|
||||||
|
source=source,
|
||||||
|
factor=temporal_factor,
|
||||||
|
schedule=torch.tensor([1.0, 0.0]),
|
||||||
|
t_cond=0.1,
|
||||||
|
device=torch.device("cpu"),
|
||||||
|
progress=lambda: None,
|
||||||
|
)
|
||||||
|
|
||||||
|
expected_frames = (source.shape[0] - 1) * temporal_factor + 1
|
||||||
|
assert sum(chunk.shape[0] for chunk in chunks) >= expected_frames
|
||||||
|
assert len(chunks) == model.sampling_block_count(
|
||||||
|
source.shape[0], temporal_factor
|
||||||
|
)
|
||||||
@@ -0,0 +1,69 @@
|
|||||||
|
import ast
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
NODE_SOURCE = Path(__file__).resolve().parents[1] / "nodes.py"
|
||||||
|
TIMED_CLASSES = {
|
||||||
|
"BIMVFIInterpolate",
|
||||||
|
"BIMVFISegmentInterpolate",
|
||||||
|
"EMAVFIInterpolate",
|
||||||
|
"EMAVFISegmentInterpolate",
|
||||||
|
"SGMVFIInterpolate",
|
||||||
|
"SGMVFISegmentInterpolate",
|
||||||
|
"LDFVFIInterpolate",
|
||||||
|
"SPEEDVFIInterpolate",
|
||||||
|
"SPEEDVFISegmentInterpolate",
|
||||||
|
"GIMMVFIInterpolate",
|
||||||
|
"GIMMVFISegmentInterpolate",
|
||||||
|
}
|
||||||
|
DIRECTLY_DECORATED = TIMED_CLASSES - {
|
||||||
|
"SPEEDVFIInterpolate",
|
||||||
|
"SPEEDVFISegmentInterpolate",
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _class_definitions():
|
||||||
|
tree = ast.parse(NODE_SOURCE.read_text(encoding="utf-8"))
|
||||||
|
return {
|
||||||
|
node.name: node
|
||||||
|
for node in tree.body
|
||||||
|
if isinstance(node, ast.ClassDef) and node.name in TIMED_CLASSES
|
||||||
|
}
|
||||||
|
|
||||||
|
|
||||||
|
def _literal_assignment(class_node, name):
|
||||||
|
for statement in class_node.body:
|
||||||
|
if (
|
||||||
|
isinstance(statement, ast.Assign)
|
||||||
|
and len(statement.targets) == 1
|
||||||
|
and isinstance(statement.targets[0], ast.Name)
|
||||||
|
and statement.targets[0].id == name
|
||||||
|
):
|
||||||
|
return ast.literal_eval(statement.value)
|
||||||
|
raise AssertionError(f"{class_node.name} does not define {name}")
|
||||||
|
|
||||||
|
|
||||||
|
def test_all_interpolation_nodes_expose_elapsed_seconds_last():
|
||||||
|
classes = _class_definitions()
|
||||||
|
assert classes.keys() == TIMED_CLASSES
|
||||||
|
|
||||||
|
for class_node in classes.values():
|
||||||
|
assert _literal_assignment(class_node, "RETURN_TYPES")[-1] == "FLOAT"
|
||||||
|
assert _literal_assignment(class_node, "RETURN_NAMES")[-1] == "elapsed_seconds"
|
||||||
|
|
||||||
|
|
||||||
|
def test_direct_interpolation_methods_append_timing_output():
|
||||||
|
classes = _class_definitions()
|
||||||
|
for class_name in DIRECTLY_DECORATED:
|
||||||
|
interpolate = next(
|
||||||
|
statement
|
||||||
|
for statement in classes[class_name].body
|
||||||
|
if isinstance(statement, ast.FunctionDef)
|
||||||
|
and statement.name == "interpolate"
|
||||||
|
)
|
||||||
|
assert any(
|
||||||
|
isinstance(decorator, ast.Call)
|
||||||
|
and isinstance(decorator.func, ast.Name)
|
||||||
|
and decorator.func.id == "_with_elapsed_seconds"
|
||||||
|
for decorator in interpolate.decorator_list
|
||||||
|
)
|
||||||
Reference in New Issue
Block a user