From 8023ed7258507f69caeba8d71fedd8ef81c9944b Mon Sep 17 00:00:00 2001 From: Korbinian <23-23enterprise@gmx.de> Date: Tue, 7 Jul 2026 13:49:58 +0200 Subject: [PATCH] commented code --- ARCHITECTURE.md | 155 +++++++++++++ LICENSE | 1 + README.md | 11 + nodes.py | 74 ++++++ normalcrafter_clean/engine.py | 360 ++++++++++++++++++++++++++++-- normalcrafter_clean/preprocess.py | 100 ++++++++- normalcrafter_clean/unet.py | 71 +++++- normalcrafter_clean/windows.py | 46 +++- tests/test_windows.py | 12 + 9 files changed, 795 insertions(+), 35 deletions(-) create mode 100644 ARCHITECTURE.md diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md new file mode 100644 index 0000000..21b06d0 --- /dev/null +++ b/ARCHITECTURE.md @@ -0,0 +1,155 @@ +# 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. diff --git a/LICENSE b/LICENSE index 261eeb9..d645695 100644 --- a/LICENSE +++ b/LICENSE @@ -1,3 +1,4 @@ + Apache License Version 2.0, January 2004 http://www.apache.org/licenses/ diff --git a/README.md b/README.md index eee2491..394f39a 100644 --- a/README.md +++ b/README.md @@ -72,3 +72,14 @@ The source is syntax-checked and the pure window scheduler has unit tests. It ha ## Attribution NormalCrafter research code: MIT License. Hugging Face Diffusers-derived UNet execution structure: Apache License 2.0. See `NOTICE`. + +## Commented edition + +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: + +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. diff --git a/nodes.py b/nodes.py index b4e3ec7..5642a3c 100644 --- a/nodes.py +++ b/nodes.py @@ -1,3 +1,15 @@ +"""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 @@ -14,23 +26,41 @@ logger = logging.getLogger(__name__) 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. + """ + @classmethod def INPUT_TYPES(cls): + # ComfyUI reads this declarative structure to construct the node UI. return { "required": { "dtype": (["auto", "float16", "bfloat16", "float32"], {"default": "auto"}), @@ -38,6 +68,8 @@ class NormalCrafterCleanLoader: "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", @@ -59,8 +91,12 @@ class NormalCrafterCleanLoader: 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``. 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(), @@ -72,24 +108,40 @@ class NormalCrafterCleanLoader: 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. + """ + @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. "max_resolution": ( "INT", {"default": 1024, "min": 256, "max": 2048, "step": 64}, ), + # 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}), } } @@ -114,6 +166,8 @@ class NormalCrafterCleanGenerate: output_size: str, renormalize_normals: bool, ): + """Validate user settings, build progress accounting, and invoke inference.""" + config = InferenceConfig( max_resolution=max_resolution, window_size=window_size, @@ -128,8 +182,14 @@ class NormalCrafterCleanGenerate: ) 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 @@ -139,16 +199,28 @@ class NormalCrafterCleanGenerate: 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() + 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",)}} @@ -164,6 +236,8 @@ class NormalCrafterCleanOffload: 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, diff --git a/normalcrafter_clean/engine.py b/normalcrafter_clean/engine.py index 5355b45..cb31391 100644 --- a/normalcrafter_clean/engine.py +++ b/normalcrafter_clean/engine.py @@ -1,3 +1,18 @@ +"""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 @@ -17,13 +32,28 @@ 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 @@ -36,12 +66,18 @@ class InferenceConfig: 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), @@ -52,13 +88,23 @@ class InferenceConfig: class NormalCrafterModel: - """Single-owner NormalCrafter inference model. + """Own and execute every model component required for inference. - No module globals, duplicate pipeline references, or hidden CUDA output tensors. - All long-lived intermediate video data is stored on CPU. GPU residency is - explicitly controlled per inference stage. + 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 @@ -74,14 +120,23 @@ class NormalCrafterModel: 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 @@ -94,13 +149,22 @@ class NormalCrafterModel: 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, @@ -108,9 +172,12 @@ class NormalCrafterModel: **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, @@ -145,11 +212,16 @@ class NormalCrafterModel: 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, @@ -157,26 +229,45 @@ class NormalCrafterModel: } 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: @@ -187,17 +278,28 @@ class NormalCrafterModel: 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(): @@ -211,13 +313,32 @@ class NormalCrafterModel: 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, @@ -225,6 +346,8 @@ class NormalCrafterModel: config, progress, ) + + # Stage 2: spatial RGB latent per frame, stored on CPU. rgb_latents = self._encode_vae( source, effective_frames, @@ -232,6 +355,8 @@ class NormalCrafterModel: config, progress, ) + + # Stage 3: sliding-window temporal prediction, stored on CPU. normal_latents = self._predict_latents( rgb_latents, clip_embeddings, @@ -239,8 +364,12 @@ class NormalCrafterModel: 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, @@ -248,9 +377,12 @@ class NormalCrafterModel: 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) @@ -262,15 +394,41 @@ class NormalCrafterModel: 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) - 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) + + # 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): + 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), @@ -278,12 +436,22 @@ class NormalCrafterModel: 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:]), @@ -291,7 +459,10 @@ class NormalCrafterModel: 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) @@ -306,26 +477,51 @@ class NormalCrafterModel: 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 - needs_upcast = original_dtype == torch.float16 and bool(getattr(self.vae.config, "force_upcast", False)) + + # 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 + 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) - frames = (frames * 2.0 - 1.0).to(device=device, dtype=compute_dtype) + + # 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:]), @@ -333,7 +529,9 @@ class NormalCrafterModel: device="cpu", ) output[start:end].copy_(latents_cpu) + del frames, latents, latents_cpu + if progress: progress("VAE encode", chunk_index, total) finally: @@ -344,13 +542,23 @@ class NormalCrafterModel: return output def _added_time_ids(self, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor: - values = [self.FPS_CONDITIONING, self.MOTION_BUCKET_ID, self.NOISE_AUG_STRENGTH] + """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( @@ -361,40 +569,78 @@ class NormalCrafterModel: 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:]), + ( + 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] + 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) @@ -405,9 +651,16 @@ class NormalCrafterModel: added_time_ids=added_time_ids, return_dict=False, )[0] - predicted = self.scheduler.step(noise_prediction, timestep, latents).prev_sample + + 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, @@ -415,13 +668,26 @@ class NormalCrafterModel: device=device, dtype=self.dtype, )[1:-1].view(1, overlap, 1, 1, 1) - predicted[:, :overlap] = old_overlap * weights + predicted[:, :overlap] * (1.0 - weights) - output_cpu[window.start : window.end].copy_(predicted[0].detach().to("cpu")) + 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)) @@ -435,9 +701,15 @@ class NormalCrafterModel: 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)) + 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 @@ -447,15 +719,28 @@ class NormalCrafterModel: 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) - forward_vae = self.vae._orig_mod.forward if hasattr(self.vae, "_orig_mod") else self.vae.forward + + # 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 + + total = ( + source.frame_count + config.vae_decode_chunk_size - 1 + ) // config.vae_decode_chunk_size try: for chunk_index, start in enumerate( @@ -463,13 +748,31 @@ class NormalCrafterModel: 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) + + 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_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): + if (out_height, out_width) != ( + source.resized_height, + source.resized_width, + ): normals = F.interpolate( normals, size=(out_height, out_width), @@ -477,13 +780,24 @@ class NormalCrafterModel: 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) - output[start:end].copy_(normals.permute(0, 2, 3, 1).contiguous()) + + # 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: diff --git a/normalcrafter_clean/preprocess.py b/normalcrafter_clean/preprocess.py index 08e2897..30111ea 100644 --- a/normalcrafter_clean/preprocess.py +++ b/normalcrafter_clean/preprocess.py @@ -1,3 +1,18 @@ +"""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 @@ -8,10 +23,19 @@ 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) @@ -19,8 +43,23 @@ def _ceil_multiple(value: int, multiple: int) -> int: class FrameSource: """Lazy CPU-side view of a ComfyUI IMAGE batch. - Frames are resized and padded only for the chunk currently needed. This avoids - materialising a second full-resolution copy of the whole video. + 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 @@ -28,6 +67,10 @@ class FrameSource: 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: @@ -35,6 +78,10 @@ class FrameSource: 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() @@ -43,14 +90,20 @@ class FrameSource: .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 @@ -60,23 +113,52 @@ class FrameSource: @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: - """Return a resized, symmetrically white-padded CPU chunk in [0, 1]. + """Materialize one processed CPU frame range in BCHW layout. - Indices beyond the real clip repeat its final frame. This is only used to - satisfy the temporal window for short clips. + 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, @@ -86,6 +168,8 @@ class FrameSource: 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, @@ -96,6 +180,12 @@ class FrameSource: 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[ :, :, diff --git a/normalcrafter_clean/unet.py b/normalcrafter_clean/unet.py index 3dcfa2c..0515810 100644 --- a/normalcrafter_clean/unet.py +++ b/normalcrafter_clean/unet.py @@ -1,3 +1,15 @@ +"""Inference-only NormalCrafter adaptation of Diffusers' 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 @@ -7,6 +19,9 @@ 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, @@ -20,10 +35,21 @@ except ImportError: # older Diffusers class NormalCrafterUNet(UNetSpatioTemporalConditionModel): """SVD UNet accepting one CLIP embedding per video frame. - Upstream SVD repeats one image embedding over the full video. NormalCrafter - instead passes ``batch * frames`` embeddings. Everything else is ordinary - inference-only UNet execution; training-only DINO and ControlNet branches are - deliberately excluded. + 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. """ def forward( @@ -34,6 +60,15 @@ 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 @@ -44,18 +79,32 @@ 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) @@ -66,7 +115,12 @@ class NormalCrafterUNet(UNetSpatioTemporalConditionModel): f"batch={batch_size}, frames={num_frames}" ) + # 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, @@ -74,6 +128,8 @@ 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): @@ -91,6 +147,8 @@ 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, @@ -98,6 +156,8 @@ 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)] @@ -118,9 +178,12 @@ 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) diff --git a/normalcrafter_clean/windows.py b/normalcrafter_clean/windows.py index 1d30b73..e2878b1 100644 --- a/normalcrafter_clean/windows.py +++ b/normalcrafter_clean/windows.py @@ -1,3 +1,11 @@ +"""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 @@ -5,20 +13,39 @@ 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]: - """Return monotonically increasing, fixed-size windows covering every frame. + """Build fixed-size windows that cover every frame without temporal gaps. - Short clips are expected to be padded by the caller to at least ``window_size``. - The final window is anchored to the end so no tail frames are skipped. + ``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: @@ -30,14 +57,27 @@ def build_windows(num_frames: int, window_size: int, step_size: int) -> list[Win 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) diff --git a/tests/test_windows.py b/tests/test_windows.py index 79fbce8..25011e7 100644 --- a/tests/test_windows.py +++ b/tests/test_windows.py @@ -1,14 +1,26 @@ +"""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