Add support for ScheduleFree Optimizer Type

bmaltais · Sep 17, 2024 · cee8636 · cee8636
1 parent 74a66d2
commit cee8636
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/kohya_gui/class_basic_training.py b/kohya_gui/class_basic_training.py
@@ -193,6 +193,7 @@ def init_lr_and_optimizer_controls(self) -> None:
                 label="Optimizer",
                 choices=[
                     "AdamW",
+                    "AdamWScheduleFree",
                     "AdamW8bit",
                     "Adafactor",
                     "DAdaptation",
@@ -211,6 +212,7 @@ def init_lr_and_optimizer_controls(self) -> None:
                     "Prodigy",
                     "SGDNesterov",
                     "SGDNesterov8bit",
+                    "SGDScheduleFree",
                 ],
                 value=self.config.get("basic.optimizer", "AdamW8bit"),
                 interactive=True,

diff --git a/requirements.txt b/requirements.txt
@@ -24,6 +24,7 @@ prodigyopt==1.0
 pytorch-lightning==1.9.0
 rich>=13.7.1
 safetensors==0.4.4
+schedulefree==1.2.7
 scipy==1.11.4
 # for T5XXL tokenizer (SD3/FLUX)
 sentencepiece==0.2.0