maybe fix slerp

CHANGELOG.md

@@ -52,6 +52,7 @@
   - Fix Face/InstantID
   - Fix CivitAI update model info for all models
   - Fix FP16/BF16 test on model load
+  - Fix variation seed possible NaNs
   - Enumerate diffusers model with multiple variants
 
 ## Update for 2024-03-19

modules/devices.py

@@ -369,7 +369,10 @@ def randn(seed, shape):
         torch.xpu.manual_seed_all(seed)
     if device.type == 'mps':
         return torch.randn(shape, device=cpu).to(device)
-    return torch.randn(shape, device=device)
+    elif shared.opts.diffusers_generator_device == "CPU":
+        return torch.randn(shape, device=cpu)
+    else:
+        return torch.randn(shape, device=device)
 
 
 def randn_without_seed(shape):

modules/processing_helpers.py

@@ -98,17 +98,22 @@ def get_sampler_index(sampler_name: str) -> int:
     return sampler_index
 
 
-def slerp(val, low, high): # from https://discuss.pytorch.org/t/help-regarding-slerp-function-for-generative-model-sampling/32475/3
-    low_norm = low/torch.norm(low, dim=1, keepdim=True)
-    high_norm = high/torch.norm(high, dim=1, keepdim=True)
-    dot = (low_norm*high_norm).sum(1)
-
-    if dot.mean() > 0.9995:
-        return low * val + high * (1 - val)
-
+def slerp(val, lo, hi): # from https://discuss.pytorch.org/t/help-regarding-slerp-function-for-generative-model-sampling/32475/3
+    lo_norm = lo / torch.norm(lo, dim=1, keepdim=True)
+    hi_norm = hi / torch.norm(hi, dim=1, keepdim=True)
+    dot = (lo_norm * hi_norm).sum(1)
+    dot_mean = dot.mean()
+    if dot_mean > 0.9995: # simplifies slerp to lerp if vectors are nearly parallel
+        return lo * val + hi * (1 - val)
+    if dot_mean < 0.0005: # also simplifies slerp to lerp to avoid division-by-zero later on
+        return lo * (1.0 - val) + hi * val
     omega = torch.acos(dot)
     so = torch.sin(omega)
-    res = (torch.sin((1.0-val)*omega)/so).unsqueeze(1)*low + (torch.sin(val*omega)/so).unsqueeze(1) * high
+    lo_res = (torch.sin((1.0 - val) * omega) / so).unsqueeze(1)
+    hi_res = (torch.sin(val * omega) / so).unsqueeze(1)
+    # lo_res[lo_res != lo_res] = 0 # replace nans with zeros, but should not happen with dot_mean filtering
+    # hi_res[hi_res != hi_res] = 0
+    res = lo * lo_res + hi * hi_res
     return res
 
 

scripts/init_latents.py

@@ -17,8 +17,7 @@ class Script(scripts.Script):
     def get_latents(p):
         generator_device = devices.cpu if shared.opts.diffusers_generator_device == "CPU" else shared.device
         generator = [torch.Generator(generator_device).manual_seed(s) for s in p.seeds]
-        shape = (len(generator), shared.sd_model.unet.config.in_channels, p.height // shared.sd_model.vae_scale_factor,
-                 p.width // shared.sd_model.vae_scale_factor)
+        shape = (len(generator), shared.sd_model.unet.config.in_channels, p.height // shared.sd_model.vae_scale_factor, p.width // shared.sd_model.vae_scale_factor)
         latents = randn_tensor(shape, generator=generator, device=shared.sd_model._execution_device, dtype=shared.sd_model.unet.dtype) # pylint: disable=protected-access
         var_generator = [torch.Generator(generator_device).manual_seed(ss) for ss in p.subseeds]
         var_latents = randn_tensor(shape, generator=var_generator, device=shared.sd_model._execution_device, dtype=shared.sd_model.unet.dtype) # pylint: disable=protected-access
@@ -26,14 +25,8 @@ class Script(scripts.Script):
 
     @staticmethod
     def set_slerp(p, latents, var_latents, generator, var_generator):
-        if p.subseed_strength < 1:
-            p.init_latent = slerp(p.subseed_strength, latents, var_latents)
-        if p.subseed_strength == 1:
-            p.init_latent = var_latents
-        if 0 < p.subseed_strength <= 0.5:
-            p.generator = generator
-        if 0.5 < p.subseed_strength <= 1:
-            p.generator = var_generator
+        p.init_latent = slerp(p.subseed_strength, latents, var_latents) if p.subseed_strength < 1 else var_latents
+        p.generator = generator if p.subseed_strength <= 0.5 else var_generator
 
 
     def process_batch(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=arguments-differ
@@ -43,3 +36,4 @@ class Script(scripts.Script):
         if p.subseed_strength != 0 and getattr(shared.sd_model, '_execution_device', None) is not None:
             latents, var_latents, generator, var_generator = self.get_latents(p)
             self.set_slerp(p, latents, var_latents, generator, var_generator)
+            # shared.log.warning(f'Init latents: start={torch.aminmax(latents)} end={torch.aminmax(var_latents)} strength={p.subseed_strength} res={torch.aminmax(p.init_latent)}')