# Psyche Avatar Engine

> **Status: retired from capability claims** (`182.C.31.06`). The entry point
> named in `project.json` and the `Dockerfile` does not exist, no weights are
> acquirable from this repository, nothing here has an inference test, and the
> default face model is research-licensed. Nothing in this README may be read as
> a shipped capability — see [CAPABILITY-STATUS.md](./CAPABILITY-STATUS.md),
> whose claims are checked against the repository by
> `tests/test_capability_status.py`.

Neural rendering engine for hyper-realistic avatar generation using 3D Gaussian
Splatting.

Part of the Psyche AI Virtual Assistant Platform.

## Overview

The Avatar Engine provides real-time, photorealistic avatar rendering using
state-of-the-art 3D Gaussian Splatting (3DGS) technology. It combines:

- **3D Gaussian Blendshapes** for real-time rendering (370+ FPS)
- **FLAME parametric model** for expression control
- **CUDA-optimized rasterization** for GPU acceleration

## Features

- **High-Quality Rendering**: Photorealistic avatars with view-dependent
  lighting
- **Real-Time Performance**: 30+ FPS at 1024x1024 resolution
- **Expression Control**: 50+ expression blendshapes via FLAME model
- **Monocular Training**: Create avatars from single-camera video
- **GPU Acceleration**: CUDA-optimized for NVIDIA GPUs

## Architecture

```
avatar_engine/
├── config.py           # Configuration management
├── engine.py           # Main orchestration class
├── models/
│   ├── avatar.py       # Avatar representation
│   ├── gaussian.py     # 3D Gaussian Splatting model
│   └── flame.py        # FLAME parametric face model
├── rendering/
│   ├── rasterizer.py   # CUDA rasterization
│   └── renderer.py     # High-level rendering API
├── preprocessing/
│   ├── video.py        # Video preprocessing
│   ├── face_detection.py
│   └── landmark.py     # Facial landmark detection
├── training/
│   ├── trainer.py      # Avatar training pipeline
│   ├── losses.py       # Loss functions (L1, SSIM, LPIPS)
│   └── scheduler.py    # Learning rate scheduling
└── utils/
    ├── camera.py       # Camera utilities
    └── transforms.py   # Geometric transforms
```

## Quick Start

### Installation

```bash
# Using Nx
nx install psyche-avatar-engine

# With CUDA support (recommended)
nx install-cuda psyche-avatar-engine

# Or directly with Poetry
cd apps/psyche/avatar-engine
poetry install
poetry install --extras cuda
```

### Download FLAME Model

The FLAME model is required for expression control. Download from:
https://flame.is.tue.mpg.de/

Place the model files in:

```
models/
└── flame/
    ├── generic_model.pkl
    └── landmark_embedding.npy
```

### Basic Usage

```python
import asyncio
from avatar_engine import AvatarEngine, AvatarConfig
from avatar_engine.models.avatar import AvatarMetadata, AvatarState
from uuid import uuid4

async def main():
    # Initialize engine
    config = AvatarConfig()
    engine = AvatarEngine(config)
    await engine.initialize()

    # Create avatar from video
    metadata = AvatarMetadata(
        name="My Avatar",
        owner_id=uuid4(),
    )
    avatar = await engine.create_avatar(
        video_path=["front.mp4", "left.mp4", "right.mp4"],  # multi-angle capture also supported
        metadata=metadata,
    )

    # Render with expression
    state = AvatarState()
    state.expression[0] = 0.5  # Smile
    output = engine.render(avatar, state)

    # Multi-resolution rendering (e.g., full + half res)
    pyramid = engine.render_multi_resolution(avatar, state, scales=[1.0, 0.5])
    output_half = pyramid[(config.render.width // 2, config.render.height // 2)]

    # Save image
    import torchvision
    torchvision.utils.save_image(
        output["color"].permute(2, 0, 1),
        "output.png"
    )

asyncio.run(main())
```

