wehub-resource-sync
450 lines
17 KiB
Python
450 lines
17 KiB
Python
# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
"""Shared utilities for chunk-wise causal operations.
|
|
|
|
This module provides utilities for managing temporal chunking in causal video generation,
|
|
supporting multiple split strategies (uniform, first_frame, first_plus_one).
|
|
"""
|
|
|
|
from typing import Any, Dict, List, Optional, Tuple
|
|
|
|
|
|
def is_chunk_causal_request(
|
|
chunk_size: Optional[int],
|
|
T_effective: int,
|
|
chunk_index: Optional[List[int]] = None,
|
|
) -> bool:
|
|
"""Decide whether a layer should run in chunk-causal (vs. fully bidirectional) mode.
|
|
|
|
Chunk-causal mode applies when EITHER:
|
|
1. ``chunk_size`` is set and strictly less than ``T_effective`` (the
|
|
standard rule used by training and most inference paths), OR
|
|
2. ``chunk_index`` is explicitly provided by the caller.
|
|
|
|
Case (2) is required for the staircase cold-start at AR step 0
|
|
phases 0 / 1, where ``T_effective`` (= ``K + G_eff``, with G_eff in
|
|
{1, 2}) can be smaller than the model's pretrained ``chunk_size``
|
|
(typically 3) but the caller still wants strict frame-causal cond
|
|
boundaries via ``chunk_index = [0, 1]``. Without this branch, the
|
|
bidirectional fallback would silently leak gen-frame information
|
|
into cond positions.
|
|
|
|
The bidirectional fallback should be taken ONLY when both
|
|
``chunk_size`` is missing/non-restrictive AND ``chunk_index`` is
|
|
not provided — i.e. the caller has not asked for any chunk
|
|
structure at all.
|
|
|
|
Args:
|
|
chunk_size: Base chunk size from model config (typically 3 for
|
|
Sana-WM); ``None`` if unset.
|
|
T_effective: Total number of frames after CP all-gather (where
|
|
applicable). Use the local ``T`` for non-CP paths.
|
|
chunk_index: Optional explicit chunk-start indices. Anything
|
|
non-``None`` is treated as the caller asking for chunk-
|
|
causal semantics, regardless of ``chunk_size``.
|
|
|
|
Returns:
|
|
``True`` if chunk-causal logic should run, ``False`` if the
|
|
layer should fall back to fully bidirectional behavior.
|
|
"""
|
|
if chunk_size is not None and chunk_size < T_effective:
|
|
return True
|
|
if chunk_index is not None:
|
|
return True
|
|
return False
|
|
|
|
|
|
def chunk_index_from_chunk_size(
|
|
T: int,
|
|
chunk_size: int,
|
|
strategy: str = "uniform",
|
|
) -> List[int]:
|
|
"""Convert chunk_size to chunk_index list with a split strategy.
|
|
|
|
Args:
|
|
T: Number of latent frames.
|
|
chunk_size: Base chunk size for the temporal dimension.
|
|
strategy: Chunk split strategy. Supported values:
|
|
- "uniform" (default): uniform chunks with optional remainder
|
|
Example: T=21, chunk_size=4 → [0,4,8,12,16,20] → sizes [4,4,4,4,4,1]
|
|
- "first_frame": first chunk is 1 frame, then uniform chunk_size
|
|
Example: T=21, chunk_size=4 → [0,1,5,9,13,17] → sizes [1,4,4,4,4,4]
|
|
- "first_plus_one": first chunk is chunk_size + 1, then uniform chunk_size
|
|
Example: T=21, chunk_size=4 → [0,5,9,13,17] → sizes [5,4,4,4,4]
|
|
|
|
Returns:
|
|
List of chunk start indices (not including the final T).
|
|
|
|
Raises:
|
|
ValueError: If chunk_size or T are invalid, or strategy is unknown.
|
|
"""
|
|
if chunk_size <= 0:
|
|
raise ValueError(f"chunk_size must be > 0, got {chunk_size}.")
|
|
if T <= 0:
|
|
raise ValueError(f"T must be > 0, got {T}.")
|
|
|
|
if strategy is None:
|
|
strategy = "uniform"
|
|
strategy = str(strategy).lower()
|
|
|
|
if strategy in ("uniform", "default"):
|
|
indices = list(range(0, T, chunk_size))
|
|
# Absorb small remainder into last chunk to avoid degenerate chunks
|
|
# (e.g., causal_conv1d crashes on length=1 sequences).
|
|
if len(indices) > 1 and (T - indices[-1]) < chunk_size:
|
|
indices.pop()
|
|
return indices
|
|
|
|
if strategy in ("first_frame", "first_frame_alone", "first_frame_only"):
|
|
if T <= 1:
|
|
return [0]
|
|
indices = [0] + list(range(1, T, chunk_size))
|
|
if len(indices) > 2 and (T - indices[-1]) < chunk_size:
|
|
indices.pop()
|
|
return indices
|
|
|
|
if strategy in ("first_plus_one", "first_chunk_plus_one"):
|
|
if T <= chunk_size + 1:
|
|
return [0]
|
|
indices = [0] + list(range(chunk_size + 1, T, chunk_size))
|
|
# Absorb small remainder into last chunk to avoid degenerate chunks
|
|
# (e.g., T_latent=41 with chunk_size=3 → last chunk would be 1 frame,
|
|
# which crashes causal_conv1d). Merge it into the previous chunk instead.
|
|
if len(indices) > 1 and (T - indices[-1]) < chunk_size:
|
|
indices.pop()
|
|
return indices
|
|
|
|
raise ValueError(f"Unknown chunk_split_strategy '{strategy}'. Supported: uniform, first_frame, first_plus_one.")
|
|
|
|
|
|
def get_chunk_index_from_config(config: Any, num_frames: Optional[int] = None) -> Optional[List[int]]:
|
|
"""Resolve chunk_index from a config, supporting chunk_size and strategy.
|
|
|
|
Priority:
|
|
1) config.model.chunk_index (explicit list)
|
|
2) config.model.chunk_size (compute with chunk_split_strategy)
|
|
3) None (no chunking)
|
|
|
|
Args:
|
|
config: Config object or dict with a "model" field.
|
|
num_frames: Number of latent frames. Required when using chunk_size.
|
|
|
|
Returns:
|
|
Chunk start indices, or None if chunking is disabled.
|
|
|
|
Raises:
|
|
ValueError: If chunk_size is set but num_frames is None.
|
|
"""
|
|
model = getattr(config, "model", None)
|
|
if model is None:
|
|
return None
|
|
|
|
def _get_model_attr(name: str, default: Any) -> Any:
|
|
if hasattr(model, "get"):
|
|
return model.get(name, default)
|
|
if isinstance(model, dict):
|
|
return model.get(name, default)
|
|
return getattr(model, name, default)
|
|
|
|
chunk_index = _get_model_attr("chunk_index", None)
|
|
chunk_size = _get_model_attr("chunk_size", None)
|
|
chunk_split_strategy = _get_model_attr("chunk_split_strategy", "uniform")
|
|
|
|
if chunk_index is not None:
|
|
if not isinstance(chunk_index, (list, tuple)):
|
|
raise TypeError(f"chunk_index must be a list, got {type(chunk_index).__name__}")
|
|
if len(chunk_index) == 0:
|
|
raise ValueError("chunk_index cannot be empty. Provide at least one chunk boundary.")
|
|
return list(chunk_index)
|
|
if chunk_size is not None:
|
|
if num_frames is None:
|
|
raise ValueError(f"num_frames must be provided when using chunk_size={chunk_size}")
|
|
return chunk_index_from_chunk_size(num_frames, chunk_size, strategy=chunk_split_strategy)
|
|
return None
|
|
|
|
|
|
def compute_chunk_sizes(chunk_index: List[int], T: int) -> List[int]:
|
|
"""Compute actual chunk sizes from chunk_index.
|
|
|
|
Args:
|
|
chunk_index: List of chunk start indices (e.g., [0, 4, 8, 12]).
|
|
T: Total number of frames.
|
|
|
|
Returns:
|
|
List of chunk sizes (e.g., [4, 4, 4, 1] if T=13).
|
|
|
|
Example:
|
|
>>> compute_chunk_sizes([0, 4, 8, 12], T=13)
|
|
[4, 4, 4, 1]
|
|
>>> compute_chunk_sizes([0, 1, 5, 9], T=13)
|
|
[1, 4, 4, 4]
|
|
"""
|
|
if not chunk_index:
|
|
return []
|
|
|
|
# Ensure chunk_index is clean
|
|
chunk_index = [idx for idx in chunk_index if 0 <= idx < T]
|
|
if not chunk_index:
|
|
return []
|
|
|
|
# Add T as the final boundary if not present
|
|
if chunk_index[-1] != T:
|
|
chunk_index = chunk_index + [T]
|
|
|
|
# Compute sizes
|
|
sizes = [chunk_index[i + 1] - chunk_index[i] for i in range(len(chunk_index) - 1)]
|
|
return sizes
|
|
|
|
|
|
def size1_chunk_position_indices(chunk_index: List[int]) -> List[int]:
|
|
"""Return frame-time positions belonging to size-1 (singleton) chunks.
|
|
|
|
A size-1 chunk has no intra-chunk lookahead, so the anti-causal
|
|
branch (backward GDN scan and the per-chunk backward conv path)
|
|
contributes nothing for these positions in a chunk-causal layer.
|
|
This helper exposes those positions so downstream code can skip
|
|
the reverse-direction compute (and zero-out the contribution).
|
|
|
|
Args:
|
|
chunk_index: Normalized chunk indices, including the trailing
|
|
``T`` boundary, e.g. ``[0, 1, 2, ..., K, K+G]`` for the
|
|
``cond_chunk_mode='frame_causal'`` layout.
|
|
|
|
Returns:
|
|
List of frame-time positions ``p`` for which ``[p, p+1)`` is a
|
|
chunk of size 1. Returns ``[]`` when no size-1 chunks exist
|
|
(e.g. uniform ``chunk_size=3`` patterns).
|
|
|
|
Examples:
|
|
>>> size1_chunk_position_indices([0, 3, 6, 9]) # uniform size 3
|
|
[]
|
|
>>> size1_chunk_position_indices([0, 1, 2, 3, 4, 7]) # frame_causal, K=4, G=3
|
|
[0, 1, 2, 3]
|
|
"""
|
|
return [s for s, e in zip(chunk_index[:-1], chunk_index[1:]) if e - s == 1]
|
|
|
|
|
|
def is_uniform_chunking(
|
|
chunk_index: List[int],
|
|
T: int,
|
|
chunk_size: int,
|
|
) -> bool:
|
|
"""Check if chunk_index represents uniform chunking.
|
|
|
|
Returns True if all chunks are equal to chunk_size except possibly the last
|
|
chunk which may be smaller (the remainder). This is the pattern that allows
|
|
safe vectorized padding with: pad_t = chunk_size - (T % chunk_size).
|
|
|
|
Uniform patterns (return True):
|
|
- [0,4,8,12,16,20] with T=21, chunk_size=4 → sizes [4,4,4,4,4,1] ✓
|
|
- [0,4,8,12,16] with T=20, chunk_size=4 → sizes [4,4,4,4,4] ✓
|
|
- [0,4,8] with T=10, chunk_size=4 → sizes [4,4,2] ✓
|
|
|
|
Non-uniform patterns (return False):
|
|
- [0,1,5,9,13,17] with T=21, chunk_size=4 → sizes [1,4,4,4,4,4] ✗
|
|
- [0,5,9,13,17] with T=21, chunk_size=4 → sizes [5,4,4,4,4] ✗
|
|
|
|
Args:
|
|
chunk_index: List of chunk start indices.
|
|
T: Total number of frames.
|
|
chunk_size: Expected uniform chunk size.
|
|
|
|
Returns:
|
|
True if chunking is uniform, False otherwise.
|
|
"""
|
|
if chunk_size <= 0:
|
|
return False
|
|
|
|
# Compute actual chunk sizes
|
|
sizes = compute_chunk_sizes(chunk_index, T)
|
|
|
|
if not sizes:
|
|
return True # Empty is trivially uniform
|
|
|
|
# Check that all chunks except possibly the last are equal to chunk_size
|
|
for i, size in enumerate(sizes):
|
|
is_last = i == len(sizes) - 1
|
|
if is_last:
|
|
# Last chunk can be <= chunk_size (remainder)
|
|
if size > chunk_size:
|
|
return False
|
|
else:
|
|
# All other chunks must be exactly chunk_size
|
|
if size != chunk_size:
|
|
return False
|
|
|
|
return True
|
|
|
|
|
|
def analyze_chunk_pattern(
|
|
chunk_index: List[int],
|
|
T: int,
|
|
chunk_size: int,
|
|
) -> Tuple[str, Dict[str, Any]]:
|
|
"""Analyze chunk pattern and return vectorization strategy.
|
|
|
|
Detects special patterns that allow hybrid vectorization:
|
|
- uniform: All chunks equal except possibly last (vectorized baseline)
|
|
- first_frame: [1, 4, 4, 4, ...] - first frame alone, then uniform tail
|
|
- first_plus_one: [5, 4, 4, 4, ...] - first chunk+1, then uniform tail
|
|
- arbitrary: Other patterns (no optimization available)
|
|
|
|
Args:
|
|
chunk_index: List of chunk start indices (e.g., [0, 4, 8, 12]).
|
|
T: Total number of frames.
|
|
chunk_size: Base chunk size for pattern detection.
|
|
|
|
Returns:
|
|
(pattern_type, metadata) where:
|
|
pattern_type: "uniform", "first_frame", "first_plus_one", or "arbitrary"
|
|
metadata: Dict with vectorization hints:
|
|
- vectorizable: bool (True if optimization available)
|
|
- first_chunk_size: int (size of first special chunk)
|
|
- tail_start_index: int (where uniform tail begins in chunk_index)
|
|
- tail_chunk_size: int (uniform size of tail chunks)
|
|
- tail_is_uniform: bool (whether tail is vectorizable)
|
|
|
|
Example:
|
|
>>> analyze_chunk_pattern([0, 1, 5, 9, 13, 17], T=21, chunk_size=4)
|
|
("first_frame", {
|
|
"vectorizable": True,
|
|
"first_chunk_size": 1,
|
|
"tail_start_index": 1,
|
|
"tail_chunk_size": 4,
|
|
"tail_is_uniform": True,
|
|
})
|
|
"""
|
|
sizes = compute_chunk_sizes(chunk_index, T)
|
|
|
|
if not sizes:
|
|
return "uniform", {"vectorizable": True}
|
|
|
|
# Check uniform: all chunks equal to chunk_size except possibly last
|
|
if is_uniform_chunking(chunk_index, T, chunk_size):
|
|
return "uniform", {"vectorizable": True}
|
|
|
|
# Check first_frame pattern: [1, 4, 4, 4, ...]
|
|
if sizes[0] == 1:
|
|
# Check if tail (sizes[1:]) is uniform
|
|
tail_is_uniform = all(s == chunk_size for s in sizes[1:-1])
|
|
# Allow last chunk to be <= chunk_size (remainder)
|
|
if len(sizes) > 1:
|
|
tail_is_uniform = tail_is_uniform and (sizes[-1] <= chunk_size)
|
|
|
|
if tail_is_uniform:
|
|
return "first_frame", {
|
|
"vectorizable": True,
|
|
"first_chunk_size": 1,
|
|
"tail_start_index": 1, # Skip first frame
|
|
"tail_chunk_size": chunk_size,
|
|
"tail_is_uniform": True,
|
|
}
|
|
|
|
# Check first_plus_one pattern: [chunk_size+1, chunk_size, chunk_size, ...]
|
|
if sizes[0] == chunk_size + 1:
|
|
# Check if tail (sizes[1:]) is uniform
|
|
tail_is_uniform = all(s == chunk_size for s in sizes[1:-1])
|
|
# Allow last chunk to be <= chunk_size (remainder)
|
|
if len(sizes) > 1:
|
|
tail_is_uniform = tail_is_uniform and (sizes[-1] <= chunk_size)
|
|
|
|
if tail_is_uniform:
|
|
return "first_plus_one", {
|
|
"vectorizable": True,
|
|
"first_chunk_size": chunk_size + 1,
|
|
"tail_start_index": chunk_size + 1, # Skip first chunk
|
|
"tail_chunk_size": chunk_size,
|
|
"tail_is_uniform": True,
|
|
}
|
|
|
|
# Arbitrary pattern - no vectorization available
|
|
return "arbitrary", {"vectorizable": False}
|
|
|
|
|
|
def normalize_chunk_index(
|
|
chunk_index: Optional[List[int]],
|
|
T: int,
|
|
chunk_size: Optional[int] = None,
|
|
chunk_split_strategy: str = "uniform",
|
|
) -> Tuple[List[int], bool]:
|
|
"""Normalize chunk_index and detect if uniform.
|
|
|
|
This function handles all the complex logic for:
|
|
1. Converting chunk_size + strategy → chunk_index (if needed)
|
|
2. Cleaning and validating chunk_index
|
|
3. Detecting if the result is uniform (safe for vectorized padding)
|
|
|
|
Args:
|
|
chunk_index: Optional pre-computed chunk indices.
|
|
T: Total number of frames.
|
|
chunk_size: Chunk size (required if chunk_index is None or for uniformity check).
|
|
chunk_split_strategy: Strategy to use if generating chunk_index from chunk_size.
|
|
|
|
Returns:
|
|
(normalized_chunk_index, is_uniform):
|
|
- normalized_chunk_index: Clean list of chunk start indices
|
|
- is_uniform: True if safe to use vectorized path with padding
|
|
|
|
Raises:
|
|
ValueError: If required parameters are missing or invalid.
|
|
"""
|
|
# Case 1: chunk_index provided explicitly
|
|
if chunk_index is not None:
|
|
normalized_chunk_index = list(chunk_index)
|
|
|
|
# Clean up: ensure starts with 0 and ends with T
|
|
if not normalized_chunk_index or normalized_chunk_index[0] != 0:
|
|
normalized_chunk_index = [0] + [idx for idx in normalized_chunk_index if idx > 0]
|
|
normalized_chunk_index = [idx for idx in normalized_chunk_index if idx < T]
|
|
if not normalized_chunk_index:
|
|
normalized_chunk_index = [0]
|
|
if normalized_chunk_index[-1] != T:
|
|
normalized_chunk_index = normalized_chunk_index + [T]
|
|
|
|
# Check if uniform (requires chunk_size for comparison)
|
|
if chunk_size is None:
|
|
# Can't verify uniformity without chunk_size, assume non-uniform (safe)
|
|
is_uniform = False
|
|
else:
|
|
is_uniform = is_uniform_chunking(normalized_chunk_index, T, chunk_size)
|
|
|
|
return normalized_chunk_index, is_uniform
|
|
|
|
# Case 2: Generate chunk_index from chunk_size + strategy
|
|
if chunk_size is None:
|
|
raise ValueError("Either chunk_index or chunk_size must be provided.")
|
|
|
|
if chunk_size <= 0:
|
|
raise ValueError(f"chunk_size must be > 0, got {chunk_size}.")
|
|
|
|
# Normalize strategy
|
|
strategy = "uniform" if chunk_split_strategy is None else str(chunk_split_strategy).lower()
|
|
|
|
# Generate chunk_index
|
|
chunk_index_gen = chunk_index_from_chunk_size(T, chunk_size, strategy=strategy)
|
|
|
|
# Add T as final boundary
|
|
if not chunk_index_gen:
|
|
chunk_index_gen = [0]
|
|
if chunk_index_gen[-1] != T:
|
|
chunk_index_gen = chunk_index_gen + [T]
|
|
|
|
# Check if uniform
|
|
is_uniform = is_uniform_chunking(chunk_index_gen, T, chunk_size)
|
|
|
|
return chunk_index_gen, is_uniform
|