Support PID 1.5 models. (Comfy-Org#14894)

Author: comfyanonymous

comfy/ldm/pixeldit/model.py

@@ -197,6 +197,9 @@ def _pre_patch_block(self, s, i, **kwargs):
         """Hook for subclasses to inject per-block state into the patch stream (e.g. PiD's LQ gate)."""
         return s
 
+    def _pre_pixel_blocks(self, s, **kwargs):
+        return s
+
     def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, **kwargs):
         H_orig, W_orig = x.shape[2], x.shape[3]
         x = comfy.ldm.common_dit.pad_to_patch_size(x, (self.patch_size, self.patch_size))
@@ -226,6 +229,7 @@ def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_
             s, y_emb = blk(s, y_emb, condition, pos_img, pos_txt, None, transformer_options=transformer_options)
         s = F.silu(t_emb + s)
 
+        s = self._pre_pixel_blocks(s, **kwargs)
         s_cond = s.view(B * L, self.hidden_size)
         x_pixels = self.pixel_embedder(x, patch_size=self.patch_size)
         for blk in self.pixel_blocks:

comfy/ldm/pixeldit/pid.py

@@ -13,15 +13,15 @@
 from .modules import precompute_freqs_cis_2d
 
 
-class SigmaAwareGatePerTokenPerDim(nn.Module):
+class SigmaAwareGate(nn.Module):
     """gate = sigmoid(content_proj(cat[x, lq]) - exp(log_alpha) * sigma); out = x + gate * lq.
 
     Trained init gives ~0.88 gate at sigma=0, ~0.05 at sigma=1.
     """
 
-    def __init__(self, dim: int, dtype=None, device=None, operations=None):
+    def __init__(self, dim: int, per_token: bool = False, dtype=None, device=None, operations=None):
         super().__init__()
-        self.content_proj = operations.Linear(dim * 2, dim, dtype=dtype, device=device)
+        self.content_proj = operations.Linear(dim * 2, 1 if per_token else dim, dtype=dtype, device=device)
         self.log_alpha = nn.Parameter(torch.empty((), dtype=dtype, device=device))
 
     def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
@@ -36,15 +36,15 @@ def forward(self, x: torch.Tensor, lq: torch.Tensor, sigma: torch.Tensor) -> tor
 class ResBlock(nn.Module):
     """Pre-activation ResNet block: GN -> SiLU -> Conv -> GN -> SiLU -> Conv + skip."""
 
