nodes/hao-ai-lab-FastVideo
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[perf]: hoist invariants out of MagiHuman denoising/latent-prep

Browse Source
This commit is contained in:
SolitaryThinker
2026-05-05 10:55:49 -07:00
parent 14c142ba9c
commit 6d0ddb44bc
4 changed files with 207 additions and 69 deletions
Download Patch File Download Diff File Expand all files Collapse all files
fastvideo/pipelines/basic/magi_human/magi_human_pipeline.py
+3
View File
@@ -177,6 +177,7 @@ class MagiHumanPipeline(ComposedPipelineBase):
t5_gemma_target_length=pc.t5_gemma_target_length,
coords_style=pc.coords_style,
text_offset=pc.text_offset,
audio_in_channels=dit_arch.audio_in_channels,
),
)
@@ -192,6 +193,8 @@ class MagiHumanPipeline(ComposedPipelineBase):
audio_txt_guidance_scale=pc.audio_txt_guidance_scale,
cfg_number=pc.cfg_number,
coords_style=pc.coords_style,
video_guidance_high_t_threshold=pc.video_guidance_high_t_threshold,
video_guidance_low_t_value=pc.video_guidance_low_t_value,
),
)
fastvideo/pipelines/basic/magi_human/pipeline_configs.py
+6
View File
@@ -93,6 +93,12 @@ class MagiHumanBaseConfig(PipelineConfig):
ref_audio_offset: int = 1000
text_offset: int = 0
# Video CFG step-dependent guidance: low-t steps use a relaxed scale.
# Upstream daVinci-MagiHuman/inference/pipeline/video_generate.py:426
# uses 5.0 for high-t and 2.0 for low-t with cutoff at t=500.
video_guidance_high_t_threshold: int = 500
video_guidance_low_t_value: float = 2.0
def __post_init__(self) -> None:
# Base text-to-AV does not need the VAE encoder (no reference-image
# conditioning). Keep decoder only to save memory.
fastvideo/pipelines/basic/magi_human/stages/denoising.py
+38 -22
View File
@@ -18,7 +18,9 @@ from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
from fastvideo.pipelines.stages.base import PipelineStage
from fastvideo.pipelines.stages.validators import VerificationResult
from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
build_packed_inputs,
StaticPackedInputs,
assemble_packed_inputs,
build_static_packed_inputs,
unpack_tokens,
)
@@ -26,29 +28,22 @@ from fastvideo.pipelines.basic.magi_human.stages.latent_preparation import (
def _dit_forward(
dit,
video_latent: torch.Tensor,
audio_latent: torch.Tensor,
audio_feat_len: int,
txt_feat: torch.Tensor,
txt_feat_len: int,
patch_size: tuple[int, int, int],
static_packed: StaticPackedInputs,
coords_style: str,
video_in_channels: int,
audio_in_channels: int,
patch_size: tuple[int, int, int],
) -> tuple[torch.Tensor, torch.Tensor]:
"""One DiT call: builds the packed token stream, runs the DiT, unpacks
the video/audio outputs back into their natural shapes.
"""
x, coords, mm = build_packed_inputs(
video_latent=video_latent,
audio_latent=audio_latent,
audio_feat_len=audio_feat_len,
x, coords, mm = assemble_packed_inputs(
static=static_packed,
txt_feat=txt_feat,
txt_feat_len=txt_feat_len,
patch_size=patch_size,
coords_style=coords_style,
)
video_token_num = x.size(0) - audio_feat_len - txt_feat_len
# DiT is bf16 internally; keep inputs fp32 (embedders dtype).
video_token_num = static_packed.video_token_num
out = dit(x, coords, mm)
return unpack_tokens(
out,
@@ -75,6 +70,8 @@ class MagiHumanDenoisingStage(PipelineStage):
audio_txt_guidance_scale: float = 5.0,
cfg_number: int = 2,
coords_style: str = "v2",
video_guidance_high_t_threshold: int = 500,
video_guidance_low_t_value: float = 2.0,
) -> None:
super().__init__()
self.transformer = transformer
@@ -86,6 +83,8 @@ class MagiHumanDenoisingStage(PipelineStage):
self.audio_txt_guidance_scale = audio_txt_guidance_scale
self.cfg_number = cfg_number
self.coords_style = coords_style
self.video_guidance_high_t_threshold = video_guidance_high_t_threshold
self.video_guidance_low_t_value = video_guidance_low_t_value
def verify_input(self, batch, fastvideo_args):
return VerificationResult()
@@ -96,7 +95,7 @@ class MagiHumanDenoisingStage(PipelineStage):
@torch.inference_mode()
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
device = batch.latents.device
shift = fastvideo_args.pipeline_config.flow_shift or 5.0
shift = fastvideo_args.pipeline_config.flow_shift
# Video and audio use independent FlowUniPC state (upstream
# inference/pipeline/video_generate.py:404-407 instantiates two
# separate schedulers). Sharing one scheduler causes the
@@ -137,38 +136,47 @@ class MagiHumanDenoisingStage(PipelineStage):
audio_feat_len = int(audio_latent.shape[1])
# Precompute step-invariant packed inputs (video+audio tokens, coords,
# modality map). Text varies per cond/uncond call so it is assembled
# inside _dit_forward via assemble_packed_inputs.
static_packed = build_static_packed_inputs(
video_latent=video_latent,
audio_latent=audio_latent,
audio_feat_len=audio_feat_len,
patch_size=self.patch_size,
coords_style=self.coords_style,
)
disable_tqdm = not getattr(fastvideo_args, "log_level_progress", True)
for idx, t in enumerate(tqdm(timesteps, disable=disable_tqdm)):
# Conditional pass
v_cond_video, v_cond_audio = _dit_forward(
self.transformer,
video_latent=video_latent,
audio_latent=audio_latent,
audio_feat_len=audio_feat_len,
txt_feat=txt_feat,
txt_feat_len=txt_feat_len,
patch_size=self.patch_size,
static_packed=static_packed,
coords_style=self.coords_style,
video_in_channels=self.video_in_channels,
audio_in_channels=self.audio_in_channels,
patch_size=self.patch_size,
)
if self.cfg_number == 2:
v_uncond_video, v_uncond_audio = _dit_forward(
self.transformer,
video_latent=video_latent,
audio_latent=audio_latent,
audio_feat_len=audio_feat_len,
txt_feat=neg_txt_feat,
txt_feat_len=neg_txt_feat_len,
patch_size=self.patch_size,
static_packed=static_packed,
coords_style=self.coords_style,
video_in_channels=self.video_in_channels,
audio_in_channels=self.audio_in_channels,
patch_size=self.patch_size,
)
# Reference trick: drop the video guidance to 2.0 for
# timesteps <= 500 in the non-SR path.
video_guidance = (self.video_txt_guidance_scale if t > 500 else 2.0)
video_guidance = (self.video_txt_guidance_scale
if t > self.video_guidance_high_t_threshold else self.video_guidance_low_t_value)
v_video = v_uncond_video + video_guidance * (v_cond_video - v_uncond_video)
v_audio = v_uncond_audio + self.audio_txt_guidance_scale * (v_cond_audio - v_uncond_audio)
else:
@@ -189,6 +197,14 @@ class MagiHumanDenoisingStage(PipelineStage):
return_dict=False,
)[0]
static_packed = build_static_packed_inputs(
video_latent=video_latent,
audio_latent=audio_latent,
audio_feat_len=audio_feat_len,
patch_size=self.patch_size,
coords_style=self.coords_style,
)
batch.latents = video_latent
batch.audio_latents = audio_latent
return batch
fastvideo/pipelines/basic/magi_human/stages/latent_preparation.py
+160 -47
View File
@@ -30,6 +30,15 @@ MODALITY_VIDEO = 0
MODALITY_AUDIO = 1
MODALITY_TEXT = 2
# Audio temporal compression ratio: 1 audio frame → 1/4 latent frame.
# Mirrors data_proxy.py:206 `(audio_feat_len - 1) // 4 + 1` where 4 is
# the audio VAE's temporal stride (same as vae_stride[0] for video).
_AUDIO_TEMPORAL_COMPRESSION = 4
# v1 text-coord reference shape: (T=2, H=1, W=1).
# Mirrors data_proxy.py:202 `ref_feat_shape=(2, 1, 1)` for coords_style=="v1".
_V1_TEXT_REF_SHAPE: tuple[int, int, int] = (2, 1, 1)
def _build_coords(
shape: tuple[int, int, int],
@@ -93,6 +102,7 @@ class MagiHumanLatentPreparationStage(PipelineStage):
t5_gemma_target_length: int = 640,
coords_style: Literal["v1", "v2"] = "v2",
text_offset: int = 0,
audio_in_channels: int = 64,
) -> None:
super().__init__()
self.vae_stride = vae_stride
@@ -102,6 +112,7 @@ class MagiHumanLatentPreparationStage(PipelineStage):
self.t5_gemma_target_length = t5_gemma_target_length
self.coords_style = coords_style
self.text_offset = text_offset
self.audio_in_channels = audio_in_channels
def verify_input(self, batch, fastvideo_args):
return VerificationResult()
@@ -139,9 +150,9 @@ class MagiHumanLatentPreparationStage(PipelineStage):
device=device,
dtype=torch.float32,
)
# Audio latent: [1, num_frames, 64]
# Audio latent: [1, num_frames, audio_in_channels]
audio_latent = torch.randn(
(1, num_frames, 64),
(1, num_frames, self.audio_in_channels),
generator=generator,
device=device,
dtype=torch.float32,
@@ -190,90 +201,192 @@ class MagiHumanLatentPreparationStage(PipelineStage):
return batch
def build_packed_inputs(
video_latent: torch.Tensor, # [1, z_dim, T, H, W]
audio_latent: torch.Tensor, # [1, num_frames, 64]
class StaticPackedInputs:
"""Step-invariant packed inputs: video+audio tokens, coords, modality map.
Computed once before the denoise loop; reused for every cond/uncond call.
Text tokens are NOT included here because cond/uncond have different lengths.
"""
__slots__ = (
"video_tokens",
"audio_tokens",
"video_coords",
"audio_coords",
"video_mm",
"audio_mm",
"video_token_num",
"audio_feat_len",
"max_ch",
)
def __init__(
self,
video_tokens: torch.Tensor,
audio_tokens: torch.Tensor,
video_coords: torch.Tensor,
audio_coords: torch.Tensor,
video_mm: torch.Tensor,
audio_mm: torch.Tensor,
max_ch: int,
) -> None:
self.video_tokens = video_tokens
self.audio_tokens = audio_tokens
self.video_coords = video_coords
self.audio_coords = audio_coords
self.video_mm = video_mm
self.audio_mm = audio_mm
self.video_token_num = video_tokens.size(0)
self.audio_feat_len = audio_tokens.size(0)
self.max_ch = max_ch
def build_static_packed_inputs(
video_latent: torch.Tensor,
audio_latent: torch.Tensor,
audio_feat_len: int,
txt_feat: torch.Tensor, # [1, target_length, 3584]
txt_feat_len: int,
patch_size: tuple[int, int, int],
coords_style: Literal["v1", "v2"] = "v2",
text_offset: int = 0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build the flat (token_sequence, coords_mapping, modality_mapping) tuple
the DiT consumes. Mirrors SingleData in inference/pipeline/data_proxy.py.
) -> StaticPackedInputs:
"""Build the step-invariant portion of the packed token stream.
Assumes batch size 1 (the reference also hard-assumes this for its
local-attention handler).
Returns video+audio tokens (padded to a common channel width), their
coords, and their modality slices. Text is excluded because cond/uncond
differ in length; call assemble_packed_inputs to attach text per call.
Mirrors SingleData.token_sequence / coords_mapping / modality_mapping in
inference/pipeline/data_proxy.py, minus the text portion.
"""
pT, pH, pW = patch_size
_, z_dim, T, H, W = video_latent.shape
assert video_latent.size(0) == 1, "batch size 1 required for MagiHuman base"
# Tokenize video: [B, L_v, pT*pH*pW*z_dim]
video_tokens = _img2tokens(video_latent, t_patch=pT, patch=pH)[0]
audio_tokens = audio_latent[0, :audio_feat_len].contiguous()
text_tokens = txt_feat[0, :txt_feat_len].contiguous()
max_ch = max(video_tokens.size(-1), audio_tokens.size(-1), text_tokens.size(-1))
max_ch = max(video_tokens.size(-1), audio_tokens.size(-1))
video_tokens = F.pad(video_tokens, (0, max_ch - video_tokens.size(-1)))
audio_tokens = F.pad(audio_tokens, (0, max_ch - audio_tokens.size(-1)))
text_tokens = F.pad(text_tokens, (0, max_ch - text_tokens.size(-1)))
token_seq = torch.cat([video_tokens, audio_tokens, text_tokens], dim=0)
device = video_tokens.device
dtype = video_tokens.dtype
video_token_num = video_tokens.size(0)
# Modality map.
device = token_seq.device
mm = torch.cat([
torch.full((video_token_num, ), MODALITY_VIDEO, dtype=torch.int64, device=device),
torch.full((audio_feat_len, ), MODALITY_AUDIO, dtype=torch.int64, device=device),
torch.full((txt_feat_len, ), MODALITY_TEXT, dtype=torch.int64, device=device),
],
dim=0)
# Coords: build per-modality then concat in (video, audio, text) order.
video_mm = torch.full((video_token_num, ), MODALITY_VIDEO, dtype=torch.int64, device=device)
audio_mm = torch.full((audio_feat_len, ), MODALITY_AUDIO, dtype=torch.int64, device=device)
video_ref_shape = (T // pT, H // pH, W // pW)
video_coords = _build_coords(
shape=(T // pT, H // pH, W // pW),
shape=video_ref_shape,
ref_feat_shape=video_ref_shape,
device=device,
dtype=token_seq.dtype,
dtype=dtype,
)
if coords_style == "v2":
audio_ref_t = (audio_feat_len - 1) // 4 + 1
audio_ref_t = (audio_feat_len - 1) // _AUDIO_TEMPORAL_COMPRESSION + 1
audio_coords = _build_coords(
shape=(audio_feat_len, 1, 1),
ref_feat_shape=(audio_ref_t // pT, 1, 1),
device=device,
dtype=token_seq.dtype,
)
text_coords = _build_coords(
shape=(txt_feat_len, 1, 1),
ref_feat_shape=(1, 1, 1),
offset_thw=(-txt_feat_len, 0, 0),
device=device,
dtype=token_seq.dtype,
dtype=dtype,
)
else:
audio_coords = _build_coords(
shape=(audio_feat_len, 1, 1),
ref_feat_shape=(T // pT, 1, 1),
device=device,
dtype=token_seq.dtype,
)
text_coords = _build_coords(
shape=(txt_feat_len, 1, 1),
ref_feat_shape=(2, 1, 1),
offset_thw=(text_offset, 0, 0),
device=device,
dtype=token_seq.dtype,
dtype=dtype,
)
coords = torch.cat([video_coords, audio_coords, text_coords], dim=0)
return StaticPackedInputs(
video_tokens=video_tokens,
audio_tokens=audio_tokens,
video_coords=video_coords,
audio_coords=audio_coords,
video_mm=video_mm,
audio_mm=audio_mm,
max_ch=max_ch,
)
def assemble_packed_inputs(
static: StaticPackedInputs,
txt_feat: torch.Tensor,
txt_feat_len: int,
coords_style: Literal["v1", "v2"] = "v2",
text_offset: int = 0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Attach per-call text tokens to the precomputed static packed inputs.
Returns (token_seq, coords, modality_map) ready for the DiT.
"""
text_tokens = txt_feat[0, :txt_feat_len].contiguous()
max_ch = max(static.max_ch, text_tokens.size(-1))
video_tokens = F.pad(static.video_tokens, (0, max_ch - static.video_tokens.size(-1)))
audio_tokens = F.pad(static.audio_tokens, (0, max_ch - static.audio_tokens.size(-1)))
text_tokens = F.pad(text_tokens, (0, max_ch - text_tokens.size(-1)))
token_seq = torch.cat([video_tokens, audio_tokens, text_tokens], dim=0)
device = token_seq.device
dtype = token_seq.dtype
text_mm = torch.full((txt_feat_len, ), MODALITY_TEXT, dtype=torch.int64, device=device)
mm = torch.cat([static.video_mm, static.audio_mm, text_mm], dim=0)
if coords_style == "v2":
text_coords = _build_coords(
shape=(txt_feat_len, 1, 1),
ref_feat_shape=(1, 1, 1),
offset_thw=(-txt_feat_len, 0, 0),
device=device,
dtype=dtype,
)
else:
text_coords = _build_coords(
shape=(txt_feat_len, 1, 1),
ref_feat_shape=_V1_TEXT_REF_SHAPE,
offset_thw=(text_offset, 0, 0),
device=device,
dtype=dtype,
)
coords = torch.cat([static.video_coords, static.audio_coords, text_coords], dim=0)
return token_seq, coords, mm
def build_packed_inputs(
video_latent: torch.Tensor,
audio_latent: torch.Tensor,
audio_feat_len: int,
txt_feat: torch.Tensor,
txt_feat_len: int,
patch_size: tuple[int, int, int],
coords_style: Literal["v1", "v2"] = "v2",
text_offset: int = 0,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Build the full packed token stream in one call (backwards-compat wrapper).
Equivalent to assemble_packed_inputs(build_static_packed_inputs(...), ...).
Prefer calling the two helpers separately when the static portion can be
reused across multiple calls (e.g. cond/uncond in the denoise loop).
"""
static = build_static_packed_inputs(
video_latent=video_latent,
audio_latent=audio_latent,
audio_feat_len=audio_feat_len,
patch_size=patch_size,
coords_style=coords_style,
)
return assemble_packed_inputs(
static=static,
txt_feat=txt_feat,
txt_feat_len=txt_feat_len,
coords_style=coords_style,
text_offset=text_offset,
)
def unpack_tokens(
output: torch.Tensor, # [L, max(V_ch, A_ch)]
video_token_num: int,