| |
| |
| |
| 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" |
| 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() |
|
|
