Add slider to control DREAM lambda

This allows adjusting the tradeoff between faster learning and better composition vs better preservation of detail.
RossM · Dec 23, 2023 · 418ba6c · 418ba6c
1 parent 92d3219
commit 418ba6c
Show file tree

Hide file tree

Showing 8 changed files with 31 additions and 4 deletions.
diff --git a/dreambooth/dataclasses/db_config.py b/dreambooth/dataclasses/db_config.py
@@ -72,6 +72,7 @@ class DreamboothConfig(BaseModel):
     max_token_length: int = 75
     min_snr_gamma: float = 0.0
     use_dream: bool = False
+    dream_detail: float = 0.5
     mixed_precision: str = "fp16"
     model_dir: str = ""
     model_name: str = ""

diff --git a/dreambooth/train_dreambooth.py b/dreambooth/train_dreambooth.py
@@ -1621,9 +1621,8 @@ def lora_save_function(weights, filename):
                                 sqrt_alpha_prod = alpha_prod ** 0.5
                                 sqrt_one_minus_alpha_prod = (1 - alpha_prod) ** 0.5
 
-                                # The paper uses lambda = sqrt(1 - alpha) ** p, with p = 1 in their experiments, but
-                                # lambda = 1 seems to give better results for fine-tuning.
-                                dream_lambda = 1
+                                # The paper uses lambda = sqrt(1 - alpha) ** p, with p = 1 in their experiments.
+                                dream_lambda = (1 - alpha_prod) ** args.dream_detail
 
                                 if args.model_type == "SDXL":
                                     with accelerator.autocast():
@@ -1649,7 +1648,7 @@ def lora_save_function(weights, filename):
                                 else:
                                     raise ValueError(f"Unknown prediction type {noise_scheduler.config.prediction_type}")
 
-                                del alpha_prod, sqrt_alpha_prod, sqrt_one_minus_alpha_prod, dream_lambda, predicted_noise, delta_noise
+                                del alpha_prod, sqrt_alpha_prod, sqrt_one_minus_alpha_prod, dream_lambda, model_pred, predicted_noise, delta_noise
 
                         if args.model_type == "SDXL":
                             with accelerator.autocast():

diff --git a/index.html b/index.html
@@ -302,6 +302,11 @@
                                                            id="use_dream" name="use_dream">
                                                     <label class="form-check-label" for="use_dream">Use DREAM</label>
                                                 </div>
+                                            </div>
+                                             <div class="form-group">
+                                                <div class="dbInput db-slider" data-min="0" data-max="1.0"
+                                                     data-step="0.01" id="dream_detail" data-value="0.5"
+                                                     data-label="DREAM detail preservation"></div>
                                             </div>
                                             <div class="form-group">
                                                 <div class="dbInput db-slider" data-min="75" data-max="300"

diff --git a/javascript/dreambooth.js b/javascript/dreambooth.js
@@ -329,6 +329,7 @@ let db_titles = {
     "Use CPU Only (SLOW)": "Guess what - this will be incredibly slow, but it will work for < 8GB GPUs.",
     "Use Concepts List": "Train multiple concepts from a JSON file or string.",
     "Use DREAM": "Enable DREAM (http://arxiv.org/abs/2312.00210). This may provide better results, but trains slower.",
+    "DREAM detail preservation": "A factor that influences how DREAM trades off composition versus detail. Low values will improve composition but may result in loss of detail. High values preserve detail but may reduce the overall effect of DREAM.",
     "Use EMA": "Enabling this will provide better results and editability, but cost more VRAM.",
     "Use EMA for prediction": "",
     "Use EMA Weights for Inference": "Enabling this will save the EMA unet weights as the 'normal' model weights and ignore the regular unet weights.",

diff --git a/scripts/main.py b/scripts/main.py
@@ -644,6 +644,13 @@ def on_ui_tabs():
                         db_use_dream = gr.Checkbox(
                             label="Use DREAM", value=False
                         )
+                        db_dream_detail_preservation = gr.Slider(
+                            label="DREAM detail preservation",
+                            minimum=0,
+                            maximum=1,
+                            step=0.01,
+                            visible=True,
+                        )
                         db_pad_tokens = gr.Checkbox(
                             label="Pad Tokens", value=True
                         )
@@ -1352,6 +1359,7 @@ def format_updates():
             db_tenc_grad_clip_norm,
             db_min_snr_gamma,
             db_use_dream,
+            db_dream_detail_preservation,
             db_pad_tokens,
             db_strict_tokens,
             db_max_token_length,
@@ -1486,6 +1494,7 @@ def toggle_advanced():
             db_max_token_length,
             db_min_snr_gamma,
             db_use_dream,
+            db_dream_detail_preservation,
             db_mixed_precision,
             db_model_name,
             db_model_path,

diff --git a/templates/defaults/defaults.json b/templates/defaults/defaults.json
@@ -155,6 +155,12 @@
   "use_dream": {
     "value": false
   },
+  "dream_detail": {
+    "value": 0.5,
+    "min": 0.0,
+    "max": 1,
+    "step": 0.01
+  },
   "max_token_length": {
     "value": 75,
     "min": 75,

diff --git a/templates/defaults/dreambooth_model_config.json b/templates/defaults/dreambooth_model_config.json
@@ -39,6 +39,7 @@
     "max_token_length": 75,
     "min_snr_gamma": 0.0,
     "use_dream": false,
+    "dream_detail": 0.5,
     "mixed_precision": "fp16",
     "noise_scheduler": "DDPM",
     "num_train_epochs": 200,

diff --git a/templates/locales/titles_en.json b/templates/locales/titles_en.json
@@ -193,6 +193,11 @@
     "title": "Enable Diffusion Rectification and Estimation-Adaptive Models (DREAM).",
     "description": "Whether or not to use DREAM. DREAM performs an additional model evaluation at each training step, which increases training time but can help improve the stability and consistency of the generated images."
   },
+  "dream_detail": {
+    "label": "DREAM detail",
+    "title": "Select how much detail DREAM preserves.",
+    "description": "A factor that influences how DREAM trades off composition versus detail. Low values will improve composition but may result in loss of detail. High values preserve detail but may reduce the overall effect of DREAM."
+  },
   "train_unet": {
     "label": "Train UNET",
     "title": "Train UNET as an additional module.",