Hugging Face's logo

Spaces:
AlayaLab
/
FloodDiffusion2-Live
Running on Zero

App Files Community
2
FloodDiffusion2-Live / space /vendor /models /diffusion_forcing_position_wan.py
caiyiyi1998's picture
caiyiyi1998
Initial commit
9a25493
Raw History Blame Contribute Delete
120 kB
"""Position-conditioned diffusion-forcing models with an isolated engine.
The original models.diffusion_forcing_wan module is intentionally not modified.
This module carries its own position-capable engine so direct per-frame root
motion and WAN-packed root conditioning can share identical offline/stream
logic without changing the base DiffForcingWanModel behavior. Root motion means
only XZ difference plus angular difference; pelvis/root pose remains in the
diffused pose state.
"""
import math
from fractions import Fraction
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
from .tools.t5 import T5EncoderModel
from .tools.wan_model import WanModel
from .normalization import load_statistics
from .quadratic_loss import load_quadratic_loss_matrix, quadratic_error_sum
EPSILON = 0.05
# Most-negative rope frame index supported by rope_params_with_negative(1024, ...).
# Clamping at this value keeps trim_stream's shifted ids inside the cached range.
ROPE_MIN_INDEX = -1024
class TriangularTimeScheduler:
def __init__(self, config):
self.steps = config["steps"]
self.chunk_size = config["chunk_size"]
self.random_epsilon = config.get("random_epsilon", 0.00) # schedule jittering
self.noise_type = config.get("noise_type", "linear")
self.sigma_type = config.get("sigma_type", "zero") # "zero", "memoryless"
if self.noise_type == "exponential" or self.noise_type == "exponential_rev":
self.exp_max = config.get("exp_max", 5.0)
elif self.noise_type == "diffusion":
self.T = config.get("T", 1000)
self.beta_start = config.get("beta_start", 0.0001)
self.beta_end = config.get("beta_end", 0.02)
if self.sigma_type == "memoryless":
self.sigma_scale = config.get("sigma_scale", 1.0)
self.content_len = config.get("content_len", None)
# True (default): t_max covers the tail ramp-down in place (pairs
# with extra_len=0). False: original L/c; cover the tail with
# extra_len padding (typically chunk_size-1) at generate time.
self.t_max_tail = bool(config.get("t_max_tail", True))
# For simplicity we require steps to be divisible by chunk_size, so that time windows align well.
def _t_max_exact(self, valid_len):
"""Exact rational t upper bound (single source of truth).
t_max_tail: 1 + (L-1)/c so the last frame's band is covered in
place; otherwise the original L/c (tail covered via extra_len)."""
if self.t_max_tail:
return Fraction(valid_len + self.chunk_size - 1, self.chunk_size)
return Fraction(valid_len, self.chunk_size)
def get_total_steps(self, seq_len):
# Exact arithmetic: deriving via float double-rounds and can lose the
# final step (e.g. steps=30, c=30, L=94).
total = self.steps * self._t_max_exact(seq_len)
assert total.denominator == 1, (
f"inference grid does not close: steps={self.steps} * t_max"
f"({seq_len}+{self.chunk_size}-1)/{self.chunk_size} is not an "
f"integer; use steps divisible by chunk_size"
)
return int(total)
def get_t_max(self, valid_len):
"""Float view of _t_max_exact (for continuous t sampling)."""
return float(self._t_max_exact(valid_len))
def get_time_steps(self, device, valid_len, current_step):
"""current_step (int) → one t scalar shared by the whole batch,
returned as a per-sample List. Inference-only (generate / stream);
training's stratified sampling bypasses this and calls the
lower-level helpers directly."""
t = current_step * (1 / self.steps)
return [torch.tensor(t, device=device) for _ in range(len(valid_len))]
def get_time_schedules(self, device, valid_len, time_steps, training=False):
time_schedules = []
time_schedules_derivative = []
for i in range(len(valid_len)):
t = time_steps[i].item()
current_time_schedules = torch.clamp(
-torch.arange(valid_len[i], device=device) / self.chunk_size + t,
min=0.0,
max=1.0,
)
current_time_schedules_derivative = torch.ones_like(
current_time_schedules
) * (1 / self.steps)
if training:
current_time_schedules = torch.clamp(
current_time_schedules
+ torch.randn_like(current_time_schedules) * self.random_epsilon,
min=0.0,
max=1.0,
)
time_schedules.append(current_time_schedules)
time_schedules_derivative.append(current_time_schedules_derivative)
return time_schedules, time_schedules_derivative
def get_windows(self, valid_len, time_steps, training=False):
# Grid t (inference) lands exactly on band boundaries; +eps makes the
# half-open [0, 1) level convention robust to float rounding there.
# Continuous t (training) never hits boundaries; eps would only shift
# the band by eps/chunk_size, so it is skipped.
eps = 0.0 if training else 0.5 * (1 / (self.steps * self.chunk_size))
input_start, input_end, output_start, output_end = [], [], [], []
for i in range(len(time_steps)):
t = time_steps[i].item()
start_index = max(
0,
math.floor((t - 1) * self.chunk_size + eps) + 1,
)
end_index = min(
valid_len[i],
math.floor(t * self.chunk_size + eps) + 1,
)
# An exact boundary hit at the sequence end gives start == end;
# the last frame is still in the band, so fall back to [end-1, end).
start_index = min(start_index, max(end_index - 1, 0))
if self.content_len is not None:
input_start.append(max(0, end_index - self.content_len))
else:
input_start.append(0)
input_end.append(end_index)
output_start.append(start_index)
output_end.append(end_index)
return input_start, input_end, output_start, output_end
def get_noise_levels(self, device, valid_len, time_schedules):
alpha = []
dalpha = []
dlog_alpha = []
beta = []
dbeta = []
dlog_beta = []
sigma = []
for i in range(len(valid_len)):
t = time_schedules[i]
if self.noise_type == "linear":
alpha_i = t
dalpha_i = torch.ones_like(alpha_i)
dlog_alpha_i = dalpha_i / torch.clamp(alpha_i, min=EPSILON)
beta_i = 1 - t
dbeta_i = -torch.ones_like(beta_i)
dlog_beta_i = dbeta_i / torch.clamp(beta_i, min=EPSILON)
elif self.noise_type == "exponential":
# "eps" prediction
k = self.exp_max
alpha_i = torch.exp(-k * (1 - t))
dalpha_i = k * alpha_i
dlog_alpha_i = k * torch.ones_like(alpha_i)
beta_i = 1 - alpha_i
dbeta_i = -dalpha_i
dlog_beta_i = dbeta_i / torch.clamp(beta_i, min=EPSILON)
elif self.noise_type == "exponential_rev":
# "x0" prediction
k = self.exp_max
beta_i = torch.exp(-k * t)
dbeta_i = -k * beta_i
dlog_beta_i = -k * torch.ones_like(beta_i)
alpha_i = 1 - beta_i
dalpha_i = -dbeta_i
dlog_alpha_i = dalpha_i / torch.clamp(alpha_i, min=EPSILON)
elif self.noise_type == "diffusion":
t_rev = 1.0 - t
beta_rate = (
self.beta_start + t_rev * (self.beta_end - self.beta_start)
) * self.T
Gamma = (
self.beta_start * t_rev
+ 0.5 * (self.beta_end - self.beta_start) * t_rev * t_rev
) * self.T
alpha_i = torch.exp(-0.5 * Gamma)
dalpha_i = 0.5 * beta_rate * alpha_i
dlog_alpha_i = 0.5 * beta_rate
beta_i = torch.sqrt(torch.clamp(1 - torch.exp(-Gamma), min=0.0))
dbeta_i = (
-0.5
* torch.exp(-Gamma)
* beta_rate
/ torch.clamp(beta_i, min=EPSILON)
)
dlog_beta_i = dbeta_i / torch.clamp(beta_i, min=EPSILON)
else:
raise ValueError(f"Unknown noise type: {self.noise_type}")
alpha.append(torch.clamp(alpha_i, min=0.0, max=1.0))
dalpha.append(dalpha_i)
dlog_alpha.append(dlog_alpha_i)
beta.append(torch.clamp(beta_i, min=0.0, max=1.0))
dbeta.append(dbeta_i)
dlog_beta.append(dlog_beta_i)
if self.sigma_type == "zero":
sigma_i = torch.zeros_like(t)
elif self.sigma_type == "memoryless":
if (
self.noise_type == "linear"
or self.noise_type == "exponential"
or self.noise_type == "exponential_rev"
):
sigma_i = self.sigma_scale * torch.sqrt(
torch.clamp(2 * dlog_alpha_i * beta_i, min=0.0)
)
elif self.noise_type == "diffusion":
sigma_i = self.sigma_scale * torch.sqrt(
torch.clamp(2 * dlog_alpha_i, min=0.0)
)
else:
sigma_i = self.sigma_scale * torch.sqrt(
torch.clamp(
2 * beta_i * (dlog_alpha_i * beta_i - dbeta_i), min=0.0
)
)
sigma.append(sigma_i)
return alpha, dalpha, beta, dbeta, sigma, dlog_alpha, dlog_beta
def slice_input(self, x, input_start, input_end):
"""Inference-side: slice each sample's current state to the input window."""
return [x[i][:, input_start[i] : input_end[i], ...] for i in range(len(x))]
def prepare(self, x, device, valid_len, current_step):
"""Inference-side helper: current_step (int) → all schedule/noise/window
coefficients sliced to their input/output windows, plus xt sliced to
input window. Training-side stratified sampling does not go through
this helper; it calls the lower-level scheduler methods directly.
Returns dict with keys:
time_schedules, time_schedules_derivative,
alpha, dalpha, beta, dbeta, sigma, dlog_alpha, dlog_beta,
input_start, input_end, output_start, output_end, xt
"""
time_steps = self.get_time_steps(device, valid_len, current_step)
time_schedules, time_schedules_derivative = self.get_time_schedules(
device, valid_len, time_steps,
)
alpha, dalpha, beta, dbeta, sigma, dlog_alpha, dlog_beta = \
self.get_noise_levels(device, valid_len, time_schedules)
input_start, input_end, output_start, output_end = \
self.get_windows(valid_len, time_steps)
xt = self.slice_input(x, input_start, input_end)
# Slice all coefficients to their respective windows
batch_size = len(valid_len)
time_schedules = [time_schedules[i][input_start[i]:input_end[i]] for i in range(batch_size)]
time_schedules_derivative = [time_schedules_derivative[i][output_start[i]:output_end[i]] for i in range(batch_size)]
alpha = [alpha[i][output_start[i]:output_end[i]] for i in range(batch_size)]
dalpha = [dalpha[i][output_start[i]:output_end[i]] for i in range(batch_size)]
beta = [beta[i][output_start[i]:output_end[i]] for i in range(batch_size)]
dbeta = [dbeta[i][output_start[i]:output_end[i]] for i in range(batch_size)]
sigma = [sigma[i][output_start[i]:output_end[i]] for i in range(batch_size)]
dlog_alpha = [dlog_alpha[i][output_start[i]:output_end[i]] for i in range(batch_size)]
dlog_beta = [dlog_beta[i][output_start[i]:output_end[i]] for i in range(batch_size)]
return {
"time_schedules": time_schedules,
"time_schedules_derivative": time_schedules_derivative,
"input_start": input_start,
"input_end": input_end,
"output_start": output_start,
"output_end": output_end,
"alpha": alpha,
"dalpha": dalpha,
"beta": beta,
"dbeta": dbeta,
"sigma": sigma,
"dlog_alpha": dlog_alpha,
"dlog_beta": dlog_beta,
"xt": xt,
}
# --- Streaming support ---
def get_committable(self, total_frames):
"""Given total accumulated conditions, return how many frames can be committed.
Currently, we suppose steps % chunk_size == 0 for simplicity."""
committable_length = max(0, total_frames - self.chunk_size + 1)
committable_steps = total_frames * (self.steps // self.chunk_size)
return committable_length, committable_steps
def get_step_rollback(self, seq_len):
"""Get the step count to subtract when wrapping the buffer by seq_len.
Corresponds to how many steps were consumed by seq_len frames."""
steps = seq_len * (self.steps // self.chunk_size)
return steps
class T5TextCrossModule(nn.Module):
"""Cross-attention module for T5 text conditioning."""
def __init__(
self,
len=512,
dim=4096,
t5_size="xxl",
checkpoint_path=None,
tokenizer_path=None,
drop_out=0.1,
cross_rope=False,
input_keys={
"text": "text",
"text_end": "text_end",
},
cache_encoded=True,
):
assert checkpoint_path is not None and tokenizer_path is not None, (
"T5 checkpoint and tokenizer paths must be provided."
)
super().__init__()
self.len = len
self.dim = dim
self.cross_attn_norm = True
self.cross_rope = cross_rope
self.drop_out = drop_out
self.input_keys = input_keys
self.text_encoder = T5EncoderModel(
text_len=len,
dtype=torch.bfloat16,
device=torch.device("cpu"),
checkpoint_path=checkpoint_path,
tokenizer_path=tokenizer_path,
shard_fn=None,
t5_size=t5_size,
)
# When False, text_cache is cleared after each encode/get_context call
# so it doesn't accumulate across training steps (~190KB per caption,
# OOM on 1M+ unique captions).
self.cache_encoded = bool(cache_encoded)
self.text_cache = {}
def encode(self, text_list, device):
"""Encode text list with cache. Returns List[Tensor]."""
# Deduplicate uncached texts
texts_to_encode = []
for text in text_list:
if text not in self.text_cache and text not in texts_to_encode:
texts_to_encode.append(text)
# Batch encode deduplicated texts
if texts_to_encode:
self.text_encoder.model.to(device)
encoded = self.text_encoder(texts_to_encode, device)
for text, feature in zip(texts_to_encode, encoded):
self.text_cache[text] = feature.cpu()
# Collect from cache
out = [self.text_cache[text].to(device) for text in text_list]
if not self.cache_encoded:
self.text_cache.clear()
return out
def get_context(self, x, valid_len, device, param_dtype, training=False):
"""Mask-based unified context.
Both uniform (text=str) and multi-segment (text=List[str] + text_end)
inputs are funneled into a single representation: for each sample,
concatenate all (unique) segments along the K dimension and produce a
(T_i, S_i) cross-attn mask telling which tokens each query frame may see.
Returns:
context: List[Tensor (S_i, D_kv)] (S_i = num_seg_i * L_text)
metadata: dict with
'full_text': List[str] (logging only)
'cross_attn_mask': List[BoolTensor (T_i, S_i)]
'text_k_rope_ids' (cross_rope=True): List[Tuple]; f_ids
has shape (S_i,) and each segment's L_text tokens share the
segment-start-frame value. Null samples carry f_ids=0
throughout (single segment, L_text tokens).
'text_q_rope_ids' (cross_rope=True): List[Tuple|None]
per-sample Q rope override; non-None only for null samples
(zeros over T_i) so cross-attn phase = 0 (matches K=0).
"""
text_key = self.input_keys.get("text", "text")
text_end_key = self.input_keys.get("text_end", "text_end")
metadata = {}
text_input = x[text_key] if text_key in x else ["" for _ in range(len(valid_len))]
text_end_input = x.get(text_end_key, None)
# Normalize each sample to (text_list, text_end_list, is_null).
# text_end_list always starts with 0 and ends with valid_len[i].
norm = []
for i in range(len(valid_len)):
is_null = training and np.random.rand() <= self.drop_out
if is_null:
norm.append(([""], [0, valid_len[i]], True))
continue
if isinstance(text_input[i], list):
tl = text_input[i]
assert text_end_input is not None, (
"text_end required when text is List[str] (multi-segment)"
)
te = [0] + [min(int(t), valid_len[i]) for t in text_end_input[i]]
else:
tl = [text_input[i]]
te = [0, valid_len[i]]
# Force last segment to extend to valid_len so batch-padding
# frames (generate passes the batch-max seq_len as every sample's
# generated_len) get text K coverage instead of an all-False mask.
# Mirrors PositionCrossModule. No-op for the single-segment
# branch where te[-1] is already valid_len[i].
te[-1] = valid_len[i]
norm.append((tl, te, False))
# Encode all unique texts (cache-aware dedup)
texts_to_encode = []
for tl, _, _ in norm:
for t in tl:
if t not in self.text_cache and t not in texts_to_encode:
texts_to_encode.append(t)
if texts_to_encode:
self.text_encoder.model.to(device)
encoded = self.text_encoder(texts_to_encode, device)
for t, f in zip(texts_to_encode, encoded):
self.text_cache[t] = f.cpu()
all_context = []
all_mask = []
all_text_k_rope_ids = [] if self.cross_rope else None
text_q_rope_override = [] if self.cross_rope else None
full_text = []
for i, (tl, te, is_null) in enumerate(norm):
seg_tensors = [self.text_cache[t].to(device).to(param_dtype) for t in tl]
# T5 returns variable-length tokens per caption (trimmed to its
# actual mask), so segments may have different K-dim lengths.
# Use cumulative offsets, NOT k * L_text (would mis-stride and
# leak/lose K tokens whenever segments have different lengths).
seg_K_lens = [t.size(0) for t in seg_tensors]
seg_K_offsets = [0]
for L in seg_K_lens:
seg_K_offsets.append(seg_K_offsets[-1] + L)
seg_lens = [t - b for t, b in zip(te[1:], te[:-1])]
context_i = torch.cat(seg_tensors, dim=0) # (sum(seg_K_lens), D_kv)
S_i = context_i.size(0)
all_context.append(context_i)
full_text.append(
" ////////// ".join([f"{u} //dur:{d}" for u, d in zip(tl, seg_lens)])
)
# Frame f belongs to segment k where te[k] <= f < te[k+1];
# mark mask[f, seg_K_offsets[k]:seg_K_offsets[k+1]] = True.
T_i = valid_len[i]
mask = torch.zeros(T_i, S_i, dtype=torch.bool, device=device)
for k, (s_f, e_f) in enumerate(zip(te[:-1], te[1:])):
s_f = max(0, s_f)
e_f = min(T_i, e_f)
if s_f >= e_f:
continue
mask[s_f:e_f, seg_K_offsets[k]:seg_K_offsets[k+1]] = True
all_mask.append(mask)
if self.cross_rope:
# K rope: per-token segment-start-frame; null samples → 0.
f_ids = torch.empty(S_i, dtype=torch.long, device=device)
if is_null:
f_ids.zero_()
else:
for k, s_f in enumerate(te[:-1]):
f_ids[seg_K_offsets[k]:seg_K_offsets[k+1]] = s_f
all_text_k_rope_ids.append((f_ids, None, None))
# Q rope override: real → None (use self-attn default arange);
# null → zeros so phase = Q(0) - K(0) = 0 (rope-off equivalent).
if is_null:
text_q_rope_override.append(
(torch.zeros(T_i, dtype=torch.long, device=device), None, None)
)
else:
text_q_rope_override.append(None)
metadata["full_text"] = full_text
metadata["cross_attn_mask"] = all_mask
if self.cross_rope:
metadata["text_k_rope_ids"] = all_text_k_rope_ids
metadata["text_q_rope_ids"] = text_q_rope_override
if not self.cache_encoded:
self.text_cache.clear()
return all_context, metadata
def get_null_context(self, real_context, device, param_dtype, valid_len):
"""Build null context with the mask-path convention.
Each sample gets a single null segment of L_text tokens (encoded("")),
an all-True (T_i, L_text) mask, K rope = 0, Q rope override = 0.
This matches the original uniform-null math (phase = 0 = rope-off
equivalent) and is the only configuration consistent with the mask
path's unified shape.
Args:
real_context: List[Tensor], used to read batch size.
valid_len: List[int], per-sample query frame count (T_i).
Returns:
null_context: List[Tensor (L_text, D_kv)]
metadata: dict with cross_attn_mask + (when self.cross_rope=True)
text_k_rope_ids and text_q_rope_ids.
"""
batch_size = len(real_context)
encoded_null = self.encode([""], device)[0].to(param_dtype)
L_text = encoded_null.size(0)
metadata = {}
null_context = [encoded_null for _ in range(batch_size)]
metadata["cross_attn_mask"] = [
torch.ones(valid_len[i], L_text, dtype=torch.bool, device=device)
for i in range(batch_size)
]
if self.cross_rope:
metadata["text_k_rope_ids"] = [
(torch.zeros(L_text, dtype=torch.long, device=device), None, None)
for _ in range(batch_size)
]
metadata["text_q_rope_ids"] = [
(torch.zeros(valid_len[i], dtype=torch.long, device=device), None, None)
for i in range(batch_size)
]
return null_context, metadata
# --- Streaming state management (mask path) ---
def init_stream(self, batch_size):
# Per-sample list of segments; each entry is (text, start_frame_in_buffer).
# End frame of segment k is start of segment k+1 (or stream_frames[i]).
self.stream_segments = [[] for _ in range(batch_size)]
# Total frames pushed via update_stream (one per call).
self.stream_frames = [0 for _ in range(batch_size)]
def update_stream(self, x, device, param_dtype):
"""Add one frame of conditioning. If text matches the latest segment's
text, extend (implicit via stream_frames++); otherwise start a new
segment whose start_frame = current frame index.
Pre-encodes any uncached texts so get_stream_context is cache-only.
"""
text_key = self.input_keys.get("text", "text")
text_input = x[text_key]
# Cache uncached texts (incl. "" so null path can pull from cache too).
to_encode = []
for t in list(text_input) + [""]:
if t not in self.text_cache and t not in to_encode:
to_encode.append(t)
if to_encode:
self.text_encoder.model.to(device)
encoded = self.text_encoder(to_encode, device)
for t, f in zip(to_encode, encoded):
self.text_cache[t] = f.cpu()
for i in range(len(self.stream_segments)):
cur_text = text_input[i] if i < len(text_input) else ""
cur_frame = self.stream_frames[i]
if (
len(self.stream_segments[i]) > 0
and self.stream_segments[i][-1][0] == cur_text
):
# Same as latest segment → frame counter handles extension.
pass
else:
self.stream_segments[i].append((cur_text, cur_frame))
self.stream_frames[i] += 1
def get_stream_context(self, start_index, end_index, device, param_dtype):
"""Return mask-form context for buffer window [start_index, end_index).
Returns:
contexts: List[Tensor (S_i, D_kv)]
metadata: {
'cross_attn_mask': List[BoolTensor (window_len, S_i)]
'text_k_rope_ids' (cross_rope=True): List[Tuple]
f_ids = segment_start_frame (buffer-absolute), repeated L_text
times within each segment.
'text_q_rope_ids' (cross_rope=True): List[None] (real path).
}
"""
window_len = end_index - start_index
contexts = []
masks = []
rope_ids = [] if self.cross_rope else None
text_q_rope_overrides = [] if self.cross_rope else None
for i in range(len(self.stream_segments)):
segs = self.stream_segments[i]
if len(segs) == 0:
raise RuntimeError(
f"sample {i}: stream not initialized (call update_stream first)"
)
seg_starts = [s for _, s in segs]
seg_ends = seg_starts[1:] + [self.stream_frames[i]]
# Active = segments overlapping [start_index, end_index).
active = []
for k, (text, s_start) in enumerate(segs):
s_end = seg_ends[k]
if s_end <= start_index or s_start >= end_index:
continue
active.append((text, s_start, s_end))
if len(active) == 0:
# Window entirely outside any segment — fall back to null seg
# covering the whole window so attention still has valid K.
active = [("", start_index, end_index)]
seg_tensors = [
self.text_cache[text].to(device).to(param_dtype)
for (text, _, _) in active
]
# T5 returns variable-length per-caption tokens; segments may
# have different K-dim lengths. Use cumulative offsets, not
# k * L_text (would mis-stride and leak/lose K tokens).
seg_K_lens = [t.size(0) for t in seg_tensors]
seg_K_offsets = [0]
for L in seg_K_lens:
seg_K_offsets.append(seg_K_offsets[-1] + L)
context_i = torch.cat(seg_tensors, dim=0)
S_i = context_i.size(0)
contexts.append(context_i)
mask = torch.zeros(window_len, S_i, dtype=torch.bool, device=device)
for k, (_, s_start, s_end) in enumerate(active):
local_start = max(0, s_start - start_index)
local_end = min(window_len, s_end - start_index)
if local_start < local_end:
mask[local_start:local_end, seg_K_offsets[k]:seg_K_offsets[k+1]] = True
masks.append(mask)
if self.cross_rope:
f_ids = torch.empty(S_i, dtype=torch.long, device=device)
for k, (_, s_start, _) in enumerate(active):
f_ids[seg_K_offsets[k]:seg_K_offsets[k+1]] = s_start
rope_ids.append((f_ids, None, None))
text_q_rope_overrides.append(None)
metadata = {"cross_attn_mask": masks}
if self.cross_rope:
metadata["text_k_rope_ids"] = rope_ids
metadata["text_q_rope_ids"] = text_q_rope_overrides
return contexts, metadata
def trim_stream(self, trim_len):
"""Shift segment starts back by trim_len; drop segments fully out of buffer.
Negative starts are kept (clamped to ROPE_MIN_INDEX) so K rope can still
encode pre-buffer positions via rope_params_with_negative.
"""
for i in range(len(self.stream_segments)):
segs = self.stream_segments[i]
if len(segs) == 0:
self.stream_frames[i] -= trim_len
continue
seg_starts = [s for _, s in segs]
seg_ends = seg_starts[1:] + [self.stream_frames[i]]
new_segs = []
for k, (text, s_start) in enumerate(segs):
s_end = seg_ends[k]
new_start = s_start - trim_len
new_end = s_end - trim_len
if new_end <= 0:
continue # fully before buffer
new_segs.append((text, max(new_start, ROPE_MIN_INDEX)))
self.stream_segments[i] = new_segs
self.stream_frames[i] -= trim_len
class DiffForcingWanModel(nn.Module):
def __init__(
self,
input_dim=256,
mean_path=None,
std_path=None,
loss_W=None,
loss_w_coefficient=1.0,
hidden_dim=1024,
ffn_dim=2048,
freq_dim=256,
num_heads=8,
num_layers=8,
time_embedding_scale=1.0,
attn_type="full", # "causal" | "full" | "partial"
rope_channel_split=[1, 0, 0],
spatial_shape=(1, 1),
prediction_type="vel", # "vel", "x0", "eps"
text_config={
"len": 512,
"dim": 4096,
"cross_rope": False,
},
schedule_config={
"noise_type": "linear",
"chunk_size": 5,
"steps": 10,
"random_epsilon": 0.00,
},
cfg_config={
"text_scale": 5.0,
"null_scale": -4.0,
},
input_keys={
"feature": "feature",
"feature_length": "feature_length",
"text": "text",
"text_end": "text_end",
},
model_input_dim=None,
):
super().__init__()
self.input_keys = input_keys
self.mean_path = mean_path
self.std_path = std_path
self.input_dim = input_dim
loss_matrix = load_quadratic_loss_matrix(
loss_W, loss_w_coefficient, self.input_dim
)
if loss_matrix is None:
self.loss_matrix = None
else:
self.register_buffer("loss_matrix", loss_matrix)
# The diffusion state/output dimensionality remains ``input_dim``.
# Conditioned variants may expose extra clean channels only to the
# denoiser (packed root + latent), without diffusing or
# scoring those channels.
self.model_input_dim = (
self.input_dim if model_input_dim is None else int(model_input_dim)
)
self.spatial_shape = tuple(spatial_shape)
self.spatial_size = math.prod(self.spatial_shape)
self.hidden_dim = hidden_dim
self.ffn_dim = ffn_dim
self.freq_dim = freq_dim
self.num_heads = num_heads
self.num_layers = num_layers
self.time_embedding_scale = time_embedding_scale
assert attn_type in ("causal", "full", "partial"), (
f"attn_type must be one of 'causal'/'full'/'partial', got {attn_type!r}"
)
self.attn_type = attn_type
self.rope_channel_split = rope_channel_split
self.prediction_type = prediction_type
self.cfg_config = cfg_config
self.schedule_config = schedule_config
# The multi-window mask is only sound for attn_type='partial': it
# forces clean x clean to be causal and W_k x clean to be a strict
# prefix, which contradicts full's "no restriction" and causal's
# "global sequence causal". Force n_windows=1 under these attn_types
# so the packed sequence degenerates to a single window and the
# multi-window mask is bypassed in forward.
if self.attn_type in ("full", "causal"):
cfg_n = int(self.schedule_config.get("train_n_windows", 10))
if cfg_n != 1:
import warnings
warnings.warn(
f"attn_type={self.attn_type!r} is incompatible with "
f"multi-window training (multi-window mask differs from "
f"full/causal inference mask). "
f"Forcing train_n_windows from {cfg_n} to 1.",
stacklevel=2,
)
self.schedule_config["train_n_windows"] = 1
self.time_scheduler = TriangularTimeScheduler(schedule_config)
# Cross-attention module (text)
self.text_module = T5TextCrossModule(**text_config)
mean, std = load_statistics(self.mean_path, self.std_path, input_dim)
self.register_buffer("mean", mean)
self.register_buffer("std", std)
self.model = WanModel(
patch_size=(1, 1, 1),
text_len=self.text_module.len,
text_dim=self.text_module.dim,
cross_attn_norm=self.text_module.cross_attn_norm,
cross_rope=self.text_module.cross_rope,
in_dim=self.model_input_dim,
dim=self.hidden_dim,
ffn_dim=self.ffn_dim,
freq_dim=self.freq_dim,
out_dim=self.input_dim,
num_heads=self.num_heads,
num_layers=self.num_layers,
window_size=(-1, -1),
qk_norm=True,
eps=1e-6,
causal=(self.attn_type == "causal"),
rope_channel_split=self.rope_channel_split,
)
self.param_dtype = torch.float32
# ------------------------------------------------------------------
# Conditioning extension points
# ------------------------------------------------------------------
# The base model deliberately keeps these hooks as no-ops. Position
# conditioning lives in models/diffusion_forcing_position_wan.py while
# this class remains bit-for-bit compatible for existing configs.
def _sample_training_noise(self, x0, inputs, batch_index):
return torch.randn_like(x0)
def _get_diffusion_state(self, inputs):
return inputs["feature"]
def _prepare_training_condition(
self, inputs, normalized_feature, valid_len, device
):
return None
def _condition_training_input(
self, noisy, clean, condition, batch_index, start, end
):
return noisy
def _prepare_generation_condition(self, inputs, valid_len, seq_len, device):
return None
def _condition_generation_input(self, noisy, condition, starts, ends):
return noisy
def _apply_generation_state_condition(self, generated, condition):
return generated
def _finalize_generation_output(self, generated, condition, valid_len):
return generated
def _generation_output_metadata(self, condition, valid_len):
return {}
def _init_stream_condition(
self, batch_size, buf_len, device, history_position=None,
history_latent=None,
):
return None
def _update_stream_condition(self, inputs, device):
return None
def _commit_stream_condition(self, index):
return None
def _rollback_stream_condition(self, trim_len):
return None
def _condition_stream_input(self, noisy, starts, ends):
return noisy
def _apply_stream_state_condition(self):
return None
def _finalize_stream_output(self, generated, start, end):
return generated
def _stream_output_metadata(self, start, end):
return {}
def _extract_inputs(self, x):
"""Extract inputs from x using input_keys mapping."""
inputs = {}
for internal_key, external_key in self.input_keys.items():
if external_key in x:
inputs[internal_key] = x[external_key]
return inputs
def preprocess(self, x):
"""Convert last-channel format to channel-first, padding to 4D (C, T, H, W).
(T, C) -> (C, T, 1, 1)
(T, H, C) -> (C, T, H, 1)
(T, H, W, C) -> (C, T, H, W)
"""
for i in range(len(x)):
ndim = x[i].ndim
if ndim == 2: # (T, C)
x[i] = x[i].permute(1, 0)[:, :, None, None]
elif ndim == 3: # (T, H, C)
x[i] = x[i].permute(2, 0, 1)[:, :, :, None]
elif ndim == 4: # (T, H, W, C)
x[i] = x[i].permute(3, 0, 1, 2)
return x
def postprocess(self, x):
"""Reverse of preprocess: channel-first 4D back to last-channel, stripping padding dims.
(C, T, 1, 1) -> (T, C)
(C, T, H, 1) -> (T, H, C)
(C, T, H, W) -> (T, H, W, C)
"""
for i in range(len(x)):
shape = x[i].shape # (C, T, H, W)
if shape[2] == 1 and shape[3] == 1: # (C, T, 1, 1) -> (T, C)
x[i] = x[i][:, :, 0, 0].permute(1, 0)
elif shape[3] == 1: # (C, T, H, 1) -> (T, H, C)
x[i] = x[i][:, :, :, 0].permute(1, 2, 0)
else: # (C, T, H, W) -> (T, H, W, C)
x[i] = x[i].permute(1, 2, 3, 0)
return x
@staticmethod
def _shift_rope_ids(rope_ids, starts):
"""Subtract starts[i] from f_ids of each tuple, making K rope window-relative
so cross-attn phase matches Q's default-arange convention even when input
window doesn't start at frame 0 (content_len truncation).
Handles per-frame (List[List[Tuple]]), uniform (List[Tuple]), or per-sample
optional (List[Tuple|None]). None per-sample entries pass through.
Tuples with f_ids=None (null/CFG placeholders) are also passed through.
"""
if rope_ids is None:
return None
def shift_tuple(tup, s):
f, h, w = tup
return ((f - s) if f is not None else None, h, w)
result = []
for sample_ids, s in zip(rope_ids, starts):
if sample_ids is None:
result.append(None)
elif isinstance(sample_ids, list):
result.append([shift_tuple(t, s) for t in sample_ids])
else:
result.append(shift_tuple(sample_ids, s))
return result
def _patch_seq_len(self, frame_seq_len):
"""Frame-count seq_len → patch-count seq_len."""
return frame_seq_len * self.spatial_size
def _expand_schedule_to_patch(self, time_schedules_input):
"""Per-frame schedule → per-patch by repeating each frame value spatial_size times.
No-op when spatial_size == 1.
"""
if self.spatial_size == 1:
return time_schedules_input
return [s.repeat_interleave(self.spatial_size) for s in time_schedules_input]
def _expand_context_to_patch(self, context):
"""Per-frame nested context → per-patch by repeating each frame's context spatial_size times.
Uniform per-sample context (List[Tensor]) is returned unchanged.
No-op when spatial_size == 1.
"""
if self.spatial_size == 1:
return context
if len(context) > 0 and isinstance(context[0], (list, tuple)):
return [
[u for u in sample_ctx for _ in range(self.spatial_size)]
for sample_ctx in context
]
return context
def _expand_mask_to_patch(self, masks):
"""Cross-attn mask Q dim (frames) → patch dim by repeating each row
spatial_size times. K dim untouched. No-op when spatial_size == 1
or masks is None.
"""
if masks is None or self.spatial_size == 1:
return masks
return [m.repeat_interleave(self.spatial_size, dim=0) for m in masks]
def _build_partial_mask(self, input_start, input_end, output_start, output_end, device):
"""Per-sample partial-causal attention mask (used when attn_type='partial').
Layout per sample:
[input_start, output_start) → history: causal sliding window of size history_len
[output_start, output_end) → target : bidirectional + sees all history
When spatial > 1, repeat-interleave to patch level (F·H·W, F·H·W).
"""
spatial = self.spatial_size
masks = []
for is_, ie_, os_, _oe in zip(input_start, input_end, output_start, output_end):
L_frame = ie_ - is_
history_len = os_ - is_
sliding_window = history_len
m = torch.zeros(L_frame, L_frame, dtype=torch.bool, device=device)
for j in range(history_len):
m[j, max(0, j - sliding_window + 1):j + 1] = True
m[history_len:, :history_len] = True # target -> history (fully visible)
m[history_len:, history_len:] = True # target x target (bidirectional)
if spatial > 1:
m = m.repeat_interleave(spatial, dim=0).repeat_interleave(spatial, dim=1)
masks.append(m)
return masks
def _build_multi_window_attn_mask(self, clean_len, windows, clean_frame_start, device):
"""Self-attn mask for multi-window packed training.
Packed sequence layout per sample: [clean (clean_len tokens)] + concat over k
of W_k tokens, where W_k = frames [windows[k][0], windows[k][1]).
The clean section corresponds to frames
[clean_frame_start, clean_frame_start + clean_len).
- clean × clean: causal (clean[i] sees clean[0..i]).
- clean × W: blocked (clean doesn't depend on noised tokens).
- W_k × clean[0:s_k - clean_frame_start]: visible — W_k sees every
clean token in the packed sequence whose absolute frame index is
less than s_k (strict causal prefix in frame coords).
- W_k × W_k: bidirectional (denoising chunk full attention).
- W_k × W_j, j != k: blocked (independent denoising instances).
When spatial > 1, repeat-interleave to patch level
(L_frame·spatial, L_frame·spatial).
"""
sizes = [e - s for s, e in windows]
total_len = clean_len + sum(sizes)
mask = torch.zeros(total_len, total_len, dtype=torch.bool, device=device)
mask[:clean_len, :clean_len] = torch.tril(
torch.ones(clean_len, clean_len, dtype=torch.bool, device=device)
)
cursor = clean_len
for s, e in windows:
L = e - s
s_packed = s - clean_frame_start # packed col for frame s
if s_packed > 0:
mask[cursor:cursor + L, :s_packed] = True
mask[cursor:cursor + L, cursor:cursor + L] = True
cursor += L
if self.spatial_size > 1:
mask = mask.repeat_interleave(
self.spatial_size, dim=0
).repeat_interleave(self.spatial_size, dim=1)
return mask
@staticmethod
def _gather_for_packed(per_frame_tensor, clean_frame_start, clean_len, windows):
"""Gather rows from a per-frame tensor (T_i, ...) into the packed
layout [clean section | W_0 | W_1 | ...]. The clean section uses
rows [clean_frame_start, clean_frame_start + clean_len); each W_k
uses rows [s_k, e_k). Returns a tensor of shape (packed_len, ...)
with packed_len = clean_len + sum(L_k).
"""
parts = []
if clean_len > 0:
parts.append(per_frame_tensor[clean_frame_start:clean_frame_start + clean_len])
for s, e in windows:
parts.append(per_frame_tensor[s:e])
return torch.cat(parts, dim=0)
def _gather_text_q_rope_packed(self, q_rope_ids_full,
clean_frame_starts, clean_lens, all_windows):
"""Gather per-sample text_q_rope_ids tuples (per-sample List[Tuple|None])
to the packed layout. Per-sample None entries pass through; non-None
tuples have their f_ids (length T_i) gathered to packed_len following
the packed layout. h_ids / w_ids are forwarded unchanged.
"""
if q_rope_ids_full is None:
return None
result = []
for i, item in enumerate(q_rope_ids_full):
if item is None:
result.append(None)
continue
f_ids, h_ids, w_ids = item
if f_ids is None:
result.append(item)
continue
gathered = self._gather_for_packed(
f_ids, clean_frame_starts[i], clean_lens[i], all_windows[i],
)
result.append((gathered, h_ids, w_ids))
return result
def _build_multi_window_rope_ids(self, clean_len, windows, clean_frame_start, device):
"""Self-attn rope f_ids: per-token position **relative to clean_frame_start**.
Clean: arange(0, clean_len)。
W_k: arange(s_k - clean_frame_start, e_k - clean_frame_start)。
Using positions relative to clean_frame_start matches the inference
path's default arange(0, e) convention. The Q-K rope phase is still
the frame difference (the constant offset cancels out after
subtraction), so the semantics are unchanged.
"""
parts = [torch.arange(clean_len, device=device, dtype=torch.long)]
for s, e in windows:
parts.append(torch.arange(
s - clean_frame_start, e - clean_frame_start,
device=device, dtype=torch.long,
))
return (torch.cat(parts), None, None)
def forward(self, x):
"""Multi-window packed training.
Per sample (valid_len = T_i), divide t-range [0, (T_i+c-1)/c) into
`train_n_windows` equal buckets (default 10, configurable via
schedule_config); each bucket yields one uniform t_k. Windows are
derived from t_k via `time_scheduler.get_windows`, so the choice of
n_windows stays consistent with the inference path's formula and
training cost is independent of sequence length.
The packed self-attn sequence is
[clean feature | W_0 | W_1 | ... | W_{N-1}]
wired together by `_build_multi_window_attn_mask` and per-token rope
ids from `_build_multi_window_rope_ids`. Cross-attn is uniform (full
True mask, no rope override, no cross-rope). Loss is per-token MSE
averaged across all window tokens in the batch.
When spatial > 1, the metadata (time / attn_mask / cross_attn_mask /
seq_len) is expanded to patch granularity; rope_ids' f_ids stays at
frame level since rope_apply broadcasts across H/W internally.
"""
x = self._extract_inputs(x)
feature_original = self._get_diffusion_state(x)
feature_length = x["feature_length"]
feature_original = (feature_original - self.mean) / self.std
batch_size = feature_original.shape[0]
pad_seq_len = feature_original.shape[1]
device = feature_original.device
n_windows = int(self.schedule_config.get("train_n_windows", 10))
valid_len = []
feature = []
for i in range(batch_size):
length = min(feature_length[i].item(), pad_seq_len)
valid_len.append(length)
feature.append(feature_original[i, :length, ...])
feature = self.preprocess(feature) # list of (C, T_i, 1, 1)
training_condition = self._prepare_training_condition(
x, feature_original, valid_len, device,
)
# Text context + metadata. cross_attn_mask handles multi-segment Q->K
# visibility; text_k_rope_ids / text_q_rope_ids carry per-segment rope
# info when cross_rope=True (None otherwise).
context, ctx_meta = self.text_module.get_context(
x, valid_len, device, self.param_dtype, training=True,
)
cross_attn_mask_full = ctx_meta["cross_attn_mask"]
text_k_rope_ids_full = ctx_meta.get("text_k_rope_ids")
text_q_rope_ids_full = ctx_meta.get("text_q_rope_ids")
all_xt = []
all_t = []
all_rope = []
all_attn_mask = []
all_cross_attn_mask = []
all_targets = [] # per-sample list of (target_tensor, seq_start, seq_end)
# Per-sample packed layout (used after the loop to gather Q-side text rope
# and to shift K-side rope into the packed coordinate frame).
all_clean_frame_starts = []
all_clean_lens = []
all_packed_windows = []
for i in range(batch_size):
T_i = valid_len[i]
x0_i = feature[i]
eps_i = self._sample_training_noise(x0_i, x, i)
# Split t-range [0, t_max) into n_windows equal buckets; one
# uniform sample per bucket. Bucket size scales with t_max so the
# window count stays constant across sequence lengths.
t_max = self.time_scheduler.get_t_max(T_i)
bucket = t_max / n_windows
t_values = [
torch.tensor(
np.random.uniform(k * bucket, (k + 1) * bucket),
device=device,
)
for k in range(n_windows)
]
input_start_all, _, output_start, output_end = self.time_scheduler.get_windows(
[T_i] * n_windows, t_values, training=True,
)
windows = list(zip(output_start, output_end))
# t_values is non-decreasing => window starts are
# monotonically non-decreasing => clean section bounds:
# - clean_frame_start = min(input_start): the earliest window's
# history start
# - clean_frame_end = max(output_start) = windows[-1][0]
# Frames beyond clean_frame_end are not read by any W_k, and
# frames before clean_frame_start are never visible to any W_k
# (mask gates by frame index), so trimming the buffer to this
# span is sufficient.
clean_frame_start = min(input_start_all)
clean_frame_end = windows[-1][0]
clean_len = clean_frame_end - clean_frame_start
# N schedules over T_i frames, one per t_k. Reuses existing
# scheduler so noise_type / random_epsilon stay consistent with
# generate.
schedules, _ = self.time_scheduler.get_time_schedules(
device, [T_i] * n_windows, t_values, training=True,
)
alphas, dalphas, betas, dbetas, _sg, _dla, _dlb = (
self.time_scheduler.get_noise_levels(
device, [T_i] * n_windows, schedules,
)
)
# Clean section: same α·x0 + β·noise noising path as every window,
# so the latent agrees with the (possibly random_epsilon-jittered)
# schedule that's also handed to the model as time embedding.
clean_schedule, _ = self.time_scheduler.get_time_schedules(
device, [T_i], [torch.tensor(t_max, device=device)],
training=True,
)
clean_alphas, _, clean_betas, _, _, _, _ = (
self.time_scheduler.get_noise_levels(
device, [T_i], clean_schedule,
)
)
# Assemble packed sequence [clean | W_0 | ... | W_{N-1}].
ca = clean_alphas[0][clean_frame_start:clean_frame_end]
cb = clean_betas[0][clean_frame_start:clean_frame_end]
x0_c = x0_i[:, clean_frame_start:clean_frame_end, ...]
eps_c = eps_i[:, clean_frame_start:clean_frame_end, ...]
xt_c = (
x0_c * ca[None, :, None, None]
+ eps_c * cb[None, :, None, None]
)
xt_c = self._condition_training_input(
xt_c,
x0_c,
training_condition,
i,
clean_frame_start,
clean_frame_end,
)
parts_xt = [xt_c]
parts_time = [clean_schedule[0][clean_frame_start:clean_frame_end]]
targets = []
cursor = clean_len
for k in range(n_windows):
s, e = windows[k]
L = e - s
a = alphas[k][s:e]
b = betas[k][s:e]
x0_w = x0_i[:, s:e, ...]
eps_w = eps_i[:, s:e, ...]
xt_w = (
x0_w * a[None, :, None, None]
+ eps_w * b[None, :, None, None]
)
xt_w = self._condition_training_input(
xt_w, x0_w, training_condition, i, s, e,
)
parts_xt.append(xt_w)
parts_time.append(schedules[k][s:e])
if self.prediction_type == "vel":
da = dalphas[k][s:e]
db = dbetas[k][s:e]
tgt = (
x0_w * da[None, :, None, None]
+ eps_w * db[None, :, None, None]
)
elif self.prediction_type == "x0":
tgt = x0_w
elif self.prediction_type == "eps":
tgt = eps_w
else:
raise ValueError(
f"Unknown prediction_type: {self.prediction_type!r}"
)
targets.append((tgt, cursor, cursor + L))
cursor += L
xt_i = torch.cat(parts_xt, dim=1)
assert xt_i.size(0) == self.model_input_dim, (
f"conditioned denoiser input has {xt_i.size(0)} channels; "
f"expected model_input_dim={self.model_input_dim}"
)
time_i = torch.cat(parts_time) * self.time_embedding_scale
all_xt.append(xt_i)
all_t.append(time_i)
all_rope.append(
self._build_multi_window_rope_ids(
clean_len, windows, clean_frame_start, device,
)
)
all_attn_mask.append(
self._build_multi_window_attn_mask(
clean_len, windows, clean_frame_start, device,
)
)
# Cross-attn mask: gather per-frame rows from cross_attn_mask_full
# to follow the packed [clean | W_k] layout. Frame-level here;
# spatial expansion happens after the loop.
all_cross_attn_mask.append(
self._gather_for_packed(
cross_attn_mask_full[i], clean_frame_start, clean_len, windows,
)
)
all_targets.append(targets)
all_clean_frame_starts.append(clean_frame_start)
all_clean_lens.append(clean_len)
all_packed_windows.append(windows)
# Expand frame-level time / cross_attn_mask to patch-level
# (no-ops when spatial_size == 1).
all_t = self._expand_schedule_to_patch(all_t)
all_cross_attn_mask = self._expand_mask_to_patch(all_cross_attn_mask)
# seq_len = max(pad_seq_len, max_packed). In n=1 mode
# packed_length <= T_i <= pad_seq_len so pad_seq_len dominates and
# the model sees the same padded geometry as the inference path.
# With n>1 (partial only) packed can exceed pad_seq_len; the max
# keeps the upper bound correct.
seq_len_for_model = self._patch_seq_len(
max(pad_seq_len, max(xt.size(1) for xt in all_xt))
)
text_seq_len = max(int(ctx.size(0)) for ctx in context)
# Cross-attn rope (only active when cross_rope=True): K rope encodes
# per-token segment-start-frame in absolute coords, Q rope override
# exists for null samples (zeros over T_i). Both must be expressed in
# the same packed coordinate frame as self-attn rope_ids (which is
# relative to clean_frame_start), so:
# - K rope: shift by clean_frame_start (subtract) per sample
# - Q rope: gather per-frame f_ids to packed layout, then shift
text_k_rope_ids = self._shift_rope_ids(
text_k_rope_ids_full, all_clean_frame_starts,
)
text_q_rope_ids = self._gather_text_q_rope_packed(
text_q_rope_ids_full, all_clean_frame_starts, all_clean_lens,
all_packed_windows,
)
text_q_rope_ids = self._shift_rope_ids(
text_q_rope_ids, all_clean_frame_starts,
)
# The multi-window mask is only meaningful for attn_type='partial'.
# Under full/causal (where __init__ already forced n_windows=1), the
# packed sequence is a single window, so we let WanModel apply
# default full attention or its internal causal mask.
final_attn_mask = all_attn_mask if self.attn_type == "partial" else None
predicted_result = self.model(
all_xt,
all_t,
context,
seq_len_for_model,
attn_mask=final_attn_mask,
text_k_rope_ids=text_k_rope_ids,
text_q_rope_ids=text_q_rope_ids,
cross_attn_mask=all_cross_attn_mask,
text_seq_len=text_seq_len,
y=None,
rope_ids=all_rope,
)
# Per-sample mean -> batch mean. Within each sample sum SE across
# all W_k tokens and divide by the token count to get per-sample
# MSE, then average across the batch.
loss = 0.0
for i in range(batch_size):
pred_i = predicted_result[i]
sample_se_sum = pred_i.new_zeros(())
sample_n = 0
for tgt, s, e in all_targets[i]:
error = pred_i[:, s:e, ...] - tgt
if self.loss_matrix is None:
sample_se_sum = sample_se_sum + (error ** 2).sum()
else:
sample_se_sum = sample_se_sum + quadratic_error_sum(
error, self.loss_matrix
)
sample_n += error.numel()
loss = loss + sample_se_sum / max(sample_n, 1)
loss = loss / batch_size
return {"total": loss, "mse": loss}
def generate(self, x):
"""
Generation - Diffusion Forcing inference
Uses triangular noise schedule, progressively generating from left to right
Generation process:
1. Start from t=0, gradually increase t
2. Each t corresponds to a noise schedule: clean on left, noisy on right, gradient in middle
3. After each denoising step, t increases slightly and continues
"""
x = self._extract_inputs(x)
feature_length = x["feature_length"] # (B,)
batch_size = len(feature_length)
# extra_len pads the rollout with disposable frames so real frames
# finish under full-width windows; 0 relies on the tail ramp-down.
extra_len = self.schedule_config.get("extra_len", 0)
seq_len = max(feature_length).item() + extra_len
if (not getattr(self.time_scheduler, "t_max_tail", True)
and extra_len < self.time_scheduler.chunk_size - 1):
import warnings
warnings.warn(
f"t_max_tail=False with extra_len={extra_len} < chunk_size-1="
f"{self.time_scheduler.chunk_size - 1}: the last frames of "
f"every generated sequence will not finish denoising",
stacklevel=2,
)
device = next(self.parameters()).device
valid_len = []
for i in range(batch_size):
length = min(feature_length[i].item(), seq_len)
valid_len.append(length)
generated_len = [seq_len for _ in range(batch_size)]
generation_condition = self._prepare_generation_condition(
x, valid_len, seq_len, device,
)
# Initialize entire sequence as pure noise
generated = torch.randn(
batch_size, seq_len, *self.spatial_shape, self.input_dim, device=device
)
generated = [generated[i] for i in range(batch_size)]
generated = self.preprocess(generated)
self._apply_generation_state_condition(generated, generation_condition)
# Precompute real text context (null built per-step from window-sliced real).
text_context, metadata = self.text_module.get_context(
x,
generated_len,
device,
self.param_dtype,
training=False,
)
full_text = metadata["full_text"]
cross_attn_mask_full = metadata["cross_attn_mask"]
text_k_rope_ids_full = metadata.get("text_k_rope_ids")
total_steps = self.time_scheduler.get_total_steps(seq_len)
# Progressively advance from t=0 to t=max_t
for step in range(total_steps):
s = self.time_scheduler.prepare(
generated, device, generated_len, current_step=step,
)
time_schedules = s["time_schedules"]
time_schedules_derivative = s["time_schedules_derivative"]
alpha = s["alpha"]
dalpha = s["dalpha"]
beta = s["beta"]
dbeta = s["dbeta"]
sigma = s["sigma"]
dlog_alpha = s["dlog_alpha"]
dlog_beta = s["dlog_beta"]
input_start_index = s["input_start"]
input_end_index = s["input_end"]
output_start_index = s["output_start"]
output_end_index = s["output_end"]
xt = s["xt"]
xt = self._condition_generation_input(
xt, generation_condition, input_start_index, input_end_index,
)
# time_schedules already sliced to input window by prepare()
time_schedules_input = [
time_schedules[i] * self.time_embedding_scale
for i in range(batch_size)
]
# Slice cross_attn_mask Q dim to input window (K untouched).
window_cross_attn_mask = [
cross_attn_mask_full[i][input_start_index[i]:input_end_index[i], :]
for i in range(batch_size)
]
# Shift K rope to window-relative coords (lockstep with forward path).
window_text_k_rope = self._shift_rope_ids(text_k_rope_ids_full, input_start_index)
# CFG: text_scale * pred_text + null_scale * pred_null
ts_patch = self._expand_schedule_to_patch(time_schedules_input)
patch_seq_len = self._patch_seq_len(seq_len)
attn_mask = (
self._build_partial_mask(input_start_index, input_end_index,
output_start_index, output_end_index, device)
if self.attn_type == "partial" else None
)
# Build null context matching window-sliced real shape.
window_lens = [
input_end_index[i] - input_start_index[i] for i in range(batch_size)
]
window_null_context, window_null_meta = self.text_module.get_null_context(
text_context, device, self.param_dtype, valid_len=window_lens,
)
window_null_mask = window_null_meta["cross_attn_mask"]
window_null_text_k_rope = window_null_meta.get("text_k_rope_ids")
window_null_text_q_rope = window_null_meta.get("text_q_rope_ids")
# Doubled batch: [text..., null...]. One model.forward, split outputs.
combined_ctx = self._expand_context_to_patch(text_context + window_null_context)
combined_cross_attn_mask = self._expand_mask_to_patch(
window_cross_attn_mask + window_null_mask
)
if window_text_k_rope is not None and window_null_text_k_rope is not None:
combined_text_k_rope = window_text_k_rope + window_null_text_k_rope
else:
combined_text_k_rope = None
# Real branch has no Q-override (cross falls back to self_attn rope);
# null branch emits one when cross_rope=True. Concat [None]*B for real.
if window_null_text_q_rope is not None:
combined_text_q_rope = [None] * batch_size + list(window_null_text_q_rope)
else:
combined_text_q_rope = None
combined_mask = attn_mask + attn_mask if attn_mask is not None else None
combined_text_seq_len = max(
int(c.size(0)) for c in (text_context + window_null_context)
)
pred_all = self.model(
xt + xt,
ts_patch + ts_patch,
combined_ctx,
patch_seq_len,
attn_mask=combined_mask,
text_k_rope_ids=combined_text_k_rope,
text_q_rope_ids=combined_text_q_rope,
cross_attn_mask=combined_cross_attn_mask,
text_seq_len=combined_text_seq_len,
y=None,
)
pred_text = pred_all[:batch_size]
pred_null = pred_all[batch_size:]
predicted_result = [
self.cfg_config["text_scale"] * pt + self.cfg_config["null_scale"] * pn
for pt, pn in zip(pred_text, pred_null)
]
# All noise coefficients already sliced to output window by prepare()
for i in range(batch_size):
os, oe = output_start_index[i], output_end_index[i]
pred_os = os - input_start_index[i]
pred_oe = oe - input_start_index[i]
predicted_result_i = predicted_result[i][:, pred_os:pred_oe, ...]
generated_i = generated[i][:, os:oe, ...]
dt = time_schedules_derivative[i][None, :, None, None]
alpha_i = alpha[i][None, :, None, None]
dalpha_i = dalpha[i][None, :, None, None]
beta_i = beta[i][None, :, None, None]
dbeta_i = dbeta[i][None, :, None, None]
sigma_i = sigma[i][None, :, None, None]
dlog_alpha_i = dlog_alpha[i][None, :, None, None]
dlog_beta_i = dlog_beta[i][None, :, None, None]
if self.prediction_type == "vel":
vel = predicted_result_i
elif self.prediction_type == "x0":
vel = (
predicted_result_i * (-dlog_beta_i * alpha_i + dalpha_i)
+ generated_i * dlog_beta_i
)
elif self.prediction_type == "eps":
vel = (
predicted_result_i * (-dlog_alpha_i * beta_i + dbeta_i)
+ generated_i * dlog_alpha_i
)
st = (vel - generated_i * dlog_alpha_i) / (
(beta_i * dlog_alpha_i - dbeta_i) * beta_i
)
generated[i][:, os:oe, ...] += (
vel * dt
+ st * 0.5 * sigma_i**2 * dt
+ sigma_i * torch.sqrt(dt) * torch.randn_like(generated_i)
)
self._apply_generation_state_condition(generated, generation_condition)
generated = self.postprocess(generated) # list of (T, C)
y_hat_out = []
for i in range(batch_size):
single_generated = generated[i][: valid_len[i], :] * self.std + self.mean
y_hat_out.append(single_generated)
y_hat_out = self._finalize_generation_output(
y_hat_out, generation_condition, valid_len,
)
out = {}
out["generated"] = y_hat_out
out["text"] = full_text
out.update(self._generation_output_metadata(generation_condition, valid_len))
return out
def init_generated(
self,
seq_len,
batch_size=1,
schedule_config={},
history_latent=None,
history_position=None,
):
"""Initialize streaming generation state.
Args:
seq_len: Model window size (how many frames WanModel processes per step).
schedule_config: Optional schedule config overrides.
history_latent: optional list of (h, ...) un-normalized latent tensors per sample
(typically vae.encode(motion) output). If provided:
- normalize via self.mean/self.std
- apply triangular noise matching the schedule at t=1
(ts[i] = clamp(1 - i/chunk_size, 0, 1) for i in [0, h))
- inject into buffer[:, :h, ...]
- set current_step = steps so wavefront starts just past history
If None: cold start (buffer all noise, current_step=0).
Buffer is 2*seq_len. Model window is always buffer[0:seq_len].
When conditions overflow seq_len, shift buffer by seq_len and restart.
"""
self.schedule_config.update(schedule_config)
content_len = self.schedule_config.get("content_len", None)
if content_len is None:
self.schedule_config["content_len"] = seq_len
else:
self.schedule_config["content_len"] = min(seq_len, content_len)
self.time_scheduler = TriangularTimeScheduler(self.schedule_config)
self.batch_size = batch_size
self.seq_len = seq_len
self.buf_len = seq_len * 2
self.current_step = 0
self.current_commit = 0
self.condition_frames = 0
device = next(self.parameters()).device
# Initialize entire buffer as pure noise
generated = torch.randn(
batch_size, self.buf_len, *self.spatial_shape, self.input_dim, device=device
)
generated = [generated[i] for i in range(batch_size)]
self.generated = self.preprocess(generated)
# Initialize streaming state for cross module
self.text_module.init_stream(self.batch_size)
self._init_stream_condition(
batch_size,
self.buf_len,
device,
history_position=history_position,
history_latent=history_latent,
)
# Optional: inject GT history with triangular re-noising matching schedule at t=1
if history_latent is not None:
h = history_latent[0].shape[0]
chunk_size = self.time_scheduler.chunk_size
# Triangular ts: ts[0]=1 (clean), ts[chunk_size-1]=1/chunk_size (most noisy),
# ts[i>=chunk_size]=0 (full noise — beyond effective transition zone)
ts_history = torch.clamp(
1 - torch.arange(h, device=device, dtype=torch.float32) / chunk_size,
min=0.0, max=1.0,
)
alpha_h, _, beta_h, _, _, _, _ = self.time_scheduler.get_noise_levels(
device, [h], [ts_history],
)
alpha_b = alpha_h[0][None, :, None, None]
beta_b = beta_h[0][None, :, None, None]
for i in range(batch_size):
z = (history_latent[i].to(device) - self.mean) / self.std
z = self.preprocess([z])[0]
noise = torch.randn_like(z)
z_noised = z * alpha_b + noise * beta_b
self.generated[i][:, :h, ...] = z_noised
# Conditioned subclasses may need to restore channels that must
# remain clean after the base model's triangular history re-noising.
self._apply_stream_state_condition()
# t starts from 1 → wavefront just past the history transition zone
self.current_step = self.time_scheduler.steps
def _rollback(self):
"""Shift buffer by seq_len when conditions overflow the window."""
for i in range(self.batch_size):
self.generated[i][:, : self.seq_len, ...] = self.generated[i][
:, self.seq_len :, ...
].clone()
self.generated[i][:, self.seq_len :, ...] = torch.randn_like(
self.generated[i][:, self.seq_len :, ...]
)
self.current_step -= self.time_scheduler.get_step_rollback(self.seq_len)
self.condition_frames -= self.seq_len
self.current_commit -= self.seq_len
self.text_module.trim_stream(self.seq_len)
self._rollback_stream_condition(self.seq_len)
@torch.no_grad()
def stream_generate_step(self, x):
"""
Streaming generation step. Each call provides 1 frame of conditions.
The scheduler determines committable frames from accumulated conditions.
Returns:
dict with "generated": list of one (N, C) tensor, or [] if nothing to commit.
"""
x = self._extract_inputs(x)
device = next(self.parameters()).device
self.generated = [g.to(device) for g in self.generated]
# 1. Update conditions (1 frame per call)
self._update_stream_condition(x, device)
self.text_module.update_stream(x, device, self.param_dtype)
self.condition_frames += 1
# 2. Rollback if conditions overflow the window
if self.condition_frames > self.buf_len:
self._rollback()
self._commit_stream_condition(self.condition_frames - 1)
# 3. Determine how many frames can be committed
committable_length, committable_steps = self.time_scheduler.get_committable(
self.condition_frames
)
while self.current_step < committable_steps:
s = self.time_scheduler.prepare(
self.generated, device, [self.buf_len] * self.batch_size,
current_step=self.current_step,
)
time_schedules = s["time_schedules"]
time_schedules_derivative = s["time_schedules_derivative"]
alpha = s["alpha"]
dalpha = s["dalpha"]
beta = s["beta"]
dbeta = s["dbeta"]
sigma = s["sigma"]
dlog_alpha = s["dlog_alpha"]
dlog_beta = s["dlog_beta"]
is_ = s["input_start"]
ie_ = s["input_end"]
os_ = s["output_start"]
oe_ = s["output_end"]
xt = s["xt"]
xt = self._condition_stream_input(xt, is_, ie_)
# time_schedules already sliced to input window by prepare()
time_schedules_input = [
time_schedules[0] * self.time_embedding_scale
] * self.batch_size
# CFG: batch text + null in one forward pass
text_context, text_meta = self.text_module.get_stream_context(
is_[0], ie_[0], device, self.param_dtype,
)
text_mask = text_meta["cross_attn_mask"]
text_k_rope = text_meta.get("text_k_rope_ids")
# Q-rope override (per-sample List[Tuple|None]): real-text stream
# entries are None and pass through _shift unchanged.
text_q_rope = text_meta.get("text_q_rope_ids")
window_len = ie_[0] - is_[0]
null_context, null_meta = self.text_module.get_null_context(
text_context, device, self.param_dtype,
valid_len=[window_len] * self.batch_size,
)
null_mask = null_meta["cross_attn_mask"]
null_text_k_rope = null_meta.get("text_k_rope_ids")
null_text_q_rope = null_meta.get("text_q_rope_ids")
ts_patch = self._expand_schedule_to_patch(time_schedules_input)
ctx_patch = self._expand_context_to_patch(text_context + null_context)
cross_mask_patch = self._expand_mask_to_patch(text_mask + null_mask)
# K-rope ids: real is buffer-absolute → shift to window-relative; null
# comes from get_null_context already window-relative (zeros).
text_k_rope = self._shift_rope_ids(
text_k_rope, [is_[0]] * self.batch_size,
)
text_q_rope = self._shift_rope_ids(
text_q_rope, [is_[0]] * self.batch_size,
)
if text_k_rope is not None and null_text_k_rope is not None:
text_k_rope_patch = text_k_rope + null_text_k_rope
else:
text_k_rope_patch = None
# Q-override: real samples carry None (cross uses self_attn rope),
# null samples carry zeros tuple (phase=0 vs K=0).
if text_q_rope is not None or null_text_q_rope is not None:
text_q_rope_patch = (
(text_q_rope if text_q_rope is not None
else [None] * self.batch_size)
+ (null_text_q_rope if null_text_q_rope is not None
else [None] * self.batch_size)
)
else:
text_q_rope_patch = None
if self.attn_type == "partial":
masks = self._build_partial_mask(is_, ie_, os_, oe_, device)
attn_mask = masks + masks # CFG doubled batch (text + null), same mask
else:
attn_mask = None
stream_text_seq_len = max(
int(c.size(0)) for c in (text_context + null_context)
)
pred_all = self.model(
xt + xt,
ts_patch + ts_patch,
ctx_patch,
self._patch_seq_len(self.seq_len),
attn_mask=attn_mask,
text_k_rope_ids=text_k_rope_patch,
text_q_rope_ids=text_q_rope_patch,
cross_attn_mask=cross_mask_patch,
text_seq_len=stream_text_seq_len,
y=None,
)
pred_text = pred_all[: self.batch_size]
pred_null = pred_all[self.batch_size :]
predicted_result = [
self.cfg_config["text_scale"] * pt + self.cfg_config["null_scale"] * pn
for pt, pn in zip(pred_text, pred_null)
]
# All noise coefficients already sliced to output window by prepare()
os_idx, oe_idx = os_[0], oe_[0]
pred_os_idx = os_idx - is_[0]
pred_oe_idx = oe_idx - is_[0]
dt = time_schedules_derivative[0][None, :, None, None]
alpha_i = alpha[0][None, :, None, None]
dalpha_i = dalpha[0][None, :, None, None]
beta_i = beta[0][None, :, None, None]
dbeta_i = dbeta[0][None, :, None, None]
sigma_i = sigma[0][None, :, None, None]
dlog_alpha_i = dlog_alpha[0][None, :, None, None]
dlog_beta_i = dlog_beta[0][None, :, None, None]
for i in range(self.batch_size):
predicted_result_i = predicted_result[i][:, pred_os_idx:pred_oe_idx, ...]
generated_i = self.generated[i][:, os_idx:oe_idx, ...]
if self.prediction_type == "vel":
vel = predicted_result_i
elif self.prediction_type == "x0":
vel = (
predicted_result_i * (-dlog_beta_i * alpha_i + dalpha_i)
+ generated_i * dlog_beta_i
)
elif self.prediction_type == "eps":
vel = (
predicted_result_i * (-dlog_alpha_i * beta_i + dbeta_i)
+ generated_i * dlog_alpha_i
)
st = (vel - generated_i * dlog_alpha_i) / (
(beta_i * dlog_alpha_i - dbeta_i) * beta_i
)
self.generated[i][:, os_idx:oe_idx, ...] += (
vel * dt
+ st * 0.5 * sigma_i**2 * dt
+ sigma_i * torch.sqrt(dt) * torch.randn_like(generated_i)
)
self._apply_stream_state_condition()
self.current_step += 1
# 5. Extract newly committed frames
if self.current_commit < committable_length:
commit_start = self.current_commit
output = [
self.generated[i][:, commit_start:committable_length, ...]
for i in range(self.batch_size)
]
output = self.postprocess(output)
output = [o * self.std + self.mean for o in output]
output = self._finalize_stream_output(
output, commit_start, committable_length,
)
self.current_commit = committable_length
result = {"generated": output}
result.update(
self._stream_output_metadata(commit_start, committable_length)
)
return result
else:
empty = [
torch.zeros(self.input_dim, 0, *self.spatial_shape, device=device)
for _ in range(self.batch_size)
]
empty = self.postprocess(empty)
empty = self._finalize_stream_output(
empty, self.current_commit, self.current_commit,
)
result = {"generated": empty}
result.update(
self._stream_output_metadata(
self.current_commit, self.current_commit,
)
)
return result
# Position-specific root packing and model contracts.
from .vae_wan_position import (
POSITION_ROOT_DIMS,
WAN_ROOT_GROUP_SIZE,
pack_root_condition,
resolve_position_representation,
)
# Channels that are identically zero in the ground truth have std ~0. Dividing
# by that std turns float noise into an O(1) training target, so below this
# floor the channel is centred but not scaled. The floor sits in the measured
# gap observed in the supported motion statistics.
STD_NORMALIZATION_FLOOR = 1e-3
def guard_std(std, floor=STD_NORMALIZATION_FLOOR):
"""Replace near-zero standard deviations with 1 (i.e. do not divide)."""
std = torch.as_tensor(std).float()
return torch.where(std < floor, torch.ones_like(std), std)
def _to_batched_tensor(value, device, dtype=torch.float32):
if isinstance(value, (list, tuple)):
tensors = [torch.as_tensor(v, device=device, dtype=dtype) for v in value]
if not tensors:
raise ValueError("position condition cannot be empty")
if tensors[0].ndim == 0:
if any(tensor.ndim != 0 for tensor in tensors):
raise ValueError("position list mixes scalar and non-scalar values")
return torch.stack(tensors, dim=0)
if tensors[0].ndim == 1:
return torch.stack(tensors, dim=0)
lengths = {int(t.size(0)) for t in tensors}
if len(lengths) != 1:
raise ValueError(
"variable-length position lists must be padded/collated before use"
)
return torch.stack(tensors, dim=0)
return torch.as_tensor(value, device=device, dtype=dtype)
class _PositionStreamBufferMixin:
"""Fixed-size, rollback-safe raw condition buffer for stream generation."""
stream_condition_dim = 0
def _assert_stream_schedule(self):
"""Subclass hook: validated again after init_generated rebuilds it."""
return None
def init_generated(self, seq_len, *args, **kwargs):
# The base rebuilds time_scheduler from a caller-supplied
# schedule_config, so schedule invariants must be re-checked.
result = super().init_generated(seq_len, *args, **kwargs)
self._assert_stream_schedule()
return result
def _init_stream_condition(
self,
batch_size,
buf_len,
device,
history_position=None,
history_latent=None,
):
self._stream_position = torch.zeros(
batch_size,
buf_len,
self.stream_condition_dim,
device=device,
dtype=torch.float32,
)
self._stream_position_valid = torch.zeros(
batch_size, buf_len, device=device, dtype=torch.bool
)
self._pending_stream_position = None
if history_position is not None and history_latent is None:
raise ValueError("history_position requires matching history_latent")
history_length = None
if history_latent is not None:
if len(history_latent) != batch_size:
raise ValueError(
f"history latent batch={len(history_latent)}, expected {batch_size}"
)
history_lengths = [int(torch.as_tensor(v).size(0)) for v in history_latent]
if len(set(history_lengths)) != 1:
raise ValueError(
f"all history latent lengths must match, got {history_lengths}"
)
history_length = history_lengths[0]
if history_length > buf_len:
raise ValueError(
f"history latent length={history_length} exceeds buffer={buf_len}"
)
if history_position is None and history_latent is not None:
history_position = self._position_from_history(history_latent)
if history_position is not None:
history = self._coerce_stream_history(history_position, device)
if history.size(0) != batch_size:
raise ValueError(
f"history position batch={history.size(0)}, expected {batch_size}"
)
if history.size(1) > buf_len:
raise ValueError(
f"history position length={history.size(1)} exceeds buffer={buf_len}"
)
if history_length is not None and history.size(1) != history_length:
raise ValueError(
f"history position length={history.size(1)} does not match "
f"history latent length={history_length}"
)
self._stream_position[:, : history.size(1)] = history
self._stream_position_valid[:, : history.size(1)] = True
def _position_from_history(self, history_latent):
return None
def _coerce_stream_history(self, position, device):
value = _to_batched_tensor(position, device)
if value.ndim == 2:
value = value.unsqueeze(0)
if value.ndim != 3 or value.size(-1) != self.stream_condition_dim:
raise ValueError(
f"history position must be (B,T,{self.stream_condition_dim}), "
f"got {tuple(value.shape)}"
)
if not torch.isfinite(value).all():
raise ValueError("history position contains NaN or Inf")
return value.float()
def _update_stream_condition(self, inputs, device):
# Validation happens before text/counters are mutated in the base class.
value = self._parse_stream_position(inputs, device)
if value.shape != (self.batch_size, self.stream_condition_dim):
raise ValueError(
f"stream position must resolve to "
f"({self.batch_size},{self.stream_condition_dim}), "
f"got {tuple(value.shape)}"
)
if not torch.isfinite(value).all():
raise ValueError("stream position contains NaN or Inf")
self._pending_stream_position = value.float()
def _commit_stream_condition(self, index):
if self._pending_stream_position is None:
raise RuntimeError("stream position validation/commit transaction is missing")
if not 0 <= index < self.buf_len:
raise IndexError(
f"stream position index {index} outside buffer [0,{self.buf_len})"
)
self._stream_position[:, index] = self._pending_stream_position
self._stream_position_valid[:, index] = True
self._pending_stream_position = None
def _rollback_stream_condition(self, trim_len):
keep = self.buf_len - trim_len
self._stream_position[:, :keep] = self._stream_position[
:, trim_len:
].clone()
self._stream_position_valid[:, :keep] = self._stream_position_valid[
:, trim_len:
].clone()
self._stream_position[:, keep:] = 0
self._stream_position_valid[:, keep:] = False
def _stream_position_slice(self, start, end):
valid = self._stream_position_valid[:, start:end]
if valid.numel() and not bool(valid.all()):
missing = (~valid).nonzero(as_tuple=False)[0].tolist()
raise RuntimeError(
f"stream root condition missing inside active window "
f"[{start},{end}); first missing batch/local-index={missing}"
)
return self._stream_position[:, start:end]
class _RootConditionedDiffusion(
_PositionStreamBufferMixin, DiffForcingWanModel
):
"""Diffusion whose root channels are pure conditioning.
The root is never noised, never predicted and never scored: the diffused
state holds only the non-root channels, the root is concatenated in front
of the denoiser input at every training window and every denoise step, and
the emitted output is ``[supplied root, network state]``. Root output is
therefore bit-exact with the supplied root.
Subclasses define how a batch's raw root maps onto one condition vector per
state unit (``condition_dim``): the direct model uses one root vector per
frame, while the latent model packs four root vectors per latent token.
"""
condition_dim = 0
def _split_statistics(self, mean_path, std_path, expected_width):
"""Split one statistics pair into (condition, state) halves.
Both halves come from the same file, so the condition is normalized
with per-dimension statistics of the exact values it carries.
"""
if (mean_path is None) != (std_path is None):
raise ValueError("mean_path and std_path must be provided together")
mean, std = load_statistics(mean_path, std_path, expected_width)
mean, std = mean.reshape(-1), std.reshape(-1)
if mean.numel() != expected_width or std.numel() != expected_width:
raise ValueError(
f"statistics must have {expected_width} values, got "
f"{mean.numel()}/{std.numel()}"
)
std = guard_std(std)
return (
mean[: self.condition_dim].clone(),
std[: self.condition_dim].clone(),
mean[self.condition_dim :].clone(),
std[self.condition_dim :].clone(),
)
def _register_statistics(self, mean_path, std_path, expected_width):
cond_mean, cond_std, state_mean, state_std = self._split_statistics(
mean_path, std_path, expected_width
)
self.mean = state_mean
self.std = state_std
self.register_buffer("position_mean", cond_mean)
self.register_buffer("position_std", cond_std)
def _assert_stream_schedule(self):
"""Reject schedules that commit frames without ever denoising them.
``get_committable`` yields zero committable steps when
``steps < chunk_size``, so the loop body never runs while frames are
still emitted.
"""
steps = int(self.time_scheduler.steps)
chunk_size = int(self.time_scheduler.chunk_size)
if steps < chunk_size or steps % chunk_size != 0:
raise ValueError(
"streaming requires steps to be a positive multiple of "
f"chunk_size, got steps={steps}, chunk_size={chunk_size}"
)
# ------------------------------------------------------------------
# Condition extraction (subclass hooks)
# ------------------------------------------------------------------
def _split_feature(self, feature):
"""Return (condition_part, state_part) of a full-width feature."""
raise NotImplementedError
def _pack_position_batch(self, position, valid_len, seq_len, device):
"""Return (B, seq_len, condition_dim) raw conditions."""
raise NotImplementedError
def _extract_inputs(self, x):
inputs = super()._extract_inputs(x)
# Position-only and streaming callers commonly use the literal key
# without adding it to input_keys. Keep custom mappings authoritative,
# then fall back to the public ``position`` key.
if "position" not in inputs:
position_key = self.input_keys.get("position", "position")
if position_key in x:
inputs["position"] = x[position_key]
elif "position" in x:
inputs["position"] = x["position"]
feature = inputs.get("feature")
if feature is not None:
width = feature.size(-1)
if width == self.condition_dim + self.input_dim:
condition, state = self._split_feature(feature)
inputs["feature"] = state
inputs.setdefault("position", condition)
elif width != self.input_dim:
raise ValueError(
f"feature must have {self.condition_dim + self.input_dim} "
f"(root+state) or {self.input_dim} (state) channels, got {width}"
)
return inputs
def _get_raw_position(self, inputs):
position = inputs.get("position")
if position is None:
raise KeyError(
"root conditioning is required: supply input_keys.position or a "
"full-width feature whose leading channels carry the root"
)
return position
# ------------------------------------------------------------------
# Training
# ------------------------------------------------------------------
def _prepare_training_condition(
self, inputs, normalized_feature, valid_len, device
):
packed = self._pack_position_batch(
self._get_raw_position(inputs),
valid_len,
normalized_feature.size(1),
device,
)
normalized = (packed - self.position_mean) / self.position_std
return [
normalized[i].transpose(0, 1)[:, :, None, None]
for i in range(normalized.size(0))
]
def _condition_training_input(
self, noisy, clean, condition, batch_index, start, end
):
root = condition[batch_index][:, start:end]
return torch.cat([root.to(noisy), noisy], dim=0)
# ------------------------------------------------------------------
# Offline generation
# ------------------------------------------------------------------
def _prepare_generation_condition(self, inputs, valid_len, seq_len, device):
packed = self._pack_position_batch(
self._get_raw_position(inputs), valid_len, seq_len, device
)
normalized = (packed - self.position_mean) / self.position_std
return {
"raw": packed,
"normalized": [
normalized[i].transpose(0, 1)[:, :, None, None]
for i in range(normalized.size(0))
],
}
def _condition_generation_input(self, noisy, condition, starts, ends):
return [
torch.cat(
[
condition["normalized"][i][:, starts[i] : ends[i]].to(noisy[i]),
noisy[i],
],
dim=0,
)
for i in range(len(noisy))
]
def _finalize_generation_output(self, generated, condition, valid_len):
return [
torch.cat(
[condition["raw"][i, : valid_len[i]].to(generated[i]), generated[i]],
dim=-1,
)
for i in range(len(generated))
]
def _generation_output_metadata(self, condition, valid_len):
return {
"position": [
condition["raw"][i, : valid_len[i]].clone()
for i in range(len(valid_len))
]
}
# ------------------------------------------------------------------
# Streaming
# ------------------------------------------------------------------
def _condition_stream_input(self, noisy, starts, ends):
result = []
for i, item in enumerate(noisy):
raw = self._stream_position_slice(starts[i], ends[i])[i]
norm = (raw - self.position_mean) / self.position_std
root = norm.transpose(0, 1)[:, :, None, None].to(item)
result.append(torch.cat([root, item], dim=0))
return result
def _finalize_stream_output(self, generated, start, end):
raw = self._stream_position_slice(start, end)
return [
torch.cat([raw[i].to(generated[i]), generated[i]], dim=-1)
for i in range(len(generated))
]
def _stream_output_metadata(self, start, end):
# Clone: the slice is a view of the rollback-mutated ring buffer.
raw = self._stream_position_slice(start, end)
return {"position": [raw[i].clone() for i in range(self.batch_size)]}
class DiffForcingPositionWanModel(_RootConditionedDiffusion):
"""Direct full-feature diffusion conditioned on clean root motion.
``input_dim`` keeps its config meaning (the complete feature width); the
diffused state contains every non-root channel. The denoiser reads the
clean root plus noisy pose state and predicts only the pose state, so the
loss never covers root motion and emitted root channels are exactly the
supplied values.
"""
def __init__(
self,
*args,
input_dim=138,
representation=None,
root_dim=None,
mean_path=None,
std_path=None,
root_mean_path=None,
root_std_path=None,
**kwargs,
):
self.feature_dim = int(input_dim)
self.representation, self.root_dim = resolve_position_representation(
representation,
feature_dim=self.feature_dim,
root_dim=root_dim,
)
self.condition_dim = self.root_dim
if not 0 < self.root_dim < self.feature_dim:
raise ValueError(
f"root_dim must be in (0,{self.feature_dim}), got {self.root_dim}"
)
kwargs.pop("model_input_dim", None)
super().__init__(
*args,
input_dim=self.feature_dim - self.root_dim,
model_input_dim=self.feature_dim,
mean_path=None,
std_path=None,
**kwargs,
)
self._assert_stream_schedule()
self.stream_condition_dim = self.root_dim
self._register_statistics(mean_path, std_path, self.feature_dim)
if root_mean_path is not None or root_std_path is not None:
# Optional per-dimension root statistics can override the leading
# values from the full-feature statistics.
if root_mean_path is None or root_std_path is None:
raise ValueError(
"root_mean_path and root_std_path must be provided together"
)
root_mean, root_std = load_statistics(
root_mean_path, root_std_path, self.root_dim
)
root_mean, root_std = root_mean.reshape(-1), root_std.reshape(-1)
if root_mean.numel() < self.root_dim or root_std.numel() < self.root_dim:
raise ValueError(
f"root statistics must have at least {self.root_dim} values"
)
self.position_mean = root_mean[: self.root_dim].clone()
self.position_std = guard_std(root_std[: self.root_dim]).clone()
def _split_feature(self, feature):
return feature[..., : self.root_dim], feature[..., self.root_dim :]
def _pack_position_batch(self, position, valid_len, seq_len, device):
position = _to_batched_tensor(position, device)
if position.ndim == 2:
position = position.unsqueeze(0)
if position.ndim != 3 or position.size(-1) < self.root_dim:
raise ValueError(
f"direct root must be (B,T,>={self.root_dim}), got "
f"{tuple(position.shape)}"
)
if position.size(0) != len(valid_len):
raise ValueError(
f"root batch={position.size(0)}, expected {len(valid_len)}"
)
raw = position[..., : self.root_dim].float()
padded = []
for i, length in enumerate(valid_len):
if raw.size(1) < length:
raise ValueError(
f"sample {i} has {raw.size(1)} root frames but needs {length}"
)
item = raw[i, :length]
if seq_len > length:
if length == 0:
raise ValueError("cannot pad an empty valid root sequence")
item = torch.cat([item, item[-1:].expand(seq_len - length, -1)], dim=0)
padded.append(item)
packed = torch.stack(padded)
if not torch.isfinite(packed).all():
raise ValueError("root condition contains NaN or Inf")
return packed
def _position_from_history(self, history_latent):
tensors = [torch.as_tensor(v) for v in history_latent]
if not tensors:
return None
width = int(tensors[0].size(-1))
if width == self.feature_dim:
return torch.stack([v[..., : self.root_dim] for v in tensors])
return None
def init_generated(
self,
seq_len,
batch_size=1,
schedule_config={},
history_latent=None,
history_position=None,
):
if history_latent is not None:
# A full-width history carries its own root, which must be pulled
# out here: the diffused state has no root channels for the base
# class to recover it from.
split_state = []
inferred_position = []
for item in history_latent:
item = torch.as_tensor(item)
width = int(item.size(-1))
if width == self.feature_dim:
inferred_position.append(item[..., : self.root_dim])
split_state.append(item[..., self.root_dim :])
elif width == self.input_dim:
split_state.append(item)
else:
raise ValueError(
f"history must be {self.feature_dim}D (root+state) or "
f"{self.input_dim}D (state), got {width}D"
)
if inferred_position:
if len(inferred_position) != len(split_state):
raise ValueError("cannot mix full-width and state-only history")
if history_position is not None:
raise ValueError(
"history_position must be omitted for a full-width history"
)
history_position = inferred_position
elif history_position is None:
raise ValueError("a state-only history requires history_position")
history_latent = split_state
return super().init_generated(
seq_len,
batch_size=batch_size,
schedule_config=schedule_config,
history_latent=history_latent,
history_position=history_position,
)
def _parse_stream_position(self, inputs, device):
position = _to_batched_tensor(self._get_raw_position(inputs), device)
if position.ndim == 3 and position.size(1) == 1:
position = position[:, 0]
if position.ndim == 1:
position = position.unsqueeze(0)
if position.ndim != 2 or position.size(-1) < self.root_dim:
raise ValueError(
f"each direct stream step needs (B,{self.root_dim}) root, "
f"got {tuple(position.shape)}"
)
return position[..., : self.root_dim]
class LatentDiffForcingPositionWanModel(_RootConditionedDiffusion):
"""Latent diffusion conditioned on the WAN-causal packed root.
``input_dim`` keeps its config meaning (the VAE latent width); the diffused
state is exactly that latent. The denoiser reads ``position_dim + latent``
channels and predicts only the latent, so the emitted token is
``[supplied packed root, network latent]`` and its root is bit-exact.
The packed root width is inferred, in order, from ``representation``, an
explicit ``root_dim``/``position_dim``, or the token statistics width. The
last path keeps existing configs working when they do not forward the
top-level representation into model parameters.
"""
def __init__(
self,
*args,
input_dim=16,
representation=None,
position_dim=None,
root_dim=None,
root_group_size=WAN_ROOT_GROUP_SIZE,
mean_path=None,
std_path=None,
**kwargs,
):
self.root_group_size = int(root_group_size)
if self.root_group_size != WAN_ROOT_GROUP_SIZE:
raise ValueError(
"the WAN position ABI groups exactly "
f"{WAN_ROOT_GROUP_SIZE} root frames per latent, got "
f"root_group_size={self.root_group_size}"
)
self.latent_dim = int(input_dim)
if self.latent_dim <= 0:
raise ValueError(f"input_dim must be positive, got {self.latent_dim}")
resolved_root_dim = None if root_dim is None else int(root_dim)
self.representation = None
if representation is not None:
self.representation, resolved_root_dim = (
resolve_position_representation(
representation,
root_dim=resolved_root_dim,
)
)
resolved_position_dim = (
None if position_dim is None else int(position_dim)
)
if resolved_position_dim is not None:
if resolved_position_dim <= 0:
raise ValueError(
f"position_dim must be positive, got {resolved_position_dim}"
)
if resolved_position_dim % self.root_group_size != 0:
raise ValueError(
"position_dim must be divisible by root_group_size "
f"({self.root_group_size}), got {resolved_position_dim}"
)
position_root_dim = resolved_position_dim // self.root_group_size
if (
resolved_root_dim is not None
and resolved_root_dim != position_root_dim
):
raise ValueError(
"position_dim and root_dim disagree: "
f"{resolved_position_dim} != {resolved_root_dim} * "
f"{self.root_group_size}"
)
resolved_root_dim = position_root_dim
elif resolved_root_dim is not None:
resolved_position_dim = resolved_root_dim * self.root_group_size
elif (
mean_path is not None and std_path is not None
and Path(mean_path).is_file() and Path(std_path).is_file()
):
mean_width = int(np.asarray(np.load(mean_path)).size)
std_width = int(np.asarray(np.load(std_path)).size)
if mean_width != std_width:
raise ValueError(
"token statistics widths must match, got "
f"{mean_width}/{std_width}"
)
resolved_position_dim = mean_width - self.latent_dim
if (
resolved_position_dim <= 0
or resolved_position_dim % self.root_group_size != 0
):
raise ValueError(
f"cannot infer packed root width from {mean_width}D token "
f"statistics, latent_dim={self.latent_dim}, and "
f"root_group_size={self.root_group_size}"
)
resolved_root_dim = resolved_position_dim // self.root_group_size
if resolved_root_dim not in set(POSITION_ROOT_DIMS.values()):
raise ValueError(
f"cannot infer a supported root width from {mean_width}D "
f"token statistics: inferred root_dim={resolved_root_dim}, "
f"supported root widths are "
f"{sorted(set(POSITION_ROOT_DIMS.values()))}; pass "
"root_dim/position_dim explicitly for a custom schema"
)
else:
# Historical root3/root12 default for configs without token stats.
_, resolved_root_dim = resolve_position_representation()
resolved_position_dim = resolved_root_dim * self.root_group_size
# If a representation supplied the root width, also validate any width
# inferred from position_dim or token statistics.
self.representation, resolved_root_dim = resolve_position_representation(
self.representation,
root_dim=resolved_root_dim,
)
self.root_dim = resolved_root_dim
self.position_dim = resolved_position_dim
self.condition_dim = self.position_dim
if self.position_dim != self.root_dim * self.root_group_size:
raise ValueError(
"position_dim must equal root_dim * root_group_size "
f"({self.root_dim}*{self.root_group_size})"
)
kwargs.pop("model_input_dim", None)
super().__init__(
*args,
input_dim=self.latent_dim,
model_input_dim=self.position_dim + self.latent_dim,
mean_path=None,
std_path=None,
**kwargs,
)
self._assert_stream_schedule()
self.stream_condition_dim = self.position_dim
self._register_statistics(
mean_path, std_path, self.position_dim + self.latent_dim
)
def _split_feature(self, feature):
return feature[..., : self.position_dim], feature[..., self.position_dim :]
def _pack_position_batch(self, position, valid_len, seq_len, device):
position = _to_batched_tensor(position, device)
if position.ndim == 2:
position = position.unsqueeze(0)
if position.ndim != 3:
raise ValueError(
f"latent root must be (B,L,{self.position_dim}) or (B,T,>="
f"{self.root_dim}), got {tuple(position.shape)}"
)
if position.size(0) != len(valid_len):
raise ValueError(
f"root batch={position.size(0)}, expected {len(valid_len)}"
)
packed_items = []
for i, length in enumerate(valid_len):
item = position[i : i + 1]
if item.size(-1) == self.position_dim:
if item.size(1) < length:
raise ValueError(
f"sample {i} has {item.size(1)} packed roots but needs {length}"
)
packed = item[:, :length]
else:
required_frames = (
0 if length == 0 else 1 + self.root_group_size * (length - 1)
)
if item.size(-1) < self.root_dim or item.size(1) < required_frames:
raise ValueError(
f"sample {i} needs {required_frames} raw root frames for "
f"{length} latents, got {item.size(1)}"
)
# Trim to this sample's own frames: in a collated batch the
# tensor is padded to the batch maximum, and pack_root_condition
# clamps its taps to the frame count it is given, so an untrimmed
# item would let a short sample's tail taps land on pad zeros.
item = item[:, :required_frames]
packed = pack_root_condition(
item[..., : self.root_dim],
latent_length=length,
root_dim=self.root_dim,
group_size=self.root_group_size,
first_chunk=True,
)
packed = packed[0]
if seq_len > length:
if length == 0:
raise ValueError("cannot pad an empty valid root sequence")
packed = torch.cat(
[packed, packed[-1:].expand(seq_len - length, -1)], dim=0
)
packed_items.append(packed)
packed = torch.stack(packed_items).float()
if not torch.isfinite(packed).all():
raise ValueError("latent root condition contains NaN or Inf")
return packed
def _position_from_history(self, history_latent):
tensors = [torch.as_tensor(v) for v in history_latent]
if not tensors:
return None
width = int(tensors[0].size(-1))
if width == self.position_dim + self.latent_dim:
return torch.stack([v[..., : self.position_dim] for v in tensors])
return None
def init_generated(
self,
seq_len,
batch_size=1,
schedule_config={},
history_latent=None,
history_position=None,
):
if history_latent is not None:
token_dim = self.position_dim + self.latent_dim
split_state = []
inferred_position = []
for item in history_latent:
item = torch.as_tensor(item)
width = int(item.size(-1))
if width == token_dim:
inferred_position.append(item[..., : self.position_dim])
split_state.append(item[..., self.position_dim :])
elif width == self.latent_dim:
split_state.append(item)
else:
raise ValueError(
f"latent history must be {token_dim}D (root+latent) or "
f"{self.latent_dim}D (latent), got {width}D"
)
if inferred_position:
if len(inferred_position) != len(split_state):
raise ValueError("cannot mix token-width and latent-only history")
if history_position is not None:
raise ValueError(
"history_position must be omitted for a token-width history"
)
history_position = inferred_position
elif history_position is None:
raise ValueError("a latent-only history requires history_position")
history_latent = split_state
return super().init_generated(
seq_len,
batch_size=batch_size,
schedule_config=schedule_config,
history_latent=history_latent,
history_position=history_position,
)
def _parse_stream_position(self, inputs, device):
position = _to_batched_tensor(self._get_raw_position(inputs), device)
if position.ndim == 3 and position.shape[-2:] == (
self.root_group_size,
self.root_dim,
):
position = position.flatten(-2)
elif position.ndim == 3 and position.size(1) == 1:
position = position[:, 0]
if position.ndim == 1:
position = position.unsqueeze(0)
if position.ndim != 2 or position.size(-1) != self.position_dim:
raise ValueError(
f"each latent stream step needs (B,{self.position_dim}) or "
f"(B,{self.root_group_size},{self.root_dim}) root, got "
f"{tuple(position.shape)}"
)
return position
# Concise aliases for configs and downstream imports.
DiffForcingWanPositionModel = DiffForcingPositionWanModel
LatentDiffForcingWanPositionModel = LatentDiffForcingPositionWanModel