1 files changed, 4 insertions, 36 deletions

diff --git a/training/functional.py b/training/functional.py
index 06848cb..c30d1c0 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -321,45 +321,13 @@ def loss_step(
     )
 
     if offset_noise_strength != 0:
-        solid_image = partial(
-            make_solid_image,
-            shape=images.shape[1:],
-            vae=vae,
+        offset_noise = torch.randn(
+            (latents.shape[0], latents.shape[1], 1, 1),
             dtype=latents.dtype,
             device=latents.device,
             generator=generator
-        )
-
-        white_cache_key = f"img_white_{images.shape[2]}_{images.shape[3]}"
-        black_cache_key = f"img_black_{images.shape[2]}_{images.shape[3]}"
-
-        if white_cache_key not in cache:
-            img_white = solid_image(1)
-            cache[white_cache_key] = img_white
-        else:
-            img_white = cache[white_cache_key]
-
-        if black_cache_key not in cache:
-            img_black = solid_image(0)
-            cache[black_cache_key] = img_black
-        else:
-            img_black = cache[black_cache_key]
-
-        offset_strength = torch.rand(
-            (bsz, 1, 1, 1),
-            dtype=latents.dtype,
-            layout=latents.layout,
-            device=latents.device,
-            generator=generator
-        )
-        offset_strength = offset_noise_strength * (offset_strength * 2 - 1)
-        offset_images = torch.where(
-            offset_strength >= 0,
-            img_white.expand(noise.shape),
-            img_black.expand(noise.shape)
-        )
-        offset_strength = offset_strength.abs().expand(noise.shape)
-        noise = slerp(noise, offset_images, offset_strength, zdim=(-1, -2))
+        ).expand(noise.shape)
+        noise += offset_noise_strength * offset_noise
 
     # Sample a random timestep for each image
     timesteps = torch.randint(