Upload folder using huggingface_hub

fcdm_diffae/decoder.py

@@ -23,10 +23,9 @@ class Decoder(nn.Module):
         -> Start blocks (2) -> Middle blocks (4) -> Skip fuse -> End blocks (2)
         -> Conv1x1 -> PixelShuffle
 
-    Dual PDG at inference:
-    - Path drop: replace middle block output with ``path_drop_mask_feature``.
-    - Token mask: replace a fraction of upsampled latent tokens with
-      ``latent_mask_feature`` before fusion.
+    Path-Drop Guidance (PDG) at inference:
+    - Replace middle block output with ``path_drop_mask_feature`` to create
+      an unconditional prediction, then extrapolate.
     """
 
     def __init__(

fcdm_diffae/samplers.py

@@ -1,4 +1,4 @@
-"""DDIM and DPM++2M samplers for VP diffusion with dual PDG support."""
+"""DDIM and DPM++2M samplers for VP diffusion with path-drop PDG support."""
 
 from __future__ import annotations
 
@@ -15,7 +15,7 @@ from .vp_diffusion import (
 
 
 class DecoderForwardFn(Protocol):
-    """Callable that predicts x0 from (x_t, t, latents) with dual PDG flags."""
+    """Callable that predicts x0 from (x_t, t, latents) with path-drop PDG flag."""
 
     def __call__(
         self,
@@ -72,7 +72,7 @@ def _predict_with_pdg(
         state: Current noised state [B, C, H, W].
         t_vec: Timestep vector [B].
         latents: Encoder latents.
-        pdg_mode: "disabled", "path_drop", or "token_mask".
+        pdg_mode: "disabled" or "path_drop".
         pdg_strength: CFG-like strength for PDG.
 
     Returns:
@@ -86,14 +86,6 @@ def _predict_with_pdg(
             torch.float32
         )
         return x0_uncond + pdg_strength * (x0_cond - x0_uncond)
-    elif pdg_mode == "token_mask":
-        x0_uncond = forward_fn(state, t_vec, latents, mask_latent_tokens=True).to(
-            torch.float32
-        )
-        x0_cond = forward_fn(state, t_vec, latents, mask_latent_tokens=False).to(
-            torch.float32
-        )
-        return x0_uncond + pdg_strength * (x0_cond - x0_uncond)
     else:
         return forward_fn(state, t_vec, latents, drop_middle_blocks=False).to(
             torch.float32
@@ -114,7 +106,7 @@ def run_ddim(
     pdg_strength: float = 1.5,
     device: torch.device | None = None,
 ) -> Tensor:
-    """Run DDIM sampling loop with dual PDG support.
+    """Run DDIM sampling loop with path-drop PDG support.
 
     Args:
         forward_fn: Decoder forward function (x_t, t, latents) -> x0_hat.
@@ -123,7 +115,7 @@ def run_ddim(
         latents: Encoder latents [B, bottleneck_dim, h, w].
         logsnr_min, logsnr_max: VP schedule endpoints.
         log_change_high, log_change_low: Shifted-cosine schedule parameters.
-        pdg_mode: "disabled", "path_drop", or "token_mask".
+        pdg_mode: "disabled" or "path_drop".
         pdg_strength: CFG-like strength for PDG.
         device: Target device.
 
@@ -190,7 +182,7 @@ def run_dpmpp_2m(
     pdg_strength: float = 1.5,
     device: torch.device | None = None,
 ) -> Tensor:
-    """Run DPM++2M sampling loop with dual PDG support.
+    """Run DPM++2M sampling loop with path-drop PDG support.
 
     Multi-step solver using exponential integrator formulation in half-lambda space.
     """