Point scripts at checkpoints/ prefix

training/sample_job.py

@@ -46,6 +46,7 @@ BASE_MODEL = "runwayml/stable-diffusion-v1-5"
 
 # "base" renders vanilla SD 1.5 with no adapter - the comparison baseline.
 CHECKPOINT = os.environ.get("CHECKPOINT", "base")
+CHECKPOINT_DIR = "checkpoints"
 
 STEPS = int(os.environ.get("STEPS", 30))
 GUIDANCE = float(os.environ.get("GUIDANCE", 7.5))
@@ -84,7 +85,7 @@ def main() -> None:
         print(f"Loading LoRA {CHECKPOINT} ...")
         pipeline.load_lora_weights(
             MODEL_REPO, weight_name="pytorch_lora_weights.safetensors",
-            subfolder=CHECKPOINT,
+            subfolder=f"{CHECKPOINT_DIR}/{CHECKPOINT}",
         )
 
     images = []

training/train_lora_job.py

@@ -80,6 +80,9 @@ CAPTION_DROPOUT = float(os.environ.get("CAPTION_DROPOUT", 0.1))
 
 CHECKPOINT_EVERY = int(os.environ.get("CHECKPOINT_EVERY", 500))
 RESUME_FROM = os.environ.get("RESUME_FROM")                # e.g. "checkpoint-1500"
+
+# Checkpoints live under this prefix so the repo root stays readable.
+CHECKPOINT_DIR = "checkpoints"
 SEED = int(os.environ.get("SEED", 42))
 
 WORK_DIR = Path("run")
@@ -139,7 +142,7 @@ def save_checkpoint(unet, optimizer, scaler, step, epoch, api, scaling_factor):
     }, indent=2))
 
     api.upload_folder(
-        folder_path=str(out), path_in_repo=name,
+        folder_path=str(out), path_in_repo=f"{CHECKPOINT_DIR}/{name}",
         repo_id=MODEL_REPO, repo_type="model",
         commit_message=f"Checkpoint at step {step}",
     )
@@ -150,9 +153,10 @@ def save_checkpoint(unet, optimizer, scaler, step, epoch, api, scaling_factor):
 def load_checkpoint(unet, optimizer, scaler, device):
     print(f"Resuming from {RESUME_FROM} ...")
     local = snapshot_download(
-        MODEL_REPO, allow_patterns=f"{RESUME_FROM}/*", repo_type="model"
+        MODEL_REPO, allow_patterns=f"{CHECKPOINT_DIR}/{RESUME_FROM}/*",
+        repo_type="model",
     )
-    folder = Path(local) / RESUME_FROM
+    folder = Path(local) / CHECKPOINT_DIR / RESUME_FROM
 
     unet.load_lora_adapter(str(folder / "pytorch_lora_weights.safetensors"))
     blob = torch.load(folder / "optimizer.pt", map_location=device)