1 files changed, 7 insertions, 3 deletions

diff --git a/training/strategy/dreambooth.py b/training/strategy/dreambooth.py
index aeaa828..93c81cb 100644
--- a/training/strategy/dreambooth.py
+++ b/training/strategy/dreambooth.py
@@ -27,7 +27,8 @@ def dreambooth_strategy(
     sample_scheduler: DPMSolverMultistepScheduler,
     train_dataloader: DataLoader,
     val_dataloader: Optional[DataLoader],
-    output_dir: Path,
+    sample_output_dir: Path,
+    checkpoint_output_dir: Path,
     seed: int,
     train_text_encoder_epochs: int,
     max_grad_norm: float = 1.0,
@@ -47,6 +48,9 @@ def dreambooth_strategy(
             "Please set gradient_accumulation_steps to 1. This feature will be supported in the future."
         )
 
+    sample_output_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_output_dir.mkdir(parents=True, exist_ok=True)
+
     weight_dtype = torch.float32
     if accelerator.state.mixed_precision == "fp16":
         weight_dtype = torch.float16
@@ -64,7 +68,7 @@ def dreambooth_strategy(
         train_dataloader=train_dataloader,
         val_dataloader=val_dataloader,
         dtype=weight_dtype,
-        output_dir=output_dir,
+        output_dir=sample_output_dir,
         seed=seed,
         batch_size=sample_batch_size,
         num_batches=sample_num_batches,
@@ -156,7 +160,7 @@ def dreambooth_strategy(
                 tokenizer=tokenizer,
                 scheduler=sample_scheduler,
             )
-            pipeline.save_pretrained(output_dir.joinpath("model"))
+            pipeline.save_pretrained(checkpoint_output_dir)
 
         del unet_
         del text_encoder_