92 lines
5.4 KiB
Python
92 lines
5.4 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 _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 _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)
|
|
|
|
def __call__(self, text: torch.Tensor, video: torch.Tensor, audio: torch.Tensor, sigma: float) -> tuple[torch.Tensor, torch.Tensor, list[tuple[int, int, int]], tuple[int, int, int], tuple[int, int, int]]:
|
|
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.")
|
|
video_rows = functional.linear(patchify_video(video).float(), self.video_weight, self.video_bias).to(torch.bfloat16)
|
|
audio_rows = functional.linear(pack_audio(audio).float(), self.audio_weight, self.audio_bias).to(torch.bfloat16)
|
|
text_length, audio_length = text_rows.shape[0], audio_rows.shape[0]
|
|
hidden = torch.cat((text_rows, audio_rows, video_rows))
|
|
video_sigma = torch.tensor(float(sigma), device=hidden.device).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 - float(video_sigma), 1 - float(audio_sigma)
|
|
unique_times = sorted({video_time, audio_time})
|
|
row = {value: index for index, value in enumerate(unique_times)}
|
|
video_row, audio_row = row[video_time] * 3, row[audio_time] * 3
|
|
times = torch.tensor(unique_times, device=hidden.device, dtype=torch.float32)
|
|
positions = torch.cat((torch.stack((torch.arange(text_length, dtype=torch.float64), torch.zeros(text_length), torch.zeros(text_length)), dim=-1), _audio_positions(audio.shape[-1], float(text_length), video.shape[-1], video.shape[-2]), _video_positions(video.shape[2], video.shape[-2], video.shape[-1], float(text_length))))
|
|
block_video_segment = (text_length + audio_length, hidden.shape[0], video_row)
|
|
block_audio_segment = (text_length, text_length + audio_length, audio_row + 2)
|
|
final_video_segment = (text_length + audio_length, hidden.shape[0], row[video_time])
|
|
final_audio_segment = (text_length, text_length + audio_length, row[audio_time])
|
|
return hidden, times, [(0, text_length, video_row + 1), block_audio_segment, block_video_segment], positions, final_video_segment, final_audio_segment
|