File size: 3,749 Bytes
2ea2835 | 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 | """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)
# Each block has five denoising calls and a separate cache write.
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) # Cache noise never overwrites output.
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))
|