changed

ambient_diffusion.py

@@ -59,7 +59,7 @@ def make_ambient_batch(
     A = sample_inpainting_mask((B, C, H, W), p=p, device=device)
 
     # 2.  add diffusion noise → x_t , then A x_t
-    noise   = torch.randn_like(clean)
+    noise = torch.randn_like(clean)
 
     if timesteps.ndim == 0:
         timesteps = torch.full((B,), timesteps, device=device, dtype=torch.long)
@@ -114,10 +114,10 @@ class AmbientDDPMPipeline(DDPMPipeline):
         device = self.device
         h = w = self.unet.config.sample_size
 
-        mask = sample_inpainting_mask(
-            (batch_size, 1, h, w), p=self.p_mask, device=device
-        )
+        # Create a dummy mask that's all ones (no masking)
+        dummy_mask = torch.ones((batch_size, 1, h, w), device=device)
 
+        # Start with random noise
         x = torch.randn(
             (batch_size, 1, h, w),
             generator=generator,
@@ -126,29 +126,28 @@ class AmbientDDPMPipeline(DDPMPipeline):
 
         self.scheduler.set_timesteps(num_inference_steps, device=device)
 
-
+        # Set up conditioning labels
         if class_labels is None:
-                cond_lbls = torch.zeros(
-                    (batch_size,
-                    getattr(self.unet.config, "multihot_dim", 1)),
-                    dtype=torch.long if not hasattr(self.unet.config, "multihot_dim") else torch.float,
-                    device=device,
-                )
+            cond_lbls = torch.zeros(
+                (batch_size, getattr(self.unet.config, "multihot_dim", 1)),
+                dtype=torch.long if not hasattr(self.unet.config, "multihot_dim") else torch.float,
+                device=device,
+            )
         else:
             cond_lbls = class_labels.to(device)
 
         uncond_lbls = torch.zeros_like(cond_lbls)
 
         for t in self.scheduler.timesteps:
+            # Always include mask channel to match training, but mask is all ones
+            x_with_mask = torch.cat([x, dummy_mask], dim=1)
 
             if guidance_scale == 1.0 or class_labels is None:
-                # single unconditional (or conditional) pass
-                net_inp = torch.cat([mask * x, mask[:, :1]], dim=1)
-                eps = self.unet(net_inp, t, class_labels=cond_lbls).sample
+                # Single conditional/unconditional pass
+                eps = self.unet(x_with_mask, t, class_labels=cond_lbls).sample
             else:
-                # duplicate batch: cond | uncond
-                net_inp = torch.cat([mask * x, mask[:, :1]], dim=1)  # (B,C+1,H,W)
-                inp     = torch.cat([net_inp, net_inp], dim=0)       # (2B,C+1,H,W)
+                # Classifier-free guidance with duplicate batch
+                inp = torch.cat([x_with_mask, x_with_mask], dim=0)
                 lbls = torch.cat([cond_lbls, uncond_lbls], dim=0)
                 tids = t.expand(2 * batch_size)
 
@@ -157,7 +156,7 @@ class AmbientDDPMPipeline(DDPMPipeline):
                 )
                 eps = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
 
-            # ---------- Eq. 3.3 update (must be INSIDE the loop) ----------
+            # Standard update (same as before)
             x0_hat = eps
             sigma = self.scheduler._get_variance(t).sqrt()
             gamma = sigma ** 2 / (sigma ** 2 + 1)