scripts/fetch_image_model.py

"""Fetch the local FLUX.2 Klein 4B image model snapshot.

    python scripts/fetch_image_model.py            # download / repair
    python scripts/fetch_image_model.py --verify   # verify only
    python scripts/fetch_image_model.py --load     # verify + really load it
    python scripts/fetch_image_model.py --force    # re-fetch suspect files
    python scripts/fetch_image_model.py --force-host  # stage for another PC

`--verify` is the cheap structural check (every published file present at its
published size, every text-encoder shard named by the weight index). `--load`
adds the decisive one: it builds `Flux2KleinPipeline` from the finished folder
exactly as the service does on an image job, so "the download finished" and
"the model loads" are never confused. Neither mode writes anything.

The default repository is `black-forest-labs/FLUX.2-klein-4B`; its Diffusers
components are written to `LOCAL_IMAGE_MODEL_DIR` (default
`models/flux2-klein-4b`). Generation is then local and offline. The fetch uses
the resumable range-transfer engine shared with the video model downloader.
The model card specifies an NVIDIA CUDA GPU with about 13 GB VRAM. `--force-host`
skips the local hardware check only for staging the snapshot elsewhere.

The snapshot includes a sharded Qwen3 text encoder, transformer, tokenizer,
scheduler, and VAE. The root-level single-file checkpoint and example images
are not needed. FLUX.2 Klein supports multi-reference image editing; native
reference conditioning is not a hard identity lock and does not expose a
mask/inpainting input.
"""

from __future__ import annotations

import os
import subprocess
import sys
import threading
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
sys.path.insert(0, str(Path(__file__).resolve().parent))

# Reuse the video model's transfer engine verbatim.
import fetch_i2v_model as transfer  # noqa: E402

# Exact files required by Flux2KleinPipeline. The root checkpoint and sample
# images are excluded; only the Diffusers components are fetched.
ALLOW_PATTERNS = [
    "model_index.json",
    "scheduler/scheduler_config.json",
    "tokenizer/*",
    "text_encoder/config.json",
    "text_encoder/generation_config.json",
    "text_encoder/model.safetensors.index.json",
    "text_encoder/model-00001-of-00002.safetensors",
    "text_encoder/model-00002-of-00002.safetensors",
    "transformer/config.json",
    "transformer/diffusion_pytorch_model.safetensors",
    "vae/config.json",
    "vae/diffusion_pytorch_model.safetensors",
]

# The allow-list excludes unrelated files; this also excludes the root-level
# checkpoint if the allow-list is broadened later.
IGNORE_PATTERNS = [
    "*.md",
    "*.png",
    "*.jpg",
    "*.jpeg",
    "*.gif",
    "*.workflow.json",
    "*.ckpt",
    "*.yaml",
    "flux-2-klein-4b.safetensors",
]

# Exact files that must exist AND be non-empty for the pipeline to load.
REQUIRED_FILES = [
    "model_index.json",
    "scheduler/scheduler_config.json",
    "tokenizer/added_tokens.json",
    "tokenizer/chat_template.jinja",
    "tokenizer/merges.txt",
    "tokenizer/special_tokens_map.json",
    "tokenizer/tokenizer.json",
    "tokenizer/tokenizer_config.json",
    "tokenizer/vocab.json",
    "text_encoder/config.json",
    "text_encoder/generation_config.json",
    "text_encoder/model.safetensors.index.json",
    "text_encoder/model-00001-of-00002.safetensors",
    "text_encoder/model-00002-of-00002.safetensors",
    "transformer/config.json",
    "transformer/diffusion_pytorch_model.safetensors",
    "vae/config.json",
    "vae/diffusion_pytorch_model.safetensors",
]

# The text encoder's index enumerates its required shards.
SHARD_MANIFESTS = ["text_encoder/model.safetensors.index.json"]

DEFAULT_REPO = "black-forest-labs/FLUX.2-klein-4B"
DEFAULT_DIR = "models/flux2-klein-4b"

