Join paths, update default scheduler

Fix "subfolder not recognized" issue.
d8ahazard · Nov 10, 2022 · 75e40bb · 75e40bb
1 parent 4d9b93d
commit 75e40bb
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 7 deletions.
diff --git a/dreambooth/train_dreambooth.py b/dreambooth/train_dreambooth.py
@@ -807,15 +807,14 @@ def save_weights(step, save_model, save_img):
             if args.train_text_encoder:
                 text_enc_model = accelerator.unwrap_model(text_encoder)
             else:
-                text_enc_model = CLIPTextModel.from_pretrained(args.working_dir, subfolder="text_encoder")
+                text_enc_model = CLIPTextModel.from_pretrained(os.path.join(args.working_dir, "text_encoder"))
             scheduler = DDIMScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", clip_sample=False, set_alpha_to_one=False)
             pipeline = StableDiffusionPipeline.from_pretrained(
                 args.working_dir,
                 unet=accelerator.unwrap_model(unet),
                 text_encoder=text_enc_model,
                 vae=AutoencoderKL.from_pretrained(
-                    args.working_dir,
-                    subfolder="vae"
+                    os.path.join(args.working_dir, "vae")
                 ),
                 safety_checker=None,
                 scheduler=scheduler,
@@ -862,7 +861,6 @@ def save_weights(step, save_model, save_img):
     shared.state.job_no = global_step
     shared.state.textinfo = f"Training step: {global_step}/{args.max_train_steps}"
     loss_avg = AverageMeter()
-    training_complete = False
     text_enc_context = nullcontext() if args.train_text_encoder else torch.no_grad()
     for epoch in range(args.num_train_epochs):
         try:

diff --git a/scripts/main.py b/scripts/main.py
@@ -37,9 +37,7 @@ def on_ui_tabs():
                     db_new_model_name = gr.Textbox(label="Name")
                     src_checkpoint = gr.Dropdown(label='Source Checkpoint', choices=sorted(
                         sd_models.checkpoints_list.keys()))
-                    # I just randomly chose ddim here because we use it everywhere else. Not sure which of these
-                    # are ideal, or if it matters at all.
-                    diff_type = gr.Dropdown(label='Scheduler', choices=["pndm", "ddim", "lms"], value="pndm")
+                    diff_type = gr.Dropdown(label='Scheduler', choices=["ddim", "pndm", "lms"], value="ddim")
 
                     with gr.Row():
                         with gr.Column(scale=3):