Merge pull request #1656 from rockerBOO/master

Fix VAE being applied (for LoRA training)
bmaltais · Nov 11, 2023 · 2b7942f · 2b7942f
2 parents 609e184 + 02bf68a
commit 2b7942f
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/library/common_gui.py b/library/common_gui.py
@@ -881,6 +881,12 @@ def run_cmd_advanced_training(**kwargs):
     wandb_api_key = kwargs.get('wandb_api_key', '')
     if wandb_api_key:
         run_cmd += f' --wandb_api_key="{wandb_api_key}"'
+
+    vae = kwargs.get(
+        'vae'
+    )
+    if vae:
+        run_cmd += ' --vae="{vae}"'
 
     return run_cmd
 

diff --git a/lora_gui.py b/lora_gui.py
@@ -170,6 +170,7 @@ def save_configuration(
     full_bf16,
     min_timestep,
     max_timestep,
+    vae
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -323,6 +324,7 @@ def open_configuration(
     min_timestep,
     max_timestep,
     training_preset,
+    vae
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -496,6 +498,7 @@ def train_model(
     full_bf16,
     min_timestep,
     max_timestep,
+    vae
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -1001,6 +1004,7 @@ def train_model(
         scale_v_pred_loss_like_noise_pred=scale_v_pred_loss_like_noise_pred,
         min_timestep=min_timestep,
         max_timestep=max_timestep,
+        vae=vae,
     )
 
     run_cmd += run_cmd_sample(
@@ -1658,6 +1662,7 @@ def update_LoRA_settings(LoRA_type):
             advanced_training.full_bf16,
             advanced_training.min_timestep,
             advanced_training.max_timestep,
+            advanced_training.vae,
         ]
 
         config.button_open_config.click(