| 286 | cfg_scale = 1. |
| 287 | |
| 288 | def make_cfg_model_fn(model): |
| 289 | def cfg_model_fn(x, sigma, class_cond): |
| 290 | x_in = torch.cat([x, x]) |
| 291 | sigma_in = torch.cat([sigma, sigma]) |
| 292 | class_uncond = torch.full_like(class_cond, num_classes) |
| 293 | class_cond_in = torch.cat([class_uncond, class_cond]) |
| 294 | out = model(x_in, sigma_in, class_cond=class_cond_in) |
| 295 | out_uncond, out_cond = out.chunk(2) |
| 296 | return out_uncond + (out_cond - out_uncond) * cfg_scale |
| 297 | if cfg_scale != 1: |
| 298 | return cfg_model_fn |
| 299 | return model |
| 300 | |
| 301 | |
| 302 | #sample a number of random images from the mdoel |