Function to include sampling with Classifier-Free Guidance (CFG)
(x, t, model, cfg_scale=1.0, uc_cond=None, cond_key="y", **model_kwargs)
| 33 | |
| 34 | |
| 35 | def forward_with_cfg(x, t, model, cfg_scale=1.0, uc_cond=None, cond_key="y", **model_kwargs): |
| 36 | """ Function to include sampling with Classifier-Free Guidance (CFG) """ |
| 37 | if cfg_scale == 1.0: # without CFG |
| 38 | model_output = model(x, t, **model_kwargs) |
| 39 | |
| 40 | else: # with CFG |
| 41 | assert cond_key in model_kwargs, f"Condition key '{cond_key}' for CFG not found in model_kwargs" |
| 42 | assert uc_cond is not None, "Unconditional condition not provided for CFG" |
| 43 | kwargs = model_kwargs.copy() |
| 44 | c = kwargs[cond_key] |
| 45 | x_in = torch.cat([x] * 2) |
| 46 | t_in = torch.cat([t] * 2) |
| 47 | if uc_cond.shape[0] == 1: |
| 48 | uc_cond = einops.repeat(uc_cond, '1 ... -> bs ...', bs=x.shape[0]) |
| 49 | c_in = torch.cat([uc_cond, c]) |
| 50 | kwargs[cond_key] = c_in |
| 51 | model_uc, model_c = model(x_in, t_in, **kwargs).chunk(2) |
| 52 | model_output = model_uc + cfg_scale * (model_c - model_uc) |
| 53 | |
| 54 | return model_output |
| 55 | |
| 56 | |
| 57 | """ Schedules """ |