File size: 7,300 Bytes
5f0e4a2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 | from __future__ import annotations
from types import SimpleNamespace
import torch
from hyvideo.commons.infer_state import initialize_infer_state
from models import HYWorldPlayPredictorV2, PredictorV2Config
from predictor_data.v2_schema import validate_case_tensors_v2, validate_vision_context_v2
def _init_infer_state() -> None:
initialize_infer_state(
SimpleNamespace(
sage_blocks_range="0-1",
use_sageattn=False,
enable_torch_compile=False,
use_fp8_gemm=False,
quant_type="fp8-per-block",
include_patterns="double_blocks",
use_vae_parallel=False,
)
)
def test_v2_zero_init_ar_context_and_backward() -> None:
_init_infer_state()
cfg = PredictorV2Config(
hidden_size=128,
heads_num=1,
mlp_width_ratio=1.0,
attn_mode="torch",
source_block_ids=(0, 1),
teacher_block_count=2,
latent_frames=1,
latent_height=2,
latent_width=2,
)
model = HYWorldPlayPredictorV2(cfg)
torch.nn.init.normal_(model.final_layer.linear.weight, std=0.01)
model.train()
anchor = torch.randn(1, 4, 128)
context = {
block_id: {
"k_txt": torch.randn(1, 1, 3, 128),
"v_txt": torch.randn(1, 1, 3, 128),
"k_vision": torch.empty(2, 1, 0, 128),
"v_vision": torch.empty(2, 1, 0, 128),
}
for block_id in (0, 1)
}
output = model(
target_model_input=torch.randn(1, 65, 1, 2, 2),
anchor_hidden=anchor,
target_frame_condition=torch.randn(1, 1, 128),
target_viewmats=torch.eye(4).reshape(1, 1, 4, 4),
target_Ks=torch.eye(3).reshape(1, 1, 3, 3),
rope_temporal_size=1,
start_rope_start_idx=0,
context_kv=context,
)
assert torch.equal(output["pred_hidden"], anchor)
assert output["pred_velocity"].shape == (1, 32, 1, 2, 2)
loss = output["pred_velocity"].float().square().mean()
loss.backward()
assert model.residual_out.weight.grad is not None
assert model.img_in.proj.weight.grad is None
assert model.final_layer.linear.weight.grad is None
assert all(
not parameter.requires_grad
for block in model.double_blocks
for name, parameter in block.named_parameters()
if name.startswith("txt_")
)
def test_v2_bf16_context_schema() -> None:
case = {"image_condition_latent": torch.zeros(1, 32, 1, 30, 52, dtype=torch.bfloat16)}
for block_id in (0, 1, 52, 53):
case[f"block_{block_id:02d}_k_txt"] = torch.zeros(
1, 16, 7, 128, dtype=torch.bfloat16
)
case[f"block_{block_id:02d}_v_txt"] = torch.zeros(
1, 16, 7, 128, dtype=torch.bfloat16
)
case["text_valid_mask"] = torch.ones(1, 7, dtype=torch.bool)
validate_case_tensors_v2(case)
empty_context = {
"k_vision": torch.empty(2, 16, 0, 128, dtype=torch.bfloat16),
"v_vision": torch.empty(2, 16, 0, 128, dtype=torch.bfloat16),
"context_valid_mask": torch.empty(1, 0, dtype=torch.bool),
}
assert validate_vision_context_v2(0, empty_context) == 0
def test_v2_padded_masked_batch_matches_compact_single_samples() -> None:
_init_infer_state()
torch.manual_seed(7)
cfg = PredictorV2Config(
hidden_size=128,
heads_num=1,
mlp_width_ratio=1.0,
attn_mode="torch",
source_block_ids=(0, 1),
teacher_block_count=2,
latent_frames=1,
latent_height=2,
latent_width=2,
)
model = HYWorldPlayPredictorV2(cfg).eval()
torch.nn.init.normal_(model.residual_out.weight, std=0.01)
torch.nn.init.normal_(model.residual_out.bias, std=0.01)
samples = []
for sample_id, (text_tokens, context_tokens, rope_size, rope_start) in enumerate(
((2, 0, 1, 0), (3, 2, 2, 1))
):
context = {}
for block_id in (0, 1):
context[block_id] = {
"k_txt": torch.randn(1, 1, text_tokens, 128),
"v_txt": torch.randn(1, 1, text_tokens, 128),
"k_vision": torch.randn(2, 1, context_tokens, 128),
"v_vision": torch.randn(2, 1, context_tokens, 128),
}
samples.append(
{
"target_model_input": torch.randn(1, 65, 1, 2, 2),
"anchor_hidden": torch.randn(1, 4, 128),
"target_frame_condition": torch.randn(1, 1, 128),
"target_viewmats": torch.eye(4).reshape(1, 1, 4, 4),
"target_Ks": torch.eye(3).reshape(1, 1, 3, 3),
"rope_temporal_size": rope_size,
"start_rope_start_idx": rope_start,
"context_kv": context,
}
)
with torch.no_grad():
singles = [model(**sample) for sample in samples]
batched_context = {}
for block_id in (0, 1):
first, second = (sample["context_kv"][block_id] for sample in samples)
def pad_text(tensor: torch.Tensor) -> torch.Tensor:
result = tensor.new_zeros(1, 1, 3, 128)
result[:, :, : tensor.shape[2]] = tensor
return result
def pad_vision(tensor: torch.Tensor) -> torch.Tensor:
result = tensor.new_zeros(2, 1, 2, 128)
result[:, :, : tensor.shape[2]] = tensor
return result
first_vision_k, second_vision_k = pad_vision(first["k_vision"]), pad_vision(
second["k_vision"]
)
first_vision_v, second_vision_v = pad_vision(first["v_vision"]), pad_vision(
second["v_vision"]
)
batched_context[block_id] = {
"k_txt": torch.cat([pad_text(first["k_txt"]), pad_text(second["k_txt"])], 0),
"v_txt": torch.cat([pad_text(first["v_txt"]), pad_text(second["v_txt"])], 0),
"text_valid_mask": torch.tensor([[1, 1, 0], [1, 1, 1]], dtype=torch.bool),
"k_vision": torch.cat(
[first_vision_k[:1], second_vision_k[:1], first_vision_k[1:], second_vision_k[1:]],
0,
),
"v_vision": torch.cat(
[first_vision_v[:1], second_vision_v[:1], first_vision_v[1:], second_vision_v[1:]],
0,
),
"context_valid_mask": torch.tensor([[0, 0], [1, 1]], dtype=torch.bool),
}
with torch.no_grad():
batched = model(
target_model_input=torch.cat([sample["target_model_input"] for sample in samples]),
anchor_hidden=torch.cat([sample["anchor_hidden"] for sample in samples]),
target_frame_condition=torch.cat(
[sample["target_frame_condition"] for sample in samples]
),
target_viewmats=torch.cat([sample["target_viewmats"] for sample in samples]),
target_Ks=torch.cat([sample["target_Ks"] for sample in samples]),
rope_temporal_size=torch.tensor([1, 2]),
start_rope_start_idx=torch.tensor([0, 1]),
context_kv=batched_context,
)
for index, single in enumerate(singles):
torch.testing.assert_close(
batched["pred_hidden"][index : index + 1],
single["pred_hidden"],
rtol=2e-5,
atol=2e-5,
)
|