xsl
24c526af68
- 模型 17/17 文件经镜像全量缓存（gated+token，直连本机不通；sfx.md/setup.md/asset-pipeline.md 镜像结论同步更新） - 修复迁移遗留：stable_audio_3/models/ 源码子包被 .gitignore models/ 规则误伤丢失，从上游恢复并加白名单 - radio：arg_muffled/brake_screech 等叙事件落地（接线原有 L111/L154 零改动）+ phone_ring SYN→SA3 升级；dungeon 13 源 A/B 对齐（RMS 偏差 ≤0.3dB）全量同名替换，SYN 原版备份 tools/asset_fix_backup/sa3_ab_20260902/ - 关键玩法语义达标：echo 距离档 0.83<1.43<2.18<2.98s 严格递增；one-shot 真峰值全 ≤-3dBFS（解码回读法） - check_all_assets 8 游戏全绿（radio 254/254、dungeon 303/303）；playtest radio 221/dungeon 1610 全绿
643 lines
26 KiB
Python
643 lines
26 KiB
Python
import typing as tp
|
|
import math
|
|
import torch
|
|
|
|
from einops import rearrange
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
|
|
from .blocks import FourierFeatures, ExpoFourierFeatures
|
|
from .transformer import ContinuousTransformer
|
|
from .lora import LoRAParametrization, set_lora_strength, has_lora, enable_lora, disable_lora, filter_lora_layers
|
|
|
|
class DiffusionTransformer(nn.Module):
|
|
def __init__(self,
|
|
io_channels=32,
|
|
patch_size=1,
|
|
embed_dim=768,
|
|
cond_token_dim=0,
|
|
project_cond_tokens=True,
|
|
global_cond_dim=0,
|
|
project_global_cond=True,
|
|
input_concat_dim=0,
|
|
prepend_cond_dim=0,
|
|
depth=12,
|
|
num_heads=8,
|
|
transformer_type: tp.Literal["continuous_transformer", "mm_transformer"] = "continuous_transformer",
|
|
global_cond_type: tp.Literal["prepend", "adaLN"] = "prepend",
|
|
timestep_cond_type: tp.Literal["global", "input_concat"] = "global",
|
|
timestep_embed_dim=None,
|
|
diffusion_objective: tp.Literal["v", "rectified_flow", "rf_denoiser"] = "v",
|
|
timestep_features_type: tp.Literal["learned", "expo"] = "learned",
|
|
timestep_features_dim = 256,
|
|
timestep_features_logsnr: bool = False,
|
|
modular_local_cond_configs = None,
|
|
**kwargs):
|
|
|
|
super().__init__()
|
|
|
|
self.cond_token_dim = cond_token_dim
|
|
|
|
# Timestep embeddings
|
|
self.timestep_cond_type = timestep_cond_type
|
|
self.timestep_features_logsnr = timestep_features_logsnr
|
|
|
|
timestep_features_dim = timestep_features_dim
|
|
|
|
if timestep_features_type == "expo":
|
|
self.timestep_features = ExpoFourierFeatures(timestep_features_dim, 0.5, 10000.0)
|
|
else:
|
|
self.timestep_features = FourierFeatures(1, timestep_features_dim)
|
|
|
|
if timestep_cond_type == "global":
|
|
timestep_embed_dim = embed_dim
|
|
elif timestep_cond_type == "input_concat":
|
|
assert timestep_embed_dim is not None, "timestep_embed_dim must be specified if timestep_cond_type is input_concat"
|
|
input_concat_dim += timestep_embed_dim
|
|
|
|
self.to_timestep_embed = nn.Sequential(
|
|
nn.Linear(timestep_features_dim, timestep_embed_dim, bias=True),
|
|
nn.SiLU(),
|
|
nn.Linear(timestep_embed_dim, timestep_embed_dim, bias=True),
|
|
)
|
|
|
|
self.diffusion_objective = diffusion_objective
|
|
|
|
if cond_token_dim > 0:
|
|
# Conditioning tokens
|
|
|
|
cond_embed_dim = cond_token_dim if not project_cond_tokens else embed_dim
|
|
self.to_cond_embed = nn.Sequential(
|
|
nn.Linear(cond_token_dim, cond_embed_dim, bias=False),
|
|
nn.SiLU(),
|
|
nn.Linear(cond_embed_dim, cond_embed_dim, bias=False)
|
|
)
|
|
else:
|
|
cond_embed_dim = 0
|
|
|
|
if global_cond_dim > 0:
|
|
# Global conditioning
|
|
global_embed_dim = global_cond_dim if not project_global_cond else embed_dim
|
|
self.to_global_embed = nn.Sequential(
|
|
nn.Linear(global_cond_dim, global_embed_dim, bias=False),
|
|
nn.SiLU(),
|
|
nn.Linear(global_embed_dim, global_embed_dim, bias=False)
|
|
)
|
|
|
|
if prepend_cond_dim > 0:
|
|
# Prepend conditioning
|
|
self.to_prepend_embed = nn.Sequential(
|
|
nn.Linear(prepend_cond_dim, embed_dim, bias=False),
|
|
nn.SiLU(),
|
|
nn.Linear(embed_dim, embed_dim, bias=False)
|
|
)
|
|
|
|
self.input_concat_dim = input_concat_dim
|
|
|
|
dim_in = io_channels + self.input_concat_dim
|
|
|
|
self.patch_size = patch_size
|
|
|
|
# Transformer
|
|
|
|
self.transformer_type = transformer_type
|
|
|
|
self.global_cond_type = global_cond_type
|
|
|
|
transformer_dim_out = io_channels * patch_size
|
|
|
|
if self.transformer_type == "continuous_transformer":
|
|
|
|
global_dim = None
|
|
|
|
if self.global_cond_type == "adaLN":
|
|
# The global conditioning is projected to the embed_dim already at this point
|
|
global_dim = embed_dim
|
|
|
|
self.transformer = ContinuousTransformer(
|
|
dim=embed_dim,
|
|
depth=depth,
|
|
dim_heads=embed_dim // num_heads,
|
|
dim_in=dim_in * patch_size,
|
|
dim_out=transformer_dim_out,
|
|
cross_attend = cond_token_dim > 0,
|
|
cond_token_dim = cond_embed_dim,
|
|
global_cond_dim=global_dim,
|
|
modular_local_cond_configs=modular_local_cond_configs,
|
|
**kwargs
|
|
)
|
|
|
|
else:
|
|
raise ValueError(f"Unknown transformer type: {self.transformer_type}")
|
|
|
|
self.preprocess_conv = nn.Conv1d(dim_in, dim_in, 1, bias=False)
|
|
nn.init.zeros_(self.preprocess_conv.weight)
|
|
self.postprocess_conv = nn.Conv1d(io_channels, io_channels, 1, bias=False)
|
|
nn.init.zeros_(self.postprocess_conv.weight)
|
|
|
|
# Fixed logsnr normalization range: maps logsnr to [0, 1] preserving direction (t=0→0, t=1→1)
|
|
_LOGSNR_MIN = -12.0
|
|
_LOGSNR_MAX = 5.0
|
|
_LOGSNR_RANGE = _LOGSNR_MAX - _LOGSNR_MIN
|
|
|
|
def _t_to_logsnr_cond(self, t: torch.Tensor) -> torch.Tensor:
|
|
"""Convert t to normalized logsnr in [0, 1] for timestep conditioning.
|
|
|
|
Maps t through logsnr = log((1-t)/t), clamps to fixed range,
|
|
then normalizes to [0, 1] preserving direction (t=0→0, t=1→1).
|
|
"""
|
|
t_clamped = t.float().clamp(1e-7, 1 - 1e-7)
|
|
logsnr = torch.log((1 - t_clamped) / t_clamped)
|
|
logsnr = logsnr.clamp(self._LOGSNR_MIN, self._LOGSNR_MAX)
|
|
return ((self._LOGSNR_MAX - logsnr) / self._LOGSNR_RANGE).to(t.dtype)
|
|
|
|
def _call_transformer(self, x, *, prepend_inputs=None, cross_attn_cond=None,
|
|
mask=None, prepend_mask=None, return_info=False,
|
|
exit_layer_ix=None, local_add_cond=None,
|
|
modular_local_cond=None, padding_mask=None,
|
|
extra_args=None, **kwargs):
|
|
"""Helper method to call transformer and handle early exit logic."""
|
|
|
|
output = self.transformer(x, prepend_embeds=prepend_inputs, context=cross_attn_cond,
|
|
return_info=return_info, exit_layer_ix=exit_layer_ix,
|
|
local_add_cond=local_add_cond, modular_local_cond=modular_local_cond,
|
|
padding_mask=padding_mask,
|
|
**(extra_args or {}), **kwargs)
|
|
|
|
if return_info:
|
|
output, info = output
|
|
|
|
# Avoid postprocessing on early exit
|
|
if exit_layer_ix is not None:
|
|
if return_info:
|
|
return output, info
|
|
else:
|
|
return output
|
|
|
|
return (output, info) if return_info and 'info' in locals() else output
|
|
|
|
def _forward(
|
|
self,
|
|
x,
|
|
t,
|
|
mask=None,
|
|
cross_attn_cond=None,
|
|
cross_attn_cond_mask=None,
|
|
input_concat_cond=None,
|
|
local_add_cond=None,
|
|
modular_local_cond=None,
|
|
global_embed=None,
|
|
prepend_cond=None,
|
|
prepend_cond_mask=None,
|
|
padding_mask=None,
|
|
return_info=False,
|
|
exit_layer_ix=None,
|
|
**kwargs):
|
|
|
|
if cross_attn_cond is not None:
|
|
cross_attn_cond = self.to_cond_embed(cross_attn_cond)
|
|
|
|
if global_embed is not None:
|
|
# Project the global conditioning to the embedding dimension
|
|
global_embed = self.to_global_embed(global_embed)
|
|
|
|
prepend_inputs = None
|
|
prepend_mask = None
|
|
prepend_length = 0
|
|
if prepend_cond is not None:
|
|
# Project the prepend conditioning to the embedding dimension
|
|
prepend_cond = self.to_prepend_embed(prepend_cond)
|
|
|
|
prepend_inputs = prepend_cond
|
|
if prepend_cond_mask is not None:
|
|
prepend_mask = prepend_cond_mask
|
|
|
|
prepend_length = prepend_cond.shape[1]
|
|
|
|
if input_concat_cond is not None:
|
|
# Interpolate input_concat_cond to the same length as x
|
|
if input_concat_cond.shape[2] != x.shape[2]:
|
|
input_concat_cond = F.interpolate(input_concat_cond, (x.shape[2], ), mode='nearest')
|
|
|
|
x = torch.cat([x, input_concat_cond], dim=1)
|
|
|
|
if local_add_cond is not None:
|
|
local_add_cond = rearrange(local_add_cond, "b c t -> b t c")
|
|
|
|
# Rearrange modular_local_cond tensors
|
|
if modular_local_cond is not None:
|
|
modular_local_cond = {
|
|
k: rearrange(v, "b c t -> b t c")
|
|
for k, v in modular_local_cond.items()
|
|
}
|
|
|
|
# Get the batch of timestep embeddings
|
|
t_cond = self._t_to_logsnr_cond(t) if self.timestep_features_logsnr else t
|
|
# Convert to model dtype for linear layers (t itself is kept in float32 for precision)
|
|
# x has already been converted to model dtype in the outer forward() method
|
|
t_cond = t_cond.to(x.dtype)
|
|
timestep_embed = self.to_timestep_embed(self.timestep_features(t_cond[:, None])) # (b, embed_dim)
|
|
|
|
# Timestep embedding is considered a global embedding. Add to the global conditioning if it exists
|
|
|
|
if self.timestep_cond_type == "global":
|
|
if global_embed is not None:
|
|
global_embed = global_embed + timestep_embed
|
|
else:
|
|
global_embed = timestep_embed
|
|
elif self.timestep_cond_type == "input_concat":
|
|
x = torch.cat([x, timestep_embed.unsqueeze(2).expand(-1, -1, x.shape[2])], dim=1)
|
|
|
|
# Add the global_embed to the prepend inputs if there is no global conditioning support in the transformer
|
|
if self.global_cond_type == "prepend" and global_embed is not None:
|
|
if prepend_inputs is None:
|
|
# Prepend inputs are just the global embed, and the mask is all ones
|
|
prepend_inputs = global_embed.unsqueeze(1)
|
|
prepend_mask = torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)
|
|
else:
|
|
# Prepend inputs are the prepend conditioning + the global embed
|
|
prepend_inputs = torch.cat([prepend_inputs, global_embed.unsqueeze(1)], dim=1)
|
|
prepend_mask = torch.cat([prepend_mask, torch.ones((x.shape[0], 1), device=x.device, dtype=torch.bool)], dim=1)
|
|
|
|
prepend_length = prepend_inputs.shape[1]
|
|
|
|
x = self.preprocess_conv(x) + x
|
|
|
|
x = rearrange(x, "b c t -> b t c")
|
|
|
|
extra_args = {}
|
|
|
|
if self.global_cond_type == "adaLN":
|
|
extra_args["global_cond"] = global_embed
|
|
|
|
if self.patch_size > 1:
|
|
x = rearrange(x, "b (t p) c -> b t (c p)", p=self.patch_size)
|
|
|
|
result = self._call_transformer(
|
|
x,
|
|
prepend_inputs=prepend_inputs,
|
|
cross_attn_cond=cross_attn_cond,
|
|
mask=mask,
|
|
prepend_mask=prepend_mask,
|
|
return_info=return_info,
|
|
exit_layer_ix=exit_layer_ix,
|
|
local_add_cond=local_add_cond,
|
|
modular_local_cond=modular_local_cond,
|
|
padding_mask=padding_mask,
|
|
extra_args=extra_args,
|
|
**kwargs,
|
|
)
|
|
|
|
# Handle early exit (result contains both output and info)
|
|
if exit_layer_ix is not None:
|
|
return result
|
|
|
|
output = result[0] if return_info else result
|
|
if return_info:
|
|
info = result[1]
|
|
|
|
output = rearrange(output, "b t c -> b c t")[:,:,prepend_length:]
|
|
|
|
if self.patch_size > 1:
|
|
output = rearrange(output, "b (c p) t -> b c (t p)", p=self.patch_size)
|
|
|
|
output = self.postprocess_conv(output) + output
|
|
|
|
if return_info:
|
|
return output, info
|
|
|
|
return output
|
|
|
|
def apg_project(self, v0, v1, padding_mask=None):
|
|
"""
|
|
Project v0 into components parallel and orthogonal to v1.
|
|
|
|
Args:
|
|
v0: Tensor to project (B, C, T)
|
|
v1: Reference direction (B, C, T)
|
|
padding_mask: Optional mask (B, T) where True = valid, False = padding.
|
|
If provided, only valid positions contribute to the projection.
|
|
"""
|
|
dtype = v0.dtype
|
|
v0, v1 = v0.float(), v1.float()
|
|
|
|
if padding_mask is not None:
|
|
# Expand mask to match tensor shape: (B, T) -> (B, 1, T)
|
|
mask = padding_mask.unsqueeze(1).float()
|
|
# Zero out padding positions for projection computation
|
|
v0_masked = v0 * mask
|
|
v1_masked = v1 * mask
|
|
# Normalize only over valid positions
|
|
v1_norm = v1_masked.norm(dim=[-1, -2], keepdim=True).clamp(min=1e-8)
|
|
v1_normalized = v1_masked / v1_norm
|
|
# Compute projection using masked values
|
|
v0_parallel = (v0_masked * v1_normalized).sum(dim=[-1, -2], keepdim=True) * v1_normalized
|
|
# Orthogonal component: subtract parallel from original (not masked) v0
|
|
# but apply mask to ensure padding stays zero
|
|
v0_orthogonal = (v0 - (v0 * v1_normalized).sum(dim=[-1, -2], keepdim=True) * v1_normalized) * mask
|
|
else:
|
|
v1 = torch.nn.functional.normalize(v1, dim=[-1, -2])
|
|
v0_parallel = (v0 * v1).sum(dim=[-1, -2], keepdim=True) * v1
|
|
v0_orthogonal = v0 - v0_parallel
|
|
|
|
return v0_parallel.to(dtype), v0_orthogonal.to(dtype)
|
|
|
|
def forward(
|
|
self,
|
|
x,
|
|
t,
|
|
cross_attn_cond=None,
|
|
cross_attn_cond_mask=None,
|
|
negative_cross_attn_cond=None,
|
|
negative_cross_attn_mask=None,
|
|
input_concat_cond=None,
|
|
local_add_cond=None,
|
|
modular_local_cond=None,
|
|
global_embed=None,
|
|
negative_global_embed=None,
|
|
prepend_cond=None,
|
|
prepend_cond_mask=None,
|
|
padding_mask=None,
|
|
cfg_scale=1.0,
|
|
cfg_dropout_prob=0.0,
|
|
cfg_interval = (0, 1),
|
|
lora_interval = (0, 1),
|
|
lora_layer_filter = "",
|
|
lora_configs = None,
|
|
causal=False,
|
|
scale_phi=0.0,
|
|
cfg_norm_threshold=0.0,
|
|
apg_scale=1.0,
|
|
mask=None,
|
|
return_info=False,
|
|
exit_layer_ix=None,
|
|
**kwargs):
|
|
|
|
assert not causal, "Causal mode is not supported for DiffusionTransformer"
|
|
|
|
model_dtype = next(self.parameters()).dtype
|
|
|
|
x = x.to(model_dtype)
|
|
|
|
# Keep t in float32: the logsnr transform log((1-t)/t) amplifies bf16
|
|
# quantization error ~380x near t=1, causing catastrophic conditioning errors.
|
|
# t is a 1D batch-size tensor so float32 has zero memory impact.
|
|
t = t.float()
|
|
|
|
if cross_attn_cond is not None:
|
|
cross_attn_cond = cross_attn_cond.to(model_dtype)
|
|
|
|
if negative_cross_attn_cond is not None:
|
|
negative_cross_attn_cond = negative_cross_attn_cond.to(model_dtype)
|
|
|
|
if input_concat_cond is not None:
|
|
input_concat_cond = input_concat_cond.to(model_dtype)
|
|
|
|
if local_add_cond is not None:
|
|
local_add_cond = local_add_cond.to(model_dtype)
|
|
|
|
if modular_local_cond is not None:
|
|
modular_local_cond = {k: v.to(model_dtype) for k, v in modular_local_cond.items()}
|
|
|
|
if global_embed is not None:
|
|
global_embed = global_embed.to(model_dtype)
|
|
|
|
if negative_global_embed is not None:
|
|
negative_global_embed = negative_global_embed.to(model_dtype)
|
|
|
|
if prepend_cond is not None:
|
|
prepend_cond = prepend_cond.to(model_dtype)
|
|
|
|
if cross_attn_cond_mask is not None:
|
|
cross_attn_cond_mask = cross_attn_cond_mask.bool()
|
|
|
|
cross_attn_cond_mask = None # Temporarily disabling conditioning masks due to kernel issue for flash attention
|
|
|
|
if prepend_cond_mask is not None:
|
|
prepend_cond_mask = prepend_cond_mask.bool()
|
|
|
|
# Early exit bypasses CFG processing
|
|
if exit_layer_ix is not None:
|
|
assert self.transformer_type == "continuous_transformer", "exit_layer_ix is only supported for continuous_transformer"
|
|
return self._forward(
|
|
x,
|
|
t,
|
|
cross_attn_cond=cross_attn_cond,
|
|
cross_attn_cond_mask=cross_attn_cond_mask,
|
|
input_concat_cond=input_concat_cond,
|
|
local_add_cond=local_add_cond,
|
|
modular_local_cond=modular_local_cond,
|
|
global_embed=global_embed,
|
|
prepend_cond=prepend_cond,
|
|
prepend_cond_mask=prepend_cond_mask,
|
|
padding_mask=padding_mask,
|
|
mask=mask,
|
|
return_info=return_info,
|
|
exit_layer_ix=exit_layer_ix,
|
|
**kwargs
|
|
)
|
|
|
|
# CFG dropout
|
|
if cfg_dropout_prob > 0.0 and cfg_scale == 1.0:
|
|
if cross_attn_cond is not None:
|
|
null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
|
|
dropout_mask = torch.bernoulli(torch.full((cross_attn_cond.shape[0], 1, 1), cfg_dropout_prob, device=cross_attn_cond.device)).to(torch.bool)
|
|
cross_attn_cond = torch.where(dropout_mask, null_embed, cross_attn_cond)
|
|
|
|
if prepend_cond is not None:
|
|
null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
|
|
dropout_mask = torch.bernoulli(torch.full((prepend_cond.shape[0], 1, 1), cfg_dropout_prob, device=prepend_cond.device)).to(torch.bool)
|
|
prepend_cond = torch.where(dropout_mask, null_embed, prepend_cond)
|
|
|
|
if self.diffusion_objective == "v":
|
|
sigma = torch.sin(t * math.pi / 2)
|
|
alpha = torch.cos(t * math.pi / 2)
|
|
elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
|
|
sigma = t
|
|
|
|
# LoRA interval
|
|
if has_lora(self):
|
|
if lora_configs is not None:
|
|
# Multi-LoRA: per-LoRA interval and layer filter
|
|
for lora_config in lora_configs:
|
|
idx = lora_config["lora_index"]
|
|
interval = lora_config.get("interval", (0, 1))
|
|
layer_filter = lora_config.get("layer_filter", "")
|
|
if interval[0] <= sigma[0] <= interval[1]:
|
|
enable_lora(self, lora_index=idx)
|
|
filter_lora_layers(self, layer_filter, lora_index=idx)
|
|
else:
|
|
disable_lora(self, lora_index=idx)
|
|
else:
|
|
# Legacy single-LoRA path
|
|
if lora_interval[0] <= sigma[0] <= lora_interval[1]:
|
|
enable_lora(self)
|
|
filter_lora_layers(self, lora_layer_filter)
|
|
else:
|
|
disable_lora(self)
|
|
|
|
if cfg_scale != 1.0 and (cross_attn_cond is not None or prepend_cond is not None) and (cfg_interval[0] <= sigma[0] <= cfg_interval[1]):
|
|
|
|
# Classifier-free guidance
|
|
# Concatenate conditioned and unconditioned inputs on the batch dimension
|
|
batch_inputs = torch.cat([x, x], dim=0)
|
|
batch_timestep = torch.cat([t, t], dim=0)
|
|
|
|
if global_embed is not None:
|
|
batch_global_cond = torch.cat([global_embed, global_embed], dim=0)
|
|
else:
|
|
batch_global_cond = None
|
|
|
|
if input_concat_cond is not None:
|
|
batch_input_concat_cond = torch.cat([input_concat_cond, input_concat_cond], dim=0)
|
|
else:
|
|
batch_input_concat_cond = None
|
|
|
|
if local_add_cond is not None:
|
|
batch_local_add_cond = torch.cat([local_add_cond, local_add_cond], dim=0)
|
|
else:
|
|
batch_local_add_cond = None
|
|
|
|
if modular_local_cond is not None:
|
|
batch_modular_local_cond = {k: torch.cat([v, v], dim=0) for k, v in modular_local_cond.items()}
|
|
else:
|
|
batch_modular_local_cond = None
|
|
|
|
batch_cond = None
|
|
batch_cond_masks = None
|
|
|
|
# Handle CFG for cross-attention conditioning
|
|
if cross_attn_cond is not None:
|
|
|
|
null_embed = torch.zeros_like(cross_attn_cond, device=cross_attn_cond.device)
|
|
|
|
# For negative cross-attention conditioning, replace the null embed with the negative cross-attention conditioning
|
|
if negative_cross_attn_cond is not None:
|
|
|
|
# If there's a negative cross-attention mask, set the masked tokens to the null embed
|
|
if negative_cross_attn_mask is not None:
|
|
negative_cross_attn_mask = negative_cross_attn_mask.to(torch.bool).unsqueeze(2)
|
|
|
|
negative_cross_attn_cond = torch.where(negative_cross_attn_mask, negative_cross_attn_cond, null_embed)
|
|
|
|
batch_cond = torch.cat([cross_attn_cond, negative_cross_attn_cond], dim=0)
|
|
|
|
else:
|
|
batch_cond = torch.cat([cross_attn_cond, null_embed], dim=0)
|
|
|
|
if cross_attn_cond_mask is not None:
|
|
batch_cond_masks = torch.cat([cross_attn_cond_mask, cross_attn_cond_mask], dim=0)
|
|
|
|
batch_prepend_cond = None
|
|
batch_prepend_cond_mask = None
|
|
|
|
if prepend_cond is not None:
|
|
|
|
null_embed = torch.zeros_like(prepend_cond, device=prepend_cond.device)
|
|
|
|
batch_prepend_cond = torch.cat([prepend_cond, null_embed], dim=0)
|
|
|
|
if prepend_cond_mask is not None:
|
|
batch_prepend_cond_mask = torch.cat([prepend_cond_mask, prepend_cond_mask], dim=0)
|
|
|
|
|
|
if mask is not None:
|
|
batch_masks = torch.cat([mask, mask], dim=0)
|
|
else:
|
|
batch_masks = None
|
|
|
|
if padding_mask is not None:
|
|
batch_padding_mask = torch.cat([padding_mask, padding_mask], dim=0)
|
|
else:
|
|
batch_padding_mask = None
|
|
|
|
batch_output = self._forward(
|
|
batch_inputs,
|
|
batch_timestep,
|
|
cross_attn_cond=batch_cond,
|
|
cross_attn_cond_mask=batch_cond_masks,
|
|
mask = batch_masks,
|
|
input_concat_cond = batch_input_concat_cond,
|
|
local_add_cond = batch_local_add_cond,
|
|
modular_local_cond=batch_modular_local_cond,
|
|
global_embed = batch_global_cond,
|
|
prepend_cond = batch_prepend_cond,
|
|
prepend_cond_mask = batch_prepend_cond_mask,
|
|
padding_mask = batch_padding_mask,
|
|
return_info = return_info,
|
|
**kwargs)
|
|
|
|
if return_info:
|
|
batch_output, info = batch_output
|
|
|
|
cond_output, uncond_output = torch.chunk(batch_output, 2, dim=0)
|
|
|
|
if self.diffusion_objective == "v":
|
|
cond_denoised = x * alpha[:, None, None] - cond_output * sigma[:, None, None]
|
|
uncond_denoised = x * alpha[:, None, None] - uncond_output * sigma[:, None, None]
|
|
|
|
elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
|
|
cond_denoised = x - cond_output * sigma[:, None, None]
|
|
uncond_denoised = x - uncond_output * sigma[:, None, None]
|
|
|
|
diff = cond_denoised - uncond_denoised
|
|
|
|
if cfg_norm_threshold > 0:
|
|
if padding_mask is not None:
|
|
# Only compute norm over valid positions
|
|
mask = padding_mask.unsqueeze(1).float() # (B, 1, T)
|
|
diff_masked = diff * mask
|
|
diff_norm = diff_masked.norm(p=2, dim=[-1, -2], keepdim=True)
|
|
else:
|
|
diff_norm = diff.norm(p=2, dim=[-1, -2], keepdim=True)
|
|
scale_factor = torch.minimum(torch.ones_like(diff), cfg_norm_threshold / diff_norm)
|
|
diff *= scale_factor
|
|
|
|
if apg_scale == 0.0:
|
|
# Vanilla CFG: use full diff
|
|
cfg_diff = diff
|
|
elif apg_scale == 1.0:
|
|
# Full APG: use only orthogonal component
|
|
_, diff_orthogonal = self.apg_project(diff, cond_denoised, padding_mask=padding_mask)
|
|
cfg_diff = diff_orthogonal
|
|
else:
|
|
# Blended APG: interpolate between full diff and orthogonal
|
|
diff_parallel, diff_orthogonal = self.apg_project(diff, cond_denoised, padding_mask=padding_mask)
|
|
cfg_diff = apg_scale * diff_orthogonal + (1 - apg_scale) * diff
|
|
|
|
cfg_denoised = cond_denoised + (cfg_scale - 1) * cfg_diff
|
|
|
|
if self.diffusion_objective == "v":
|
|
output = (x * alpha[:, None, None] - cfg_denoised) / sigma[:, None, None]
|
|
elif self.diffusion_objective in ["rectified_flow", "rf_denoiser"]:
|
|
output = (x - cfg_denoised) / sigma[:, None, None]
|
|
|
|
# CFG Rescale
|
|
if scale_phi != 0.0:
|
|
cond_out_std = cond_output.std(dim=1, keepdim=True)
|
|
out_cfg_std = output.std(dim=1, keepdim=True)
|
|
output = scale_phi * (output * (cond_out_std/out_cfg_std)) + (1-scale_phi) * output
|
|
|
|
if return_info:
|
|
info["uncond_output"] = uncond_output
|
|
return output, info
|
|
|
|
return output
|
|
|
|
else:
|
|
return self._forward(
|
|
x,
|
|
t,
|
|
cross_attn_cond=cross_attn_cond,
|
|
cross_attn_cond_mask=cross_attn_cond_mask,
|
|
input_concat_cond=input_concat_cond,
|
|
local_add_cond=local_add_cond,
|
|
modular_local_cond=modular_local_cond,
|
|
global_embed=global_embed,
|
|
prepend_cond=prepend_cond,
|
|
prepend_cond_mask=prepend_cond_mask,
|
|
padding_mask=padding_mask,
|
|
mask=mask,
|
|
return_info=return_info,
|
|
**kwargs
|
|
) |