File size: 4,516 Bytes
dd1362a | 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 | # bora_finalize.py -- runs on the VM after training: pushes everything to the Hub model repo and verifies it.
# Waits for "TRAIN_EXIT=0" in /ephemeral/bora/train.log, uploads, then checks every remote file size against the local one.
# Writes FINALIZE_OK or FINALIZE_FAILED (with the reason) to /ephemeral/bora/finalize.log.
import sys
import time
from pathlib import Path
from huggingface_hub import HfApi, CommitOperationAdd
from huggingface_hub.hf_api import RepoFile
REPO = "Ami1ne/bora-v1"
RUN = Path("/ephemeral/bora_runs/v1")
OVF = Path("/ephemeral/bora_runs/overfit")
CODE = Path("/ephemeral/bora/code")
TRAIN_LOG = Path("/ephemeral/bora/train.log")
LOG = Path("/ephemeral/bora/finalize.log")
def log(msg):
line = f"[{time.strftime('%m-%d %H:%M:%S')}] {msg}"
print(line, flush=True)
with open(LOG, "a") as f:
f.write(line + "\n")
def main():
log("waiting for TRAIN_EXIT in train.log")
while True:
text = TRAIN_LOG.read_text() if TRAIN_LOG.exists() else ""
exits = [l for l in text.splitlines() if l.startswith("TRAIN_EXIT=")]
if exits:
break
time.sleep(60)
if exits[-1] != "TRAIN_EXIT=0" or "done at step 50000" not in text:
log(f"FINALIZE_FAILED: training did not finish cleanly ({exits[-1]}); nothing pushed")
sys.exit(1)
emas = sorted(RUN.glob("ema_step*.safetensors"), key=lambda p: int(p.stem[8:]))
files = {"ema_latest.safetensors": emas[-1]}
for p in emas:
files[f"checkpoints/{p.name}"] = p
files["checkpoints/ckpt_last.pt"] = RUN / "ckpt_last.pt" # model + EMA + optimizer, to resume training later
for n in ("config.json", "log.csv", "val.csv", "train.log"):
files[f"logs/{n}"] = RUN / n
for p in sorted((RUN / "samples").glob("*")):
files[f"samples/{p.name}"] = p
for p in sorted((OVF / "samples").glob("*.mp4")):
files[f"overfit/{p.name}"] = p
for n in ("log.csv", "train.log"):
files[f"overfit/{n}"] = OVF / n
for p in sorted(CODE.glob("*.py")):
files[f"code/{p.name}"] = p
readme = RUN / "README.md"
readme.write_text(f"""# Bora V1
Smallest video generation model from the @Codeically series. Flow-matching video DiT (336M parameters, dim 1024, depth 20)
over frozen Wan2.1-VAE latents. 256x256, 8 fps, blocks of 4 latent frames with 1 frame of overlap, 8 sampling steps.
Trained 50,000 steps at batch 32 on one A100-80G, on `Ami1ne/Bora` (Pexels / Pixabay / NASA clips).
Text encoder: google/t5-v1_1-base, 32 tokens.
- `ema_latest.safetensors`: EMA weights ({emas[-1].name}), bf16
- `checkpoints/ckpt_last.pt`: full training state to resume
- `code/`: bora_dit.py (model + sampler), bora_train.py, bora_dataset.py, bora_run.py
- `samples/`: sample videos every 2,000 steps; `logs/`: training and validation curves
Generate: `python code/bora_run.py sample --runs <dir with ema file> --prompts "ocean waves crashing"`.
Decode latents with the Wan2.1 VAE: `vae.decode(z * latents_std + latents_mean)`.
""")
files["README.md"] = readme
files = {k: v for k, v in files.items() if v.exists()}
total = sum(v.stat().st_size for v in files.values())
log(f"pushing {len(files)} files, {total / 1024**3:.2f} GB to {REPO}")
api = HfApi()
log(f"logged in as {api.whoami()['name']}")
api.create_repo(REPO, repo_type="model", exist_ok=True)
for attempt in range(1, 5):
try:
remote = {it.path: it.size for it in api.list_repo_tree(REPO, repo_type="model", recursive=True) if isinstance(it, RepoFile)}
ops = [CommitOperationAdd(path_in_repo=k, path_or_fileobj=str(v)) for k, v in files.items() if remote.get(k) != v.stat().st_size]
if ops:
api.create_commit(REPO, repo_type="model", operations=ops, commit_message=f"Bora V1 final: {emas[-1].stem}")
break
except Exception as e:
log(f"push attempt {attempt} failed: {type(e).__name__}: {str(e)[:200]}")
time.sleep(60 * attempt)
remote = {it.path: it.size for it in api.list_repo_tree(REPO, repo_type="model", recursive=True) if isinstance(it, RepoFile)}
bad = [k for k, v in files.items() if remote.get(k) != v.stat().st_size]
if bad:
log(f"FINALIZE_FAILED: {len(bad)} files missing or wrong size on the Hub: {bad[:10]}")
sys.exit(1)
log(f"FINALIZE_OK: {len(files)} files verified on https://huggingface.co/{REPO} ({total / 1024**3:.2f} GB)")
if __name__ == "__main__":
main()
|