39 lines
2 KiB
Python
39 lines
2 KiB
Python
"""Curve-form AdaLN used by the pruned H3 NVFP4 checkpoint."""
|
|
|
|
import torch
|
|
import torch.nn.functional as functional
|
|
from torch import nn
|
|
|
|
from .checkpoint import H3Checkpoint
|
|
|
|
|
|
class H3CurveAdaLN(nn.Module):
|
|
"""Interpolate the H3 timestep curve and emit six modality-specific AdaLN tensors."""
|
|
|
|
def __init__(self, curve_table: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor, *, hidden_size: int = 5376):
|
|
super().__init__()
|
|
if curve_table.shape != (1025, 8):
|
|
raise ValueError(f"Unexpected H3 AdaLN curve shape: {tuple(curve_table.shape)}")
|
|
if weight.shape != (18 * hidden_size, curve_table.shape[1]) or bias.shape != (18 * hidden_size,):
|
|
raise ValueError("Unexpected H3 AdaLN projection dimensions.")
|
|
self.hidden_size = hidden_size
|
|
self.register_buffer("curve_table", curve_table.to(torch.float32), persistent=False)
|
|
self.register_buffer("weight", weight.to(torch.float32), persistent=False)
|
|
self.register_buffer("bias", bias.to(torch.float32), persistent=False)
|
|
|
|
@classmethod
|
|
def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str):
|
|
return cls(
|
|
checkpoint.tensor("adaln_t_table", dtype=torch.float32),
|
|
checkpoint.tensor(f"{prefix}.linear.weight", dtype=torch.bfloat16),
|
|
checkpoint.tensor(f"{prefix}.linear.bias", dtype=torch.bfloat16),
|
|
)
|
|
|
|
def forward(self, timesteps: torch.Tensor) -> tuple[torch.Tensor, ...]:
|
|
"""Return shift/scale/gate tensors ordered as MSA then MLP for all modalities."""
|
|
table = self.curve_table
|
|
position = timesteps.float().clamp(0, 1) * (table.shape[0] - 1)
|
|
lower = position.floor().long().clamp(max=table.shape[0] - 2)
|
|
embedding = torch.lerp(table[lower], table[lower + 1], (position - lower).unsqueeze(1))
|
|
values = functional.linear(embedding, self.weight, self.bias)
|
|
return values.view(values.shape[0] * 3, 6 * self.hidden_size).chunk(6, dim=-1)
|