Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -43,19 +43,20 @@ class GraphDiT(nn.Module):
         self.hidden_size = dm_cfg.hidden_size
         self.mol_visualizer = MolecularVisualization(self.atom_decoder)
 
-        self.denoiser = Transformer(
-            max_n_nodes=self.max_n_nodes,
-            hidden_size=dm_cfg.hidden_size,
-            depth=dm_cfg.depth,
-            num_heads=dm_cfg.num_heads,
-            mlp_ratio=dm_cfg.mlp_ratio,
-            drop_condition=dm_cfg.drop_condition,
-            Xdim=self.Xdim,
-            Edim=self.Edim,
-            ydim=self.ydim,
-        )
+        # self.denoiser = Transformer(
+        #     max_n_nodes=self.max_n_nodes,
+        #     hidden_size=dm_cfg.hidden_size,
+        #     depth=dm_cfg.depth,
+        #     num_heads=dm_cfg.num_heads,
+        #     mlp_ratio=dm_cfg.mlp_ratio,
+        #     drop_condition=dm_cfg.drop_condition,
+        #     Xdim=self.Xdim,
+        #     Edim=self.Edim,
+        #     ydim=self.ydim,
+        # )
+        self.denoiser = None
+
         self.model_dtype = model_dtype
-        
         self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
             dm_cfg.diffusion_noise_schedule, timesteps=dm_cfg.diffusion_steps
         )
@@ -86,53 +87,17 @@ class GraphDiT(nn.Module):
         )
         self.limit_dist = utils.PlaceHolder(X=x_marginals, E=e_marginals, y=None)
 
-    # def to(self, *args, **kwargs):
-    #     self = super().to(*args, **kwargs)
-    #     self.model_dtype = next(self.denoiser.parameters()).dtype
-    #     return self
-
-    def init_model(self, model_dir, verbose=False):
+    def init_model(self, model_dir):
         model_file = os.path.join(model_dir, 'model.pt')
         if os.path.exists(model_file):
-            self.denoiser.load_state_dict(torch.load(model_file, map_location='cpu', weights_only=True))
+            # self.denoiser.load_state_dict(torch.load(model_file, map_location='cpu', weights_only=True))
+            pass
         else:
             raise FileNotFoundError(f"Model file not found: {model_file}")
-        
-        if verbose:
-            print('GraphDiT Denoiser Model initialized.')
-            print('Denoiser model:\n', self.denoiser)
-
-    def save_pretrained(self, output_dir):
-        if not os.path.exists(output_dir):
-            os.makedirs(output_dir)
-        
-        # Save model
-        model_path = os.path.join(output_dir, 'model.pt')
-        torch.save(self.denoiser.state_dict(), model_path)
-        
-        # Save model config
-        config_path = os.path.join(output_dir, 'model_config.yaml')
-        with open(config_path, 'w') as f:
-            yaml.dump(vars(self.model_config), f)
-        
-        # Save data info
-        data_info_path = os.path.join(output_dir, 'data.meta.json')
-        data_info_dict = {
-            "active_atoms": self.data_info.active_atoms,
-            "max_node": self.data_info.max_n_nodes,
-            "n_atoms_per_mol_dist": self.data_info.n_nodes.tolist(),
-            "bond_type_dist": self.data_info.edge_types.tolist(),
-            "transition_E": self.data_info.transition_E.tolist(),
-            "atom_type_dist": self.data_info.node_types.tolist(),
-            "valencies": self.data_info.valency_distribution.tolist()
-        }
-        with open(data_info_path, 'w') as f:
-            json.dump(data_info_dict, f, indent=2)
-        
-        print('GraphDiT Model and configurations saved to:', output_dir)
 
     def disable_grads(self):
-        self.denoiser.disable_grads()
+        pass
+        # self.denoiser.disable_grads()
     
     def forward(
         self, x, edge_index, edge_attr, graph_batch, properties, no_label_index
@@ -215,7 +180,7 @@ class GraphDiT(nn.Module):
         }
         return noisy_data
 
-    # @torch.no_grad()
+    @torch.no_grad()
     def generate(
         self,
         properties,
@@ -307,7 +272,7 @@ class GraphDiT(nn.Module):
 
     def check_valid(self, smiles):
         return check_valid(smiles)
-        
+    
     def sample_p_zs_given_zt(
         self, s, t, X_t, E_t, properties, node_mask, guide_scale, device
     ):