(
diffusion,
model,
shape,
steps,
clip_denoised=True,
progress=False,
model_kwargs=None,
device=None,
sigma_min=0.002,
sigma_max=80, # higher for highres?
rho=7.0,
sampler="heun",
s_churn=0.0,
s_tmin=0.0,
s_tmax=float("inf"),
s_noise=1.0,
guidance_scale=0.0,
)
| 116 | |
| 117 | |
| 118 | def karras_sample_progressive( |
| 119 | diffusion, |
| 120 | model, |
| 121 | shape, |
| 122 | steps, |
| 123 | clip_denoised=True, |
| 124 | progress=False, |
| 125 | model_kwargs=None, |
| 126 | device=None, |
| 127 | sigma_min=0.002, |
| 128 | sigma_max=80, # higher for highres? |
| 129 | rho=7.0, |
| 130 | sampler="heun", |
| 131 | s_churn=0.0, |
| 132 | s_tmin=0.0, |
| 133 | s_tmax=float("inf"), |
| 134 | s_noise=1.0, |
| 135 | guidance_scale=0.0, |
| 136 | ): |
| 137 | sigmas = get_sigmas_karras(steps, sigma_min, sigma_max, rho, device=device) |
| 138 | x_T = th.randn(*shape, device=device) * sigma_max |
| 139 | sample_fn = {"heun": sample_heun, "dpm": sample_dpm, "ancestral": sample_euler_ancestral}[ |
| 140 | sampler |
| 141 | ] |
| 142 | |
| 143 | if sampler != "ancestral": |
| 144 | sampler_args = dict(s_churn=s_churn, s_tmin=s_tmin, s_tmax=s_tmax, s_noise=s_noise) |
| 145 | else: |
| 146 | sampler_args = {} |
| 147 | |
| 148 | if isinstance(diffusion, KarrasDenoiser): |
| 149 | |
| 150 | def denoiser(x_t, sigma): |
| 151 | _, denoised = diffusion.denoise(model, x_t, sigma, **model_kwargs) |
| 152 | if clip_denoised: |
| 153 | denoised = denoised.clamp(-1, 1) |
| 154 | return denoised |
| 155 | |
| 156 | elif isinstance(diffusion, GaussianDiffusion): |
| 157 | model = GaussianToKarrasDenoiser(model, diffusion) |
| 158 | |
| 159 | def denoiser(x_t, sigma): |
| 160 | _, denoised = model.denoise( |
| 161 | x_t, sigma, clip_denoised=clip_denoised, model_kwargs=model_kwargs |
| 162 | ) |
| 163 | return denoised |
| 164 | |
| 165 | else: |
| 166 | raise NotImplementedError |
| 167 | |
| 168 | if guidance_scale != 0 and guidance_scale != 1: |
| 169 | |
| 170 | def guided_denoiser(x_t, sigma): |
| 171 | x_t = th.cat([x_t, x_t], dim=0) |
| 172 | sigma = th.cat([sigma, sigma], dim=0) |
| 173 | x_0 = denoiser(x_t, sigma) |
| 174 | cond_x_0, uncond_x_0 = th.split(x_0, len(x_0) // 2, dim=0) |
| 175 | x_0 = uncond_x_0 + guidance_scale * (cond_x_0 - uncond_x_0) |
no test coverage detected