update defaults

train_configs/training_args_face_sdxl.json

@@ -1,21 +1,20 @@
 {
     "name": "xander_sdxl",
     "sd_model_version": "sdxl",
-    "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/xander.zip",
+    "lora_training_urls": "https://storage.googleapis.com/public-assets-xander/A_workbox/lora_training_sets/mira.zip",
     "concept_mode": "face",
     "sample_imgs_lora_scale": 0.7,
-    "seed": 3,
+    "seed": 0,
     "resolution": 512,
     "train_batch_size": 4,
     "n_sample_imgs": 8,
     "max_train_steps": 300,
-    "token_warmup_steps": 0,
     "checkpointing_steps": 200,
 
     "disable_ti": false,
     
     "ti_lr": 0.001,
-    "unet_lr": 0.0005,
+    "unet_lr": 0.0003,
     "lora_rank": 16,
 
     "debug": true

train_configs/training_args_style_sdxl.json

@@ -3,19 +3,18 @@
     "sd_model_version": "sdxl",
     "lora_training_urls": "https://edenartlab-lfs.s3.amazonaws.com/datasets/clipx.zip",
     "concept_mode": "style",
-    "sample_imgs_lora_scale": 0.8,
+    "sample_imgs_lora_scale": 0.75,
     "seed": 0,
     "resolution": 512,
     "train_batch_size": 4,
-    "n_sample_imgs": 6,
+    "n_sample_imgs": 8,
     "max_train_steps": 300,
-    "token_warmup_steps": 0,
-    "checkpointing_steps": 150,
-
-    "n_tokens": 2,
+    "checkpointing_steps": 200,
+    
+    "disable_ti": false,
 
     "ti_lr": 0.001,
-    "unet_lr": 0.001,
+    "unet_lr": 0.0003,
     "lora_rank": 16,
 
     "debug": true

trainer/config.py

@@ -48,7 +48,7 @@ class TrainingConfig(BaseModel):
     train_img_size: List[int] = None
     train_aspect_ratio: float = None
     train_batch_size: int = 4
-    max_train_steps: int = 360
+    max_train_steps: int = 300
     num_train_epochs: int = None
     checkpointing_steps: int = 10000
     gradient_accumulation_steps: int = 1
@@ -56,7 +56,7 @@ class TrainingConfig(BaseModel):
 
     unet_optimizer_type: Literal["adamw", "prodigy", "AdamW8bit"] = "adamw"
     unet_lr_warmup_steps: int = None  # slowly increase the learning rate of the adamw unet optimizer
-    unet_lr: float = 0.0005
+    unet_lr: float = 0.0003
     prodigy_d_coef: float = 1.0
     unet_prodigy_growth_factor: float = 1.05  # lower values make the lr go up slower (1.01 is for 1k step runs, 1.02 is for 500 step runs)
     lora_weight_decay: float = 0.002
@@ -66,7 +66,7 @@ class TrainingConfig(BaseModel):
     ti_weight_decay: float = 0.0
     ti_optimizer: Literal["adamw", "prodigy"] = "adamw"
     freeze_ti_after_completion_f: float = 0.6     # freeze the TI after this fraction of the training is done
-    freeze_unet_before_completion_f: float = 0.3  # freeze the UNET before this fraction of the training is done
+    freeze_unet_before_completion_f: float = 0.0  # freeze the UNET before this fraction of the training is done
     
     token_attention_loss_w: float = 3e-7
     cond_reg_w: float = 0.0e-5

trainer/loss.py

@@ -193,7 +193,7 @@ class ConditioningRegularizer:
             self.distribution_regularizers[f'txt_encoder_{idx}'] = DistributionLoss(pretrained_token_embeddings, outdir = self.config.output_dir if config.debug else None)
             idx += 1
 
-    def apply_regularization(self, loss, losses, prompt_embeds_norms, prompt_embeds, std_loss_w = 0.003, pipe=None):
+    def apply_regularization(self, loss, losses, prompt_embeds_norms, prompt_embeds, std_loss_w = 0.01, pipe=None):
         noise_sigma = 0.0
         if noise_sigma > 0.0: # experimental: apply random noise to the conditioning vectors as a form of regularization
             prompt_embeds[0,1:-2,:] += torch.randn_like(prompt_embeds[0,2:-2,:]) * noise_sigma