Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -19,72 +19,73 @@ class GraphDiT(nn.Module):
         model_dtype,
     ):
         super().__init__()
-
-        dm_cfg, data_info = utils.load_config(model_config_path, data_info_path)
-
-        input_dims = data_info.input_dims
-        output_dims = data_info.output_dims
-        nodes_dist = data_info.nodes_dist
-        active_index = data_info.active_index
-
-        self.model_config = dm_cfg
-        self.data_info = data_info
-        self.T = dm_cfg.diffusion_steps
-        self.Xdim = input_dims["X"]
-        self.Edim = input_dims["E"]
-        self.ydim = input_dims["y"]
-        self.Xdim_output = output_dims["X"]
-        self.Edim_output = output_dims["E"]
-        self.ydim_output = output_dims["y"]
-        self.node_dist = nodes_dist
-        self.active_index = active_index
-        self.max_n_nodes = data_info.max_n_nodes
-        self.atom_decoder = data_info.atom_decoder
-        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.model_dtype = model_dtype
-        self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
-            dm_cfg.diffusion_noise_schedule, timesteps=dm_cfg.diffusion_steps
-        )
-        x_marginals = data_info.node_types.to(self.model_dtype) / torch.sum(
-            data_info.node_types.to(self.model_dtype)
-        )
-        e_marginals = data_info.edge_types.to(self.model_dtype) / torch.sum(
-            data_info.edge_types.to(self.model_dtype)
-        )
-        x_marginals = x_marginals / x_marginals.sum()
-        e_marginals = e_marginals / e_marginals.sum()
-
-        xe_conditions = data_info.transition_E.to(self.model_dtype)
-        xe_conditions = xe_conditions[self.active_index][:, self.active_index]
-
-        xe_conditions = xe_conditions.sum(dim=1)
-        ex_conditions = xe_conditions.t()
-        xe_conditions = xe_conditions / xe_conditions.sum(dim=-1, keepdim=True)
-        ex_conditions = ex_conditions / ex_conditions.sum(dim=-1, keepdim=True)
-
-        self.transition_model = utils.MarginalTransition(
-            x_marginals=x_marginals,
-            e_marginals=e_marginals,
-            xe_conditions=xe_conditions,
-            ex_conditions=ex_conditions,
-            y_classes=self.ydim_output,
-            n_nodes=self.max_n_nodes,
-        )
-        self.limit_dist = utils.PlaceHolder(X=x_marginals, E=e_marginals, y=None)
+        pass
+
+        # dm_cfg, data_info = utils.load_config(model_config_path, data_info_path)
+
+        # input_dims = data_info.input_dims
+        # output_dims = data_info.output_dims
+        # nodes_dist = data_info.nodes_dist
+        # active_index = data_info.active_index
+
+        # self.model_config = dm_cfg
+        # self.data_info = data_info
+        # self.T = dm_cfg.diffusion_steps
+        # self.Xdim = input_dims["X"]
+        # self.Edim = input_dims["E"]
+        # self.ydim = input_dims["y"]
+        # self.Xdim_output = output_dims["X"]
+        # self.Edim_output = output_dims["E"]
+        # self.ydim_output = output_dims["y"]
+        # self.node_dist = nodes_dist
+        # self.active_index = active_index
+        # self.max_n_nodes = data_info.max_n_nodes
+        # self.atom_decoder = data_info.atom_decoder
+        # 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.model_dtype = model_dtype
+        # self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
+        #     dm_cfg.diffusion_noise_schedule, timesteps=dm_cfg.diffusion_steps
+        # )
+        # x_marginals = data_info.node_types.to(self.model_dtype) / torch.sum(
+        #     data_info.node_types.to(self.model_dtype)
+        # )
+        # e_marginals = data_info.edge_types.to(self.model_dtype) / torch.sum(
+        #     data_info.edge_types.to(self.model_dtype)
+        # )
+        # x_marginals = x_marginals / x_marginals.sum()
+        # e_marginals = e_marginals / e_marginals.sum()
+
+        # xe_conditions = data_info.transition_E.to(self.model_dtype)
+        # xe_conditions = xe_conditions[self.active_index][:, self.active_index]
+
+        # xe_conditions = xe_conditions.sum(dim=1)
+        # ex_conditions = xe_conditions.t()
+        # xe_conditions = xe_conditions / xe_conditions.sum(dim=-1, keepdim=True)
+        # ex_conditions = ex_conditions / ex_conditions.sum(dim=-1, keepdim=True)
+
+        # self.transition_model = utils.MarginalTransition(
+        #     x_marginals=x_marginals,
+        #     e_marginals=e_marginals,
+        #     xe_conditions=xe_conditions,
+        #     ex_conditions=ex_conditions,
+        #     y_classes=self.ydim_output,
+        #     n_nodes=self.max_n_nodes,
+        # )
+        # self.limit_dist = utils.PlaceHolder(X=x_marginals, E=e_marginals, y=None)
 
     def init_model(self, model_dir):
         model_file = os.path.join(model_dir, 'model.pt')