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()