config

backend/huggingface/nvidia/PiD/model_index.json

@@ -1,7 +1,4 @@
 {
-  "_class_name": "Lumina2Pipeline",
-  "_diffusers_version": "0.35.0.dev0",
-  "_name_or_path": "neta-lumina-v1.0",
   "scheduler": [
     "diffusers",
     "FlowMatchEulerDiscreteScheduler"
@@ -16,10 +13,10 @@
   ],
   "transformer": [
     "diffusers",
-    "Lumina2Transformer2DModel"
+    "PiDTransformer2DModel"
   ],
   "vae": [
     "diffusers",
-    "AutoencoderKL"
+    "PiDAutoVAE"
   ]
 }

backend/huggingface/nvidia/PiD/text_encoder/config.json

@@ -1,14 +1,16 @@
 {
-  "_name_or_path": "google/gemma-2-2b",
   "architectures": [
-    "Gemma2Model"
+    "Gemma2ForCausalLM"
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
   "attn_logit_softcapping": 50.0,
   "bos_token_id": 2,
   "cache_implementation": "hybrid",
-  "eos_token_id": 1,
+  "eos_token_id": [
+    1,
+    107
+  ],
   "final_logit_softcapping": 30.0,
   "head_dim": 256,
   "hidden_act": "gelu_pytorch_tanh",
@@ -26,8 +28,8 @@
   "rms_norm_eps": 1e-06,
   "rope_theta": 10000.0,
   "sliding_window": 4096,
-  "torch_dtype": "float32",
-  "transformers_version": "4.44.2",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.42.4",
   "use_cache": true,
   "vocab_size": 256000
 }

modules_forge/packages/huggingface_guess/model_list.py

@@ -587,7 +587,7 @@ class ErnieImage(BASE):
 
 
 class PiD(BASE):
-    huggingface_repo = ""
+    huggingface_repo = "nvidia/PiD"
 
     unet_config = {
         "image_model": "pid",
@@ -609,29 +609,14 @@ class PiD(BASE):
 
     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
+    def clip_target(self, state_dict: dict):
+        pref = self.text_encoder_key_prefix[0]
+        if "{}gemma2_2b.transformer.model.embed_tokens.weight".format(pref) in state_dict:
+            state_dict.pop("{}gemma2_2b.logit_scale".format(pref), None)
+            state_dict.pop("{}spiece_model".format(pref), None)
+            return {"gemma2_2b.transformer": "text_encoder"}
+        else:
+            return {"gemma2_2b": "text_encoder"}
 
 
 models = [