217 lines
11 KiB
Python
217 lines
11 KiB
Python
"""Direct prompt-only H3 packed-token construction."""
|
|
|
|
import math
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
|
|
|
|
FRAME_PER_TOKEN = (1, 4, 4, 4, 4)
|
|
FRAME_RESCALE = 5.0 / 3.0
|
|
|
|
|
|
def patchify_video(latent: torch.Tensor) -> torch.Tensor:
|
|
batch, channels, frames, height, width = latent.shape
|
|
if batch != 1:
|
|
raise ValueError("H3 supports batch size one.")
|
|
if height % 2 or width % 2:
|
|
raise ValueError("H3 video latent dimensions must be divisible by two.")
|
|
return latent.reshape(batch, channels, frames, 1, height // 2, 2, width // 2, 2).permute(0, 2, 4, 6, 1, 3, 5, 7).reshape(-1, channels * 4)
|
|
|
|
|
|
def pack_audio(latent: torch.Tensor) -> torch.Tensor:
|
|
return latent[0].permute(1, 2, 0).reshape(-1, latent.shape[1]).transpose(0, 1).contiguous().transpose(0, 1)
|
|
|
|
|
|
def unpatchify_video(rows: torch.Tensor, frames: int, latent_height: int, latent_width: int) -> torch.Tensor:
|
|
height, width = latent_height // 2, latent_width // 2
|
|
x = rows.reshape(1, frames, height, width, 24, 1, 2, 2).permute(0, 4, 1, 5, 2, 6, 3, 7)
|
|
return x.reshape(1, 24, frames, latent_height, latent_width)
|
|
|
|
|
|
def _axis(dim: int, area: float) -> torch.Tensor:
|
|
ratio, count = dim / area, dim // 2
|
|
return (torch.arange(count, dtype=torch.float64) * ratio / count + (1 - ratio) / 2) * 32
|
|
|
|
|
|
def _frame_positions(height: int, width: int) -> torch.Tensor:
|
|
"""(t ignored) area-normalized (h, w) grid of one latent frame's 2x2-patch rows."""
|
|
area = math.sqrt(height * width)
|
|
ys, xs = torch.meshgrid(_axis(height, area), _axis(width, area), indexing="ij")
|
|
return torch.stack((ys.flatten(), xs.flatten()), dim=-1) # [frame_rows, 2]
|
|
|
|
|
|
def _video_positions(frames: int, height: int, width: int, offset: float) -> torch.Tensor:
|
|
area = math.sqrt(height * width)
|
|
ys, xs = torch.meshgrid(_axis(height, area), _axis(width, area), indexing="ij")
|
|
spatial = torch.stack((ys.flatten(), xs.flatten()), dim=-1)
|
|
spans = torch.tensor([FRAME_RESCALE * FRAME_PER_TOKEN[index % 5] for index in range(frames)], dtype=torch.float64)
|
|
times = offset + torch.cat((torch.zeros(1, dtype=torch.float64), spans[:-1].cumsum(0)))
|
|
result = torch.empty(frames, spatial.shape[0], 3, dtype=torch.float64)
|
|
result[:, :, 0], result[:, :, 1:] = times[:, None], spatial[None]
|
|
return result.reshape(-1, 3)
|
|
|
|
|
|
def _cond_positions(frames: int, cond_t: float, height: int, width: int) -> torch.Tensor:
|
|
spatial = _frame_positions(height, width)
|
|
g = torch.empty(spatial.shape[0], 3, dtype=torch.float64)
|
|
g[:, 0] = cond_t
|
|
g[:, 1:] = spatial
|
|
return g
|
|
|
|
|
|
def _video_t_spans(n: int) -> list[float]:
|
|
return [FRAME_RESCALE * FRAME_PER_TOKEN[k % 5] for k in range(n)]
|
|
|
|
|
|
def _audio_positions(steps: int, offset: float, width: int, height: int) -> torch.Tensor:
|
|
area = math.sqrt(height * width)
|
|
x_axis = _axis(width, area)
|
|
result = torch.zeros(steps * 2, 3, dtype=torch.float64)
|
|
result[:, 0] = (offset + torch.arange(steps, dtype=torch.float64)).repeat(2)
|
|
result[:steps, 2], result[steps:, 2] = x_axis[0], x_axis[-1]
|
|
return result
|
|
|
|
|
|
class H3PromptPacker:
|
|
"""Build `[text | audio | video]` tokens for prompt-only H3 T2V."""
|
|
|
|
def __init__(self, checkpoint):
|
|
self.video_weight = checkpoint.tensor("video_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
|
self.video_bias = checkpoint.tensor("video_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
|
|
self.audio_weight = checkpoint.tensor("audio_patch_proj.weight", dtype=torch.bfloat16).to(torch.float32)
|
|
self.audio_bias = checkpoint.tensor("audio_patch_proj.bias", dtype=torch.bfloat16).to(torch.float32)
|
|
self.text_weight = checkpoint.tensor("condition_proj.weight", dtype=torch.bfloat16)
|
|
self.text_bias = checkpoint.tensor("condition_proj.bias", dtype=torch.bfloat16)
|
|
|
|
VISUAL_COND_TIMESTEP = 0.999
|
|
|
|
def __call__(
|
|
self,
|
|
text: torch.Tensor,
|
|
video: torch.Tensor,
|
|
audio: torch.Tensor,
|
|
sigma: float | torch.Tensor,
|
|
model_timesteps: torch.Tensor | None = None,
|
|
*,
|
|
text_token_tags: torch.Tensor | None = None,
|
|
cond_latents: list[torch.Tensor] | None = None,
|
|
cond_frame_indices: list[int] | None = None,
|
|
frame_count: int | None = None,
|
|
seed: int = 0,
|
|
) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], torch.Tensor, tuple[int, int, int], tuple[int, int, int]]:
|
|
"""Build ``[text | cond | audio | video]`` packed rows for (fl2va) H3.
|
|
|
|
``text`` is the refined text span (width 5376 when already refined, 5120
|
|
for raw Qwen states); ``text_token_tags`` is the per-token DiT modality tag
|
|
(1=text, 0=video over vision pads). ``cond_latents`` are normalized keyframe
|
|
latents ``[1,24,1,H/16,W/16]`` and ``cond_frame_indices`` preserves each
|
|
keyframe's resolved first/last pixel index. They are spliced right after
|
|
the text as non-denoised cond rows with their own near-1 timestep. Returns
|
|
``(hidden, times, segments, positions, video_seg, audio_seg)`` where
|
|
``segments`` rows are ``t_row*3 + modality_tag``.
|
|
"""
|
|
if text.shape[-1] == 5120:
|
|
text_rows = functional.linear(text[0].to(self.text_weight.dtype), self.text_weight, self.text_bias).to(torch.bfloat16)
|
|
elif text.shape[-1] == 5376:
|
|
text_rows = text[0].to(torch.bfloat16)
|
|
else:
|
|
raise ValueError("H3 text states must be Qwen 5120-wide or refined 5376-wide.")
|
|
text_length = text_rows.shape[0]
|
|
latent_t, latent_h, latent_w = video.shape[2], video.shape[-2], video.shape[-1]
|
|
frame_rows = (latent_h // 2) * (latent_w // 2)
|
|
|
|
video_rows = functional.linear(patchify_video(video.to(torch.bfloat16)).float(), self.video_weight, self.video_bias).to(torch.bfloat16)
|
|
audio_rows = functional.linear(pack_audio(audio.to(torch.bfloat16)).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16)
|
|
cond_rows = None
|
|
cond_lengths = []
|
|
if cond_latents:
|
|
if cond_frame_indices is None or len(cond_frame_indices) != len(cond_latents):
|
|
raise ValueError("cond_frame_indices must match cond_latents")
|
|
cond_patches = []
|
|
# every cond video restarts the same CPU RNG stream (Comfy _cond_video_rows)
|
|
for idx, z in enumerate(cond_latents):
|
|
r = patchify_video(z.to(torch.float32))
|
|
if self.VISUAL_COND_TIMESTEP < 1.0:
|
|
gen = torch.Generator("cpu").manual_seed(int(seed))
|
|
noise = torch.randn(r.shape, generator=gen, dtype=torch.float32)
|
|
r = self.VISUAL_COND_TIMESTEP * r + (1.0 - self.VISUAL_COND_TIMESTEP) * noise.to(r.device)
|
|
cond_patches.append(r)
|
|
cond_lengths.append(r.shape[0])
|
|
cond_rows = functional.linear(torch.cat(cond_patches, dim=0), self.video_weight.to(torch.float32), self.video_bias.to(torch.float32)).to(torch.bfloat16)
|
|
|
|
if model_timesteps is None:
|
|
video_sigma = torch.as_tensor(sigma, device=text_rows.device, dtype=torch.float32).clamp(min=1e-6)
|
|
base = video_sigma / (12.0 + video_sigma * (1.0 - 12.0))
|
|
audio_sigma = 3.0 * base / (1.0 + (3.0 - 1.0) * base)
|
|
video_time, audio_time = (1.0 - video_sigma).item(), (1.0 - audio_sigma).item()
|
|
else:
|
|
times_override = model_timesteps.to(device=text_rows.device, dtype=torch.float32).flatten()
|
|
unique_override = sorted(times_override.tolist())
|
|
video_time, audio_time = unique_override[0], unique_override[-1]
|
|
|
|
has_vis_cond = cond_rows is not None
|
|
cond_time = max(video_time, self.VISUAL_COND_TIMESTEP)
|
|
unique_times = sorted({video_time, audio_time} | ({cond_time} if has_vis_cond else set()))
|
|
t_row = {value: index for index, value in enumerate(unique_times)}
|
|
times = torch.tensor(unique_times, device=text_rows.device, dtype=torch.float32)
|
|
|
|
# Assemble hidden in segment order: text | cond | audio | video.
|
|
parts = [text_rows]
|
|
offsets = [0]
|
|
for block in ((cond_rows, "cond"), (audio_rows, "audio"), (video_rows, "video")):
|
|
if block[0] is None:
|
|
continue
|
|
parts.append(block[0])
|
|
offsets.append(offsets[-1] + block[0].shape[0])
|
|
hidden = torch.cat(parts, dim=0)
|
|
|
|
audio_length = audio_rows.shape[0]
|
|
cond_length = cond_rows.shape[0] if cond_rows is not None else 0
|
|
|
|
# Positions: text rows, cond (first/last t anchors), audio, video.
|
|
text_positions = torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1)
|
|
position_blocks = [text_positions]
|
|
if cond_rows is not None and cond_latents:
|
|
spans = _video_t_spans(latent_t)
|
|
cond_t_values = []
|
|
for pixel_index in cond_frame_indices:
|
|
if pixel_index == 0:
|
|
cond_t_values.append(float(text_length))
|
|
elif frame_count is not None and pixel_index == frame_count - 1:
|
|
cond_t_values.append(float(text_length) + sum(spans) - FRAME_RESCALE)
|
|
else:
|
|
raise ValueError("only first/last keyframe anchors are supported")
|
|
position_blocks.append(torch.cat([_cond_positions(frame_rows, cond_t, latent_h, latent_w) for cond_t in cond_t_values], dim=0))
|
|
position_blocks.append(_audio_positions(audio.shape[-1], float(text_length), latent_w, latent_h))
|
|
position_blocks.append(_video_positions(latent_t, latent_h, latent_w, float(text_length)))
|
|
positions = torch.cat(position_blocks, dim=0)
|
|
|
|
# mod_segments: (start, stop, t_row*3 + tag).
|
|
segments: list[tuple[int, int, int]] = []
|
|
if text_token_tags is not None:
|
|
tags = text_token_tags.view(-1).tolist()
|
|
row_base = t_row[video_time] * 3
|
|
run_start = 0
|
|
for i in range(1, text_length + 1):
|
|
if i == text_length or tags[i] != tags[run_start]:
|
|
segments.append((run_start, i, row_base + int(tags[run_start])))
|
|
run_start = i
|
|
else:
|
|
segments.append((0, text_length, t_row[video_time] * 3 + 1))
|
|
|
|
cursor_start = text_length
|
|
if cond_rows is not None:
|
|
for length in cond_lengths:
|
|
segments.append((cursor_start, cursor_start + length, t_row[cond_time] * 3 + 0))
|
|
cursor_start += length
|
|
segments.append((cursor_start, cursor_start + audio_length, t_row[audio_time] * 3 + 2))
|
|
cursor_start += audio_length
|
|
video_start = cursor_start
|
|
video_stop = video_start + video_rows.shape[0]
|
|
segments.append((video_start, video_stop, t_row[video_time] * 3 + 0))
|
|
|
|
video_segment = (video_start, video_stop, t_row[video_time])
|
|
audio_start = text_length + cond_length
|
|
audio_segment = (audio_start, audio_start + audio_length, t_row[audio_time])
|
|
return hidden, times, segments, positions, video_segment, audio_segment
|