-    def __init__(self, channels: int, num_groups: int = 4, dtype=None, device=None, operations=None):
+    def __init__(self, channels: int, num_groups: int = 4, conv_padding_mode: str = "zeros", dtype=None, device=None, operations=None):
         super().__init__()
         self.block = nn.Sequential(
             operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
             operations.GroupNorm(num_groups, channels, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Conv2d(channels, channels, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(channels, channels, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -62,9 +62,13 @@ def __init__(
         patch_size: int = 16,
         sr_scale: int = 4,
         latent_spatial_down_factor: int = 8,
+        latent_unpatchify_factor: int = 1,
         num_res_blocks: int = 4,
         num_outputs: int = 7,
         interval: int = 2,
+        conv_padding_mode: str = "zeros",
+        gate_per_token: bool = False,
+        pit_output: bool = False,
         dtype=None, device=None, operations=None,
     ):
         super().__init__()
@@ -74,34 +78,38 @@ def __init__(
         self.patch_size = patch_size
         self.sr_scale = sr_scale
         self.latent_spatial_down_factor = latent_spatial_down_factor
+        self.latent_unpatchify_factor = latent_unpatchify_factor
         self.num_outputs = num_outputs
         self.interval = interval
 
-        z_to_patch_ratio = (sr_scale * latent_spatial_down_factor) / patch_size
+        effective_latent_channels = latent_channels // (latent_unpatchify_factor * latent_unpatchify_factor)
+        effective_spatial_down_factor = latent_spatial_down_factor // latent_unpatchify_factor
+        z_to_patch_ratio = (sr_scale * effective_spatial_down_factor) / patch_size
         self.z_to_patch_ratio = z_to_patch_ratio
         if z_to_patch_ratio >= 1:
             self.latent_fold_factor = 0
-            latent_proj_in_ch = latent_channels
+            latent_proj_in_ch = effective_latent_channels
         else:
             fold_factor = int(1 / z_to_patch_ratio)
             assert fold_factor * z_to_patch_ratio == 1.0
             self.latent_fold_factor = fold_factor
-            latent_proj_in_ch = latent_channels * fold_factor * fold_factor
+            latent_proj_in_ch = effective_latent_channels * fold_factor * fold_factor
 
         layers = [
-            operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(latent_proj_in_ch, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
             nn.SiLU(),
-            operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, dtype=dtype, device=device),
+            operations.Conv2d(hidden_dim, hidden_dim, kernel_size=3, padding=1, padding_mode=conv_padding_mode, dtype=dtype, device=device),
         ]
         for _ in range(num_res_blocks):
-            layers.append(ResBlock(hidden_dim, dtype=dtype, device=device, operations=operations))
+            layers.append(ResBlock(hidden_dim, conv_padding_mode=conv_padding_mode, dtype=dtype, device=device, operations=operations))
         self.latent_proj = nn.Sequential(*layers)
 
         self.output_heads = nn.ModuleList(
             [operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) for _ in range(num_outputs)]
         )
+        self.pit_head = operations.Linear(hidden_dim, out_dim, dtype=dtype, device=device) if pit_output else None
         self.gate_modules = nn.ModuleList(
-            [SigmaAwareGatePerTokenPerDim(out_dim, dtype=dtype, device=device, operations=operations)
+            [SigmaAwareGate(out_dim, per_token=gate_per_token, dtype=dtype, device=device, operations=operations)
              for _ in range(num_outputs)]
         )
 
@@ -115,6 +123,11 @@ def gate(self, x: torch.Tensor, lq_feature: torch.Tensor, sigma: torch.Tensor, o
         return self.gate_modules[out_idx](x, lq_feature, sigma)
 
     def _align_latent_to_patch_grid(self, lq_latent: torch.Tensor, pH: int, pW: int) -> torch.Tensor:
+        f = self.latent_unpatchify_factor
+        if f > 1:
+            B, C, H, W = lq_latent.shape
+            lq_latent = lq_latent.reshape(B, C // (f * f), f, f, H, W)
+            lq_latent = lq_latent.permute(0, 1, 4, 2, 5, 3).reshape(B, C // (f * f), H * f, W * f)
         B, z_dim = lq_latent.shape[:2]
         if self.z_to_patch_ratio >= 1:
             if lq_latent.shape[2] != pH or lq_latent.shape[3] != pW:
@@ -134,7 +147,10 @@ def forward(self, lq_latent: torch.Tensor, target_pH: int, target_pW: int) -> Li
         feat = self._align_latent_to_patch_grid(lq_latent, target_pH, target_pW)
         B, C, H, W = feat.shape
         tokens = feat.permute(0, 2, 3, 1).contiguous().view(B, H * W, C)
-        return [head(tokens) for head in self.output_heads]
+        outputs = [head(tokens) for head in self.output_heads]
+        if self.pit_head is not None:
+            outputs.append(self.pit_head(tokens))
+        return outputs
 
 
 class PidNet(PixDiT_T2I):
@@ -148,6 +164,10 @@ def __init__(
         lq_interval: int = 2,
         sr_scale: int = 4,
         latent_spatial_down_factor: int = 8,
+        lq_latent_unpatchify_factor: int = 1,
+        lq_conv_padding_mode: str = "zeros",
+        lq_gate_per_token: bool = False,
+        pit_lq_inject: bool = False,
         rope_ref_h: int = 1024, # NTK ref resolution in PIXEL units: 1024px / patch=16 -> grid_ref=64.
         rope_ref_w: int = 1024,
         image_model=None,
@@ -165,6 +185,8 @@ def _pit_rope_fn(head_dim, h, w, device=None, dtype=torch.float32, **rope_opts):
         for blk in self.pixel_blocks:
             blk._rope_fn = _pit_rope_fn
 
+        self.pit_lq_inject = pit_lq_inject
+
         num_lq_outputs = (self.patch_depth + lq_interval - 1) // lq_interval
         self.lq_proj = LQProjection2D(
             latent_channels=lq_latent_channels,
@@ -173,13 +195,20 @@ def _pit_rope_fn(head_dim, h, w, device=None, dtype=torch.float32, **rope_opts):
             patch_size=self.patch_size,
             sr_scale=sr_scale,
             latent_spatial_down_factor=latent_spatial_down_factor,
+            latent_unpatchify_factor=lq_latent_unpatchify_factor,
             num_res_blocks=lq_num_res_blocks,
             num_outputs=num_lq_outputs,
             interval=lq_interval,
+            conv_padding_mode=lq_conv_padding_mode,
+            gate_per_token=lq_gate_per_token,
+            pit_output=pit_lq_inject,
             dtype=dtype,
             device=device,
             operations=operations,
         )
+        self.pit_lq_gate = SigmaAwareGate(
+            self.hidden_size, per_token=lq_gate_per_token, dtype=dtype, device=device, operations=operations
+        ) if pit_lq_inject else None
 
     def _fetch_patch_pos(self, height, width, device, dtype, **rope_opts):
         return precompute_freqs_cis_2d(
@@ -197,6 +226,11 @@ def _pre_patch_block(self, s, i, pid_lq_features, pid_degrade_sigma, **kwargs):
             return s
         return self.lq_proj.gate(s, pid_lq_features[out_idx], pid_degrade_sigma, out_idx)
 
+    def _pre_pixel_blocks(self, s, pid_pit_lq_feature=None, pid_degrade_sigma=None, **kwargs):
+        if pid_pit_lq_feature is None:
+            return s
+        return self.pit_lq_gate(s, pid_pit_lq_feature, pid_degrade_sigma)
+
     def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_options={}, lq_latent=None, degrade_sigma=None, **kwargs):
         if lq_latent is None:
             raise ValueError("PidNet requires lq_latent — attach via PiDConditioning")
@@ -216,12 +250,14 @@ def _forward(self, x, timesteps, context=None, attention_mask=None, transformer_
             degrade_sigma = degrade_sigma.expand(B).contiguous()
 
         lq_features = self.lq_proj(lq_latent=lq_latent.to(x), target_pH=Hs, target_pW=Ws)
+        pit_lq_feature = lq_features.pop() if self.pit_lq_inject else None
 
         return super()._forward(
             x, timesteps,
             context=context, attention_mask=attention_mask,
             transformer_options=transformer_options,
             pid_lq_features=lq_features,
+            pid_pit_lq_feature=pit_lq_feature,
             pid_degrade_sigma=degrade_sigma,
             **kwargs,
         )

comfy/model_detection.py

@@ -470,15 +470,46 @@ def detect_unet_config(state_dict, key_prefix, metadata=None):
     # PiD (Pixel Diffusion Decoder). Must check BEFORE plain PixelDiT_T2I.
     _lq_w_key = '{}lq_proj.latent_proj.0.weight'.format(key_prefix)
     if _lq_w_key in state_dict_keys:
-        in_ch = int(state_dict[_lq_w_key].shape[1])
+        latent_proj_in_channels = int(state_dict[_lq_w_key].shape[1])
+        hidden_dim = int(state_dict[_lq_w_key].shape[0])
         _gate_prefix = '{}lq_proj.gate_modules.'.format(key_prefix)
         num_gates = len({k[len(_gate_prefix):].split('.')[0]
                          for k in state_dict_keys if k.startswith(_gate_prefix)})
+        pid_v1_5 = '{}lq_proj.pit_head.weight'.format(key_prefix) in state_dict_keys
         dit_config = {"image_model": "pid",
-                      "lq_latent_channels": in_ch,
-                      "latent_spatial_down_factor": 16 if in_ch >= 64 else 8}
+                      "lq_hidden_dim": hidden_dim}
         if num_gates > 0:
             dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
+        if pid_v1_5:
+            pid_v1_5_variants = {
+                16: {  # Flux and QwenImage
+                    "lq_latent_channels": 16,
+                    "latent_spatial_down_factor": 8,
+                    "lq_latent_unpatchify_factor": 1,
+                },
+                32: {  # Flux2 after 2x latent unpatchify
+                    "lq_latent_channels": 128,
+                    "latent_spatial_down_factor": 16,
+                    "lq_latent_unpatchify_factor": 2,
+                },
+            }
+            variant = pid_v1_5_variants.get(latent_proj_in_channels)
+            if variant is None:
+                raise ValueError(f"Unsupported PiD v1.5 latent projection with {latent_proj_in_channels} input channels")
+            gate_weight = state_dict['{}lq_proj.gate_modules.0.content_proj.weight'.format(key_prefix)]
+            dit_config.update(variant)
+            dit_config.update({
+                "lq_conv_padding_mode": "replicate",
+                "lq_gate_per_token": gate_weight.shape[0] == 1,
+                "pit_lq_inject": True,
+                "rope_ref_h": 2048,
+                "rope_ref_w": 2048,
+            })
+        else:
+            dit_config.update({
+                "lq_latent_channels": latent_proj_in_channels,
+                "latent_spatial_down_factor": 16 if latent_proj_in_channels >= 64 else 8,
+            })
         return dit_config
 
     if '{}core.pixel_embedder.proj.weight'.format(key_prefix) in state_dict_keys:  # PixelDiT T2I

tests-unit/comfy_test/model_detection_test.py

@@ -97,6 +97,21 @@ def _make_seedvr2_3b_shared_mm_sd():
     }
 
 
+def _make_pid_v1_5_sd(latent_proj_channels=16):
+    sd = {
+        "pixel_embedder.proj.weight": torch.empty(16, 3, device="meta"),
+        "lq_proj.latent_proj.0.weight": torch.empty(1024, latent_proj_channels, 3, 3, device="meta"),
+        "lq_proj.pit_head.weight": torch.empty(1536, 1024, device="meta"),
+        "lq_proj.gate_modules.0.content_proj.weight": torch.empty(1, 3072, device="meta"),
+        "pixel_blocks.0.attn.q_norm.weight": torch.empty(72, device="meta"),
+        "pixel_blocks.0.adaLN_modulation.0.weight": torch.empty(24576, 1536, device="meta"),
+        "pixel_blocks.0.adaLN_modulation.0.bias": torch.empty(24576, device="meta"),
+    }
+    for i in range(7):
+        sd[f"lq_proj.gate_modules.{i}.log_alpha"] = torch.empty((), device="meta")
+    return sd
+
+
 def _add_model_diffusion_prefix(sd):
     return {f"model.diffusion_model.{k}": v for k, v in sd.items()}
 
@@ -206,6 +221,43 @@ def test_seedvr2_model_match_accepts_full_checkpoint_prefix(self):
 
         assert type(model_config_from_unet(sd, "model.diffusion_model.")).__name__ == "SeedVR2"
 
+    def test_pid_v1_5_detection(self):
+        sd = _make_pid_v1_5_sd()
+        unet_config = detect_unet_config(sd, "")
+
+        assert unet_config == {
+            "image_model": "pid",
+            "lq_latent_channels": 16,
+            "lq_hidden_dim": 1024,
+            "latent_spatial_down_factor": 8,
+            "lq_interval": 2,
+            "lq_latent_unpatchify_factor": 1,
+            "lq_conv_padding_mode": "replicate",
+            "lq_gate_per_token": True,
+            "pit_lq_inject": True,
+            "rope_ref_h": 2048,
+            "rope_ref_w": 2048,
+        }
+        assert type(model_config_from_unet_config(unet_config, sd)).__name__ == "PiD"
+
+    def test_pid_v1_5_flux2_detection(self):
+        unet_config = detect_unet_config(_make_pid_v1_5_sd(latent_proj_channels=32), "")
+
+        assert unet_config["lq_latent_channels"] == 128
+        assert unet_config["latent_spatial_down_factor"] == 16
+        assert unet_config["lq_latent_unpatchify_factor"] == 2
+
+    def test_pid_v1_5_pixel_adaln_conversion(self):
+        sd = _make_pid_v1_5_sd()
+        model_config = model_config_from_unet_config(detect_unet_config(sd, ""), sd)
+        processed = model_config.process_unet_state_dict(sd)
+
+        assert processed["pixel_blocks.0.attn.q_norm.weight"].shape == (72,)
+        assert processed["pixel_blocks.0.adaLN_modulation_msa.weight"].shape == (12288, 1536)
+        assert processed["pixel_blocks.0.adaLN_modulation_mlp.weight"].shape == (12288, 1536)
+        assert processed["pixel_blocks.0.adaLN_modulation_msa.bias"].shape == (12288,)
+        assert processed["pixel_blocks.0.adaLN_modulation_mlp.bias"].shape == (12288,)
+
     def test_unet_config_and_required_keys_combination_is_unique(self):
         """Each model in the registry must have a unique combination of
         ``unet_config`` and ``required_keys``. If two models share the same