## NeRF Backup (Instant-NGP)

3DGS is the primary runtime rendering path. For research, bootstrapping, and as
a fallback, the Avatar Engine also includes a pure-PyTorch Instant-NGP-style
NeRF implementation under `avatar_engine.nerf`:

- Hash grid encoding: `HashGridEncoding`
- Model: `InstantNGPModel` (density + color MLPs)
- Acceleration: `OccupancyGrid` (optional)
- Rendering/training: `NerfRenderer`, `NerfTrainer`

Minimal training loop:

```python
import torch
from avatar_engine.nerf import InstantNGPModel, NerfDataset, NerfTrainer, NerfTrainingConfig

device = "cuda"
model = InstantNGPModel().to(device)

dataset = NerfDataset(
    images=images,                 # [N, H, W, 3] in [0,1] (torch or numpy)
    view_matrices=view_matrices,   # [N, 4, 4] world->camera
    proj_matrices=proj_matrices,   # [N, 4, 4] camera->clip (OpenGL-style)
    device="cpu",                  # keep dataset on CPU; batches are moved to `device`
)

trainer = NerfTrainer(
    model=model,
    dataset=dataset,
    config=NerfTrainingConfig(device=device),
)
trainer.train()
trainer.save_model("models/nerf/instant_ngp.pt.gz")
```

NeRF → 3DGS bootstrap (for seeding the Gaussian pipeline):

```python
from avatar_engine.nerf import NerfTo3DGSConfig, nerf_to_gaussian_model

gaussians = nerf_to_gaussian_model(model, config=NerfTo3DGSConfig(num_gaussians=50_000))
```

Hybrid rendering (NeRF as background, 3DGS as foreground):

```python
from avatar_engine.nerf import HybridRenderer

hybrid = HybridRenderer(engine=engine)
out = hybrid.render(avatar=avatar, nerf_model=model)
```

Dynamic NeRF (motion via time-conditioning):

```python
from avatar_engine.nerf import DynamicInstantNGPModel

dyn = DynamicInstantNGPModel()
sigma, rgb = dyn(positions, dirs, time=torch.tensor(0.5))  # time in [0,1]
```

## Configuration

### Environment Variables

```bash
# Device settings
AVATAR_DEVICE=cuda
AVATAR_GPU_ID=0
AVATAR_ENABLE_MIXED_PRECISION=true

# Rendering
AVATAR_RENDER_WIDTH=1024
AVATAR_RENDER_HEIGHT=1024
AVATAR_RENDER_TARGET_FPS=30
AVATAR_RENDER_ENABLE_LOD=true
AVATAR_RENDER_LOD_MIN_PIXEL_RADIUS=0.0
AVATAR_RENDER_LOD_MAX_GAUSSIANS=0
AVATAR_RENDER_ENABLE_TEMPORAL_AA=true
AVATAR_RENDER_ENABLE_FRAME_INTERPOLATION=false
AVATAR_RENDER_TAA_HISTORY_WEIGHT=0.9
AVATAR_RENDER_TAA_JITTER_PIXELS=0.5
AVATAR_RENDER_TAA_ENABLE_HISTORY_CLAMP=true
AVATAR_RENDER_TAA_ALPHA_THRESHOLD=0.001
AVATAR_RENDER_ENABLE_AMBIENT_OCCLUSION=false
AVATAR_RENDER_AMBIENT_OCCLUSION_NUM_SAMPLES=16
AVATAR_RENDER_AMBIENT_OCCLUSION_RADIUS=0.25
AVATAR_RENDER_AMBIENT_OCCLUSION_INTENSITY=1.0
AVATAR_RENDER_AMBIENT_OCCLUSION_BIAS=0.01
AVATAR_RENDER_AMBIENT_OCCLUSION_RESOLUTION_SCALE=0.5
AVATAR_RENDER_AMBIENT_OCCLUSION_ALPHA_THRESHOLD=0.001
AVATAR_RENDER_AMBIENT_OCCLUSION_SEED=1337
AVATAR_RENDER_ENABLE_DYNAMIC_RESOLUTION=false
AVATAR_RENDER_DYNAMIC_RESOLUTION_MIN_SCALE=0.5
AVATAR_RENDER_DYNAMIC_RESOLUTION_MAX_SCALE=1.0
AVATAR_RENDER_DYNAMIC_RESOLUTION_SCALE_STEP=0.05
AVATAR_RENDER_DYNAMIC_RESOLUTION_HYSTERESIS=0.1
AVATAR_RENDER_DYNAMIC_RESOLUTION_EMA_ALPHA=0.2
AVATAR_RENDER_DYNAMIC_RESOLUTION_COOLDOWN_FRAMES=5
AVATAR_RENDER_DYNAMIC_RESOLUTION_ALIGN_TO=16
AVATAR_RENDER_MULTI_RESOLUTION_SCALES=1.0,0.5

# Training
AVATAR_TRAIN_ITERATIONS=30000
AVATAR_TRAIN_LEARNING_RATE=0.0001

# Storage
AVATAR_STORAGE_MODELS_DIR=models
AVATAR_STORAGE_S3_BUCKET=my-bucket
AVATAR_STORAGE_ENABLE_COMPRESSION=true
AVATAR_STORAGE_COMPRESSION_LEVEL=6
AVATAR_STORAGE_PROGRESSIVE_GAUSSIANS=true
```

