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)
2026-02-24 05:44:04 -05:00
parent a0c24406bd
commit 202b4e1d61
11 changed files with 1138 additions and 0 deletions
--- a/scripts/ray-join.sh
+++ 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 <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"
--- a/scripts/setup.sh
+++ 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=<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"
--- a/scripts/vrm_export.py
+++ 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()