| """Exercise the actual sampler without loading the 14B weights.""" |
| from types import SimpleNamespace |
| import torch |
| from pipeline.causal_inference import CausalInferencePipeline |
| from utils.scheduler import FlowMatchScheduler |
|
|
|
|
| class RecordingGenerator(torch.nn.Module): |
| def __init__(self): |
| super().__init__() |
| self.model = SimpleNamespace(blocks=[SimpleNamespace(self_attn=SimpleNamespace())], num_heads=1, dim=1, local_attn_size=1) |
| self.attention_type = 'softmax' |
| self.calls = [] |
| self.scheduler = FlowMatchScheduler(shift=5, sigma_min=0, extra_one_step=True) |
| self.scheduler.set_timesteps(1000, training=True) |
|
|
| def get_scheduler(self): |
| return self.scheduler |
|
|
| def forward(self, noisy_image_or_video, timestep, **kwargs): |
| self.calls.append((timestep.clone(), noisy_image_or_video.clone())) |
| prediction = noisy_image_or_video * 0.5 |
| return torch.zeros_like(prediction), prediction |
|
|
|
|
| def make_pipeline(): |
| args = SimpleNamespace(denoising_step_list=[1000, 750, 500, 250], warp_denoising_step=True, |
| independent_first_frame=False, num_frame_per_block=3, context_noise=0, |
| inference_terminal_zero_cleanup=True, inference_train_matched_context_cache=True, |
| inference_aux_noise_seed=12345) |
| generator = RecordingGenerator() |
| pipe = CausalInferencePipeline(args, 'cpu', generator=generator, |
| text_encoder=lambda **kw: {}, vae=SimpleNamespace()) |
| pipe.frame_seq_length = 1 |
| return pipe, generator |
|
|
|
|
| def test_terminal_cleanup_and_noised_cache_are_executed(): |
| pipe, generator = make_pipeline() |
| assert len(pipe.denoising_step_list) == 5 |
| assert pipe.denoising_step_list[-1] == 0 |
| noise = torch.ones(1, 6, 1, 1, 1) |
| result = pipe.inference(noise=noise, text_prompts=['test'], decode=False) |
| |
| assert len(generator.calls) == 12 |
| for block in range(2): |
| calls = generator.calls[block * 6:(block + 1) * 6] |
| assert [float(t[0, 0]) for t, _ in calls] == pipe.denoising_step_list.tolist() + [0.0] |
| clean = result[:, block * 3:(block + 1) * 3] |
| assert torch.equal(clean, calls[4][1] * 0.5) |
| assert not torch.equal(calls[5][1], clean) |
|
|
|
|
|
|
| def test_microbatch_adapter_preserves_five_sample_noise_order(): |
| from utils.inference_microbatch import install_microbatch_adapter |
| records = [] |
| def record(noise, text_prompts, standard_noise_override, aux_noise_override, **kwargs): |
| records.append((noise.clone(), list(standard_noise_override), list(aux_noise_override))) |
| return noise.clone(), noise.clone() |
| pipe = SimpleNamespace(inference=record, num_frame_per_block=3, inference_aux_noise_seed=12345, |
| vae=SimpleNamespace(model=SimpleNamespace(clear_cache=lambda: None))) |
| install_microbatch_adapter(pipe) |
| torch.manual_seed(0) |
| noise = torch.randn(5, 21, 1, 2, 2) |
| state = torch.get_rng_state().clone() |
| videos, latents = pipe.inference(noise, ['test'] * 5) |
| final_state = torch.get_rng_state().clone() |
| torch.set_rng_state(state) |
| main = [torch.randn(15, 1, 2, 2) for _ in range(21)] |
| rng = torch.Generator().manual_seed(12345) |
| aux = [torch.randn(15, 1, 2, 2, generator=rng) for _ in range(14)] |
| assert torch.equal(final_state, torch.get_rng_state()) |
| assert torch.equal(videos, noise) and torch.equal(latents, noise) |
| for i, (initial, actual_main, actual_aux) in enumerate(records): |
| assert torch.equal(initial, noise[i:i+1]) |
| assert all(torch.equal(a, b[i*3:(i+1)*3]) for a, b in zip(actual_main, main)) |
| assert all(torch.equal(a, b[i*3:(i+1)*3]) for a, b in zip(actual_aux, aux)) |
|
|
