Removed offset noise from training, added init offset to pipeline

1 files changed, 2 insertions, 10 deletions

diff --git a/training/functional.py b/training/functional.py
index 36269f0..1c38635 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -253,7 +253,6 @@ def loss_step(
     text_encoder: CLIPTextModel,
     with_prior_preservation: bool,
     prior_loss_weight: float,
-    noise_offset: float,
     seed: int,
     step: int,
     batch: dict[str, Any],
@@ -268,17 +267,12 @@ def loss_step(
     generator = torch.Generator(device=latents.device).manual_seed(seed + step) if eval else None
 
     # Sample noise that we'll add to the latents
-    offsets = noise_offset * torch.randn(
-        latents.shape[0], 1, 1, 1,
+    noise = torch.randn(
+        latents.shape,
         dtype=latents.dtype,
         layout=latents.layout,
         device=latents.device,
         generator=generator
-    ).expand(latents.shape)
-    noise = torch.normal(
-        mean=offsets,
-        std=1,
-        generator=generator,
     )
 
     # Sample a random timestep for each image
@@ -565,7 +559,6 @@ def train(
     global_step_offset: int = 0,
     with_prior_preservation: bool = False,
     prior_loss_weight: float = 1.0,
-    noise_offset: float = 0.2,
     **kwargs,
 ):
     text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, extra = strategy.prepare(
@@ -600,7 +593,6 @@ def train(
         text_encoder,
         with_prior_preservation,
         prior_loss_weight,
-        noise_offset,
         seed,
     )