[IPEX] Support SDE samplers

This is a W/A since `torch.Generator()` API doesn't support `xpu`
backend at the moment. So replacing it with `torch.xpu.Generator()` API
provided by IPEX.
This commit is contained in:
Nuullll
2023-07-05 15:45:53 +08:00
parent 45d50bd106
commit 860bf8e2bf
3 changed files with 5 additions and 8 deletions
+3
View File
@@ -206,6 +206,9 @@ if backend == 'ipex':
args[4].to("cpu") if args[4] is not None else args[4],
args[5], args[6], args[7], args[8]).to(get_cuda_device_string()),
lambda *args, **kwargs: args[1].device != torch.device("cpu"))
CondFunc('torchsde._brownian.brownian_interval._randn',
lambda _, size, dtype, device, seed: torch.randn(size, dtype=dtype, device=device, generator=torch.xpu.Generator(device).manual_seed(int(seed))),
lambda _, size, dtype, device, seed: device != torch.device("cpu"))
cpu = torch.device("cpu")
device = device_interrogate = device_gfpgan = device_esrgan = device_codeformer = None
+1 -8
View File
@@ -326,14 +326,7 @@ class KDiffusionSampler:
sigma_max = sigmas.max()
current_iter_seeds = p.all_seeds[p.iteration * p.batch_size:(p.iteration + 1) * p.batch_size]
if devices.backend == 'ipex': #Remove this after Intel adds support for torch.Generator()
try:
return BrownianTreeNoiseSampler(x.to("cpu"), sigma_min, sigma_max, seed=current_iter_seeds, transform=lambda x: x.to("cpu"), transform_last=lambda x: x.to(shared.device)) # pylint: disable=E1123
except Exception:
shared.log.error("Please apply this patch to repositories/k-diffusion/k_diffusion/sampling.py: https://github.com/crowsonkb/k-diffusion/pull/68/files")
return None
else:
return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds)
return BrownianTreeNoiseSampler(x, sigma_min, sigma_max, seed=current_iter_seeds)
def sample_img2img(self, p, x, noise, conditioning, unconditional_conditioning, steps=None, image_conditioning=None):
steps, t_enc = sd_samplers_common.setup_img2img_steps(p, steps)