4 files changed, 3 insertions, 9 deletions

diff --git a/training/functional.py b/training/functional.py
index 8f47734..ccbb4ad 100644
--- a/training/functional.py
+++ b/training/functional.py
@@ -261,8 +261,8 @@ def loss_step(
     eval: bool = False
 ):
     # Convert images to latent space
-    latents = vae.encode(batch["pixel_values"]).latent_dist.sample().detach()
-    latents = latents * 0.18215
+    latents = vae.encode(batch["pixel_values"]).latent_dist.sample()
+    latents = latents * vae.config.scaling_factor
 
     generator = torch.Generator(device=latents.device).manual_seed(seed + step) if eval else None
 
diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index b4c77f3..8aaed3a 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -1,4 +1,3 @@
-from contextlib import nullcontext
 from typing import Optional
 from functools import partial
 from contextlib import contextmanager, nullcontext
@@ -6,7 +5,6 @@ from pathlib import Path
 import itertools
 
 import torch
-import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator
diff --git a/training/strategy/lora.py b/training/strategy/lora.py
index 88d1824..92abaa6 100644
--- a/training/strategy/lora.py
+++ b/training/strategy/lora.py
@@ -1,11 +1,9 @@
-from contextlib import nullcontext
 from typing import Optional
 from functools import partial
-from contextlib import contextmanager, nullcontext
+from contextlib import contextmanager
 from pathlib import Path
 
 import torch
-import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator
diff --git a/training/strategy/ti.py b/training/strategy/ti.py
index d306f18..da2b81c 100644
--- a/training/strategy/ti.py
+++ b/training/strategy/ti.py
@@ -1,11 +1,9 @@
-from contextlib import nullcontext
 from typing import Optional
 from functools import partial
 from contextlib import contextmanager, nullcontext
 from pathlib import Path
 
 import torch
-import torch.nn as nn
 from torch.utils.data import DataLoader
 
 from accelerate import Accelerator