[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.

Author: Nuullll

.gitignore

@@ -1,6 +1,7 @@
 # defaults
 __pycache__
 .ruff_cache
+.vscode
 /cache.json
 /metadata.json
 /config.json

modules/devices.py

@@ -206,6 +206,9 @@ def set_cuda_params():
         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

modules/sd_samplers_kdiffusion.py

@@ -326,14 +326,7 @@ def create_noise_sampler(self, x, sigmas, p):
         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)