93 lines
3.8 KiB
Python
93 lines
3.8 KiB
Python
|
import unittest
|
||
|
|
from unittest.mock import patch
|
||
|
|
|
||
|
|
import torch
|
||
|
|
from torch import nn
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.lora import DynamicLoraMixin
|
||
|
|
from h3_blackwell_runtime.sampler import sample_video_turbo, turbo_sigmas
|
||
|
|
|
||
|
|
|
||
|
|
class _Linear(DynamicLoraMixin, nn.Module):
|
||
|
|
def __init__(self, weight):
|
||
|
|
super().__init__()
|
||
|
|
self.register_buffer("weight", weight)
|
||
|
|
self.in_features = weight.shape[1]
|
||
|
|
self.out_features = weight.shape[0]
|
||
|
|
self._init_dynamic_lora()
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
return self._apply_lora(x, torch.nn.functional.linear(x, self.weight))
|
||
|
|
|
||
|
|
|
||
|
|
class TurboLoraContracts(unittest.TestCase):
|
||
|
|
def test_dynamic_lora_matches_unmerged_formula(self):
|
||
|
|
base_weight = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
||
|
|
down = torch.tensor([[2.0, -1.0]])
|
||
|
|
up = torch.tensor([[3.0], [-2.0]])
|
||
|
|
x = torch.tensor([[0.5, -1.0]])
|
||
|
|
linear = _Linear(base_weight)
|
||
|
|
linear.add_lora("turbo", down, up, alpha=0.5)
|
||
|
|
linear.set_lora("turbo", strength=0.75)
|
||
|
|
expected = torch.nn.functional.linear(x, base_weight)
|
||
|
|
expected += 0.75 * 0.5 * torch.nn.functional.linear(torch.nn.functional.linear(x, down), up)
|
||
|
|
torch.testing.assert_close(linear(x), expected)
|
||
|
|
|
||
|
|
def test_disabled_lora_returns_exact_base_result(self):
|
||
|
|
weight = torch.randn(3, 2)
|
||
|
|
x = torch.randn(4, 2)
|
||
|
|
linear = _Linear(weight)
|
||
|
|
linear.add_lora("turbo", torch.randn(1, 2), torch.randn(3, 1), alpha=1.0)
|
||
|
|
expected = torch.nn.functional.linear(x, weight)
|
||
|
|
self.assertTrue(torch.equal(linear(x), expected))
|
||
|
|
linear.set_lora("turbo", strength=0.0)
|
||
|
|
self.assertTrue(torch.equal(linear(x), expected))
|
||
|
|
|
||
|
|
def test_unknown_lora_is_rejected(self):
|
||
|
|
linear = _Linear(torch.randn(3, 2))
|
||
|
|
with self.assertRaisesRegex(ValueError, "not attached"):
|
||
|
|
linear.set_lora("missing")
|
||
|
|
|
||
|
|
def test_four_step_shift_six_schedule(self):
|
||
|
|
actual = turbo_sigmas(4, 6.0, device="cpu")
|
||
|
|
expected = torch.tensor([1.0, 0.9473684211, 0.8571428571, 0.6666666667, 0.0])
|
||
|
|
torch.testing.assert_close(actual, expected)
|
||
|
|
|
||
|
|
def test_eight_step_shift_twelve_schedule_has_nine_grid_points(self):
|
||
|
|
actual = turbo_sigmas(8, 12.0, device="cpu")
|
||
|
|
self.assertEqual(actual.shape, (9,))
|
||
|
|
self.assertEqual(actual[0].item(), 1.0)
|
||
|
|
self.assertEqual(actual[-1].item(), 0.0)
|
||
|
|
self.assertTrue(bool(torch.all(actual[:-1] > actual[1:])))
|
||
|
|
|
||
|
|
def test_turbo_sampler_uses_positive_dataward_updates_and_independent_clocks(self):
|
||
|
|
video = torch.zeros(1, 1, 1, 1, 1)
|
||
|
|
audio = torch.zeros(1, 1, 1, 1)
|
||
|
|
calls = []
|
||
|
|
|
||
|
|
def packer(text, current_video, current_audio, sigma, model_timesteps, **kwargs):
|
||
|
|
calls.append((current_video.clone(), current_audio.clone(), model_timesteps.clone()))
|
||
|
|
return (None, None, None, None, None, None)
|
||
|
|
|
||
|
|
def model(*args):
|
||
|
|
return torch.ones(1), torch.ones(1)
|
||
|
|
|
||
|
|
with (
|
||
|
|
patch("h3_blackwell_runtime.sampler.unpatchify_video", return_value=torch.full_like(video, 2.0)),
|
||
|
|
patch("h3_blackwell_runtime.sampler._unpack_audio", return_value=torch.full_like(audio, 3.0)),
|
||
|
|
):
|
||
|
|
result_video, result_audio = sample_video_turbo(
|
||
|
|
model, packer, torch.empty(0), video, audio,
|
||
|
|
steps=2, video_shift=6.0, audio_shift=3.0, return_audio=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
torch.testing.assert_close(result_video, torch.full_like(video, 2.0))
|
||
|
|
torch.testing.assert_close(result_audio, torch.full_like(audio, 3.0))
|
||
|
|
torch.testing.assert_close(calls[1][0], torch.full_like(video, 2.0 / 7.0))
|
||
|
|
torch.testing.assert_close(calls[1][1], torch.full_like(audio, 0.75))
|
||
|
|
torch.testing.assert_close(calls[1][2], torch.tensor([1.0 / 7.0, 0.25]))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
unittest.main()
|