### Configuration Classes

```python
from avatar_engine.config import (
    AvatarConfig,
    RenderConfig,
    TrainingConfig,
    FLAMEConfig,
)

config = AvatarConfig(
    device="cuda",
    render=RenderConfig(
        width=1024,
        height=1024,
        quality="high",
    ),
    training=TrainingConfig(
        iterations=30000,
        batch_size=1,
    ),
)
```

## Expression Control

The avatar supports 50 expression coefficients from the FLAME model:

| Index | Expression | Description                   |
| ----- | ---------- | ----------------------------- |
| 0-9   | Jaw        | Jaw open, left/right movement |
| 10-19 | Mouth      | Smile, pucker, funnel         |
| 20-29 | Eyes       | Blink, squint, wide           |
| 30-39 | Brows      | Raise, furrow, inner up       |
| 40-49 | Cheeks     | Puff, dimple                  |

### Viseme Support

For lip sync, use the `viseme_weights` field in `AvatarState`:

| Index | Viseme | Phoneme   |
| ----- | ------ | --------- |
| 0     | sil    | Silence   |
| 1     | PP     | p, b, m   |
| 2     | FF     | f, v      |
| 3     | TH     | th        |
| 4     | DD     | t, d      |
| 5     | kk     | k, g      |
| 6     | CH     | ch, j, sh |
| 7     | SS     | s, z      |
| 8     | nn     | n, l      |
| 9     | RR     | r         |
| 10    | aa     | a         |
| 11    | E      | e         |
| 12    | ih     | i         |
| 13    | oh     | o         |
| 14    | ou     | u         |

## Training Pipeline

### Data Requirements

- **Video**: 5-60 seconds, 30 FPS minimum
- **Resolution**: 512x512 minimum (1080p recommended)
- **Lighting**: Even, frontal lighting
- **Movement**: Varied expressions, head poses

### Training Process

1. **Preprocessing**: Face detection, landmark extraction, background removal
2. **FLAME Fitting**: Estimate shape and expression parameters
3. **Gaussian Initialization**: Initialize from FLAME mesh vertices
4. **Optimization**: Jointly optimize Gaussians and blendshapes
5. **Densification**: Add/remove Gaussians based on gradients

### Loss Functions

- **L1 Loss**: Pixel-wise reconstruction (weight: 0.8)
- **SSIM Loss**: Structural similarity (weight: 0.2)
- **LPIPS Loss**: Perceptual similarity (optional)
- **Regularization**: Gaussian scale/opacity regularization

