first comfy-manager-commit of betterNormalCrafterWrapper
This commit is contained in:
Korbinian
@@ -0,0 +1,20 @@
|
||||
name: Publish to Comfy registry
|
||||
on:
|
||||
workflow_dispatch:
|
||||
push:
|
||||
branches:
|
||||
- main
|
||||
paths:
|
||||
- "pyproject.toml"
|
||||
|
||||
jobs:
|
||||
publish-node:
|
||||
name: Publish Custom Node to registry
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Check out code
|
||||
uses: actions/checkout@v7
|
||||
- name: Publish Custom Node
|
||||
uses: Comfy-Org/publish-node-action@main
|
||||
with:
|
||||
personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} ## Add your own personal access token to your Github Repository secrets and reference it here.
|
||||
-155
@@ -1,155 +0,0 @@
|
||||
# NormalCrafter Clean — Architecture Notes
|
||||
|
||||
This document explains the runtime architecture of the commented implementation. The executable behavior is unchanged from the validated clean build; the added material is documentation only.
|
||||
|
||||
## 1. Design goal
|
||||
|
||||
The implementation is built around one ownership rule:
|
||||
|
||||
> Full-video intermediate tensors live on CPU. CUDA only receives the component and chunk/window currently being processed.
|
||||
|
||||
This prevents ComfyUI's graph cache from retaining a complete CUDA output and prevents total clip length from directly determining peak VRAM.
|
||||
|
||||
## 2. File map
|
||||
|
||||
- `nodes.py` — thin ComfyUI integration and UI declarations.
|
||||
- `normalcrafter_clean/engine.py` — model loading, device lifecycle, and the four inference stages.
|
||||
- `normalcrafter_clean/preprocess.py` — lazy frame resizing, temporal tail repetition, and spatial padding.
|
||||
- `normalcrafter_clean/unet.py` — inference-only SVD UNet adaptation for one CLIP embedding per frame.
|
||||
- `normalcrafter_clean/windows.py` — pure deterministic temporal window scheduler.
|
||||
- `tests/test_windows.py` — unit tests for coverage and overlap behavior.
|
||||
|
||||
## 3. End-to-end flow
|
||||
|
||||
```text
|
||||
ComfyUI IMAGE [F,H,W,C], CPU or GPU
|
||||
│
|
||||
▼
|
||||
FrameSource
|
||||
- RGB only
|
||||
- detached CPU FP32
|
||||
- aspect-preserving resize metadata
|
||||
- symmetric spatial padding metadata
|
||||
│
|
||||
├──────────────────────────────────────────────┐
|
||||
│ │
|
||||
▼ ▼
|
||||
CLIP encoder VAE encoder
|
||||
chunked RGB frames chunked RGB frames
|
||||
│ │
|
||||
▼ ▼
|
||||
CPU CLIP embeddings CPU RGB latents
|
||||
│ │
|
||||
└──────────────────────┬───────────────────────┘
|
||||
▼
|
||||
Temporal UNet
|
||||
one overlapping window at a time
|
||||
│
|
||||
▼
|
||||
CPU normal latents
|
||||
│
|
||||
▼
|
||||
VAE decoder
|
||||
chunked decode
|
||||
│
|
||||
▼
|
||||
CPU ComfyUI IMAGE [F,H,W,3]
|
||||
```
|
||||
|
||||
## 4. Model components
|
||||
|
||||
`NormalCrafterModel` owns each heavy module exactly once:
|
||||
|
||||
- `image_encoder`: CLIP vision encoder from Stable Video Diffusion.
|
||||
- `vae`: NormalCrafter's fine-tuned temporal VAE.
|
||||
- `unet`: NormalCrafter's spatio-temporal UNet.
|
||||
- `scheduler`: Euler scheduler used for the released one-step inference path.
|
||||
|
||||
There is no global pipeline variable and no second wrapper object containing duplicate strong references.
|
||||
|
||||
## 5. Device modes
|
||||
|
||||
### `staged`
|
||||
|
||||
Only the active heavy component is on the inference device:
|
||||
|
||||
1. CLIP on CUDA, VAE and UNet on CPU.
|
||||
2. VAE on CUDA, CLIP and UNet on CPU.
|
||||
3. UNet on CUDA, CLIP and VAE on CPU.
|
||||
4. VAE on CUDA, CLIP and UNet on CPU.
|
||||
|
||||
This minimizes VRAM but pays for several model transfers.
|
||||
|
||||
### `resident`
|
||||
|
||||
All heavy components are moved to the selected device and remain there during inference. This avoids model transfer overhead and is faster when VRAM permits it.
|
||||
|
||||
## 6. Temporal behavior
|
||||
|
||||
The released defaults are:
|
||||
|
||||
- window size: 14 frames
|
||||
- step size: 10 frames
|
||||
- regular overlap: 4 frames
|
||||
|
||||
Long clips are processed by fixed-size windows. The last window is anchored to the end of the clip so no tail frame is omitted. If that anchor is close to the previous window, the final overlap may be larger than four frames.
|
||||
|
||||
The next window starts with the previous prediction in its overlapping prefix. After inference, old and new overlap predictions are linearly crossfaded.
|
||||
|
||||
Only the immediately previous GPU prediction window is retained. Temporal GPU state is therefore bounded by `window_size`, not total clip length.
|
||||
|
||||
## 7. Short clips
|
||||
|
||||
A clip shorter than one temporal window is logically extended by repeating its final frame. This happens lazily in `FrameSource.get_bchw`; no permanently padded video tensor is created.
|
||||
|
||||
Only real frames are decoded and returned.
|
||||
|
||||
## 8. Spatial behavior
|
||||
|
||||
Input is never enlarged. If the longest side exceeds `max_resolution`, the clip is downscaled while preserving aspect ratio.
|
||||
|
||||
The resized image is symmetrically padded to a multiple of 64 for the VAE/UNet hierarchy. Padding is removed after decoding. With `output_size="original"`, the decoded normal map is then resized back to the original input dimensions.
|
||||
|
||||
Because interpolation blends XYZ components, optional normal renormalization restores unit vector length before mapping values from `[-1,1]` to `[0,1]`.
|
||||
|
||||
## 9. One-step inference
|
||||
|
||||
The runtime path is deterministic:
|
||||
|
||||
1. Normal latents start at zero.
|
||||
2. Existing overlap predictions seed the next window prefix.
|
||||
3. RGB latents are concatenated channel-wise with the normal state.
|
||||
4. The UNet evaluates one scheduler timestep.
|
||||
5. Euler produces the predicted normal latent.
|
||||
|
||||
There is intentionally no seed input because the published runtime does not start from random noise.
|
||||
|
||||
## 10. Memory invariants
|
||||
|
||||
The implementation relies on these invariants:
|
||||
|
||||
- ComfyUI receives a CPU output tensor.
|
||||
- Full CLIP embeddings are CPU tensors.
|
||||
- Full RGB latent video is a CPU tensor.
|
||||
- Full normal latent video is a CPU tensor.
|
||||
- Only one encoder/decoder chunk or UNet window is transferred to CUDA at once.
|
||||
- Output tensors are preallocated rather than grown through repeated `torch.cat`.
|
||||
- Temporary local tensors are explicitly dereferenced after each stage iteration.
|
||||
- Exceptions still execute model offloading through `finally`.
|
||||
|
||||
`torch.cuda.empty_cache()` can release allocator-reserved blocks only after live references disappear. It is cleanup support, not the mechanism that fixes ownership.
|
||||
|
||||
## 11. Performance knobs
|
||||
|
||||
- `offload_mode="resident"` is the largest speed gain when VRAM is sufficient.
|
||||
- Larger CLIP/VAE chunk sizes reduce transfer and launch overhead but increase peak VRAM.
|
||||
- A larger `step_size` creates less overlap and fewer UNet windows, but changes temporal behavior and must never exceed `window_size`.
|
||||
- Lower `max_resolution` reduces both VAE and UNet cost significantly.
|
||||
|
||||
The safest optimization strategy is to alter one knob at a time while measuring output consistency, runtime, `torch.cuda.max_memory_allocated()`, and repeated-run stability.
|
||||
|
||||
## 12. Concurrency
|
||||
|
||||
One `NormalCrafterModel` instance has a re-entrant lock around inference and device changes. Two graph executions cannot simultaneously move the same modules between CPU and CUDA or mutate the shared scheduler.
|
||||
|
||||
Parallelism should use separate model instances only when the hardware has enough memory for them.
|
||||
@@ -1,15 +0,0 @@
|
||||
ComfyUI NormalCrafter Clean
|
||||
|
||||
This project contains an inference-only adaptation of ideas and small portions of
|
||||
code structure from:
|
||||
|
||||
1. NormalCrafter by Yanrui Bin et al.
|
||||
https://github.com/Binyr/NormalCrafter
|
||||
Licensed under the MIT License.
|
||||
|
||||
2. Hugging Face Diffusers, including the Stable Video Diffusion pipeline and
|
||||
UNetSpatioTemporalConditionModel execution structure.
|
||||
https://github.com/huggingface/diffusers
|
||||
Licensed under the Apache License, Version 2.0.
|
||||
|
||||
No model weights are redistributed by this package.
|
||||
@@ -1,85 +1,86 @@
|
||||
# ComfyUI NormalCrafter Clean
|
||||
# Better NormalCrafter Wrapper
|
||||
|
||||
A ground-up ComfyUI inference implementation for **NormalCrafter**. It uses the released NormalCrafter UNet/VAE weights and the Stable Video Diffusion image encoder/scheduler, but does not reuse the research pipeline or the existing ComfyUI wrapper.
|
||||
|
||||
## Why this exists
|
||||
|
||||
The existing wrapper mixes node UI, downloads, global pipeline caching, device state, inference, retry logic, resizing, and cleanup. It can also return live CUDA tensors into ComfyUI's cache. This implementation makes ownership explicit:
|
||||
|
||||
- one model object owns each component exactly once;
|
||||
- no module-level model globals;
|
||||
- ComfyUI receives a CPU tensor;
|
||||
- CLIP embeddings, RGB latents, and predicted normal latents are stored on CPU between stages;
|
||||
- only the active chunk/window is moved to CUDA;
|
||||
- the output latent video is preallocated instead of repeatedly grown with `torch.cat`;
|
||||
- short clips are actually padded and then trimmed;
|
||||
- cleanup runs through a single lifecycle path.
|
||||
A small ComfyUI wrapper around NormalCrafter with two nodes, explicit VRAM controls, and staged GPU inference.
|
||||
|
||||
## Nodes
|
||||
|
||||
### NormalCrafter Clean - Load
|
||||
### NormalCrafter - Load
|
||||
|
||||
Loads:
|
||||
Selects two model sources through dropdowns:
|
||||
|
||||
- `Yanrui95/NormalCrafter`: NormalCrafter UNet and normal VAE;
|
||||
- `stabilityai/stable-video-diffusion-img2vid-xt`: CLIP vision encoder and Euler scheduler.
|
||||
- **NormalCrafter model** — UNet + VAE
|
||||
- **Base model** — CLIP image encoder + feature extractor + scheduler
|
||||
|
||||
The model is initially kept on CPU. `auto` dtype selects FP16 on CUDA and FP32 on CPU.
|
||||
The wrapper registers and uses ComfyUI's model category `normalcrafter`. It does **not** rely on one hardcoded filesystem root. Model discovery goes through:
|
||||
|
||||
### NormalCrafter Clean - Generate Normals
|
||||
|
||||
`staged` offload mode runs the model in four explicit phases:
|
||||
|
||||
1. CLIP frame embeddings;
|
||||
2. RGB VAE encoding;
|
||||
3. sliding-window UNet inference;
|
||||
4. normal VAE decoding.
|
||||
|
||||
Only one heavy component is resident on CUDA at a time. `resident` keeps all components on the selected device and is faster, but needs substantially more VRAM.
|
||||
|
||||
The released model uses a deterministic one-step scheduler path. There is therefore intentionally no seed input.
|
||||
|
||||
### NormalCrafter Clean - Offload
|
||||
|
||||
Moves every model component to CPU and clears the CUDA allocator cache.
|
||||
|
||||
## Installation
|
||||
|
||||
Place the folder in `ComfyUI/custom_nodes/` and install dependencies inside the same Python environment as ComfyUI:
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```python
|
||||
folder_paths.get_folder_paths("normalcrafter")
|
||||
```
|
||||
|
||||
Restart ComfyUI afterward. The first loader execution downloads the Hugging Face weights unless `local_files_only` is enabled.
|
||||
This means the normal ComfyUI location and any `normalcrafter` directories configured through `extra_model_paths.yaml` are searched together.
|
||||
|
||||
## Defaults
|
||||
The standard location is registered as:
|
||||
|
||||
- window size: 14 frames
|
||||
- step size: 10 frames
|
||||
- temporal overlap: 4 frames in the regular case
|
||||
- conditioning FPS: 7
|
||||
- motion bucket: 127
|
||||
- noise augmentation: 0
|
||||
- output: original input resolution
|
||||
```text
|
||||
<ComfyUI models dir>/normalcrafter/
|
||||
```
|
||||
|
||||
The final normal vectors are renormalized after resizing by default. Disable this only for strict comparisons with older glue code.
|
||||
A corresponding extra path can be configured like this:
|
||||
|
||||
## Validation status
|
||||
```yaml
|
||||
my_models:
|
||||
base_path: D:/AI/
|
||||
normalcrafter: models/normalcrafter/
|
||||
```
|
||||
|
||||
The source is syntax-checked and the pure window scheduler has unit tests. It has **not** been executed here against the multi-gigabyte model weights on an NVIDIA GPU. The first real validation should compare a fixed clip against the research repository and record both visual output and `torch.cuda.memory_allocated()` across repeated runs.
|
||||
Every compatible direct child folder found below any registered `normalcrafter` root appears in the dropdown.
|
||||
|
||||
## Attribution
|
||||
A local NormalCrafter repo is detected by:
|
||||
|
||||
NormalCrafter research code: MIT License. Hugging Face Diffusers-derived UNet execution structure: Apache License 2.0. See `NOTICE`.
|
||||
```text
|
||||
unet/
|
||||
vae/
|
||||
```
|
||||
|
||||
## Commented edition
|
||||
A local SVD base repo is detected by:
|
||||
|
||||
This archive is the documentation-heavy edition of the clean implementation.
|
||||
Executable statements were checked for AST equivalence against the working build;
|
||||
only comments, docstrings, and `ARCHITECTURE.md` were added. Start with:
|
||||
```text
|
||||
image_encoder/
|
||||
feature_extractor/
|
||||
scheduler/
|
||||
```
|
||||
|
||||
1. `ARCHITECTURE.md` for the complete runtime and ownership model;
|
||||
2. `nodes.py` for the ComfyUI-facing surface;
|
||||
3. `normalcrafter_clean/engine.py` for the inference pipeline;
|
||||
4. `preprocess.py`, `windows.py`, and `unet.py` for the isolated subsystems.
|
||||
The official Hugging Face sources remain available in the dropdown. If selected, only the required subfolders are downloaded. Existing copies are searched across all registered `normalcrafter` roots first; otherwise the first writable root in ComfyUI's path order is used.
|
||||
|
||||
### NormalCrafter - Generate Normals
|
||||
|
||||
Inputs:
|
||||
|
||||
- `max_resolution` — maximum long edge processed by the model
|
||||
- `window_size` — temporal UNet context size; official default is 14
|
||||
- `step_size` — distance between temporal windows; official default is 10
|
||||
- `chunk_size` — shared CLIP/VAE encode/decode batch size; lower values use less VRAM
|
||||
|
||||
`step_size` must be less than or equal to `window_size`.
|
||||
|
||||
## Memory model
|
||||
|
||||
Inference is deliberately staged:
|
||||
|
||||
1. CLIP encoder on GPU
|
||||
2. VAE encoder on GPU
|
||||
3. temporal UNet on GPU
|
||||
4. VAE decoder on GPU
|
||||
|
||||
Only the active heavy model is GPU-resident. Full-video CLIP embeddings and latent videos stay on CPU between stages. After every inference call all heavy modules are moved back to CPU automatically; there is no separate offload node.
|
||||
|
||||
`window_size` mainly controls temporal-UNet VRAM. `chunk_size` mainly controls CLIP/VAE VRAM. `max_resolution` affects all stages strongly.
|
||||
|
||||
## Structure
|
||||
|
||||
```text
|
||||
nodes.py # ComfyUI surface: exactly two nodes
|
||||
normalcrafter_clean/model.py # ComfyUI path discovery + loading + GPU residency
|
||||
normalcrafter_clean/inference.py # preprocessing + four-stage inference
|
||||
normalcrafter_clean/unet.py # NormalCrafter-specific SVD UNet forward
|
||||
```
|
||||
|
||||
+1
-7
@@ -1,9 +1,3 @@
|
||||
# ComfyUI imports custom-node folders as packages. Keeping direct-file imports
|
||||
# inert also allows the pure utility tests to run outside a ComfyUI checkout.
|
||||
if __package__:
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
else: # pragma: no cover - only used by standalone test collection
|
||||
NODE_CLASS_MAPPINGS = {}
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {}
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
|
||||
@@ -1,251 +1,219 @@
|
||||
"""ComfyUI-facing nodes for the clean NormalCrafter implementation.
|
||||
|
||||
This module deliberately contains only the thin UI/integration layer. The actual
|
||||
model lifecycle, inference stages, windowing, preprocessing, and tensor ownership
|
||||
live in ``normalcrafter_clean``.
|
||||
|
||||
That separation is important: ComfyUI may cache node outputs and keep node
|
||||
instances alive longer than one execution. Keeping model logic out of this file
|
||||
reduces the risk of accidental global state, duplicate references, or live CUDA
|
||||
tensors being retained by the graph cache.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
import torch
|
||||
"""The two ComfyUI nodes exposed by Better NormalCrafter Wrapper."""
|
||||
|
||||
import comfy.model_management
|
||||
import comfy.utils
|
||||
|
||||
from .normalcrafter_clean.engine import InferenceConfig, NormalCrafterModel
|
||||
from .normalcrafter_clean.windows import build_windows
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
from .normalcrafter_clean.inference import estimate_progress_steps, run_inference
|
||||
from .normalcrafter_clean.model import (
|
||||
base_model_sources,
|
||||
load_models,
|
||||
normalcrafter_sources,
|
||||
)
|
||||
|
||||
|
||||
def _resolve_dtype(name: str, device: torch.device) -> torch.dtype:
|
||||
"""Translate the UI dtype choice into a concrete PyTorch dtype.
|
||||
|
||||
``auto`` uses FP16 on CUDA because that is the practical default for the
|
||||
released model weights. CPU inference is forced to FP32: many PyTorch CPU
|
||||
kernels either do not implement FP16 efficiently or are numerically fragile.
|
||||
"""
|
||||
|
||||
if name == "auto":
|
||||
return torch.float16 if device.type == "cuda" else torch.float32
|
||||
|
||||
mapping = {
|
||||
"float16": torch.float16,
|
||||
"bfloat16": torch.bfloat16,
|
||||
"float32": torch.float32,
|
||||
}
|
||||
dtype = mapping[name]
|
||||
|
||||
if device.type == "cpu" and dtype != torch.float32:
|
||||
logger.warning("CPU inference requested; forcing float32 instead of %s", name)
|
||||
return torch.float32
|
||||
|
||||
return dtype
|
||||
|
||||
|
||||
class NormalCrafterCleanLoader:
|
||||
"""Load all NormalCrafter components into one explicitly-owned model object.
|
||||
|
||||
The loader does not move the model to CUDA. ``NormalCrafterModel`` is created
|
||||
on CPU and the generate node later decides whether to stage components one by
|
||||
one or keep all of them resident on the selected inference device.
|
||||
"""
|
||||
class NormalCrafterLoader:
|
||||
"""Choose and load the NormalCrafter and SVD model components."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
# ComfyUI reads this declarative structure to construct the node UI.
|
||||
return {
|
||||
"required": {
|
||||
"dtype": (["auto", "float16", "bfloat16", "float32"], {"default": "auto"}),
|
||||
"attention": (["auto", "default", "xformers"], {"default": "auto"}),
|
||||
"local_files_only": ("BOOLEAN", {"default": False}),
|
||||
},
|
||||
"optional": {
|
||||
# These fields are configurable for mirrors, forks, or local repos,
|
||||
# while preserving the official public defaults.
|
||||
"model_repo": ("STRING", {"default": "Yanrui95/NormalCrafter"}),
|
||||
"base_repo": (
|
||||
"STRING",
|
||||
{"default": "stabilityai/stable-video-diffusion-img2vid-xt"},
|
||||
"normalcrafter_model": (
|
||||
normalcrafter_sources(),
|
||||
{
|
||||
"default": "HF: Yanrui95/NormalCrafter",
|
||||
"tooltip": (
|
||||
"NormalCrafter model containing the UNet and VAE. "
|
||||
"Local models are discovered from every ComfyUI "
|
||||
"'normalcrafter' model path, including paths from "
|
||||
"extra_model_paths.yaml. Selecting an HF entry downloads "
|
||||
"the required files into a registered normalcrafter path."
|
||||
),
|
||||
},
|
||||
),
|
||||
},
|
||||
"base_model": (
|
||||
base_model_sources(),
|
||||
{
|
||||
"default": "HF: stabilityai/stable-video-diffusion-img2vid-xt",
|
||||
"tooltip": (
|
||||
"Stable Video Diffusion base components used by "
|
||||
"NormalCrafter: CLIP image encoder, feature extractor "
|
||||
"and scheduler. Local models are discovered from the "
|
||||
"same registered ComfyUI 'normalcrafter' paths."
|
||||
),
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("NORMALCRAFTER_CLEAN_MODEL",)
|
||||
RETURN_TYPES = ("NORMALCRAFTER_MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"Loaded NormalCrafter model bundle. Heavy components stay on CPU until inference needs them.",
|
||||
)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "NormalCrafter/Clean"
|
||||
CATEGORY = "NormalCrafter"
|
||||
DESCRIPTION = (
|
||||
"Loads the NormalCrafter UNet/VAE and the required Stable Video Diffusion "
|
||||
"base components. Model discovery follows ComfyUI's registered "
|
||||
"'normalcrafter' paths, including extra_model_paths.yaml."
|
||||
)
|
||||
|
||||
def load(
|
||||
self,
|
||||
dtype: str,
|
||||
attention: str,
|
||||
local_files_only: bool,
|
||||
model_repo: str = "Yanrui95/NormalCrafter",
|
||||
base_repo: str = "stabilityai/stable-video-diffusion-img2vid-xt",
|
||||
):
|
||||
"""Instantiate the model and return it as a ComfyUI custom object type."""
|
||||
|
||||
# Respect ComfyUI's own device selection rather than hard-coding ``cuda``.
|
||||
def load(self, normalcrafter_model: str, base_model: str):
|
||||
device = comfy.model_management.get_torch_device()
|
||||
resolved_dtype = _resolve_dtype(dtype, device)
|
||||
|
||||
model = NormalCrafterModel.from_pretrained(
|
||||
model_repo=model_repo.strip(),
|
||||
base_repo=base_repo.strip(),
|
||||
dtype=resolved_dtype,
|
||||
attention_mode=attention,
|
||||
local_files_only=local_files_only,
|
||||
return (
|
||||
load_models(
|
||||
device=device,
|
||||
normalcrafter_source=normalcrafter_model,
|
||||
base_model_source=base_model,
|
||||
),
|
||||
)
|
||||
return (model,)
|
||||
|
||||
|
||||
class NormalCrafterCleanGenerate:
|
||||
"""Run the four-stage NormalCrafter inference pipeline.
|
||||
|
||||
ComfyUI IMAGE tensors use ``[frames, height, width, channels]`` and values in
|
||||
``[0, 1]``. The engine accepts that exact format and also returns the same
|
||||
format on CPU, so the graph never caches a live CUDA output tensor.
|
||||
"""
|
||||
class NormalCrafterGenerate:
|
||||
"""Run NormalCrafter with explicit temporal and VRAM controls."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"model": ("NORMALCRAFTER_CLEAN_MODEL",),
|
||||
"images": ("IMAGE",),
|
||||
# The input aspect ratio is preserved. This setting only limits
|
||||
# the longest side before symmetric padding to a VAE-safe size.
|
||||
"model": (
|
||||
"NORMALCRAFTER_MODEL",
|
||||
{
|
||||
"tooltip": (
|
||||
"NormalCrafter model bundle from the NormalCrafter - Load node."
|
||||
)
|
||||
},
|
||||
),
|
||||
"images": (
|
||||
"IMAGE",
|
||||
{
|
||||
"tooltip": (
|
||||
"Input image batch or video frames in temporal order. "
|
||||
"The output contains one normal map per input frame."
|
||||
)
|
||||
},
|
||||
),
|
||||
"max_resolution": (
|
||||
"INT",
|
||||
{"default": 1024, "min": 256, "max": 2048, "step": 64},
|
||||
{
|
||||
"default": 1024,
|
||||
"min": 256,
|
||||
"max": 2048,
|
||||
"step": 64,
|
||||
"tooltip": (
|
||||
"Maximum length of the input's longest side during inference. "
|
||||
"Images are never enlarged. Lower values reduce VRAM use and "
|
||||
"increase speed, at the cost of spatial detail. Output is "
|
||||
"resized back to the original input resolution."
|
||||
),
|
||||
},
|
||||
),
|
||||
"window_size": (
|
||||
"INT",
|
||||
{
|
||||
"default": 14,
|
||||
"min": 2,
|
||||
"max": 64,
|
||||
"step": 1,
|
||||
"tooltip": (
|
||||
"Number of consecutive frames processed by the temporal UNet "
|
||||
"at once. Larger windows provide more temporal context but use "
|
||||
"more VRAM. NormalCrafter's released default is 14."
|
||||
),
|
||||
},
|
||||
),
|
||||
"step_size": (
|
||||
"INT",
|
||||
{
|
||||
"default": 10,
|
||||
"min": 1,
|
||||
"max": 64,
|
||||
"step": 1,
|
||||
"tooltip": (
|
||||
"How many frames the temporal window advances each pass. Must "
|
||||
"be <= window_size. Smaller values create more overlap and "
|
||||
"usually smoother temporal blending, but require more UNet "
|
||||
"passes. With window_size 14, the released default is 10 "
|
||||
"(4-frame regular overlap)."
|
||||
),
|
||||
},
|
||||
),
|
||||
"chunk_size": (
|
||||
"INT",
|
||||
{
|
||||
"default": 7,
|
||||
"min": 1,
|
||||
"max": 64,
|
||||
"step": 1,
|
||||
"tooltip": (
|
||||
"Batch size used for CLIP encoding, VAE encoding and VAE "
|
||||
"decoding. Lower values reduce peak VRAM; higher values can "
|
||||
"improve throughput if enough VRAM is available. This does "
|
||||
"not change the temporal UNet window size."
|
||||
),
|
||||
},
|
||||
),
|
||||
# NormalCrafter was trained around 14-frame temporal windows.
|
||||
"window_size": ("INT", {"default": 14, "min": 2, "max": 32}),
|
||||
# 14 / 10 produces the original regular overlap of four frames.
|
||||
"step_size": ("INT", {"default": 10, "min": 1, "max": 32}),
|
||||
# Chunk sizes trade transfer overhead for peak VRAM usage.
|
||||
"clip_chunk_size": ("INT", {"default": 16, "min": 1, "max": 128}),
|
||||
"vae_encode_chunk_size": ("INT", {"default": 7, "min": 1, "max": 32}),
|
||||
"vae_decode_chunk_size": ("INT", {"default": 4, "min": 1, "max": 32}),
|
||||
# ``staged`` keeps only one heavy component on CUDA at a time.
|
||||
# ``resident`` is faster but requires enough VRAM for all models.
|
||||
"offload_mode": (["staged", "resident"], {"default": "staged"}),
|
||||
"offload_after": ("BOOLEAN", {"default": True}),
|
||||
# ``original`` resizes decoded normals back to input dimensions.
|
||||
"output_size": (["original", "processed"], {"default": "original"}),
|
||||
# Interpolation can shorten vectors; this restores unit length.
|
||||
"renormalize_normals": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE",)
|
||||
RETURN_NAMES = ("normal_maps",)
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"Per-frame RGB normal maps in ComfyUI IMAGE format, restored to the input resolution.",
|
||||
)
|
||||
FUNCTION = "generate"
|
||||
CATEGORY = "NormalCrafter/Clean"
|
||||
CATEGORY = "NormalCrafter"
|
||||
DESCRIPTION = (
|
||||
"Runs NormalCrafter in four stages: CLIP encode, VAE encode, temporal UNet "
|
||||
"inference and VAE decode. Full-video intermediates stay on CPU and only "
|
||||
"the active component/chunk/window is moved to the inference device."
|
||||
)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
model: NormalCrafterModel,
|
||||
images: torch.Tensor,
|
||||
model,
|
||||
images,
|
||||
max_resolution: int,
|
||||
window_size: int,
|
||||
step_size: int,
|
||||
clip_chunk_size: int,
|
||||
vae_encode_chunk_size: int,
|
||||
vae_decode_chunk_size: int,
|
||||
offload_mode: str,
|
||||
offload_after: bool,
|
||||
output_size: str,
|
||||
renormalize_normals: bool,
|
||||
chunk_size: int,
|
||||
):
|
||||
"""Validate user settings, build progress accounting, and invoke inference."""
|
||||
device = comfy.model_management.get_torch_device()
|
||||
total_steps = estimate_progress_steps(
|
||||
int(images.shape[0]),
|
||||
window_size=window_size,
|
||||
step_size=step_size,
|
||||
chunk_size=chunk_size,
|
||||
)
|
||||
progress_bar = comfy.utils.ProgressBar(total_steps)
|
||||
|
||||
config = InferenceConfig(
|
||||
def progress() -> None:
|
||||
progress_bar.update(1)
|
||||
|
||||
output = run_inference(
|
||||
model,
|
||||
images,
|
||||
device=device,
|
||||
max_resolution=max_resolution,
|
||||
window_size=window_size,
|
||||
step_size=step_size,
|
||||
clip_chunk_size=clip_chunk_size,
|
||||
vae_encode_chunk_size=vae_encode_chunk_size,
|
||||
vae_decode_chunk_size=vae_decode_chunk_size,
|
||||
offload_mode=offload_mode,
|
||||
offload_after=offload_after,
|
||||
output_size=output_size,
|
||||
renormalize_normals=renormalize_normals,
|
||||
chunk_size=chunk_size,
|
||||
progress=progress,
|
||||
)
|
||||
config.validate()
|
||||
|
||||
# Short clips are temporally extended by repeating the final frame until
|
||||
# one full window is available. The engine removes those synthetic frames
|
||||
# again before returning the result.
|
||||
frame_count = int(images.shape[0])
|
||||
effective_count = max(frame_count, window_size)
|
||||
|
||||
# Progress is counted in actual chunks/windows, not arbitrary percentages.
|
||||
# This keeps the bar responsive for both short and long clips.
|
||||
total_steps = (
|
||||
(effective_count + clip_chunk_size - 1) // clip_chunk_size
|
||||
+ (effective_count + vae_encode_chunk_size - 1) // vae_encode_chunk_size
|
||||
+ len(build_windows(effective_count, window_size, step_size))
|
||||
+ (frame_count + vae_decode_chunk_size - 1) // vae_decode_chunk_size
|
||||
)
|
||||
progress_bar = comfy.utils.ProgressBar(max(total_steps, 1))
|
||||
|
||||
def progress(_stage: str, _current: int, _total: int) -> None:
|
||||
# The callback keeps stage names available for future logging/UI use,
|
||||
# while ComfyUI's built-in bar currently only needs one increment.
|
||||
progress_bar.update(1)
|
||||
|
||||
device = comfy.model_management.get_torch_device()
|
||||
output = model.infer(images, device=device, config=config, progress=progress)
|
||||
|
||||
# The engine already offloads when requested. Asking ComfyUI to release its
|
||||
# allocator cache as well makes freed blocks available to other nodes.
|
||||
if offload_after:
|
||||
comfy.model_management.soft_empty_cache()
|
||||
|
||||
# The model lifecycle already returned all modules to CPU. This only lets
|
||||
# ComfyUI release allocator blocks that are now unused.
|
||||
comfy.model_management.soft_empty_cache()
|
||||
return (output,)
|
||||
|
||||
|
||||
class NormalCrafterCleanOffload:
|
||||
"""Explicitly move all model components back to CPU.
|
||||
|
||||
This is useful in long graphs where the user wants precise control over when
|
||||
NormalCrafter relinquishes VRAM, independent of the generate node setting.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"model": ("NORMALCRAFTER_CLEAN_MODEL",)}}
|
||||
|
||||
RETURN_TYPES = ("NORMALCRAFTER_CLEAN_MODEL",)
|
||||
RETURN_NAMES = ("model",)
|
||||
FUNCTION = "offload"
|
||||
CATEGORY = "NormalCrafter/Clean"
|
||||
|
||||
def offload(self, model: NormalCrafterModel):
|
||||
model.offload_to_cpu(empty_cache=True)
|
||||
comfy.model_management.soft_empty_cache()
|
||||
return (model,)
|
||||
|
||||
|
||||
# ComfyUI discovers nodes through these mappings. Internal class names remain
|
||||
# stable identifiers for serialized workflows; display names may be more readable.
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"NormalCrafterCleanLoader": NormalCrafterCleanLoader,
|
||||
"NormalCrafterCleanGenerate": NormalCrafterCleanGenerate,
|
||||
"NormalCrafterCleanOffload": NormalCrafterCleanOffload,
|
||||
"NormalCrafterLoader": NormalCrafterLoader,
|
||||
"NormalCrafterGenerate": NormalCrafterGenerate,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"NormalCrafterCleanLoader": "NormalCrafter Clean - Load",
|
||||
"NormalCrafterCleanGenerate": "NormalCrafter Clean - Generate Normals",
|
||||
"NormalCrafterCleanOffload": "NormalCrafter Clean - Offload",
|
||||
"NormalCrafterLoader": "NormalCrafter - Load",
|
||||
"NormalCrafterGenerate": "NormalCrafter - Generate Normals",
|
||||
}
|
||||
|
||||
@@ -1 +1 @@
|
||||
"""Clean NormalCrafter inference package."""
|
||||
"""NormalCrafter runtime package."""
|
||||
|
||||
@@ -1,807 +0,0 @@
|
||||
"""Memory-conscious NormalCrafter inference engine.
|
||||
|
||||
The engine is intentionally organized as four explicit stages:
|
||||
|
||||
1. CLIP image encoding: RGB frames -> semantic frame embeddings.
|
||||
2. VAE encoding: RGB frames -> spatial RGB latents.
|
||||
3. UNet inference: RGB latents + CLIP embeddings -> normal-map latents.
|
||||
4. VAE decoding: normal-map latents -> three-channel normal maps.
|
||||
|
||||
The central ownership rule is simple: large, long-lived intermediate videos stay
|
||||
on CPU. Only the currently active chunk or temporal window is moved to the
|
||||
inference device. This avoids retaining an entire video on CUDA and makes peak
|
||||
VRAM mostly independent of total clip length.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import gc
|
||||
import inspect
|
||||
import logging
|
||||
import threading
|
||||
from typing import Callable, Literal
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from diffusers import AutoencoderKLTemporalDecoder, EulerDiscreteScheduler
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
|
||||
from .preprocess import FrameSource
|
||||
from .unet import NormalCrafterUNet
|
||||
from .windows import Window, build_windows, overlap_length
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Called after every completed chunk/window. The strings make progress reporting
|
||||
# descriptive without tying the engine to ComfyUI itself.
|
||||
ProgressCallback = Callable[[str, int, int], None]
|
||||
|
||||
# ``staged`` minimizes VRAM by moving only one model component to the device.
|
||||
# ``resident`` avoids repeated device transfers but keeps every component loaded.
|
||||
OffloadMode = Literal["staged", "resident"]
|
||||
|
||||
# ``original`` restores the source dimensions after decoding. ``processed`` keeps
|
||||
# the aspect-preserving, max-resolution-limited dimensions used by the model.
|
||||
OutputSize = Literal["original", "processed"]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InferenceConfig:
|
||||
"""Immutable runtime settings for one inference call.
|
||||
|
||||
Keeping these values in a frozen dataclass prevents an inference run from
|
||||
observing settings that change halfway through execution.
|
||||
"""
|
||||
|
||||
max_resolution: int = 1024
|
||||
window_size: int = 14
|
||||
step_size: int = 10
|
||||
clip_chunk_size: int = 16
|
||||
vae_encode_chunk_size: int = 7
|
||||
vae_decode_chunk_size: int = 4
|
||||
offload_mode: OffloadMode = "staged"
|
||||
offload_after: bool = True
|
||||
output_size: OutputSize = "original"
|
||||
renormalize_normals: bool = True
|
||||
|
||||
def validate(self) -> None:
|
||||
"""Reject combinations that would produce gaps or invalid tensors."""
|
||||
|
||||
if self.max_resolution < 64:
|
||||
raise ValueError("max_resolution must be at least 64")
|
||||
if self.window_size < 2:
|
||||
raise ValueError("window_size must be at least 2")
|
||||
|
||||
# A step larger than the window would leave uncovered frames between
|
||||
# windows. A step of zero would never advance.
|
||||
if not 1 <= self.step_size <= self.window_size:
|
||||
raise ValueError("step_size must be between 1 and window_size")
|
||||
|
||||
for name, value in (
|
||||
("clip_chunk_size", self.clip_chunk_size),
|
||||
("vae_encode_chunk_size", self.vae_encode_chunk_size),
|
||||
("vae_decode_chunk_size", self.vae_decode_chunk_size),
|
||||
):
|
||||
if value < 1:
|
||||
raise ValueError(f"{name} must be positive")
|
||||
|
||||
|
||||
class NormalCrafterModel:
|
||||
"""Own and execute every model component required for inference.
|
||||
|
||||
Ownership guarantees:
|
||||
|
||||
* No module-level model globals.
|
||||
* No duplicate pipeline object pointing at the same modules.
|
||||
* No CUDA tensor is returned to ComfyUI.
|
||||
* Intermediate full-video tensors live on CPU.
|
||||
* A re-entrant lock serializes device moves and inference on this instance.
|
||||
|
||||
The object may stay alive in a ComfyUI graph cache, but its resource state is
|
||||
explicit and controllable through ``offload_to_cpu``.
|
||||
"""
|
||||
|
||||
# These values reproduce the conditioning used by the released pipeline.
|
||||
# They are inherited from the Stable Video Diffusion interface, even though
|
||||
# NormalCrafter performs a deterministic one-step prediction.
|
||||
FPS_CONDITIONING = 7
|
||||
MOTION_BUCKET_ID = 127
|
||||
NOISE_AUG_STRENGTH = 0.0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
unet: NormalCrafterUNet,
|
||||
vae: AutoencoderKLTemporalDecoder,
|
||||
image_encoder: CLIPVisionModelWithProjection,
|
||||
feature_extractor: CLIPImageProcessor,
|
||||
scheduler: EulerDiscreteScheduler,
|
||||
dtype: torch.dtype,
|
||||
attention_mode: str,
|
||||
) -> None:
|
||||
# ``eval`` disables training behavior such as dropout and documents that
|
||||
# these modules are inference-only in this implementation.
|
||||
self.unet = unet.eval()
|
||||
self.vae = vae.eval()
|
||||
self.image_encoder = image_encoder.eval()
|
||||
|
||||
# The feature extractor only provides normalization metadata here; image
|
||||
# resizing is done directly in Torch to avoid PIL/NumPy round-trips.
|
||||
self.feature_extractor = feature_extractor
|
||||
self.scheduler = scheduler
|
||||
self.dtype = dtype
|
||||
self.attention_mode = attention_mode
|
||||
|
||||
# Device moves are not thread-safe when two graph executions share one
|
||||
# model instance. The lock also protects scheduler mutation.
|
||||
self._lock = threading.RLock()
|
||||
|
||||
self._configure_attention(attention_mode)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained(
|
||||
cls,
|
||||
*,
|
||||
model_repo: str = "Yanrui95/NormalCrafter",
|
||||
base_repo: str = "stabilityai/stable-video-diffusion-img2vid-xt",
|
||||
dtype: torch.dtype = torch.float16,
|
||||
attention_mode: str = "auto",
|
||||
local_files_only: bool = False,
|
||||
) -> "NormalCrafterModel":
|
||||
"""Load the released weights and assemble one clean model owner.
|
||||
|
||||
NormalCrafter supplies its fine-tuned UNet and VAE. The CLIP vision
|
||||
encoder, preprocessing statistics, and Euler scheduler come from the SVD
|
||||
base repository used during training.
|
||||
"""
|
||||
|
||||
common = {
|
||||
"torch_dtype": dtype,
|
||||
"low_cpu_mem_usage": True,
|
||||
"local_files_only": local_files_only,
|
||||
}
|
||||
|
||||
logger.info("Loading NormalCrafter UNet from %s", model_repo)
|
||||
unet = NormalCrafterUNet.from_pretrained(model_repo, subfolder="unet", **common)
|
||||
|
||||
logger.info("Loading NormalCrafter VAE from %s", model_repo)
|
||||
vae = AutoencoderKLTemporalDecoder.from_pretrained(
|
||||
model_repo,
|
||||
subfolder="vae",
|
||||
**common,
|
||||
)
|
||||
|
||||
# Some Hugging Face repos expose an explicit fp16 variant. Try it first
|
||||
# for half-precision execution, then fall back to the default files.
|
||||
encoder_kwargs = dict(common)
|
||||
if dtype in (torch.float16, torch.bfloat16):
|
||||
encoder_kwargs["variant"] = "fp16"
|
||||
|
||||
try:
|
||||
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
base_repo,
|
||||
subfolder="image_encoder",
|
||||
**encoder_kwargs,
|
||||
)
|
||||
except (OSError, ValueError, TypeError):
|
||||
encoder_kwargs.pop("variant", None)
|
||||
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
base_repo,
|
||||
subfolder="image_encoder",
|
||||
**encoder_kwargs,
|
||||
)
|
||||
|
||||
feature_extractor = CLIPImageProcessor.from_pretrained(
|
||||
base_repo,
|
||||
subfolder="feature_extractor",
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
scheduler = EulerDiscreteScheduler.from_pretrained(
|
||||
base_repo,
|
||||
subfolder="scheduler",
|
||||
local_files_only=local_files_only,
|
||||
)
|
||||
|
||||
model = cls(
|
||||
unet=unet,
|
||||
vae=vae,
|
||||
image_encoder=image_encoder,
|
||||
feature_extractor=feature_extractor,
|
||||
scheduler=scheduler,
|
||||
dtype=dtype,
|
||||
attention_mode=attention_mode,
|
||||
)
|
||||
|
||||
# Loading libraries may choose a device based on environment/config. The
|
||||
# public lifecycle contract is that a newly loaded model starts on CPU.
|
||||
model.offload_to_cpu(empty_cache=False)
|
||||
return model
|
||||
|
||||
@property
|
||||
def modules(self) -> dict[str, torch.nn.Module]:
|
||||
"""Return the heavy components participating in device management."""
|
||||
|
||||
return {
|
||||
"image_encoder": self.image_encoder,
|
||||
"vae": self.vae,
|
||||
"unet": self.unet,
|
||||
}
|
||||
|
||||
def _configure_attention(self, mode: str) -> None:
|
||||
"""Enable xFormers when available or explicitly requested."""
|
||||
|
||||
if mode not in {"auto", "default", "xformers"}:
|
||||
raise ValueError(f"unsupported attention mode: {mode}")
|
||||
|
||||
if mode == "default":
|
||||
return
|
||||
|
||||
try:
|
||||
self.unet.enable_xformers_memory_efficient_attention()
|
||||
logger.info("NormalCrafter xFormers attention enabled")
|
||||
except Exception as exc:
|
||||
# Explicit requests should fail loudly. ``auto`` is allowed to fall
|
||||
# back because current PyTorch/Diffusers attention may be sufficient.
|
||||
if mode == "xformers":
|
||||
raise RuntimeError("xFormers was requested but could not be enabled") from exc
|
||||
logger.info("xFormers unavailable; using Diffusers default attention: %s", exc)
|
||||
|
||||
@staticmethod
|
||||
def _device_of(module: torch.nn.Module) -> torch.device:
|
||||
"""Read a module's current device from its first parameter."""
|
||||
|
||||
try:
|
||||
return next(module.parameters()).device
|
||||
except StopIteration:
|
||||
# Defensive fallback for parameterless modules.
|
||||
return torch.device("cpu")
|
||||
|
||||
def _activate(self, name: str, device: torch.device, offload_mode: OffloadMode) -> None:
|
||||
"""Place model components according to the selected residency policy.
|
||||
|
||||
``resident``:
|
||||
Move all heavy modules to the inference device and leave them there.
|
||||
|
||||
``staged``:
|
||||
Move only the named component to the inference device and force the
|
||||
others to CPU. After those moves, release now-unused CUDA cache blocks.
|
||||
"""
|
||||
|
||||
if offload_mode == "resident":
|
||||
for module in self.modules.values():
|
||||
if self._device_of(module) != device:
|
||||
module.to(device)
|
||||
return
|
||||
|
||||
for module_name, module in self.modules.items():
|
||||
target = device if module_name == name else torch.device("cpu")
|
||||
if self._device_of(module) != target:
|
||||
module.to(target)
|
||||
|
||||
self._release_cuda_cache(device)
|
||||
|
||||
@staticmethod
|
||||
def _release_cuda_cache(device: torch.device) -> None:
|
||||
"""Release unreachable Python objects and unused CUDA allocator blocks.
|
||||
|
||||
``empty_cache`` does not free live tensors. Its usefulness depends on the
|
||||
preceding ownership discipline: temporary tensors must first go out of
|
||||
scope or be deleted, and inactive modules must already have moved to CPU.
|
||||
"""
|
||||
|
||||
gc.collect()
|
||||
if device.type == "cuda" and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def offload_to_cpu(self, *, empty_cache: bool = True) -> None:
|
||||
"""Move all owned model weights to CPU and optionally empty CUDA cache."""
|
||||
|
||||
for module in self.modules.values():
|
||||
module.to("cpu")
|
||||
|
||||
if empty_cache:
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
def infer(
|
||||
self,
|
||||
images_bhwc: torch.Tensor,
|
||||
*,
|
||||
device: torch.device,
|
||||
config: InferenceConfig,
|
||||
progress: ProgressCallback | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run complete inference and return CPU normal maps in ComfyUI format.
|
||||
|
||||
Input shape:
|
||||
``[frames, height, width, channels]``, values expected in ``[0, 1]``.
|
||||
|
||||
Output shape:
|
||||
``[real_frames, output_height, output_width, 3]`` on CPU, values in
|
||||
``[0, 1]`` where encoded RGB corresponds to XYZ normal direction.
|
||||
"""
|
||||
|
||||
config.validate()
|
||||
|
||||
# FrameSource owns the canonical CPU input and lazily prepares only the
|
||||
# chunk needed by the current stage.
|
||||
source = FrameSource(images_bhwc, max_resolution=config.max_resolution)
|
||||
|
||||
# The temporal UNet requires at least one full window. Short clips repeat
|
||||
# their final frame internally; decoding is still limited to real frames.
|
||||
effective_frames = source.effective_frame_count(config.window_size)
|
||||
completed = False
|
||||
|
||||
# ``torch.inference_mode`` removes autograd bookkeeping more aggressively
|
||||
# than ``no_grad`` and prevents accidental graph retention.
|
||||
with self._lock, torch.inference_mode():
|
||||
try:
|
||||
# Stage 1: semantic embedding per frame, stored on CPU.
|
||||
clip_embeddings = self._encode_clip(
|
||||
source,
|
||||
effective_frames,
|
||||
device,
|
||||
config,
|
||||
progress,
|
||||
)
|
||||
|
||||
# Stage 2: spatial RGB latent per frame, stored on CPU.
|
||||
rgb_latents = self._encode_vae(
|
||||
source,
|
||||
effective_frames,
|
||||
device,
|
||||
config,
|
||||
progress,
|
||||
)
|
||||
|
||||
# Stage 3: sliding-window temporal prediction, stored on CPU.
|
||||
normal_latents = self._predict_latents(
|
||||
rgb_latents,
|
||||
clip_embeddings,
|
||||
device,
|
||||
config,
|
||||
progress,
|
||||
)
|
||||
|
||||
# These full-video intermediates are no longer needed before the
|
||||
# decoder stage. Delete references early to reduce host memory.
|
||||
del rgb_latents, clip_embeddings
|
||||
|
||||
# Stage 4: decode only real input frames into final normals.
|
||||
output = self._decode_normals(
|
||||
normal_latents,
|
||||
source,
|
||||
device,
|
||||
config,
|
||||
progress,
|
||||
)
|
||||
|
||||
completed = True
|
||||
return output
|
||||
finally:
|
||||
# On failure we always offload. On success this follows the user's
|
||||
# setting, allowing ``resident`` mode to remain fast across runs.
|
||||
if config.offload_after or not completed:
|
||||
self.offload_to_cpu(empty_cache=True)
|
||||
|
||||
def _encode_clip(
|
||||
self,
|
||||
source: FrameSource,
|
||||
frame_count: int,
|
||||
device: torch.device,
|
||||
config: InferenceConfig,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
"""Encode every frame into one CLIP vision embedding on CPU.
|
||||
|
||||
Shape flow, typically:
|
||||
|
||||
``[chunk, 3, H, W]`` -> ``[chunk, 3, 224, 224]`` ->
|
||||
``[chunk, 1, embedding_dim]``.
|
||||
"""
|
||||
|
||||
self._activate("image_encoder", device, config.offload_mode)
|
||||
|
||||
# Reproduce the CLIP normalization configured by the SVD repository.
|
||||
# Mean/std remain on CPU because subtraction occurs before device transfer.
|
||||
mean = torch.tensor(
|
||||
self.feature_extractor.image_mean,
|
||||
dtype=torch.float32,
|
||||
).view(1, 3, 1, 1)
|
||||
std = torch.tensor(
|
||||
self.feature_extractor.image_std,
|
||||
dtype=torch.float32,
|
||||
).view(1, 3, 1, 1)
|
||||
|
||||
output: torch.Tensor | None = None
|
||||
total = (frame_count + config.clip_chunk_size - 1) // config.clip_chunk_size
|
||||
|
||||
for chunk_index, start in enumerate(
|
||||
range(0, frame_count, config.clip_chunk_size),
|
||||
start=1,
|
||||
):
|
||||
end = min(start + config.clip_chunk_size, frame_count)
|
||||
|
||||
# FrameSource resizes/pads lazily and repeats the final real frame for
|
||||
# synthetic short-clip indices.
|
||||
frames = source.get_bchw(start, end)
|
||||
|
||||
# CLIP vision expects its own fixed spatial resolution.
|
||||
frames = F.interpolate(
|
||||
frames,
|
||||
size=(224, 224),
|
||||
mode="bicubic",
|
||||
align_corners=False,
|
||||
antialias=True,
|
||||
).clamp_(0.0, 1.0)
|
||||
|
||||
# Only the current normalized chunk crosses to the inference device.
|
||||
pixels = ((frames - mean) / std).to(device=device, dtype=self.dtype)
|
||||
embeddings = self.image_encoder(pixel_values=pixels).image_embeds
|
||||
|
||||
# Diffusers cross-attention expects a sequence dimension. Depending on
|
||||
# Transformers version, ``image_embeds`` may arrive as rank two.
|
||||
if embeddings.ndim == 2:
|
||||
embeddings = embeddings.unsqueeze(1)
|
||||
|
||||
# Copy the result back immediately so total video length does not grow
|
||||
# CUDA residency. Half precision is retained to limit host memory.
|
||||
embeddings_cpu = embeddings.detach().to(device="cpu", dtype=self.dtype)
|
||||
|
||||
# Allocate the final CPU tensor once the runtime embedding shape is
|
||||
# known, then fill slices instead of accumulating and concatenating.
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
(frame_count, *embeddings_cpu.shape[1:]),
|
||||
dtype=self.dtype,
|
||||
device="cpu",
|
||||
)
|
||||
output[start:end].copy_(embeddings_cpu)
|
||||
|
||||
# Drop all references to current chunk/device tensors before moving on.
|
||||
del frames, pixels, embeddings, embeddings_cpu
|
||||
|
||||
if progress:
|
||||
progress("CLIP", chunk_index, total)
|
||||
|
||||
assert output is not None
|
||||
return output
|
||||
|
||||
def _encode_vae(
|
||||
self,
|
||||
source: FrameSource,
|
||||
frame_count: int,
|
||||
device: torch.device,
|
||||
config: InferenceConfig,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
"""Encode RGB frames into VAE latents and store them on CPU.
|
||||
|
||||
Unlike generative sampling, this path uses the latent distribution mode,
|
||||
making encoding deterministic and avoiding random sampling noise.
|
||||
"""
|
||||
|
||||
self._activate("vae", device, config.offload_mode)
|
||||
|
||||
original_dtype = next(self.vae.parameters()).dtype
|
||||
|
||||
# Some Diffusers VAEs declare ``force_upcast`` because selected operations
|
||||
# are unstable in FP16. The dtype is restored in ``finally`` even on error.
|
||||
needs_upcast = (
|
||||
original_dtype == torch.float16
|
||||
and bool(getattr(self.vae.config, "force_upcast", False))
|
||||
)
|
||||
if needs_upcast:
|
||||
self.vae.to(dtype=torch.float32)
|
||||
compute_dtype = next(self.vae.parameters()).dtype
|
||||
|
||||
output: torch.Tensor | None = None
|
||||
total = (
|
||||
frame_count + config.vae_encode_chunk_size - 1
|
||||
) // config.vae_encode_chunk_size
|
||||
|
||||
try:
|
||||
for chunk_index, start in enumerate(
|
||||
range(0, frame_count, config.vae_encode_chunk_size),
|
||||
start=1,
|
||||
):
|
||||
end = min(start + config.vae_encode_chunk_size, frame_count)
|
||||
|
||||
frames = source.get_bchw(start, end)
|
||||
|
||||
# VAE input convention is [-1, 1], whereas ComfyUI uses [0, 1].
|
||||
frames = (frames * 2.0 - 1.0).to(
|
||||
device=device,
|
||||
dtype=compute_dtype,
|
||||
)
|
||||
|
||||
latents = self.vae.encode(frames).latent_dist.mode()
|
||||
latents_cpu = latents.detach().to(device="cpu", dtype=self.dtype)
|
||||
|
||||
# Runtime allocation avoids assumptions about latent channel count
|
||||
# or spatial scale while still eliminating list + torch.cat churn.
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
(frame_count, *latents_cpu.shape[1:]),
|
||||
dtype=self.dtype,
|
||||
device="cpu",
|
||||
)
|
||||
output[start:end].copy_(latents_cpu)
|
||||
|
||||
del frames, latents, latents_cpu
|
||||
|
||||
if progress:
|
||||
progress("VAE encode", chunk_index, total)
|
||||
finally:
|
||||
if needs_upcast:
|
||||
self.vae.to(dtype=original_dtype)
|
||||
|
||||
assert output is not None
|
||||
return output
|
||||
|
||||
def _added_time_ids(self, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
|
||||
"""Construct the three Stable Video Diffusion conditioning values."""
|
||||
|
||||
values = [
|
||||
self.FPS_CONDITIONING,
|
||||
self.MOTION_BUCKET_ID,
|
||||
self.NOISE_AUG_STRENGTH,
|
||||
]
|
||||
|
||||
# Fail explicitly if a different UNet config expects a different added
|
||||
# conditioning width. Silent shape coercion here would produce bad output.
|
||||
expected = self.unet.add_embedding.linear_1.in_features
|
||||
actual = int(self.unet.config.addition_time_embed_dim) * len(values)
|
||||
if expected != actual:
|
||||
raise ValueError(
|
||||
f"UNet added-time embedding mismatch: expected {expected}, constructed {actual}"
|
||||
)
|
||||
|
||||
return torch.tensor([values], dtype=dtype, device=device)
|
||||
|
||||
def _predict_latents(
|
||||
self,
|
||||
rgb_latents_cpu: torch.Tensor,
|
||||
clip_embeddings_cpu: torch.Tensor,
|
||||
device: torch.device,
|
||||
config: InferenceConfig,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
"""Predict normal latents over overlapping temporal windows.
|
||||
|
||||
Full-video input/output tensors remain on CPU. Each iteration moves only
|
||||
one window of RGB latents and CLIP embeddings to the inference device.
|
||||
|
||||
The released inference path is a deterministic single scheduler step:
|
||||
normal latents start at zero, RGB latents are concatenated as conditioning,
|
||||
and the UNet predicts the one-step update.
|
||||
"""
|
||||
|
||||
self._activate("unet", device, config.offload_mode)
|
||||
|
||||
frame_count = int(rgb_latents_cpu.shape[0])
|
||||
windows = build_windows(frame_count, config.window_size, config.step_size)
|
||||
|
||||
# The output shape matches the VAE latent shape. It is allocated once and
|
||||
# overwritten by each anchored window instead of repeatedly concatenated.
|
||||
output_cpu = torch.empty_like(rgb_latents_cpu, device="cpu")
|
||||
|
||||
previous_window: Window | None = None
|
||||
previous_latents: torch.Tensor | None = None
|
||||
|
||||
for index, window in enumerate(windows, start=1):
|
||||
# RGB latent window: [frames, channels, h, w] ->
|
||||
# [batch=1, frames, channels, h, w].
|
||||
image_latents = rgb_latents_cpu[window.start : window.end].to(
|
||||
device=device,
|
||||
dtype=self.dtype,
|
||||
).unsqueeze(0)
|
||||
|
||||
# One CLIP embedding is supplied for every video frame.
|
||||
embeddings = clip_embeddings_cpu[window.start : window.end].to(
|
||||
device=device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
|
||||
# NormalCrafter's prediction state starts from zero rather than random
|
||||
# diffusion noise. This is why the node intentionally has no seed.
|
||||
latents = torch.zeros(
|
||||
(
|
||||
1,
|
||||
window.length,
|
||||
int(self.unet.config.out_channels),
|
||||
*image_latents.shape[-2:],
|
||||
),
|
||||
device=device,
|
||||
dtype=self.dtype,
|
||||
)
|
||||
|
||||
overlap = overlap_length(previous_window, window)
|
||||
old_overlap: torch.Tensor | None = None
|
||||
|
||||
if overlap and previous_window is not None and previous_latents is not None:
|
||||
# Determine where the current window begins inside the previous
|
||||
# window, then seed the overlapping prefix with prior predictions.
|
||||
previous_offset = window.start - previous_window.start
|
||||
old_overlap = previous_latents[
|
||||
:,
|
||||
previous_offset : previous_offset + overlap,
|
||||
]
|
||||
latents[:, :overlap].copy_(old_overlap)
|
||||
|
||||
# NormalCrafter uses exactly one Euler scheduler timestep.
|
||||
self.scheduler.set_timesteps(1, device=device)
|
||||
timestep = self.scheduler.timesteps[0]
|
||||
|
||||
# Scheduler scaling is kept even for one step because it is part of
|
||||
# the trained SVD inference convention.
|
||||
model_input = self.scheduler.scale_model_input(latents, timestep)
|
||||
|
||||
# Channel concatenation: four predicted-normal channels plus four RGB
|
||||
# latent conditioning channels -> eight UNet input channels.
|
||||
model_input = torch.cat((model_input, image_latents), dim=2)
|
||||
added_time_ids = self._added_time_ids(dtype=self.dtype, device=device)
|
||||
|
||||
noise_prediction = self.unet(
|
||||
model_input,
|
||||
timestep,
|
||||
encoder_hidden_states=embeddings,
|
||||
added_time_ids=added_time_ids,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
predicted = self.scheduler.step(
|
||||
noise_prediction,
|
||||
timestep,
|
||||
latents,
|
||||
).prev_sample
|
||||
|
||||
if old_overlap is not None:
|
||||
# Crossfade old -> new prediction over the overlap. Endpoint values
|
||||
# are excluded so neither side receives a fully duplicated weight.
|
||||
weights = torch.linspace(
|
||||
1.0,
|
||||
0.0,
|
||||
overlap + 2,
|
||||
device=device,
|
||||
dtype=self.dtype,
|
||||
)[1:-1].view(1, overlap, 1, 1, 1)
|
||||
|
||||
predicted[:, :overlap] = (
|
||||
old_overlap * weights
|
||||
+ predicted[:, :overlap] * (1.0 - weights)
|
||||
)
|
||||
|
||||
# Later windows intentionally overwrite their entire covered region.
|
||||
# The overlap has already been blended, and the final anchored window
|
||||
# guarantees coverage through the last frame.
|
||||
output_cpu[window.start : window.end].copy_(
|
||||
predicted[0].detach().to("cpu")
|
||||
)
|
||||
|
||||
# Retain only the immediately previous GPU window for overlap seeding.
|
||||
# This bounds temporal state by ``window_size``, not total clip length.
|
||||
previous_window = window
|
||||
previous_latents = predicted.detach()
|
||||
|
||||
del image_latents, embeddings, latents, model_input, added_time_ids, noise_prediction
|
||||
|
||||
if progress:
|
||||
progress("UNet", index, len(windows))
|
||||
|
||||
return output_cpu
|
||||
|
||||
def _decode_normals(
|
||||
self,
|
||||
normal_latents_cpu: torch.Tensor,
|
||||
source: FrameSource,
|
||||
device: torch.device,
|
||||
config: InferenceConfig,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
"""Decode normal latents into final CPU ComfyUI IMAGE tensors."""
|
||||
|
||||
self._activate("vae", device, config.offload_mode)
|
||||
|
||||
original_dtype = next(self.vae.parameters()).dtype
|
||||
needs_upcast = (
|
||||
original_dtype == torch.float16
|
||||
and bool(getattr(self.vae.config, "force_upcast", False))
|
||||
)
|
||||
if needs_upcast:
|
||||
self.vae.to(dtype=torch.float32)
|
||||
compute_dtype = next(self.vae.parameters()).dtype
|
||||
|
||||
if config.output_size == "original":
|
||||
out_height, out_width = source.original_height, source.original_width
|
||||
else:
|
||||
out_height, out_width = source.resized_height, source.resized_width
|
||||
|
||||
# Allocate only the real frame count. Any synthetic padding frames used to
|
||||
# satisfy a short temporal window are never decoded or returned.
|
||||
output = torch.empty(
|
||||
(source.frame_count, out_height, out_width, 3),
|
||||
dtype=torch.float32,
|
||||
device="cpu",
|
||||
)
|
||||
|
||||
scaling_factor = float(self.vae.config.scaling_factor)
|
||||
|
||||
# Diffusers versions differ in whether temporal decoders accept an
|
||||
# explicit ``num_frames`` argument. Inspect once and adapt safely.
|
||||
forward_vae = (
|
||||
self.vae._orig_mod.forward
|
||||
if hasattr(self.vae, "_orig_mod")
|
||||
else self.vae.forward
|
||||
)
|
||||
accepts_num_frames = "num_frames" in inspect.signature(forward_vae).parameters
|
||||
|
||||
total = (
|
||||
source.frame_count + config.vae_decode_chunk_size - 1
|
||||
) // config.vae_decode_chunk_size
|
||||
|
||||
try:
|
||||
for chunk_index, start in enumerate(
|
||||
range(0, source.frame_count, config.vae_decode_chunk_size),
|
||||
start=1,
|
||||
):
|
||||
end = min(start + config.vae_decode_chunk_size, source.frame_count)
|
||||
|
||||
latents = normal_latents_cpu[start:end].to(
|
||||
device=device,
|
||||
dtype=compute_dtype,
|
||||
)
|
||||
|
||||
# Undo the latent scaling expected by the VAE decoder.
|
||||
latents = latents / scaling_factor
|
||||
|
||||
decode_kwargs = (
|
||||
{"num_frames": end - start}
|
||||
if accepts_num_frames
|
||||
else {}
|
||||
)
|
||||
# Decode, promote to FP32, and transfer to CPU in one expression.
|
||||
# The chained calls are kept identical to the validated version.
|
||||
normals = self.vae.decode(latents, **decode_kwargs).sample.float().to("cpu")
|
||||
|
||||
# Remove the white spatial padding introduced before encoding.
|
||||
normals = source.crop_padding(normals)
|
||||
|
||||
if (out_height, out_width) != (
|
||||
source.resized_height,
|
||||
source.resized_width,
|
||||
):
|
||||
normals = F.interpolate(
|
||||
normals,
|
||||
size=(out_height, out_width),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
# Decoder output represents XYZ components in [-1, 1].
|
||||
normals = normals.clamp(-1.0, 1.0)
|
||||
|
||||
# Bilinear resizing blends vector components and can reduce vector
|
||||
# magnitude. Renormalization restores unit-length surface normals.
|
||||
if config.renormalize_normals:
|
||||
normals = F.normalize(normals, p=2.0, dim=1, eps=1e-6)
|
||||
|
||||
# Encode signed XYZ into display/image range [0, 1].
|
||||
normals = (normals * 0.5 + 0.5).clamp_(0.0, 1.0)
|
||||
|
||||
# Convert BCHW -> BHWC, the layout required by ComfyUI IMAGE.
|
||||
output[start:end].copy_(
|
||||
normals.permute(0, 2, 3, 1).contiguous()
|
||||
)
|
||||
|
||||
del latents, normals
|
||||
|
||||
if progress:
|
||||
progress("VAE decode", chunk_index, total)
|
||||
finally:
|
||||
if needs_upcast:
|
||||
self.vae.to(dtype=original_dtype)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,441 @@
|
||||
"""NormalCrafter inference pipeline.
|
||||
|
||||
Full-video intermediates stay on CPU. Only the active CLIP/VAE chunk or UNet
|
||||
window is moved to the inference device.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import inspect
|
||||
import math
|
||||
from typing import Callable
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .model import NormalCrafterModels
|
||||
|
||||
FPS_CONDITIONING = 7
|
||||
MOTION_BUCKET_ID = 127
|
||||
NOISE_AUG_STRENGTH = 0.0
|
||||
|
||||
ProgressCallback = Callable[[], None]
|
||||
|
||||
|
||||
@dataclass
|
||||
class FrameSource:
|
||||
"""Lazy CPU-side preprocessing view over a ComfyUI IMAGE batch."""
|
||||
|
||||
frames: torch.Tensor
|
||||
max_resolution: int
|
||||
window_size: int
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if self.frames.ndim != 4 or self.frames.shape[-1] < 3:
|
||||
raise ValueError("images must have shape [frames, height, width, channels>=3]")
|
||||
if self.frames.shape[0] == 0:
|
||||
raise ValueError("images contains no frames")
|
||||
|
||||
self.frames = (
|
||||
self.frames[..., :3]
|
||||
.detach()
|
||||
.to(device="cpu", dtype=torch.float32)
|
||||
.contiguous()
|
||||
.clamp_(0.0, 1.0)
|
||||
)
|
||||
|
||||
_, self.original_height, self.original_width, _ = self.frames.shape
|
||||
scale = min(
|
||||
1.0,
|
||||
self.max_resolution / max(self.original_height, self.original_width),
|
||||
)
|
||||
self.height = max(1, int(round(self.original_height * scale)))
|
||||
self.width = max(1, int(round(self.original_width * scale)))
|
||||
|
||||
padded_height = math.ceil(self.height / 64) * 64
|
||||
padded_width = math.ceil(self.width / 64) * 64
|
||||
pad_h = padded_height - self.height
|
||||
pad_w = padded_width - self.width
|
||||
self.pad_top = pad_h // 2
|
||||
self.pad_bottom = pad_h - self.pad_top
|
||||
self.pad_left = pad_w // 2
|
||||
self.pad_right = pad_w - self.pad_left
|
||||
|
||||
@property
|
||||
def frame_count(self) -> int:
|
||||
return int(self.frames.shape[0])
|
||||
|
||||
@property
|
||||
def effective_frame_count(self) -> int:
|
||||
return max(self.frame_count, self.window_size)
|
||||
|
||||
def get(self, start: int, end: int) -> torch.Tensor:
|
||||
"""Return processed BCHW frames; short clips repeat their final frame."""
|
||||
|
||||
indices = torch.arange(start, end, dtype=torch.long).clamp_max(
|
||||
self.frame_count - 1
|
||||
)
|
||||
chunk = self.frames.index_select(0, indices).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
if (self.height, self.width) != (self.original_height, self.original_width):
|
||||
chunk = F.interpolate(
|
||||
chunk,
|
||||
size=(self.height, self.width),
|
||||
mode="bicubic",
|
||||
align_corners=False,
|
||||
antialias=True,
|
||||
).clamp_(0.0, 1.0)
|
||||
|
||||
if self.pad_top or self.pad_bottom or self.pad_left or self.pad_right:
|
||||
chunk = F.pad(
|
||||
chunk,
|
||||
(self.pad_left, self.pad_right, self.pad_top, self.pad_bottom),
|
||||
value=1.0,
|
||||
)
|
||||
return chunk
|
||||
|
||||
def crop(self, frames: torch.Tensor) -> torch.Tensor:
|
||||
return frames[
|
||||
:,
|
||||
:,
|
||||
self.pad_top : self.pad_top + self.height,
|
||||
self.pad_left : self.pad_left + self.width,
|
||||
]
|
||||
|
||||
|
||||
def _validate_settings(window_size: int, step_size: int, chunk_size: int) -> None:
|
||||
if window_size < 2:
|
||||
raise ValueError("window_size must be at least 2")
|
||||
if not 1 <= step_size <= window_size:
|
||||
raise ValueError("step_size must be between 1 and window_size")
|
||||
if chunk_size < 1:
|
||||
raise ValueError("chunk_size must be at least 1")
|
||||
|
||||
|
||||
def _build_windows(
|
||||
frame_count: int,
|
||||
window_size: int,
|
||||
step_size: int,
|
||||
) -> list[tuple[int, int]]:
|
||||
"""Build the same full-size, tail-anchored windows as NormalCrafter."""
|
||||
|
||||
if frame_count < window_size:
|
||||
raise ValueError("frame_count must already be padded to window_size")
|
||||
|
||||
last_start = frame_count - window_size
|
||||
starts = list(range(0, last_start + 1, step_size))
|
||||
if starts[-1] != last_start:
|
||||
starts.append(last_start)
|
||||
return [(start, start + window_size) for start in starts]
|
||||
|
||||
|
||||
def estimate_progress_steps(
|
||||
frame_count: int,
|
||||
*,
|
||||
window_size: int,
|
||||
step_size: int,
|
||||
chunk_size: int,
|
||||
) -> int:
|
||||
_validate_settings(window_size, step_size, chunk_size)
|
||||
effective = max(frame_count, window_size)
|
||||
return max(
|
||||
1,
|
||||
math.ceil(effective / chunk_size)
|
||||
+ math.ceil(effective / chunk_size)
|
||||
+ len(_build_windows(effective, window_size, step_size))
|
||||
+ math.ceil(frame_count / chunk_size),
|
||||
)
|
||||
|
||||
|
||||
def run_inference(
|
||||
models: NormalCrafterModels,
|
||||
images: torch.Tensor,
|
||||
*,
|
||||
device: torch.device,
|
||||
max_resolution: int,
|
||||
window_size: int,
|
||||
step_size: int,
|
||||
chunk_size: int,
|
||||
progress: ProgressCallback | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run CLIP -> VAE encode -> temporal UNet -> VAE decode."""
|
||||
|
||||
_validate_settings(window_size, step_size, chunk_size)
|
||||
source = FrameSource(
|
||||
images,
|
||||
max_resolution=max_resolution,
|
||||
window_size=window_size,
|
||||
)
|
||||
|
||||
try:
|
||||
with torch.inference_mode():
|
||||
embeddings = _encode_clip(models, source, device, chunk_size, progress)
|
||||
rgb_latents = _encode_rgb(models, source, device, chunk_size, progress)
|
||||
normal_latents = _predict_normals(
|
||||
models,
|
||||
rgb_latents,
|
||||
embeddings,
|
||||
device,
|
||||
window_size,
|
||||
step_size,
|
||||
progress,
|
||||
)
|
||||
del rgb_latents, embeddings
|
||||
return _decode_normals(
|
||||
models,
|
||||
normal_latents,
|
||||
source,
|
||||
device,
|
||||
chunk_size,
|
||||
progress,
|
||||
)
|
||||
finally:
|
||||
models.offload_all()
|
||||
|
||||
|
||||
def _tick(progress: ProgressCallback | None) -> None:
|
||||
if progress is not None:
|
||||
progress()
|
||||
|
||||
|
||||
def _encode_clip(
|
||||
models: NormalCrafterModels,
|
||||
source: FrameSource,
|
||||
device: torch.device,
|
||||
chunk_size: int,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
models.activate("image_encoder", device)
|
||||
|
||||
mean = torch.tensor(
|
||||
models.feature_extractor.image_mean,
|
||||
dtype=torch.float32,
|
||||
).view(1, 3, 1, 1)
|
||||
std = torch.tensor(
|
||||
models.feature_extractor.image_std,
|
||||
dtype=torch.float32,
|
||||
).view(1, 3, 1, 1)
|
||||
output = None
|
||||
|
||||
for start in range(0, source.effective_frame_count, chunk_size):
|
||||
end = min(start + chunk_size, source.effective_frame_count)
|
||||
frames = source.get(start, end)
|
||||
frames = F.interpolate(
|
||||
frames,
|
||||
size=(224, 224),
|
||||
mode="bicubic",
|
||||
align_corners=False,
|
||||
antialias=True,
|
||||
).clamp_(0.0, 1.0)
|
||||
|
||||
pixels = ((frames - mean) / std).to(device=device, dtype=models.dtype)
|
||||
embeddings = models.image_encoder(pixel_values=pixels).image_embeds
|
||||
if embeddings.ndim == 2:
|
||||
embeddings = embeddings.unsqueeze(1)
|
||||
embeddings = embeddings.detach().to(device="cpu", dtype=models.dtype)
|
||||
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
(source.effective_frame_count, *embeddings.shape[1:]),
|
||||
dtype=models.dtype,
|
||||
device="cpu",
|
||||
)
|
||||
output[start:end].copy_(embeddings)
|
||||
_tick(progress)
|
||||
|
||||
assert output is not None
|
||||
return output
|
||||
|
||||
|
||||
def _encode_rgb(
|
||||
models: NormalCrafterModels,
|
||||
source: FrameSource,
|
||||
device: torch.device,
|
||||
chunk_size: int,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
models.activate("vae", device)
|
||||
vae_dtype = next(models.vae.parameters()).dtype
|
||||
output = None
|
||||
|
||||
for start in range(0, source.effective_frame_count, chunk_size):
|
||||
end = min(start + chunk_size, source.effective_frame_count)
|
||||
frames = (source.get(start, end) * 2.0 - 1.0).to(
|
||||
device=device,
|
||||
dtype=vae_dtype,
|
||||
)
|
||||
latents = models.vae.encode(frames).latent_dist.mode()
|
||||
latents = latents.detach().to(device="cpu", dtype=models.dtype)
|
||||
|
||||
if output is None:
|
||||
output = torch.empty(
|
||||
(source.effective_frame_count, *latents.shape[1:]),
|
||||
dtype=models.dtype,
|
||||
device="cpu",
|
||||
)
|
||||
output[start:end].copy_(latents)
|
||||
_tick(progress)
|
||||
|
||||
assert output is not None
|
||||
return output
|
||||
|
||||
|
||||
def _added_time_ids(
|
||||
models: NormalCrafterModels,
|
||||
device: torch.device,
|
||||
) -> torch.Tensor:
|
||||
values = [FPS_CONDITIONING, MOTION_BUCKET_ID, NOISE_AUG_STRENGTH]
|
||||
expected = models.unet.add_embedding.linear_1.in_features
|
||||
actual = int(models.unet.config.addition_time_embed_dim) * len(values)
|
||||
if expected != actual:
|
||||
raise ValueError(
|
||||
f"UNet added-time embedding mismatch: expected {expected}, got {actual}"
|
||||
)
|
||||
return torch.tensor([values], dtype=models.dtype, device=device)
|
||||
|
||||
|
||||
def _predict_normals(
|
||||
models: NormalCrafterModels,
|
||||
rgb_latents: torch.Tensor,
|
||||
embeddings: torch.Tensor,
|
||||
device: torch.device,
|
||||
window_size: int,
|
||||
step_size: int,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
models.activate("unet", device)
|
||||
windows = _build_windows(
|
||||
int(rgb_latents.shape[0]),
|
||||
window_size,
|
||||
step_size,
|
||||
)
|
||||
output = torch.empty_like(rgb_latents, device="cpu")
|
||||
|
||||
previous_window: tuple[int, int] | None = None
|
||||
previous_prediction: torch.Tensor | None = None
|
||||
|
||||
for start, end in windows:
|
||||
image_latents = rgb_latents[start:end].to(
|
||||
device=device,
|
||||
dtype=models.dtype,
|
||||
).unsqueeze(0)
|
||||
frame_embeddings = embeddings[start:end].to(
|
||||
device=device,
|
||||
dtype=models.dtype,
|
||||
)
|
||||
latents = torch.zeros(
|
||||
(
|
||||
1,
|
||||
end - start,
|
||||
int(models.unet.config.out_channels),
|
||||
*image_latents.shape[-2:],
|
||||
),
|
||||
device=device,
|
||||
dtype=models.dtype,
|
||||
)
|
||||
|
||||
overlap = 0 if previous_window is None else max(0, previous_window[1] - start)
|
||||
old_overlap = None
|
||||
if overlap and previous_window is not None and previous_prediction is not None:
|
||||
previous_offset = start - previous_window[0]
|
||||
old_overlap = previous_prediction[
|
||||
:,
|
||||
previous_offset : previous_offset + overlap,
|
||||
]
|
||||
latents[:, :overlap].copy_(old_overlap)
|
||||
|
||||
models.scheduler.set_timesteps(1, device=device)
|
||||
timestep = models.scheduler.timesteps[0]
|
||||
model_input = models.scheduler.scale_model_input(latents, timestep)
|
||||
model_input = torch.cat((model_input, image_latents), dim=2)
|
||||
|
||||
noise_prediction = models.unet(
|
||||
model_input,
|
||||
timestep,
|
||||
encoder_hidden_states=frame_embeddings,
|
||||
added_time_ids=_added_time_ids(models, device),
|
||||
return_dict=False,
|
||||
)[0]
|
||||
predicted = models.scheduler.step(
|
||||
noise_prediction,
|
||||
timestep,
|
||||
latents,
|
||||
).prev_sample
|
||||
|
||||
if old_overlap is not None:
|
||||
weights = torch.linspace(
|
||||
1.0,
|
||||
0.0,
|
||||
overlap + 2,
|
||||
device=device,
|
||||
dtype=models.dtype,
|
||||
)[1:-1].view(1, overlap, 1, 1, 1)
|
||||
predicted[:, :overlap] = (
|
||||
old_overlap * weights
|
||||
+ predicted[:, :overlap] * (1.0 - weights)
|
||||
)
|
||||
|
||||
output[start:end].copy_(predicted[0].detach().to("cpu"))
|
||||
previous_window = (start, end)
|
||||
previous_prediction = predicted.detach()
|
||||
_tick(progress)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
def _decode_normals(
|
||||
models: NormalCrafterModels,
|
||||
normal_latents: torch.Tensor,
|
||||
source: FrameSource,
|
||||
device: torch.device,
|
||||
chunk_size: int,
|
||||
progress: ProgressCallback | None,
|
||||
) -> torch.Tensor:
|
||||
models.activate("vae", device)
|
||||
vae_dtype = next(models.vae.parameters()).dtype
|
||||
|
||||
output = torch.empty(
|
||||
(source.frame_count, source.original_height, source.original_width, 3),
|
||||
dtype=torch.float32,
|
||||
device="cpu",
|
||||
)
|
||||
scaling_factor = float(models.vae.config.scaling_factor)
|
||||
|
||||
forward = (
|
||||
models.vae._orig_mod.forward
|
||||
if hasattr(models.vae, "_orig_mod")
|
||||
else models.vae.forward
|
||||
)
|
||||
accepts_num_frames = "num_frames" in inspect.signature(forward).parameters
|
||||
|
||||
for start in range(0, source.frame_count, chunk_size):
|
||||
end = min(start + chunk_size, source.frame_count)
|
||||
latents = normal_latents[start:end].to(
|
||||
device=device,
|
||||
dtype=vae_dtype,
|
||||
) / scaling_factor
|
||||
kwargs = {"num_frames": end - start} if accepts_num_frames else {}
|
||||
normals = models.vae.decode(latents, **kwargs).sample.float().to("cpu")
|
||||
normals = source.crop(normals)
|
||||
|
||||
if (source.height, source.width) != (
|
||||
source.original_height,
|
||||
source.original_width,
|
||||
):
|
||||
normals = F.interpolate(
|
||||
normals,
|
||||
size=(source.original_height, source.original_width),
|
||||
mode="bilinear",
|
||||
align_corners=False,
|
||||
)
|
||||
|
||||
normals = normals.clamp(-1.0, 1.0)
|
||||
normals = F.normalize(normals, p=2.0, dim=1, eps=1e-6)
|
||||
normals = (normals * 0.5 + 0.5).clamp_(0.0, 1.0)
|
||||
output[start:end].copy_(
|
||||
normals.permute(0, 2, 3, 1).contiguous()
|
||||
)
|
||||
_tick(progress)
|
||||
|
||||
return output
|
||||
@@ -0,0 +1,375 @@
|
||||
"""Model discovery, loading, and GPU residency for NormalCrafter.
|
||||
|
||||
NormalCrafter uses its own ComfyUI model category: ``normalcrafter``.
|
||||
All path discovery goes through ComfyUI's ``folder_paths`` registry, so paths
|
||||
registered through ``extra_model_paths.yaml`` are visible automatically.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import gc
|
||||
import logging
|
||||
import os
|
||||
|
||||
import folder_paths
|
||||
import torch
|
||||
from diffusers import AutoencoderKLTemporalDecoder, EulerDiscreteScheduler
|
||||
from huggingface_hub import snapshot_download
|
||||
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
||||
|
||||
from .unet import NormalCrafterUNet
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
MODEL_CATEGORY = "normalcrafter"
|
||||
NORMALCRAFTER_REPO = "Yanrui95/NormalCrafter"
|
||||
SVD_REPO = "stabilityai/stable-video-diffusion-img2vid-xt"
|
||||
|
||||
HF_PREFIX = "HF: "
|
||||
LOCAL_PREFIX = "local: "
|
||||
|
||||
|
||||
def _register_model_category() -> None:
|
||||
"""Ensure ComfyUI knows the default ``models/normalcrafter`` location.
|
||||
|
||||
If ``extra_model_paths.yaml`` already registered one or more ``normalcrafter``
|
||||
paths, they stay in ComfyUI's ordering and the standard ComfyUI model path is
|
||||
simply added as another search location.
|
||||
"""
|
||||
|
||||
default_path = os.path.join(folder_paths.models_dir, MODEL_CATEGORY)
|
||||
|
||||
try:
|
||||
existing = folder_paths.get_folder_paths(MODEL_CATEGORY)
|
||||
except KeyError:
|
||||
existing = []
|
||||
|
||||
if default_path in existing:
|
||||
return
|
||||
|
||||
# Current ComfyUI supports ``is_default``. The fallback keeps compatibility
|
||||
# with older versions whose function accepted only two arguments.
|
||||
try:
|
||||
folder_paths.add_model_folder_path(
|
||||
MODEL_CATEGORY,
|
||||
default_path,
|
||||
is_default=not existing,
|
||||
)
|
||||
except TypeError:
|
||||
folder_paths.add_model_folder_path(MODEL_CATEGORY, default_path)
|
||||
|
||||
|
||||
_register_model_category()
|
||||
|
||||
|
||||
def _model_roots() -> list[str]:
|
||||
"""Return all NormalCrafter roots currently registered in ComfyUI."""
|
||||
|
||||
roots = folder_paths.get_folder_paths(MODEL_CATEGORY)
|
||||
# Preserve ComfyUI ordering while removing duplicate spellings of a path.
|
||||
seen: set[str] = set()
|
||||
result: list[str] = []
|
||||
for root in roots:
|
||||
normalized = os.path.normcase(os.path.abspath(root))
|
||||
if normalized in seen:
|
||||
continue
|
||||
seen.add(normalized)
|
||||
result.append(os.path.abspath(root))
|
||||
return result
|
||||
|
||||
|
||||
def _is_within_registered_root(path: str) -> bool:
|
||||
candidate = os.path.normcase(os.path.abspath(path))
|
||||
for root in _model_roots():
|
||||
registered = os.path.normcase(os.path.abspath(root))
|
||||
try:
|
||||
if os.path.commonpath((candidate, registered)) == registered:
|
||||
return True
|
||||
except ValueError:
|
||||
# Different Windows drives cannot share a common path.
|
||||
continue
|
||||
return False
|
||||
|
||||
|
||||
def _local_model_dirs(required_subfolders: tuple[str, ...]) -> list[str]:
|
||||
"""Find compatible model directories across every ComfyUI-registered root."""
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
for root in _model_roots():
|
||||
if not os.path.isdir(root):
|
||||
continue
|
||||
|
||||
try:
|
||||
names = sorted(os.listdir(root), key=str.lower)
|
||||
except OSError:
|
||||
continue
|
||||
|
||||
for name in names:
|
||||
path = os.path.join(root, name)
|
||||
if not os.path.isdir(path):
|
||||
continue
|
||||
if all(
|
||||
os.path.isdir(os.path.join(path, subfolder))
|
||||
for subfolder in required_subfolders
|
||||
):
|
||||
results.append(f"{LOCAL_PREFIX}{path}")
|
||||
|
||||
return results
|
||||
|
||||
|
||||
def normalcrafter_sources() -> list[str]:
|
||||
"""Dropdown values suitable for NormalCrafter UNet + VAE loading."""
|
||||
|
||||
return [
|
||||
f"{HF_PREFIX}{NORMALCRAFTER_REPO}",
|
||||
*_local_model_dirs(("unet", "vae")),
|
||||
]
|
||||
|
||||
|
||||
def base_model_sources() -> list[str]:
|
||||
"""Dropdown values suitable for SVD CLIP + scheduler loading."""
|
||||
|
||||
return [
|
||||
f"{HF_PREFIX}{SVD_REPO}",
|
||||
*_local_model_dirs(("image_encoder", "feature_extractor", "scheduler")),
|
||||
]
|
||||
|
||||
|
||||
def _existing_repo_dir(repo_id: str) -> str | None:
|
||||
"""Return an already-downloaded repo from any registered model root."""
|
||||
|
||||
folder_name = repo_id.replace("/", "--")
|
||||
for root in _model_roots():
|
||||
candidate = os.path.join(root, folder_name)
|
||||
if os.path.isdir(candidate):
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def _download_root() -> str:
|
||||
"""Choose the first writable NormalCrafter root in ComfyUI's own order."""
|
||||
|
||||
errors: list[str] = []
|
||||
for root in _model_roots():
|
||||
try:
|
||||
os.makedirs(root, exist_ok=True)
|
||||
if os.access(root, os.W_OK):
|
||||
return root
|
||||
errors.append(f"{root} (not writable)")
|
||||
except OSError as exc:
|
||||
errors.append(f"{root} ({exc})")
|
||||
|
||||
details = "; ".join(errors) or "no paths registered"
|
||||
raise RuntimeError(
|
||||
"No writable ComfyUI 'normalcrafter' model path is available: " + details
|
||||
)
|
||||
|
||||
|
||||
def _download_repo(repo_id: str, required_subfolders: tuple[str, ...]) -> str:
|
||||
"""Download only required components into a registered NormalCrafter root."""
|
||||
|
||||
existing = _existing_repo_dir(repo_id)
|
||||
if existing is not None:
|
||||
return existing
|
||||
|
||||
target = os.path.join(_download_root(), repo_id.replace("/", "--"))
|
||||
allow_patterns = [f"{subfolder}/**" for subfolder in required_subfolders]
|
||||
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
local_dir=target,
|
||||
allow_patterns=allow_patterns,
|
||||
)
|
||||
return target
|
||||
|
||||
|
||||
def _resolve_source(source: str, required_subfolders: tuple[str, ...]) -> str:
|
||||
"""Resolve a dropdown value to a validated local model directory."""
|
||||
|
||||
if source.startswith(HF_PREFIX):
|
||||
path = _download_repo(source[len(HF_PREFIX) :], required_subfolders)
|
||||
elif source.startswith(LOCAL_PREFIX):
|
||||
path = source[len(LOCAL_PREFIX) :]
|
||||
if not _is_within_registered_root(path):
|
||||
raise ValueError(
|
||||
"Selected model path is not inside a ComfyUI-registered "
|
||||
f"'{MODEL_CATEGORY}' directory: {path}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unsupported model source: {source}")
|
||||
|
||||
if not os.path.isdir(path):
|
||||
raise FileNotFoundError(f"NormalCrafter model folder not found: {path}")
|
||||
|
||||
missing = [
|
||||
subfolder
|
||||
for subfolder in required_subfolders
|
||||
if not os.path.isdir(os.path.join(path, subfolder))
|
||||
]
|
||||
if missing:
|
||||
raise FileNotFoundError(
|
||||
f"Model folder {path} is missing required subfolders: {', '.join(missing)}"
|
||||
)
|
||||
|
||||
return path
|
||||
|
||||
|
||||
def _clear_cuda_cache() -> None:
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def _auto_dtype(device: torch.device) -> torch.dtype:
|
||||
return torch.float16 if device.type == "cuda" else torch.float32
|
||||
|
||||
|
||||
@dataclass
|
||||
class NormalCrafterModels:
|
||||
"""Own the heavy modules and keep GPU residency explicit."""
|
||||
|
||||
unet: NormalCrafterUNet
|
||||
vae: AutoencoderKLTemporalDecoder
|
||||
image_encoder: CLIPVisionModelWithProjection
|
||||
feature_extractor: CLIPImageProcessor
|
||||
scheduler: EulerDiscreteScheduler
|
||||
dtype: torch.dtype
|
||||
|
||||
def _modules(self) -> dict[str, torch.nn.Module]:
|
||||
return {
|
||||
"image_encoder": self.image_encoder,
|
||||
"vae": self.vae,
|
||||
"unet": self.unet,
|
||||
}
|
||||
|
||||
def activate(self, name: str, device: torch.device) -> torch.nn.Module:
|
||||
"""Keep only one heavy component on the inference device."""
|
||||
|
||||
modules = self._modules()
|
||||
if name not in modules:
|
||||
raise ValueError(f"Unknown NormalCrafter component: {name}")
|
||||
|
||||
cpu = torch.device("cpu")
|
||||
for module_name, module in modules.items():
|
||||
target = device if module_name == name else cpu
|
||||
try:
|
||||
current = next(module.parameters()).device
|
||||
except StopIteration:
|
||||
current = cpu
|
||||
|
||||
if current != target:
|
||||
module.to(target)
|
||||
|
||||
_clear_cuda_cache()
|
||||
return modules[name]
|
||||
|
||||
def offload_all(self) -> None:
|
||||
"""Return every heavy module to CPU and release unused CUDA blocks."""
|
||||
|
||||
for module in self._modules().values():
|
||||
module.to("cpu")
|
||||
_clear_cuda_cache()
|
||||
|
||||
|
||||
def _enable_memory_efficient_attention(unet: NormalCrafterUNet) -> None:
|
||||
try:
|
||||
unet.enable_xformers_memory_efficient_attention()
|
||||
logger.info("NormalCrafter: xFormers attention enabled")
|
||||
except Exception as exc:
|
||||
logger.info("NormalCrafter: using Diffusers default attention (%s)", exc)
|
||||
|
||||
|
||||
def load_models(
|
||||
*,
|
||||
device: torch.device,
|
||||
normalcrafter_source: str,
|
||||
base_model_source: str,
|
||||
) -> NormalCrafterModels:
|
||||
"""Load the selected model components on CPU.
|
||||
|
||||
UNet and CLIP use FP16 on CUDA systems and FP32 on CPU. The temporal VAE is
|
||||
intentionally loaded in FP32, matching its force-upcast inference behavior.
|
||||
Inference moves CLIP, VAE and UNet to the GPU one stage at a time.
|
||||
"""
|
||||
|
||||
normalcrafter_path = _resolve_source(
|
||||
normalcrafter_source,
|
||||
("unet", "vae"),
|
||||
)
|
||||
base_model_path = _resolve_source(
|
||||
base_model_source,
|
||||
("image_encoder", "feature_extractor", "scheduler"),
|
||||
)
|
||||
|
||||
dtype = _auto_dtype(device)
|
||||
common = {
|
||||
"torch_dtype": dtype,
|
||||
"low_cpu_mem_usage": True,
|
||||
"local_files_only": True,
|
||||
}
|
||||
|
||||
unet = NormalCrafterUNet.from_pretrained(
|
||||
normalcrafter_path,
|
||||
subfolder="unet",
|
||||
**common,
|
||||
)
|
||||
|
||||
# The temporal VAE declares force_upcast and is evaluated in FP32 in the
|
||||
# reference NormalCrafter pipeline. Load it in FP32 from the start instead of
|
||||
# repeatedly calling vae.to(dtype=...) during inference. Recent Diffusers
|
||||
# versions explicitly warn against those runtime dtype casts because selected
|
||||
# VAE modules may need to remain in FP32.
|
||||
vae = AutoencoderKLTemporalDecoder.from_pretrained(
|
||||
normalcrafter_path,
|
||||
subfolder="vae",
|
||||
torch_dtype=torch.float32,
|
||||
low_cpu_mem_usage=True,
|
||||
local_files_only=True,
|
||||
)
|
||||
|
||||
encoder_kwargs = dict(common)
|
||||
if dtype == torch.float16:
|
||||
encoder_kwargs["variant"] = "fp16"
|
||||
|
||||
try:
|
||||
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
base_model_path,
|
||||
subfolder="image_encoder",
|
||||
**encoder_kwargs,
|
||||
)
|
||||
except (OSError, ValueError, TypeError):
|
||||
encoder_kwargs.pop("variant", None)
|
||||
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
||||
base_model_path,
|
||||
subfolder="image_encoder",
|
||||
**encoder_kwargs,
|
||||
)
|
||||
|
||||
feature_extractor = CLIPImageProcessor.from_pretrained(
|
||||
base_model_path,
|
||||
subfolder="feature_extractor",
|
||||
local_files_only=True,
|
||||
)
|
||||
scheduler = EulerDiscreteScheduler.from_pretrained(
|
||||
base_model_path,
|
||||
subfolder="scheduler",
|
||||
local_files_only=True,
|
||||
)
|
||||
|
||||
unet.eval()
|
||||
vae.eval()
|
||||
image_encoder.eval()
|
||||
_enable_memory_efficient_attention(unet)
|
||||
|
||||
models = NormalCrafterModels(
|
||||
unet=unet,
|
||||
vae=vae,
|
||||
image_encoder=image_encoder,
|
||||
feature_extractor=feature_extractor,
|
||||
scheduler=scheduler,
|
||||
dtype=dtype,
|
||||
)
|
||||
models.offload_all()
|
||||
return models
|
||||
@@ -1,194 +0,0 @@
|
||||
"""CPU-side frame preparation for NormalCrafter inference.
|
||||
|
||||
ComfyUI supplies video-like image batches as a single ``BHWC`` tensor:
|
||||
|
||||
[frame, height, width, channel]
|
||||
|
||||
The model stack, however, expects ``BCHW`` chunks whose spatial dimensions are
|
||||
compatible with the VAE. This module bridges those conventions without creating
|
||||
a second processed copy of the complete clip.
|
||||
|
||||
The crucial design choice is laziness: ``FrameSource`` stores the original clip
|
||||
once on CPU and performs resize, padding, layout conversion, and short-clip frame
|
||||
repetition only for the range requested by the current inference stage.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
def _round_size(value: float) -> int:
|
||||
"""Round a scaled spatial dimension while guaranteeing at least one pixel."""
|
||||
|
||||
return max(1, int(round(value)))
|
||||
|
||||
|
||||
def _ceil_multiple(value: int, multiple: int) -> int:
|
||||
"""Return the smallest multiple of ``multiple`` not smaller than ``value``.
|
||||
|
||||
NormalCrafter's VAE/UNet stack repeatedly downsamples spatial dimensions.
|
||||
Padding to 64 is conservative and prevents incompatible shapes deep inside
|
||||
the network while keeping the original aspect ratio intact.
|
||||
"""
|
||||
|
||||
return int(math.ceil(value / multiple) * multiple)
|
||||
|
||||
|
||||
@dataclass
|
||||
class FrameSource:
|
||||
"""Lazy CPU-side view of a ComfyUI IMAGE batch.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
frames_bhwc:
|
||||
ComfyUI ``IMAGE`` tensor in ``[frames, height, width, channels]`` layout.
|
||||
Values are expected to represent display-linear image data in ``[0, 1]``.
|
||||
max_resolution:
|
||||
Upper bound for the longest spatial side used by the model. Inputs are
|
||||
never enlarged; smaller clips retain their native dimensions.
|
||||
pad_multiple:
|
||||
Spatial alignment required by the latent model. Padding is symmetric,
|
||||
except for the unavoidable one-pixel asymmetry when the total is odd.
|
||||
|
||||
Memory behavior
|
||||
---------------
|
||||
Only the source clip is retained. A resized/padded tensor is created solely
|
||||
for the requested chunk and becomes collectible after that stage iteration.
|
||||
This is central to keeping peak RAM and VRAM independent of total clip length.
|
||||
"""
|
||||
|
||||
frames_bhwc: torch.Tensor
|
||||
max_resolution: int
|
||||
pad_multiple: int = 64
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""Validate input and precompute immutable geometry metadata."""
|
||||
|
||||
# ComfyUI IMAGE is BHWC. Alpha or auxiliary channels are allowed, but the
|
||||
# released model consumes exactly RGB, so at least three channels exist.
|
||||
if self.frames_bhwc.ndim != 4 or self.frames_bhwc.shape[-1] < 3:
|
||||
raise ValueError("images must have shape [frames, height, width, channels>=3]")
|
||||
if self.frames_bhwc.shape[0] == 0:
|
||||
raise ValueError("images contains no frames")
|
||||
if self.max_resolution < 64:
|
||||
raise ValueError("max_resolution must be at least 64")
|
||||
|
||||
# Break any autograd relationship, discard non-RGB channels, and establish
|
||||
# one canonical long-lived representation: contiguous CPU FP32 in [0, 1].
|
||||
# Returning to CPU here also prevents a CUDA input tensor from becoming a
|
||||
# hidden full-video allocation retained throughout inference.
|
||||
self.frames_bhwc = (
|
||||
self.frames_bhwc[..., :3]
|
||||
.detach()
|
||||
.to(device="cpu", dtype=torch.float32)
|
||||
.contiguous()
|
||||
.clamp_(0.0, 1.0)
|
||||
)
|
||||
|
||||
# Compute a single aspect-preserving scale. ``min(1.0, ...)`` explicitly
|
||||
# forbids upscaling: max_resolution is a ceiling, not a target size.
|
||||
_, self.original_height, self.original_width, _ = self.frames_bhwc.shape
|
||||
scale = min(1.0, self.max_resolution / max(self.original_height, self.original_width))
|
||||
self.resized_height = _round_size(self.original_height * scale)
|
||||
self.resized_width = _round_size(self.original_width * scale)
|
||||
|
||||
# The network receives padded dimensions, while output restoration later
|
||||
# removes precisely these margins before optional resize to source size.
|
||||
self.padded_height = _ceil_multiple(self.resized_height, self.pad_multiple)
|
||||
self.padded_width = _ceil_multiple(self.resized_width, self.pad_multiple)
|
||||
|
||||
# Split padding around the image. When an odd number of pixels is needed,
|
||||
# bottom/right receive the extra pixel; crop_padding uses the same offsets.
|
||||
total_pad_h = self.padded_height - self.resized_height
|
||||
total_pad_w = self.padded_width - self.resized_width
|
||||
self.pad_top = total_pad_h // 2
|
||||
self.pad_bottom = total_pad_h - self.pad_top
|
||||
self.pad_left = total_pad_w // 2
|
||||
self.pad_right = total_pad_w - self.pad_left
|
||||
|
||||
@property
|
||||
def frame_count(self) -> int:
|
||||
"""Number of real input frames, excluding any temporal repetition."""
|
||||
|
||||
return int(self.frames_bhwc.shape[0])
|
||||
|
||||
def effective_frame_count(self, window_size: int) -> int:
|
||||
"""Frame count seen by encoders after minimum-window temporal padding.
|
||||
|
||||
NormalCrafter's temporal UNet requires at least one complete window. For a
|
||||
shorter clip, ``get_bchw`` repeats the final real frame up to this count.
|
||||
The decoder still emits only ``frame_count`` real output frames.
|
||||
"""
|
||||
|
||||
return max(self.frame_count, window_size)
|
||||
|
||||
def get_bchw(self, start: int, end: int) -> torch.Tensor:
|
||||
"""Materialize one processed CPU frame range in BCHW layout.
|
||||
|
||||
The returned tensor is:
|
||||
|
||||
* CPU resident;
|
||||
* FP32;
|
||||
* RGB only;
|
||||
* in ``[0, 1]``;
|
||||
* resized without changing aspect ratio;
|
||||
* symmetrically white-padded to ``pad_multiple``;
|
||||
* temporally extended by repeating the final frame when ``end`` exceeds
|
||||
the real clip length.
|
||||
|
||||
White padding matches the neutral image-space convention used by the
|
||||
original pipeline more closely than zero/black padding.
|
||||
"""
|
||||
|
||||
if start < 0 or end <= start:
|
||||
raise ValueError(f"invalid frame range [{start}, {end})")
|
||||
|
||||
# Clamp synthetic indices to the final real frame. ``index_select`` then
|
||||
# handles ordinary ranges and repeated-tail ranges through one code path.
|
||||
indices = torch.arange(start, end, dtype=torch.long).clamp_max(self.frame_count - 1)
|
||||
|
||||
# Select BHWC frames, then convert once to the BCHW layout expected by
|
||||
# interpolation, CLIP preprocessing, and the VAE encoder.
|
||||
chunk = self.frames_bhwc.index_select(0, indices).permute(0, 3, 1, 2).contiguous()
|
||||
|
||||
# Bicubic plus antialiasing is appropriate for RGB downscaling. Clamp after
|
||||
# interpolation because bicubic kernels may overshoot slightly outside
|
||||
# the nominal [0, 1] input range near high-contrast edges.
|
||||
if (self.resized_height, self.resized_width) != (self.original_height, self.original_width):
|
||||
chunk = F.interpolate(
|
||||
chunk,
|
||||
size=(self.resized_height, self.resized_width),
|
||||
mode="bicubic",
|
||||
align_corners=False,
|
||||
antialias=True,
|
||||
).clamp_(0.0, 1.0)
|
||||
|
||||
# torch.nn.functional.pad expects (left, right, top, bottom) for BCHW.
|
||||
# Padding is skipped entirely when the resized dimensions are already safe.
|
||||
if self.pad_top or self.pad_bottom or self.pad_left or self.pad_right:
|
||||
chunk = F.pad(
|
||||
chunk,
|
||||
(self.pad_left, self.pad_right, self.pad_top, self.pad_bottom),
|
||||
mode="constant",
|
||||
value=1.0,
|
||||
)
|
||||
return chunk
|
||||
|
||||
def crop_padding(self, tensor_bchw: torch.Tensor) -> torch.Tensor:
|
||||
"""Remove the exact spatial padding previously introduced by ``get_bchw``.
|
||||
|
||||
Channel and batch/frame dimensions are preserved. The method intentionally
|
||||
returns a view where possible; callers may later interpolate or copy it.
|
||||
"""
|
||||
|
||||
return tensor_bchw[
|
||||
:,
|
||||
:,
|
||||
self.pad_top : self.pad_top + self.resized_height,
|
||||
self.pad_left : self.pad_left + self.resized_width,
|
||||
]
|
||||
@@ -1,56 +1,19 @@
|
||||
"""Inference-only NormalCrafter adaptation of Diffusers' temporal SVD UNet.
|
||||
"""NormalCrafter adaptation of Diffusers' spatio-temporal SVD UNet."""
|
||||
|
||||
The released checkpoint uses the Stable Video Diffusion UNet architecture but
|
||||
conditions each frame with its own CLIP embedding. Standard SVD commonly starts
|
||||
from one image embedding and repeats it across the generated video. The override
|
||||
below keeps Diffusers' block execution intact while accepting either convention.
|
||||
|
||||
Training-only machinery from the research project, including DINO supervision,
|
||||
is deliberately absent: it shaped the learned weights but is not part of runtime
|
||||
inference.
|
||||
"""
|
||||
|
||||
# Adapted for inference from Hugging Face Diffusers (Apache-2.0) and
|
||||
# Binyr/NormalCrafter (MIT). See NOTICE.
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Union
|
||||
|
||||
import torch
|
||||
from diffusers import UNetSpatioTemporalConditionModel
|
||||
|
||||
# Diffusers moved this output class between package paths. Supporting both keeps
|
||||
# the node compatible across a broader range of ComfyUI environments without
|
||||
# changing the actual UNet computation.
|
||||
try:
|
||||
from diffusers.models.unets.unet_spatio_temporal_condition import (
|
||||
UNetSpatioTemporalConditionOutput,
|
||||
)
|
||||
except ImportError: # older Diffusers
|
||||
from diffusers.models.unet_spatio_temporal_condition import ( # type: ignore
|
||||
UNetSpatioTemporalConditionOutput,
|
||||
)
|
||||
from diffusers.models.unets.unet_spatio_temporal_condition import UNetSpatioTemporalConditionOutput
|
||||
except ImportError:
|
||||
from diffusers.models.unet_spatio_temporal_condition import UNetSpatioTemporalConditionOutput # type: ignore
|
||||
|
||||
|
||||
class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
"""SVD UNet accepting one CLIP embedding per video frame.
|
||||
|
||||
Input shape
|
||||
-----------
|
||||
``sample`` is ``[batch, frames, channels, latent_h, latent_w]``. In the clean
|
||||
engine, channels are eight: four zero/previous normal-latent channels plus
|
||||
four encoded RGB-conditioning channels.
|
||||
|
||||
Conditioning shape
|
||||
------------------
|
||||
``encoder_hidden_states`` may contain either:
|
||||
|
||||
* one embedding per batch item, which is repeated over frames like SVD; or
|
||||
* one embedding per frame, already flattened as ``batch * frames``.
|
||||
|
||||
Output shape mirrors the temporal input layout and contains the four predicted
|
||||
normal-latent channels expected by the Euler scheduler step.
|
||||
"""
|
||||
"""SVD UNet variant that accepts one CLIP embedding per frame."""
|
||||
|
||||
def forward(
|
||||
self,
|
||||
@@ -60,15 +23,6 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
added_time_ids: torch.Tensor,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
"""Run one spatio-temporal UNet evaluation.
|
||||
|
||||
This follows Diffusers' native forward structure closely. The important
|
||||
NormalCrafter-specific behavior is the frame-wise conditioning validation
|
||||
and preservation rather than unconditional repetition of one embedding.
|
||||
"""
|
||||
|
||||
# Normalize scalar Python or zero-dimensional Torch timesteps into the
|
||||
# one-dimensional tensor form expected by the projection layer.
|
||||
timesteps = timestep
|
||||
if not torch.is_tensor(timesteps):
|
||||
dtype = torch.float64 if isinstance(timestep, float) else torch.int64
|
||||
@@ -79,48 +33,26 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
batch_size, num_frames = sample.shape[:2]
|
||||
|
||||
# One scheduler timestep applies to every item in the inference batch.
|
||||
timesteps = timesteps.expand(batch_size)
|
||||
|
||||
# Standard diffusion timestep embedding. Cast to sample dtype so half or
|
||||
# bfloat inference does not accidentally promote subsequent activations.
|
||||
t_emb = self.time_proj(timesteps).to(dtype=sample.dtype)
|
||||
emb = self.time_embedding(t_emb)
|
||||
|
||||
# SVD's added conditioning carries FPS, motion bucket, and noise strength.
|
||||
# The projection output is added to the ordinary diffusion-time embedding.
|
||||
time_embeds = self.add_time_proj(added_time_ids.flatten())
|
||||
time_embeds = time_embeds.reshape(batch_size, -1).to(emb.dtype)
|
||||
emb = emb + self.add_embedding(time_embeds)
|
||||
|
||||
# Diffusers' spatial blocks operate on batch*frames as the leading axis;
|
||||
# therefore every frame receives a copy of its video's time embedding.
|
||||
emb = emb.repeat_interleave(num_frames, dim=0)
|
||||
|
||||
# [B, F, C, H, W] -> [B*F, C, H, W]. Temporal awareness is preserved by
|
||||
# image_only_indicator and temporal modules inside the SVD blocks.
|
||||
sample = sample.flatten(0, 1)
|
||||
|
||||
# Accept standard SVD conditioning ([B, ...]) or NormalCrafter's per-frame
|
||||
# conditioning ([B*F, ...]). Reject every other count before entering deep
|
||||
# attention layers, where the resulting error would be far less readable.
|
||||
expected_frame_batch = batch_size * num_frames
|
||||
if encoder_hidden_states.shape[0] == batch_size:
|
||||
encoder_hidden_states = encoder_hidden_states.repeat_interleave(num_frames, dim=0)
|
||||
elif encoder_hidden_states.shape[0] != expected_frame_batch:
|
||||
raise ValueError(
|
||||
"encoder_hidden_states must contain either one embedding per batch "
|
||||
f"or one per frame; got {encoder_hidden_states.shape[0]} for "
|
||||
f"batch={batch_size}, frames={num_frames}"
|
||||
"encoder_hidden_states must contain one embedding per batch item or per frame"
|
||||
)
|
||||
|
||||
# Initial projection maps the eight-channel latent input into the UNet's
|
||||
# first feature width defined by the downloaded checkpoint configuration.
|
||||
sample = self.conv_in(sample)
|
||||
|
||||
# All frames are video frames, not isolated images. SVD temporal blocks use
|
||||
# this indicator to control image/video-specific blending behavior.
|
||||
image_only_indicator = torch.zeros(
|
||||
batch_size,
|
||||
num_frames,
|
||||
@@ -128,8 +60,6 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
device=sample.device,
|
||||
)
|
||||
|
||||
# Encoder/down path. Residual feature maps are accumulated for the matching
|
||||
# decoder/up blocks, reproducing the ordinary UNet skip-connection topology.
|
||||
down_block_res_samples = (sample,)
|
||||
for downsample_block in self.down_blocks:
|
||||
if getattr(downsample_block, "has_cross_attention", False):
|
||||
@@ -147,8 +77,6 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
)
|
||||
down_block_res_samples += residuals
|
||||
|
||||
# Bottleneck with cross-attention and temporal processing at the lowest
|
||||
# spatial resolution, where receptive field is largest.
|
||||
sample = self.mid_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
@@ -156,8 +84,6 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
|
||||
# Decoder/up path. Each block consumes exactly as many stored residuals as
|
||||
# it has ResNet sub-blocks, then removes them from the residual stack.
|
||||
for upsample_block in self.up_blocks:
|
||||
residuals = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[: -len(upsample_block.resnets)]
|
||||
@@ -178,12 +104,9 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel):
|
||||
image_only_indicator=image_only_indicator,
|
||||
)
|
||||
|
||||
# Final normalization, activation, and convolution restore the checkpoint's
|
||||
# configured output channel count. Then recover explicit video dimensions.
|
||||
sample = self.conv_out(self.conv_act(self.conv_norm_out(sample)))
|
||||
sample = sample.reshape(batch_size, num_frames, *sample.shape[1:])
|
||||
|
||||
# Match Diffusers' conventional tuple-or-dataclass API.
|
||||
if not return_dict:
|
||||
return (sample,)
|
||||
return UNetSpatioTemporalConditionOutput(sample=sample)
|
||||
|
||||
@@ -1,83 +0,0 @@
|
||||
"""Deterministic temporal window scheduling for arbitrary-length clips.
|
||||
|
||||
NormalCrafter's spatio-temporal UNet processes a fixed number of frames at once.
|
||||
Long clips are therefore divided into overlapping windows. The scheduler in this
|
||||
module has no Torch or ComfyUI dependency, making it simple to reason about and
|
||||
unit-test independently from the model stack.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Window:
|
||||
"""Half-open frame interval ``[start, end)``.
|
||||
|
||||
Frozen instances are safe to share as descriptive values, and ``slots`` keeps
|
||||
this tiny object free of an unnecessary per-instance attribute dictionary.
|
||||
"""
|
||||
|
||||
start: int
|
||||
end: int
|
||||
|
||||
@property
|
||||
def length(self) -> int:
|
||||
"""Number of frames contained in the interval."""
|
||||
|
||||
return self.end - self.start
|
||||
|
||||
|
||||
def build_windows(num_frames: int, window_size: int, step_size: int) -> list[Window]:
|
||||
"""Build fixed-size windows that cover every frame without temporal gaps.
|
||||
|
||||
``step_size`` controls how far each regular window advances. Consequently,
|
||||
regular overlap is ``window_size - step_size``. The final window is anchored
|
||||
exactly to the clip end, even when this creates a larger-than-regular overlap.
|
||||
|
||||
Example for 25 frames, window 14, step 10::
|
||||
|
||||
[0, 14), [10, 24), [11, 25)
|
||||
|
||||
The last start moves from the regular candidate 20 back to 11 so frames 24 and
|
||||
25 are covered by a full-size model window. Short clips must already have an
|
||||
effective frame count of at least ``window_size``; ``FrameSource`` provides
|
||||
that temporal extension by repeating the final real frame.
|
||||
"""
|
||||
|
||||
if num_frames <= 0:
|
||||
raise ValueError("num_frames must be positive")
|
||||
if window_size <= 0:
|
||||
raise ValueError("window_size must be positive")
|
||||
if step_size <= 0:
|
||||
raise ValueError("step_size must be positive")
|
||||
if step_size > window_size:
|
||||
raise ValueError("step_size cannot exceed window_size; gaps would be created")
|
||||
if num_frames < window_size:
|
||||
raise ValueError("caller must pad short clips to window_size")
|
||||
|
||||
# A full-size final window beginning here ends exactly at num_frames.
|
||||
last_start = num_frames - window_size
|
||||
|
||||
# Generate the regular cadence as far as a full window still fits.
|
||||
starts = list(range(0, last_start + 1, step_size))
|
||||
|
||||
# A cadence rarely lands exactly on the clip tail. Append one anchored window
|
||||
# instead of accepting an uncovered remainder or creating a shorter window.
|
||||
if starts[-1] != last_start:
|
||||
starts.append(last_start)
|
||||
|
||||
return [Window(start=s, end=s + window_size) for s in starts]
|
||||
|
||||
|
||||
def overlap_length(previous: Window | None, current: Window) -> int:
|
||||
"""Return how many leading frames of ``current`` overlap ``previous``.
|
||||
|
||||
Windows are generated in increasing order, so only the previous end and the
|
||||
current start are needed. Non-overlapping intervals correctly produce zero.
|
||||
"""
|
||||
|
||||
if previous is None:
|
||||
return 0
|
||||
return max(0, previous.end - current.start)
|
||||
@@ -0,0 +1,22 @@
|
||||
[project]
|
||||
name = "BetterNormalCrafterWrapper"
|
||||
version = "1.0.0"
|
||||
description = "A better and memory-safe implementation of the beloved NormalCrafterWrapper"
|
||||
requires-python = ">=3.13"
|
||||
license = "Apache License"
|
||||
license-files = ["LICENSE"]
|
||||
dependencies = [
|
||||
"diffusers>=0.30.0,<1.0",
|
||||
"transformers>=4.39.0,<5.0",
|
||||
"accelerate>=0.28.0",
|
||||
"huggingface_hub>=0.23.0",
|
||||
"safetensors>=0.4.2",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Repository = "https://github.com/kaski23/BetterNormalCrafterWrapper"
|
||||
|
||||
[tool.comfy]
|
||||
PublisherId = "kaski"
|
||||
DisplayName = "Better NormalCrafter Wrapper"
|
||||
Icon = ""
|
||||
@@ -1,26 +0,0 @@
|
||||
"""Fast unit tests for temporal coverage rules.
|
||||
|
||||
These tests intentionally avoid importing Torch, Diffusers, or ComfyUI. Window
|
||||
scheduling is pure infrastructure and should remain testable even on a machine
|
||||
without model dependencies or CUDA.
|
||||
"""
|
||||
|
||||
from normalcrafter_clean.windows import Window, build_windows, overlap_length
|
||||
|
||||
|
||||
def test_exact_single_window():
|
||||
# A clip exactly as long as the model context needs one window and no padding.
|
||||
assert build_windows(14, 14, 10) == [Window(0, 14)]
|
||||
|
||||
|
||||
def test_tail_is_anchored():
|
||||
# The regular cadence would leave frame 24 uncovered. The scheduler appends a
|
||||
# final full-size window beginning at 11 so its end lands exactly on frame 25.
|
||||
assert build_windows(25, 14, 10) == [Window(0, 14), Window(10, 24), Window(11, 25)]
|
||||
|
||||
|
||||
def test_overlap():
|
||||
# Regular overlap is 14 - 10 = 4. The tail-anchored final window may overlap
|
||||
# more strongly; here [10,24) and [11,25) share thirteen frames.
|
||||
assert overlap_length(Window(0, 14), Window(10, 24)) == 4
|
||||
assert overlap_length(Window(10, 24), Window(11, 25)) == 13
|
||||
Reference in New Issue
Block a user