| import unittest |
|
|
| import torch |
| from torch import nn |
|
|
| from algorithms.dememwm.models.diffusion import Diffusion |
|
|
|
|
| class FakeDenoiser(nn.Module): |
| def __init__(self, output_frames=None): |
| super().__init__() |
| self.output_frames = output_frames |
| self.calls = [] |
|
|
| def forward(self, x, t, action_cond, **kwargs): |
| self.calls.append({ |
| "x_shape": tuple(x.shape), |
| "t_shape": tuple(t.shape), |
| "kwargs": kwargs, |
| }) |
| frames = self.output_frames if self.output_frames is not None else x.shape[1] |
| return torch.zeros((x.shape[0], frames, *x.shape[2:]), device=x.device, dtype=x.dtype) |
|
|
|
|
| def _make_diffusion(output_frames=None): |
| diffusion = Diffusion.__new__(Diffusion) |
| nn.Module.__init__(diffusion) |
| diffusion.x_shape = torch.Size((1, 1, 1)) |
| diffusion.timesteps = 4 |
| diffusion.sampling_timesteps = 4 |
| diffusion.is_ddim_sampling = False |
| diffusion.objective = "pred_noise" |
| diffusion.use_fused_snr = False |
| diffusion.snr_clip = 5.0 |
| diffusion.cum_snr_decay = 0.9 |
| diffusion.ddim_sampling_eta = 0.0 |
| diffusion.clip_noise = 10.0 |
| diffusion.stabilization_level = 1 |
| diffusion.model = FakeDenoiser(output_frames=output_frames) |
|
|
| betas = torch.tensor([0.05, 0.10, 0.15, 0.20], dtype=torch.float32) |
| alphas = 1.0 - betas |
| alphas_cumprod = torch.cumprod(alphas, dim=0) |
| alphas_cumprod_prev = torch.nn.functional.pad(alphas_cumprod[:-1], (1, 0), value=1.0) |
| posterior_variance = betas * (1.0 - alphas_cumprod_prev) / (1.0 - alphas_cumprod) |
| snr = alphas_cumprod / (1.0 - alphas_cumprod) |
|
|
| diffusion.register_buffer("betas", betas) |
| diffusion.register_buffer("alphas_cumprod", alphas_cumprod) |
| diffusion.register_buffer("alphas_cumprod_prev", alphas_cumprod_prev) |
| diffusion.register_buffer("sqrt_alphas_cumprod", torch.sqrt(alphas_cumprod)) |
| diffusion.register_buffer("sqrt_one_minus_alphas_cumprod", torch.sqrt(1.0 - alphas_cumprod)) |
| diffusion.register_buffer("log_one_minus_alphas_cumprod", torch.log(1.0 - alphas_cumprod)) |
| diffusion.register_buffer("sqrt_recip_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod)) |
| diffusion.register_buffer("sqrt_recipm1_alphas_cumprod", torch.sqrt(1.0 / alphas_cumprod - 1.0)) |
| diffusion.register_buffer("posterior_variance", posterior_variance) |
| diffusion.register_buffer("posterior_log_variance_clipped", torch.log(posterior_variance.clamp(min=1e-20))) |
| diffusion.register_buffer("posterior_mean_coef1", betas * torch.sqrt(alphas_cumprod_prev) / (1.0 - alphas_cumprod)) |
| diffusion.register_buffer("posterior_mean_coef2", (1.0 - alphas_cumprod_prev) * torch.sqrt(alphas) / (1.0 - alphas_cumprod)) |
| diffusion.register_buffer("snr", snr) |
| diffusion.register_buffer("clipped_snr", snr.clamp(max=diffusion.snr_clip)) |
| return diffusion |
|
|
|
|
| def _packed_inputs(): |
| x = torch.arange(5, dtype=torch.float32).view(5, 1, 1, 1, 1) |
| action_cond = torch.zeros((5, 1, 3), dtype=torch.float32) |
| noise_levels = torch.tensor([[1], [2], [0], [0], [0]], dtype=torch.long) |
| segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1} |
| return x, action_cond, noise_levels, segments |
|
|
|
|
| class DeMemWMDiffusionTargetOnlyTests(unittest.TestCase): |
| def test_forward_noises_packed_input_but_returns_target_loss(self): |
| torch.manual_seed(0) |
| diffusion = _make_diffusion(output_frames=2) |
| x, action_cond, noise_levels, segments = _packed_inputs() |
|
|
| x_pred, loss = diffusion( |
| x, |
| action_cond, |
| None, |
| noise_levels=noise_levels, |
| reference_length=0, |
| frame_memory_segments=segments, |
| ) |
|
|
| self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1)) |
| self.assertEqual(tuple(x_pred.shape), (2, 1, 1, 1, 1)) |
| self.assertEqual(tuple(loss.shape), (2, 1, 1, 1, 1)) |
|
|
| def test_forward_without_frame_memory_keeps_full_length(self): |
| torch.manual_seed(0) |
| diffusion = _make_diffusion() |
| x, action_cond, noise_levels, _ = _packed_inputs() |
|
|
| x_pred, loss = diffusion( |
| x, |
| action_cond, |
| None, |
| noise_levels=noise_levels, |
| reference_length=0, |
| ) |
|
|
| self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1)) |
| self.assertEqual(tuple(x_pred.shape), tuple(x.shape)) |
| self.assertEqual(tuple(loss.shape), tuple(x.shape)) |
|
|
| def test_padded_frame_memory_masks_keep_target_prediction_and_loss_shapes(self): |
| torch.manual_seed(0) |
| diffusion = _make_diffusion(output_frames=2) |
| x, action_cond, noise_levels, segments = _packed_inputs() |
| masks = { |
| "target": torch.ones((1, 2), dtype=torch.bool), |
| "anchor": torch.ones((1, 1), dtype=torch.bool), |
| "dynamic": torch.zeros((1, 1), dtype=torch.bool), |
| "revisit": torch.ones((1, 1), dtype=torch.bool), |
| } |
|
|
| x_pred, loss = diffusion( |
| x, |
| action_cond, |
| None, |
| noise_levels=noise_levels, |
| reference_length=0, |
| frame_memory_segments=segments, |
| frame_memory_masks=masks, |
| ) |
|
|
| self.assertIs(diffusion.model.calls[0]["kwargs"]["frame_memory_masks"], masks) |
| self.assertEqual(tuple(x_pred.shape), (2, 1, 1, 1, 1)) |
| self.assertEqual(tuple(loss.shape), (2, 1, 1, 1, 1)) |
|
|
| def test_frame_memory_pose_is_batch_first_and_separate_from_pose_cond(self): |
| torch.manual_seed(0) |
| diffusion = _make_diffusion(output_frames=2) |
| x = torch.zeros((5, 2, 1, 1, 1), dtype=torch.float32) |
| action_cond = torch.zeros((5, 2, 3), dtype=torch.float32) |
| noise_levels = torch.zeros((5, 2), dtype=torch.long) |
| pose_cond = torch.full((5, 2, 5), 99.0, dtype=torch.float32) |
| frame_memory_pose = torch.arange(5 * 2 * 5, dtype=torch.float32).view(5, 2, 5) |
| image_hw = torch.tensor([[360, 640], [720, 1280]], dtype=torch.long) |
| segments = {"target": 2, "anchor": 1, "dynamic": 1, "revisit": 1} |
|
|
| diffusion( |
| x, |
| action_cond, |
| pose_cond, |
| noise_levels=noise_levels, |
| reference_length=0, |
| frame_memory_segments=segments, |
| frame_memory_pose=frame_memory_pose, |
| image_hw=image_hw, |
| ) |
|
|
| kwargs = diffusion.model.calls[0]["kwargs"] |
| self.assertIsNone(kwargs["pose_cond"]) |
| self.assertEqual(tuple(kwargs["frame_memory_pose"].shape), (2, 5, 5)) |
| self.assertTrue(torch.equal(kwargs["frame_memory_pose"], frame_memory_pose.permute(1, 0, 2))) |
| self.assertIs(kwargs["image_hw"], image_hw) |
|
|
| def test_baseline_pose_cond_still_reaches_dit_batch_first(self): |
| torch.manual_seed(0) |
| diffusion = _make_diffusion() |
| x = torch.zeros((5, 2, 1, 1, 1), dtype=torch.float32) |
| action_cond = torch.zeros((5, 2, 3), dtype=torch.float32) |
| noise_levels = torch.zeros((5, 2), dtype=torch.long) |
| pose_cond = torch.arange(5 * 2 * 5, dtype=torch.float32).view(5, 2, 5) |
|
|
| diffusion( |
| x, |
| action_cond, |
| pose_cond, |
| noise_levels=noise_levels, |
| reference_length=0, |
| ) |
|
|
| kwargs = diffusion.model.calls[0]["kwargs"] |
| self.assertTrue(torch.equal(kwargs["pose_cond"], pose_cond.permute(1, 0, 2))) |
| self.assertNotIn("frame_memory_pose", kwargs) |
|
|
| def test_posterior_and_sample_steps_return_target_frames_for_frame_memory(self): |
| torch.manual_seed(0) |
| x, action_cond, _, segments = _packed_inputs() |
| curr = torch.tensor([[2], [1], [-1], [-1], [-1]], dtype=torch.long) |
| next_level = torch.tensor([[1], [0], [-1], [-1], [-1]], dtype=torch.long) |
|
|
| diffusion = _make_diffusion(output_frames=2) |
| mean, variance, log_variance = diffusion.p_mean_variance( |
| x, |
| curr, |
| action_cond=action_cond, |
| pose_cond=None, |
| reference_length=0, |
| frame_memory_segments=segments, |
| ) |
| self.assertEqual(tuple(mean.shape), (2, 1, 1, 1, 1)) |
| self.assertEqual(tuple(variance.shape), (2, 1, 1, 1, 1)) |
| self.assertEqual(tuple(log_variance.shape), (2, 1, 1, 1, 1)) |
|
|
| diffusion = _make_diffusion(output_frames=2) |
| ddpm = diffusion.ddpm_sample_step( |
| x, |
| action_cond, |
| None, |
| curr_noise_level=curr, |
| reference_length=0, |
| frame_memory_segments=segments, |
| ) |
| self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1)) |
| self.assertEqual(tuple(ddpm.shape), (2, 1, 1, 1, 1)) |
|
|
| diffusion = _make_diffusion(output_frames=2) |
| ddim = diffusion.ddim_sample_step( |
| x, |
| action_cond, |
| None, |
| curr_noise_level=curr, |
| next_noise_level=next_level, |
| reference_length=0, |
| frame_memory_segments=segments, |
| ) |
| self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1)) |
| self.assertEqual(tuple(ddim.shape), (2, 1, 1, 1, 1)) |
|
|
| diffusion = _make_diffusion(output_frames=2) |
| sample = diffusion.sample_step( |
| x, |
| action_cond, |
| None, |
| curr_noise_level=torch.tensor([[3], [2], [0], [0], [0]], dtype=torch.long), |
| next_noise_level=torch.tensor([[2], [1], [0], [0], [0]], dtype=torch.long), |
| reference_length=0, |
| frame_memory_segments=segments, |
| ) |
| self.assertEqual(diffusion.model.calls[0]["x_shape"], (1, 5, 1, 1, 1)) |
| self.assertEqual(tuple(sample.shape), (2, 1, 1, 1, 1)) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