# `python-ai/` itself - the folder holding `app/` and `models/`.
#
# A relative LOCAL_IMAGE_MODEL_DIR is anchored HERE and never at the process
# working directory. `Path.cwd()` made the target depend on where the script was
# launched from: from the repository root it wrote `DevAtlas/models/flux2-klein-4b`
# while the service looked in `DevAtlas/python-ai/models/flux2-klein-4b`, so a
# finished 16 GB download still produced
# "local FLUX.2 Klein model directory 'models/flux2-klein-4b' does not exist".
# The project root is also what the child worker below runs in.
PROJECT_ROOT = Path(__file__).resolve().parent.parent


def resolve_repo() -> str:
    from app.config.settings import get_settings

    return get_settings().local_image_model_id or DEFAULT_REPO


# ---------------------------------------------------------------------------
# PRE-FLIGHT
# ---------------------------------------------------------------------------
#
# Hub access and hardware are checked before the multi-gigabyte transfer starts:
#
#   1. HUB ACCESS. The manifest read above succeeds even when the file requests
#      will not (proxies, a rate limit, a mirror that only serves metadata), so
#      the first real file request is where a broken run used to die - with a
#      bare `HTTP 401`, a partially-allocated `.part` file left behind, and
#      nothing said about which of the two it was. Probing the smallest real
#      file first turns that into an instant, actionable message.
#
#   2. HARDWARE. Disk is checked by the shared engine; CUDA availability and
#      VRAM are checked here against the model card's approximate 13 GB need.


def _hub_token() -> str:
    """The Hugging Face token, from the env or the standard cache file.

    `huggingface_hub` accepts the same two sources this checks, so anything
    found here is genuinely usable by the downloader below.
    """
    for name in ("HF_TOKEN", "HUGGING_FACE_HUB_TOKEN"):
        value = (os.environ.get(name) or "").strip()
        if value:
            return value
    try:
        from huggingface_hub import constants as hf_constants

        token_path = getattr(hf_constants, "HF_TOKEN_PATH", None)
        if token_path and Path(token_path).is_file():
            return Path(token_path).read_text(encoding="utf-8").strip()
    except Exception:
        pass
    return ""


def dedupe(problems: list) -> list:
    """The same problem once, in the order it was first seen.

    The required-file pass and the Hub-manifest pass overlap, so an incomplete
    snapshot otherwise prints every missing file twice and buries the
    actionable list - the report has to name each missing file exactly once.
    """
    seen: set = set()
    unique: list = []
    for problem in problems:
        if problem not in seen:
            seen.add(problem)
            unique.append(problem)
    return unique


def reject(path: Path, problems: list) -> int:
    """Print an incomplete-snapshot report and return a failure exit code."""
    print(f"Snapshot is INCOMPLETE at {path}:")
    for problem in dedupe(problems):
        print(f"  - {problem}")
    print("\nFinish it with: python scripts/fetch_image_model.py")
    return 1


def resolve_target() -> Path:
    """The exact folder the snapshot is written to and read back from.

    Relative values (the documented `models/flux2-klein-4b`) are resolved
    against `python-ai/`, so the answer is the same absolute path the service
    uses no matter which directory this script or the server was started from.
    """
    from app.config.settings import get_settings

    target = Path(get_settings().local_image_model_dir or DEFAULT_DIR)
    if not target.is_absolute():
        target = PROJECT_ROOT / target
    return Path(os.path.normpath(str(target)))


