808 lines
30 KiB
Python
808 lines
30 KiB
Python
"""Memory-conscious NormalCrafter inference engine.
|
|
|
|
The engine is intentionally organized as four explicit stages:
|
|
|
|
1. CLIP image encoding: RGB frames -> semantic frame embeddings.
|
|
2. VAE encoding: RGB frames -> spatial RGB latents.
|
|
3. UNet inference: RGB latents + CLIP embeddings -> normal-map latents.
|
|
4. VAE decoding: normal-map latents -> three-channel normal maps.
|
|
|
|
The central ownership rule is simple: large, long-lived intermediate videos stay
|
|
on CPU. Only the currently active chunk or temporal window is moved to the
|
|
inference device. This avoids retaining an entire video on CUDA and makes peak
|
|
VRAM mostly independent of total clip length.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import gc
|
|
import inspect
|
|
import logging
|
|
import threading
|
|
from typing import Callable, Literal
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from diffusers import AutoencoderKLTemporalDecoder, EulerDiscreteScheduler
|
|
from transformers import CLIPImageProcessor, CLIPVisionModelWithProjection
|
|
|
|
from .preprocess import FrameSource
|
|
from .unet import NormalCrafterUNet
|
|
from .windows import Window, build_windows, overlap_length
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Called after every completed chunk/window. The strings make progress reporting
|
|
# descriptive without tying the engine to ComfyUI itself.
|
|
ProgressCallback = Callable[[str, int, int], None]
|
|
|
|
# ``staged`` minimizes VRAM by moving only one model component to the device.
|
|
# ``resident`` avoids repeated device transfers but keeps every component loaded.
|
|
OffloadMode = Literal["staged", "resident"]
|
|
|
|
# ``original`` restores the source dimensions after decoding. ``processed`` keeps
|
|
# the aspect-preserving, max-resolution-limited dimensions used by the model.
|
|
OutputSize = Literal["original", "processed"]
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class InferenceConfig:
|
|
"""Immutable runtime settings for one inference call.
|
|
|
|
Keeping these values in a frozen dataclass prevents an inference run from
|
|
observing settings that change halfway through execution.
|
|
"""
|
|
|
|
max_resolution: int = 1024
|
|
window_size: int = 14
|
|
step_size: int = 10
|
|
clip_chunk_size: int = 16
|
|
vae_encode_chunk_size: int = 7
|
|
vae_decode_chunk_size: int = 4
|
|
offload_mode: OffloadMode = "staged"
|
|
offload_after: bool = True
|
|
output_size: OutputSize = "original"
|
|
renormalize_normals: bool = True
|
|
|
|
def validate(self) -> None:
|
|
"""Reject combinations that would produce gaps or invalid tensors."""
|
|
|
|
if self.max_resolution < 64:
|
|
raise ValueError("max_resolution must be at least 64")
|
|
if self.window_size < 2:
|
|
raise ValueError("window_size must be at least 2")
|
|
|
|
# A step larger than the window would leave uncovered frames between
|
|
# windows. A step of zero would never advance.
|
|
if not 1 <= self.step_size <= self.window_size:
|
|
raise ValueError("step_size must be between 1 and window_size")
|
|
|
|
for name, value in (
|
|
("clip_chunk_size", self.clip_chunk_size),
|
|
("vae_encode_chunk_size", self.vae_encode_chunk_size),
|
|
("vae_decode_chunk_size", self.vae_decode_chunk_size),
|
|
):
|
|
if value < 1:
|
|
raise ValueError(f"{name} must be positive")
|
|
|
|
|
|
class NormalCrafterModel:
|
|
"""Own and execute every model component required for inference.
|
|
|
|
Ownership guarantees:
|
|
|
|
* No module-level model globals.
|
|
* No duplicate pipeline object pointing at the same modules.
|
|
* No CUDA tensor is returned to ComfyUI.
|
|
* Intermediate full-video tensors live on CPU.
|
|
* A re-entrant lock serializes device moves and inference on this instance.
|
|
|
|
The object may stay alive in a ComfyUI graph cache, but its resource state is
|
|
explicit and controllable through ``offload_to_cpu``.
|
|
"""
|
|
|
|
# These values reproduce the conditioning used by the released pipeline.
|
|
# They are inherited from the Stable Video Diffusion interface, even though
|
|
# NormalCrafter performs a deterministic one-step prediction.
|
|
FPS_CONDITIONING = 7
|
|
MOTION_BUCKET_ID = 127
|
|
NOISE_AUG_STRENGTH = 0.0
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
unet: NormalCrafterUNet,
|
|
vae: AutoencoderKLTemporalDecoder,
|
|
image_encoder: CLIPVisionModelWithProjection,
|
|
feature_extractor: CLIPImageProcessor,
|
|
scheduler: EulerDiscreteScheduler,
|
|
dtype: torch.dtype,
|
|
attention_mode: str,
|
|
) -> None:
|
|
# ``eval`` disables training behavior such as dropout and documents that
|
|
# these modules are inference-only in this implementation.
|
|
self.unet = unet.eval()
|
|
self.vae = vae.eval()
|
|
self.image_encoder = image_encoder.eval()
|
|
|
|
# The feature extractor only provides normalization metadata here; image
|
|
# resizing is done directly in Torch to avoid PIL/NumPy round-trips.
|
|
self.feature_extractor = feature_extractor
|
|
self.scheduler = scheduler
|
|
self.dtype = dtype
|
|
self.attention_mode = attention_mode
|
|
|
|
# Device moves are not thread-safe when two graph executions share one
|
|
# model instance. The lock also protects scheduler mutation.
|
|
self._lock = threading.RLock()
|
|
|
|
self._configure_attention(attention_mode)
|
|
|
|
@classmethod
|
|
def from_pretrained(
|
|
cls,
|
|
*,
|
|
model_repo: str = "Yanrui95/NormalCrafter",
|
|
base_repo: str = "stabilityai/stable-video-diffusion-img2vid-xt",
|
|
dtype: torch.dtype = torch.float16,
|
|
attention_mode: str = "auto",
|
|
local_files_only: bool = False,
|
|
) -> "NormalCrafterModel":
|
|
"""Load the released weights and assemble one clean model owner.
|
|
|
|
NormalCrafter supplies its fine-tuned UNet and VAE. The CLIP vision
|
|
encoder, preprocessing statistics, and Euler scheduler come from the SVD
|
|
base repository used during training.
|
|
"""
|
|
|
|
common = {
|
|
"torch_dtype": dtype,
|
|
"low_cpu_mem_usage": True,
|
|
"local_files_only": local_files_only,
|
|
}
|
|
|
|
logger.info("Loading NormalCrafter UNet from %s", model_repo)
|
|
unet = NormalCrafterUNet.from_pretrained(model_repo, subfolder="unet", **common)
|
|
|
|
logger.info("Loading NormalCrafter VAE from %s", model_repo)
|
|
vae = AutoencoderKLTemporalDecoder.from_pretrained(
|
|
model_repo,
|
|
subfolder="vae",
|
|
**common,
|
|
)
|
|
|
|
# Some Hugging Face repos expose an explicit fp16 variant. Try it first
|
|
# for half-precision execution, then fall back to the default files.
|
|
encoder_kwargs = dict(common)
|
|
if dtype in (torch.float16, torch.bfloat16):
|
|
encoder_kwargs["variant"] = "fp16"
|
|
|
|
try:
|
|
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
|
base_repo,
|
|
subfolder="image_encoder",
|
|
**encoder_kwargs,
|
|
)
|
|
except (OSError, ValueError, TypeError):
|
|
encoder_kwargs.pop("variant", None)
|
|
image_encoder = CLIPVisionModelWithProjection.from_pretrained(
|
|
base_repo,
|
|
subfolder="image_encoder",
|
|
**encoder_kwargs,
|
|
)
|
|
|
|
feature_extractor = CLIPImageProcessor.from_pretrained(
|
|
base_repo,
|
|
subfolder="feature_extractor",
|
|
local_files_only=local_files_only,
|
|
)
|
|
scheduler = EulerDiscreteScheduler.from_pretrained(
|
|
base_repo,
|
|
subfolder="scheduler",
|
|
local_files_only=local_files_only,
|
|
)
|
|
|
|
model = cls(
|
|
unet=unet,
|
|
vae=vae,
|
|
image_encoder=image_encoder,
|
|
feature_extractor=feature_extractor,
|
|
scheduler=scheduler,
|
|
dtype=dtype,
|
|
attention_mode=attention_mode,
|
|
)
|
|
|
|
# Loading libraries may choose a device based on environment/config. The
|
|
# public lifecycle contract is that a newly loaded model starts on CPU.
|
|
model.offload_to_cpu(empty_cache=False)
|
|
return model
|
|
|
|
@property
|
|
def modules(self) -> dict[str, torch.nn.Module]:
|
|
"""Return the heavy components participating in device management."""
|
|
|
|
return {
|
|
"image_encoder": self.image_encoder,
|
|
"vae": self.vae,
|
|
"unet": self.unet,
|
|
}
|
|
|
|
def _configure_attention(self, mode: str) -> None:
|
|
"""Enable xFormers when available or explicitly requested."""
|
|
|
|
if mode not in {"auto", "default", "xformers"}:
|
|
raise ValueError(f"unsupported attention mode: {mode}")
|
|
|
|
if mode == "default":
|
|
return
|
|
|
|
try:
|
|
self.unet.enable_xformers_memory_efficient_attention()
|
|
logger.info("NormalCrafter xFormers attention enabled")
|
|
except Exception as exc:
|
|
# Explicit requests should fail loudly. ``auto`` is allowed to fall
|
|
# back because current PyTorch/Diffusers attention may be sufficient.
|
|
if mode == "xformers":
|
|
raise RuntimeError("xFormers was requested but could not be enabled") from exc
|
|
logger.info("xFormers unavailable; using Diffusers default attention: %s", exc)
|
|
|
|
@staticmethod
|
|
def _device_of(module: torch.nn.Module) -> torch.device:
|
|
"""Read a module's current device from its first parameter."""
|
|
|
|
try:
|
|
return next(module.parameters()).device
|
|
except StopIteration:
|
|
# Defensive fallback for parameterless modules.
|
|
return torch.device("cpu")
|
|
|
|
def _activate(self, name: str, device: torch.device, offload_mode: OffloadMode) -> None:
|
|
"""Place model components according to the selected residency policy.
|
|
|
|
``resident``:
|
|
Move all heavy modules to the inference device and leave them there.
|
|
|
|
``staged``:
|
|
Move only the named component to the inference device and force the
|
|
others to CPU. After those moves, release now-unused CUDA cache blocks.
|
|
"""
|
|
|
|
if offload_mode == "resident":
|
|
for module in self.modules.values():
|
|
if self._device_of(module) != device:
|
|
module.to(device)
|
|
return
|
|
|
|
for module_name, module in self.modules.items():
|
|
target = device if module_name == name else torch.device("cpu")
|
|
if self._device_of(module) != target:
|
|
module.to(target)
|
|
|
|
self._release_cuda_cache(device)
|
|
|
|
@staticmethod
|
|
def _release_cuda_cache(device: torch.device) -> None:
|
|
"""Release unreachable Python objects and unused CUDA allocator blocks.
|
|
|
|
``empty_cache`` does not free live tensors. Its usefulness depends on the
|
|
preceding ownership discipline: temporary tensors must first go out of
|
|
scope or be deleted, and inactive modules must already have moved to CPU.
|
|
"""
|
|
|
|
gc.collect()
|
|
if device.type == "cuda" and torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
def offload_to_cpu(self, *, empty_cache: bool = True) -> None:
|
|
"""Move all owned model weights to CPU and optionally empty CUDA cache."""
|
|
|
|
for module in self.modules.values():
|
|
module.to("cpu")
|
|
|
|
if empty_cache:
|
|
gc.collect()
|
|
if torch.cuda.is_available():
|
|
torch.cuda.empty_cache()
|
|
|
|
def infer(
|
|
self,
|
|
images_bhwc: torch.Tensor,
|
|
*,
|
|
device: torch.device,
|
|
config: InferenceConfig,
|
|
progress: ProgressCallback | None = None,
|
|
) -> torch.Tensor:
|
|
"""Run complete inference and return CPU normal maps in ComfyUI format.
|
|
|
|
Input shape:
|
|
``[frames, height, width, channels]``, values expected in ``[0, 1]``.
|
|
|
|
Output shape:
|
|
``[real_frames, output_height, output_width, 3]`` on CPU, values in
|
|
``[0, 1]`` where encoded RGB corresponds to XYZ normal direction.
|
|
"""
|
|
|
|
config.validate()
|
|
|
|
# FrameSource owns the canonical CPU input and lazily prepares only the
|
|
# chunk needed by the current stage.
|
|
source = FrameSource(images_bhwc, max_resolution=config.max_resolution)
|
|
|
|
# The temporal UNet requires at least one full window. Short clips repeat
|
|
# their final frame internally; decoding is still limited to real frames.
|
|
effective_frames = source.effective_frame_count(config.window_size)
|
|
completed = False
|
|
|
|
# ``torch.inference_mode`` removes autograd bookkeeping more aggressively
|
|
# than ``no_grad`` and prevents accidental graph retention.
|
|
with self._lock, torch.inference_mode():
|
|
try:
|
|
# Stage 1: semantic embedding per frame, stored on CPU.
|
|
clip_embeddings = self._encode_clip(
|
|
source,
|
|
effective_frames,
|
|
device,
|
|
config,
|
|
progress,
|
|
)
|
|
|
|
# Stage 2: spatial RGB latent per frame, stored on CPU.
|
|
rgb_latents = self._encode_vae(
|
|
source,
|
|
effective_frames,
|
|
device,
|
|
config,
|
|
progress,
|
|
)
|
|
|
|
# Stage 3: sliding-window temporal prediction, stored on CPU.
|
|
normal_latents = self._predict_latents(
|
|
rgb_latents,
|
|
clip_embeddings,
|
|
device,
|
|
config,
|
|
progress,
|
|
)
|
|
|
|
# These full-video intermediates are no longer needed before the
|
|
# decoder stage. Delete references early to reduce host memory.
|
|
del rgb_latents, clip_embeddings
|
|
|
|
# Stage 4: decode only real input frames into final normals.
|
|
output = self._decode_normals(
|
|
normal_latents,
|
|
source,
|
|
device,
|
|
config,
|
|
progress,
|
|
)
|
|
|
|
completed = True
|
|
return output
|
|
finally:
|
|
# On failure we always offload. On success this follows the user's
|
|
# setting, allowing ``resident`` mode to remain fast across runs.
|
|
if config.offload_after or not completed:
|
|
self.offload_to_cpu(empty_cache=True)
|
|
|
|
def _encode_clip(
|
|
self,
|
|
source: FrameSource,
|
|
frame_count: int,
|
|
device: torch.device,
|
|
config: InferenceConfig,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
"""Encode every frame into one CLIP vision embedding on CPU.
|
|
|
|
Shape flow, typically:
|
|
|
|
``[chunk, 3, H, W]`` -> ``[chunk, 3, 224, 224]`` ->
|
|
``[chunk, 1, embedding_dim]``.
|
|
"""
|
|
|
|
self._activate("image_encoder", device, config.offload_mode)
|
|
|
|
# Reproduce the CLIP normalization configured by the SVD repository.
|
|
# Mean/std remain on CPU because subtraction occurs before device transfer.
|
|
mean = torch.tensor(
|
|
self.feature_extractor.image_mean,
|
|
dtype=torch.float32,
|
|
).view(1, 3, 1, 1)
|
|
std = torch.tensor(
|
|
self.feature_extractor.image_std,
|
|
dtype=torch.float32,
|
|
).view(1, 3, 1, 1)
|
|
|
|
output: torch.Tensor | None = None
|
|
total = (frame_count + config.clip_chunk_size - 1) // config.clip_chunk_size
|
|
|
|
for chunk_index, start in enumerate(
|
|
range(0, frame_count, config.clip_chunk_size),
|
|
start=1,
|
|
):
|
|
end = min(start + config.clip_chunk_size, frame_count)
|
|
|
|
# FrameSource resizes/pads lazily and repeats the final real frame for
|
|
# synthetic short-clip indices.
|
|
frames = source.get_bchw(start, end)
|
|
|
|
# CLIP vision expects its own fixed spatial resolution.
|
|
frames = F.interpolate(
|
|
frames,
|
|
size=(224, 224),
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
antialias=True,
|
|
).clamp_(0.0, 1.0)
|
|
|
|
# Only the current normalized chunk crosses to the inference device.
|
|
pixels = ((frames - mean) / std).to(device=device, dtype=self.dtype)
|
|
embeddings = self.image_encoder(pixel_values=pixels).image_embeds
|
|
|
|
# Diffusers cross-attention expects a sequence dimension. Depending on
|
|
# Transformers version, ``image_embeds`` may arrive as rank two.
|
|
if embeddings.ndim == 2:
|
|
embeddings = embeddings.unsqueeze(1)
|
|
|
|
# Copy the result back immediately so total video length does not grow
|
|
# CUDA residency. Half precision is retained to limit host memory.
|
|
embeddings_cpu = embeddings.detach().to(device="cpu", dtype=self.dtype)
|
|
|
|
# Allocate the final CPU tensor once the runtime embedding shape is
|
|
# known, then fill slices instead of accumulating and concatenating.
|
|
if output is None:
|
|
output = torch.empty(
|
|
(frame_count, *embeddings_cpu.shape[1:]),
|
|
dtype=self.dtype,
|
|
device="cpu",
|
|
)
|
|
output[start:end].copy_(embeddings_cpu)
|
|
|
|
# Drop all references to current chunk/device tensors before moving on.
|
|
del frames, pixels, embeddings, embeddings_cpu
|
|
|
|
if progress:
|
|
progress("CLIP", chunk_index, total)
|
|
|
|
assert output is not None
|
|
return output
|
|
|
|
def _encode_vae(
|
|
self,
|
|
source: FrameSource,
|
|
frame_count: int,
|
|
device: torch.device,
|
|
config: InferenceConfig,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
"""Encode RGB frames into VAE latents and store them on CPU.
|
|
|
|
Unlike generative sampling, this path uses the latent distribution mode,
|
|
making encoding deterministic and avoiding random sampling noise.
|
|
"""
|
|
|
|
self._activate("vae", device, config.offload_mode)
|
|
|
|
original_dtype = next(self.vae.parameters()).dtype
|
|
|
|
# Some Diffusers VAEs declare ``force_upcast`` because selected operations
|
|
# are unstable in FP16. The dtype is restored in ``finally`` even on error.
|
|
needs_upcast = (
|
|
original_dtype == torch.float16
|
|
and bool(getattr(self.vae.config, "force_upcast", False))
|
|
)
|
|
if needs_upcast:
|
|
self.vae.to(dtype=torch.float32)
|
|
compute_dtype = next(self.vae.parameters()).dtype
|
|
|
|
output: torch.Tensor | None = None
|
|
total = (
|
|
frame_count + config.vae_encode_chunk_size - 1
|
|
) // config.vae_encode_chunk_size
|
|
|
|
try:
|
|
for chunk_index, start in enumerate(
|
|
range(0, frame_count, config.vae_encode_chunk_size),
|
|
start=1,
|
|
):
|
|
end = min(start + config.vae_encode_chunk_size, frame_count)
|
|
|
|
frames = source.get_bchw(start, end)
|
|
|
|
# VAE input convention is [-1, 1], whereas ComfyUI uses [0, 1].
|
|
frames = (frames * 2.0 - 1.0).to(
|
|
device=device,
|
|
dtype=compute_dtype,
|
|
)
|
|
|
|
latents = self.vae.encode(frames).latent_dist.mode()
|
|
latents_cpu = latents.detach().to(device="cpu", dtype=self.dtype)
|
|
|
|
# Runtime allocation avoids assumptions about latent channel count
|
|
# or spatial scale while still eliminating list + torch.cat churn.
|
|
if output is None:
|
|
output = torch.empty(
|
|
(frame_count, *latents_cpu.shape[1:]),
|
|
dtype=self.dtype,
|
|
device="cpu",
|
|
)
|
|
output[start:end].copy_(latents_cpu)
|
|
|
|
del frames, latents, latents_cpu
|
|
|
|
if progress:
|
|
progress("VAE encode", chunk_index, total)
|
|
finally:
|
|
if needs_upcast:
|
|
self.vae.to(dtype=original_dtype)
|
|
|
|
assert output is not None
|
|
return output
|
|
|
|
def _added_time_ids(self, *, dtype: torch.dtype, device: torch.device) -> torch.Tensor:
|
|
"""Construct the three Stable Video Diffusion conditioning values."""
|
|
|
|
values = [
|
|
self.FPS_CONDITIONING,
|
|
self.MOTION_BUCKET_ID,
|
|
self.NOISE_AUG_STRENGTH,
|
|
]
|
|
|
|
# Fail explicitly if a different UNet config expects a different added
|
|
# conditioning width. Silent shape coercion here would produce bad output.
|
|
expected = self.unet.add_embedding.linear_1.in_features
|
|
actual = int(self.unet.config.addition_time_embed_dim) * len(values)
|
|
if expected != actual:
|
|
raise ValueError(
|
|
f"UNet added-time embedding mismatch: expected {expected}, constructed {actual}"
|
|
)
|
|
|
|
return torch.tensor([values], dtype=dtype, device=device)
|
|
|
|
def _predict_latents(
|
|
self,
|
|
rgb_latents_cpu: torch.Tensor,
|
|
clip_embeddings_cpu: torch.Tensor,
|
|
device: torch.device,
|
|
config: InferenceConfig,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
"""Predict normal latents over overlapping temporal windows.
|
|
|
|
Full-video input/output tensors remain on CPU. Each iteration moves only
|
|
one window of RGB latents and CLIP embeddings to the inference device.
|
|
|
|
The released inference path is a deterministic single scheduler step:
|
|
normal latents start at zero, RGB latents are concatenated as conditioning,
|
|
and the UNet predicts the one-step update.
|
|
"""
|
|
|
|
self._activate("unet", device, config.offload_mode)
|
|
|
|
frame_count = int(rgb_latents_cpu.shape[0])
|
|
windows = build_windows(frame_count, config.window_size, config.step_size)
|
|
|
|
# The output shape matches the VAE latent shape. It is allocated once and
|
|
# overwritten by each anchored window instead of repeatedly concatenated.
|
|
output_cpu = torch.empty_like(rgb_latents_cpu, device="cpu")
|
|
|
|
previous_window: Window | None = None
|
|
previous_latents: torch.Tensor | None = None
|
|
|
|
for index, window in enumerate(windows, start=1):
|
|
# RGB latent window: [frames, channels, h, w] ->
|
|
# [batch=1, frames, channels, h, w].
|
|
image_latents = rgb_latents_cpu[window.start : window.end].to(
|
|
device=device,
|
|
dtype=self.dtype,
|
|
).unsqueeze(0)
|
|
|
|
# One CLIP embedding is supplied for every video frame.
|
|
embeddings = clip_embeddings_cpu[window.start : window.end].to(
|
|
device=device,
|
|
dtype=self.dtype,
|
|
)
|
|
|
|
# NormalCrafter's prediction state starts from zero rather than random
|
|
# diffusion noise. This is why the node intentionally has no seed.
|
|
latents = torch.zeros(
|
|
(
|
|
1,
|
|
window.length,
|
|
int(self.unet.config.out_channels),
|
|
*image_latents.shape[-2:],
|
|
),
|
|
device=device,
|
|
dtype=self.dtype,
|
|
)
|
|
|
|
overlap = overlap_length(previous_window, window)
|
|
old_overlap: torch.Tensor | None = None
|
|
|
|
if overlap and previous_window is not None and previous_latents is not None:
|
|
# Determine where the current window begins inside the previous
|
|
# window, then seed the overlapping prefix with prior predictions.
|
|
previous_offset = window.start - previous_window.start
|
|
old_overlap = previous_latents[
|
|
:,
|
|
previous_offset : previous_offset + overlap,
|
|
]
|
|
latents[:, :overlap].copy_(old_overlap)
|
|
|
|
# NormalCrafter uses exactly one Euler scheduler timestep.
|
|
self.scheduler.set_timesteps(1, device=device)
|
|
timestep = self.scheduler.timesteps[0]
|
|
|
|
# Scheduler scaling is kept even for one step because it is part of
|
|
# the trained SVD inference convention.
|
|
model_input = self.scheduler.scale_model_input(latents, timestep)
|
|
|
|
# Channel concatenation: four predicted-normal channels plus four RGB
|
|
# latent conditioning channels -> eight UNet input channels.
|
|
model_input = torch.cat((model_input, image_latents), dim=2)
|
|
added_time_ids = self._added_time_ids(dtype=self.dtype, device=device)
|
|
|
|
noise_prediction = self.unet(
|
|
model_input,
|
|
timestep,
|
|
encoder_hidden_states=embeddings,
|
|
added_time_ids=added_time_ids,
|
|
return_dict=False,
|
|
)[0]
|
|
|
|
predicted = self.scheduler.step(
|
|
noise_prediction,
|
|
timestep,
|
|
latents,
|
|
).prev_sample
|
|
|
|
if old_overlap is not None:
|
|
# Crossfade old -> new prediction over the overlap. Endpoint values
|
|
# are excluded so neither side receives a fully duplicated weight.
|
|
weights = torch.linspace(
|
|
1.0,
|
|
0.0,
|
|
overlap + 2,
|
|
device=device,
|
|
dtype=self.dtype,
|
|
)[1:-1].view(1, overlap, 1, 1, 1)
|
|
|
|
predicted[:, :overlap] = (
|
|
old_overlap * weights
|
|
+ predicted[:, :overlap] * (1.0 - weights)
|
|
)
|
|
|
|
# Later windows intentionally overwrite their entire covered region.
|
|
# The overlap has already been blended, and the final anchored window
|
|
# guarantees coverage through the last frame.
|
|
output_cpu[window.start : window.end].copy_(
|
|
predicted[0].detach().to("cpu")
|
|
)
|
|
|
|
# Retain only the immediately previous GPU window for overlap seeding.
|
|
# This bounds temporal state by ``window_size``, not total clip length.
|
|
previous_window = window
|
|
previous_latents = predicted.detach()
|
|
|
|
del image_latents, embeddings, latents, model_input, added_time_ids, noise_prediction
|
|
|
|
if progress:
|
|
progress("UNet", index, len(windows))
|
|
|
|
return output_cpu
|
|
|
|
def _decode_normals(
|
|
self,
|
|
normal_latents_cpu: torch.Tensor,
|
|
source: FrameSource,
|
|
device: torch.device,
|
|
config: InferenceConfig,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
"""Decode normal latents into final CPU ComfyUI IMAGE tensors."""
|
|
|
|
self._activate("vae", device, config.offload_mode)
|
|
|
|
original_dtype = next(self.vae.parameters()).dtype
|
|
needs_upcast = (
|
|
original_dtype == torch.float16
|
|
and bool(getattr(self.vae.config, "force_upcast", False))
|
|
)
|
|
if needs_upcast:
|
|
self.vae.to(dtype=torch.float32)
|
|
compute_dtype = next(self.vae.parameters()).dtype
|
|
|
|
if config.output_size == "original":
|
|
out_height, out_width = source.original_height, source.original_width
|
|
else:
|
|
out_height, out_width = source.resized_height, source.resized_width
|
|
|
|
# Allocate only the real frame count. Any synthetic padding frames used to
|
|
# satisfy a short temporal window are never decoded or returned.
|
|
output = torch.empty(
|
|
(source.frame_count, out_height, out_width, 3),
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
)
|
|
|
|
scaling_factor = float(self.vae.config.scaling_factor)
|
|
|
|
# Diffusers versions differ in whether temporal decoders accept an
|
|
# explicit ``num_frames`` argument. Inspect once and adapt safely.
|
|
forward_vae = (
|
|
self.vae._orig_mod.forward
|
|
if hasattr(self.vae, "_orig_mod")
|
|
else self.vae.forward
|
|
)
|
|
accepts_num_frames = "num_frames" in inspect.signature(forward_vae).parameters
|
|
|
|
total = (
|
|
source.frame_count + config.vae_decode_chunk_size - 1
|
|
) // config.vae_decode_chunk_size
|
|
|
|
try:
|
|
for chunk_index, start in enumerate(
|
|
range(0, source.frame_count, config.vae_decode_chunk_size),
|
|
start=1,
|
|
):
|
|
end = min(start + config.vae_decode_chunk_size, source.frame_count)
|
|
|
|
latents = normal_latents_cpu[start:end].to(
|
|
device=device,
|
|
dtype=compute_dtype,
|
|
)
|
|
|
|
# Undo the latent scaling expected by the VAE decoder.
|
|
latents = latents / scaling_factor
|
|
|
|
decode_kwargs = (
|
|
{"num_frames": end - start}
|
|
if accepts_num_frames
|
|
else {}
|
|
)
|
|
# Decode, promote to FP32, and transfer to CPU in one expression.
|
|
# The chained calls are kept identical to the validated version.
|
|
normals = self.vae.decode(latents, **decode_kwargs).sample.float().to("cpu")
|
|
|
|
# Remove the white spatial padding introduced before encoding.
|
|
normals = source.crop_padding(normals)
|
|
|
|
if (out_height, out_width) != (
|
|
source.resized_height,
|
|
source.resized_width,
|
|
):
|
|
normals = F.interpolate(
|
|
normals,
|
|
size=(out_height, out_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
|
|
# Decoder output represents XYZ components in [-1, 1].
|
|
normals = normals.clamp(-1.0, 1.0)
|
|
|
|
# Bilinear resizing blends vector components and can reduce vector
|
|
# magnitude. Renormalization restores unit-length surface normals.
|
|
if config.renormalize_normals:
|
|
normals = F.normalize(normals, p=2.0, dim=1, eps=1e-6)
|
|
|
|
# Encode signed XYZ into display/image range [0, 1].
|
|
normals = (normals * 0.5 + 0.5).clamp_(0.0, 1.0)
|
|
|
|
# Convert BCHW -> BHWC, the layout required by ComfyUI IMAGE.
|
|
output[start:end].copy_(
|
|
normals.permute(0, 2, 3, 1).contiguous()
|
|
)
|
|
|
|
del latents, normals
|
|
|
|
if progress:
|
|
progress("VAE decode", chunk_index, total)
|
|
finally:
|
|
if needs_upcast:
|
|
self.vae.to(dtype=original_dtype)
|
|
|
|
return output
|