Update graph_decoder/diffusion_model.py

graph_decoder/diffusion_model.py

@@ -1,4 +1,3 @@
-import spaces
 import os 
 import yaml
 import json
@@ -20,73 +19,72 @@ class GraphDiT(nn.Module):
         model_dtype,
     ):
         super().__init__()
-        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)
+        
+        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')
@@ -179,8 +177,7 @@ class GraphDiT(nn.Module):
         }
         return noisy_data
 
-    # @torch.no_grad()
-    @spaces.GPU
+    @torch.no_grad()
     def generate(
         self,
         properties,