def check_hub_access(repo: str, revision: str) -> list:
    """Problems that would make every file request fail. Cheap by design.

    Fetches the smallest file in the snapshot (`model_index.json`, a few
    hundred bytes) rather than a weight file. A `200` here means the Hub is
    reachable and the transfer can start.
    """
    import requests

    url = f"https://huggingface.co/{repo}/resolve/{revision}/model_index.json"
    headers = {"User-Agent": "devatlas-fetch-image/1.0"}
    token = _hub_token()
    if token:
        headers["Authorization"] = f"Bearer {token}"

    try:
        response = requests.get(url, headers=headers, timeout=30)
    except Exception as exc:
        return [
            f"could not reach the Hugging Face Hub: {exc}\n"
            "    Check the network/proxy, then re-run."
        ]

    if response.status_code in (401, 403):
        detail = (
            f"the access token is missing, or has not been granted access "
            f"(HTTP {response.status_code})."
            if token
            else f"no access token was found (HTTP {response.status_code})."
        )
        return [
            f"{repo} refused the request, so this download cannot start: {detail}\n"
            f"    1. Check the repo page is reachable without a token:\n"
            f"       https://huggingface.co/{repo}\n"
            "    2. If it now requires one, create a read token at\n"
            "       https://huggingface.co/settings/tokens\n"
            "    3. Export it, then re-run this script:\n"
            "         PowerShell: $env:HF_TOKEN = \"hf_...\"\n"
            "         bash:       export HF_TOKEN=hf_...\n"
            "       (or run `huggingface-cli login` once, which caches it)\n"
            "    Nothing was downloaded and no bytes were wasted."
        ]
    if response.status_code != 200:
        return [
            f"the Hub returned HTTP {response.status_code} for "
            f"{repo}/model_index.json; cannot start the transfer."
        ]
    return []


def check_hardware() -> list:
    """Check the model card's approximate CUDA VRAM requirement before fetching."""
    try:
        import torch
    except ImportError:
        return ["PyTorch is missing; install the CUDA-enabled PyTorch build first."]

    if not torch.cuda.is_available():
        return [
            "FLUX.2 Klein requires an NVIDIA CUDA GPU with about 13 GB VRAM; "
            "this machine's PyTorch has no CUDA device. Use --force-host only "
            "when staging the model for a different machine."
        ]

    device = torch.device("cuda:0")
    total_bytes = int(torch.cuda.get_device_properties(device).total_memory)
    free_bytes, _ = torch.cuda.mem_get_info(device)
    print(
        f"  Device       : cuda:{torch.cuda.get_device_name(device)} - "
        f"{free_bytes / 1024**3:.1f} GB VRAM free of {total_bytes / 1024**3:.1f} GB"
    )
    if total_bytes < 13 * 1024**3:
        return [
            f"FLUX.2 Klein needs about 13 GB VRAM; this GPU has "
            f"{total_bytes / 1024**3:.1f} GB total."
        ]
    return []


def load_model(target: Path) -> int:
    """Build the real pipeline from the finished snapshot. Renders nothing.

    The provider is constructed exactly the way `deps.py` constructs it at
    runtime, so a PASS here means the service's own load path works: the
    16.0 GB of weights are mapped by Diffusers, the dtype/offload setup is
    applied, and `Flux2KleinPipeline` comes up. A FAIL prints the reason
    Diffusers or the hardware gave, unchanged.
    """
    from app.config.settings import get_settings

    from app.services.providers.image_provider import LocalDiffusionImageProvider

    settings = get_settings()
    provider = LocalDiffusionImageProvider(
        str(target),
        steps=settings.local_image_steps,
        strength=settings.local_image_strength,
        guidance_scale=settings.local_image_guidance,
        size=settings.local_image_size,
        negative_prompt=settings.local_image_negative_prompt,
        threads=settings.local_image_threads,
        use_fp16=settings.local_image_fp16,
    )
    print(f"Loading FLUX.2 Klein from: {target}")
    started = time.perf_counter()
    try:
        pipeline = provider.verify_load()
    except Exception as exc:
        print(f"\nLOAD FAILED after {time.perf_counter() - started:.1f}s:")
        print(f"  {exc}")
        print(
            "\nImage jobs are refused until this load succeeds. There is no "
            "smaller-model fallback."
        )
        return 1
    print(
        f"\nLOAD OK: {type(pipeline).__name__} is loaded and ready "
        f"({time.perf_counter() - started:.1f}s)."
    )
    return 0


def reclaim_unusable_partials(target: Path, manifest: dict) -> int:
    """Delete `.part` scratch files that cannot be resumed. Returns bytes freed.

    A `.part` is only resumable when its `<name>.parts.json` ledger sits next to
    it AND its length already equals the published size. Without that pairing
    the engine cannot tell which byte ranges actually landed, so it truncates
    the file to the published size and re-downloads it from zero - the 1.9 GB
    on disk buys nothing.

    This is what a killed run leaves behind: a pre-allocated `.part` whose
    ledger was never flushed. Reclaiming it stops a dead scratch file from
    sitting on the volume indefinitely, and it is strictly safe because the
    bytes are unrecoverable by definition. Files with a valid ledger are left
    alone - that ledger is the whole point of the resumable downloader.
    """
    freed = 0
    files = manifest.get("files", {})
    for part in sorted(target.rglob("*.part")):
        ledger = part.with_name(part.name[: -len(".part")] + ".parts.json")
        if ledger.is_file():
            continue  # resumable state - never touch it
        rel = part.relative_to(target).as_posix()[: -len(".part")]
        published = files.get(rel)
        # A correctly-sized, ledger-less part is a completed-but-unrenamed
        # file. Leave it: the next run re-verifies it by hash and promotes it.
        if published is not None and part.stat().st_size == published:
            continue
        try:
            freed += part.stat().st_size
            part.unlink()
        except OSError:
            continue
    return freed


def _worker(repo: str, target: str, stall_seconds: int, force: bool = False) -> int:
    """One download attempt in a CHILD process, with the stall watchdog."""
    root = Path(target)
    root.mkdir(parents=True, exist_ok=True)
    stall = {"last": time.time(), "bytes": -1}

    def watch() -> None:
        while True:
            time.sleep(10)
            done = transfer._PROGRESS[0]
            if done != stall["bytes"]:
                stall["bytes"] = done
                stall["last"] = time.time()
                print(f"    {done / 1e9:.2f} GB transferred", flush=True)
            elif time.time() - stall["last"] > stall_seconds:
                print(f"    no transfer progress for {stall_seconds}s - wedged", flush=True)
                os._exit(3)

    threading.Thread(target=watch, daemon=True).start()

    try:
        manifest = transfer.fetch_manifest(repo, ALLOW_PATTERNS, IGNORE_PATTERNS)
    except Exception as exc:
        print(f"    could not read the Hub manifest: {exc}", flush=True)
        return 2

    reclaimed = transfer._clean_partial_downloads(root, keep_parts=True)
    if reclaimed:
        print(f"    reclaimed {reclaimed / 1e9:.2f} GB of stale partials", flush=True)

    # Drop `.part` files with no ledger: they cannot be resumed, so the engine
    # would truncate and re-fetch them from zero anyway.
    unusable = reclaim_unusable_partials(root, manifest)
    if unusable:
        print(
            f"    freed {unusable / 1e9:.2f} GB of un-resumable partials",
            flush=True,
        )

    errors = transfer.download_snapshot(manifest, root, force=force)
    if errors:
        print(f"    {len(errors)} file(s) did not complete", flush=True)
        return 2
    return 0


def main() -> int:
    if len(sys.argv) > 1 and sys.argv[1] == "--_worker":
        # 4 or 5 arguments: the optional 5th is `--force`, re-fetching every
        # file instead of keeping the ones already on disk at the right size.
        if len(sys.argv) not in (5, 6):
            print(
                "usage: fetch_image_model.py --_worker <repo> <target> "
                "<stall-seconds> [--force]",
                file=sys.stderr,
            )
            return 64
        return _worker(
            sys.argv[2],
            sys.argv[3],
            int(sys.argv[4]),
            force="--force" in sys.argv[5:],
        )

    verify_only = "--verify" in sys.argv
    load_check = "--load" in sys.argv
    force = "--force" in sys.argv
    # Deliberate escape hatch to stage the snapshot for another CUDA-capable
    # machine.
    # Access and disk are NOT overridable - a 401 or a full volume always fails.
    force_host = "--force-host" in sys.argv

    repo = resolve_repo()
    target = resolve_target()
    # `--verify` / `--load` never create the folder: an empty
    # `flux2-klein-4b` directory that makes the snapshot look "present" is
    # exactly what must not be produced. A missing folder is reported file by
    # file below instead.
    if not (verify_only or load_check):
        target.mkdir(parents=True, exist_ok=True)

    print(f"Fetching {repo} -> {target}")
    try:
        manifest = transfer.fetch_manifest(repo, ALLOW_PATTERNS, IGNORE_PATTERNS)
    except Exception as exc:
        if not verify_only:
            print(f"\nCould not reach the Hugging Face Hub: {exc}")
            print("Check the network/proxy, then re-run.")
            return 1
        # A local completeness check must not depend on the network: the machine
        # that most needs to know the snapshot is short is often the one whose
        # Hub access is broken. Fall back to the structural checks.
        manifest = None
        print(f"\n  ! the Hub's file list could not be read ({exc})")
        print("    checking the files the pipeline requires, offline\n")
    if manifest is not None:
        print(
            f"Revision {manifest['revision']} - {len(manifest['files'])} files, "
            f"{sum(manifest['files'].values()) / 1e9:.1f} GB total."
        )
    print("Re-running is safe: complete files are skipped, partial ones resume.\n")

    # `--verify` is deliberately a LOCAL check: it reads this folder and the
    # Hub's file list, downloads nothing, and must keep working on a machine
    # that has no token (which is exactly the machine that needs to know the
    # snapshot is incomplete). So the download pre-flight runs AFTER it.
    if verify_only:
        problems = transfer.verify_snapshot(target, manifest, REQUIRED_FILES, SHARD_MANIFESTS)
        if problems:
            return reject(target, problems)
        print(f"Snapshot is complete and verified at: {target}")
        return 0

    # `--load` = `--verify` plus a REAL pipeline build. The structural check
    # runs first on purpose: a truncated shard must be reported as "missing /
    # wrong size: <file>" by name, not as whatever safetensors says when a
    # half-written file is handed to it.
    if load_check:
        problems = transfer.verify_snapshot(target, manifest, REQUIRED_FILES, SHARD_MANIFESTS)
        if problems:
            return reject(target, problems)
        print(f"Snapshot is complete and verified at: {target}\n")
        return load_model(target)

    # --- Preflight: refuse to spend hours on a transfer that cannot work ----
    # Access first, then storage, then hardware. A failure returns BEFORE the first byte is
    # written, so nothing is half-downloaded and no disk is consumed.
    print("Pre-flight:")
    access_problems = check_hub_access(repo, manifest["revision"])
    if access_problems:
        for problem in access_problems:
            print(f"  - {problem}")
        return 1

    space_problems = transfer._check_disk_space(target, manifest)
    hardware_problems = [] if force_host else check_hardware()
    if force_host:
        print("  Host         : --force-host given, skipping the CUDA hardware gate.")
    blocking = space_problems + hardware_problems
    if blocking:
        for problem in blocking:
            print(f"  - {problem}")
        if hardware_problems and not space_problems:
            print(
                "\nNothing was downloaded. Re-run with --force-host to stage the "
                "weights anyway\n(for example to copy them to a larger machine)."
            )
        return 1
    print("  Preflight OK.\n")

    attempts = 6
    stall_seconds = 300
    script = os.path.abspath(__file__)
    failed = False

    for attempt in range(1, attempts + 1):
        print(f"\nDownload attempt {attempt}/{attempts}...", flush=True)
        proc = subprocess.Popen(
            [sys.executable, script, "--_worker", repo, str(target),
             str(stall_seconds)] + (["--force"] if force else []),
            cwd=str(PROJECT_ROOT),
        )
        code = proc.wait()

        if code == 0:
            break
        if code == 3:
            print("\nAttempt: transfer stalled (no disk progress).", flush=True)
        else:
            print(f"\nAttempt failed (exit {code}).", flush=True)

        if attempt == attempts:
            print(f"\nDownload failed after {attempts} attempts.")
            print("Re-run this script to resume. Nothing already fetched is lost.")
            failed = True
            break
        print("Retrying in 5s...", flush=True)
        time.sleep(5)

    if failed:
        return 1

    problems = transfer.verify_snapshot(target, manifest, REQUIRED_FILES, SHARD_MANIFESTS)
    if problems:
        print("\nSnapshot is STILL INCOMPLETE:")
        for problem in problems:
            print(f"  - {problem}")
        print("\nRe-run this script to finish the download.")
        return 1

    print(f"\nDone. Image model snapshot is complete and verified at: {target}")
    print("Set IMAGE_PROVIDER=local_diffusion in python-ai/.env to use it.")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())