33 lines
1.1 KiB
Python
33 lines
1.1 KiB
Python
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
from runtime import _scale_audio_velocity_for_core
|
|
|
|
|
|
def test_current_flow_av_core_keeps_inner_audio_velocity_unscaled():
|
|
h3 = SimpleNamespace(time_shift_sigma=lambda *args: torch.tensor(0.5))
|
|
audio_out = torch.tensor([1.0, -2.0], dtype=torch.float16)
|
|
|
|
result = _scale_audio_velocity_for_core(h3, audio_out, torch.tensor(0.5), 12.0, 3.0)
|
|
|
|
assert result is audio_out
|
|
|
|
|
|
def test_legacy_core_scales_inner_audio_velocity_with_schedule_slope():
|
|
calls = []
|
|
|
|
def time_shift_slope(sigma_v, shift_v, shift_a):
|
|
calls.append((sigma_v, shift_v, shift_a))
|
|
return torch.tensor(0.25, dtype=torch.float32)
|
|
|
|
h3 = SimpleNamespace(time_shift_slope=time_shift_slope)
|
|
sigma_v = torch.tensor(0.5)
|
|
audio_out = torch.tensor([4.0, -8.0], dtype=torch.float16)
|
|
|
|
result = _scale_audio_velocity_for_core(h3, audio_out, sigma_v, 12.0, 3.0)
|
|
|
|
assert calls == [(sigma_v, 12.0, 3.0)]
|
|
assert result.dtype == torch.float16
|
|
assert torch.equal(result, torch.tensor([1.0, -2.0], dtype=torch.float16))
|