239 lines
7.7 KiB
Python
239 lines
7.7 KiB
Python
import torch
|
|
import pytest
|
|
|
|
from fastvideo.pipelines import TrainingBatch
|
|
from fastvideo.train.methods.rl.common import (
|
|
DiffusionSampler,
|
|
SamplingConfig,
|
|
distributed_k_repeat_indices,
|
|
media_to_video_array,
|
|
validation_caption,
|
|
validation_shard_indices,
|
|
)
|
|
from fastvideo.train.utils.config import load_run_config
|
|
|
|
|
|
class _FakeScheduler:
|
|
|
|
def __init__(self):
|
|
self.num_train_timesteps = 1000
|
|
self.set_timesteps_calls = []
|
|
self.timesteps = torch.empty(0)
|
|
self.sigmas = torch.empty(0)
|
|
self.step_calls = 0
|
|
|
|
def set_timesteps(self, num_inference_steps=None, device=None, timesteps=None, sigmas=None):
|
|
self.set_timesteps_calls.append({
|
|
"num_inference_steps": num_inference_steps,
|
|
"timesteps": timesteps,
|
|
"sigmas": sigmas,
|
|
})
|
|
if timesteps is not None:
|
|
self.timesteps = torch.tensor(timesteps, device=device, dtype=torch.float32)
|
|
else:
|
|
self.timesteps = torch.linspace(1000, 0, int(num_inference_steps), device=device)
|
|
if sigmas is not None:
|
|
self.sigmas = torch.tensor(sigmas, device=device, dtype=torch.float32)
|
|
else:
|
|
self.sigmas = torch.cat([self.timesteps / 1000.0, torch.zeros(1, device=device)])
|
|
|
|
def step(self, model_output, timestep, sample, return_dict=False):
|
|
del timestep
|
|
self.step_calls += 1
|
|
prev = sample + model_output
|
|
return (prev, ) if not return_dict else {"prev_sample": prev}
|
|
|
|
|
|
class _FakeModel:
|
|
|
|
def __init__(self):
|
|
self.noise_scheduler = _FakeScheduler()
|
|
self.add_noise_calls = 0
|
|
self.timestep_shapes = []
|
|
|
|
def predict_noise(self, noisy_latents, timestep, batch, *, conditional, attn_kind):
|
|
del conditional, attn_kind
|
|
self.timestep_shapes.append(tuple(timestep.shape))
|
|
assert batch.timesteps is timestep
|
|
return torch.zeros_like(noisy_latents)
|
|
|
|
def predict_x0(self, noisy_latents, timestep, batch, *, conditional, attn_kind):
|
|
del conditional, attn_kind
|
|
self.timestep_shapes.append(tuple(timestep.shape))
|
|
assert batch.timesteps is timestep
|
|
return noisy_latents
|
|
|
|
def add_noise(self, clean_latents, noise, timestep):
|
|
del timestep
|
|
self.add_noise_calls += 1
|
|
return clean_latents + noise
|
|
|
|
|
|
def _batch():
|
|
batch = TrainingBatch()
|
|
batch.latents = torch.zeros(2, 1, 3, 4, 4)
|
|
return batch
|
|
|
|
|
|
def test_sampler_preserves_latent_dtype():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(num_steps=4))
|
|
batch = _batch()
|
|
batch.latents = batch.latents.to(torch.bfloat16)
|
|
|
|
result = sampler.sample(model, batch, generator=torch.Generator().manual_seed(0))
|
|
|
|
assert result.latents.dtype is torch.bfloat16
|
|
|
|
|
|
def test_sampler_uses_scheduler_generated_timesteps_by_default():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(num_steps=4))
|
|
|
|
result = sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0))
|
|
|
|
assert result.timesteps.tolist() == [1000.0, 666.6666259765625, 333.3333435058594, 0.0]
|
|
|
|
|
|
def test_sampler_honors_explicit_timestep_override():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(num_steps=3, timesteps=[900, 300, 10]))
|
|
|
|
result = sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0))
|
|
|
|
assert result.timesteps.tolist() == [900.0, 300.0, 10.0]
|
|
|
|
|
|
def test_sampler_honors_explicit_timesteps_without_matching_num_steps():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(timesteps=[900, 300, 10]))
|
|
|
|
result = sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0))
|
|
|
|
assert result.timesteps.tolist() == [900.0, 300.0, 10.0]
|
|
assert model.noise_scheduler.set_timesteps_calls == []
|
|
|
|
|
|
def test_sampling_config_rejects_unknown_keys():
|
|
with pytest.raises(ValueError, match="Unsupported method.sampling key"):
|
|
SamplingConfig.from_mapping({"solver": "dpm2"})
|
|
|
|
|
|
def test_sampler_restores_original_batch_timestep_after_sampling():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(num_steps=2))
|
|
batch = _batch()
|
|
original_timesteps = torch.tensor([123.0])
|
|
batch.timesteps = original_timesteps
|
|
|
|
sampler.sample(batch=batch, model=model, generator=torch.Generator().manual_seed(0))
|
|
|
|
assert batch.timesteps is original_timesteps
|
|
|
|
|
|
def test_euler_sampler_does_not_renoise_between_steps():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(num_steps=4))
|
|
|
|
sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0))
|
|
|
|
assert model.add_noise_calls == 0
|
|
assert model.timestep_shapes == [(2, ), (2, ), (2, ), (2, )]
|
|
|
|
|
|
def test_sde_reflow_sampler_renoises_between_steps():
|
|
model = _FakeModel()
|
|
sampler = DiffusionSampler(SamplingConfig(num_steps=4, trajectory="sde_reflow"))
|
|
|
|
sampler.sample(model, _batch(), generator=torch.Generator().manual_seed(0))
|
|
|
|
assert model.add_noise_calls == 3
|
|
|
|
|
|
def test_diffusion_nft_config_uses_rl_sampler_not_dmd_pipeline():
|
|
config_path = "examples/train/configs/rl/wan/diffusion_nft_pick_clip.yaml"
|
|
|
|
cfg = load_run_config(config_path)
|
|
raw_text = open(config_path, encoding="utf-8").read()
|
|
|
|
assert cfg.method["_target_"] == "fastvideo.train.methods.rl.diffusion_nft.DiffusionNFTMethod"
|
|
assert cfg.training.optimizer.learning_rate == 3.0e-5
|
|
assert cfg.training.data.num_latent_t == 1
|
|
assert cfg.training.data.num_frames == 1
|
|
assert "sampling_timesteps" not in raw_text
|
|
assert "WanDMDPipeline" not in raw_text
|
|
assert "solver" not in cfg.method["sampling"]
|
|
assert cfg.method["sampling"]["scheduler"] == "flow_match_euler"
|
|
assert cfg.method["sampling"]["trajectory"] == "ode"
|
|
assert cfg.method["sampling"]["flow_shift"] == "inherit"
|
|
assert "deterministic" not in cfg.method["sampling"]
|
|
assert "noise_level" not in cfg.method["sampling"]
|
|
assert cfg.method["validation"]["every_steps"] == 10
|
|
assert cfg.method["validation"]["num_steps"] == 40
|
|
assert cfg.method["validation"]["num_prompts"] == 16
|
|
assert cfg.method["validation"]["log_samples"] is True
|
|
|
|
|
|
def test_validation_shard_indices_are_stable_and_padded():
|
|
rank0 = validation_shard_indices(5, rank=0, world_size=2)
|
|
rank1 = validation_shard_indices(5, rank=1, world_size=2)
|
|
|
|
assert rank0 == [(0, True), (2, True), (4, True)]
|
|
assert rank1 == [(1, True), (3, True), (0, False)]
|
|
|
|
|
|
def test_distributed_k_repeat_indices_repeats_prompts_globally():
|
|
rank0 = distributed_k_repeat_indices(
|
|
dataset_length=100,
|
|
batch_size=6,
|
|
repeats_per_prompt=24,
|
|
world_size=4,
|
|
rank=0,
|
|
seed=123,
|
|
)
|
|
all_indices = []
|
|
for rank in range(4):
|
|
sample = distributed_k_repeat_indices(
|
|
dataset_length=100,
|
|
batch_size=6,
|
|
repeats_per_prompt=24,
|
|
world_size=4,
|
|
rank=rank,
|
|
seed=123,
|
|
)
|
|
all_indices.extend(sample.local_indices)
|
|
|
|
assert rank0.unique_prompt_count == 1
|
|
assert len(all_indices) == 24
|
|
assert len(set(all_indices)) == 1
|
|
|
|
|
|
def test_validation_caption_puts_rewards_first():
|
|
caption = validation_caption(
|
|
"a small blue cube",
|
|
{
|
|
"avg": 0.75,
|
|
"pickscore": 0.5,
|
|
},
|
|
)
|
|
|
|
assert caption.startswith("avg: 0.7500 | pickscore: 0.5000 | ")
|
|
assert caption.endswith("a small blue cube")
|
|
|
|
|
|
def test_media_to_video_array_treats_frame_as_single_frame_video():
|
|
frame = torch.ones(3, 4, 5)
|
|
|
|
video = media_to_video_array(frame)
|
|
|
|
assert video.shape == (1, 3, 4, 5)
|
|
assert video.dtype.name == "uint8"
|
|
|
|
|
|
def test_media_to_video_array_preserves_video_frames():
|
|
media = torch.ones(3, 2, 4, 5)
|
|
|
|
video = media_to_video_array(media)
|
|
|
|
assert video.shape == (2, 3, 4, 5)
|