(
self,
net_cfg: dict,
timesteps: int = 1000,
beta_schedule: str = 'linear',
loss_type: str = 'l2',
parameterization: str = 'v',
linear_start: float = 0.00085,
linear_end: float = 0.0120,
ddim_steps: int = 10, # 10 timesteps default for FlowModel
)
| 70 | |
| 71 | class DiffusionFlow(nn.Module): |
| 72 | def __init__( |
| 73 | self, |
| 74 | net_cfg: dict, |
| 75 | timesteps: int = 1000, |
| 76 | beta_schedule: str = 'linear', |
| 77 | loss_type: str = 'l2', |
| 78 | parameterization: str = 'v', |
| 79 | linear_start: float = 0.00085, |
| 80 | linear_end: float = 0.0120, |
| 81 | ddim_steps: int = 10, # 10 timesteps default for FlowModel |
| 82 | ): |
| 83 | super().__init__() |
| 84 | self.net = instantiate_from_config(net_cfg) |
| 85 | |
| 86 | self.diffusion_cfg = dict( |
| 87 | timesteps=timesteps, |
| 88 | beta_schedule=beta_schedule, |
| 89 | zero_terminal_snr=False, |
| 90 | loss_type=loss_type, |
| 91 | parameterization=parameterization, |
| 92 | linear_start=linear_start, |
| 93 | linear_end=linear_end, |
| 94 | cosine_s=8e-3, |
| 95 | original_elbo_weight=0., |
| 96 | v_posterior=0., |
| 97 | l_simple_weight=1.0 |
| 98 | ) |
| 99 | self.diffusion = GaussianDiffusion(**self.diffusion_cfg) |
| 100 | |
| 101 | self.ddim_steps = ddim_steps |
| 102 | self.ddim_sampler = DDIMSampler(self.diffusion) |
| 103 | |
| 104 | def forward(self, x: torch.Tensor, t: torch.Tensor, **kwargs): |
| 105 | return self.net(x, t, **kwargs) |
no test coverage detected