From 9cedf18a974cd272c260da5bbe8d4c585f31a301 Mon Sep 17 00:00:00 2001
From: rockerBOO <rockerboo@gmail.com>
Date: Tue, 15 Jul 2025 19:53:54 -0400
Subject: [PATCH] Add import hint

---
 library/train_util.py | 3 +++
 train_network.py      | 6 +++++-
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/library/train_util.py b/library/train_util.py
index d7243e07..09c101a8 100644
--- a/library/train_util.py
+++ b/library/train_util.py
@@ -6029,6 +6029,9 @@ def get_noise_noisy_latents_and_timesteps(
     else:
         noisy_latents = noise_scheduler.add_noise(latents, noise, timesteps)
 
+
+    noise_scheduler.alphas_cumprod = noise_scheduler.alphas_cumprod.to("cpu")
+
     return noise, noisy_latents, timesteps
 
 
diff --git a/train_network.py b/train_network.py
index 2eabcdf3..02995c95 100644
--- a/train_network.py
+++ b/train_network.py
@@ -45,7 +45,11 @@ from library.custom_train_functions import (
     apply_masked_loss,
 )
 from library.utils import setup_logging, add_logging_arguments
-from wavelet_loss import WaveletLoss
+
+try:
+    from wavelet_loss import WaveletLoss
+except:
+    raise ImportError("WaveletLoss is not installed. Please install it with `pip install git+https://github.com/rockerBOO/wavelet-loss`")
 
 setup_logging()
 import logging