diff --git a/modules_forge/packages/huggingface_guess/detection.py b/modules_forge/packages/huggingface_guess/detection.py index a99fdad7..aa012038 100644 --- a/modules_forge/packages/huggingface_guess/detection.py +++ b/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} diff --git a/modules_forge/packages/huggingface_guess/latent.py b/modules_forge/packages/huggingface_guess/latent.py index 3df146df..45dbb5e5 100644 --- a/modules_forge/packages/huggingface_guess/latent.py +++ b/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 diff --git a/modules_forge/packages/huggingface_guess/model_list.py b/modules_forge/packages/huggingface_guess/model_list.py index 504c9931..766e803b 100644 --- a/modules_forge/packages/huggingface_guess/model_list.py +++ b/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, ]