diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..4cba47b --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.py[cod] +.pytest_cache/ +.vscode/ +.idea/ diff --git a/NOTICE b/NOTICE new file mode 100644 index 0000000..02010c4 --- /dev/null +++ b/NOTICE @@ -0,0 +1,15 @@ +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. diff --git a/README.md b/README.md new file mode 100644 index 0000000..eee2491 --- /dev/null +++ b/README.md @@ -0,0 +1,74 @@ +# ComfyUI NormalCrafter Clean + +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. + +## Nodes + +### NormalCrafter Clean - Load + +Loads: + +- `Yanrui95/NormalCrafter`: NormalCrafter UNet and normal VAE; +- `stabilityai/stable-video-diffusion-img2vid-xt`: CLIP vision encoder and Euler scheduler. + +The model is initially kept on CPU. `auto` dtype selects FP16 on CUDA and FP32 on CPU. + +### 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 +``` + +Restart ComfyUI afterward. The first loader execution downloads the Hugging Face weights unless `local_files_only` is enabled. + +## Defaults + +- 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 + +The final normal vectors are renormalized after resizing by default. Disable this only for strict comparisons with older glue code. + +## Validation status + +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. + +## Attribution + +NormalCrafter research code: MIT License. Hugging Face Diffusers-derived UNet execution structure: Apache License 2.0. See `NOTICE`. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..b51af66 --- /dev/null +++ b/__init__.py @@ -0,0 +1,9 @@ +# 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 = {} + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..b4e3ec7 --- /dev/null +++ b/nodes.py @@ -0,0 +1,177 @@ +from __future__ import annotations + +import logging + +import torch + +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__) + + +def _resolve_dtype(name: str, device: torch.device) -> torch.dtype: + 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: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "dtype": (["auto", "float16", "bfloat16", "float32"], {"default": "auto"}), + "attention": (["auto", "default", "xformers"], {"default": "auto"}), + "local_files_only": ("BOOLEAN", {"default": False}), + }, + "optional": { + "model_repo": ("STRING", {"default": "Yanrui95/NormalCrafter"}), + "base_repo": ( + "STRING", + {"default": "stabilityai/stable-video-diffusion-img2vid-xt"}, + ), + }, + } + + RETURN_TYPES = ("NORMALCRAFTER_CLEAN_MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "load" + CATEGORY = "NormalCrafter/Clean" + + 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", + ): + 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 (model,) + + +class NormalCrafterCleanGenerate: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("NORMALCRAFTER_CLEAN_MODEL",), + "images": ("IMAGE",), + "max_resolution": ( + "INT", + {"default": 1024, "min": 256, "max": 2048, "step": 64}, + ), + "window_size": ("INT", {"default": 14, "min": 2, "max": 32}), + "step_size": ("INT", {"default": 10, "min": 1, "max": 32}), + "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}), + "offload_mode": (["staged", "resident"], {"default": "staged"}), + "offload_after": ("BOOLEAN", {"default": True}), + "output_size": (["original", "processed"], {"default": "original"}), + "renormalize_normals": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("normal_maps",) + FUNCTION = "generate" + CATEGORY = "NormalCrafter/Clean" + + def generate( + self, + model: NormalCrafterModel, + images: torch.Tensor, + 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, + ): + config = InferenceConfig( + 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, + ) + config.validate() + + frame_count = int(images.shape[0]) + effective_count = max(frame_count, window_size) + 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: + progress_bar.update(1) + + device = comfy.model_management.get_torch_device() + output = model.infer(images, device=device, config=config, progress=progress) + if offload_after: + comfy.model_management.soft_empty_cache() + return (output,) + + +class NormalCrafterCleanOffload: + @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,) + + +NODE_CLASS_MAPPINGS = { + "NormalCrafterCleanLoader": NormalCrafterCleanLoader, + "NormalCrafterCleanGenerate": NormalCrafterCleanGenerate, + "NormalCrafterCleanOffload": NormalCrafterCleanOffload, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "NormalCrafterCleanLoader": "NormalCrafter Clean - Load", + "NormalCrafterCleanGenerate": "NormalCrafter Clean - Generate Normals", + "NormalCrafterCleanOffload": "NormalCrafter Clean - Offload", +} diff --git a/normalcrafter_clean/__init__.py b/normalcrafter_clean/__init__.py new file mode 100644 index 0000000..1233944 --- /dev/null +++ b/normalcrafter_clean/__init__.py @@ -0,0 +1 @@ +"""Clean NormalCrafter inference package.""" diff --git a/normalcrafter_clean/engine.py b/normalcrafter_clean/engine.py new file mode 100644 index 0000000..5355b45 --- /dev/null +++ b/normalcrafter_clean/engine.py @@ -0,0 +1,493 @@ +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__) +ProgressCallback = Callable[[str, int, int], None] +OffloadMode = Literal["staged", "resident"] +OutputSize = Literal["original", "processed"] + + +@dataclass(frozen=True, slots=True) +class InferenceConfig: + 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: + 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") + 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: + """Single-owner NormalCrafter inference model. + + 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. + """ + + 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: + self.unet = unet.eval() + self.vae = vae.eval() + self.image_encoder = image_encoder.eval() + self.feature_extractor = feature_extractor + self.scheduler = scheduler + self.dtype = dtype + self.attention_mode = attention_mode + 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": + 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, + ) + + 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, + ) + model.offload_to_cpu(empty_cache=False) + return model + + @property + def modules(self) -> dict[str, torch.nn.Module]: + return { + "image_encoder": self.image_encoder, + "vae": self.vae, + "unet": self.unet, + } + + def _configure_attention(self, mode: str) -> None: + 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: + 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: + try: + return next(module.parameters()).device + except StopIteration: + return torch.device("cpu") + + def _activate(self, name: str, device: torch.device, offload_mode: OffloadMode) -> None: + 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: + 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: + 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: + config.validate() + source = FrameSource(images_bhwc, max_resolution=config.max_resolution) + effective_frames = source.effective_frame_count(config.window_size) + completed = False + + with self._lock, torch.inference_mode(): + try: + clip_embeddings = self._encode_clip( + source, + effective_frames, + device, + config, + progress, + ) + rgb_latents = self._encode_vae( + source, + effective_frames, + device, + config, + progress, + ) + normal_latents = self._predict_latents( + rgb_latents, + clip_embeddings, + device, + config, + progress, + ) + del rgb_latents, clip_embeddings + + output = self._decode_normals( + normal_latents, + source, + device, + config, + progress, + ) + completed = True + return output + finally: + 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: + 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) + 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) + frames = source.get_bchw(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=self.dtype) + embeddings = self.image_encoder(pixel_values=pixels).image_embeds + if embeddings.ndim == 2: + embeddings = embeddings.unsqueeze(1) + embeddings_cpu = embeddings.detach().to(device="cpu", dtype=self.dtype) + + if output is None: + output = torch.empty( + (frame_count, *embeddings_cpu.shape[1:]), + dtype=self.dtype, + device="cpu", + ) + output[start:end].copy_(embeddings_cpu) + 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: + 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 + + 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) + 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) + + 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: + values = [self.FPS_CONDITIONING, self.MOTION_BUCKET_ID, self.NOISE_AUG_STRENGTH] + 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: + 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) + 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): + image_latents = rgb_latents_cpu[window.start : window.end].to( + device=device, + dtype=self.dtype, + ).unsqueeze(0) + embeddings = clip_embeddings_cpu[window.start : window.end].to( + device=device, + dtype=self.dtype, + ) + + 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: + previous_offset = window.start - previous_window.start + old_overlap = previous_latents[:, previous_offset : previous_offset + overlap] + latents[:, :overlap].copy_(old_overlap) + + self.scheduler.set_timesteps(1, device=device) + timestep = self.scheduler.timesteps[0] + model_input = self.scheduler.scale_model_input(latents, timestep) + 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: + 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) + + output_cpu[window.start : window.end].copy_(predicted[0].detach().to("cpu")) + 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: + 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 + + 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 + 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) + latents = latents / scaling_factor + decode_kwargs = {"num_frames": end - start} if accepts_num_frames else {} + normals = self.vae.decode(latents, **decode_kwargs).sample.float().to("cpu") + 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, + ) + + normals = normals.clamp(-1.0, 1.0) + if config.renormalize_normals: + 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()) + + del latents, normals + if progress: + progress("VAE decode", chunk_index, total) + finally: + if needs_upcast: + self.vae.to(dtype=original_dtype) + + return output diff --git a/normalcrafter_clean/preprocess.py b/normalcrafter_clean/preprocess.py new file mode 100644 index 0000000..08e2897 --- /dev/null +++ b/normalcrafter_clean/preprocess.py @@ -0,0 +1,104 @@ +from __future__ import annotations + +from dataclasses import dataclass +import math + +import torch +import torch.nn.functional as F + + +def _round_size(value: float) -> int: + return max(1, int(round(value))) + + +def _ceil_multiple(value: int, multiple: int) -> int: + return int(math.ceil(value / multiple) * multiple) + + +@dataclass +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. + """ + + frames_bhwc: torch.Tensor + max_resolution: int + pad_multiple: int = 64 + + def __post_init__(self) -> None: + 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") + + self.frames_bhwc = ( + self.frames_bhwc[..., :3] + .detach() + .to(device="cpu", dtype=torch.float32) + .contiguous() + .clamp_(0.0, 1.0) + ) + + _, 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) + + self.padded_height = _ceil_multiple(self.resized_height, self.pad_multiple) + self.padded_width = _ceil_multiple(self.resized_width, self.pad_multiple) + + 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: + return int(self.frames_bhwc.shape[0]) + + def effective_frame_count(self, window_size: int) -> int: + 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]. + + Indices beyond the real clip repeat its final frame. This is only used to + satisfy the temporal window for short clips. + """ + if start < 0 or end <= start: + raise ValueError(f"invalid frame range [{start}, {end})") + + indices = torch.arange(start, end, dtype=torch.long).clamp_max(self.frame_count - 1) + chunk = self.frames_bhwc.index_select(0, indices).permute(0, 3, 1, 2).contiguous() + + 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) + + 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: + return tensor_bchw[ + :, + :, + self.pad_top : self.pad_top + self.resized_height, + self.pad_left : self.pad_left + self.resized_width, + ] diff --git a/normalcrafter_clean/unet.py b/normalcrafter_clean/unet.py new file mode 100644 index 0000000..3dcfa2c --- /dev/null +++ b/normalcrafter_clean/unet.py @@ -0,0 +1,126 @@ +# 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 + +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, + ) + + +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. + """ + + def forward( + self, + sample: torch.Tensor, + timestep: Union[torch.Tensor, float, int], + encoder_hidden_states: torch.Tensor, + added_time_ids: torch.Tensor, + return_dict: bool = True, + ): + timesteps = timestep + if not torch.is_tensor(timesteps): + dtype = torch.float64 if isinstance(timestep, float) else torch.int64 + if sample.device.type == "mps": + dtype = torch.float32 if isinstance(timestep, float) else torch.int32 + timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device) + elif timesteps.ndim == 0: + timesteps = timesteps[None].to(sample.device) + + batch_size, num_frames = sample.shape[:2] + timesteps = timesteps.expand(batch_size) + + t_emb = self.time_proj(timesteps).to(dtype=sample.dtype) + emb = self.time_embedding(t_emb) + + 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) + emb = emb.repeat_interleave(num_frames, dim=0) + + sample = sample.flatten(0, 1) + + 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}" + ) + + sample = self.conv_in(sample) + image_only_indicator = torch.zeros( + batch_size, + num_frames, + dtype=sample.dtype, + device=sample.device, + ) + + down_block_res_samples = (sample,) + for downsample_block in self.down_blocks: + if getattr(downsample_block, "has_cross_attention", False): + sample, residuals = downsample_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + else: + sample, residuals = downsample_block( + hidden_states=sample, + temb=emb, + image_only_indicator=image_only_indicator, + ) + down_block_res_samples += residuals + + sample = self.mid_block( + hidden_states=sample, + temb=emb, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + + 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)] + + if getattr(upsample_block, "has_cross_attention", False): + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=residuals, + encoder_hidden_states=encoder_hidden_states, + image_only_indicator=image_only_indicator, + ) + else: + sample = upsample_block( + hidden_states=sample, + temb=emb, + res_hidden_states_tuple=residuals, + image_only_indicator=image_only_indicator, + ) + + sample = self.conv_out(self.conv_act(self.conv_norm_out(sample))) + sample = sample.reshape(batch_size, num_frames, *sample.shape[1:]) + + if not return_dict: + return (sample,) + return UNetSpatioTemporalConditionOutput(sample=sample) diff --git a/normalcrafter_clean/windows.py b/normalcrafter_clean/windows.py new file mode 100644 index 0000000..1d30b73 --- /dev/null +++ b/normalcrafter_clean/windows.py @@ -0,0 +1,43 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class Window: + start: int + end: int + + @property + def length(self) -> int: + 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. + + 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. + """ + 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") + + last_start = num_frames - window_size + starts = list(range(0, last_start + 1, step_size)) + 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: + if previous is None: + return 0 + return max(0, previous.end - current.start) diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..4cd3152 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +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 diff --git a/tests/test_windows.py b/tests/test_windows.py new file mode 100644 index 0000000..79fbce8 --- /dev/null +++ b/tests/test_windows.py @@ -0,0 +1,14 @@ +from normalcrafter_clean.windows import Window, build_windows, overlap_length + + +def test_exact_single_window(): + assert build_windows(14, 14, 10) == [Window(0, 14)] + + +def test_tail_is_anchored(): + assert build_windows(25, 14, 10) == [Window(0, 14), Window(10, 24), Window(11, 25)] + + +def test_overlap(): + assert overlap_length(Window(0, 14), Window(10, 24)) == 4 + assert overlap_length(Window(10, 24), Window(11, 25)) == 13