## API Reference

### AvatarEngine

```python
class AvatarEngine:
    async def initialize() -> None
    async def create_avatar(video_path, metadata) -> Avatar
    async def load_avatar(model_path) -> Avatar
    async def save_avatar(avatar, model_path) -> None
    def render(avatar, state, view_matrix, proj_matrix) -> dict
    def render_batch(avatar, states) -> list[dict]
```

### Avatar

```python
class Avatar:
    @property
    def id() -> UUID
    @property
    def is_ready() -> bool
    @property
    def num_gaussians() -> int
    def set_state(state: AvatarState) -> None
    def get_state() -> AvatarState
    async def load(model_path: Path) -> None
    async def save(model_path: Path) -> None
```

### AvatarState

```python
@dataclass
class AvatarState:
    expression: np.ndarray  # [50] expression coefficients
    rotation: np.ndarray    # [3] head rotation (axis-angle)
    translation: np.ndarray # [3] head translation
    jaw_pose: np.ndarray    # [3] jaw rotation
    left_eye_rotation: np.ndarray  # [3]
    right_eye_rotation: np.ndarray # [3]
    left_blink: float
    right_blink: float
    viseme_weights: np.ndarray  # [15] lip sync weights
```

## Performance

### Benchmarks (RTX 4090)

| Resolution | Gaussians | FPS  | Memory |
| ---------- | --------- | ---- | ------ |
| 512x512    | 50K       | 200+ | 1.5 GB |
| 768x768    | 75K       | 120+ | 2.5 GB |
| 1024x1024  | 100K      | 60+  | 4.0 GB |

### Optimization Tips

1. **Reduce Gaussians**: Lower `max_gaussians` for faster rendering
2. **Lower Resolution**: Use 768x768 for real-time streaming
3. **Disable Features**: Turn off temporal AA if not needed
4. **Mixed Precision**: Enable FP16 for training

## Development

### Running the Server

```bash
# Development mode with auto-reload
nx serve psyche-avatar-engine

# Production mode
nx serve-prod psyche-avatar-engine
```

### Running Tests

```bash
# All tests
nx test psyche-avatar-engine

# Skip GPU tests
nx test-no-gpu psyche-avatar-engine

# With coverage
nx test-cov psyche-avatar-engine
```

### Linting & Formatting

```bash
nx lint psyche-avatar-engine
nx format psyche-avatar-engine
```

### Docker

```bash
# Build GPU-enabled image
nx docker-build psyche-avatar-engine

# Run with GPU support
nx docker-run psyche-avatar-engine
```

## Nx Integration

This service is integrated with the Oshun Nx monorepo:

```bash
# Available targets
nx serve psyche-avatar-engine      # Development server
nx serve-prod psyche-avatar-engine # Production server
nx build psyche-avatar-engine      # Build package
nx install psyche-avatar-engine    # Install dependencies
nx install-cuda psyche-avatar-engine # Install with CUDA
nx lint psyche-avatar-engine       # Run linters
nx format psyche-avatar-engine     # Format code
nx test psyche-avatar-engine       # Run tests
nx test-no-gpu psyche-avatar-engine # Tests without GPU
nx test-cov psyche-avatar-engine   # Tests with coverage
nx docker-build psyche-avatar-engine # Build Docker image
nx docker-run psyche-avatar-engine # Run Docker container
nx train psyche-avatar-engine      # Run training
nx download-models psyche-avatar-engine # Download models
```

## License

Proprietary - Oshun Platform

## References

- [3D Gaussian Splatting](https://repo-sam.inria.fr/fungraph/3d-gaussian-splatting/)
- [3D Gaussian Blendshapes](https://github.com/zjumsj/GaussianBlendshapes)
- [FLAME Model](https://flame.is.tue.mpg.de/)
- [GaussianAvatars](https://github.com/ShenhanQian/GaussianAvatars)