(
self,
model,
x,
t,
index,
use_original_steps=False,
temperature=1.,
noise_dropout=0.,
model_kwargs=None,
clip_denoised=False,
cfg_scale=1.,
uc_cond=None,
cond_key='y'
)
| 148 | return img, intermediates |
| 149 | |
| 150 | def p_sample_ddim( |
| 151 | self, |
| 152 | model, |
| 153 | x, |
| 154 | t, |
| 155 | index, |
| 156 | use_original_steps=False, |
| 157 | temperature=1., |
| 158 | noise_dropout=0., |
| 159 | model_kwargs=None, |
| 160 | clip_denoised=False, |
| 161 | cfg_scale=1., |
| 162 | uc_cond=None, |
| 163 | cond_key='y' |
| 164 | ): |
| 165 | model_kwargs = model_kwargs or {} |
| 166 | |
| 167 | # model prediction |
| 168 | if cfg_scale == 1.: # without CFG |
| 169 | model_output = model(x, t, **model_kwargs) |
| 170 | else: # with CFG |
| 171 | assert cond_key in model_kwargs, f"Condition key '{cond_key}' for CFG not found in model_kwargs" |
| 172 | assert uc_cond is not None, "Unconditional condition not provided for CFG" |
| 173 | kwargs = model_kwargs.copy() |
| 174 | c = kwargs[cond_key] |
| 175 | x_in = torch.cat([x] * 2) |
| 176 | t_in = torch.cat([t] * 2) |
| 177 | if uc_cond.shape[0] == 1: |
| 178 | uc_cond = einops.repeat(uc_cond, '1 ... -> bs ...', bs=x.shape[0]) |
| 179 | c_in = torch.cat([uc_cond, c]) |
| 180 | kwargs[cond_key] = c_in |
| 181 | model_uc, model_c = model(x_in, t_in, **kwargs).chunk(2) |
| 182 | model_output = model_uc + cfg_scale * (model_c - model_uc) |
| 183 | |
| 184 | if self.ddpm.parameterization == "v": |
| 185 | e_t = self.ddpm.predict_eps_from_z_and_v(x, t, model_output) |
| 186 | else: |
| 187 | e_t = model_output |
| 188 | |
| 189 | # inference schedule |
| 190 | alphas = self.ddpm.alphas_cumprod if use_original_steps else self.ddim_alphas |
| 191 | alphas_prev = self.ddpm.alphas_cumprod_prev if use_original_steps else self.ddim_alphas_prev |
| 192 | sqrt_one_minus_alphas = self.ddpm.sqrt_one_minus_alphas_cumprod if use_original_steps else self.ddim_sqrt_one_minus_alphas |
| 193 | sigmas = self.ddim_sigmas_for_original_num_steps if use_original_steps else self.ddim_sigmas |
| 194 | |
| 195 | # select parameters corresponding to the currently considered timestep |
| 196 | bs, dev = x.shape[0], x.device |
| 197 | a_t = torch.full((bs, 1, 1, 1), alphas[index], device=dev) |
| 198 | a_prev = torch.full((bs, 1, 1, 1), alphas_prev[index], device=dev) |
| 199 | sigma_t = torch.full((bs, 1, 1, 1), sigmas[index], device=dev) |
| 200 | sqrt_one_minus_at = torch.full((bs, 1, 1, 1), sqrt_one_minus_alphas[index], device=dev) |
| 201 | |
| 202 | # current prediction for x_0 |
| 203 | if self.ddpm.parameterization != "v": |
| 204 | pred_x0 = (x - sqrt_one_minus_at * e_t) / a_t.sqrt() |
| 205 | else: |
| 206 | pred_x0 = self.ddpm.predict_start_from_z_and_v(x, t, model_output) |
| 207 |
no test coverage detected