Added additional config options for custom plugins I needed

2024-01-15 08:31:09 -07:00
parent e190fbaeb8
commit 5276975fb0
7 changed files with 37 additions and 31 deletions
--- a/toolkit/config_modules.py
+++ b/toolkit/config_modules.py
@@ -222,6 +222,7 @@ class TrainConfig:
        self.negative_prompt = kwargs.get('negative_prompt', None)
        # multiplier applied to loos on regularization images
        self.reg_weight = kwargs.get('reg_weight', 1.0)
+        self.num_train_timesteps = kwargs.get('num_train_timesteps', 1000)

        # dropout that happens before encoding. It functions independently per text encoder
        self.prompt_dropout_prob = kwargs.get('prompt_dropout_prob', 0.0)