| 223 | |
| 224 | |
| 225 | class FlowSDE: |
| 226 | def __init__(self, schedule, sample_eps=0): |
| 227 | """ Sampler class for the FlowModel """ |
| 228 | self.sample_eps = sample_eps # velocity & [GVP, LINEAR] is stable everywhere, hence 0 |
| 229 | self.schedule = schedule |
| 230 | |
| 231 | def drift(self, x, t, model, **model_kwargs): |
| 232 | model_output = model(x, t, **model_kwargs) |
| 233 | assert model_output.shape == x.shape, "Output shape from ODE solver must match input shape" |
| 234 | return model_output |
| 235 | |
| 236 | def score(self, x, t, model, **model_kwargs): |
| 237 | # we only train velocity, hence only need to compute score from velocity |
| 238 | score_out = self.schedule.get_score_from_velocity(model(x, t, **model_kwargs), x, t) |
| 239 | return score_out |
| 240 | |
| 241 | def check_interval(self, diffusion_form="sigma", reverse=False, last_step_size=0.04): |
| 242 | t0 = 0 |
| 243 | t1 = 1 |
| 244 | eps = self.sample_eps |
| 245 | if (isinstance(self.schedule, GVPSchedule) or isinstance(self.schedule, LinearSchedule)): |
| 246 | # avoid numerical issue by taking a first semi-implicit step |
| 247 | t0 = eps if diffusion_form == "SBDM" else 0 |
| 248 | t1 = 1 - eps if last_step_size == 0 else 1 - last_step_size |
| 249 | |
| 250 | if reverse: |
| 251 | t0, t1 = 1 - t0, 1 - t1 |
| 252 | |
| 253 | return t0, t1 |
| 254 | |
| 255 | def __get_sde_diffusion_and_drift( |
| 256 | self, |
| 257 | diffusion_form="SBDM", |
| 258 | diffusion_norm=1.0, |
| 259 | ): |
| 260 | |
| 261 | def diffusion_fn(x, t): |
| 262 | diffusion = self.schedule.compute_diffusion(x, t, form=diffusion_form, norm=diffusion_norm) |
| 263 | return diffusion |
| 264 | |
| 265 | sde_drift = \ |
| 266 | lambda x, t, model, **kwargs: \ |
| 267 | self.drift(x, t, model, **kwargs) + diffusion_fn(x, t) * self.score(x, t, model, **kwargs) |
| 268 | |
| 269 | sde_diffusion = diffusion_fn |
| 270 | |
| 271 | return sde_drift, sde_diffusion |
| 272 | |
| 273 | def last_step( |
| 274 | self, |
| 275 | x, |
| 276 | t, |
| 277 | model, |
| 278 | sde_drift, |
| 279 | last_step, |
| 280 | last_step_size, |
| 281 | **model_kwargs |
| 282 | ): |