Function to include sampling with Classifier-Free Guidance (CFG)
(x, t, model, cfg_scale=1.0, uc_cond=None, cond_key="y", **model_kwargs)
| 191 | |
| 192 | |
| 193 | def forward_with_cfg(x, t, model, cfg_scale=1.0, uc_cond=None, cond_key="y", **model_kwargs): |
| 194 | """ Function to include sampling with Classifier-Free Guidance (CFG) """ |
| 195 | if cfg_scale == 1.0: # without CFG |
| 196 | model_output = model(x, t, **model_kwargs) |
| 197 | |
| 198 | else: # with CFG |
| 199 | assert cond_key in model_kwargs, f"Condition key '{cond_key}' for CFG not found in model_kwargs" |
| 200 | assert uc_cond is not None, "Unconditional condition not provided for CFG" |
| 201 | kwargs = model_kwargs.copy() |
| 202 | c = kwargs[cond_key] |
| 203 | x_in = torch.cat([x] * 2) |
| 204 | t_in = torch.cat([t] * 2) |
| 205 | if uc_cond.shape[0] == 1: |
| 206 | uc_cond = einops.repeat(uc_cond, '1 ... -> bs ...', bs=x.shape[0]) |
| 207 | c_in = torch.cat([uc_cond, c]) |
| 208 | kwargs[cond_key] = c_in |
| 209 | model_uc, model_c = model(x_in, t_in, **kwargs).chunk(2) |
| 210 | model_output = model_uc + cfg_scale * (model_c - model_uc) |
| 211 | |
| 212 | return model_output |
| 213 | |
| 214 | |
| 215 | if __name__ == "__main__": |
nothing calls this directly
no outgoing calls
no test coverage detected