chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,51 @@
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from utils.i2v_conditioning import (
|
||||
_get_i2v_context_frames,
|
||||
_i2v_loss_mask_like,
|
||||
_overwrite_i2v_context,
|
||||
)
|
||||
|
||||
|
||||
class DMDI2VConditioningTest(unittest.TestCase):
|
||||
def test_overwrite_i2v_context_keeps_only_initial_frames_clean(self):
|
||||
tensor = torch.randn(2, 4, 3, 2, 2)
|
||||
original = tensor.clone()
|
||||
initial_latent = torch.full((2, 1, 3, 2, 2), 7.0)
|
||||
|
||||
context_frames = _get_i2v_context_frames(tensor, initial_latent)
|
||||
conditioned = _overwrite_i2v_context(tensor, initial_latent, context_frames)
|
||||
|
||||
self.assertEqual(context_frames, 1)
|
||||
self.assertTrue(torch.equal(conditioned[:, :1], initial_latent))
|
||||
self.assertTrue(torch.equal(conditioned[:, 1:], original[:, 1:]))
|
||||
|
||||
def test_i2v_loss_mask_excludes_context_frames(self):
|
||||
tensor = torch.randn(2, 4, 3, 2, 2)
|
||||
mask = _i2v_loss_mask_like(tensor, context_frames=1)
|
||||
|
||||
self.assertFalse(mask[:, :1].any())
|
||||
self.assertTrue(mask[:, 1:].all())
|
||||
|
||||
def test_overwrite_i2v_context_keeps_chunk_length_unchanged(self):
|
||||
initial_latent = torch.full((2, 1, 3, 2, 2), 5.0)
|
||||
noise_chunk = torch.randn(2, 4, 3, 2, 2)
|
||||
|
||||
first_chunk = _overwrite_i2v_context(noise_chunk, initial_latent, context_frames=1)
|
||||
|
||||
self.assertEqual(first_chunk.shape[1], 4)
|
||||
self.assertTrue(torch.equal(first_chunk[:, :1], initial_latent))
|
||||
self.assertTrue(torch.equal(first_chunk[:, 1:], noise_chunk[:, 1:]))
|
||||
|
||||
def test_i2v_context_must_not_cover_entire_clip(self):
|
||||
tensor = torch.randn(2, 1, 3, 2, 2)
|
||||
initial_latent = torch.randn(2, 1, 3, 2, 2)
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
_get_i2v_context_frames(tensor, initial_latent)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,29 @@
|
||||
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from utils.fp8 import quantize_model_fp8
|
||||
from utils.torch_compile_utils import SafeCompiledCallable
|
||||
|
||||
|
||||
def test_fp8_quantization_requires_cuda(monkeypatch):
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
|
||||
with pytest.raises(RuntimeError, match="requires a CUDA GPU"):
|
||||
quantize_model_fp8(nn.Linear(16, 16, dtype=torch.bfloat16))
|
||||
|
||||
|
||||
def test_strict_compile_wrapper_does_not_shadow_torch(monkeypatch):
|
||||
monkeypatch.setattr(torch, "compile", lambda fn, **_kwargs: fn)
|
||||
|
||||
compiled = SafeCompiledCallable(
|
||||
lambda value: value + 1,
|
||||
name="test",
|
||||
mode=None,
|
||||
suppress_errors=False,
|
||||
)
|
||||
|
||||
assert compiled(1) == 2
|
||||
@@ -0,0 +1,40 @@
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
from utils.dataset import MultiVideoConcatDataset
|
||||
|
||||
|
||||
class I2VDatasetFrameAccountingTest(unittest.TestCase):
|
||||
def _dataset(self, root, latent_frames):
|
||||
(root / "video" / "sample").mkdir(parents=True)
|
||||
(root / "caption" / "sample").mkdir(parents=True)
|
||||
total_raw_frames = 1 + (latent_frames - 1) * 4
|
||||
return MultiVideoConcatDataset(
|
||||
data_dir=str(root),
|
||||
video_size=(704, 1280),
|
||||
total_frames=total_raw_frames,
|
||||
independent_first_frame=True,
|
||||
num_frame_per_block=8,
|
||||
temporal_compression_ratio=4,
|
||||
)
|
||||
|
||||
def test_96_frame_i2v_uses_regular_8_frame_blocks(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
dataset = self._dataset(Path(tmp), latent_frames=96)
|
||||
|
||||
self.assertEqual(dataset.first_chunk_latent_frames, 8)
|
||||
self.assertEqual(dataset.first_chunk_frames, 29)
|
||||
self.assertEqual(dataset.total_segments, 12)
|
||||
|
||||
def test_33_frame_i2v_keeps_legacy_one_plus_block_layout(self):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
dataset = self._dataset(Path(tmp), latent_frames=33)
|
||||
|
||||
self.assertEqual(dataset.first_chunk_latent_frames, 9)
|
||||
self.assertEqual(dataset.first_chunk_frames, 33)
|
||||
self.assertEqual(dataset.total_segments, 4)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,63 @@
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
|
||||
from wan_5b.distributed.sp_training import (
|
||||
sp_training_sequence_frame_count,
|
||||
validate_sequence_parallel_training_config,
|
||||
)
|
||||
|
||||
|
||||
class I2VSequenceParallelConfigTest(unittest.TestCase):
|
||||
def test_i2v_training_frames_use_full_configured_sequence(self):
|
||||
cfg = SimpleNamespace(
|
||||
i2v=True,
|
||||
independent_first_frame=True,
|
||||
image_or_video_shape=[1, 96, 48, 44, 80],
|
||||
)
|
||||
|
||||
self.assertEqual(sp_training_sequence_frame_count(cfg), 96)
|
||||
|
||||
def test_i2v_sequence_parallel_partitions_the_96_frame_sequence(self):
|
||||
cfg = SimpleNamespace(
|
||||
i2v=True,
|
||||
independent_first_frame=True,
|
||||
image_or_video_shape=[1, 96, 48, 44, 80],
|
||||
)
|
||||
|
||||
validate_sequence_parallel_training_config(
|
||||
cfg,
|
||||
sp_size=4,
|
||||
num_frame_per_block=8,
|
||||
)
|
||||
|
||||
def test_i2v_sequence_parallel_rejects_non_block_aligned_length(self):
|
||||
cfg = SimpleNamespace(
|
||||
i2v=True,
|
||||
independent_first_frame=True,
|
||||
image_or_video_shape=[1, 97, 48, 44, 80],
|
||||
)
|
||||
|
||||
with self.assertRaisesRegex(ValueError, r"training latent frames"):
|
||||
validate_sequence_parallel_training_config(
|
||||
cfg,
|
||||
sp_size=4,
|
||||
num_frame_per_block=8,
|
||||
)
|
||||
|
||||
def test_t2v_sequence_parallel_uses_full_sequence(self):
|
||||
cfg = SimpleNamespace(
|
||||
i2v=False,
|
||||
independent_first_frame=False,
|
||||
image_or_video_shape=[1, 96, 48, 44, 80],
|
||||
)
|
||||
|
||||
self.assertEqual(sp_training_sequence_frame_count(cfg), 96)
|
||||
validate_sequence_parallel_training_config(
|
||||
cfg,
|
||||
sp_size=4,
|
||||
num_frame_per_block=8,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,192 @@
|
||||
import unittest
|
||||
import importlib.util
|
||||
import pathlib
|
||||
import sys
|
||||
import types
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class _FakeScheduler:
|
||||
def __init__(self):
|
||||
self.num_train_timesteps = 1000
|
||||
self.timesteps = torch.arange(1000, dtype=torch.float32)
|
||||
self.sigmas = torch.linspace(1.0, 0.0, 1000)
|
||||
|
||||
def add_noise(self, clean, noise, timestep):
|
||||
return clean + 10.0
|
||||
|
||||
def training_target(self, clean, noise, timestep):
|
||||
return torch.zeros_like(clean)
|
||||
|
||||
def training_weight(self, timestep):
|
||||
return torch.ones(timestep.numel(), device=timestep.device, dtype=torch.float32)
|
||||
|
||||
|
||||
class _FakeGenerator:
|
||||
def __init__(self):
|
||||
self.recorded = None
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
*,
|
||||
noisy_image_or_video,
|
||||
conditional_dict,
|
||||
timestep,
|
||||
clean_x,
|
||||
aug_t,
|
||||
):
|
||||
self.recorded = {
|
||||
"noisy_image_or_video": noisy_image_or_video.detach().clone(),
|
||||
"timestep": timestep.detach().clone(),
|
||||
"clean_x": clean_x.detach().clone(),
|
||||
"aug_t": aug_t.detach().clone() if aug_t is not None else None,
|
||||
}
|
||||
return torch.zeros_like(noisy_image_or_video), torch.zeros_like(noisy_image_or_video)
|
||||
|
||||
|
||||
class _FakeBuffer:
|
||||
num_blocks = 0
|
||||
|
||||
def is_empty(self):
|
||||
return False
|
||||
|
||||
def add(self, error_block, timestep_index, block_pos=None):
|
||||
pass
|
||||
|
||||
def stats(self):
|
||||
return {
|
||||
"total_added": 0,
|
||||
"filled_buckets": "0/0",
|
||||
"total_entries": 0,
|
||||
}
|
||||
|
||||
|
||||
class _StubBaseModel(torch.nn.Module):
|
||||
def _get_timestep(
|
||||
self,
|
||||
min_timestep,
|
||||
max_timestep,
|
||||
batch_size,
|
||||
num_frame,
|
||||
num_frame_per_block,
|
||||
uniform_timestep=False,
|
||||
):
|
||||
if uniform_timestep:
|
||||
return torch.randint(
|
||||
min_timestep,
|
||||
max_timestep,
|
||||
[batch_size, 1],
|
||||
device=self.device,
|
||||
dtype=torch.long,
|
||||
).repeat(1, num_frame)
|
||||
timestep = torch.randint(
|
||||
min_timestep,
|
||||
max_timestep,
|
||||
[batch_size, num_frame],
|
||||
device=self.device,
|
||||
dtype=torch.long,
|
||||
)
|
||||
timestep = timestep.reshape(timestep.shape[0], -1, num_frame_per_block)
|
||||
timestep[:, :, 1:] = timestep[:, :, 0:1]
|
||||
return timestep.reshape(timestep.shape[0], -1)
|
||||
|
||||
|
||||
def _load_causal_diffusion_with_stubs():
|
||||
repo_root = pathlib.Path(__file__).resolve().parents[1]
|
||||
module_path = repo_root / "model" / "diffusion.py"
|
||||
saved = {
|
||||
name: sys.modules.get(name)
|
||||
for name in (
|
||||
"model",
|
||||
"model.base",
|
||||
"pipeline",
|
||||
"utils.wan_5b_wrapper",
|
||||
)
|
||||
}
|
||||
|
||||
fake_model = types.ModuleType("model")
|
||||
fake_model.__path__ = []
|
||||
fake_base = types.ModuleType("model.base")
|
||||
fake_base.BaseModel = _StubBaseModel
|
||||
|
||||
fake_pipeline = types.ModuleType("pipeline")
|
||||
fake_pipeline.CausalDiffusionInferencePipeline = object
|
||||
|
||||
fake_wrapper = types.ModuleType("utils.wan_5b_wrapper")
|
||||
fake_wrapper.WanDiffusionWrapper = object
|
||||
fake_wrapper.WanTextEncoder = object
|
||||
fake_wrapper.WanVAEWrapper = object
|
||||
|
||||
sys.modules["model"] = fake_model
|
||||
sys.modules["model.base"] = fake_base
|
||||
sys.modules["pipeline"] = fake_pipeline
|
||||
sys.modules["utils.wan_5b_wrapper"] = fake_wrapper
|
||||
try:
|
||||
spec = importlib.util.spec_from_file_location(
|
||||
"_diffusion_under_test", module_path
|
||||
)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(module)
|
||||
return module.CausalDiffusion
|
||||
finally:
|
||||
for name, value in saved.items():
|
||||
if value is None:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = value
|
||||
|
||||
|
||||
class I2VTeacherForcingContextTest(unittest.TestCase):
|
||||
def test_teacher_forcing_keeps_i2v_context_clean_after_augmentation(self):
|
||||
CausalDiffusion = _load_causal_diffusion_with_stubs()
|
||||
model = CausalDiffusion.__new__(CausalDiffusion)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.device = torch.device("cpu")
|
||||
model.dtype = torch.float32
|
||||
model.args = SimpleNamespace(i2v=True)
|
||||
model.independent_first_frame = True
|
||||
model.num_frame_per_block = 2
|
||||
model.scheduler = _FakeScheduler()
|
||||
model.teacher_forcing = True
|
||||
model.noise_augmentation_max_timestep = 10
|
||||
model.generator = _FakeGenerator()
|
||||
model.error_buffer = _FakeBuffer()
|
||||
model.noise_error_buffer = _FakeBuffer()
|
||||
model.er_start_step = 0
|
||||
model.er_clean_prob = 0.0
|
||||
model.er_latent_inject_prob = 0.0
|
||||
model.er_noise_inject_prob = 0.0
|
||||
model.er_context_inject_prob = 1.0
|
||||
model.er_buffer_warmup_iter = -1
|
||||
|
||||
def inject_context_error(clean_latent_aug, index, batch_size, num_frame):
|
||||
return clean_latent_aug + 100.0
|
||||
|
||||
model._inject_error_buffer = inject_context_error
|
||||
|
||||
clean_latent = torch.zeros(1, 4, 1, 1, 1)
|
||||
initial_latent = torch.full((1, 1, 1, 1, 1), 7.0)
|
||||
|
||||
model.generator_loss(
|
||||
image_or_video_shape=[1, 4, 1, 1, 1],
|
||||
conditional_dict={"prompt_embeds": torch.zeros(1, 1)},
|
||||
unconditional_dict={},
|
||||
clean_latent=clean_latent,
|
||||
initial_latent=initial_latent,
|
||||
global_step=0,
|
||||
)
|
||||
|
||||
recorded = model.generator.recorded
|
||||
self.assertTrue(torch.equal(recorded["noisy_image_or_video"][:, :1], initial_latent))
|
||||
self.assertTrue(torch.equal(recorded["clean_x"][:, :1], initial_latent))
|
||||
self.assertTrue((recorded["timestep"][:, :1] == 0).all())
|
||||
self.assertTrue((recorded["aug_t"][:, :1] == 0).all())
|
||||
|
||||
# Later frames still receive clean-side augmentation and context error.
|
||||
self.assertTrue(torch.equal(recorded["clean_x"][:, 1:], torch.full((1, 3, 1, 1, 1), 110.0)))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user