Blame for runtime.py

"""Space adapter: CPU prepare, separate GPU detect/generate, CPU-only tickets.

The entry point must import ``spaces`` before importing this module and decorate
detect/generate callbacks with ``spaces.GPU``. Construct once at app startup; both
callbacks must share a Gradio concurrency limit of one. No Flask worker is used.
"""

from __future__ import annotations

import json
import math
import os
from pathlib import Path
import re
import shutil
import subprocess
import sys
import uuid

import cv2
import numpy as np
import torch


def _integer(value, name, low, high):
    try:
        number = float(value)
        result = int(number)
    except (TypeError, ValueError, OverflowError):
        raise ValueError(f"{name} must be an integer between {low} and {high}") from None
    if isinstance(value, bool) or not math.isfinite(number) or result != number or not low <= result <= high:
        raise ValueError(f"{name} must be an integer between {low} and {high}")
    return result


def _write_json(path, value):
    temporary = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp")
    temporary.write_text(json.dumps(value, ensure_ascii=False, allow_nan=False), encoding="utf-8")
    temporary.replace(path)


def _link(source, destination):
    try:
        os.link(source, destination)
    except OSError:
        shutil.copyfile(source, destination)


class SpaceRuntime:
    def __init__(self, root: Path, output_root: Path, enable_camera=True, enable_postprocess=True):
        self.root = Path(root).resolve()
        self.output_root = Path(output_root).resolve()
        if self.output_root == self.root or self.output_root in self.root.parents:
            raise ValueError("Space outputs must use a dedicated directory")
        if self.output_root.is_relative_to(self.root) and self.output_root.relative_to(self.root).parts[0] != "output":
            raise ValueError("Outputs inside the repository must be under output/")
        self.output_root.mkdir(parents=True, exist_ok=True)
        self.enable_camera = bool(enable_camera)
        self.enable_postprocess = bool(enable_postprocess)
        sys.path.insert(0, str(self.root))
        os.chdir(self.root)  # Existing asset paths are relative to the repository.
        os.environ["FLOWHMR_SAM_COMPILE"] = "0"
        from flowhmr.pipeline.v2m_inference import V2MInferenceEngine

        checkpoint = self.root / "checkpoints/flowhmr_latest/flowhmr_latest.ckpt"
        config = checkpoint.parent / "config.yml"
        self.engine = V2MInferenceEngine.build(
            ckpt=str(checkpoint), model_cfg_path=str(config), device="cuda", postprocess=False,
        )
        if self.enable_camera:
            from flowhmr.utils.runtime_vggt_camera import DEFAULT_VGGT_OMEGA_CKPT, _load_vggt_model

            _load_vggt_model(DEFAULT_VGGT_OMEGA_CKPT, self.engine.device)

    def _job(self, ticket):
        job_id = ticket.get("job_id") if isinstance(ticket, dict) else None
        if not isinstance(job_id, str) or re.fullmatch(r"[0-9a-f]{32}", job_id) is None:
            raise ValueError("Invalid or expired video session")
        directory = (self.output_root / job_id).resolve()
        if directory.parent != self.output_root or not (directory / "job.json").is_file():
            raise ValueError("Invalid or expired video session")
        metadata = json.loads((directory / "job.json").read_text(encoding="utf-8"))
        metadata["max_frames"] = _integer(metadata["max_frames"], "max_frames", 30, 900)
        return directory, metadata

    @staticmethod
    def _ticket(metadata):
        return {key: metadata[key] for key in ("job_id", "max_frames", "num_frames", "fps")}

    def prepare(self, video_path: str, max_frames: int = 900) -> dict:
        """CPU callback: transcode and truncate before detection consumes GPU time."""
        from flowhmr.utils.runtime_video_processing import VideoProcessor, find_ffmpeg

        max_frames = _integer(max_frames, "max_frames", 30, 900)
        source = Path(video_path).resolve()
        if not source.is_file():
            raise ValueError("Upload a video first")
        job_id = uuid.uuid4().hex
        directory = self.output_root / job_id
        input_dir = directory / "input"
        input_dir.mkdir(parents=True)
        video = input_dir / "input_30fps.mp4"
        command = [
            find_ffmpeg(), "-nostdin", "-y", "-i", str(source),
            "-map", "0:v:0", "-map", "0:a:0?", "-vf", "fps=30,pad=ceil(iw/2)*2:ceil(ih/2)*2",
            "-frames:v", str(max_frames), "-t", str(max_frames / 30),
            "-c:v", "libx264", "-threads", "2", "-preset", "fast", "-crf", "22",
            "-pix_fmt", "yuv420p", "-c:a", "aac", "-b:a", "128k",
            "-movflags", "+faststart", str(video),
        ]
        try:
            result = subprocess.run(command, capture_output=True, text=True, timeout=180)
            if result.returncode:
                raise ValueError(f"Video conversion failed: {result.stderr[-1500:]}")
            meta = VideoProcessor.get_video_meta(str(video))
            if meta.total_frames < 1 or meta.total_frames > max_frames:
                raise ValueError("The uploaded video could not be limited to the requested frame count")
            _link(video, directory / "input.mp4")
            metadata = {
                "job_id": job_id, "max_frames": max_frames, "num_frames": meta.total_frames,
                "fps": 30, "width": meta.width, "height": meta.height,
            }
            _write_json(directory / "job.json", metadata)
            return self._ticket(metadata)
        except Exception:
            shutil.rmtree(directory, ignore_errors=True)
            raise

    def detect(self, ticket) -> dict:
        """GPU callback: finish detection and release GPU before the user selects."""
        from flowhmr.utils.runtime_video_processing import VideoProcessor

        directory, metadata = self._job(ticket)
        video = directory / "input/input_30fps.mp4"
        result = self.engine.proc.detect_and_track(str(video))
        if not 0 < result.total_frames <= metadata["max_frames"]:
            raise ValueError("Decoded video length exceeds the prepared frame limit")
        best = VideoProcessor.select_best_track(result)
        offered = [i for i, track in enumerate(result.tracks) if track.length >= 15 or i == best]
        if not offered:
            offered = [-1]
        candidates = self._thumbnails(video, directory, result, offered)
        tracking = {
            "total_frames": result.total_frames, "width": result.width, "height": result.height,
            "tracks": [
                {"track_id": int(track.track_id), "frames": track.frames.astype(int).tolist(),
                 "bboxes": track.bboxes.astype(float).tolist(), "length": int(track.length),
                 "avg_area": float(track.avg_area)} for track in result.tracks
            ],
            "candidates": candidates, "default_index": offered.index(best),
        }
        _write_json(directory / "tracking.json", tracking)
        metadata["num_frames"] = result.total_frames
        _write_json(directory / "job.json", metadata)
        return {**self._ticket(metadata), "candidates": candidates, "default_index": tracking["default_index"]}

    @staticmethod
    def _thumbnails(video, directory, result, offered):
        """CPU image crops; paths refer only to generated session files."""
        thumbnails = directory / "persons"
        thumbnails.mkdir(exist_ok=True)
        capture = cv2.VideoCapture(str(video))
        candidates = []
        try:
            for index, track_index in enumerate(offered):
                track = result.tracks[track_index] if track_index >= 0 else None
                if track is not None:
                    boxes = track.bboxes
                    largest = int(np.argmax((boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])))
                    frame_index = int(track.frames[largest])
                    box = boxes[largest]
                else:
                    frame_index = 0
                    box = [0, 0, result.width, result.height]
                capture.set(cv2.CAP_PROP_POS_FRAMES, frame_index)
                ok, frame = capture.read()
                if not ok:
                    raise ValueError("Could not read a person-selection thumbnail")
                x1, y1, x2, y2 = map(float, box)
                pad_x, pad_y = (x2 - x1) * .15, (y2 - y1) * .08
                left, top = max(0, int(x1 - pad_x)), max(0, int(y1 - pad_y))
                right, bottom = min(result.width, int(x2 + pad_x)), min(result.height, int(y2 + pad_y))
                crop = frame[top:bottom, left:right]
                if crop.size == 0:
                    crop = frame
                scale = 192 / crop.shape[0]
                crop = cv2.resize(crop, (max(1, round(crop.shape[1] * scale)), 192), interpolation=cv2.INTER_AREA)
                thumbnail = thumbnails / f"{index}.jpg"
                if not cv2.imwrite(str(thumbnail), crop, [cv2.IMWRITE_JPEG_QUALITY, 88]):
                    raise OSError("Could not save a person-selection thumbnail")
                candidates.append({
                    "index": index, "track_index": track_index,
                    "track_id": int(track.track_id) if track is not None else -1,
                    "label": f"Person {index}" if track is not None else "Full frame (no person detected)",
                    "start": int(track.frames.min()) if track is not None else 0,
                    "end": int(track.frames.max()) if track is not None else result.total_frames - 1,
                    "num_frames": int(track.length) if track is not None else result.total_frames,
                    "avg_area": float(track.avg_area) if track is not None else float(result.width * result.height),
                    "thumbnail_path": str(thumbnail),
                })
        finally:
            capture.release()
        return candidates

    def generate(self, ticket, person_index, static_camera, vggt_interval, sam_chunk_size,
                 steps, cfg_scale, seed, postprocess, progress_callback=None) -> dict:
        """GPU callback: each request gets an isolated run; never wait for input here."""
        from flowhmr.utils.runtime_video_processing import TrackInfo, TrackingResult, VideoProcessor

        directory, metadata = self._job(ticket)
        if not isinstance(static_camera, bool) or not isinstance(postprocess, bool):
            raise ValueError("Camera and post-processing settings must be boolean")
        if not static_camera and not self.enable_camera:
            raise ValueError("Camera estimation is disabled in this Space")
        if postprocess and not self.enable_postprocess:
            raise ValueError("Post-processing is disabled in this Space")
        interval = _integer(vggt_interval, "vggt_interval", 1, 60)
        chunk = _integer(sam_chunk_size, "sam_chunk_size", 4, 32)
        if chunk not in (4, 8, 16, 32):
            raise ValueError("sam_chunk_size must be 4, 8, 16 or 32")
        steps = _integer(steps, "steps", 1, 50)
        seed = _integer(seed, "seed", 0, 2 ** 31 - 1)
        try:
            guidance = float(cfg_scale)
        except (ValueError, TypeError):
            raise ValueError("cfg_scale must be between 0 and 10") from None
        if isinstance(cfg_scale, bool) or not math.isfinite(guidance) or not 0 <= guidance <= 10:
            raise ValueError("cfg_scale must be between 0 and 10")
        tracking_path = directory / "tracking.json"
        if not tracking_path.is_file():
            raise ValueError("Run person detection first")
        tracking = json.loads(tracking_path.read_text(encoding="utf-8"))
        candidates = tracking["candidates"]
        if person_index is None:
            person_index = tracking["default_index"]
        person_index = _integer(person_index, "person_index", 0, len(candidates) - 1)
        result = TrackingResult(
            tracks=[TrackInfo(
                track_id=track["track_id"], frames=np.asarray(track["frames"], dtype=np.int64),
                bboxes=np.asarray(track["bboxes"], dtype=np.float32), length=track["length"],
                avg_area=track["avg_area"],
            ) for track in tracking["tracks"]],
            total_frames=tracking["total_frames"], width=tracking["width"], height=tracking["height"],
        )
        boxes, start_end = VideoProcessor.track_bboxes(result, candidates[person_index]["track_index"])
        run_id = uuid.uuid4().hex
        run = directory / "runs" / run_id
        (run / "input").mkdir(parents=True)
        _link(directory / "input.mp4", run / "input.mp4")
        _link(directory / "input/input_30fps.mp4", run / "input/input_30fps.mp4")
        np.savez(run / "input/input_bbox.npz", bbox=boxes, start_end=start_end)
        params = {
            "person_index": person_index, "static_camera": static_camera, "vggt_interval": interval,
            "sam_chunk_size": chunk, "steps": steps, "cfg_scale": guidance, "seed": seed,
            "postprocess": postprocess, "max_frames": metadata["max_frames"],
        }
        _write_json(run / "params.json", params)
        correction = None
        if postprocess:
            from flowhmr.utils.runtime_v2m_postprocess_simple import PostprocessPipeline

            correction = PostprocessPipeline.default(self.engine.body_model, self.engine.smpl_mesh)
        previous_steps = self.engine.pipeline.validation_steps
        previous_correction = self.engine.postprocess
        self.engine.pipeline.validation_steps = steps
        self.engine.postprocess = correction
        try:
            sub_dir = self.engine.process_video(
                str(run / "input.mp4"), str(run), seed=seed, cfg_scale=guidance,
                max_frames=metadata["max_frames"], vggt_frame_interval=interval,
                static_camera=static_camera, sam_chunk_size=chunk,
                progress_callback=progress_callback,
            )
        finally:
            self.engine.pipeline.validation_steps = previous_steps
            self.engine.postprocess = previous_correction
        suffix = "" if postprocess else "_origin"
        generated = Path(sub_dir) / f"input_seed{seed}{suffix}"
        npz_path = generated.with_suffix(".npz")
        if not npz_path.is_file():
            raise FileNotFoundError("Motion generation did not produce its NPZ output")
        name = f"motion_steps{steps}_cfg{guidance:g}_seed{seed}"
        output_npz = run / f"{name}.npz"
        _link(npz_path, output_npz)
        generated_fbx = generated.with_suffix(".fbx")
        output_fbx = run / f"{name}.fbx"
        if generated_fbx.is_file():
            _link(generated_fbx, output_fbx)
        mesh = self._export_mesh(output_npz, run / "mesh", start_end)
        output = {
            "job_id": metadata["job_id"], "run_id": run_id, **mesh,
            "video_path": str(run / "input.mp4"), "npz_path": str(output_npz),
            "fbx_path": str(output_fbx) if output_fbx.is_file() else None, "params": params,
        }
        _write_json(run / "result.json", output)
        return output

    @torch.no_grad()
    def _export_mesh(self, motion_npz, directory, start_end):
        """Preserve exported world coordinates; publish only 30-frame mesh chunks."""
        directory.mkdir(parents=True)
        with np.load(motion_npz, allow_pickle=False) as data:
            poses = torch.from_numpy(np.asarray(data["poses"], dtype=np.float32).copy())
            trans = torch.from_numpy(np.asarray(data["trans"], dtype=np.float32).copy())
            betas = torch.from_numpy(np.asarray(data["betas"], dtype=np.float32).copy()).reshape(1, -1)
        count = poses.shape[0]
        poses = poses.reshape(count, -1, 3)
        model = self.engine.smpl_mesh
        device = model.v_template.device
        betas = betas[:, :model.shapedirs.shape[-1]].to(device)
        faces = np.asarray(model.faces, dtype="<u4")
        faces_path = directory / "faces.bin"
        faces.tofile(faces_path)
        chunk_paths = []
        lower, upper = np.full(3, np.inf), np.full(3, -np.inf)
        for start in range(0, count, 30):
            vertices = model({
                "poses": poses[start:start + 30].to(device), "shapes": betas,
                "trans": trans[start:start + 30].to(device),
            })["vertices"].float().cpu().numpy()
            if not np.isfinite(vertices).all():
                raise ValueError("Generated mesh contains non-finite coordinates")
            lower = np.minimum(lower, vertices.reshape(-1, 3).min(0))
            upper = np.maximum(upper, vertices.reshape(-1, 3).max(0))
            chunk_path = directory / f"vertices_{start:06d}.bin"
            vertices.astype("<f2").tofile(chunk_path)
            chunk_paths.append(str(chunk_path))
        meta = {
            "num_frames": count, "num_vertices": int(model.v_template.shape[-2]),
            "num_faces": len(faces), "fps": 30, "vertex_dtype": "float16", "face_dtype": "uint32",
            "up_axis": "y", "bbox_min": lower.tolist(), "bbox_max": upper.tolist(),
            "root_trajectory": trans.numpy()[:, [0, 2]].round(4).tolist(),
            "track_start": int(start_end[0]), "track_end": int(start_end[0]) + count - 1,
            "chunk_frames": 30,
        }
        meta_path = directory / "meta.json"
        _write_json(meta_path, meta)
        return {"meta": meta, "meta_path": str(meta_path), "faces_path": str(faces_path), "chunk_paths": chunk_paths}