detection

modules_forge/packages/huggingface_guess/detection.py

@@ -1,4 +1,4 @@
-# reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.11.0/comfy/model_detection.py
+# reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.24.1/comfy/model_detection.py
 
 import logging
 
@@ -220,6 +220,17 @@ def detect_unet_config(state_dict: dict, key_prefix: str) -> dict:
 
         return dit_config
 
+    if (_lq_w_key := "{}lq_proj.latent_proj.0.weight".format(key_prefix)) in state_dict_keys:  # PiD
+        _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)})
+        in_ch = int(state_dict[_lq_w_key].shape[1])
+        dit_config = {"image_model": "pid"}
+        dit_config["lq_latent_channels"] = in_ch
+        dit_config["latent_spatial_down_factor"] = 16 if in_ch >= 64 else 8
+        if num_gates > 0:
+            dit_config["lq_interval"] = (14 + num_gates - 1) // num_gates
+        return dit_config
+
     if "{}txt_norm.weight".format(key_prefix) in state_dict_keys:  # Qwen Image
         _qweight: bool = "{}transformer_blocks.0.attn.to_qkv.qweight".format(key_prefix) in state_dict_keys
         dit_config = {"nunchaku": _qweight}

modules_forge/packages/huggingface_guess/latent.py

@@ -1,4 +1,4 @@
-# reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.9.0/comfy/latent_formats.py
+# reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.24.1/comfy/latent_formats.py
 
 import torch
 
@@ -169,3 +169,20 @@ class SDXL_Flux2(Flux2):
         super().__init__()
         self.latent_rgb_factors_reshape = None
         self.latent_channels = 32
+
+
+class RGB(LatentFormat):
+    def __init__(self):
+        self.latent_channels = 3
+        self.latent_rgb_factors = [
+            #  R    G    B
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+        ]
+
+    def process_in(self, latent):
+        return latent
+
+    def process_out(self, latent):
+        return latent

modules_forge/packages/huggingface_guess/model_list.py

@@ -1,4 +1,4 @@
-# reference: https://github.com/Comfy-Org/ComfyUI/blob/master/comfy/supported_models.py
+# reference: https://github.com/Comfy-Org/ComfyUI/blob/v0.24.1/comfy/supported_models.py
 
 from enum import Enum
 
@@ -586,6 +586,54 @@ class ErnieImage(BASE):
         return {"ministral3_3b.transformer": "text_encoder"}
 
 
+class PiD(BASE):
+    huggingface_repo = ""
+
+    unet_config = {
+        "image_model": "pid",
+    }
+
+    sampling_settings = {
+        "shift": 1.5,
+    }
+
+    memory_usage_factor = 0.04
+
+    unet_extra_config = {}
+    latent_format = latent.RGB
+
+    supported_inference_dtypes = [torch.bfloat16, torch.float32]
+
+    vae_key_prefix = ["vae."]
+    text_encoder_key_prefix = ["text_encoders."]
+
+    unet_target = "transformer"
+
+    def process_unet_state_dict(self, state_dict: dict[str, torch.Tensor]):
+        pixel_dim = next(v for k, v in state_dict.items() if k.endswith("pixel_embedder.proj.weight")).shape[0]
+        marker = ".adaLN_modulation.0."
+
+        out = {}
+        for k, v in state_dict.items():
+            if k.startswith("_repa_projector") or k.startswith("net_ema."):
+                continue
+            if k.startswith("core."):
+                k = k[len("core.") :]
+            elif k.startswith("net."):
+                k = k[len("net.") :]
+            if "pixel_blocks." in k and marker in k:
+                p2 = v.shape[0] // (6 * pixel_dim)
+                trail = v.shape[1:]
+                vv = v.view(p2, 6, pixel_dim, *trail)
+                base, suffix = k.split(marker)
+                out[f"{base}.adaLN_modulation_msa.{suffix}"] = vv[:, 0:3].reshape(3 * p2 * pixel_dim, *trail).contiguous()
+                out[f"{base}.adaLN_modulation_mlp.{suffix}"] = vv[:, 3:6].reshape(3 * p2 * pixel_dim, *trail).contiguous()
+            else:
+                out[k] = v
+
+        return out
+
+
 models = [
     SD15,
     SDXL,
@@ -603,4 +651,5 @@ models = [
     WAN21_I2V,
     QwenImage,
     ErnieImage,
+    PiD,
 ]