New Euler_a scheduler

1 files changed, 3 insertions, 3 deletions

diff --git a/textual_inversion.py b/textual_inversion.py
index bcdfd3a..dd7c3bd 100644
--- a/textual_inversion.py
+++ b/textual_inversion.py
@@ -22,7 +22,7 @@ from tqdm.auto import tqdm
 from transformers import CLIPTextModel, CLIPTokenizer
 from slugify import slugify
 
-from schedulers.scheduling_euler_a import EulerAScheduler
+from schedulers.scheduling_euler_ancestral_discrete import EulerAncestralDiscreteScheduler
 from pipelines.stable_diffusion.vlpn_stable_diffusion import VlpnStableDiffusion
 from data.csv import CSVDataModule
 from models.clip.prompt import PromptProcessor
@@ -398,7 +398,7 @@ class Checkpointer:
         samples_path = Path(self.output_dir).joinpath("samples")
 
         unwrapped = self.accelerator.unwrap_model(self.text_encoder)
-        scheduler = EulerAScheduler(
+        scheduler = EulerAncestralDiscreteScheduler(
             beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
         )
 
@@ -639,7 +639,7 @@ def main():
             batched_data = [missing_data[i:i+args.sample_batch_size]
                             for i in range(0, len(missing_data), args.sample_batch_size)]
 
-            scheduler = EulerAScheduler(
+            scheduler = EulerAncestralDiscreteScheduler(
                 beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear"
             )