diff --git a/train_configs/training_args_face_sdxl.json b/train_configs/training_args_face_sdxl.json index 18da4ae..a3d29f7 100644 --- a/train_configs/training_args_face_sdxl.json +++ b/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 diff --git a/train_configs/training_args_style_sdxl.json b/train_configs/training_args_style_sdxl.json index e05e739..2a9236c 100644 --- a/train_configs/training_args_style_sdxl.json +++ b/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 diff --git a/trainer/config.py b/trainer/config.py index 5174080..4f62048 100644 --- a/trainer/config.py +++ b/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 diff --git a/trainer/loss.py b/trainer/loss.py index e5c621f..c63cb70 100644 --- a/trainer/loss.py +++ b/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