seeker
50 lines
1.5 KiB
Python
50 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import torch
|
|
|
|
from pi06mini.recap import (
|
|
apply_condition_dropout,
|
|
classifier_free_guidance,
|
|
condition_dropout_mask,
|
|
)
|
|
|
|
|
|
def test_default_condition_dropout_is_statistically_thirty_percent() -> None:
|
|
generator = torch.Generator().manual_seed(11)
|
|
mask = condition_dropout_mask(20_000, generator=generator)
|
|
|
|
assert mask.dtype == torch.bool
|
|
assert abs(float(mask.float().mean()) - 0.3) < 0.015
|
|
|
|
|
|
def test_condition_dropout_replaces_whole_examples() -> None:
|
|
conditions = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
|
dropped, mask = apply_condition_dropout(
|
|
conditions,
|
|
null_condition=-1.0,
|
|
mask=torch.tensor([False, True, False]),
|
|
)
|
|
|
|
assert mask.tolist() == [False, True, False]
|
|
torch.testing.assert_close(dropped[0], conditions[0])
|
|
torch.testing.assert_close(dropped[1], torch.tensor([-1.0, -1.0]))
|
|
torch.testing.assert_close(dropped[2], conditions[2])
|
|
|
|
|
|
def test_cfg_interpolates_and_extrapolates_flow_predictions() -> None:
|
|
conditional = torch.tensor([[2.0, 4.0]])
|
|
unconditional = torch.tensor([[0.0, 2.0]])
|
|
|
|
torch.testing.assert_close(
|
|
classifier_free_guidance(conditional, unconditional, guidance_scale=0.0),
|
|
unconditional,
|
|
)
|
|
torch.testing.assert_close(
|
|
classifier_free_guidance(conditional, unconditional, guidance_scale=1.0),
|
|
conditional,
|
|
)
|
|
torch.testing.assert_close(
|
|
classifier_free_guidance(conditional, unconditional, guidance_scale=1.5),
|
|
torch.tensor([[3.0, 5.0]]),
|
|
)
|