42 lines
1.7 KiB
Python
42 lines
1.7 KiB
Python
|
"""Direct 50-block H3 denoiser backbone with standalone Sage3 attention."""
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
from .adaln import H3CurveAdaLN
|
||
|
|
from .block import H3DiTBlock
|
||
|
|
from .checkpoint import H3Checkpoint
|
||
|
|
from .rope import h3_rope_rotation
|
||
|
|
|
||
|
|
|
||
|
|
class H3DenoiserBackbone(nn.Module):
|
||
|
|
"""Execute H3 transformer blocks over an already packed Ref2VA hidden sequence."""
|
||
|
|
|
||
|
|
def __init__(self, blocks: list[H3DiTBlock], adaln: list[H3CurveAdaLN], inv_freq: torch.Tensor):
|
||
|
|
super().__init__()
|
||
|
|
if len(blocks) != 50 or len(adaln) != 50:
|
||
|
|
raise ValueError("The released H3 denoiser has exactly 50 transformer blocks.")
|
||
|
|
self.blocks = nn.ModuleList(blocks)
|
||
|
|
self.adaln = nn.ModuleList(adaln)
|
||
|
|
self.register_buffer("inv_freq", inv_freq.to(torch.float32), persistent=False)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"):
|
||
|
|
return cls(
|
||
|
|
[H3DiTBlock.from_checkpoint(checkpoint, index, output_dtype=output_dtype, attention_backend=attention_backend) for index in range(50)],
|
||
|
|
[H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{index}.adaln_proj") for index in range(50)],
|
||
|
|
checkpoint.tensor("rope.inv_freq", dtype=torch.float32),
|
||
|
|
)
|
||
|
|
|
||
|
|
def forward(
|
||
|
|
self,
|
||
|
|
hidden: torch.Tensor,
|
||
|
|
timesteps: torch.Tensor,
|
||
|
|
position_ids: torch.Tensor,
|
||
|
|
segments: list[tuple[int, int, int]],
|
||
|
|
) -> torch.Tensor:
|
||
|
|
rotation = h3_rope_rotation(position_ids.to(hidden.device), self.inv_freq, hidden.dtype)
|
||
|
|
for block, adaln in zip(self.blocks, self.adaln, strict=True):
|
||
|
|
hidden = block(hidden, rotation, *adaln(timesteps), segments)
|
||
|
|
return hidden
|