442 lines
13 KiB
Python
442 lines
13 KiB
Python
"""NormalCrafter inference pipeline.
|
|
|
|
Full-video intermediates stay on CPU. Only the active CLIP/VAE chunk or UNet
|
|
window is moved to the inference device.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import inspect
|
|
import math
|
|
from typing import Callable
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from .model import NormalCrafterModels
|
|
|
|
FPS_CONDITIONING = 7
|
|
MOTION_BUCKET_ID = 127
|
|
NOISE_AUG_STRENGTH = 0.0
|
|
|
|
ProgressCallback = Callable[[], None]
|
|
|
|
|
|
@dataclass
|
|
class FrameSource:
|
|
"""Lazy CPU-side preprocessing view over a ComfyUI IMAGE batch."""
|
|
|
|
frames: torch.Tensor
|
|
max_resolution: int
|
|
window_size: int
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.frames.ndim != 4 or self.frames.shape[-1] < 3:
|
|
raise ValueError("images must have shape [frames, height, width, channels>=3]")
|
|
if self.frames.shape[0] == 0:
|
|
raise ValueError("images contains no frames")
|
|
|
|
self.frames = (
|
|
self.frames[..., :3]
|
|
.detach()
|
|
.to(device="cpu", dtype=torch.float32)
|
|
.contiguous()
|
|
.clamp_(0.0, 1.0)
|
|
)
|
|
|
|
_, self.original_height, self.original_width, _ = self.frames.shape
|
|
scale = min(
|
|
1.0,
|
|
self.max_resolution / max(self.original_height, self.original_width),
|
|
)
|
|
self.height = max(1, int(round(self.original_height * scale)))
|
|
self.width = max(1, int(round(self.original_width * scale)))
|
|
|
|
padded_height = math.ceil(self.height / 64) * 64
|
|
padded_width = math.ceil(self.width / 64) * 64
|
|
pad_h = padded_height - self.height
|
|
pad_w = padded_width - self.width
|
|
self.pad_top = pad_h // 2
|
|
self.pad_bottom = pad_h - self.pad_top
|
|
self.pad_left = pad_w // 2
|
|
self.pad_right = pad_w - self.pad_left
|
|
|
|
@property
|
|
def frame_count(self) -> int:
|
|
return int(self.frames.shape[0])
|
|
|
|
@property
|
|
def effective_frame_count(self) -> int:
|
|
return max(self.frame_count, self.window_size)
|
|
|
|
def get(self, start: int, end: int) -> torch.Tensor:
|
|
"""Return processed BCHW frames; short clips repeat their final frame."""
|
|
|
|
indices = torch.arange(start, end, dtype=torch.long).clamp_max(
|
|
self.frame_count - 1
|
|
)
|
|
chunk = self.frames.index_select(0, indices).permute(0, 3, 1, 2).contiguous()
|
|
|
|
if (self.height, self.width) != (self.original_height, self.original_width):
|
|
chunk = F.interpolate(
|
|
chunk,
|
|
size=(self.height, self.width),
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
antialias=True,
|
|
).clamp_(0.0, 1.0)
|
|
|
|
if self.pad_top or self.pad_bottom or self.pad_left or self.pad_right:
|
|
chunk = F.pad(
|
|
chunk,
|
|
(self.pad_left, self.pad_right, self.pad_top, self.pad_bottom),
|
|
value=1.0,
|
|
)
|
|
return chunk
|
|
|
|
def crop(self, frames: torch.Tensor) -> torch.Tensor:
|
|
return frames[
|
|
:,
|
|
:,
|
|
self.pad_top : self.pad_top + self.height,
|
|
self.pad_left : self.pad_left + self.width,
|
|
]
|
|
|
|
|
|
def _validate_settings(window_size: int, step_size: int, chunk_size: int) -> None:
|
|
if window_size < 2:
|
|
raise ValueError("window_size must be at least 2")
|
|
if not 1 <= step_size <= window_size:
|
|
raise ValueError("step_size must be between 1 and window_size")
|
|
if chunk_size < 1:
|
|
raise ValueError("chunk_size must be at least 1")
|
|
|
|
|
|
def _build_windows(
|
|
frame_count: int,
|
|
window_size: int,
|
|
step_size: int,
|
|
) -> list[tuple[int, int]]:
|
|
"""Build the same full-size, tail-anchored windows as NormalCrafter."""
|
|
|
|
if frame_count < window_size:
|
|
raise ValueError("frame_count must already be padded to window_size")
|
|
|
|
last_start = frame_count - window_size
|
|
starts = list(range(0, last_start + 1, step_size))
|
|
if starts[-1] != last_start:
|
|
starts.append(last_start)
|
|
return [(start, start + window_size) for start in starts]
|
|
|
|
|
|
def estimate_progress_steps(
|
|
frame_count: int,
|
|
*,
|
|
window_size: int,
|
|
step_size: int,
|
|
chunk_size: int,
|
|
) -> int:
|
|
_validate_settings(window_size, step_size, chunk_size)
|
|
effective = max(frame_count, window_size)
|
|
return max(
|
|
1,
|
|
math.ceil(effective / chunk_size)
|
|
+ math.ceil(effective / chunk_size)
|
|
+ len(_build_windows(effective, window_size, step_size))
|
|
+ math.ceil(frame_count / chunk_size),
|
|
)
|
|
|
|
|
|
def run_inference(
|
|
models: NormalCrafterModels,
|
|
images: torch.Tensor,
|
|
*,
|
|
device: torch.device,
|
|
max_resolution: int,
|
|
window_size: int,
|
|
step_size: int,
|
|
chunk_size: int,
|
|
progress: ProgressCallback | None = None,
|
|
) -> torch.Tensor:
|
|
"""Run CLIP -> VAE encode -> temporal UNet -> VAE decode."""
|
|
|
|
_validate_settings(window_size, step_size, chunk_size)
|
|
source = FrameSource(
|
|
images,
|
|
max_resolution=max_resolution,
|
|
window_size=window_size,
|
|
)
|
|
|
|
try:
|
|
with torch.inference_mode():
|
|
embeddings = _encode_clip(models, source, device, chunk_size, progress)
|
|
rgb_latents = _encode_rgb(models, source, device, chunk_size, progress)
|
|
normal_latents = _predict_normals(
|
|
models,
|
|
rgb_latents,
|
|
embeddings,
|
|
device,
|
|
window_size,
|
|
step_size,
|
|
progress,
|
|
)
|
|
del rgb_latents, embeddings
|
|
return _decode_normals(
|
|
models,
|
|
normal_latents,
|
|
source,
|
|
device,
|
|
chunk_size,
|
|
progress,
|
|
)
|
|
finally:
|
|
models.offload_all()
|
|
|
|
|
|
def _tick(progress: ProgressCallback | None) -> None:
|
|
if progress is not None:
|
|
progress()
|
|
|
|
|
|
def _encode_clip(
|
|
models: NormalCrafterModels,
|
|
source: FrameSource,
|
|
device: torch.device,
|
|
chunk_size: int,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
models.activate("image_encoder", device)
|
|
|
|
mean = torch.tensor(
|
|
models.feature_extractor.image_mean,
|
|
dtype=torch.float32,
|
|
).view(1, 3, 1, 1)
|
|
std = torch.tensor(
|
|
models.feature_extractor.image_std,
|
|
dtype=torch.float32,
|
|
).view(1, 3, 1, 1)
|
|
output = None
|
|
|
|
for start in range(0, source.effective_frame_count, chunk_size):
|
|
end = min(start + chunk_size, source.effective_frame_count)
|
|
frames = source.get(start, end)
|
|
frames = F.interpolate(
|
|
frames,
|
|
size=(224, 224),
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
antialias=True,
|
|
).clamp_(0.0, 1.0)
|
|
|
|
pixels = ((frames - mean) / std).to(device=device, dtype=models.dtype)
|
|
embeddings = models.image_encoder(pixel_values=pixels).image_embeds
|
|
if embeddings.ndim == 2:
|
|
embeddings = embeddings.unsqueeze(1)
|
|
embeddings = embeddings.detach().to(device="cpu", dtype=models.dtype)
|
|
|
|
if output is None:
|
|
output = torch.empty(
|
|
(source.effective_frame_count, *embeddings.shape[1:]),
|
|
dtype=models.dtype,
|
|
device="cpu",
|
|
)
|
|
output[start:end].copy_(embeddings)
|
|
_tick(progress)
|
|
|
|
assert output is not None
|
|
return output
|
|
|
|
|
|
def _encode_rgb(
|
|
models: NormalCrafterModels,
|
|
source: FrameSource,
|
|
device: torch.device,
|
|
chunk_size: int,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
models.activate("vae", device)
|
|
vae_dtype = next(models.vae.parameters()).dtype
|
|
output = None
|
|
|
|
for start in range(0, source.effective_frame_count, chunk_size):
|
|
end = min(start + chunk_size, source.effective_frame_count)
|
|
frames = (source.get(start, end) * 2.0 - 1.0).to(
|
|
device=device,
|
|
dtype=vae_dtype,
|
|
)
|
|
latents = models.vae.encode(frames).latent_dist.mode()
|
|
latents = latents.detach().to(device="cpu", dtype=models.dtype)
|
|
|
|
if output is None:
|
|
output = torch.empty(
|
|
(source.effective_frame_count, *latents.shape[1:]),
|
|
dtype=models.dtype,
|
|
device="cpu",
|
|
)
|
|
output[start:end].copy_(latents)
|
|
_tick(progress)
|
|
|
|
assert output is not None
|
|
return output
|
|
|
|
|
|
def _added_time_ids(
|
|
models: NormalCrafterModels,
|
|
device: torch.device,
|
|
) -> torch.Tensor:
|
|
values = [FPS_CONDITIONING, MOTION_BUCKET_ID, NOISE_AUG_STRENGTH]
|
|
expected = models.unet.add_embedding.linear_1.in_features
|
|
actual = int(models.unet.config.addition_time_embed_dim) * len(values)
|
|
if expected != actual:
|
|
raise ValueError(
|
|
f"UNet added-time embedding mismatch: expected {expected}, got {actual}"
|
|
)
|
|
return torch.tensor([values], dtype=models.dtype, device=device)
|
|
|
|
|
|
def _predict_normals(
|
|
models: NormalCrafterModels,
|
|
rgb_latents: torch.Tensor,
|
|
embeddings: torch.Tensor,
|
|
device: torch.device,
|
|
window_size: int,
|
|
step_size: int,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
models.activate("unet", device)
|
|
windows = _build_windows(
|
|
int(rgb_latents.shape[0]),
|
|
window_size,
|
|
step_size,
|
|
)
|
|
output = torch.empty_like(rgb_latents, device="cpu")
|
|
|
|
previous_window: tuple[int, int] | None = None
|
|
previous_prediction: torch.Tensor | None = None
|
|
|
|
for start, end in windows:
|
|
image_latents = rgb_latents[start:end].to(
|
|
device=device,
|
|
dtype=models.dtype,
|
|
).unsqueeze(0)
|
|
frame_embeddings = embeddings[start:end].to(
|
|
device=device,
|
|
dtype=models.dtype,
|
|
)
|
|
latents = torch.zeros(
|
|
(
|
|
1,
|
|
end - start,
|
|
int(models.unet.config.out_channels),
|
|
*image_latents.shape[-2:],
|
|
),
|
|
device=device,
|
|
dtype=models.dtype,
|
|
)
|
|
|
|
overlap = 0 if previous_window is None else max(0, previous_window[1] - start)
|
|
old_overlap = None
|
|
if overlap and previous_window is not None and previous_prediction is not None:
|
|
previous_offset = start - previous_window[0]
|
|
old_overlap = previous_prediction[
|
|
:,
|
|
previous_offset : previous_offset + overlap,
|
|
]
|
|
latents[:, :overlap].copy_(old_overlap)
|
|
|
|
models.scheduler.set_timesteps(1, device=device)
|
|
timestep = models.scheduler.timesteps[0]
|
|
model_input = models.scheduler.scale_model_input(latents, timestep)
|
|
model_input = torch.cat((model_input, image_latents), dim=2)
|
|
|
|
noise_prediction = models.unet(
|
|
model_input,
|
|
timestep,
|
|
encoder_hidden_states=frame_embeddings,
|
|
added_time_ids=_added_time_ids(models, device),
|
|
return_dict=False,
|
|
)[0]
|
|
predicted = models.scheduler.step(
|
|
noise_prediction,
|
|
timestep,
|
|
latents,
|
|
).prev_sample
|
|
|
|
if old_overlap is not None:
|
|
weights = torch.linspace(
|
|
1.0,
|
|
0.0,
|
|
overlap + 2,
|
|
device=device,
|
|
dtype=models.dtype,
|
|
)[1:-1].view(1, overlap, 1, 1, 1)
|
|
predicted[:, :overlap] = (
|
|
old_overlap * weights
|
|
+ predicted[:, :overlap] * (1.0 - weights)
|
|
)
|
|
|
|
output[start:end].copy_(predicted[0].detach().to("cpu"))
|
|
previous_window = (start, end)
|
|
previous_prediction = predicted.detach()
|
|
_tick(progress)
|
|
|
|
return output
|
|
|
|
|
|
def _decode_normals(
|
|
models: NormalCrafterModels,
|
|
normal_latents: torch.Tensor,
|
|
source: FrameSource,
|
|
device: torch.device,
|
|
chunk_size: int,
|
|
progress: ProgressCallback | None,
|
|
) -> torch.Tensor:
|
|
models.activate("vae", device)
|
|
vae_dtype = next(models.vae.parameters()).dtype
|
|
|
|
output = torch.empty(
|
|
(source.frame_count, source.original_height, source.original_width, 3),
|
|
dtype=torch.float32,
|
|
device="cpu",
|
|
)
|
|
scaling_factor = float(models.vae.config.scaling_factor)
|
|
|
|
forward = (
|
|
models.vae._orig_mod.forward
|
|
if hasattr(models.vae, "_orig_mod")
|
|
else models.vae.forward
|
|
)
|
|
accepts_num_frames = "num_frames" in inspect.signature(forward).parameters
|
|
|
|
for start in range(0, source.frame_count, chunk_size):
|
|
end = min(start + chunk_size, source.frame_count)
|
|
latents = normal_latents[start:end].to(
|
|
device=device,
|
|
dtype=vae_dtype,
|
|
) / scaling_factor
|
|
kwargs = {"num_frames": end - start} if accepts_num_frames else {}
|
|
normals = models.vae.decode(latents, **kwargs).sample.float().to("cpu")
|
|
normals = source.crop(normals)
|
|
|
|
if (source.height, source.width) != (
|
|
source.original_height,
|
|
source.original_width,
|
|
):
|
|
normals = F.interpolate(
|
|
normals,
|
|
size=(source.original_height, source.original_width),
|
|
mode="bilinear",
|
|
align_corners=False,
|
|
)
|
|
|
|
normals = normals.clamp(-1.0, 1.0)
|
|
normals = F.normalize(normals, p=2.0, dim=1, eps=1e-6)
|
|
normals = (normals * 0.5 + 0.5).clamp_(0.0, 1.0)
|
|
output[start:end].copy_(
|
|
normals.permute(0, 2, 3, 1).contiguous()
|
|
)
|
|
_tick(progress)
|
|
|
|
return output
|