Upload 3 files

config.json

@@ -4,7 +4,7 @@
   ],
   "auto_map": {
     "AutoConfig": "configuration_tinyimagegen.TinyImageGenConfig",
-    "AutoModel": "modeling_tinyimagegen.TinyImageGenModel",
+    "AutoModel": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion",
     "AutoModelForImageDiffusion": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion"
   },
   "dtype": "float32",

configuration_tinyimagegen.py

@@ -44,7 +44,7 @@ class TinyImageGenConfig(PretrainedConfig):
 
         self.auto_map = {
             "AutoConfig": "configuration_tinyimagegen.TinyImageGenConfig",
-            "AutoModel": "modeling_tinyimagegen.TinyImageGenModel",
+            "AutoModel": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion",
             "AutoModelForImageDiffusion": "modeling_tinyimagegen.TinyImageGenModelForImageDiffusion",
         }
 

modeling_tinyimagegen.py

@@ -292,13 +292,21 @@ class TinyImageGenPreTrainedModel(PreTrainedModel):
 
         if os.path.exists(st_file):
             state_dict = load_file(st_file)
-            model.load_state_dict(state_dict, strict=False)
         elif os.path.exists(bin_file):
             state_dict = torch.load(bin_file, map_location="cpu")
-            model.load_state_dict(state_dict, strict=False)
         else:
             return super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
 
+        model_keys = set(model.state_dict().keys())
+        st_keys = set(state_dict.keys())
+        if not model_keys.intersection(st_keys):
+            if any(k.startswith("model.") for k in st_keys):
+                state_dict = {k[6:] if k.startswith("model.") else k: v for k, v in state_dict.items()}
+            elif any(k.startswith("model.") for k in model_keys):
+                state_dict = {f"model.{k}": v for k, v in state_dict.items()}
+
+        model.load_state_dict(state_dict, strict=True)
+
         if torch_dtype is not None:
             model.to(dtype=torch_dtype)