Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -43,18 +43,17 @@ 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 = None
+        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.model_dtype = model_dtype
         self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
@@ -90,14 +89,12 @@ class GraphDiT(nn.Module):
     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))
-            pass
+            self.denoiser.load_state_dict(torch.load(model_file, map_location='cpu', weights_only=True))
         else:
             raise FileNotFoundError(f"Model file not found: {model_file}")
 
     def disable_grads(self):
-        pass
-        # self.denoiser.disable_grads()
+        self.denoiser.disable_grads()
     
     def forward(
         self, x, edge_index, edge_attr, graph_batch, properties, no_label_index
@@ -193,7 +190,6 @@ class GraphDiT(nn.Module):
         properties = torch.tensor(properties, dtype=torch.float).reshape(1, -1).to(device)
         batch_size = properties.size(0)
         assert batch_size == 1
-        # print('self.denoiser.dtype', self.model_dtype)
         if num_nodes is None:
             num_nodes = self.node_dist.sample_n(batch_size, device)
         else: