Upload 3 files

modeling_tinyimagegen.py

@@ -81,6 +81,10 @@ class RotaryEmbedding2D(nn.Module):
         self.register_buffer("inv_freq_w", inv_freq_w, persistent=False)
 
     def forward(self, grid_h: int, grid_w: int, device: torch.device, dtype: torch.dtype = torch.float32):
+        if self.inv_freq_h is None or self.inv_freq_h.device.type == "meta" or (self.inv_freq_h == 0).all():
+            self.inv_freq_h = 1.0 / (self.base ** (torch.arange(0, self.dim_h, 2, dtype=torch.float32, device=device) / self.dim_h))
+            self.inv_freq_w = 1.0 / (self.base ** (torch.arange(0, self.dim_w, 2, dtype=torch.float32, device=device) / self.dim_w))
+
         t_h = torch.arange(grid_h, device=device, dtype=torch.float32)
         t_w = torch.arange(grid_w, device=device, dtype=torch.float32)
 
@@ -129,6 +133,9 @@ class HadamardMLP(nn.Module):
         self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not hasattr(self, "hadamard_mat") or self.hadamard_mat is None or self.hadamard_mat.device.type == "meta" or (self.hadamard_mat == 0).all() or self.hadamard_mat.abs().max() > 10.0:
+            hadamard_mat = get_hadamard_matrix(self.dim, dtype=torch.float32).to(x.device)
+            self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)
         mat = self.hadamard_mat.type_as(x)
         h = (x * self.scale1) @ mat
         g = F.silu(x * self.gate)
@@ -295,7 +302,14 @@ class TinyImageGenPreTrainedModel(PreTrainedModel):
         elif os.path.exists(bin_file):
             state_dict = torch.load(bin_file, map_location="cpu")
         else:
-            return super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+            model = super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+            for module in model.modules():
+                if type(module).__name__ == "RotaryEmbedding2D":
+                    module.inv_freq_h = 1.0 / (module.base ** (torch.arange(0, module.dim_h, 2, dtype=torch.float32) / module.dim_h))
+                    module.inv_freq_w = 1.0 / (module.base ** (torch.arange(0, module.dim_w, 2, dtype=torch.float32) / module.dim_w))
+                elif type(module).__name__ == "HadamardMLP":
+                    module.register_buffer("hadamard_mat", get_hadamard_matrix(module.dim, dtype=torch.float32), persistent=False)
+            return model
 
         model_keys = set(model.state_dict().keys())
         st_keys = set(state_dict.keys())
@@ -307,6 +321,13 @@ class TinyImageGenPreTrainedModel(PreTrainedModel):
 
         model.load_state_dict(state_dict, strict=True)
 
+        for module in model.modules():
+            if type(module).__name__ == "RotaryEmbedding2D":
+                module.inv_freq_h = 1.0 / (module.base ** (torch.arange(0, module.dim_h, 2, dtype=torch.float32) / module.dim_h))
+                module.inv_freq_w = 1.0 / (module.base ** (torch.arange(0, module.dim_w, 2, dtype=torch.float32) / module.dim_w))
+            elif type(module).__name__ == "HadamardMLP":
+                module.register_buffer("hadamard_mat", get_hadamard_matrix(module.dim, dtype=torch.float32), persistent=False)
+
         if torch_dtype is not None:
             model.to(dtype=torch_dtype)