File size: 3,755 Bytes
1ab3690 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 | """Launches the training stages on Hugging Face Jobs.
python3 main.py baseline render vanilla SD 1.5 on the eval prompts
python3 main.py latents VAE-encode the dataset once
python3 main.py train LoRA fine-tune on the cached latents
python3 main.py sample checkpoint-2000
Everything runs on HF infrastructure; the M2 has no usable GPU. No -v volume
is passed - each job pulls what it needs from the Hub itself.
Order matters: `baseline` first (so there is something to compare against),
then `latents`, then `train`, then `sample` per checkpoint.
Per-run overrides go through the environment rather than flags here, since
each job script already reads its own constants that way:
EPOCHS=20 python3 main.py train
"""
import os
import subprocess
import sys
from pathlib import Path
HERE = Path(__file__).resolve().parent
# Job scripts are self-contained PEP-723 files. They ship to HF standalone
# and cannot import from this repo, so each carries its own constants and
# writes them into a manifest - no shared config module to drift out of sync.
STAGES = {
"latents": (HERE / "cache_latents_job.py", "a10g-small", "2h"),
"train": (HERE / "train_lora_job.py", "a10g-small", "6h"),
"sample": (HERE / "sample_job.py", "a10g-small", "1h"),
}
# Variables worth forwarding to the job if they are set locally.
FORWARDED = [
"RESOLUTION", "EPOCHS", "BATCH_SIZE", "GRAD_ACCUM", "LEARNING_RATE",
"LORA_RANK", "LORA_ALPHA", "CAPTION_DROPOUT", "CHECKPOINT_EVERY",
"RESUME_FROM", "SEED", "CHECKPOINT", "STEPS", "GUIDANCE",
"LATENTS_REVISION", "SOURCE_REVISION", "TARGET_REVISION", "MODEL_REPO",
]
DATASET_REPO = "whosouravsharma/text-to-image-diffusiondb-2M"
def require_latents() -> None:
"""Refuse to launch training before stage 1 has produced latents.
Checked here rather than only inside the job: a GPU job that dies on a
missing branch still costs a scheduling round-trip, and the failure shows
up as a 404 traceback rather than as the one line that explains it.
"""
from huggingface_hub import HfApi
revision = os.environ.get("LATENTS_REVISION", "latents-512")
branches = [
b.name for b in
HfApi().list_repo_refs(DATASET_REPO, repo_type="dataset").branches
]
if revision not in branches:
raise SystemExit(
f"No '{revision}' branch on {DATASET_REPO}.\n"
f"Existing branches: {', '.join(branches)}\n\n"
f"Cache the latents first:\n"
f" python3 main.py latents"
)
def run_stage(stage: str, checkpoint: str | None = None) -> int:
if stage == "baseline":
stage, checkpoint = "sample", "base"
if stage not in STAGES:
raise SystemExit(
f"Unknown stage {stage!r}. "
f"Choose from: baseline, {', '.join(STAGES)}"
)
if stage == "train":
require_latents()
script, flavor, timeout = STAGES[stage]
command = [
"hf", "jobs", "uv", "run",
"--flavor", flavor,
"--timeout", timeout,
"--secrets", "HF_TOKEN",
]
if checkpoint:
command += ["--env", f"CHECKPOINT={checkpoint}"]
for key in FORWARDED:
if key in os.environ and not (checkpoint and key == "CHECKPOINT"):
command += ["--env", f"{key}={os.environ[key]}"]
command.append(str(script))
print("=" * 60)
print(f"TRAINING STAGE: {stage}" + (f" ({checkpoint})" if checkpoint else ""))
print("=" * 60)
print("\nRunning:", " ".join(command), "\n")
return subprocess.call(command)
if __name__ == "__main__":
if len(sys.argv) < 2:
raise SystemExit(__doc__)
sys.exit(run_stage(sys.argv[1], sys.argv[2] if len(sys.argv) > 2 else None))
|