feat: scaffold avatar pipeline with ComfyUI driver, MLflow logging, and rclone promotion

- 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)
This commit is contained in:
2026-02-24 05:44:04 -05:00
parent a0c24406bd
commit 202b4e1d61
11 changed files with 1138 additions and 0 deletions

189
README.md
View File

@@ -1,2 +1,191 @@
# avatar-pipeline # 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

View File

@@ -0,0 +1,2 @@
# ComfyUI image-to-VRM avatar generation pipeline
# with TRELLIS + UniRig on desktop Ray worker

227
avatar_pipeline/generate.py Normal file
View File

@@ -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()

View File

@@ -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

131
avatar_pipeline/promote.py Normal file
View File

@@ -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()

40
pyproject.toml Normal file
View File

@@ -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"

7
renovate.json Normal file
View File

@@ -0,0 +1,7 @@
{
"$schema": "https://docs.renovatebot.com/renovate-schema.json",
"extends": [
"local>daviestechlabs/renovate-config",
"local>daviestechlabs/renovate-config:python"
]
}

44
scripts/ray-join.sh Executable file
View File

@@ -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 <ray-head-ip>:<port>"
echo " or: RAY_HEAD_ADDRESS=<ip>:<port> $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"

228
scripts/setup.sh Executable file
View File

@@ -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=<key> secret_access_key=<secret>"
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"

111
scripts/vrm_export.py Normal file
View File

@@ -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()

9
workflows/README.md Normal file
View File

@@ -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.