1 files changed, 5 insertions, 6 deletions

diff --git a/training/functional.py b/training/functional.py
index ee73ab2..87bb339 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -277,7 +277,7 @@ def loss_step(
     with_prior_preservation: bool,
     prior_loss_weight: float,
     seed: int,
-    perlin_strength: float,
+    offset_noise_strength: float,
     step: int,
     batch: dict[str, Any],
     eval: bool = False,
@@ -300,11 +300,10 @@ def loss_step(
         generator=generator
     )
 
-    if perlin_strength != 0:
-        noise += perlin_strength * perlin_noise(
+    if offset_noise_strength != 0:
+        noise += offset_noise_strength * perlin_noise(
             latents.shape,
             res=1,
-            octaves=4,
             dtype=latents.dtype,
             device=latents.device,
             generator=generator
@@ -610,7 +609,7 @@ def train(
     global_step_offset: int = 0,
     with_prior_preservation: bool = False,
     prior_loss_weight: float = 1.0,
-    perlin_strength: float = 0.1,
+    offset_noise_strength: float = 0.1,
     **kwargs,
 ):
     text_encoder, unet, optimizer, train_dataloader, val_dataloader, lr_scheduler, extra = strategy.prepare(
@@ -642,7 +641,7 @@ def train(
         with_prior_preservation,
         prior_loss_weight,
         seed,
-        perlin_strength,
+        offset_noise_strength,
     )
 
     if accelerator.is_main_process: