settled-reach/tooling/domains/assets/trellis.py

"""Trellis 3D model generator connector — Gradio API wrapper.

Talks to the Trellis Gradio app on tower-of-joy :11510 (URL from config.json).
Pipeline: start session → upload + preprocess → seed → image_to_3d →
extract_glb → download .glb.

Formerly tooling/db/trellis_connector.py plus the tooling/trellis-batch.sh loop
(T-1290). `batch` is that loop in Python, and no longer hardcoded to the 16
character bodies — it takes a directory, or explicit names.

Gradio API parameter reference (TRELLIS v1, microsoft/TRELLIS):

    /image_to_3d — 9 inputs:
        0: image           (Image)   preprocessed image from /preprocess_image_1
        1: multiimages     (Gallery) [] for single-image mode
        2: is_multiimage   (State)   False for single-image, True for multi-image
        3: seed            (Slider)  int, 0-2147483647
        4: ss_guidance      (Slider)  float, sparse structure guidance strength (default 7.5)
        5: ss_steps         (Slider)  int, sparse structure sampling steps (default 12)
        6: slat_guidance    (Slider)  float, structured latent guidance strength (default 3.0)
        7: slat_steps       (Slider)  int, structured latent sampling steps (default 12)
        8: multiimage_algo  (Radio)   "stochastic" or "multidiffusion"

    /extract_glb — 3 inputs:
        0: output_buf      (State)   None — server uses internal state from image_to_3d
        1: simplify         (Slider)  float, mesh simplification ratio (default 0.95)
        2: texture_size     (Slider)  int, texture resolution (default 1024)

    Common failure modes:
        - "needed 9, got 8": missing is_multiimage (position 2) — must pass False
        - "needed 3, got 2": missing output_buf (position 0) — must pass None
        - "'float' cannot be interpreted as int": numpy version issue on server,
          Gradio Sliders send all values as float. Fix: patch flow_euler.py on
          the server to cast steps to int, or pin numpy < 2.0
        - CUDA device mismatch after crash: restart the container to clear GPU state
"""

from __future__ import annotations

import json
import os
import random
import string
import time
import urllib.request
from pathlib import Path

from tooling.core import config, console
from tooling.core.errors import ReachError
from tooling.domains.assets import endpoints

SERVICE = "trellis"

# The batch defaults are the character-body run the bash script was written for.
BODIES_INPUT = ".tmp/image-gen/characters/bodies"
BODIES_OUTPUT = ".tmp/glb-gen/characters/bodies"


def get_base_url() -> str:
    return endpoints.base_url(SERVICE)


def health() -> dict:
    """Is the Trellis API reachable, and what endpoints does it name?"""
    base = get_base_url()
    data = endpoints.call_json(
        urllib.request.Request(f"{base}/info", method="GET"),
        service=SERVICE,
        what="the health check",
        timeout=10,
    )
    return {"ok": True, "url": base, "endpoints": list(data.get("named_endpoints", {}).keys())}


def _call_api(base: str, endpoint: str, data: list, timeout: float = 600, session_hash: str | None = None):
    """POST one Gradio endpoint, sharing a session so gr.State survives between calls.

    image_to_3d stores its output in server-side State keyed by session_hash;
    extract_glb reads it back. Without a shared session the state is lost.
    """
    body: dict = {"data": data}
    if session_hash:
        body["session_hash"] = session_hash
    console.event(f"Calling {endpoint}...")
    result = endpoints.call_json(
        urllib.request.Request(
            f"{base}/api{endpoint}",
            data=json.dumps(body).encode(),
            headers={"Content-Type": "application/json"},
            method="POST",
        ),
        service=SERVICE,
        what=endpoint,
        timeout=timeout,
    )
    if isinstance(result, dict) and "data" in result:
        return result["data"]
    return result


def multipart_body(filename: str, payload: bytes, boundary: str = "----TrellisConnectorBoundary") -> bytes:
    """The multipart/form-data body Gradio's /upload expects — a 'files' field."""
    head = (
        f"--{boundary}\r\n"
        f'Content-Disposition: form-data; name="files"; filename="{filename}"\r\n'
        "Content-Type: image/png\r\n"
        "\r\n"
    ).encode()
    return head + payload + f"\r\n--{boundary}--\r\n".encode()


def _upload_image(base: str, image_path: str):
    filename = os.path.basename(image_path)
    with open(image_path, "rb") as f:
        body = multipart_body(filename, f.read())
    console.event(f"Uploading {filename}...")
    result = endpoints.call_json(
        urllib.request.Request(
            f"{base}/upload",
            data=body,
            headers={"Content-Type": "multipart/form-data; boundary=----TrellisConnectorBoundary"},
            method="POST",
        ),
        service=SERVICE,
        what="the image upload",
        timeout=30,
    )
    if isinstance(result, list) and result:
        return result[0]
    raise ReachError(
        f"Trellis accepted the upload but returned no file reference: {result}",
        fix="the Gradio app's upload API may have changed — check `reach assets trellis health`",
    )


def _download_file(url: str, output_path: str, base: str) -> int:
    if url.startswith("/"):
        url = f"{base}{url}"
    elif not url.startswith("http"):
        url = f"{base}/file={url}"
    console.event(f"Downloading to {output_path}...")
    body = endpoints.call(
        urllib.request.Request(url, method="GET"), service=SERVICE, what="the GLB download", timeout=120
    )
    with open(output_path, "wb") as f:
        f.write(body)
    return os.path.getsize(output_path)


def glb_url_from(glb_result) -> str | None:
    """extract_glb returns [model_viewer_data, download_button_data]; take the first URL."""
    if isinstance(glb_result, list):
        for item in glb_result:
            if isinstance(item, dict):
                url = item.get("url") or item.get("path")
                if url:
                    return url
    return None


def generate(
    image_path: str,
    output: str | None = None,
    simplify: float = 0.95,
    texture_size: int = 1024,
    seed: int = 0,
    timeout: int = 600,
) -> dict:
    """Image → .glb. Returns the result dict."""
    # The input is checked before the service, so a typo does not read as an
    # outage. (The old check also crashed: it passed indent= to print().)
    if not os.path.isfile(image_path):
        raise ReachError(f"image not found: {image_path}", fix="pass an existing .png")

    base = get_base_url()
    endpoints.call(
        urllib.request.Request(f"{base}/info", method="GET"),
        service=SERVICE,
        what="the availability check",
        timeout=5,
    )
    start_time = time.time()
    output = output or f"{os.path.splitext(os.path.basename(image_path))[0]}.glb"

    session = "".join(random.choices(string.ascii_lowercase + string.digits, k=12))
    console.event(f"Session: {session}")

    console.event("Starting session...", phase="1/5")
    _call_api(base, "/start_session", [], timeout=30, session_hash=session)

    console.event("Uploading and preprocessing image...", phase="2/5")
    uploaded = _upload_image(base, image_path)
    file_ref = {"path": uploaded, "meta": {"_type": "gradio.FileData"}}
    preprocessed = _call_api(base, "/preprocess_image_1", [file_ref], timeout=60, session_hash=session)
    preprocessed_ref = preprocessed[0] if isinstance(preprocessed, list) and preprocessed else preprocessed

    console.event("Generating 3D model...", phase="3/5")
    seed_result = _call_api(base, "/get_seed", [True, seed], timeout=10, session_hash=session)
    actual_seed = seed_result[0] if isinstance(seed_result, list) and seed_result else seed

    # Nine inputs, positions per the reference above. The server needs
    # flow_euler.py patched to cast steps to int (numpy >= 2).
    _call_api(
        base,
        "/image_to_3d",
        [preprocessed_ref, [], False, actual_seed, 7.5, 12, 3.0, 12, "stochastic"],
        timeout=timeout,
        session_hash=session,
    )

    console.event("Extracting GLB...", phase="4/5")
    glb_result = _call_api(base, "/extract_glb", [None, simplify, texture_size], timeout=120, session_hash=session)
    glb_url = glb_url_from(glb_result)
    if not glb_url:
        raise ReachError(
            f"could not find the GLB URL in the extract_glb response: {str(glb_result)[:300]}",
            fix="the Gradio app's response shape may have changed — see trellis.glb_url_from",
        )

    console.event("Downloading GLB...", phase="5/5")
    os.makedirs(os.path.dirname(os.path.abspath(output)), exist_ok=True)
    size = _download_file(glb_url, output, base)

    return {
        "ok": True,
        "file": output,
        "size_bytes": size,
        "simplify": simplify,
        "texture_size": texture_size,
        "seed": actual_seed,
        "generation_time_s": round(time.time() - start_time, 1),
        "source_image": image_path,
    }


def batch_plan(input_dir: Path, output_dir: Path, names: list[str] | None) -> list[tuple[str, Path, Path]]:
    """What a batch would do: (name, input png, output glb) per item — pure."""
    if names:
        chosen = names
    else:
        chosen = sorted(p.stem for p in input_dir.glob("*.png"))
    return [(n, input_dir / f"{n}.png", output_dir / f"{n}.glb") for n in chosen]


def batch(
    input_dir: str = BODIES_INPUT,
    output_dir: str = BODIES_OUTPUT,
    names: list[str] | None = None,
    simplify: float = 0.95,
    texture_size: int = 1024,
    cooldown: float = 15,
    max_retries: int = 3,
    retry_delay: float = 60,
) -> dict:
    """Generate one .glb per input image, one at a time, gently.

    Skips outputs that already exist. Retries each failure, and cools the GPU
    down between successful jobs — the Trellis box is shared (VRAM is scarce).
    Returns the summary; the router prints it and then fails the command if
    anything failed. The bash original exited 0 regardless.
    """
    root = config.repo_root()
    in_dir = (root / input_dir) if not Path(input_dir).is_absolute() else Path(input_dir)
    out_dir = (root / output_dir) if not Path(output_dir).is_absolute() else Path(output_dir)
    plan = batch_plan(in_dir, out_dir, names)
    if not plan:
        raise ReachError(f"no .png inputs in {in_dir}", fix="pass --input-dir with images, or --names")

    out_dir.mkdir(parents=True, exist_ok=True)
    total = len(plan)
    results = []
    console.event(f"Trellis batch: {total} item(s), cooldown {cooldown}s, {max_retries} tries each")

    for i, (name, src, dst) in enumerate(plan, 1):
        phase = f"{i}/{total}"
        if dst.exists():
            console.event(f"{name} — already exists, skipping", phase=phase)
            results.append({"name": name, "status": "skipped"})
            continue
        if not src.exists():
            console.event(f"{name} — input not found: {src}", phase=phase, level="warn")
            results.append({"name": name, "status": "failed", "error": f"input not found: {src}"})
            continue

        error = None
        for attempt in range(1, max_retries + 1):
            console.event(f"{name} (attempt {attempt}/{max_retries})...", phase=phase, progress=i / total)
            try:
                generate(str(src), output=str(dst), simplify=simplify, texture_size=texture_size)
                error = None
                break
            except ReachError as exc:
                error = exc.message
                console.event(f"{name} failed: {exc.message}", phase=phase, level="warn")
                if attempt < max_retries:
                    console.event(f"waiting {retry_delay}s before retry...", phase=phase)
                    time.sleep(retry_delay)

        if error is None:
            results.append({"name": name, "status": "success", "file": str(dst)})
            if i < total:
                console.event(f"cooling down {cooldown}s...", phase=phase)
                time.sleep(cooldown)
        else:
            results.append({"name": name, "status": "failed", "error": error})

    summary = {
        "ok": not any(r["status"] == "failed" for r in results),
        "total": total,
        "success": sum(r["status"] == "success" for r in results),
        "skipped": sum(r["status"] == "skipped" for r in results),
        "failed": sum(r["status"] == "failed" for r in results),
        "results": results,
    }
    return summary