"""HF Jobs backend for the LTX-2 trainer Space. Per training request: 1. stage the dataset (videos + dataset.json) + config.yaml + job_config.json locally, 2. sync them to a per-run HF bucket, 3. generate a self-contained UV job script, 4. submit it with `hf jobs uv run --flavor a100-large --secrets HF_TOKEN --detach`. On the Job, the script syncs the source bucket + run bucket, runs `uv sync --frozen` (reproducing the working trainer env from the lockfile), downloads the base checkpoint and Gemma, then runs process_dataset.py → train.py, which pushes the trained LoRA to the Hub. For IC-LoRA, references are user-supplied (paired `*_reference` videos) — no auto-derivation. The Space itself only needs gradio + huggingface_hub + pyyaml (no torch). """ from __future__ import annotations import json import os import re import shutil import subprocess import tempfile import zipfile from pathlib import Path import yaml SRC_BUCKET = os.environ.get("LTX_SRC_BUCKET", "linoyts/ltx2-trainer-src") DEFAULT_FLAVOR = "a100-large" # the largest single-GPU flavor `hf jobs uv run` accepts (80GB) VIDEO_EXTS = {".mp4", ".mov", ".mkv", ".webm", ".avi", ".m4v"} # Deterministic on-Job paths (the Space bakes these into config.yaml). JOB_ROOT = "/tmp/ltxjob" JOB_MODEL = f"{JOB_ROOT}/models/ltx-2.3-22b-dev.safetensors" JOB_GEMMA = f"{JOB_ROOT}/gemma" JOB_RUN = f"{JOB_ROOT}/run" MODES = { "IC-LoRA (in-context control)": { "needs_reference": True, "target_modules": [ "attn1.to_k", "attn1.to_q", "attn1.to_v", "attn1.to_out.0", "attn2.to_k", "attn2.to_q", "attn2.to_v", "attn2.to_out.0", "ff.net.0.proj", "ff.net.2", ], }, "Text-to-Video LoRA": {"needs_reference": False, "first_frame_prob": 0.0, "target_modules": ["to_k", "to_q", "to_v", "to_out.0"]}, "Image-to-Video LoRA": {"needs_reference": False, "first_frame_prob": 0.5, "target_modules": ["to_k", "to_q", "to_v", "to_out.0"]}, } def parse_resolution(resolution: str) -> tuple[int, int, int]: parts = resolution.lower().replace(" ", "").split("x") if len(parts) != 3: raise ValueError(f"Resolution must be 'WxHxF', got {resolution!r}") w, h, f = (int(p) for p in parts) if w % 32 or h % 32: raise ValueError(f"Width and height must be divisible by 32 (got {w}x{h}).") if f % 8 != 1: raise ValueError(f"Frame count must satisfy frames % 8 == 1 (got {f}).") return w, h, f def _collect_videos(uploaded: list[str], dest: Path) -> list[Path]: dest.mkdir(parents=True, exist_ok=True) for p in uploaded or []: src = Path(p) if src.suffix.lower() == ".zip": with zipfile.ZipFile(src) as zf: for m in zf.namelist(): if Path(m).suffix.lower() in VIDEO_EXTS and not m.startswith("__MACOSX"): with zf.open(m) as s, open(dest / Path(m).name, "wb") as d: shutil.copyfileobj(s, d) elif src.suffix.lower() in VIDEO_EXTS: shutil.copy2(src, dest / src.name) return sorted(p for p in dest.glob("*") if p.suffix.lower() in VIDEO_EXTS) def _is_reference(p: Path) -> bool: return p.stem.endswith("_reference") def build_dataset_items( videos: list[Path], captions: list[str], needs_reference: bool ) -> tuple[list[dict], list[Path]]: """Build dataset.json rows. For IC-LoRA, pair each target `X.ext` with `X_reference.ext` (user-supplied — no auto-derivation). Captions align to the sorted target clips. Returns (items, targets). Raises ValueError on missing references or no targets. """ vids = [v for v in videos if v.suffix.lower() in VIDEO_EXTS] if needs_reference: targets = sorted(v for v in vids if not _is_reference(v)) refs = {v.stem[: -len("_reference")]: v for v in vids if _is_reference(v)} else: targets, refs = sorted(vids), {} items, missing = [], [] for i, v in enumerate(targets): cap = captions[i] if i < len(captions) and captions[i].strip() else "a video" row = {"media_path": f"videos/{v.name}", "caption": cap} if needs_reference: ref = refs.get(v.stem) if ref is None: missing.append(v.name) continue row["reference_video"] = f"videos/{ref.name}" items.append(row) if needs_reference and missing: raise ValueError( "Missing reference video(s) for: " + ", ".join(missing) + ". For IC-LoRA, every target `X.mp4` needs a paired `X_reference.mp4`." ) if not items: raise ValueError("No target videos found in the upload.") return items, targets def build_config_dict(params: dict, videos: list[Path]) -> dict: w, h, f = parse_resolution(params["resolution"]) mode_cfg = MODES[params["mode"]] conditions: list[dict] = [] if mode_cfg["needs_reference"]: conditions.append({"type": "reference", "latents_dir": "reference_latents", "probability": 1.0}) conditions.append({"type": "first_frame", "probability": 0.2}) elif mode_cfg.get("first_frame_prob", 0.0) > 0: conditions.append({"type": "first_frame", "probability": mode_cfg["first_frame_prob"]}) val_sample: dict = {"prompt": params["captions"][0] if params.get("captions") else "a video"} if mode_cfg["needs_reference"] and videos: stem = videos[0].stem ref_name = f"{stem}_reference{videos[0].suffix}" val_sample["conditions"] = [ {"type": "reference", "video": f"{JOB_RUN}/dataset/videos/{ref_name}", "include_in_output": True} ] return { "model": {"model_path": JOB_MODEL, "text_encoder_path": JOB_GEMMA, "training_mode": "lora", "load_checkpoint": None}, "lora": {"rank": int(params["rank"]), "alpha": int(params["alpha"]), "dropout": 0.0, "target_modules": mode_cfg["target_modules"]}, "training_strategy": {"name": "flexible", "video": {"is_generated": True, "latents_dir": "latents", "conditions": conditions}}, "optimization": {"learning_rate": float(params["learning_rate"]), "steps": int(params["steps"]), "batch_size": int(params["batch_size"]), "gradient_accumulation_steps": int(params["gradient_accumulation_steps"]), "max_grad_norm": 1.0, "optimizer_type": params["optimizer_type"], "scheduler_type": "linear", "scheduler_params": {}, "enable_gradient_checkpointing": True}, "acceleration": {"mixed_precision_mode": "bf16", "quantization": params["quantization"] or None, "load_text_encoder_in_8bit": bool(params["load_text_encoder_in_8bit"]), "offload_optimizer_during_validation": True}, "data": {"preprocessed_data_root": f"{JOB_RUN}/dataset/.precomputed", "num_dataloader_workers": 2}, "validation": {"samples": [val_sample], "negative_prompt": "worst quality, inconsistent motion, blurry, jittery, distorted", "video_dims": [w, h, f], "frame_rate": 25.0, "seed": int(params["seed"]), "inference_steps": 16, "interval": max(int(params["steps"]), 1), "guidance_scale": 4.0, "stg_scale": 1.0, "stg_blocks": [29], "stg_mode": "stg_v", "generate_audio": False, "skip_initial_validation": True}, # keep_last_n: -1 avoids a trainer bug — when steps is a multiple of the checkpoint # interval, the final step is saved twice and cleanup (keep_last_n>0) deletes the file # before push_to_hub. Keeping all checkpoints (tiny for LoRA) sidesteps it. "checkpoints": {"interval": max(int(params["steps"]), 1), "keep_last_n": -1, "precision": "bfloat16"}, "flow_matching": {"timestep_sampling_mode": "shifted_logit_normal", "timestep_sampling_params": {}}, "hub": {"push_to_hub": bool(params["push_to_hub"]), "hub_model_id": params["hub_model_id"] or None}, "wandb": {"enabled": False, "project": "ltx-2-trainer", "entity": None, "tags": ["ltx2", "jobs"], "log_validation_videos": True}, "seed": int(params["seed"]), "output_dir": f"{JOB_RUN}/outputs", } # -------------------------------------------------------------------------------------- # UV job script (runs on HF Jobs hardware) # -------------------------------------------------------------------------------------- JOB_SCRIPT_TEMPLATE = '''# /// script # requires-python = ">=3.10" # dependencies = ["huggingface_hub[hf-xet]>=1.5", "hf_transfer"] # /// """Auto-generated LTX-2 training job. Reproduces the trainer env via uv sync --frozen.""" import json, os, subprocess, sys from pathlib import Path os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1" JOB = Path("{JOB_ROOT}"); SRC = JOB / "src"; RUN = JOB / "run" MODEL = "{JOB_MODEL}"; GEMMA = "{JOB_GEMMA}" SRC_BUCKET = "{SRC_BUCKET}"; RUN_BUCKET = "{RUN_BUCKET}" def sh(cmd, cwd=None): print(">>>", " ".join(cmd), flush=True) subprocess.run(cmd, cwd=str(cwd) if cwd else None, check=True) def main(): JOB.mkdir(parents=True, exist_ok=True) print("=== 0/6 system libs (opencv needs libGL) ===", flush=True) subprocess.run("apt-get update && apt-get install -y --no-install-recommends libgl1 libglib2.0-0", shell=True, check=False) print("=== 1/6 sync source bucket ===", flush=True) sh(["hf", "buckets", "sync", f"hf://buckets/{{SRC_BUCKET}}", str(SRC)]) print("=== 2/6 sync run bucket (dataset + config) ===", flush=True) sh(["hf", "buckets", "sync", f"hf://buckets/{{RUN_BUCKET}}", str(RUN)]) jc = json.loads((RUN / "job_config.json").read_text()) print("=== 3/6 uv sync (reproduce trainer env) ===", flush=True) sh(["uv", "sync", "--frozen"], cwd=SRC) print("=== 4/6 download base checkpoint + Gemma ===", flush=True) from huggingface_hub import hf_hub_download, snapshot_download hf_hub_download("Lightricks/LTX-2.3", "ltx-2.3-22b-dev.safetensors", local_dir=str(JOB / "models")) snapshot_download("google/gemma-3-12b-it-qat-q4_0-unquantized", local_dir=GEMMA) tr = SRC / "packages" / "ltx-trainer" def uvrun(args): sh(["uv", "run", "python", *args], cwd=tr) ds_json = RUN / "dataset" / "dataset.json" print("=== 5/6 preprocess dataset ===", flush=True) uvrun(["scripts/process_dataset.py", str(ds_json), "--resolution-buckets", jc["resolution"], "--model-path", MODEL, "--text-encoder-path", GEMMA, "--skip-audio"]) print("=== 6/6 train (pushes LoRA to the Hub) ===", flush=True) uvrun(["scripts/train.py", str(RUN / "config.yaml"), "--disable-progress-bars"]) print("=== DONE ===", flush=True) if __name__ == "__main__": main() ''' def _run_bucket_name(run_name: str) -> str: safe = re.sub(r"[^a-zA-Z0-9-]+", "-", run_name).strip("-").lower() or "run" return f"ltx2-train-{safe}" def _namespace(token: str | None) -> str: from huggingface_hub import whoami # noqa: PLC0415 return whoami(token=token or os.environ.get("HF_TOKEN"))["name"] def submit(params: dict, uploaded_videos: list[str], flavor: str, timeout: str) -> dict: """Stage data → bucket → generate UV script → submit job. Returns {job_id, url, bucket, log}.""" token = (params.get("hf_token") or os.environ.get("HF_TOKEN") or "").strip() ns = _namespace(token) bucket = f"{ns}/{_run_bucket_name(params['run_name'])}" tmp = Path(tempfile.mkdtemp(prefix="ltxrun-")) try: videos = _collect_videos(uploaded_videos, tmp / "dataset" / "videos") if not videos: raise ValueError("No valid video files in the upload.") # dataset.json — pairs targets with user-supplied references for IC-LoRA needs_reference = MODES[params["mode"]]["needs_reference"] items, targets = build_dataset_items(videos, params.get("captions", []), needs_reference) (tmp / "dataset" / "dataset.json").write_text(json.dumps(items, indent=2)) # config.yaml + job_config.json (validation sample references the first target's reference) cfg = build_config_dict(params, targets) (tmp / "config.yaml").write_text(yaml.safe_dump(cfg, sort_keys=False)) # job script reads job_config.json + config.yaml at the run root (sibling of dataset/) (tmp / "job_config.json").write_text(json.dumps({"resolution": params["resolution"]}, indent=2)) env = os.environ.copy() if token: env["HF_TOKEN"] = token # create + sync the per-run bucket subprocess.run(["hf", "buckets", "create", bucket.split("/", 1)[1]], env=env, capture_output=True, text=True) # ok if exists sync = subprocess.run(["hf", "buckets", "sync", str(tmp), f"hf://buckets/{bucket}"], env=env, capture_output=True, text=True) if sync.returncode != 0: raise RuntimeError(f"bucket sync failed:\n{sync.stdout}\n{sync.stderr}") # render + write the job script script = JOB_SCRIPT_TEMPLATE.format( JOB_ROOT=JOB_ROOT, JOB_MODEL=JOB_MODEL, JOB_GEMMA=JOB_GEMMA, SRC_BUCKET=SRC_BUCKET, RUN_BUCKET=bucket, ) script_path = tmp / "job_train.py" script_path.write_text(script) cmd = ["hf", "jobs", "uv", "run", "--flavor", flavor, "--timeout", timeout, "--secrets", "HF_TOKEN", "--detach", str(script_path)] res = subprocess.run(cmd, env=env, capture_output=True, text=True) out = (res.stdout or "") + (res.stderr or "") if res.returncode != 0: raise RuntimeError(f"job submission failed:\n{out}") job_id = _parse_job_id(out) url = f"https://huggingface.co/jobs/{ns}/{job_id}" if job_id else "" return {"job_id": job_id or "", "url": url, "bucket": bucket, "log": out.strip()} finally: shutil.rmtree(tmp, ignore_errors=True) def _parse_job_id(text: str) -> str | None: m = re.search(r"\b([0-9a-f]{24,})\b", text) # job ids are long hex if m: return m.group(1) m = re.search(r"jobs/[^/]+/(\S+)", text) return m.group(1) if m else None def job_logs(job_id: str, token: str = "") -> str: env = os.environ.copy() if token: env["HF_TOKEN"] = token res = subprocess.run(["hf", "jobs", "logs", job_id], env=env, capture_output=True, text=True) return (res.stdout or "") + (res.stderr or "") def job_status(job_id: str, token: str = "") -> str: env = os.environ.copy() if token: env["HF_TOKEN"] = token res = subprocess.run(["hf", "jobs", "inspect", job_id], env=env, capture_output=True, text=True) try: data = json.loads(res.stdout) if isinstance(data, list) and data: data = data[0] return str(data.get("status", {}).get("stage", "UNKNOWN")) except Exception: # noqa: BLE001 return (res.stdout or res.stderr or "UNKNOWN").strip()[:200]