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