Update graph_decoder/diffusion_model.py
Browse files
graph_decoder/diffusion_model.py
CHANGED
|
@@ -1,5 +1,3 @@
|
|
| 1 |
-
import spaces
|
| 2 |
-
|
| 3 |
import os
|
| 4 |
import yaml
|
| 5 |
import json
|
|
@@ -57,10 +55,6 @@ class GraphDiT(nn.Module):
|
|
| 57 |
ydim=self.ydim,
|
| 58 |
)
|
| 59 |
self.model_dtype = model_dtype
|
| 60 |
-
# self.device = next(self.denoiser.parameters()).device
|
| 61 |
-
|
| 62 |
-
# model_params = torch.load(model_params_path, map_location='cpu')
|
| 63 |
-
# self.denoiser.load_state_dict(model_params)
|
| 64 |
|
| 65 |
self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
|
| 66 |
dm_cfg.diffusion_noise_schedule, timesteps=dm_cfg.diffusion_steps
|
|
@@ -221,7 +215,7 @@ class GraphDiT(nn.Module):
|
|
| 221 |
}
|
| 222 |
return noisy_data
|
| 223 |
|
| 224 |
-
@torch.no_grad()
|
| 225 |
def generate(
|
| 226 |
self,
|
| 227 |
properties,
|
|
|
|
|
|
|
|
|
|
| 1 |
import os
|
| 2 |
import yaml
|
| 3 |
import json
|
|
|
|
| 55 |
ydim=self.ydim,
|
| 56 |
)
|
| 57 |
self.model_dtype = model_dtype
|
|
|
|
|
|
|
|
|
|
|
|
|
| 58 |
|
| 59 |
self.noise_schedule = utils.PredefinedNoiseScheduleDiscrete(
|
| 60 |
dm_cfg.diffusion_noise_schedule, timesteps=dm_cfg.diffusion_steps
|
|
|
|
| 215 |
}
|
| 216 |
return noisy_data
|
| 217 |
|
| 218 |
+
# @torch.no_grad()
|
| 219 |
def generate(
|
| 220 |
self,
|
| 221 |
properties,
|