73 lines
3.6 KiB
Python
73 lines
3.6 KiB
Python
"""Final H3 curve-AdaLN and video/audio output heads."""
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
from torch import nn
|
|
|
|
from .attention import rms_norm
|
|
from .checkpoint import H3Checkpoint
|
|
|
|
|
|
class H3FinalLayer(nn.Module):
|
|
def __init__(
|
|
self,
|
|
curve_table: torch.Tensor,
|
|
norm_weight: torch.Tensor,
|
|
adaln_weight: torch.Tensor,
|
|
adaln_bias: torch.Tensor,
|
|
video_weight: torch.Tensor,
|
|
video_bias: torch.Tensor,
|
|
audio_weight: torch.Tensor,
|
|
audio_bias: torch.Tensor,
|
|
*,
|
|
hidden_size: int = 5376,
|
|
eps: float = 1e-5,
|
|
):
|
|
super().__init__()
|
|
if adaln_weight.shape != (2 * hidden_size, curve_table.shape[1]):
|
|
raise ValueError("Unexpected final H3 AdaLN projection dimensions.")
|
|
self.hidden_size = hidden_size
|
|
self.eps = eps
|
|
self.register_buffer("curve_table", curve_table.to(torch.float32), persistent=False)
|
|
self.register_buffer("norm_weight", norm_weight, persistent=False)
|
|
self.register_buffer("adaln_weight", adaln_weight.to(torch.float32), persistent=False)
|
|
self.register_buffer("adaln_bias", adaln_bias.to(torch.float32), persistent=False)
|
|
self.register_buffer("video_weight", video_weight.to(torch.float32), persistent=False)
|
|
self.register_buffer("video_bias", video_bias.to(torch.float32), persistent=False)
|
|
self.register_buffer("audio_weight", audio_weight.to(torch.float32), persistent=False)
|
|
self.register_buffer("audio_bias", audio_bias.to(torch.float32), persistent=False)
|
|
|
|
@classmethod
|
|
def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16):
|
|
return cls(
|
|
checkpoint.tensor("adaln_t_table", dtype=torch.float32),
|
|
checkpoint.tensor("final_layer.norm.weight", dtype=output_dtype),
|
|
checkpoint.tensor("final_layer.adaln_proj.linear.weight", dtype=torch.float32),
|
|
checkpoint.tensor("final_layer.adaln_proj.linear.bias", dtype=torch.float32),
|
|
checkpoint.tensor("final_layer.video_out.weight", dtype=torch.float32),
|
|
checkpoint.tensor("final_layer.video_out.bias", dtype=torch.float32),
|
|
checkpoint.tensor("final_layer.audio_out.weight", dtype=torch.float32),
|
|
checkpoint.tensor("final_layer.audio_out.bias", dtype=torch.float32),
|
|
)
|
|
|
|
def forward(
|
|
self,
|
|
hidden: torch.Tensor,
|
|
timesteps: torch.Tensor,
|
|
video_segment: tuple[int, int, int],
|
|
audio_segment: tuple[int, int, int],
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
position = timesteps.float().clamp(0, 1) * (self.curve_table.shape[0] - 1)
|
|
lower = position.floor().long().clamp(max=self.curve_table.shape[0] - 2)
|
|
embedding = torch.lerp(self.curve_table[lower], self.curve_table[lower + 1], (position - lower).unsqueeze(1))
|
|
shift, scale = functional.linear(embedding, self.adaln_weight, self.adaln_bias).chunk(2, dim=-1)
|
|
|
|
video_start, video_stop, video_row = video_segment
|
|
audio_start, audio_stop, audio_row = audio_segment
|
|
normalized = rms_norm(hidden, self.norm_weight, self.eps)
|
|
video_hidden = normalized[video_start:video_stop] * (1 + scale[video_row].to(hidden.dtype)) + shift[video_row].to(hidden.dtype)
|
|
audio_hidden = normalized[audio_start:audio_stop] * (1 + scale[audio_row].to(hidden.dtype)) + shift[audio_row].to(hidden.dtype)
|
|
return (
|
|
functional.linear(video_hidden.float(), self.video_weight, self.video_bias),
|
|
functional.linear(audio_hidden.float(), self.audio_weight, self.audio_bias),
|
|
)
|