Upload v2/workflow.py with huggingface_hub

v2/workflow.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import json
+import traceback
+from dataclasses import asdict
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+from PIL import Image
+
+from .config import WorkflowConfig, ensure_job_dirs
+from .stages import (
+    materialize_artifacts,
+    run_rig_facefusion_stage,
+    run_shape_stage,
+    run_texture_stage,
+)
+
+
+class UnifiedWorkflow:
+    def __init__(
+        self, config: WorkflowConfig, output_dir: Path, seed: int, tex_seed: int
+    ):
+        self.config = config
+        self.output_dir = output_dir
+        self.seed = int(seed)
+        self.tex_seed = int(tex_seed)
+        self.artifacts_dir, self.logs_dir = ensure_job_dirs(output_dir)
+
+    def run(self, image_path: Path) -> dict[str, Any]:
+        started = datetime.now(timezone.utc)
+        manifest: dict[str, Any] = {
+            "workflow": "meshforge-v2-unified",
+            "started_at": started.isoformat(),
+            "input": {"image": str(image_path)},
+            "config": asdict(self.config),
+            "seed": self.seed,
+            "tex_seed": self.tex_seed,
+            "stages": [],
+            "artifacts": {},
+            "status": "running",
+        }
+
+        all_outputs: dict[str, Any] = {}
+
+        try:
+            image_array = np.array(Image.open(image_path).convert("RGB"))
+
+            shape = run_shape_stage(image_array, self.config, self.seed)
+            all_outputs.update(shape)
+            manifest["stages"].append(
+                {"name": "shape", "status": "ok", "detail": shape["status"]}
+            )
+
+            texture = run_texture_stage(
+                shape["shape_glb"], image_array, self.config, self.tex_seed
+            )
+            all_outputs.update(texture)
+            manifest["stages"].append(
+                {"name": "texture", "status": "ok", "detail": texture["status"]}
+            )
+
+            fused = run_rig_facefusion_stage(
+                texture["textured_glb"], image_array, self.config
+            )
+            all_outputs.update(fused)
+            manifest["stages"].append(
+                {
+                    "name": "rig_pshuman_fusion",
+                    "status": "ok",
+                    "detail": fused["status"],
+                }
+            )
+
+            manifest["artifacts"] = materialize_artifacts(
+                all_outputs, self.artifacts_dir
+            )
+            manifest["status"] = "completed"
+        except Exception as exc:
+            manifest["status"] = "failed"
+            manifest["error"] = {
+                "message": str(exc),
+                "traceback": traceback.format_exc(),
+            }
+
+        manifest["finished_at"] = datetime.now(timezone.utc).isoformat()
+        self._write_manifest(manifest)
+        return manifest
+
+    def _write_manifest(self, manifest: dict[str, Any]) -> None:
+        out = self.output_dir / "manifest.json"
+        out.parent.mkdir(parents=True, exist_ok=True)
+        out.write_text(json.dumps(manifest, indent=2), encoding="utf-8")