(
self,
model,
noise,
ddim_steps,
eta=0.,
model_kwargs=None,
progress=True,
temperature=1.,
noise_dropout=0.,
clip_denoised=False,
log_every_t=100,
cfg_scale=1.,
uc_cond=None,
cond_key='y'
)
| 96 | |
| 97 | @torch.no_grad() |
| 98 | def sample( |
| 99 | self, |
| 100 | model, |
| 101 | noise, |
| 102 | ddim_steps, |
| 103 | eta=0., |
| 104 | model_kwargs=None, |
| 105 | progress=True, |
| 106 | temperature=1., |
| 107 | noise_dropout=0., |
| 108 | clip_denoised=False, |
| 109 | log_every_t=100, |
| 110 | cfg_scale=1., |
| 111 | uc_cond=None, |
| 112 | cond_key='y' |
| 113 | ): |
| 114 | bs, dev = noise.shape[0], noise.device |
| 115 | |
| 116 | self.make_schedule(ddim_num_steps=ddim_steps, device=dev, ddim_eta=eta, verbose=False) |
| 117 | |
| 118 | timesteps = self.ddim_timesteps |
| 119 | time_range = np.flip(timesteps) |
| 120 | total_steps = timesteps.shape[0] |
| 121 | iterator = tqdm(time_range, desc='DDIM Sampler', total=total_steps, disable=not progress) |
| 122 | |
| 123 | img = noise |
| 124 | intermediates = {'x_inter': [img], 'pred_x0': [img]} |
| 125 | for i, step in enumerate(iterator): |
| 126 | index = total_steps - i - 1 |
| 127 | ts = torch.full((bs,), step, device=dev, dtype=torch.long) |
| 128 | |
| 129 | outs = self.p_sample_ddim( |
| 130 | model=model, |
| 131 | x=img, |
| 132 | t=ts, |
| 133 | index=index, |
| 134 | temperature=temperature, |
| 135 | noise_dropout=noise_dropout, |
| 136 | model_kwargs=model_kwargs, |
| 137 | clip_denoised=clip_denoised, |
| 138 | cfg_scale=cfg_scale, |
| 139 | uc_cond=uc_cond, |
| 140 | cond_key=cond_key |
| 141 | ) |
| 142 | img, pred_x0 = outs |
| 143 | |
| 144 | if index % log_every_t == 0 or index == total_steps - 1: |
| 145 | intermediates['x_inter'].append(img) |
| 146 | intermediates['pred_x0'].append(pred_x0) |
| 147 | |
| 148 | return img, intermediates |
| 149 | |
| 150 | def p_sample_ddim( |
| 151 | self, |
nothing calls this directly
no test coverage detected