From e9b7143c790ebc8b7b18c41f51d0e799ac84a337 Mon Sep 17 00:00:00 2001
From: Volpeon <git@volpeon.ink>
Date: Mon, 13 Feb 2023 21:19:46 +0100
Subject: Better noise generation during training:
 https://www.crosslabs.org/blog/diffusion-with-offset-noise

---
 training/functional.py | 7 +++++++
 1 file changed, 7 insertions(+)

(limited to 'training')

diff --git a/training/functional.py b/training/functional.py
index 83e70e2..62b8260 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -157,6 +157,7 @@ def save_samples(
                 width=image_size,
                 generator=gen,
                 guidance_scale=guidance_scale,
+                sag_scale=0,
                 num_inference_steps=num_steps,
                 output_type='pil'
             ).images
@@ -273,6 +274,12 @@ def loss_step(
         layout=latents.layout,
         device=latents.device,
         generator=generator
+    ) + 0.1 * torch.randn(
+        latents.shape[0], latents.shape[1], 1, 1,
+        dtype=latents.dtype,
+        layout=latents.layout,
+        device=latents.device,
+        generator=generator
     )
     bsz = latents.shape[0]
     # Sample a random timestep for each image
-- 
cgit v1.2.3-70-g09d2