From 202b4e1d61aa553f4514defb6756418ad4c87dc8 Mon Sep 17 00:00:00 2001 From: "Billy D." Date: Tue, 24 Feb 2026 05:44:04 -0500 Subject: [PATCH] feat: scaffold avatar pipeline with ComfyUI driver, MLflow logging, and rclone promotion MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - setup.sh: automated desktop env setup (ComfyUI, 3D-Pack, UniRig, Blender, Ray) - ray-join.sh: join Ray cluster as external worker with 3d_gen resource label - vrm_export.py: headless Blender GLB→VRM conversion script - generate.py: ComfyUI API driver (submit workflow JSON, poll, download outputs) - log_mlflow.py: REST-only MLflow experiment tracking (no SDK dependency) - promote.py: rclone promotion of VRM files to gravenhollow S3 - CLI entry points: avatar-generate, avatar-promote - workflows/ placeholder for ComfyUI exported workflow JSONs Implements ADR-0063 (ComfyUI + TRELLIS + UniRig 3D avatar pipeline) --- README.md | 189 ++++++++++++++++++++++++++++ avatar_pipeline/__init__.py | 2 + avatar_pipeline/generate.py | 227 +++++++++++++++++++++++++++++++++ avatar_pipeline/log_mlflow.py | 150 ++++++++++++++++++++++ avatar_pipeline/promote.py | 131 +++++++++++++++++++ pyproject.toml | 40 ++++++ renovate.json | 7 ++ scripts/ray-join.sh | 44 +++++++ scripts/setup.sh | 228 ++++++++++++++++++++++++++++++++++ scripts/vrm_export.py | 111 +++++++++++++++++ workflows/README.md | 9 ++ 11 files changed, 1138 insertions(+) create mode 100644 avatar_pipeline/__init__.py create mode 100644 avatar_pipeline/generate.py create mode 100644 avatar_pipeline/log_mlflow.py create mode 100644 avatar_pipeline/promote.py create mode 100644 pyproject.toml create mode 100644 renovate.json create mode 100755 scripts/ray-join.sh create mode 100755 scripts/setup.sh create mode 100644 scripts/vrm_export.py create mode 100644 workflows/README.md diff --git a/README.md b/README.md index 7ab7623..7e50f98 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,191 @@ # avatar-pipeline +ComfyUI-based image-to-VRM avatar generation pipeline using [TRELLIS](https://github.com/microsoft/TRELLIS) (image → 3D mesh) and [UniRig](https://github.com/VAST-AI-Research/UniRig) (automatic rigging), running on a desktop NVIDIA GPU as an on-demand [Ray](https://www.ray.io/) worker. + +See [ADR-0063](https://git.daviestechlabs.io/daviestechlabs/homelab-design/src/branch/main/decisions/0063-comfyui-3d-avatar-pipeline.md) for the full design decision. + +## Overview + +``` +Reference Image → TRELLIS (1.2B) → Textured GLB → UniRig → Rigged GLB → Blender CLI → VRM + │ + rclone → gravenhollow + MLflow → experiment tracking +``` + +**Hardware target:** Arch Linux desktop with Ryzen 9 7950X, 64 GB DDR5, NVIDIA RTX 4070 (12 GB VRAM). + +## Project Structure + +``` +avatar-pipeline/ +├── pyproject.toml # uv project, CLI entry points +├── renovate.json # Renovate dependency updates +├── scripts/ +│ ├── setup.sh # Desktop environment setup (ComfyUI, 3D-Pack, UniRig, Blender, Ray) +│ ├── ray-join.sh # Join the Ray cluster as external worker +│ └── vrm_export.py # Blender headless GLB → VRM conversion +├── workflows/ +│ └── (ComfyUI workflow JSONs — exported from UI) +├── avatar_pipeline/ +│ ├── __init__.py +│ ├── generate.py # ComfyUI API driver: submit workflow → poll → download +│ ├── log_mlflow.py # Log params/metrics/artifacts to MLflow REST API +│ └── promote.py # rclone promote VRM files to gravenhollow storage +├── LICENSE +└── README.md +``` + +## Setup + +### Prerequisites + +- NVIDIA drivers + CUDA toolkit (`sudo pacman -S nvidia nvidia-utils cuda cudnn`) +- [uv](https://astral.sh/uv) installed +- [Blender](https://www.blender.org/) 4.x with [VRM Add-on](https://vrm-addon-for-blender.info/en/) +- [rclone](https://rclone.org/) configured with `gravenhollow` remote + +### Install Everything + +```bash +./scripts/setup.sh +``` + +This installs ComfyUI, ComfyUI-3D-Pack (includes TRELLIS nodes), UniRig, Ray, and the avatar-pipeline package. + +For partial installs: +```bash +./scripts/setup.sh --comfyui-only # ComfyUI + 3D-Pack only +./scripts/setup.sh --unirig-only # UniRig only +``` + +### Join Ray Cluster + +```bash +# Set the Ray head address (exposed via NodePort from the Talos cluster) +./scripts/ray-join.sh 192.168.100.50:6379 + +# Or via env var +RAY_HEAD_ADDRESS=192.168.100.50:6379 ./scripts/ray-join.sh + +# Verify +ray status +``` + +The desktop joins with resource labels `{"3d_gen": 1, "rtx4070": 1}` so only 3D generation workloads get scheduled here. + +## Usage + +### 1. Build a Workflow in ComfyUI + +```bash +cd ComfyUI && source .venv/bin/activate +python main.py # Starts at http://localhost:8188 +``` + +Build a node graph: Load Image → TRELLIS Image-to-3D → Mesh Simplify → UniRig Skeleton → UniRig Skinning → Save GLB. + +Export the workflow in API format and save to `workflows/`. + +### 2. Generate via CLI + +```bash +# Run a saved workflow +avatar-generate \ + --workflow workflows/image-to-vrm.json \ + --image reference.png \ + --seed 42 \ + --output-dir exports/ + +# Verbose output +avatar-generate -v --workflow workflows/image-to-vrm.json --image photo.png +``` + +### 3. Convert to VRM + +```bash +blender --background --python scripts/vrm_export.py -- \ + --input exports/model.glb \ + --output exports/Silver-Mage.vrm \ + --name "Silver Mage" +``` + +### 4. Log to MLflow + +Generation parameters and metrics are logged to the cluster's MLflow instance. Set the tracking URI: + +```bash +export MLFLOW_TRACKING_URI=http://mlflow.lab.daviestechlabs.io:5000 +``` + +Logging is integrated into the generate workflow, or can be called directly: + +```python +from avatar_pipeline.log_mlflow import log_generation + +log_generation( + avatar_name="Silver-Mage", + params={"trellis_seed": 42, "trellis_simplify": 0.95, "texture_size": 1024}, + metrics={"vertex_count": 12345, "face_count": 8000, "duration_s": 92.5}, +) +``` + +### 5. Promote to Production + +```bash +# Preview what would be copied +avatar-promote --dry-run exports/Silver-Mage.vrm + +# Promote to gravenhollow +avatar-promote exports/Silver-Mage.vrm + +# Promote all VRM files +avatar-promote exports/*.vrm +``` + +After promotion, register the model in companions-frontend by adding it to `AllowedAvatarModels` in the Go + JS allowlists. + +## Key Parameters + +| Parameter | Default | Description | +|-----------|---------|-------------| +| `trellis_seed` | random | Reproducibility seed for TRELLIS generation | +| `trellis_steps` | 12 | Sampling steps (sparse structure + SLAT) | +| `trellis_cfg_strength` | 7.5 | Classifier-free guidance strength | +| `trellis_simplify` | 0.95 | Triangle reduction ratio (1.0 = no reduction) | +| `texture_size` | 1024 | Output texture resolution | + +## VRAM Budget (RTX 4070 — 12 GB) + +Models run sequentially, not concurrently: + +| Step | VRAM | Time | +|------|------|------| +| TRELLIS image-large (1.2B, fp16) | ~10 GB | ~30s | +| UniRig skeleton prediction | ~4 GB | ~30s | +| UniRig skinning weights | ~4 GB | ~30s | +| Blender CLI (VRM export) | CPU only | ~10s | + +Peak: ~10 GB during TRELLIS. 64 GB system RAM handles model loading overhead. + +## Development + +```bash +# Install in dev mode +uv pip install -e ".[dev]" + +# Lint +uv run ruff check . +uv run ruff format --check . + +# Auto-fix +uv run ruff check --fix . && uv run ruff format . +``` + +## Related + +- [ADR-0063](https://git.daviestechlabs.io/daviestechlabs/homelab-design/src/branch/main/decisions/0063-comfyui-3d-avatar-pipeline.md) — Design decision +- [ADR-0011](https://git.daviestechlabs.io/daviestechlabs/homelab-design/src/branch/main/decisions/0011-kuberay-unified-gpu-backend.md) — KubeRay GPU backend +- [TRELLIS](https://github.com/microsoft/TRELLIS) — Image-to-3D generation (CVPR'25) +- [UniRig](https://github.com/VAST-AI-Research/UniRig) — Automatic rigging (SIGGRAPH'25) +- [ComfyUI-3D-Pack](https://github.com/MrForExample/ComfyUI-3D-Pack) — 3D nodes for ComfyUI \ No newline at end of file diff --git a/avatar_pipeline/__init__.py b/avatar_pipeline/__init__.py new file mode 100644 index 0000000..f887501 --- /dev/null +++ b/avatar_pipeline/__init__.py @@ -0,0 +1,2 @@ +# ComfyUI image-to-VRM avatar generation pipeline +# with TRELLIS + UniRig on desktop Ray worker diff --git a/avatar_pipeline/generate.py b/avatar_pipeline/generate.py new file mode 100644 index 0000000..7e0c23b --- /dev/null +++ b/avatar_pipeline/generate.py @@ -0,0 +1,227 @@ +"""Submit a ComfyUI workflow and collect the output. + +ComfyUI exposes a REST API at http://localhost:8188: + POST /prompt — queue a workflow + GET /history/{id} — poll execution status + GET /view?filename= — download output files + +This module loads a workflow JSON, injects runtime parameters +(image path, seed, etc.), submits it, and waits for completion. + +Usage: + avatar-generate --workflow workflows/image-to-vrm.json \\ + --image reference.png \\ + --name "My Avatar" \\ + [--seed 42] \\ + [--comfyui-url http://localhost:8188] +""" + +from __future__ import annotations + +import argparse +import json +import logging +import sys +import time +import uuid +from pathlib import Path + +import httpx + +logger = logging.getLogger(__name__) + +DEFAULT_COMFYUI_URL = "http://localhost:8188" +POLL_INTERVAL_S = 2.0 +MAX_WAIT_S = 600 # 10 minutes + + +def load_workflow(path: Path) -> dict: + """Load a ComfyUI workflow JSON (API format).""" + with open(path) as f: + return json.load(f) + + +def inject_params(workflow: dict, image_path: str | None, seed: int | None) -> dict: + """Inject runtime parameters into the workflow. + + This is a best-effort approach — it scans nodes for known types + and overrides their values. The exact node IDs depend on the + exported workflow JSON, so this may need adjustment per workflow. + """ + for _node_id, node in workflow.items(): + if not isinstance(node, dict): + continue + + class_type = node.get("class_type", "") + inputs = node.get("inputs", {}) + + # Override image path in LoadImage nodes + if class_type == "LoadImage" and image_path: + inputs["image"] = image_path + logger.info("Injected image path: %s", image_path) + + # Override seed in any node that has a seed input + if seed is not None and "seed" in inputs: + inputs["seed"] = seed + logger.info("Injected seed=%d into node %s (%s)", seed, _node_id, class_type) + + return workflow + + +def submit_prompt( + client: httpx.Client, + workflow: dict, + client_id: str, +) -> str: + """Submit a workflow to ComfyUI and return the prompt ID.""" + payload = { + "prompt": workflow, + "client_id": client_id, + } + resp = client.post("/prompt", json=payload) + resp.raise_for_status() + data = resp.json() + prompt_id = data["prompt_id"] + logger.info("Submitted prompt: %s", prompt_id) + return prompt_id + + +def wait_for_completion( + client: httpx.Client, + prompt_id: str, + timeout: float = MAX_WAIT_S, +) -> dict: + """Poll /history/{prompt_id} until the workflow completes.""" + start = time.monotonic() + + while time.monotonic() - start < timeout: + resp = client.get(f"/history/{prompt_id}") + resp.raise_for_status() + history = resp.json() + + if prompt_id in history: + entry = history[prompt_id] + status = entry.get("status", {}) + + if status.get("completed", False): + elapsed = time.monotonic() - start + logger.info("Workflow completed in %.1fs", elapsed) + return entry + + if status.get("status_str") == "error": + logger.error("Workflow failed: %s", status) + raise RuntimeError(f"ComfyUI workflow failed: {status}") + + time.sleep(POLL_INTERVAL_S) + + raise TimeoutError(f"Workflow did not complete within {timeout}s") + + +def collect_outputs(entry: dict) -> list[dict]: + """Extract output file info from a completed history entry.""" + outputs = [] + for _node_id, node_output in entry.get("outputs", {}).items(): + for output_type in ("images", "gltf", "meshes"): + for item in node_output.get(output_type, []): + outputs.append(item) + return outputs + + +def download_output( + client: httpx.Client, + filename: str, + subfolder: str, + output_dir: Path, + file_type: str = "output", +) -> Path: + """Download an output file from ComfyUI.""" + params = {"filename": filename, "subfolder": subfolder, "type": file_type} + resp = client.get("/view", params=params) + resp.raise_for_status() + + dest = output_dir / filename + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(resp.content) + logger.info("Downloaded: %s (%d bytes)", dest, len(resp.content)) + return dest + + +def generate( + workflow_path: Path, + image_path: str | None = None, + seed: int | None = None, + output_dir: Path | None = None, + comfyui_url: str = DEFAULT_COMFYUI_URL, +) -> list[Path]: + """Run the full generation pipeline: load → inject → submit → wait → download.""" + output_dir = output_dir or Path("exports") + output_dir.mkdir(parents=True, exist_ok=True) + + client_id = str(uuid.uuid4()) + workflow = load_workflow(workflow_path) + workflow = inject_params(workflow, image_path, seed) + + with httpx.Client(base_url=comfyui_url, timeout=30.0) as client: + prompt_id = submit_prompt(client, workflow, client_id) + entry = wait_for_completion(client, prompt_id) + outputs = collect_outputs(entry) + + downloaded = [] + for item in outputs: + path = download_output( + client, + filename=item["filename"], + subfolder=item.get("subfolder", ""), + output_dir=output_dir, + file_type=item.get("type", "output"), + ) + downloaded.append(path) + + return downloaded + + +def main() -> None: + """CLI entry point.""" + parser = argparse.ArgumentParser( + description="Submit a ComfyUI workflow and collect outputs" + ) + parser.add_argument( + "--workflow", + type=Path, + required=True, + help="Path to ComfyUI workflow JSON (API format)", + ) + parser.add_argument("--image", help="Reference image path to inject into LoadImage nodes") + parser.add_argument("--seed", type=int, help="Seed to inject into generation nodes") + parser.add_argument("--output-dir", type=Path, default=Path("exports"), help="Output directory") + parser.add_argument( + "--comfyui-url", + default=DEFAULT_COMFYUI_URL, + help="ComfyUI server URL (default: %(default)s)", + ) + parser.add_argument("--verbose", "-v", action="store_true") + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + + try: + outputs = generate( + workflow_path=args.workflow, + image_path=args.image, + seed=args.seed, + output_dir=args.output_dir, + comfyui_url=args.comfyui_url, + ) + print(f"\nGenerated {len(outputs)} output(s):") + for p in outputs: + print(f" {p}") + except Exception: + logger.exception("Generation failed") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/avatar_pipeline/log_mlflow.py b/avatar_pipeline/log_mlflow.py new file mode 100644 index 0000000..d1ae1f3 --- /dev/null +++ b/avatar_pipeline/log_mlflow.py @@ -0,0 +1,150 @@ +"""Log avatar generation results to MLflow via REST API. + +Uses the same lightweight REST-only approach as ray-serve's +mlflow_logger.py — no heavyweight mlflow SDK dependency. + +Usage: + from avatar_pipeline.log_mlflow import log_generation + + log_generation( + avatar_name="Silver-Mage", + params={"trellis_seed": 42, "trellis_simplify": 0.95}, + metrics={"vertex_count": 12345, "face_count": 8000, "duration_s": 45.2}, + artifacts={"vrm": Path("exports/Silver-Mage.vrm")}, + ) +""" + +from __future__ import annotations + +import json +import logging +import os +import time +import urllib.error +import urllib.request +from pathlib import Path + +logger = logging.getLogger(__name__) + +DEFAULT_TRACKING_URI = "http://mlflow.mlflow.svc.cluster.local:80" +EXPERIMENT_NAME = "3d-avatar-generation" + + +def _base_url() -> str: + return os.environ.get("MLFLOW_TRACKING_URI", DEFAULT_TRACKING_URI).rstrip("/") + + +def _post(path: str, body: dict) -> dict: + url = f"{_base_url()}/api/2.0/mlflow/{path}" + data = json.dumps(body).encode() + req = urllib.request.Request( + url, data=data, headers={"Content-Type": "application/json"}, method="POST" + ) + with urllib.request.urlopen(req, timeout=15) as resp: + return json.loads(resp.read().decode()) + + +def _get_or_create_experiment(name: str) -> str: + try: + resp = _post("experiments/get-by-name", {"experiment_name": name}) + return resp["experiment"]["experiment_id"] + except urllib.error.HTTPError: + resp = _post("experiments/create", {"name": name}) + return resp["experiment_id"] + + +def _create_run(experiment_id: str, run_name: str, tags: dict[str, str]) -> str: + tag_list = [ + {"key": k, "value": v} + for k, v in { + "mlflow.runName": run_name, + "mlflow.source.type": "LOCAL", + "mlflow.source.name": "avatar-pipeline", + "hostname": os.environ.get("HOSTNAME", "desktop"), + **tags, + }.items() + ] + resp = _post( + "runs/create", + { + "experiment_id": experiment_id, + "run_name": run_name, + "start_time": int(time.time() * 1000), + "tags": tag_list, + }, + ) + return resp["run"]["info"]["run_id"] + + +def _log_params(run_id: str, params: dict[str, str | int | float]) -> None: + param_list = [{"key": k, "value": str(v)[:500]} for k, v in params.items()] + _post("runs/log-batch", {"run_id": run_id, "params": param_list}) + + +def _log_metrics(run_id: str, metrics: dict[str, float]) -> None: + ts = int(time.time() * 1000) + metric_list = [ + {"key": k, "value": float(v), "timestamp": ts, "step": 0} for k, v in metrics.items() + ] + _post("runs/log-batch", {"run_id": run_id, "metrics": metric_list}) + + +def _log_artifact(run_id: str, key: str, path: Path) -> None: + """Log an artifact path as a tag (actual artifact upload requires artifact store). + + For local files, we record the path. For S3-promoted files, the caller + should include the S3 URI in params instead. + """ + _post( + "runs/log-batch", + { + "run_id": run_id, + "tags": [{"key": f"artifact.{key}", "value": str(path)}], + }, + ) + + +def _end_run(run_id: str) -> None: + _post( + "runs/update", + { + "run_id": run_id, + "status": "FINISHED", + "end_time": int(time.time() * 1000), + }, + ) + + +def log_generation( + avatar_name: str, + params: dict[str, str | int | float] | None = None, + metrics: dict[str, float] | None = None, + artifacts: dict[str, Path] | None = None, + tags: dict[str, str] | None = None, + experiment_name: str = EXPERIMENT_NAME, +) -> str | None: + """Log a complete avatar generation run to MLflow. + + Returns the MLflow run_id on success, None on failure. + """ + try: + exp_id = _get_or_create_experiment(experiment_name) + run_id = _create_run(exp_id, run_name=avatar_name, tags=tags or {}) + + if params: + _log_params(run_id, params) + + if metrics: + _log_metrics(run_id, metrics) + + if artifacts: + for key, path in artifacts.items(): + _log_artifact(run_id, key, path) + + _end_run(run_id) + logger.info("Logged to MLflow: run_id=%s, experiment=%s", run_id, experiment_name) + return run_id + + except Exception: + logger.warning("MLflow logging failed — generation still succeeded", exc_info=True) + return None diff --git a/avatar_pipeline/promote.py b/avatar_pipeline/promote.py new file mode 100644 index 0000000..b65fb44 --- /dev/null +++ b/avatar_pipeline/promote.py @@ -0,0 +1,131 @@ +"""Promote VRM files to gravenhollow via rclone. + +Usage: + avatar-promote exports/Silver-Mage.vrm + avatar-promote exports/Silver-Mage.vrm --bucket companion-avatars + avatar-promote --dry-run exports/*.vrm +""" + +from __future__ import annotations + +import argparse +import logging +import shutil +import subprocess +import sys +from pathlib import Path + +logger = logging.getLogger(__name__) + +DEFAULT_REMOTE = "gravenhollow" +DEFAULT_BUCKET = "avatar-models" + + +def check_rclone() -> bool: + """Verify rclone is installed and the remote is configured.""" + if not shutil.which("rclone"): + logger.error("rclone not found. Install: sudo pacman -S rclone") + return False + + result = subprocess.run( + ["rclone", "listremotes"], + capture_output=True, + text=True, + check=False, + ) + remotes = result.stdout.strip().split("\n") + + if f"{DEFAULT_REMOTE}:" not in remotes: + logger.error( + "rclone remote '%s' not configured. Run scripts/setup.sh or configure manually.", + DEFAULT_REMOTE, + ) + return False + + return True + + +def promote( + files: list[Path], + remote: str = DEFAULT_REMOTE, + bucket: str = DEFAULT_BUCKET, + dry_run: bool = False, +) -> list[str]: + """Copy VRM files to gravenhollow S3 via rclone. + + Returns list of promoted remote paths. + """ + if not check_rclone(): + raise RuntimeError("rclone not available") + + promoted = [] + for file_path in files: + if not file_path.exists(): + logger.warning("File not found, skipping: %s", file_path) + continue + + if file_path.suffix.lower() not in (".vrm", ".glb", ".fbx"): + logger.warning("Unexpected file type, skipping: %s", file_path) + continue + + dest = f"{remote}:{bucket}/{file_path.name}" + cmd = ["rclone", "copy", str(file_path), f"{remote}:{bucket}/"] + + if dry_run: + cmd.append("--dry-run") + + logger.info("Promoting: %s → %s", file_path, dest) + + result = subprocess.run( + cmd, + capture_output=True, + text=True, + check=False, + ) + + if result.returncode != 0: + logger.error("rclone failed for %s: %s", file_path, result.stderr) + continue + + if dry_run: + logger.info(" (dry-run) Would copy %s", file_path.name) + else: + logger.info(" Promoted: %s", dest) + + promoted.append(dest) + + return promoted + + +def main() -> None: + """CLI entry point.""" + parser = argparse.ArgumentParser(description="Promote VRM files to gravenhollow storage") + parser.add_argument("files", nargs="+", type=Path, help="VRM/GLB files to promote") + parser.add_argument("--remote", default=DEFAULT_REMOTE, help="rclone remote name") + parser.add_argument("--bucket", default=DEFAULT_BUCKET, help="S3 bucket name") + parser.add_argument("--dry-run", action="store_true", help="Show what would be copied") + parser.add_argument("--verbose", "-v", action="store_true") + args = parser.parse_args() + + logging.basicConfig( + level=logging.DEBUG if args.verbose else logging.INFO, + format="%(asctime)s %(levelname)s %(name)s: %(message)s", + ) + + try: + promoted = promote( + files=args.files, + remote=args.remote, + bucket=args.bucket, + dry_run=args.dry_run, + ) + print(f"\nPromoted {len(promoted)} file(s)") + for p in promoted: + print(f" {p}") + except Exception: + logger.exception("Promotion failed") + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..98b0a04 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,40 @@ +[project] +name = "avatar-pipeline" +version = "0.1.0" +description = "ComfyUI image-to-VRM avatar generation pipeline with TRELLIS + UniRig" +requires-python = ">=3.11" +license = { text = "MIT" } +authors = [{ name = "Davies Tech Labs" }] + +dependencies = [ + # HTTP client for ComfyUI API + MLflow REST + "httpx>=0.27.0", + + # Image handling + "Pillow>=10.0.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=8.0.0", + "ruff>=0.4.0", +] + +[project.scripts] +avatar-generate = "avatar_pipeline.generate:main" +avatar-promote = "avatar_pipeline.promote:main" + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.ruff] +line-length = 100 +target-version = "py311" + +[tool.ruff.lint] +select = ["E", "F", "W", "I", "UP", "B", "C4", "SIM"] +ignore = ["E501"] + +[tool.ruff.format] +quote-style = "double" diff --git a/renovate.json b/renovate.json new file mode 100644 index 0000000..00171e1 --- /dev/null +++ b/renovate.json @@ -0,0 +1,7 @@ +{ + "$schema": "https://docs.renovatebot.com/renovate-schema.json", + "extends": [ + "local>daviestechlabs/renovate-config", + "local>daviestechlabs/renovate-config:python" + ] +} diff --git a/scripts/ray-join.sh b/scripts/ray-join.sh new file mode 100755 index 0000000..655b853 --- /dev/null +++ b/scripts/ray-join.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# Join the Ray cluster as an external worker. +# +# The Ray head's GCS port must be reachable from this machine. +# Set RAY_HEAD_ADDRESS or pass it as the first argument. +# +# Usage: +# ./scripts/ray-join.sh # uses RAY_HEAD_ADDRESS env +# ./scripts/ray-join.sh 192.168.100.50:6379 # explicit address +# RAY_HEAD_ADDRESS=192.168.100.50:6379 ./scripts/ray-join.sh +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "$SCRIPT_DIR")" +COMFYUI_DIR="$PROJECT_DIR/ComfyUI" + +RAY_HEAD="${1:-${RAY_HEAD_ADDRESS:-}}" + +if [[ -z "$RAY_HEAD" ]]; then + echo "ERROR: Ray head address required." + echo "Usage: $0 :" + echo " or: RAY_HEAD_ADDRESS=: $0" + exit 1 +fi + +# Activate ComfyUI venv (where Ray is installed) +# shellcheck disable=SC1091 +source "$COMFYUI_DIR/.venv/bin/activate" + +# Stop any existing Ray worker on this machine +ray stop --force 2>/dev/null || true + +echo "Joining Ray cluster at $RAY_HEAD..." + +ray start \ + --address="$RAY_HEAD" \ + --num-cpus=16 \ + --num-gpus=1 \ + --resources='{"3d_gen": 1, "rtx4070": 1}' \ + --node-name=desktop + +echo "" +echo "Desktop joined Ray cluster. Verify with: ray status" +echo "To disconnect: ray stop" diff --git a/scripts/setup.sh b/scripts/setup.sh new file mode 100755 index 0000000..1087721 --- /dev/null +++ b/scripts/setup.sh @@ -0,0 +1,228 @@ +#!/usr/bin/env bash +# Desktop environment setup for ComfyUI + TRELLIS + UniRig + Blender +# Target: Arch Linux with NVIDIA RTX 4070 (CUDA 12.x) +# +# Prerequisites: +# - NVIDIA drivers installed (nvidia, nvidia-utils) +# - CUDA toolkit installed (cuda, cudnn) +# - uv installed (https://astral.sh/uv) +# +# Usage: ./scripts/setup.sh [--comfyui-only | --unirig-only | --all] +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +PROJECT_DIR="$(dirname "$SCRIPT_DIR")" +COMFYUI_DIR="$PROJECT_DIR/ComfyUI" + +RED='\033[0;31m' +GREEN='\033[0;32m' +YELLOW='\033[1;33m' +NC='\033[0m' + +info() { echo -e "${GREEN}[INFO]${NC} $*"; } +warn() { echo -e "${YELLOW}[WARN]${NC} $*"; } +error() { echo -e "${RED}[ERROR]${NC} $*" >&2; } + +check_prerequisites() { + info "Checking prerequisites..." + + if ! command -v nvidia-smi &>/dev/null; then + error "nvidia-smi not found. Install NVIDIA drivers: sudo pacman -S nvidia nvidia-utils" + exit 1 + fi + info "NVIDIA driver: $(nvidia-smi --query-gpu=driver_version --format=csv,noheader | head -1)" + + if ! command -v nvcc &>/dev/null; then + error "nvcc not found. Install CUDA: sudo pacman -S cuda cudnn" + exit 1 + fi + info "CUDA: $(nvcc --version | grep release | awk '{print $6}')" + + if ! command -v uv &>/dev/null; then + error "uv not found. Install: curl -LsSf https://astral.sh/uv/install.sh | sh" + exit 1 + fi + info "uv: $(uv --version)" + + if ! command -v git &>/dev/null; then + error "git not found. Install: sudo pacman -S git" + exit 1 + fi +} + +install_comfyui() { + info "Installing ComfyUI..." + + if [[ -d "$COMFYUI_DIR" ]]; then + info "ComfyUI directory exists, pulling latest..." + git -C "$COMFYUI_DIR" pull + else + git clone https://github.com/comfyanonymous/ComfyUI.git "$COMFYUI_DIR" + fi + + cd "$COMFYUI_DIR" + + if [[ ! -d ".venv" ]]; then + uv venv --python 3.11 + fi + + # shellcheck disable=SC1091 + source .venv/bin/activate + uv pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124 + uv pip install -r requirements.txt + + info "ComfyUI installed at $COMFYUI_DIR" +} + +install_3d_pack() { + info "Installing ComfyUI-3D-Pack (includes TRELLIS nodes)..." + + local nodes_dir="$COMFYUI_DIR/custom_nodes" + mkdir -p "$nodes_dir" + + if [[ -d "$nodes_dir/ComfyUI-3D-Pack" ]]; then + info "ComfyUI-3D-Pack exists, pulling latest..." + git -C "$nodes_dir/ComfyUI-3D-Pack" pull + else + git clone https://github.com/MrForExample/ComfyUI-3D-Pack.git "$nodes_dir/ComfyUI-3D-Pack" + fi + + cd "$nodes_dir/ComfyUI-3D-Pack" + + # Activate ComfyUI venv for shared deps + # shellcheck disable=SC1091 + source "$COMFYUI_DIR/.venv/bin/activate" + uv pip install -r requirements.txt + + if [[ -f "install.py" ]]; then + python install.py + fi + + info "ComfyUI-3D-Pack installed" +} + +install_unirig() { + info "Installing UniRig..." + + local unirig_dir="$PROJECT_DIR/UniRig" + + if [[ -d "$unirig_dir" ]]; then + info "UniRig exists, pulling latest..." + git -C "$unirig_dir" pull + else + git clone https://github.com/VAST-AI-Research/UniRig.git "$unirig_dir" + fi + + cd "$unirig_dir" + + # UniRig uses its own venv (different torch/dep requirements possible) + if [[ ! -d ".venv" ]]; then + uv venv --python 3.11 + fi + + # shellcheck disable=SC1091 + source .venv/bin/activate + uv pip install torch torchvision --index-url https://download.pytorch.org/whl/cu124 + uv pip install -r requirements.txt + + # spconv for sparse convolutions + uv pip install spconv-cu124 || warn "spconv-cu124 install failed — may need manual build" + + # flash-attn for efficient attention + uv pip install flash-attn --no-build-isolation || warn "flash-attn install failed — TRELLIS will fall back to xformers" + + info "UniRig installed at $unirig_dir" +} + +install_blender() { + info "Checking Blender..." + + if command -v blender &>/dev/null; then + info "Blender: $(blender --version 2>/dev/null | head -1)" + else + warn "Blender not found. Install: sudo pacman -S blender" + warn "Then install VRM Add-on: https://vrm-addon-for-blender.info/en/" + fi +} + +install_rclone() { + info "Checking rclone..." + + if command -v rclone &>/dev/null; then + info "rclone: $(rclone --version | head -1)" + + if rclone listremotes | grep -q "^gravenhollow:"; then + info "rclone remote 'gravenhollow' already configured" + else + warn "rclone remote 'gravenhollow' not configured." + warn "Run: rclone config create gravenhollow s3 provider=Other endpoint=https://gravenhollow.lab.daviestechlabs.io:30292 access_key_id= secret_access_key=" + fi + else + warn "rclone not found. Install: sudo pacman -S rclone" + fi +} + +install_ray() { + info "Checking Ray..." + + # Install ray into ComfyUI venv + # shellcheck disable=SC1091 + source "$COMFYUI_DIR/.venv/bin/activate" + + if python -c "import ray" 2>/dev/null; then + info "Ray: $(python -c 'import ray; print(ray.__version__)')" + else + info "Installing Ray..." + uv pip install "ray[default]>=2.53.0" + info "Ray installed" + fi +} + +install_avatar_pipeline() { + info "Installing avatar-pipeline package..." + cd "$PROJECT_DIR" + + # shellcheck disable=SC1091 + source "$COMFYUI_DIR/.venv/bin/activate" + uv pip install -e ".[dev]" + + info "avatar-pipeline installed (editable)" +} + +# ── Main ────────────────────────────────────────────────────────────── + +MODE="${1:---all}" + +check_prerequisites + +case "$MODE" in + --comfyui-only) + install_comfyui + install_3d_pack + ;; + --unirig-only) + install_unirig + ;; + --all) + install_comfyui + install_3d_pack + install_unirig + install_blender + install_rclone + install_ray + install_avatar_pipeline + ;; + *) + echo "Usage: $0 [--comfyui-only | --unirig-only | --all]" + exit 1 + ;; +esac + +info "Setup complete!" +info "" +info "Next steps:" +info " 1. Start ComfyUI: cd $COMFYUI_DIR && source .venv/bin/activate && python main.py" +info " 2. Open browser: http://localhost:8188" +info " 3. Build your image-to-VRM workflow in the ComfyUI node graph" +info " 4. Export workflow JSON to: $PROJECT_DIR/workflows/" +info " 5. Join Ray cluster: ./scripts/ray-join.sh" diff --git a/scripts/vrm_export.py b/scripts/vrm_export.py new file mode 100644 index 0000000..84249fe --- /dev/null +++ b/scripts/vrm_export.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Headless Blender script for GLB → VRM conversion. + +Invoked via: + blender --background --python scripts/vrm_export.py -- \\ + --input rigged_model.glb \\ + --output avatar.vrm \\ + --name "Avatar Name" \\ + --author "DaviesTechLabs Pipeline" + +Requires: + - Blender 4.x + - VRM Add-on for Blender (https://vrm-addon-for-blender.info/en/) +""" + +import argparse +import sys + +# Blender's bpy is only available when running inside Blender +try: + import bpy # type: ignore[import-untyped] +except ImportError: + print("ERROR: This script must be run inside Blender.") + print("Usage: blender --background --python scripts/vrm_export.py -- --input model.glb --output model.vrm") + sys.exit(1) + + +def parse_args() -> argparse.Namespace: + """Parse arguments after the '--' separator in Blender CLI.""" + argv = sys.argv[sys.argv.index("--") + 1 :] if "--" in sys.argv else [] + parser = argparse.ArgumentParser(description="Convert rigged GLB to VRM") + parser.add_argument("--input", required=True, help="Path to rigged GLB file") + parser.add_argument("--output", required=True, help="Output VRM file path") + parser.add_argument("--name", default="Generated Avatar", help="Avatar display name") + parser.add_argument("--author", default="DaviesTechLabs Pipeline", help="Author name") + return parser.parse_args(argv) + + +def clear_scene() -> None: + """Remove all objects from the scene.""" + bpy.ops.wm.read_factory_settings(use_empty=True) + + +def import_glb(path: str) -> None: + """Import a GLB file into the current scene.""" + bpy.ops.import_scene.gltf(filepath=path) + print(f"Imported: {path}") + print(f" Objects: {[obj.name for obj in bpy.data.objects]}") + + +def find_armature() -> "bpy.types.Object": + """Find the armature object in the scene.""" + armatures = [obj for obj in bpy.data.objects if obj.type == "ARMATURE"] + if not armatures: + print("ERROR: No armature found in imported model.") + print(" Available objects:", [obj.name for obj in bpy.data.objects]) + sys.exit(1) + if len(armatures) > 1: + print(f"WARNING: Multiple armatures found, using first: {armatures[0].name}") + return armatures[0] + + +def configure_vrm_metadata(armature: "bpy.types.Object", name: str, author: str) -> None: + """Set VRM 1.0 metadata on the armature.""" + # The VRM Add-on stores metadata as a custom property + armature["vrm_addon_extension"] = { + "spec_version": "1.0", + "vrm0": { + "meta": { + "title": name, + "author": author, + "allowedUserName": "Everyone", + "violentUsageName": "Disallow", + "sexualUsageName": "Disallow", + "commercialUsageName": "Allow", + "licenseName": "MIT", + } + }, + } + print(f" VRM metadata: name={name}, author={author}") + + +def export_vrm(path: str) -> None: + """Export the scene as VRM.""" + bpy.ops.export_scene.vrm(filepath=path) + print(f"Exported VRM: {path}") + + +def main() -> None: + args = parse_args() + + print(f"Converting GLB to VRM:") + print(f" Input: {args.input}") + print(f" Output: {args.output}") + print(f" Name: {args.name}") + + clear_scene() + import_glb(args.input) + + armature = find_armature() + bpy.context.view_layer.objects.active = armature + armature.select_set(True) + + configure_vrm_metadata(armature, args.name, args.author) + export_vrm(args.output) + + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/workflows/README.md b/workflows/README.md new file mode 100644 index 0000000..3144c32 --- /dev/null +++ b/workflows/README.md @@ -0,0 +1,9 @@ +# Workflow JSON files exported from ComfyUI go here. +# +# To export a workflow in API format: +# 1. Open ComfyUI at http://localhost:8188 +# 2. Build your node graph (TRELLIS → mesh processing → UniRig → export) +# 3. Click "Save (API Format)" or use Developer Mode → "Save API workflow" +# 4. Save the JSON file to this directory +# +# The generate.py driver loads these and submits them to ComfyUI's /prompt API.