195 lines
8.0 KiB
Python
195 lines
8.0 KiB
Python
"""CPU-side frame preparation for NormalCrafter inference.
|
|
|
|
ComfyUI supplies video-like image batches as a single ``BHWC`` tensor:
|
|
|
|
[frame, height, width, channel]
|
|
|
|
The model stack, however, expects ``BCHW`` chunks whose spatial dimensions are
|
|
compatible with the VAE. This module bridges those conventions without creating
|
|
a second processed copy of the complete clip.
|
|
|
|
The crucial design choice is laziness: ``FrameSource`` stores the original clip
|
|
once on CPU and performs resize, padding, layout conversion, and short-clip frame
|
|
repetition only for the range requested by the current inference stage.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
|
|
def _round_size(value: float) -> int:
|
|
"""Round a scaled spatial dimension while guaranteeing at least one pixel."""
|
|
|
|
return max(1, int(round(value)))
|
|
|
|
|
|
def _ceil_multiple(value: int, multiple: int) -> int:
|
|
"""Return the smallest multiple of ``multiple`` not smaller than ``value``.
|
|
|
|
NormalCrafter's VAE/UNet stack repeatedly downsamples spatial dimensions.
|
|
Padding to 64 is conservative and prevents incompatible shapes deep inside
|
|
the network while keeping the original aspect ratio intact.
|
|
"""
|
|
|
|
return int(math.ceil(value / multiple) * multiple)
|
|
|
|
|
|
@dataclass
|
|
class FrameSource:
|
|
"""Lazy CPU-side view of a ComfyUI IMAGE batch.
|
|
|
|
Parameters
|
|
----------
|
|
frames_bhwc:
|
|
ComfyUI ``IMAGE`` tensor in ``[frames, height, width, channels]`` layout.
|
|
Values are expected to represent display-linear image data in ``[0, 1]``.
|
|
max_resolution:
|
|
Upper bound for the longest spatial side used by the model. Inputs are
|
|
never enlarged; smaller clips retain their native dimensions.
|
|
pad_multiple:
|
|
Spatial alignment required by the latent model. Padding is symmetric,
|
|
except for the unavoidable one-pixel asymmetry when the total is odd.
|
|
|
|
Memory behavior
|
|
---------------
|
|
Only the source clip is retained. A resized/padded tensor is created solely
|
|
for the requested chunk and becomes collectible after that stage iteration.
|
|
This is central to keeping peak RAM and VRAM independent of total clip length.
|
|
"""
|
|
|
|
frames_bhwc: torch.Tensor
|
|
max_resolution: int
|
|
pad_multiple: int = 64
|
|
|
|
def __post_init__(self) -> None:
|
|
"""Validate input and precompute immutable geometry metadata."""
|
|
|
|
# ComfyUI IMAGE is BHWC. Alpha or auxiliary channels are allowed, but the
|
|
# released model consumes exactly RGB, so at least three channels exist.
|
|
if self.frames_bhwc.ndim != 4 or self.frames_bhwc.shape[-1] < 3:
|
|
raise ValueError("images must have shape [frames, height, width, channels>=3]")
|
|
if self.frames_bhwc.shape[0] == 0:
|
|
raise ValueError("images contains no frames")
|
|
if self.max_resolution < 64:
|
|
raise ValueError("max_resolution must be at least 64")
|
|
|
|
# Break any autograd relationship, discard non-RGB channels, and establish
|
|
# one canonical long-lived representation: contiguous CPU FP32 in [0, 1].
|
|
# Returning to CPU here also prevents a CUDA input tensor from becoming a
|
|
# hidden full-video allocation retained throughout inference.
|
|
self.frames_bhwc = (
|
|
self.frames_bhwc[..., :3]
|
|
.detach()
|
|
.to(device="cpu", dtype=torch.float32)
|
|
.contiguous()
|
|
.clamp_(0.0, 1.0)
|
|
)
|
|
|
|
# Compute a single aspect-preserving scale. ``min(1.0, ...)`` explicitly
|
|
# forbids upscaling: max_resolution is a ceiling, not a target size.
|
|
_, self.original_height, self.original_width, _ = self.frames_bhwc.shape
|
|
scale = min(1.0, self.max_resolution / max(self.original_height, self.original_width))
|
|
self.resized_height = _round_size(self.original_height * scale)
|
|
self.resized_width = _round_size(self.original_width * scale)
|
|
|
|
# The network receives padded dimensions, while output restoration later
|
|
# removes precisely these margins before optional resize to source size.
|
|
self.padded_height = _ceil_multiple(self.resized_height, self.pad_multiple)
|
|
self.padded_width = _ceil_multiple(self.resized_width, self.pad_multiple)
|
|
|
|
# Split padding around the image. When an odd number of pixels is needed,
|
|
# bottom/right receive the extra pixel; crop_padding uses the same offsets.
|
|
total_pad_h = self.padded_height - self.resized_height
|
|
total_pad_w = self.padded_width - self.resized_width
|
|
self.pad_top = total_pad_h // 2
|
|
self.pad_bottom = total_pad_h - self.pad_top
|
|
self.pad_left = total_pad_w // 2
|
|
self.pad_right = total_pad_w - self.pad_left
|
|
|
|
@property
|
|
def frame_count(self) -> int:
|
|
"""Number of real input frames, excluding any temporal repetition."""
|
|
|
|
return int(self.frames_bhwc.shape[0])
|
|
|
|
def effective_frame_count(self, window_size: int) -> int:
|
|
"""Frame count seen by encoders after minimum-window temporal padding.
|
|
|
|
NormalCrafter's temporal UNet requires at least one complete window. For a
|
|
shorter clip, ``get_bchw`` repeats the final real frame up to this count.
|
|
The decoder still emits only ``frame_count`` real output frames.
|
|
"""
|
|
|
|
return max(self.frame_count, window_size)
|
|
|
|
def get_bchw(self, start: int, end: int) -> torch.Tensor:
|
|
"""Materialize one processed CPU frame range in BCHW layout.
|
|
|
|
The returned tensor is:
|
|
|
|
* CPU resident;
|
|
* FP32;
|
|
* RGB only;
|
|
* in ``[0, 1]``;
|
|
* resized without changing aspect ratio;
|
|
* symmetrically white-padded to ``pad_multiple``;
|
|
* temporally extended by repeating the final frame when ``end`` exceeds
|
|
the real clip length.
|
|
|
|
White padding matches the neutral image-space convention used by the
|
|
original pipeline more closely than zero/black padding.
|
|
"""
|
|
|
|
if start < 0 or end <= start:
|
|
raise ValueError(f"invalid frame range [{start}, {end})")
|
|
|
|
# Clamp synthetic indices to the final real frame. ``index_select`` then
|
|
# handles ordinary ranges and repeated-tail ranges through one code path.
|
|
indices = torch.arange(start, end, dtype=torch.long).clamp_max(self.frame_count - 1)
|
|
|
|
# Select BHWC frames, then convert once to the BCHW layout expected by
|
|
# interpolation, CLIP preprocessing, and the VAE encoder.
|
|
chunk = self.frames_bhwc.index_select(0, indices).permute(0, 3, 1, 2).contiguous()
|
|
|
|
# Bicubic plus antialiasing is appropriate for RGB downscaling. Clamp after
|
|
# interpolation because bicubic kernels may overshoot slightly outside
|
|
# the nominal [0, 1] input range near high-contrast edges.
|
|
if (self.resized_height, self.resized_width) != (self.original_height, self.original_width):
|
|
chunk = F.interpolate(
|
|
chunk,
|
|
size=(self.resized_height, self.resized_width),
|
|
mode="bicubic",
|
|
align_corners=False,
|
|
antialias=True,
|
|
).clamp_(0.0, 1.0)
|
|
|
|
# torch.nn.functional.pad expects (left, right, top, bottom) for BCHW.
|
|
# Padding is skipped entirely when the resized dimensions are already safe.
|
|
if self.pad_top or self.pad_bottom or self.pad_left or self.pad_right:
|
|
chunk = F.pad(
|
|
chunk,
|
|
(self.pad_left, self.pad_right, self.pad_top, self.pad_bottom),
|
|
mode="constant",
|
|
value=1.0,
|
|
)
|
|
return chunk
|
|
|
|
def crop_padding(self, tensor_bchw: torch.Tensor) -> torch.Tensor:
|
|
"""Remove the exact spatial padding previously introduced by ``get_bchw``.
|
|
|
|
Channel and batch/frame dimensions are preserved. The method intentionally
|
|
returns a view where possible; callers may later interpolate or copy it.
|
|
"""
|
|
|
|
return tensor_bchw[
|
|
:,
|
|
:,
|
|
self.pad_top : self.pad_top + self.resized_height,
|
|
self.pad_left : self.pad_left + self.resized_width,
|
|
]
|