multimodalart HF staff commited on
Commit
6a97866
1 Parent(s): 31024de

Update training code

Browse files
Files changed (2) hide show
  1. .gitignore +2 -1
  2. train_dreambooth.py +2 -4
.gitignore CHANGED
@@ -3,4 +3,5 @@ output_model
3
  instance_images
4
  __pycache__
5
  intraining.lock
6
- hastrained.success
 
 
3
  instance_images
4
  __pycache__
5
  intraining.lock
6
+ hastrained.success
7
+ diffusers_model.tar
train_dreambooth.py CHANGED
@@ -572,9 +572,7 @@ def run_training(args_imported):
572
  eps=args.adam_epsilon,
573
  )
574
 
575
- noise_scheduler = DDPMScheduler(
576
- beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000
577
- )
578
 
579
  train_dataset = DreamBoothDataset(
580
  instance_data_root=args.instance_data_dir,
@@ -760,7 +758,7 @@ def run_training(args_imported):
760
  # Add the prior loss to the instance loss.
761
  loss = loss + args.prior_loss_weight * prior_loss
762
  else:
763
- loss = F.mse_loss(model_pred.float(), target.float(), reduction="none").mean([1, 2, 3]).mean()
764
 
765
  accelerator.backward(loss)
766
  if accelerator.sync_gradients:
 
572
  eps=args.adam_epsilon,
573
  )
574
 
575
+ noise_scheduler = DDPMScheduler.from_config(args.pretrained_model_name_or_path, subfolder="scheduler")
 
 
576
 
577
  train_dataset = DreamBoothDataset(
578
  instance_data_root=args.instance_data_dir,
 
758
  # Add the prior loss to the instance loss.
759
  loss = loss + args.prior_loss_weight * prior_loss
760
  else:
761
+ loss = F.mse_loss(model_pred.float(), target.float(), reduction="mean")
762
 
763
  accelerator.backward(loss)
764
  if accelerator.sync_gradients: