initial commit

.dockerignore

@@ -0,0 +1,17 @@
+# The .dockerignore file excludes files from the container build process.
+#
+# https://docs.docker.com/engine/reference/builder/#dockerignore-file
+
+# Exclude Git files
+.git
+.github
+.gitignore
+
+# Exclude Python cache files
+__pycache__
+.mypy_cache
+.pytest_cache
+.ruff_cache
+
+# Exclude Python virtual environment
+/venv

.gitattributes

@@ -0,0 +1,13 @@
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.xml filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.pdf filter=lfs diff=lfs merge=lfs -text
+*.mp3 filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.mpg filter=lfs diff=lfs merge=lfs -text
+*.webp filter=lfs diff=lfs merge=lfs -text
+*.webm filter=lfs diff=lfs merge=lfs -text
+*.gif filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,31 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+**/__pycache__/
+*.py[cod]
+**/*.py[cod]
+*$py.class
+
+# Model weights
+**/*.pth
+**/*.onnx
+
+# Ipython notebook
+*.ipynb
+
+# Temporary files or benchmark resources
+animations/*
+tmp/*
+
+# more ignores
+.DS_Store
+*.log
+.idea/
+.vscode/
+*.pyc
+.ipynb_checkpoints
+results/
+data/audio/*.wav
+data/video/*.mp4
+ffmpeg-7.0-amd64-static
+venv/
+.cog/

Dockerfile

@@ -0,0 +1,67 @@
+FROM nvidia/cuda:12.4.0-devel-ubuntu22.04
+
+ARG DEBIAN_FRONTEND=noninteractive
+
+ENV PYTHONUNBUFFERED=1
+
+RUN apt-get update && apt-get install --no-install-recommends -y \
+  build-essential \
+  python3.11 \
+  python3-pip \
+  python3-dev \
+  git \
+  curl \
+  ffmpeg \
+  libglib2.0-0 \
+  libsm6 \
+  libxrender1 \
+  libxext6 \
+  && apt-get clean && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+# Install pget as root
+RUN echo "Installing pget" && \
+    curl -o /usr/local/bin/pget -L 'https://github.com/replicate/pget/releases/download/v0.2.1/pget' && \
+    chmod +x /usr/local/bin/pget
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+# Switch to the "user" user
+USER user
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH
+
+
+# Set home to the user's home directory
+ENV PYTHONPATH=$HOME/app \
+  PYTHONUNBUFFERED=1 \
+  DATA_ROOT=/tmp/data
+
+RUN echo "Installing requirements.txt"
+RUN pip3 install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# yeah.. this is manual for now
+#RUN cd client
+#RUN bun i
+#RUN bun build ./src/index.tsx --outdir ../public/
+
+RUN echo "Installing openmim and mim dependencies"
+RUN pip3 install --no-cache-dir -U openmim
+RUN mim install mmengine
+RUN mim install "mmcv>=2.0.1"
+RUN mim install "mmdet>=3.3.0"
+RUN mim install "mmpose>=1.3.2"
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+EXPOSE 8080
+
+ENV PORT 8080
+
+CMD python3 app.py

LICENSE

@@ -0,0 +1,56 @@
+## For FacePoke (the modifications I made + the server itself)
+
+MIT License
+
+Copyright (c) 2024 Julian Bilcke
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+## For LivePortrait
+
+MIT License
+
+Copyright (c) 2024 Kuaishou Visual Generation and Interaction Center
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+---
+
+The code of InsightFace is released under the MIT License.
+The models of InsightFace are for non-commercial research purposes only.
+
+If you want to use the LivePortrait project for commercial purposes, you
+should remove and replace InsightFace’s detection models to fully comply with
+the MIT license.

README.md

@@ -0,0 +1,119 @@
+---
+title: FacePoke
+emoji: 💬
+colorFrom: yellow
+colorTo: red
+sdk: docker
+pinned: true
+license: mit
+header: mini
+app_file: app.py
+app_port: 8080
+---
+
+# FacePoke
+
+![FacePoke Demo](https://your-demo-image-url-here.gif)
+
+## Table of Contents
+
+- [Introduction](#introduction)
+- [Acknowledgements](#acknowledgements)
+- [Installation](#installation)
+  - [Local Setup](#local-setup)
+  - [Docker Deployment](#docker-deployment)
+- [Development](#development)
+- [Contributing](#contributing)
+- [License](#license)
+
+## Introduction
+
+A real-time head transformation app.
+
+For best performance please run the app from your own machine (local or in the cloud).
+
+**Repository**: [GitHub - jbilcke-hf/FacePoke](https://github.com/jbilcke-hf/FacePoke)
+
+You can try the demo but it is a shared space, latency may be high if there are multiple users or if you live far from the datacenter hosting the Hugging Face Space.
+
+**Live Demo**: [FacePoke on Hugging Face Spaces](https://huggingface.co/spaces/jbilcke-hf/FacePoke)
+
+## Acknowledgements
+
+This project is based on LivePortrait: https://arxiv.org/abs/2407.03168
+
+It uses the face transformation routines from https://github.com/PowerHouseMan/ComfyUI-AdvancedLivePortrait
+
+## Installation
+
+### Local Setup
+
+1. Clone the repository:
+   ```bash
+   git clone https://github.com/jbilcke-hf/FacePoke.git
+   cd FacePoke
+   ```
+
+2. Install Python dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+
+3. Install frontend dependencies:
+   ```bash
+   cd client
+   bun install
+   ```
+
+4. Build the frontend:
+   ```bash
+   bun build ./src/index.tsx --outdir ../public/
+   ```
+
+5. Start the backend server:
+   ```bash
+   python app.py
+   ```
+
+6. Open `http://localhost:8080` in your web browser.
+
+### Docker Deployment
+
+1. Build the Docker image:
+   ```bash
+   docker build -t facepoke .
+   ```
+
+2. Run the container:
+   ```bash
+   docker run -p 8080:8080 facepoke
+   ```
+
+3. To deploy to Hugging Face Spaces:
+   - Fork the repository on GitHub.
+   - Create a new Space on Hugging Face.
+   - Connect your GitHub repository to the Space.
+   - Configure the Space to use the Docker runtime.
+
+## Development
+
+The project structure is organized as follows:
+
+- `app.py`: Main backend server handling WebSocket connections.
+- `engine.py`: Core logic.
+- `loader.py`: Initializes and loads AI models.
+- `client/`: Frontend React application.
+  - `src/`: TypeScript source files.
+  - `public/`: Static assets and built files.
+
+## Contributing
+
+Contributions to FacePoke are welcome! Please read our [Contributing Guidelines](CONTRIBUTING.md) for details on how to submit pull requests, report issues, or request features.
+
+## License
+
+FacePoke is released under the MIT License. See the [LICENSE](LICENSE) file for details.
+
+---
+
+Developed with ❤️ by Julian Bilcke at Hugging Face

app.py

@@ -0,0 +1,264 @@
+"""
+FacePoke API
+
+Author: Julian Bilcke
+Date: September 30, 2024
+"""
+
+import sys
+import asyncio
+import hashlib
+from aiohttp import web, WSMsgType
+import json
+import uuid
+import logging
+import os
+import zipfile
+import signal
+from typing import Dict, Any, List, Optional
+import base64
+import io
+from PIL import Image
+import numpy as np
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Set asyncio logger to DEBUG level
+logging.getLogger("asyncio").setLevel(logging.DEBUG)
+
+logger.debug(f"Python version: {sys.version}")
+
+# SIGSEGV handler
+def SIGSEGV_signal_arises(signalNum, stack):
+    logger.critical(f"{signalNum} : SIGSEGV arises")
+    logger.critical(f"Stack trace: {stack}")
+
+signal.signal(signal.SIGSEGV, SIGSEGV_signal_arises)
+
+from loader import initialize_models
+from engine import Engine, base64_data_uri_to_PIL_Image, create_engine
+
+# Global constants
+DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
+MODELS_DIR = os.path.join(DATA_ROOT, "models")
+
+image_cache: Dict[str, Image.Image] = {}
+
+async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
+    """
+    Handle WebSocket connections for the FacePoke application.
+
+    Args:
+        request (web.Request): The incoming request object.
+
+    Returns:
+        web.WebSocketResponse: The WebSocket response object.
+    """
+    ws = web.WebSocketResponse()
+    await ws.prepare(request)
+
+    session: Optional[FacePokeSession] = None
+    try:
+        logger.info("New WebSocket connection established")
+
+        while True:
+            msg = await ws.receive()
+
+            if msg.type == WSMsgType.TEXT:
+                data = json.loads(msg.data)
+
+                # let's not log user requests, they are heavy
+                #logger.debug(f"Received message: {data}")
+
+                if data['type'] == 'modify_image':
+                    uuid = data.get('uuid')
+                    if not uuid:
+                        logger.warning("Received message without UUID")
+
+                    await handle_modify_image(request, ws, data, uuid)
+
+
+            elif msg.type in (WSMsgType.CLOSE, WSMsgType.ERROR):
+                logger.warning(f"WebSocket connection closed: {msg.type}")
+                break
+
+    except Exception as e:
+        logger.error(f"Error in websocket_handler: {str(e)}")
+        logger.exception("Full traceback:")
+    finally:
+        if session:
+            await session.stop()
+            del active_sessions[session.session_id]
+        logger.info("WebSocket connection closed")
+    return ws
+
+async def handle_modify_image(request: web.Request, ws: web.WebSocketResponse, msg: Dict[str, Any], uuid: str):
+    """
+    Handle the 'modify_image' request.
+
+    Args:
+        request (web.Request): The incoming request object.
+        ws (web.WebSocketResponse): The WebSocket response object.
+        msg (Dict[str, Any]): The message containing the image or image_hash and modification parameters.
+        uuid: A unique identifier for the request.
+    """
+    logger.info("Received modify_image request")
+    try:
+        engine = request.app['engine']
+        image_hash = msg.get('image_hash')
+
+        if image_hash:
+            image_or_hash = image_hash
+        else:
+            image_data = msg['image']
+            image_or_hash = image_data
+
+        modified_image_base64 = await engine.modify_image(image_or_hash, msg['params'])
+
+        await ws.send_json({
+            "type": "modified_image",
+            "image": modified_image_base64,
+            "image_hash": engine.get_image_hash(image_or_hash),
+            "success": True,
+            "uuid": uuid  # Include the UUID in the response
+        })
+        logger.info("Successfully sent modified image")
+    except Exception as e:
+        logger.error(f"Error in modify_image: {str(e)}")
+        await ws.send_json({
+            "type": "modified_image",
+            "success": False,
+            "error": str(e),
+            "uuid": uuid  # Include the UUID even in error responses
+        })
+
+async def index(request: web.Request) -> web.Response:
+    """Serve the index.html file"""
+    content = open(os.path.join(os.path.dirname(__file__), "public", "index.html"), "r").read()
+    return web.Response(content_type="text/html", text=content)
+
+async def js_index(request: web.Request) -> web.Response:
+    """Serve the index.js file"""
+    content = open(os.path.join(os.path.dirname(__file__), "public", "index.js"), "r").read()
+    return web.Response(content_type="application/javascript", text=content)
+
+async def hf_logo(request: web.Request) -> web.Response:
+    """Serve the hf-logo.svg file"""
+    content = open(os.path.join(os.path.dirname(__file__), "public", "hf-logo.svg"), "r").read()
+    return web.Response(content_type="image/svg+xml", text=content)
+
+async def on_shutdown(app: web.Application):
+    """Cleanup function to be called on server shutdown."""
+    logger.info("Server shutdown initiated, cleaning up resources...")
+    for session in list(active_sessions.values()):
+        await session.stop()
+    active_sessions.clear()
+    logger.info("All active sessions have been closed")
+
+    if 'engine' in app:
+        await app['engine'].cleanup()
+        logger.info("Engine instance cleaned up")
+
+    logger.info("Server shutdown complete")
+
+async def initialize_app() -> web.Application:
+    """Initialize and configure the web application."""
+    try:
+        logger.info("Initializing application...")
+        models = await initialize_models()
+        logger.info("🚀 Creating Engine instance...")
+        engine = create_engine(models)
+        logger.info("✅ Engine instance created.")
+
+        app = web.Application()
+        app['engine'] = engine
+
+        app.on_shutdown.append(on_shutdown)
+
+        # Configure routes
+        app.router.add_get("/", index)
+        app.router.add_get("/index.js", js_index)
+        app.router.add_get("/hf-logo.svg", hf_logo)
+        app.router.add_get("/ws", websocket_handler)
+
+        logger.info("Application routes configured")
+
+        return app
+    except Exception as e:
+        logger.error(f"🚨 Error during application initialization: {str(e)}")
+        logger.exception("Full traceback:")
+        raise
+
+async def start_background_tasks(app: web.Application):
+    """
+    Start background tasks for the application.
+
+    Args:
+        app (web.Application): The web application instance.
+    """
+    app['cleanup_task'] = asyncio.create_task(periodic_cleanup(app))
+
+async def cleanup_background_tasks(app: web.Application):
+    """
+    Clean up background tasks when the application is shutting down.
+
+    Args:
+        app (web.Application): The web application instance.
+    """
+    app['cleanup_task'].cancel()
+    await app['cleanup_task']
+
+async def periodic_cleanup(app: web.Application):
+    """
+    Perform periodic cleanup tasks for the application.
+
+    Args:
+        app (web.Application): The web application instance.
+    """
+    while True:
+        try:
+            await asyncio.sleep(3600)  # Run cleanup every hour
+            await cleanup_inactive_sessions(app)
+        except asyncio.CancelledError:
+            break
+        except Exception as e:
+            logger.error(f"Error in periodic cleanup: {str(e)}")
+            logger.exception("Full traceback:")
+
+async def cleanup_inactive_sessions(app: web.Application):
+    """
+    Clean up inactive sessions.
+
+    Args:
+        app (web.Application): The web application instance.
+    """
+    logger.info("Starting cleanup of inactive sessions")
+    inactive_sessions = [
+        session_id for session_id, session in active_sessions.items()
+        if not session.is_running.is_set()
+    ]
+    for session_id in inactive_sessions:
+        session = active_sessions.pop(session_id)
+        await session.stop()
+        logger.info(f"Cleaned up inactive session: {session_id}")
+    logger.info(f"Cleaned up {len(inactive_sessions)} inactive sessions")
+
+def main():
+    """
+    Main function to start the FacePoke application.
+    """
+    try:
+        logger.info("Starting FacePoke application")
+        app = asyncio.run(initialize_app())
+        app.on_startup.append(start_background_tasks)
+        app.on_cleanup.append(cleanup_background_tasks)
+        logger.info("Application initialized, starting web server")
+        web.run_app(app, host="0.0.0.0", port=8080)
+    except Exception as e:
+        logger.critical(f"🚨 FATAL: Failed to start the app: {str(e)}")
+        logger.exception("Full traceback:")
+
+if __name__ == "__main__":
+    main()

build.sh

@@ -0,0 +1,3 @@
+cd client
+bun i
+bun build ./src/index.tsx --outdir ../public/

client/.gitignore

@@ -0,0 +1,175 @@
+# Based on https://raw.githubusercontent.com/github/gitignore/main/Node.gitignore
+
+# Logs
+
+logs
+_.log
+npm-debug.log_
+yarn-debug.log*
+yarn-error.log*
+lerna-debug.log*
+.pnpm-debug.log*
+
+# Caches
+
+.cache
+
+# Diagnostic reports (https://nodejs.org/api/report.html)
+
+report.[0-9]_.[0-9]_.[0-9]_.[0-9]_.json
+
+# Runtime data
+
+pids
+_.pid
+_.seed
+*.pid.lock
+
+# Directory for instrumented libs generated by jscoverage/JSCover
+
+lib-cov
+
+# Coverage directory used by tools like istanbul
+
+coverage
+*.lcov
+
+# nyc test coverage
+
+.nyc_output
+
+# Grunt intermediate storage (https://gruntjs.com/creating-plugins#storing-task-files)
+
+.grunt
+
+# Bower dependency directory (https://bower.io/)
+
+bower_components
+
+# node-waf configuration
+
+.lock-wscript
+
+# Compiled binary addons (https://nodejs.org/api/addons.html)
+
+build/Release
+
+# Dependency directories
+
+node_modules/
+jspm_packages/
+
+# Snowpack dependency directory (https://snowpack.dev/)
+
+web_modules/
+
+# TypeScript cache
+
+*.tsbuildinfo
+
+# Optional npm cache directory
+
+.npm
+
+# Optional eslint cache
+
+.eslintcache
+
+# Optional stylelint cache
+
+.stylelintcache
+
+# Microbundle cache
+
+.rpt2_cache/
+.rts2_cache_cjs/
+.rts2_cache_es/
+.rts2_cache_umd/
+
+# Optional REPL history
+
+.node_repl_history
+
+# Output of 'npm pack'
+
+*.tgz
+
+# Yarn Integrity file
+
+.yarn-integrity
+
+# dotenv environment variable files
+
+.env
+.env.development.local
+.env.test.local
+.env.production.local
+.env.local
+
+# parcel-bundler cache (https://parceljs.org/)
+
+.parcel-cache
+
+# Next.js build output
+
+.next
+out
+
+# Nuxt.js build / generate output
+
+.nuxt
+dist
+
+# Gatsby files
+
+# Comment in the public line in if your project uses Gatsby and not Next.js
+
+# https://nextjs.org/blog/next-9-1#public-directory-support
+
+# public
+
+# vuepress build output
+
+.vuepress/dist
+
+# vuepress v2.x temp and cache directory
+
+.temp
+
+# Docusaurus cache and generated files
+
+.docusaurus
+
+# Serverless directories
+
+.serverless/
+
+# FuseBox cache
+
+.fusebox/
+
+# DynamoDB Local files
+
+.dynamodb/
+
+# TernJS port file
+
+.tern-port
+
+# Stores VSCode versions used for testing VSCode extensions
+
+.vscode-test
+
+# yarn v2
+
+.yarn/cache
+.yarn/unplugged
+.yarn/build-state.yml
+.yarn/install-state.gz
+.pnp.*
+
+# IntelliJ based IDEs
+.idea
+
+# Finder (MacOS) folder config
+.DS_Store

client/README.md

@@ -0,0 +1,13 @@
+# FacePoke.js
+
+To install dependencies:
+
+```bash
+bun i
+```
+
+To build:
+
+```bash
+bun build ./src/index.tsx --outdir ../public
+```

client/package.json

@@ -0,0 +1,35 @@
+{
+  "name": "@aitube/facepoke",
+  "module": "src/index.ts",
+  "type": "module",
+  "scripts": {
+    "build": "bun build ./src/index.tsx --outdir ../public/"
+  },
+  "devDependencies": {
+    "@types/bun": "latest"
+  },
+  "peerDependencies": {
+    "typescript": "^5.0.0"
+  },
+  "dependencies": {
+    "@mediapipe/tasks-vision": "^0.10.16",
+    "@radix-ui/react-icons": "^1.3.0",
+    "@types/lodash": "^4.17.10",
+    "@types/react": "^18.3.9",
+    "@types/react-dom": "^18.3.0",
+    "@types/uuid": "^10.0.0",
+    "beautiful-react-hooks": "^5.0.2",
+    "class-variance-authority": "^0.7.0",
+    "clsx": "^2.1.1",
+    "lodash": "^4.17.21",
+    "lucide-react": "^0.446.0",
+    "react": "^18.3.1",
+    "react-dom": "^18.3.1",
+    "tailwind-merge": "^2.5.2",
+    "tailwindcss": "^3.4.13",
+    "tailwindcss-animate": "^1.0.7",
+    "usehooks-ts": "^3.1.0",
+    "uuid": "^10.0.0",
+    "zustand": "^5.0.0-rc.2"
+  }
+}

client/src/app.tsx

@@ -0,0 +1,190 @@
+import React, { useState, useEffect, useRef, useCallback, useMemo } from 'react';
+import { RotateCcw } from 'lucide-react';
+
+import { Alert, AlertDescription, AlertTitle } from '@/components/ui/alert';
+import { truncateFileName } from './lib/utils';
+import { useFaceLandmarkDetection } from './hooks/useFaceLandmarkDetection';
+import { PoweredBy } from './components/PoweredBy';
+import { Spinner } from './components/Spinner';
+import { DoubleCard } from './components/DoubleCard';
+import { useFacePokeAPI } from './hooks/useFacePokeAPI';
+import { Layout } from './layout';
+import { useMainStore } from './hooks/useMainStore';
+import { convertImageToBase64 } from './lib/convertImageToBase64';
+
+export function App() {
+  const error = useMainStore(s => s.error);
+  const setError = useMainStore(s => s.setError);
+  const imageFile = useMainStore(s => s.imageFile);
+  const setImageFile = useMainStore(s => s.setImageFile);
+  const originalImage = useMainStore(s => s.originalImage);
+  const setOriginalImage = useMainStore(s => s.setOriginalImage);
+  const previewImage = useMainStore(s => s.previewImage);
+  const setPreviewImage = useMainStore(s => s.setPreviewImage);
+  const resetImage = useMainStore(s => s.resetImage);
+
+  const {
+    status,
+    setStatus,
+    isDebugMode,
+    setIsDebugMode,
+    interruptMessage,
+  } = useFacePokeAPI()
+
+  // State for face detection
+  const {
+    canvasRef,
+    canvasRefCallback,
+    mediaPipeRef,
+    faceLandmarks,
+    isMediaPipeReady,
+    blendShapes,
+
+    setFaceLandmarks,
+    setBlendShapes,
+
+    handleMouseDown,
+    handleMouseUp,
+    handleMouseMove,
+    handleMouseEnter,
+    handleMouseLeave,
+    currentOpacity
+  } = useFaceLandmarkDetection()
+
+  // Refs
+  const videoRef = useRef<HTMLDivElement>(null);
+
+  // Handle file change
+  const handleFileChange = useCallback(async (event: React.ChangeEvent<HTMLInputElement>) => {
+    const files = event.target.files;
+    if (files && files[0]) {
+      setImageFile(files[0]);
+      setStatus(`File selected: ${truncateFileName(files[0].name, 16)}`);
+
+      try {
+        const image = await convertImageToBase64(files[0]);
+        setPreviewImage(image);
+        setOriginalImage(image);
+      } catch (err) {
+        console.log(`failed to convert the image: `, err);
+        setImageFile(null);
+        setStatus('');
+        setPreviewImage('');
+        setOriginalImage('');
+        setFaceLandmarks([]);
+        setBlendShapes([]);
+      }
+    } else {
+      setImageFile(null);
+      setStatus('');
+      setPreviewImage('');
+      setOriginalImage('');
+      setFaceLandmarks([]);
+      setBlendShapes([]);
+    }
+  }, [isMediaPipeReady, setImageFile, setPreviewImage, setOriginalImage, setFaceLandmarks, setBlendShapes, setStatus]);
+
+  const canDisplayBlendShapes = false
+
+  // Display blend shapes
+  const displayBlendShapes = useMemo(() => (
+      <div className="mt-4">
+        <h3 className="text-lg font-semibold mb-2">Blend Shapes</h3>
+        <ul className="space-y-1">
+          {(blendShapes?.[0]?.categories || []).map((shape, index) => (
+            <li key={index} className="flex items-center">
+              <span className="w-32 text-sm">{shape.categoryName || shape.displayName}</span>
+              <div className="w-full bg-gray-200 rounded-full h-2.5">
+                <div
+                  className="bg-blue-600 h-2.5 rounded-full"
+                  style={{ width: `${shape.score * 100}%` }}
+                ></div>
+              </div>
+              <span className="ml-2 text-sm">{shape.score.toFixed(2)}</span>
+            </li>
+          ))}
+        </ul>
+      </div>
+  ), [JSON.stringify(blendShapes)])
+
+  // JSX
+  return (
+    <Layout>
+        {error && (
+          <Alert variant="destructive">
+            <AlertTitle>Error</AlertTitle>
+            <AlertDescription>{error}</AlertDescription>
+          </Alert>
+        )}
+        {interruptMessage && (
+          <Alert>
+            <AlertTitle>Notice</AlertTitle>
+            <AlertDescription>{interruptMessage}</AlertDescription>
+          </Alert>
+        )}
+        <div className="mb-5 relative">
+          <div className="flex flex-row items-center justify-between w-full">
+            <div className="relative">
+              <input
+                id="imageInput"
+                type="file"
+                accept="image/*"
+                onChange={handleFileChange}
+                className="hidden"
+                disabled={!isMediaPipeReady}
+              />
+              <label
+                htmlFor="imageInput"
+                className={`cursor-pointer inline-flex items-center px-3 py-1.5 border border-transparent text-sm font-medium rounded-md text-white ${
+                  isMediaPipeReady ? 'bg-gray-600 hover:bg-gray-500' : 'bg-gray-500 cursor-not-allowed'
+                } focus:outline-none focus:ring-2 focus:ring-offset-2 focus:ring-gray-500 shadow-xl`}
+              >
+                <Spinner />
+                {imageFile ? truncateFileName(imageFile.name, 32) : (isMediaPipeReady ? 'Choose an image' : 'Initializing...')}
+              </label>
+            </div>
+            {previewImage && <label className="mt-4 flex items-center">
+              <input
+                type="checkbox"
+                checked={isDebugMode}
+                onChange={(e) => setIsDebugMode(e.target.checked)}
+                className="mr-2"
+              />
+              Show face landmarks on hover
+            </label>}
+          </div>
+          {previewImage && (
+            <div className="mt-5 relative shadow-2xl rounded-xl overflow-hidden">
+              <img
+                src={previewImage}
+                alt="Preview"
+                className="w-full"
+              />
+              <canvas
+                ref={canvasRefCallback}
+                className="absolute top-0 left-0 w-full h-full select-none"
+                onMouseEnter={handleMouseEnter}
+                onMouseLeave={handleMouseLeave}
+                onMouseDown={handleMouseDown}
+                onMouseUp={handleMouseUp}
+                onMouseMove={handleMouseMove}
+                style={{
+                  position: 'absolute',
+                  top: 0,
+                  left: 0,
+                  width: '100%',
+                  height: '100%',
+                  opacity: isDebugMode ? currentOpacity : 0.0,
+                  transition: 'opacity 0.2s ease-in-out'
+                }}
+
+              />
+            </div>
+          )}
+          {canDisplayBlendShapes && displayBlendShapes}
+        </div>
+        <PoweredBy />
+
+    </Layout>
+  );
+}

client/src/components/DoubleCard.tsx

@@ -0,0 +1,18 @@
+import React, { type ReactNode } from 'react';
+
+export function DoubleCard({ children }: { children: ReactNode }) {
+  return (
+    <>
+      <div className="absolute inset-0 bg-gradient-to-r from-cyan-200 to-sky-300 shadow-2xl transform -skew-y-6 sm:skew-y-0 sm:-rotate-6 sm:rounded-3xl" style={{ borderTop: "solid 2px rgba(255, 255, 255, 0.2)" }}></div>
+      <div className="relative px-5 py-8 bg-gradient-to-r from-cyan-100 to-sky-200 shadow-2xl sm:rounded-3xl sm:p-12" style={{ borderTop: "solid 2px #ffffff33" }}>
+        <div className="max-w-lg mx-auto">
+          <div className="divide-y divide-gray-200">
+            <div className="text-lg leading-7 space-y-5 text-gray-700 sm:text-xl sm:leading-8">
+              {children}
+            </div>
+          </div>
+        </div>
+      </div>
+    </>
+  );
+}

client/src/components/PoweredBy.tsx

@@ -0,0 +1,17 @@
+export function PoweredBy() {
+  return (
+    <div className="flex flex-row items-center justify-center font-sans mt-4 w-full">
+      {/*<span className="text-neutral-900 text-sm"
+        style={{ textShadow: "rgb(255 255 255 / 80%) 0px 0px 2px" }}>
+        Best hosted on
+      </span>*/}
+      <span className="ml-2 mr-1">
+        <img src="/hf-logo.svg" alt="Hugging Face" className="w-5 h-5" />
+      </span>
+      <span className="text-neutral-900 text-sm font-semibold"
+        style={{ textShadow: "rgb(255 255 255 / 80%) 0px 0px 2px" }}>
+        Hugging Face
+      </span>
+    </div>
+  )
+}

client/src/components/Spinner.tsx

@@ -0,0 +1,7 @@
+export function Spinner() {
+  return (
+    <svg className="mr-3 h-6 w-6" fill="none" viewBox="0 0 24 24" stroke="currentColor">
+      <path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M4 16l4.586-4.586a2 2 0 012.828 0L16 16m-2-2l1.586-1.586a2 2 0 012.828 0L20 14m-6-6h.01M6 20h12a2 2 0 002-2V6a2 2 0 00-2-2H6a2 2 0 00-2 2v12a2 2 0 002 2z" />
+    </svg>
+  )
+}

client/src/components/Title.tsx

@@ -0,0 +1,8 @@
+export function Title() {
+  return (
+    <h2 className="bg-gradient-to-bl from-sky-500 to-sky-800 bg-clip-text text-5xl font-extrabold text-transparent leading-normal text-center"
+      style={{ textShadow: "rgb(176 229 255 / 61%) 0px 0px 2px" }}>
+      💬 FacePoke
+    </h2>
+  )
+}

client/src/components/ui/alert.tsx

@@ -0,0 +1,59 @@
+import * as React from "react"
+import { cva, type VariantProps } from "class-variance-authority"
+
+import { cn } from "@/lib/utils"
+
+const alertVariants = cva(
+  "relative w-full rounded-lg border p-4 [&>svg~*]:pl-7 [&>svg+div]:translate-y-[-3px] [&>svg]:absolute [&>svg]:left-4 [&>svg]:top-4 [&>svg]:text-foreground",
+  {
+    variants: {
+      variant: {
+        default: "bg-background text-foreground",
+        destructive:
+          "border-destructive/50 text-destructive dark:border-destructive [&>svg]:text-destructive",
+      },
+    },
+    defaultVariants: {
+      variant: "default",
+    },
+  }
+)
+
+const Alert = React.forwardRef<
+  HTMLDivElement,
+  React.HTMLAttributes<HTMLDivElement> & VariantProps<typeof alertVariants>
+>(({ className, variant, ...props }, ref) => (
+  <div
+    ref={ref}
+    role="alert"
+    className={cn(alertVariants({ variant }), className)}
+    {...props}
+  />
+))
+Alert.displayName = "Alert"
+
+const AlertTitle = React.forwardRef<
+  HTMLParagraphElement,
+  React.HTMLAttributes<HTMLHeadingElement>
+>(({ className, ...props }, ref) => (
+  <h5
+    ref={ref}
+    className={cn("mb-1 font-medium leading-none tracking-tight", className)}
+    {...props}
+  />
+))
+AlertTitle.displayName = "AlertTitle"
+
+const AlertDescription = React.forwardRef<
+  HTMLParagraphElement,
+  React.HTMLAttributes<HTMLParagraphElement>
+>(({ className, ...props }, ref) => (
+  <div
+    ref={ref}
+    className={cn("text-sm [&_p]:leading-relaxed", className)}
+    {...props}
+  />
+))
+AlertDescription.displayName = "AlertDescription"
+
+export { Alert, AlertTitle, AlertDescription }

client/src/hooks/landmarks.ts

@@ -0,0 +1,520 @@
+import * as vision from '@mediapipe/tasks-vision';
+
+// Define unique colors for each landmark group
+export const landmarkColors: { [key: string]: string } = {
+  lips: '#FF0000',
+  leftEye: '#00FF00',
+  leftEyebrow: '#0000FF',
+  leftIris: '#FFFF00',
+  rightEye: '#FF00FF',
+  rightEyebrow: '#00FFFF',
+  rightIris: '#FFA500',
+  faceOval: '#800080',
+  tesselation: '#C0C0C0',
+};
+
+// Define landmark groups with their semantic names
+export const landmarkGroups: { [key: string]: any } = {
+  lips: vision.FaceLandmarker.FACE_LANDMARKS_LIPS,
+  leftEye: vision.FaceLandmarker.FACE_LANDMARKS_LEFT_EYE,
+  leftEyebrow: vision.FaceLandmarker.FACE_LANDMARKS_LEFT_EYEBROW,
+  leftIris: vision.FaceLandmarker.FACE_LANDMARKS_LEFT_IRIS,
+  rightEye: vision.FaceLandmarker.FACE_LANDMARKS_RIGHT_EYE,
+  rightEyebrow: vision.FaceLandmarker.FACE_LANDMARKS_RIGHT_EYEBROW,
+  rightIris: vision.FaceLandmarker.FACE_LANDMARKS_RIGHT_IRIS,
+  faceOval: vision.FaceLandmarker.FACE_LANDMARKS_FACE_OVAL,
+  // tesselation: vision.FaceLandmarker.FACE_LANDMARKS_TESSELATION,
+};
+
+export const FACEMESH_LIPS = Object.freeze(new Set([[61, 146], [146, 91], [91, 181], [181, 84], [84, 17],
+                           [17, 314], [314, 405], [405, 321], [321, 375],
+                           [375, 291], [61, 185], [185, 40], [40, 39], [39, 37],
+                           [37, 0], [0, 267],
+                           [267, 269], [269, 270], [270, 409], [409, 291],
+                           [78, 95], [95, 88], [88, 178], [178, 87], [87, 14],
+                           [14, 317], [317, 402], [402, 318], [318, 324],
+                           [324, 308], [78, 191], [191, 80], [80, 81], [81, 82],
+                           [82, 13], [13, 312], [312, 311], [311, 310],
+                           [310, 415], [415, 308]]))
+
+export const FACEMESH_LEFT_EYE = Object.freeze(new Set([[263, 249], [249, 390], [390, 373], [373, 374],
+                               [374, 380], [380, 381], [381, 382], [382, 362],
+                               [263, 466], [466, 388], [388, 387], [387, 386],
+                               [386, 385], [385, 384], [384, 398], [398, 362]]))
+
+export const FACEMESH_LEFT_IRIS = Object.freeze(new Set([[474, 475], [475, 476], [476, 477],
+                                 [477, 474]]))
+
+export const FACEMESH_LEFT_EYEBROW = Object.freeze(new Set([[276, 283], [283, 282], [282, 295],
+                                   [295, 285], [300, 293], [293, 334],
+                                   [334, 296], [296, 336]]))
+
+export const FACEMESH_RIGHT_EYE = Object.freeze(new Set([[33, 7], [7, 163], [163, 144], [144, 145],
+                                [145, 153], [153, 154], [154, 155], [155, 133],
+                                [33, 246], [246, 161], [161, 160], [160, 159],
+                                [159, 158], [158, 157], [157, 173], [173, 133]]))
+
+export const FACEMESH_RIGHT_EYEBROW = Object.freeze(new Set([[46, 53], [53, 52], [52, 65], [65, 55],
+                                    [70, 63], [63, 105], [105, 66], [66, 107]]))
+
+export const FACEMESH_RIGHT_IRIS = Object.freeze(new Set([[469, 470], [470, 471], [471, 472],
+                                 [472, 469]]))
+
+export const FACEMESH_FACE_OVAL = Object.freeze(new Set([[10, 338], [338, 297], [297, 332], [332, 284],
+                                [284, 251], [251, 389], [389, 356], [356, 454],
+                                [454, 323], [323, 361], [361, 288], [288, 397],
+                                [397, 365], [365, 379], [379, 378], [378, 400],
+                                [400, 377], [377, 152], [152, 148], [148, 176],
+                                [176, 149], [149, 150], [150, 136], [136, 172],
+                                [172, 58], [58, 132], [132, 93], [93, 234],
+                                [234, 127], [127, 162], [162, 21], [21, 54],
+                                [54, 103], [103, 67], [67, 109], [109, 10]]))
+
+export const FACEMESH_NOSE = Object.freeze(new Set([[168, 6], [6, 197], [197, 195], [195, 5],
+                           [5, 4], [4, 1], [1, 19], [19, 94], [94, 2], [98, 97],
+                           [97, 2], [2, 326], [326, 327], [327, 294],
+                           [294, 278], [278, 344], [344, 440], [440, 275],
+                           [275, 4], [4, 45], [45, 220], [220, 115], [115, 48],
+                           [48, 64], [64, 98]]))
+
+export const FACEMESH_CONTOURS = Object.freeze(new Set([
+  ...FACEMESH_LIPS,
+  ...FACEMESH_LEFT_EYE,
+  ...FACEMESH_LEFT_EYEBROW,
+  ...FACEMESH_RIGHT_EYE,
+  ...FACEMESH_RIGHT_EYEBROW,
+  ...FACEMESH_FACE_OVAL
+]));
+
+export const FACEMESH_IRISES = Object.freeze(new Set([
+  ...FACEMESH_LEFT_IRIS,
+  ...FACEMESH_RIGHT_IRIS
+]));
+
+export const FACEMESH_TESSELATION = Object.freeze(new Set([
+    [127, 34],  [34, 139],  [139, 127], [11, 0],    [0, 37],    [37, 11],
+    [232, 231], [231, 120], [120, 232], [72, 37],   [37, 39],   [39, 72],
+    [128, 121], [121, 47],  [47, 128],  [232, 121], [121, 128], [128, 232],
+    [104, 69],  [69, 67],   [67, 104],  [175, 171], [171, 148], [148, 175],
+    [118, 50],  [50, 101],  [101, 118], [73, 39],   [39, 40],   [40, 73],
+    [9, 151],   [151, 108], [108, 9],   [48, 115],  [115, 131], [131, 48],
+    [194, 204], [204, 211], [211, 194], [74, 40],   [40, 185],  [185, 74],
+    [80, 42],   [42, 183],  [183, 80],  [40, 92],   [92, 186],  [186, 40],
+    [230, 229], [229, 118], [118, 230], [202, 212], [212, 214], [214, 202],
+    [83, 18],   [18, 17],   [17, 83],   [76, 61],   [61, 146],  [146, 76],
+    [160, 29],  [29, 30],   [30, 160],  [56, 157],  [157, 173], [173, 56],
+    [106, 204], [204, 194], [194, 106], [135, 214], [214, 192], [192, 135],
+    [203, 165], [165, 98],  [98, 203],  [21, 71],   [71, 68],   [68, 21],
+    [51, 45],   [45, 4],    [4, 51],    [144, 24],  [24, 23],   [23, 144],
+    [77, 146],  [146, 91],  [91, 77],   [205, 50],  [50, 187],  [187, 205],
+    [201, 200], [200, 18],  [18, 201],  [91, 106],  [106, 182], [182, 91],
+    [90, 91],   [91, 181],  [181, 90],  [85, 84],   [84, 17],   [17, 85],
+    [206, 203], [203, 36],  [36, 206],  [148, 171], [171, 140], [140, 148],
+    [92, 40],   [40, 39],   [39, 92],   [193, 189], [189, 244], [244, 193],
+    [159, 158], [158, 28],  [28, 159],  [247, 246], [246, 161], [161, 247],
+    [236, 3],   [3, 196],   [196, 236], [54, 68],   [68, 104],  [104, 54],
+    [193, 168], [168, 8],   [8, 193],   [117, 228], [228, 31],  [31, 117],
+    [189, 193], [193, 55],  [55, 189],  [98, 97],   [97, 99],   [99, 98],
+    [126, 47],  [47, 100],  [100, 126], [166, 79],  [79, 218],  [218, 166],
+    [155, 154], [154, 26],  [26, 155],  [209, 49],  [49, 131],  [131, 209],
+    [135, 136], [136, 150], [150, 135], [47, 126],  [126, 217], [217, 47],
+    [223, 52],  [52, 53],   [53, 223],  [45, 51],   [51, 134],  [134, 45],
+    [211, 170], [170, 140], [140, 211], [67, 69],   [69, 108],  [108, 67],
+    [43, 106],  [106, 91],  [91, 43],   [230, 119], [119, 120], [120, 230],
+    [226, 130], [130, 247], [247, 226], [63, 53],   [53, 52],   [52, 63],
+    [238, 20],  [20, 242],  [242, 238], [46, 70],   [70, 156],  [156, 46],
+    [78, 62],   [62, 96],   [96, 78],   [46, 53],   [53, 63],   [63, 46],
+    [143, 34],  [34, 227],  [227, 143], [123, 117], [117, 111], [111, 123],
+    [44, 125],  [125, 19],  [19, 44],   [236, 134], [134, 51],  [51, 236],
+    [216, 206], [206, 205], [205, 216], [154, 153], [153, 22],  [22, 154],
+    [39, 37],   [37, 167],  [167, 39],  [200, 201], [201, 208], [208, 200],
+    [36, 142],  [142, 100], [100, 36],  [57, 212],  [212, 202], [202, 57],
+    [20, 60],   [60, 99],   [99, 20],   [28, 158],  [158, 157], [157, 28],
+    [35, 226],  [226, 113], [113, 35],  [160, 159], [159, 27],  [27, 160],
+    [204, 202], [202, 210], [210, 204], [113, 225], [225, 46],  [46, 113],
+    [43, 202],  [202, 204], [204, 43],  [62, 76],   [76, 77],   [77, 62],
+    [137, 123], [123, 116], [116, 137], [41, 38],   [38, 72],   [72, 41],
+    [203, 129], [129, 142], [142, 203], [64, 98],   [98, 240],  [240, 64],
+    [49, 102],  [102, 64],  [64, 49],   [41, 73],   [73, 74],   [74, 41],
+    [212, 216], [216, 207], [207, 212], [42, 74],   [74, 184],  [184, 42],
+    [169, 170], [170, 211], [211, 169], [170, 149], [149, 176], [176, 170],
+    [105, 66],  [66, 69],   [69, 105],  [122, 6],   [6, 168],   [168, 122],
+    [123, 147], [147, 187], [187, 123], [96, 77],   [77, 90],   [90, 96],
+    [65, 55],   [55, 107],  [107, 65],  [89, 90],   [90, 180],  [180, 89],
+    [101, 100], [100, 120], [120, 101], [63, 105],  [105, 104], [104, 63],
+    [93, 137],  [137, 227], [227, 93],  [15, 86],   [86, 85],   [85, 15],
+    [129, 102], [102, 49],  [49, 129],  [14, 87],   [87, 86],   [86, 14],
+    [55, 8],    [8, 9],     [9, 55],    [100, 47],  [47, 121],  [121, 100],
+    [145, 23],  [23, 22],   [22, 145],  [88, 89],   [89, 179],  [179, 88],
+    [6, 122],   [122, 196], [196, 6],   [88, 95],   [95, 96],   [96, 88],
+    [138, 172], [172, 136], [136, 138], [215, 58],  [58, 172],  [172, 215],
+    [115, 48],  [48, 219],  [219, 115], [42, 80],   [80, 81],   [81, 42],
+    [195, 3],   [3, 51],    [51, 195],  [43, 146],  [146, 61],  [61, 43],
+    [171, 175], [175, 199], [199, 171], [81, 82],   [82, 38],   [38, 81],
+    [53, 46],   [46, 225],  [225, 53],  [144, 163], [163, 110], [110, 144],
+    [52, 65],   [65, 66],   [66, 52],   [229, 228], [228, 117], [117, 229],
+    [34, 127],  [127, 234], [234, 34],  [107, 108], [108, 69],  [69, 107],
+    [109, 108], [108, 151], [151, 109], [48, 64],   [64, 235],  [235, 48],
+    [62, 78],   [78, 191],  [191, 62],  [129, 209], [209, 126], [126, 129],
+    [111, 35],  [35, 143],  [143, 111], [117, 123], [123, 50],  [50, 117],
+    [222, 65],  [65, 52],   [52, 222],  [19, 125],  [125, 141], [141, 19],
+    [221, 55],  [55, 65],   [65, 221],  [3, 195],   [195, 197], [197, 3],
+    [25, 7],    [7, 33],    [33, 25],   [220, 237], [237, 44],  [44, 220],
+    [70, 71],   [71, 139],  [139, 70],  [122, 193], [193, 245], [245, 122],
+    [247, 130], [130, 33],  [33, 247],  [71, 21],   [21, 162],  [162, 71],
+    [170, 169], [169, 150], [150, 170], [188, 174], [174, 196], [196, 188],
+    [216, 186], [186, 92],  [92, 216],  [2, 97],    [97, 167],  [167, 2],
+    [141, 125], [125, 241], [241, 141], [164, 167], [167, 37],  [37, 164],
+    [72, 38],   [38, 12],   [12, 72],   [38, 82],   [82, 13],   [13, 38],
+    [63, 68],   [68, 71],   [71, 63],   [226, 35],  [35, 111],  [111, 226],
+    [101, 50],  [50, 205],  [205, 101], [206, 92],  [92, 165],  [165, 206],
+    [209, 198], [198, 217], [217, 209], [165, 167], [167, 97],  [97, 165],
+    [220, 115], [115, 218], [218, 220], [133, 112], [112, 243], [243, 133],
+    [239, 238], [238, 241], [241, 239], [214, 135], [135, 169], [169, 214],
+    [190, 173], [173, 133], [133, 190], [171, 208], [208, 32],  [32, 171],
+    [125, 44],  [44, 237],  [237, 125], [86, 87],   [87, 178],  [178, 86],
+    [85, 86],   [86, 179],  [179, 85],  [84, 85],   [85, 180],  [180, 84],
+    [83, 84],   [84, 181],  [181, 83],  [201, 83],  [83, 182],  [182, 201],
+    [137, 93],  [93, 132],  [132, 137], [76, 62],   [62, 183],  [183, 76],
+    [61, 76],   [76, 184],  [184, 61],  [57, 61],   [61, 185],  [185, 57],
+    [212, 57],  [57, 186],  [186, 212], [214, 207], [207, 187], [187, 214],
+    [34, 143],  [143, 156], [156, 34],  [79, 239],  [239, 237], [237, 79],
+    [123, 137], [137, 177], [177, 123], [44, 1],    [1, 4],     [4, 44],
+    [201, 194], [194, 32],  [32, 201],  [64, 102],  [102, 129], [129, 64],
+    [213, 215], [215, 138], [138, 213], [59, 166],  [166, 219], [219, 59],
+    [242, 99],  [99, 97],   [97, 242],  [2, 94],    [94, 141],  [141, 2],
+    [75, 59],   [59, 235],  [235, 75],  [24, 110],  [110, 228], [228, 24],
+    [25, 130],  [130, 226], [226, 25],  [23, 24],   [24, 229],  [229, 23],
+    [22, 23],   [23, 230],  [230, 22],  [26, 22],   [22, 231],  [231, 26],
+    [112, 26],  [26, 232],  [232, 112], [189, 190], [190, 243], [243, 189],
+    [221, 56],  [56, 190],  [190, 221], [28, 56],   [56, 221],  [221, 28],
+    [27, 28],   [28, 222],  [222, 27],  [29, 27],   [27, 223],  [223, 29],
+    [30, 29],   [29, 224],  [224, 30],  [247, 30],  [30, 225],  [225, 247],
+    [238, 79],  [79, 20],   [20, 238],  [166, 59],  [59, 75],   [75, 166],
+    [60, 75],   [75, 240],  [240, 60],  [147, 177], [177, 215], [215, 147],
+    [20, 79],   [79, 166],  [166, 20],  [187, 147], [147, 213], [213, 187],
+    [112, 233], [233, 244], [244, 112], [233, 128], [128, 245], [245, 233],
+    [128, 114], [114, 188], [188, 128], [114, 217], [217, 174], [174, 114],
+    [131, 115], [115, 220], [220, 131], [217, 198], [198, 236], [236, 217],
+    [198, 131], [131, 134], [134, 198], [177, 132], [132, 58],  [58, 177],
+    [143, 35],  [35, 124],  [124, 143], [110, 163], [163, 7],   [7, 110],
+    [228, 110], [110, 25],  [25, 228],  [356, 389], [389, 368], [368, 356],
+    [11, 302],  [302, 267], [267, 11],  [452, 350], [350, 349], [349, 452],
+    [302, 303], [303, 269], [269, 302], [357, 343], [343, 277], [277, 357],
+    [452, 453], [453, 357], [357, 452], [333, 332], [332, 297], [297, 333],
+    [175, 152], [152, 377], [377, 175], [347, 348], [348, 330], [330, 347],
+    [303, 304], [304, 270], [270, 303], [9, 336],   [336, 337], [337, 9],
+    [278, 279], [279, 360], [360, 278], [418, 262], [262, 431], [431, 418],
+    [304, 408], [408, 409], [409, 304], [310, 415], [415, 407], [407, 310],
+    [270, 409], [409, 410], [410, 270], [450, 348], [348, 347], [347, 450],
+    [422, 430], [430, 434], [434, 422], [313, 314], [314, 17],  [17, 313],
+    [306, 307], [307, 375], [375, 306], [387, 388], [388, 260], [260, 387],
+    [286, 414], [414, 398], [398, 286], [335, 406], [406, 418], [418, 335],
+    [364, 367], [367, 416], [416, 364], [423, 358], [358, 327], [327, 423],
+    [251, 284], [284, 298], [298, 251], [281, 5],   [5, 4],     [4, 281],
+    [373, 374], [374, 253], [253, 373], [307, 320], [320, 321], [321, 307],
+    [425, 427], [427, 411], [411, 425], [421, 313], [313, 18],  [18, 421],
+    [321, 405], [405, 406], [406, 321], [320, 404], [404, 405], [405, 320],
+    [315, 16],  [16, 17],   [17, 315],  [426, 425], [425, 266], [266, 426],
+    [377, 400], [400, 369], [369, 377], [322, 391], [391, 269], [269, 322],
+    [417, 465], [465, 464], [464, 417], [386, 257], [257, 258], [258, 386],
+    [466, 260], [260, 388], [388, 466], [456, 399], [399, 419], [419, 456],
+    [284, 332], [332, 333], [333, 284], [417, 285], [285, 8],   [8, 417],
+    [346, 340], [340, 261], [261, 346], [413, 441], [441, 285], [285, 413],
+    [327, 460], [460, 328], [328, 327], [355, 371], [371, 329], [329, 355],
+    [392, 439], [439, 438], [438, 392], [382, 341], [341, 256], [256, 382],
+    [429, 420], [420, 360], [360, 429], [364, 394], [394, 379], [379, 364],
+    [277, 343], [343, 437], [437, 277], [443, 444], [444, 283], [283, 443],
+    [275, 440], [440, 363], [363, 275], [431, 262], [262, 369], [369, 431],
+    [297, 338], [338, 337], [337, 297], [273, 375], [375, 321], [321, 273],
+    [450, 451], [451, 349], [349, 450], [446, 342], [342, 467], [467, 446],
+    [293, 334], [334, 282], [282, 293], [458, 461], [461, 462], [462, 458],
+    [276, 353], [353, 383], [383, 276], [308, 324], [324, 325], [325, 308],
+    [276, 300], [300, 293], [293, 276], [372, 345], [345, 447], [447, 372],
+    [352, 345], [345, 340], [340, 352], [274, 1],   [1, 19],    [19, 274],
+    [456, 248], [248, 281], [281, 456], [436, 427], [427, 425], [425, 436],
+    [381, 256], [256, 252], [252, 381], [269, 391], [391, 393], [393, 269],
+    [200, 199], [199, 428], [428, 200], [266, 330], [330, 329], [329, 266],
+    [287, 273], [273, 422], [422, 287], [250, 462], [462, 328], [328, 250],
+    [258, 286], [286, 384], [384, 258], [265, 353], [353, 342], [342, 265],
+    [387, 259], [259, 257], [257, 387], [424, 431], [431, 430], [430, 424],
+    [342, 353], [353, 276], [276, 342], [273, 335], [335, 424], [424, 273],
+    [292, 325], [325, 307], [307, 292], [366, 447], [447, 345], [345, 366],
+    [271, 303], [303, 302], [302, 271], [423, 266], [266, 371], [371, 423],
+    [294, 455], [455, 460], [460, 294], [279, 278], [278, 294], [294, 279],
+    [271, 272], [272, 304], [304, 271], [432, 434], [434, 427], [427, 432],
+    [272, 407], [407, 408], [408, 272], [394, 430], [430, 431], [431, 394],
+    [395, 369], [369, 400], [400, 395], [334, 333], [333, 299], [299, 334],
+    [351, 417], [417, 168], [168, 351], [352, 280], [280, 411], [411, 352],
+    [325, 319], [319, 320], [320, 325], [295, 296], [296, 336], [336, 295],
+    [319, 403], [403, 404], [404, 319], [330, 348], [348, 349], [349, 330],
+    [293, 298], [298, 333], [333, 293], [323, 454], [454, 447], [447, 323],
+    [15, 16],   [16, 315],  [315, 15],  [358, 429], [429, 279], [279, 358],
+    [14, 15],   [15, 316],  [316, 14],  [285, 336], [336, 9],   [9, 285],
+    [329, 349], [349, 350], [350, 329], [374, 380], [380, 252], [252, 374],
+    [318, 402], [402, 403], [403, 318], [6, 197],   [197, 419], [419, 6],
+    [318, 319], [319, 325], [325, 318], [367, 364], [364, 365], [365, 367],
+    [435, 367], [367, 397], [397, 435], [344, 438], [438, 439], [439, 344],
+    [272, 271], [271, 311], [311, 272], [195, 5],   [5, 281],   [281, 195],
+    [273, 287], [287, 291], [291, 273], [396, 428], [428, 199], [199, 396],
+    [311, 271], [271, 268], [268, 311], [283, 444], [444, 445], [445, 283],
+    [373, 254], [254, 339], [339, 373], [282, 334], [334, 296], [296, 282],
+    [449, 347], [347, 346], [346, 449], [264, 447], [447, 454], [454, 264],
+    [336, 296], [296, 299], [299, 336], [338, 10],  [10, 151],  [151, 338],
+    [278, 439], [439, 455], [455, 278], [292, 407], [407, 415], [415, 292],
+    [358, 371], [371, 355], [355, 358], [340, 345], [345, 372], [372, 340],
+    [346, 347], [347, 280], [280, 346], [442, 443], [443, 282], [282, 442],
+    [19, 94],   [94, 370],  [370, 19],  [441, 442], [442, 295], [295, 441],
+    [248, 419], [419, 197], [197, 248], [263, 255], [255, 359], [359, 263],
+    [440, 275], [275, 274], [274, 440], [300, 383], [383, 368], [368, 300],
+    [351, 412], [412, 465], [465, 351], [263, 467], [467, 466], [466, 263],
+    [301, 368], [368, 389], [389, 301], [395, 378], [378, 379], [379, 395],
+    [412, 351], [351, 419], [419, 412], [436, 426], [426, 322], [322, 436],
+    [2, 164],   [164, 393], [393, 2],   [370, 462], [462, 461], [461, 370],
+    [164, 0],   [0, 267],   [267, 164], [302, 11],  [11, 12],   [12, 302],
+    [268, 12],  [12, 13],   [13, 268],  [293, 300], [300, 301], [301, 293],
+    [446, 261], [261, 340], [340, 446], [330, 266], [266, 425], [425, 330],
+    [426, 423], [423, 391], [391, 426], [429, 355], [355, 437], [437, 429],
+    [391, 327], [327, 326], [326, 391], [440, 457], [457, 438], [438, 440],
+    [341, 382], [382, 362], [362, 341], [459, 457], [457, 461], [461, 459],
+    [434, 430], [430, 394], [394, 434], [414, 463], [463, 362], [362, 414],
+    [396, 369], [369, 262], [262, 396], [354, 461], [461, 457], [457, 354],
+    [316, 403], [403, 402], [402, 316], [315, 404], [404, 403], [403, 315],
+    [314, 405], [405, 404], [404, 314], [313, 406], [406, 405], [405, 313],
+    [421, 418], [418, 406], [406, 421], [366, 401], [401, 361], [361, 366],
+    [306, 408], [408, 407], [407, 306], [291, 409], [409, 408], [408, 291],
+    [287, 410], [410, 409], [409, 287], [432, 436], [436, 410], [410, 432],
+    [434, 416], [416, 411], [411, 434], [264, 368], [368, 383], [383, 264],
+    [309, 438], [438, 457], [457, 309], [352, 376], [376, 401], [401, 352],
+    [274, 275], [275, 4],   [4, 274],   [421, 428], [428, 262], [262, 421],
+    [294, 327], [327, 358], [358, 294], [433, 416], [416, 367], [367, 433],
+    [289, 455], [455, 439], [439, 289], [462, 370], [370, 326], [326, 462],
+    [2, 326],   [326, 370], [370, 2],   [305, 460], [460, 455], [455, 305],
+    [254, 449], [449, 448], [448, 254], [255, 261], [261, 446], [446, 255],
+    [253, 450], [450, 449], [449, 253], [252, 451], [451, 450], [450, 252],
+    [256, 452], [452, 451], [451, 256], [341, 453], [453, 452], [452, 341],
+    [413, 464], [464, 463], [463, 413], [441, 413], [413, 414], [414, 441],
+    [258, 442], [442, 441], [441, 258], [257, 443], [443, 442], [442, 257],
+    [259, 444], [444, 443], [443, 259], [260, 445], [445, 444], [444, 260],
+    [467, 342], [342, 445], [445, 467], [459, 458], [458, 250], [250, 459],
+    [289, 392], [392, 290], [290, 289], [290, 328], [328, 460], [460, 290],
+    [376, 433], [433, 435], [435, 376], [250, 290], [290, 392], [392, 250],
+    [411, 416], [416, 433], [433, 411], [341, 463], [463, 464], [464, 341],
+    [453, 464], [464, 465], [465, 453], [357, 465], [465, 412], [412, 357],
+    [343, 412], [412, 399], [399, 343], [360, 363], [363, 440], [440, 360],
+    [437, 399], [399, 456], [456, 437], [420, 456], [456, 363], [363, 420],
+    [401, 435], [435, 288], [288, 401], [372, 383], [383, 353], [353, 372],
+    [339, 255], [255, 249], [249, 339], [448, 261], [261, 255], [255, 448],
+    [133, 243], [243, 190], [190, 133], [133, 155], [155, 112], [112, 133],
+    [33, 246],  [246, 247], [247, 33],  [33, 130],  [130, 25],  [25, 33],
+    [398, 384], [384, 286], [286, 398], [362, 398], [398, 414], [414, 362],
+    [362, 463], [463, 341], [341, 362], [263, 359], [359, 467], [467, 263],
+    [263, 249], [249, 255], [255, 263], [466, 467], [467, 260], [260, 466],
+    [75, 60],   [60, 166],  [166, 75],  [238, 239], [239, 79],  [79, 238],
+    [162, 127], [127, 139], [139, 162], [72, 11],   [11, 37],   [37, 72],
+    [121, 232], [232, 120], [120, 121], [73, 72],   [72, 39],   [39, 73],
+    [114, 128], [128, 47],  [47, 114],  [233, 232], [232, 128], [128, 233],
+    [103, 104], [104, 67],  [67, 103],  [152, 175], [175, 148], [148, 152],
+    [119, 118], [118, 101], [101, 119], [74, 73],   [73, 40],   [40, 74],
+    [107, 9],   [9, 108],   [108, 107], [49, 48],   [48, 131],  [131, 49],
+    [32, 194],  [194, 211], [211, 32],  [184, 74],  [74, 185],  [185, 184],
+    [191, 80],  [80, 183],  [183, 191], [185, 40],  [40, 186],  [186, 185],
+    [119, 230], [230, 118], [118, 119], [210, 202], [202, 214], [214, 210],
+    [84, 83],   [83, 17],   [17, 84],   [77, 76],   [76, 146],  [146, 77],
+    [161, 160], [160, 30],  [30, 161],  [190, 56],  [56, 173],  [173, 190],
+    [182, 106], [106, 194], [194, 182], [138, 135], [135, 192], [192, 138],
+    [129, 203], [203, 98],  [98, 129],  [54, 21],   [21, 68],   [68, 54],
+    [5, 51],    [51, 4],    [4, 5],     [145, 144], [144, 23],  [23, 145],
+    [90, 77],   [77, 91],   [91, 90],   [207, 205], [205, 187], [187, 207],
+    [83, 201],  [201, 18],  [18, 83],   [181, 91],  [91, 182],  [182, 181],
+    [180, 90],  [90, 181],  [181, 180], [16, 85],   [85, 17],   [17, 16],
+    [205, 206], [206, 36],  [36, 205],  [176, 148], [148, 140], [140, 176],
+    [165, 92],  [92, 39],   [39, 165],  [245, 193], [193, 244], [244, 245],
+    [27, 159],  [159, 28],  [28, 27],   [30, 247],  [247, 161], [161, 30],
+    [174, 236], [236, 196], [196, 174], [103, 54],  [54, 104],  [104, 103],
+    [55, 193],  [193, 8],   [8, 55],    [111, 117], [117, 31],  [31, 111],
+    [221, 189], [189, 55],  [55, 221],  [240, 98],  [98, 99],   [99, 240],
+    [142, 126], [126, 100], [100, 142], [219, 166], [166, 218], [218, 219],
+    [112, 155], [155, 26],  [26, 112],  [198, 209], [209, 131], [131, 198],
+    [169, 135], [135, 150], [150, 169], [114, 47],  [47, 217],  [217, 114],
+    [224, 223], [223, 53],  [53, 224],  [220, 45],  [45, 134],  [134, 220],
+    [32, 211],  [211, 140], [140, 32],  [109, 67],  [67, 108],  [108, 109],
+    [146, 43],  [43, 91],   [91, 146],  [231, 230], [230, 120], [120, 231],
+    [113, 226], [226, 247], [247, 113], [105, 63],  [63, 52],   [52, 105],
+    [241, 238], [238, 242], [242, 241], [124, 46],  [46, 156],  [156, 124],
+    [95, 78],   [78, 96],   [96, 95],   [70, 46],   [46, 63],   [63, 70],
+    [116, 143], [143, 227], [227, 116], [116, 123], [123, 111], [111, 116],
+    [1, 44],    [44, 19],   [19, 1],    [3, 236],   [236, 51],  [51, 3],
+    [207, 216], [216, 205], [205, 207], [26, 154],  [154, 22],  [22, 26],
+    [165, 39],  [39, 167],  [167, 165], [199, 200], [200, 208], [208, 199],
+    [101, 36],  [36, 100],  [100, 101], [43, 57],   [57, 202],  [202, 43],
+    [242, 20],  [20, 99],   [99, 242],  [56, 28],   [28, 157],  [157, 56],
+    [124, 35],  [35, 113],  [113, 124], [29, 160],  [160, 27],  [27, 29],
+    [211, 204], [204, 210], [210, 211], [124, 113], [113, 46],  [46, 124],
+    [106, 43],  [43, 204],  [204, 106], [96, 62],   [62, 77],   [77, 96],
+    [227, 137], [137, 116], [116, 227], [73, 41],   [41, 72],   [72, 73],
+    [36, 203],  [203, 142], [142, 36],  [235, 64],  [64, 240],  [240, 235],
+    [48, 49],   [49, 64],   [64, 48],   [42, 41],   [41, 74],   [74, 42],
+    [214, 212], [212, 207], [207, 214], [183, 42],  [42, 184],  [184, 183],
+    [210, 169], [169, 211], [211, 210], [140, 170], [170, 176], [176, 140],
+    [104, 105], [105, 69],  [69, 104],  [193, 122], [122, 168], [168, 193],
+    [50, 123],  [123, 187], [187, 50],  [89, 96],   [96, 90],   [90, 89],
+    [66, 65],   [65, 107],  [107, 66],  [179, 89],  [89, 180],  [180, 179],
+    [119, 101], [101, 120], [120, 119], [68, 63],   [63, 104],  [104, 68],
+    [234, 93],  [93, 227],  [227, 234], [16, 15],   [15, 85],   [85, 16],
+    [209, 129], [129, 49],  [49, 209],  [15, 14],   [14, 86],   [86, 15],
+    [107, 55],  [55, 9],    [9, 107],   [120, 100], [100, 121], [121, 120],
+    [153, 145], [145, 22],  [22, 153],  [178, 88],  [88, 179],  [179, 178],
+    [197, 6],   [6, 196],   [196, 197], [89, 88],   [88, 96],   [96, 89],
+    [135, 138], [138, 136], [136, 135], [138, 215], [215, 172], [172, 138],
+    [218, 115], [115, 219], [219, 218], [41, 42],   [42, 81],   [81, 41],
+    [5, 195],   [195, 51],  [51, 5],    [57, 43],   [43, 61],   [61, 57],
+    [208, 171], [171, 199], [199, 208], [41, 81],   [81, 38],   [38, 41],
+    [224, 53],  [53, 225],  [225, 224], [24, 144],  [144, 110], [110, 24],
+    [105, 52],  [52, 66],   [66, 105],  [118, 229], [229, 117], [117, 118],
+    [227, 34],  [34, 234],  [234, 227], [66, 107],  [107, 69],  [69, 66],
+    [10, 109],  [109, 151], [151, 10],  [219, 48],  [48, 235],  [235, 219],
+    [183, 62],  [62, 191],  [191, 183], [142, 129], [129, 126], [126, 142],
+    [116, 111], [111, 143], [143, 116], [118, 117], [117, 50],  [50, 118],
+    [223, 222], [222, 52],  [52, 223],  [94, 19],   [19, 141],  [141, 94],
+    [222, 221], [221, 65],  [65, 222],  [196, 3],   [3, 197],   [197, 196],
+    [45, 220],  [220, 44],  [44, 45],   [156, 70],  [70, 139],  [139, 156],
+    [188, 122], [122, 245], [245, 188], [139, 71],  [71, 162],  [162, 139],
+    [149, 170], [170, 150], [150, 149], [122, 188], [188, 196], [196, 122],
+    [206, 216], [216, 92],  [92, 206],  [164, 2],   [2, 167],   [167, 164],
+    [242, 141], [141, 241], [241, 242], [0, 164],   [164, 37],  [37, 0],
+    [11, 72],   [72, 12],   [12, 11],   [12, 38],   [38, 13],   [13, 12],
+    [70, 63],   [63, 71],   [71, 70],   [31, 226],  [226, 111], [111, 31],
+    [36, 101],  [101, 205], [205, 36],  [203, 206], [206, 165], [165, 203],
+    [126, 209], [209, 217], [217, 126], [98, 165],  [165, 97],  [97, 98],
+    [237, 220], [220, 218], [218, 237], [237, 239], [239, 241], [241, 237],
+    [210, 214], [214, 169], [169, 210], [140, 171], [171, 32],  [32, 140],
+    [241, 125], [125, 237], [237, 241], [179, 86],  [86, 178],  [178, 179],
+    [180, 85],  [85, 179],  [179, 180], [181, 84],  [84, 180],  [180, 181],
+    [182, 83],  [83, 181],  [181, 182], [194, 201], [201, 182], [182, 194],
+    [177, 137], [137, 132], [132, 177], [184, 76],  [76, 183],  [183, 184],
+    [185, 61],  [61, 184],  [184, 185], [186, 57],  [57, 185],  [185, 186],
+    [216, 212], [212, 186], [186, 216], [192, 214], [214, 187], [187, 192],
+    [139, 34],  [34, 156],  [156, 139], [218, 79],  [79, 237],  [237, 218],
+    [147, 123], [123, 177], [177, 147], [45, 44],   [44, 4],    [4, 45],
+    [208, 201], [201, 32],  [32, 208],  [98, 64],   [64, 129],  [129, 98],
+    [192, 213], [213, 138], [138, 192], [235, 59],  [59, 219],  [219, 235],
+    [141, 242], [242, 97],  [97, 141],  [97, 2],    [2, 141],   [141, 97],
+    [240, 75],  [75, 235],  [235, 240], [229, 24],  [24, 228],  [228, 229],
+    [31, 25],   [25, 226],  [226, 31],  [230, 23],  [23, 229],  [229, 230],
+    [231, 22],  [22, 230],  [230, 231], [232, 26],  [26, 231],  [231, 232],
+    [233, 112], [112, 232], [232, 233], [244, 189], [189, 243], [243, 244],
+    [189, 221], [221, 190], [190, 189], [222, 28],  [28, 221],  [221, 222],
+    [223, 27],  [27, 222],  [222, 223], [224, 29],  [29, 223],  [223, 224],
+    [225, 30],  [30, 224],  [224, 225], [113, 247], [247, 225], [225, 113],
+    [99, 60],   [60, 240],  [240, 99],  [213, 147], [147, 215], [215, 213],
+    [60, 20],   [20, 166],  [166, 60],  [192, 187], [187, 213], [213, 192],
+    [243, 112], [112, 244], [244, 243], [244, 233], [233, 245], [245, 244],
+    [245, 128], [128, 188], [188, 245], [188, 114], [114, 174], [174, 188],
+    [134, 131], [131, 220], [220, 134], [174, 217], [217, 236], [236, 174],
+    [236, 198], [198, 134], [134, 236], [215, 177], [177, 58],  [58, 215],
+    [156, 143], [143, 124], [124, 156], [25, 110],  [110, 7],   [7, 25],
+    [31, 228],  [228, 25],  [25, 31],   [264, 356], [356, 368], [368, 264],
+    [0, 11],    [11, 267],  [267, 0],   [451, 452], [452, 349], [349, 451],
+    [267, 302], [302, 269], [269, 267], [350, 357], [357, 277], [277, 350],
+    [350, 452], [452, 357], [357, 350], [299, 333], [333, 297], [297, 299],
+    [396, 175], [175, 377], [377, 396], [280, 347], [347, 330], [330, 280],
+    [269, 303], [303, 270], [270, 269], [151, 9],   [9, 337],   [337, 151],
+    [344, 278], [278, 360], [360, 344], [424, 418], [418, 431], [431, 424],
+    [270, 304], [304, 409], [409, 270], [272, 310], [310, 407], [407, 272],
+    [322, 270], [270, 410], [410, 322], [449, 450], [450, 347], [347, 449],
+    [432, 422], [422, 434], [434, 432], [18, 313],  [313, 17],  [17, 18],
+    [291, 306], [306, 375], [375, 291], [259, 387], [387, 260], [260, 259],
+    [424, 335], [335, 418], [418, 424], [434, 364], [364, 416], [416, 434],
+    [391, 423], [423, 327], [327, 391], [301, 251], [251, 298], [298, 301],
+    [275, 281], [281, 4],   [4, 275],   [254, 373], [373, 253], [253, 254],
+    [375, 307], [307, 321], [321, 375], [280, 425], [425, 411], [411, 280],
+    [200, 421], [421, 18],  [18, 200],  [335, 321], [321, 406], [406, 335],
+    [321, 320], [320, 405], [405, 321], [314, 315], [315, 17],  [17, 314],
+    [423, 426], [426, 266], [266, 423], [396, 377], [377, 369], [369, 396],
+    [270, 322], [322, 269], [269, 270], [413, 417], [417, 464], [464, 413],
+    [385, 386], [386, 258], [258, 385], [248, 456], [456, 419], [419, 248],
+    [298, 284], [284, 333], [333, 298], [168, 417], [417, 8],   [8, 168],
+    [448, 346], [346, 261], [261, 448], [417, 413], [413, 285], [285, 417],
+    [326, 327], [327, 328], [328, 326], [277, 355], [355, 329], [329, 277],
+    [309, 392], [392, 438], [438, 309], [381, 382], [382, 256], [256, 381],
+    [279, 429], [429, 360], [360, 279], [365, 364], [364, 379], [379, 365],
+    [355, 277], [277, 437], [437, 355], [282, 443], [443, 283], [283, 282],
+    [281, 275], [275, 363], [363, 281], [395, 431], [431, 369], [369, 395],
+    [299, 297], [297, 337], [337, 299], [335, 273], [273, 321], [321, 335],
+    [348, 450], [450, 349], [349, 348], [359, 446], [446, 467], [467, 359],
+    [283, 293], [293, 282], [282, 283], [250, 458], [458, 462], [462, 250],
+    [300, 276], [276, 383], [383, 300], [292, 308], [308, 325], [325, 292],
+    [283, 276], [276, 293], [293, 283], [264, 372], [372, 447], [447, 264],
+    [346, 352], [352, 340], [340, 346], [354, 274], [274, 19],  [19, 354],
+    [363, 456], [456, 281], [281, 363], [426, 436], [436, 425], [425, 426],
+    [380, 381], [381, 252], [252, 380], [267, 269], [269, 393], [393, 267],
+    [421, 200], [200, 428], [428, 421], [371, 266], [266, 329], [329, 371],
+    [432, 287], [287, 422], [422, 432], [290, 250], [250, 328], [328, 290],
+    [385, 258], [258, 384], [384, 385], [446, 265], [265, 342], [342, 446],
+    [386, 387], [387, 257], [257, 386], [422, 424], [424, 430], [430, 422],
+    [445, 342], [342, 276], [276, 445], [422, 273], [273, 424], [424, 422],
+    [306, 292], [292, 307], [307, 306], [352, 366], [366, 345], [345, 352],
+    [268, 271], [271, 302], [302, 268], [358, 423], [423, 371], [371, 358],
+    [327, 294], [294, 460], [460, 327], [331, 279], [279, 294], [294, 331],
+    [303, 271], [271, 304], [304, 303], [436, 432], [432, 427], [427, 436],
+    [304, 272], [272, 408], [408, 304], [395, 394], [394, 431], [431, 395],
+    [378, 395], [395, 400], [400, 378], [296, 334], [334, 299], [299, 296],
+    [6, 351],   [351, 168], [168, 6],   [376, 352], [352, 411], [411, 376],
+    [307, 325], [325, 320], [320, 307], [285, 295], [295, 336], [336, 285],
+    [320, 319], [319, 404], [404, 320], [329, 330], [330, 349], [349, 329],
+    [334, 293], [293, 333], [333, 334], [366, 323], [323, 447], [447, 366],
+    [316, 15],  [15, 315],  [315, 316], [331, 358], [358, 279], [279, 331],
+    [317, 14],  [14, 316],  [316, 317], [8, 285],   [285, 9],   [9, 8],
+    [277, 329], [329, 350], [350, 277], [253, 374], [374, 252], [252, 253],
+    [319, 318], [318, 403], [403, 319], [351, 6],   [6, 419],   [419, 351],
+    [324, 318], [318, 325], [325, 324], [397, 367], [367, 365], [365, 397],
+    [288, 435], [435, 397], [397, 288], [278, 344], [344, 439], [439, 278],
+    [310, 272], [272, 311], [311, 310], [248, 195], [195, 281], [281, 248],
+    [375, 273], [273, 291], [291, 375], [175, 396], [396, 199], [199, 175],
+    [312, 311], [311, 268], [268, 312], [276, 283], [283, 445], [445, 276],
+    [390, 373], [373, 339], [339, 390], [295, 282], [282, 296], [296, 295],
+    [448, 449], [449, 346], [346, 448], [356, 264], [264, 454], [454, 356],
+    [337, 336], [336, 299], [299, 337], [337, 338], [338, 151], [151, 337],
+    [294, 278], [278, 455], [455, 294], [308, 292], [292, 415], [415, 308],
+    [429, 358], [358, 355], [355, 429], [265, 340], [340, 372], [372, 265],
+    [352, 346], [346, 280], [280, 352], [295, 442], [442, 282], [282, 295],
+    [354, 19],  [19, 370],  [370, 354], [285, 441], [441, 295], [295, 285],
+    [195, 248], [248, 197], [197, 195], [457, 440], [440, 274], [274, 457],
+    [301, 300], [300, 368], [368, 301], [417, 351], [351, 465], [465, 417],
+    [251, 301], [301, 389], [389, 251], [394, 395], [395, 379], [379, 394],
+    [399, 412], [412, 419], [419, 399], [410, 436], [436, 322], [322, 410],
+    [326, 2],   [2, 393],   [393, 326], [354, 370], [370, 461], [461, 354],
+    [393, 164], [164, 267], [267, 393], [268, 302], [302, 12],  [12, 268],
+    [312, 268], [268, 13],  [13, 312],  [298, 293], [293, 301], [301, 298],
+    [265, 446], [446, 340], [340, 265], [280, 330], [330, 425], [425, 280],
+    [322, 426], [426, 391], [391, 322], [420, 429], [429, 437], [437, 420],
+    [393, 391], [391, 326], [326, 393], [344, 440], [440, 438], [438, 344],
+    [458, 459], [459, 461], [461, 458], [364, 434], [434, 394], [394, 364],
+    [428, 396], [396, 262], [262, 428], [274, 354], [354, 457], [457, 274],
+    [317, 316], [316, 402], [402, 317], [316, 315], [315, 403], [403, 316],
+    [315, 314], [314, 404], [404, 315], [314, 313], [313, 405], [405, 314],
+    [313, 421], [421, 406], [406, 313], [323, 366], [366, 361], [361, 323],
+    [292, 306], [306, 407], [407, 292], [306, 291], [291, 408], [408, 306],
+    [291, 287], [287, 409], [409, 291], [287, 432], [432, 410], [410, 287],
+    [427, 434], [434, 411], [411, 427], [372, 264], [264, 383], [383, 372],
+    [459, 309], [309, 457], [457, 459], [366, 352], [352, 401], [401, 366],
+    [1, 274],   [274, 4],   [4, 1],     [418, 421], [421, 262], [262, 418],
+    [331, 294], [294, 358], [358, 331], [435, 433], [433, 367], [367, 435],
+    [392, 289], [289, 439], [439, 392], [328, 462], [462, 326], [326, 328],
+    [94, 2],    [2, 370],   [370, 94],  [289, 305], [305, 455], [455, 289],
+    [339, 254], [254, 448], [448, 339], [359, 255], [255, 446], [446, 359],
+    [254, 253], [253, 449], [449, 254], [253, 252], [252, 450], [450, 253],
+    [252, 256], [256, 451], [451, 252], [256, 341], [341, 452], [452, 256],
+    [414, 413], [413, 463], [463, 414], [286, 441], [441, 414], [414, 286],
+    [286, 258], [258, 441], [441, 286], [258, 257], [257, 442], [442, 258],
+    [257, 259], [259, 443], [443, 257], [259, 260], [260, 444], [444, 259],
+    [260, 467], [467, 445], [445, 260], [309, 459], [459, 250], [250, 309],
+    [305, 289], [289, 290], [290, 305], [305, 290], [290, 460], [460, 305],
+    [401, 376], [376, 435], [435, 401], [309, 250], [250, 392], [392, 309],
+    [376, 411], [411, 433], [433, 376], [453, 341], [341, 464], [464, 453],
+    [357, 453], [453, 465], [465, 357], [343, 357], [357, 412], [412, 343],
+    [437, 343], [343, 399], [399, 437], [344, 360], [360, 440], [440, 344],
+    [420, 437], [437, 456], [456, 420], [360, 420], [420, 363], [363, 360],
+    [361, 401], [401, 288], [288, 361], [265, 372], [372, 353], [353, 265],
+    [390, 339], [339, 249], [249, 390], [339, 448], [448, 255], [255, 339]]))

client/src/hooks/useFaceLandmarkDetection.tsx

@@ -0,0 +1,632 @@
+import React, { useState, useEffect, useRef, useCallback } from 'react';
+import * as vision from '@mediapipe/tasks-vision';
+
+import { facePoke } from '@/lib/facePoke';
+import { useMainStore } from './useMainStore';
+import useThrottledCallback from 'beautiful-react-hooks/useThrottledCallback';
+
+import { landmarkGroups, FACEMESH_LIPS, FACEMESH_LEFT_EYE, FACEMESH_LEFT_EYEBROW, FACEMESH_RIGHT_EYE, FACEMESH_RIGHT_EYEBROW, FACEMESH_FACE_OVAL } from './landmarks';
+
+// New types for improved type safety
+export type LandmarkGroup = 'lips' | 'leftEye' | 'leftEyebrow' | 'rightEye' | 'rightEyebrow' | 'faceOval' | 'background';
+export type LandmarkCenter = { x: number; y: number; z: number };
+export type ClosestLandmark = { group: LandmarkGroup; distance: number; vector: { x: number; y: number; z: number } };
+
+export type MediaPipeResources = {
+  faceLandmarker: vision.FaceLandmarker | null;
+  drawingUtils: vision.DrawingUtils | null;
+};
+
+export function useFaceLandmarkDetection() {
+  const error = useMainStore(s => s.error);
+  const setError = useMainStore(s => s.setError);
+  const imageFile = useMainStore(s => s.imageFile);
+  const setImageFile = useMainStore(s => s.setImageFile);
+  const originalImage = useMainStore(s => s.originalImage);
+  const originalImageHash = useMainStore(s => s.originalImageHash);
+  const setOriginalImageHash = useMainStore(s => s.setOriginalImageHash);
+  const previewImage = useMainStore(s => s.previewImage);
+  const setPreviewImage = useMainStore(s => s.setPreviewImage);
+  const resetImage = useMainStore(s => s.resetImage);
+
+  ;(window as any).debugJuju = useMainStore;
+  ////////////////////////////////////////////////////////////////////////
+  // ok so apparently I cannot vary the latency, or else there is a bug
+  // const averageLatency = useMainStore(s => s.averageLatency);
+  const averageLatency = 220
+  ////////////////////////////////////////////////////////////////////////
+
+  // State for face detection
+  const [faceLandmarks, setFaceLandmarks] = useState<vision.NormalizedLandmark[][]>([]);
+  const [isMediaPipeReady, setIsMediaPipeReady] = useState(false);
+  const [isDrawingUtilsReady, setIsDrawingUtilsReady] = useState(false);
+  const [blendShapes, setBlendShapes] = useState<vision.Classifications[]>([]);
+
+  // State for mouse interaction
+  const [dragStart, setDragStart] = useState<{ x: number; y: number } | null>(null);
+  const [dragEnd, setDragEnd] = useState<{ x: number; y: number } | null>(null);
+
+  const [isDragging, setIsDragging] = useState(false);
+  const [isWaitingForResponse, setIsWaitingForResponse] = useState(false);
+  const dragStartRef = useRef<{ x: number; y: number } | null>(null);
+  const currentMousePosRef = useRef<{ x: number; y: number } | null>(null);
+  const lastModifiedImageHashRef = useRef<string | null>(null);
+
+  const [currentLandmark, setCurrentLandmark] = useState<ClosestLandmark | null>(null);
+  const [previousLandmark, setPreviousLandmark] = useState<ClosestLandmark | null>(null);
+  const [currentOpacity, setCurrentOpacity] = useState(0);
+  const [previousOpacity, setPreviousOpacity] = useState(0);
+
+  const [isHovering, setIsHovering] = useState(false);
+
+  // Refs
+  const canvasRef = useRef<HTMLCanvasElement>(null);
+  const mediaPipeRef = useRef<MediaPipeResources>({
+    faceLandmarker: null,
+    drawingUtils: null,
+  });
+
+  const setActiveLandmark = useCallback((newLandmark: ClosestLandmark | undefined) => {
+    //if (newLandmark && (!currentLandmark || newLandmark.group !== currentLandmark.group)) {
+      setPreviousLandmark(currentLandmark || null);
+      setCurrentLandmark(newLandmark || null);
+      setCurrentOpacity(0);
+      setPreviousOpacity(1);
+    //}
+  }, [currentLandmark, setPreviousLandmark, setCurrentLandmark, setCurrentOpacity, setPreviousOpacity]);
+
+  // Initialize MediaPipe
+  useEffect(() => {
+    console.log('Initializing MediaPipe...');
+    let isMounted = true;
+
+    const initializeMediaPipe = async () => {
+      const { FaceLandmarker, FilesetResolver, DrawingUtils } = vision;
+
+      try {
+        console.log('Initializing FilesetResolver...');
+        const filesetResolver = await FilesetResolver.forVisionTasks(
+          "https://cdn.jsdelivr.net/npm/@mediapipe/tasks-vision@0.10.3/wasm"
+        );
+
+        console.log('Creating FaceLandmarker...');
+        const faceLandmarker = await FaceLandmarker.createFromOptions(filesetResolver, {
+          baseOptions: {
+            modelAssetPath: `https://storage.googleapis.com/mediapipe-models/face_landmarker/face_landmarker/float16/1/face_landmarker.task`,
+            delegate: "GPU"
+          },
+          outputFaceBlendshapes: true,
+          runningMode: "IMAGE",
+          numFaces: 1
+        });
+
+        if (isMounted) {
+          console.log('FaceLandmarker created successfully.');
+          mediaPipeRef.current.faceLandmarker = faceLandmarker;
+          setIsMediaPipeReady(true);
+        } else {
+          faceLandmarker.close();
+        }
+      } catch (error) {
+        console.error('Error during MediaPipe initialization:', error);
+        setError('Failed to initialize face detection. Please try refreshing the page.');
+      }
+    };
+
+    initializeMediaPipe();
+
+
+    return () => {
+      isMounted = false;
+      if (mediaPipeRef.current.faceLandmarker) {
+        mediaPipeRef.current.faceLandmarker.close();
+      }
+    };
+  }, []);
+
+  // New state for storing landmark centers
+  const [landmarkCenters, setLandmarkCenters] = useState<Record<LandmarkGroup, LandmarkCenter>>({} as Record<LandmarkGroup, LandmarkCenter>);
+
+  // Function to compute the center of each landmark group
+  const computeLandmarkCenters = useCallback((landmarks: vision.NormalizedLandmark[]) => {
+    const centers: Record<LandmarkGroup, LandmarkCenter> = {} as Record<LandmarkGroup, LandmarkCenter>;
+
+    const computeGroupCenter = (group: Readonly<Set<number[]>>): LandmarkCenter => {
+      let sumX = 0, sumY = 0, sumZ = 0, count = 0;
+      group.forEach(([index]) => {
+        if (landmarks[index]) {
+          sumX += landmarks[index].x;
+          sumY += landmarks[index].y;
+          sumZ += landmarks[index].z || 0;
+          count++;
+        }
+      });
+      return { x: sumX / count, y: sumY / count, z: sumZ / count };
+    };
+
+    centers.lips = computeGroupCenter(FACEMESH_LIPS);
+    centers.leftEye = computeGroupCenter(FACEMESH_LEFT_EYE);
+    centers.leftEyebrow = computeGroupCenter(FACEMESH_LEFT_EYEBROW);
+    centers.rightEye = computeGroupCenter(FACEMESH_RIGHT_EYE);
+    centers.rightEyebrow = computeGroupCenter(FACEMESH_RIGHT_EYEBROW);
+    centers.faceOval = computeGroupCenter(FACEMESH_FACE_OVAL);
+    centers.background = { x: 0.5, y: 0.5, z: 0 };
+
+    setLandmarkCenters(centers);
+    // console.log('Landmark centers computed:', centers);
+  }, []);
+
+   // Function to find the closest landmark to the mouse position
+   const findClosestLandmark = useCallback((mouseX: number, mouseY: number, isGroup?: LandmarkGroup): ClosestLandmark => {
+    const defaultLandmark: ClosestLandmark = {
+      group: 'background',
+      distance: 0,
+      vector: {
+        x: mouseX,
+        y: mouseY,
+        z: 0
+      }
+    }
+
+    if (Object.keys(landmarkCenters).length === 0) {
+      console.warn('Landmark centers not computed yet');
+      return defaultLandmark;
+    }
+
+    let closestGroup: LandmarkGroup | null = null;
+    let minDistance = Infinity;
+    let closestVector = { x: 0, y: 0, z: 0 };
+    let faceOvalDistance = Infinity;
+    let faceOvalVector = { x: 0, y: 0, z: 0 };
+
+    Object.entries(landmarkCenters).forEach(([group, center]) => {
+      const dx = mouseX - center.x;
+      const dy = mouseY - center.y;
+      const distance = Math.sqrt(dx * dx + dy * dy);
+
+      if (group === 'faceOval') {
+        faceOvalDistance = distance;
+        faceOvalVector = { x: dx, y: dy, z: 0 };
+      }
+
+      // filter to keep the group if it is belonging to  `ofGroup`
+      if (isGroup) {
+        if (group !== isGroup) {
+          return
+        }
+      }
+
+      if (distance < minDistance) {
+        minDistance = distance;
+        closestGroup = group as LandmarkGroup;
+        closestVector = { x: dx, y: dy, z: 0 }; // Z is 0 as mouse interaction is 2D
+      }
+    });
+
+    // Fallback to faceOval if no group found or distance is too large
+    if (minDistance > 0.05) {
+      // console.log('Distance is too high, so we use the faceOval group');
+      closestGroup = 'background';
+      minDistance = faceOvalDistance;
+      closestVector = faceOvalVector;
+    }
+
+    if (closestGroup) {
+      // console.log(`Closest landmark: ${closestGroup}, distance: ${minDistance.toFixed(4)}`);
+      return { group: closestGroup, distance: minDistance, vector: closestVector };
+    } else {
+      // console.log('No group found, returning fallback');
+      return defaultLandmark
+    }
+  }, [landmarkCenters]);
+
+  // Detect face landmarks
+  const detectFaceLandmarks = useCallback(async (imageDataUrl: string) => {
+    // console.log('Attempting to detect face landmarks...');
+    if (!isMediaPipeReady) {
+      console.log('MediaPipe not ready. Skipping detection.');
+      return;
+    }
+
+    const faceLandmarker = mediaPipeRef.current.faceLandmarker;
+
+    if (!faceLandmarker) {
+      console.error('FaceLandmarker is not initialized.');
+      return;
+    }
+
+    const drawingUtils = mediaPipeRef.current.drawingUtils;
+
+    const image = new Image();
+    image.src = imageDataUrl;
+    await new Promise((resolve) => { image.onload = resolve; });
+
+    const faceLandmarkerResult = faceLandmarker.detect(image);
+    // console.log("Face landmarks detected:", faceLandmarkerResult);
+
+    setFaceLandmarks(faceLandmarkerResult.faceLandmarks);
+    setBlendShapes(faceLandmarkerResult.faceBlendshapes || []);
+
+    if (faceLandmarkerResult.faceLandmarks && faceLandmarkerResult.faceLandmarks[0]) {
+      computeLandmarkCenters(faceLandmarkerResult.faceLandmarks[0]);
+    }
+
+    if (canvasRef.current && drawingUtils) {
+      drawLandmarks(faceLandmarkerResult.faceLandmarks[0], canvasRef.current, drawingUtils);
+    }
+  }, [isMediaPipeReady, isDrawingUtilsReady, computeLandmarkCenters]);
+
+  const drawLandmarks = useCallback((
+    landmarks: vision.NormalizedLandmark[],
+    canvas: HTMLCanvasElement,
+    drawingUtils: vision.DrawingUtils
+  ) => {
+    const ctx = canvas.getContext('2d');
+    if (!ctx) return;
+
+    ctx.clearRect(0, 0, canvas.width, canvas.height);
+
+    if (canvasRef.current && previewImage) {
+      const img = new Image();
+      img.onload = () => {
+        canvas.width = img.width;
+        canvas.height = img.height;
+
+        const drawLandmarkGroup = (landmark: ClosestLandmark | null, opacity: number) => {
+          if (!landmark) return;
+          const connections = landmarkGroups[landmark.group];
+          if (connections) {
+            ctx.globalAlpha = opacity;
+            drawingUtils.drawConnectors(
+              landmarks,
+              connections,
+              { color: 'orange', lineWidth: 4 }
+            );
+          }
+        };
+
+        drawLandmarkGroup(previousLandmark, previousOpacity);
+        drawLandmarkGroup(currentLandmark, currentOpacity);
+
+        ctx.globalAlpha = 1;
+      };
+      img.src = previewImage;
+    }
+  }, [previewImage, currentLandmark, previousLandmark, currentOpacity, previousOpacity]);
+
+  useEffect(() => {
+    if (isMediaPipeReady && isDrawingUtilsReady && faceLandmarks.length > 0 && canvasRef.current && mediaPipeRef.current.drawingUtils) {
+      drawLandmarks(faceLandmarks[0], canvasRef.current, mediaPipeRef.current.drawingUtils);
+    }
+  }, [isMediaPipeReady, isDrawingUtilsReady, faceLandmarks, currentLandmark, previousLandmark, currentOpacity, previousOpacity, drawLandmarks]);
+  useEffect(() => {
+    let animationFrame: number;
+    const animate = () => {
+      setCurrentOpacity((prev) => Math.min(prev + 0.2, 1));
+      setPreviousOpacity((prev) => Math.max(prev - 0.2, 0));
+
+      if (currentOpacity < 1 || previousOpacity > 0) {
+        animationFrame = requestAnimationFrame(animate);
+      }
+    };
+    animationFrame = requestAnimationFrame(animate);
+    return () => cancelAnimationFrame(animationFrame);
+  }, [currentLandmark]);
+
+  // Canvas ref callback
+  const canvasRefCallback = useCallback((node: HTMLCanvasElement | null) => {
+    if (node !== null) {
+      const ctx = node.getContext('2d');
+      if (ctx) {
+        // Get device pixel ratio
+        const pixelRatio = window.devicePixelRatio || 1;
+
+        // Scale canvas based on the pixel ratio
+        node.width = node.clientWidth * pixelRatio;
+        node.height = node.clientHeight * pixelRatio;
+        ctx.scale(pixelRatio, pixelRatio);
+
+        mediaPipeRef.current.drawingUtils = new vision.DrawingUtils(ctx);
+        setIsDrawingUtilsReady(true);
+      } else {
+        console.error('Failed to get 2D context from canvas.');
+      }
+      canvasRef.current = node;
+    }
+  }, []);
+
+
+  useEffect(() => {
+    if (!isMediaPipeReady) {
+      console.log('MediaPipe not ready. Skipping landmark detection.');
+      return
+    }
+    if (!previewImage) {
+      console.log('Preview image not ready. Skipping landmark detection.');
+      return
+    }
+    if (!isDrawingUtilsReady) {
+      console.log('DrawingUtils not ready. Skipping landmark detection.');
+      return
+    }
+    detectFaceLandmarks(previewImage);
+  }, [isMediaPipeReady, isDrawingUtilsReady, previewImage])
+
+
+
+  const modifyImage = useCallback(({ landmark, vector }: {
+      landmark: ClosestLandmark
+      vector: { x: number; y: number; z: number }
+    }) => {
+
+    const {
+      originalImage,
+      originalImageHash,
+      params: previousParams,
+      setParams,
+      setError
+    } = useMainStore.getState()
+
+
+    if (!originalImage) {
+      console.error('Image file or facePoke not available');
+      return;
+    }
+
+    const params = {
+      ...previousParams
+    }
+
+    const minX = -0.50;
+    const maxX = 0.50;
+    const minY = -0.50;
+    const maxY = 0.50;
+
+    // Function to map a value from one range to another
+    const mapRange = (value: number, inMin: number, inMax: number, outMin: number, outMax: number): number => {
+      return Math.min(outMax, Math.max(outMin, ((value - inMin) * (outMax - outMin)) / (inMax - inMin) + outMin));
+    };
+
+    console.log("modifyImage:", {
+      originalImage,
+      originalImageHash,
+      landmark,
+      vector,
+      minX,
+      maxX,
+      minY,
+      maxY,
+    })
+
+    // Map landmarks to ImageModificationParams
+    switch (landmark.group) {
+      case 'leftEye':
+      case 'rightEye':
+         // eyebrow (min: -20, max: 5, default: 0)
+        const eyesMin = 210
+        const eyesMax = 5
+        params.eyes = mapRange(vector.x, minX, maxX, eyesMin, eyesMax);
+
+        break;
+      case 'leftEyebrow':
+      case 'rightEyebrow':
+        // moving the mouse vertically for the eyebrow
+        // should make them up/down
+        // eyebrow (min: -10, max: 15, default: 0)
+        const eyebrowMin = -10
+        const eyebrowMax = 15
+        params.eyebrow = mapRange(vector.y, minY, maxY, eyebrowMin, eyebrowMax);
+
+        break;
+      case 'lips':
+        // aaa (min: -30, max: 120, default: 0)
+        //const aaaMin = -30
+        //const aaaMax = 120
+        //params.aaa = mapRange(vector.x, minY, maxY, aaaMin, aaaMax);
+
+        // eee (min: -20, max: 15, default: 0)
+        const eeeMin = -20
+        const eeeMax = 15
+        params.eee = mapRange(vector.y, minY, maxY, eeeMin, eeeMax);
+
+
+        // woo (min: -20, max: 15, default: 0)
+        const wooMin = -20
+        const wooMax = 15
+        params.woo = mapRange(vector.x, minX, maxX, wooMin, wooMax);
+
+        break;
+      case 'faceOval':
+          // displacing the face horizontally by moving the mouse on the X axis
+          // should perform a yaw rotation
+          // rotate_roll (min: -20, max: 20, default: 0)
+          const rollMin = -40
+          const rollMax = 40
+
+          // note: we invert the axis here
+          params.rotate_roll = mapRange(vector.x, minX, maxX, rollMin, rollMax);
+          break;
+
+      case 'background':
+        // displacing the face horizontally by moving the mouse on the X axis
+        // should perform a yaw rotation
+        // rotate_yaw (min: -20, max: 20, default: 0)
+        const yawMin = -40
+        const yawMax = 40
+
+        // note: we invert the axis here
+        params.rotate_yaw = mapRange(-vector.x, minX, maxX, yawMin, yawMax);
+
+        // displacing the face vertically by moving the mouse on the Y axis
+        // should perform a pitch rotation
+        // rotate_pitch (min: -20, max: 20, default: 0)
+        const pitchMin = -40
+        const pitchMax = 40
+        params.rotate_pitch = mapRange(vector.y, minY, maxY, pitchMin, pitchMax);
+        break;
+      default:
+        return
+    }
+
+    for (const [key, value] of Object.entries(params)) {
+      if (isNaN(value as any) || !isFinite(value as any)) {
+        console.log(`${key} is NaN, aborting`)
+        return
+      }
+    }
+    console.log(`PITCH=${params.rotate_pitch || 0}, YAW=${params.rotate_yaw || 0}, ROLL=${params.rotate_roll || 0}`);
+
+    setParams(params)
+    try {
+      // For the first request or when the image file changes, send the full image
+      if (!lastModifiedImageHashRef.current || lastModifiedImageHashRef.current !== originalImageHash) {
+        lastModifiedImageHashRef.current = originalImageHash;
+        facePoke.modifyImage(originalImage, null, params);
+      } else {
+        // For subsequent requests, send only the hash
+        facePoke.modifyImage(null, lastModifiedImageHashRef.current, params);
+      }
+    } catch (error) {
+      // console.error('Error modifying image:', error);
+      setError('Failed to modify image');
+    }
+  }, []);
+
+  // this is throttled by our average latency
+  const modifyImageWithRateLimit = useThrottledCallback((params: {
+    landmark: ClosestLandmark
+    vector: { x: number; y: number; z: number }
+  }) => {
+    modifyImage(params);
+  }, [modifyImage], averageLatency);
+
+  const handleMouseEnter = useCallback(() => {
+    setIsHovering(true);
+  }, []);
+
+  const handleMouseLeave = useCallback(() => {
+    setIsHovering(false);
+  }, []);
+
+  // Update mouse event handlers
+  const handleMouseDown = useCallback((event: React.MouseEvent<HTMLCanvasElement>) => {
+    if (!canvasRef.current) return;
+
+    const rect = canvasRef.current.getBoundingClientRect();
+    const x = (event.clientX - rect.left) / rect.width;
+    const y = (event.clientY - rect.top) / rect.height;
+
+    const landmark = findClosestLandmark(x, y);
+    console.log(`Mouse down on ${landmark.group}`);
+    setActiveLandmark(landmark);
+    setDragStart({ x, y });
+    dragStartRef.current = { x, y };
+  }, [findClosestLandmark, setActiveLandmark, setDragStart]);
+
+  const handleMouseMove = useCallback((event: React.MouseEvent<HTMLCanvasElement>) => {
+    if (!canvasRef.current) return;
+
+    const rect = canvasRef.current.getBoundingClientRect();
+    const x = (event.clientX - rect.left) / rect.width;
+    const y = (event.clientY - rect.top) / rect.height;
+
+    // only send an API request to modify the image if we are actively dragging
+    if (dragStart && dragStartRef.current) {
+
+      const landmark = findClosestLandmark(x, y, currentLandmark?.group);
+
+      console.log(`Dragging mouse (was over ${currentLandmark?.group || 'nothing'}, now over ${landmark.group})`);
+
+      // Compute the vector from the landmark center to the current mouse position
+      modifyImageWithRateLimit({
+        landmark: currentLandmark || landmark, // this will still use the initially selected landmark
+        vector: {
+          x: x - landmarkCenters[landmark.group].x,
+          y: y - landmarkCenters[landmark.group].y,
+          z: 0 // Z is 0 as mouse interaction is 2D
+        }
+      });
+      setIsDragging(true);
+    } else {
+      const landmark = findClosestLandmark(x, y);
+
+      //console.log(`Moving mouse over ${landmark.group}`);
+      // console.log(`Simple mouse move over ${landmark.group}`);
+
+      // we need to be careful here, we don't want to change the active
+      // landmark dynamically if we are busy dragging
+
+      if (!currentLandmark || (currentLandmark?.group !== landmark?.group)) {
+        // console.log("setting activeLandmark to ", landmark);
+        setActiveLandmark(landmark);
+      }
+      setIsHovering(true); // Ensure hovering state is maintained during movement
+    }
+  }, [currentLandmark, dragStart, setIsHovering, setActiveLandmark, setIsDragging, modifyImageWithRateLimit, landmarkCenters]);
+
+  const handleMouseUp = useCallback((event: React.MouseEvent<HTMLCanvasElement>) => {
+    if (!canvasRef.current) return;
+
+    const rect = canvasRef.current.getBoundingClientRect();
+    const x = (event.clientX - rect.left) / rect.width;
+    const y = (event.clientY - rect.top) / rect.height;
+
+    // only send an API request to modify the image if we are actively dragging
+    if (dragStart && dragStartRef.current) {
+
+      const landmark = findClosestLandmark(x, y, currentLandmark?.group);
+
+      console.log(`Mouse up (was over ${currentLandmark?.group || 'nothing'}, now over ${landmark.group})`);
+
+      // Compute the vector from the landmark center to the current mouse position
+      modifyImageWithRateLimit({
+        landmark: currentLandmark || landmark, // this will still use the initially selected landmark
+        vector: {
+          x: x - landmarkCenters[landmark.group].x,
+          y: y - landmarkCenters[landmark.group].y,
+          z: 0 // Z is 0 as mouse interaction is 2D
+        }
+      });
+    }
+
+    setIsDragging(false);
+    dragStartRef.current = null;
+    setActiveLandmark(undefined);
+  }, [currentLandmark, isDragging, modifyImageWithRateLimit, findClosestLandmark, setActiveLandmark, landmarkCenters, modifyImageWithRateLimit, setIsDragging]);
+
+  useEffect(() => {
+    facePoke.setOnModifiedImage((image: string, image_hash: string) => {
+      if (image) {
+        setPreviewImage(image);
+      }
+      setOriginalImageHash(image_hash);
+      lastModifiedImageHashRef.current = image_hash;
+    });
+  }, [setPreviewImage, setOriginalImageHash]);
+
+  return {
+    canvasRef,
+    canvasRefCallback,
+    mediaPipeRef,
+    faceLandmarks,
+    isMediaPipeReady,
+    isDrawingUtilsReady,
+    blendShapes,
+
+    //dragStart,
+    //setDragStart,
+    //dragEnd,
+    //setDragEnd,
+    setFaceLandmarks,
+    setBlendShapes,
+
+    handleMouseDown,
+    handleMouseUp,
+    handleMouseMove,
+    handleMouseEnter,
+    handleMouseLeave,
+
+    currentLandmark,
+    currentOpacity,
+  }
+}

client/src/hooks/useFacePokeAPI.ts

@@ -0,0 +1,44 @@
+import { useEffect, useState } from "react";
+
+import { facePoke } from "../lib/facePoke";
+import { useMainStore } from "./useMainStore";
+
+export function useFacePokeAPI() {
+
+   // State for FacePoke
+  const [status, setStatus] = useState('');
+  const [isDebugMode, setIsDebugMode] = useState(false);
+  const [interruptMessage, setInterruptMessage] = useState<string | null>(null);
+
+  const [isLoading, setIsLoading] = useState(false);
+
+  // Initialize FacePoke
+  useEffect(() => {
+    const urlParams = new URLSearchParams(window.location.search);
+    setIsDebugMode(urlParams.get('debug') === 'true');
+  }, []);
+
+  // Handle WebSocket interruptions
+  useEffect(() => {
+    const handleInterruption = (event: CustomEvent) => {
+      setInterruptMessage(event.detail.message);
+    };
+
+    window.addEventListener('websocketInterruption' as any, handleInterruption);
+
+    return () => {
+      window.removeEventListener('websocketInterruption' as any, handleInterruption);
+    };
+  }, []);
+
+  return {
+    facePoke,
+    status,
+    setStatus,
+    isDebugMode,
+    setIsDebugMode,
+    interruptMessage,
+    isLoading,
+    setIsLoading,
+  }
+}

client/src/hooks/useMainStore.ts

@@ -0,0 +1,58 @@
+import { create } from 'zustand'
+import type { ClosestLandmark } from './useFaceLandmarkDetection'
+import type { ImageModificationParams } from '@/lib/facePoke'
+
+interface ImageState {
+  error: string
+  imageFile: File | null
+  originalImage: string
+  previewImage: string
+  originalImageHash: string
+  minLatency: number
+  averageLatency: number
+  maxLatency: number
+  activeLandmark?: ClosestLandmark
+  params: Partial<ImageModificationParams>
+  setError: (error?: string) => void
+  setImageFile: (file: File | null) => void
+  setOriginalImage: (url: string) => void
+  setOriginalImageHash: (hash: string) => void
+  setPreviewImage: (url: string) => void
+  resetImage: () => void
+  setAverageLatency: (averageLatency: number) => void
+  setActiveLandmark: (activeLandmark?: ClosestLandmark) => void
+  setParams: (params: Partial<ImageModificationParams>) => void
+}
+
+export const useMainStore = create<ImageState>((set, get) => ({
+  error: '',
+  imageFile: null,
+  originalImage: '',
+  originalImageHash: '',
+  previewImage: '',
+  minLatency: 20, // min time between requests
+  averageLatency: 190, // this should be the average for most people
+  maxLatency: 4000, // max time between requests
+  activeLandmark: undefined,
+  params: {},
+  setError: (error: string = '') => set({ error }),
+  setImageFile: (file) => set({ imageFile: file }),
+  setOriginalImage: (url) => set({ originalImage: url }),
+  setOriginalImageHash: (originalImageHash) => set({ originalImageHash }),
+  setPreviewImage: (url) => set({ previewImage: url }),
+  resetImage: () => {
+    const { originalImage } = get()
+    if (originalImage) {
+      set({ previewImage: originalImage })
+    }
+  },
+  setAverageLatency: (averageLatency: number) => set({ averageLatency }),
+  setActiveLandmark: (activeLandmark?: ClosestLandmark) => set({ activeLandmark }),
+  setParams: (params: Partial<ImageModificationParams>) => {
+    const {params: previousParams } = get()
+    set({ params: {
+      ...previousParams,
+      ...params
+    }})
+  },
+}))

client/src/index.tsx

@@ -0,0 +1,6 @@
+import { createRoot } from 'react-dom/client';
+
+import { App } from './app';
+
+const root = createRoot(document.getElementById('root')!);
+root.render(<App />);

client/src/layout.tsx

@@ -0,0 +1,14 @@
+import React, { type ReactNode } from 'react';
+
+export function Layout({ children }: { children: ReactNode }) {
+  return (
+    <div className="fixed min-h-screen w-full flex items-center justify-center bg-gradient-to-br from-gray-300 to-stone-300"
+      style={{ boxShadow: "inset 0 0 10vh 0 rgb(0 0 0 / 30%)" }}>
+      <div className="min-h-screen w-full py-8 flex flex-col justify-center">
+        <div className="relative p-4 sm:max-w-5xl sm:mx-auto">
+          {children}
+        </div>
+      </div>
+    </div>
+  );
+}

client/src/lib/circularBuffer.ts

@@ -0,0 +1,31 @@
+
+
+/**
+ * Circular buffer for storing and managing response times.
+ */
+export class CircularBuffer<T> {
+  private buffer: T[];
+  private pointer: number;
+
+  constructor(private capacity: number) {
+    this.buffer = new Array<T>(capacity);
+    this.pointer = 0;
+  }
+
+  /**
+   * Adds an item to the buffer, overwriting the oldest item if full.
+   * @param item - The item to add to the buffer.
+   */
+  push(item: T): void {
+    this.buffer[this.pointer] = item;
+    this.pointer = (this.pointer + 1) % this.capacity;
+  }
+
+  /**
+   * Retrieves all items currently in the buffer.
+   * @returns An array of all items in the buffer.
+   */
+  getAll(): T[] {
+    return this.buffer.filter(item => item !== undefined);
+  }
+}

client/src/lib/convertImageToBase64.ts

@@ -0,0 +1,19 @@
+export async function convertImageToBase64(imageFile: File): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+
+    reader.onload = () => {
+      if (typeof reader.result === 'string') {
+        resolve(reader.result);
+      } else {
+        reject(new Error('Failed to convert image to base64'));
+      }
+    };
+
+    reader.onerror = () => {
+      reject(new Error('Error reading file'));
+    };
+
+    reader.readAsDataURL(imageFile);
+  });
+}

client/src/lib/facePoke.ts

@@ -0,0 +1,398 @@
+import { v4 as uuidv4 } from 'uuid';
+import { CircularBuffer } from './circularBuffer';
+import { useMainStore } from '@/hooks/useMainStore';
+
+/**
+ * Represents a tracked request with its UUID and timestamp.
+ */
+export interface TrackedRequest {
+  uuid: string;
+  timestamp: number;
+}
+
+/**
+ * Represents the parameters for image modification.
+ */
+export interface ImageModificationParams {
+  eyes: number;
+  eyebrow: number;
+  wink: number;
+  pupil_x: number;
+  pupil_y: number;
+  aaa: number;
+  eee: number;
+  woo: number;
+  smile: number;
+  rotate_pitch: number;
+  rotate_yaw: number;
+  rotate_roll: number;
+}
+
+/**
+ * Represents a message to modify an image.
+ */
+export interface ModifyImageMessage {
+  type: 'modify_image';
+  image?: string;
+  image_hash?: string;
+  params: Partial<ImageModificationParams>;
+}
+
+
+/**
+ * Callback type for handling modified images.
+ */
+type OnModifiedImage = (image: string, image_hash: string) => void;
+
+/**
+ * Enum representing the different states of a WebSocket connection.
+ */
+enum WebSocketState {
+  CONNECTING = 0,
+  OPEN = 1,
+  CLOSING = 2,
+  CLOSED = 3
+}
+
+/**
+ * FacePoke class manages the WebSocket connection
+ */
+export class FacePoke {
+  private ws: WebSocket | null = null;
+  private readonly connectionId: string = uuidv4();
+  private isUnloading: boolean = false;
+  private onModifiedImage: OnModifiedImage = () => {};
+  private reconnectAttempts: number = 0;
+  private readonly maxReconnectAttempts: number = 5;
+  private readonly reconnectDelay: number = 5000;
+  private readonly eventListeners: Map<string, Set<Function>> = new Map();
+
+  private requestTracker: Map<string, TrackedRequest> = new Map();
+  private responseTimeBuffer: CircularBuffer<number>;
+  private readonly MAX_TRACKED_TIMES = 5; // Number of recent response times to track
+
+  /**
+   * Creates an instance of FacePoke.
+   * Initializes the WebSocket connection.
+   */
+  constructor() {
+    console.log(`[FacePoke] Initializing FacePoke instance with connection ID: ${this.connectionId}`);
+    this.initializeWebSocket();
+    this.setupUnloadHandler();
+
+    this.responseTimeBuffer = new CircularBuffer<number>(this.MAX_TRACKED_TIMES);
+    console.log(`[FacePoke] Initialized response time tracker with capacity: ${this.MAX_TRACKED_TIMES}`);
+  }
+
+
+  /**
+   * Generates a unique UUID for a request and starts tracking it.
+   * @returns The generated UUID for the request.
+   */
+  private trackRequest(): string {
+    const uuid = uuidv4();
+    this.requestTracker.set(uuid, { uuid, timestamp: Date.now() });
+    // console.log(`[FacePoke] Started tracking request with UUID: ${uuid}`);
+    return uuid;
+  }
+
+   /**
+   * Completes tracking for a request and updates response time statistics.
+   * @param uuid - The UUID of the completed request.
+   */
+   private completeRequest(uuid: string): void {
+    const request = this.requestTracker.get(uuid);
+    if (request) {
+      const responseTime = Date.now() - request.timestamp;
+      this.responseTimeBuffer.push(responseTime);
+      this.requestTracker.delete(uuid);
+      this.updateThrottleTime();
+      console.log(`[FacePoke] Completed request ${uuid}. Response time: ${responseTime}ms`);
+    } else {
+      console.warn(`[FacePoke] Attempted to complete unknown request: ${uuid}`);
+    }
+  }
+
+    /**
+   * Calculates the average response time from recent requests.
+   * @returns The average response time in milliseconds.
+   */
+    private calculateAverageResponseTime(): number {
+      const times = this.responseTimeBuffer.getAll();
+
+      const averageLatency = useMainStore.getState().averageLatency;
+
+      if (times.length === 0) return averageLatency;
+      const sum = times.reduce((acc, time) => acc + time, 0);
+      return sum / times.length;
+    }
+
+  /**
+   * Updates the throttle time based on recent response times.
+   */
+  private updateThrottleTime(): void {
+    const { minLatency, maxLatency, averageLatency, setAverageLatency } = useMainStore.getState();
+    const avgResponseTime = this.calculateAverageResponseTime();
+    const newLatency = Math.min(minLatency, Math.max(minLatency, avgResponseTime));
+
+    if (newLatency !== averageLatency) {
+      setAverageLatency(newLatency)
+      console.log(`[FacePoke] Updated throttle time (latency is ${newLatency}ms)`);
+    }
+  }
+
+  /**
+   * Sets the callback function for handling modified images.
+   * @param handler - The function to be called when a modified image is received.
+   */
+  public setOnModifiedImage(handler: OnModifiedImage): void {
+    this.onModifiedImage = handler;
+    console.log(`[FacePoke] onModifiedImage handler set`);
+  }
+
+  /**
+   * Starts or restarts the WebSocket connection.
+   */
+  public async startWebSocket(): Promise<void> {
+    console.log(`[FacePoke] Starting WebSocket connection.`);
+    if (!this.ws || this.ws.readyState !== WebSocketState.OPEN) {
+      await this.initializeWebSocket();
+    }
+  }
+
+  /**
+   * Initializes the WebSocket connection.
+   * Implements exponential backoff for reconnection attempts.
+   */
+  private async initializeWebSocket(): Promise<void> {
+    console.log(`[FacePoke][${this.connectionId}] Initializing WebSocket connection`);
+
+    const connect = () => {
+      this.ws = new WebSocket(`wss://${window.location.host}/ws`);
+
+      this.ws.onopen = this.handleWebSocketOpen.bind(this);
+      this.ws.onmessage = this.handleWebSocketMessage.bind(this);
+      this.ws.onclose = this.handleWebSocketClose.bind(this);
+      this.ws.onerror = this.handleWebSocketError.bind(this);
+    };
+
+    // const debouncedConnect = debounce(connect, this.reconnectDelay, { leading: true, trailing: false });
+
+    connect(); // Initial connection attempt
+  }
+
+  /**
+   * Handles the WebSocket open event.
+   */
+  private handleWebSocketOpen(): void {
+    console.log(`[FacePoke][${this.connectionId}] WebSocket connection opened`);
+    this.reconnectAttempts = 0; // Reset reconnect attempts on successful connection
+    this.emitEvent('websocketOpen');
+  }
+
+  // Update handleWebSocketMessage to complete request tracking
+  private handleWebSocketMessage(event: MessageEvent): void {
+    try {
+      const data = JSON.parse(event.data);
+      // console.log(`[FacePoke][${this.connectionId}] Received JSON data:`, data);
+
+      if (data.uuid) {
+        this.completeRequest(data.uuid);
+      }
+
+      if (data.type === 'modified_image') {
+        if (data?.image) {
+          this.onModifiedImage(data.image, data.image_hash);
+        }
+      }
+
+      this.emitEvent('message', data);
+    } catch (error) {
+      console.error(`[FacePoke][${this.connectionId}] Error parsing WebSocket message:`, error);
+    }
+  }
+
+  /**
+   * Handles WebSocket close events.
+   * Implements reconnection logic with exponential backoff.
+   * @param event - The CloseEvent containing close information.
+   */
+  private handleWebSocketClose(event: CloseEvent): void {
+    if (event.wasClean) {
+      console.log(`[FacePoke][${this.connectionId}] WebSocket connection closed cleanly, code=${event.code}, reason=${event.reason}`);
+    } else {
+      console.warn(`[FacePoke][${this.connectionId}] WebSocket connection abruptly closed`);
+    }
+
+    this.emitEvent('websocketClose', event);
+
+    // Attempt to reconnect after a delay, unless the page is unloading or max attempts reached
+    if (!this.isUnloading && this.reconnectAttempts < this.maxReconnectAttempts) {
+      this.reconnectAttempts++;
+      const delay = Math.min(1000 * (2 ** this.reconnectAttempts), 30000); // Exponential backoff, max 30 seconds
+      console.log(`[FacePoke][${this.connectionId}] Attempting to reconnect in ${delay}ms (Attempt ${this.reconnectAttempts}/${this.maxReconnectAttempts})...`);
+      setTimeout(() => this.initializeWebSocket(), delay);
+    } else if (this.reconnectAttempts >= this.maxReconnectAttempts) {
+      console.error(`[FacePoke][${this.connectionId}] Max reconnect attempts reached. Please refresh the page.`);
+      this.emitEvent('maxReconnectAttemptsReached');
+    }
+  }
+
+  /**
+   * Handles WebSocket errors.
+   * @param error - The error event.
+   */
+  private handleWebSocketError(error: Event): void {
+    console.error(`[FacePoke][${this.connectionId}] WebSocket error:`, error);
+    this.emitEvent('websocketError', error);
+  }
+
+  /**
+   * Handles interruption messages from the server.
+   * @param message - The interruption message.
+   */
+  private handleInterruption(message: string): void {
+    console.warn(`[FacePoke] Interruption: ${message}`);
+    this.emitEvent('interruption', message);
+  }
+
+  /**
+   * Toggles the microphone on or off.
+   * @param isOn - Whether to turn the microphone on (true) or off (false).
+   */
+  public async toggleMicrophone(isOn: boolean): Promise<void> {
+    console.log(`[FacePoke] Attempting to ${isOn ? 'start' : 'stop'} microphone`);
+    try {
+      if (isOn) {
+        await this.startMicrophone();
+      } else {
+        this.stopMicrophone();
+      }
+      this.emitEvent('microphoneToggled', isOn);
+    } catch (error) {
+      console.error(`[FacePoke] Error toggling microphone:`, error);
+      this.emitEvent('microphoneError', error);
+      throw error;
+    }
+  }
+
+
+/**
+ * Cleans up resources and closes connections.
+ */
+public cleanup(): void {
+  console.log('[FacePoke] Starting cleanup process');
+  if (this.ws) {
+    this.ws.close();
+    this.ws = null;
+  }
+  this.eventListeners.clear();
+  console.log('[FacePoke] Cleanup completed');
+  this.emitEvent('cleanup');
+}
+
+  /**
+   * Modifies an image based on the provided parameters
+   * @param image - The data-uri base64 image to modify.
+   * @param imageHash - The hash of the image to modify.
+   * @param params - The parameters for image modification.
+   */
+    public modifyImage(image: string | null, imageHash: string | null, params: Partial<ImageModificationParams>): void {
+      try {
+        const message: ModifyImageMessage = {
+          type: 'modify_image',
+          params: params
+        };
+
+        if (image) {
+          message.image = image;
+        } else if (imageHash) {
+          message.image_hash = imageHash;
+        } else {
+          throw new Error('Either image or imageHash must be provided');
+        }
+
+        this.sendJsonMessage(message);
+        // console.log(`[FacePoke] Sent modify image request with UUID: ${uuid}`);
+      } catch (err) {
+        console.error(`[FacePoke] Failed to modify the image:`, err);
+      }
+    }
+
+  /**
+   * Sends a JSON message through the WebSocket connection with request tracking.
+   * @param message - The message to send.
+   * @throws Error if the WebSocket is not open.
+   */
+  private sendJsonMessage<T>(message: T): void {
+    if (!this.ws || this.ws.readyState !== WebSocketState.OPEN) {
+      const error = new Error('WebSocket connection is not open');
+      console.error('[FacePoke] Error sending JSON message:', error);
+      this.emitEvent('sendJsonMessageError', error);
+      throw error;
+    }
+
+    const uuid = this.trackRequest();
+    const messageWithUuid = { ...message, uuid };
+    // console.log(`[FacePoke] Sending JSON message with UUID ${uuid}:`, messageWithUuid);
+    this.ws.send(JSON.stringify(messageWithUuid));
+  }
+
+/**
+ * Sets up the unload handler to clean up resources when the page is unloading.
+ */
+private setupUnloadHandler(): void {
+  window.addEventListener('beforeunload', () => {
+    console.log('[FacePoke] Page is unloading, cleaning up resources');
+    this.isUnloading = true;
+    if (this.ws) {
+      this.ws.close(1000, 'Page is unloading');
+    }
+    this.cleanup();
+  });
+}
+
+/**
+ * Adds an event listener for a specific event type.
+ * @param eventType - The type of event to listen for.
+ * @param listener - The function to be called when the event is emitted.
+ */
+public addEventListener(eventType: string, listener: Function): void {
+  if (!this.eventListeners.has(eventType)) {
+    this.eventListeners.set(eventType, new Set());
+  }
+  this.eventListeners.get(eventType)!.add(listener);
+  console.log(`[FacePoke] Added event listener for '${eventType}'`);
+}
+
+/**
+ * Removes an event listener for a specific event type.
+ * @param eventType - The type of event to remove the listener from.
+ * @param listener - The function to be removed from the listeners.
+ */
+public removeEventListener(eventType: string, listener: Function): void {
+  const listeners = this.eventListeners.get(eventType);
+  if (listeners) {
+    listeners.delete(listener);
+    console.log(`[FacePoke] Removed event listener for '${eventType}'`);
+  }
+}
+
+/**
+ * Emits an event to all registered listeners for that event type.
+ * @param eventType - The type of event to emit.
+ * @param data - Optional data to pass to the event listeners.
+ */
+private emitEvent(eventType: string, data?: any): void {
+  const listeners = this.eventListeners.get(eventType);
+  if (listeners) {
+    console.log(`[FacePoke] Emitting event '${eventType}' with data:`, data);
+    listeners.forEach(listener => listener(data));
+  }
+}
+}
+
+/**
+* Singleton instance of the FacePoke class.
+*/
+export const facePoke = new FacePoke();

client/src/lib/throttle.ts

@@ -0,0 +1,32 @@
+
+/**
+ * Custom throttle function that allows the first call to go through immediately
+ * and then limits subsequent calls.
+ * @param func - The function to throttle.
+ * @param limit - The minimum time between function calls in milliseconds.
+ * @returns A throttled version of the function.
+ */
+export function throttle<T extends (...args: any[]) => any>(func: T, limit: number): T {
+  let lastCall = 0;
+  let timeoutId: NodeJS.Timer | null = null;
+
+  return function (this: any, ...args: Parameters<T>) {
+    const context = this;
+    const now = Date.now();
+
+    if (now - lastCall >= limit) {
+      if (timeoutId !== null) {
+        clearTimeout(timeoutId);
+        timeoutId = null;
+      }
+      lastCall = now;
+      return func.apply(context, args);
+    } else if (!timeoutId) {
+      timeoutId = setTimeout(() => {
+        lastCall = Date.now();
+        timeoutId = null;
+        func.apply(context, args);
+      }, limit - (now - lastCall));
+    }
+  } as T;
+}

client/src/lib/utils.ts

@@ -0,0 +1,15 @@
+import { clsx, type ClassValue } from "clsx"
+import { twMerge } from "tailwind-merge"
+
+export function cn(...inputs: ClassValue[]) {
+  return twMerge(clsx(inputs))
+}
+
+export function truncateFileName(fileName: string, maxLength: number = 16) {
+  if (fileName.length <= maxLength) return fileName;
+
+  const start = fileName.slice(0, maxLength / 2 - 1);
+  const end = fileName.slice(-maxLength / 2 + 2);
+
+  return `${start}...${end}`;
+};

client/src/styles/globals.css

@@ -0,0 +1,81 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;
+
+@layer base {
+  :root {
+    --background: 0 0% 100%;
+    --foreground: 222.2 47.4% 11.2%;
+
+    --muted: 210 40% 96.1%;
+    --muted-foreground: 215.4 16.3% 46.9%;
+
+    --popover: 0 0% 100%;
+    --popover-foreground: 222.2 47.4% 11.2%;
+
+    --border: 214.3 31.8% 91.4%;
+    --input: 214.3 31.8% 91.4%;
+
+    --card: 0 0% 100%;
+    --card-foreground: 222.2 47.4% 11.2%;
+
+    --primary: 222.2 47.4% 11.2%;
+    --primary-foreground: 210 40% 98%;
+
+    --secondary: 210 40% 96.1%;
+    --secondary-foreground: 222.2 47.4% 11.2%;
+
+    --accent: 210 40% 96.1%;
+    --accent-foreground: 222.2 47.4% 11.2%;
+
+    --destructive: 0 100% 50%;
+    --destructive-foreground: 210 40% 98%;
+
+    --ring: 215 20.2% 65.1%;
+
+    --radius: 0.5rem;
+  }
+
+  .dark {
+    --background: 224 71% 4%;
+    --foreground: 213 31% 91%;
+
+    --muted: 223 47% 11%;
+    --muted-foreground: 215.4 16.3% 56.9%;
+
+    --accent: 216 34% 17%;
+    --accent-foreground: 210 40% 98%;
+
+    --popover: 224 71% 4%;
+    --popover-foreground: 215 20.2% 65.1%;
+
+    --border: 216 34% 17%;
+    --input: 216 34% 17%;
+
+    --card: 224 71% 4%;
+    --card-foreground: 213 31% 91%;
+
+    --primary: 210 40% 98%;
+    --primary-foreground: 222.2 47.4% 1.2%;
+
+    --secondary: 222.2 47.4% 11.2%;
+    --secondary-foreground: 210 40% 98%;
+
+    --destructive: 0 63% 31%;
+    --destructive-foreground: 210 40% 98%;
+
+    --ring: 216 34% 17%;
+
+    --radius: 0.5rem;
+  }
+}
+
+@layer base {
+  * {
+    @apply border-border;
+  }
+  body {
+    @apply bg-background text-foreground;
+    font-feature-settings: "rlig" 1, "calt" 1;
+  }
+}

client/tailwind.config.js

@@ -0,0 +1,86 @@
+const { fontFamily } = require("tailwindcss/defaultTheme")
+
+/** @type {import('tailwindcss').Config} */
+module.exports = {
+  darkMode: ["class"],
+  content: [
+    "app/**/*.{ts,tsx}",
+    "components/**/*.{ts,tsx}",
+    '../public/index.html'
+  ],
+  theme: {
+    container: {
+      center: true,
+      padding: "2rem",
+      screens: {
+        "2xl": "1400px",
+      },
+    },
+    extend: {
+      colors: {
+        border: "hsl(var(--border))",
+        input: "hsl(var(--input))",
+        ring: "hsl(var(--ring))",
+        background: "hsl(var(--background))",
+        foreground: "hsl(var(--foreground))",
+        primary: {
+          DEFAULT: "hsl(var(--primary))",
+          foreground: "hsl(var(--primary-foreground))",
+        },
+        secondary: {
+          DEFAULT: "hsl(var(--secondary))",
+          foreground: "hsl(var(--secondary-foreground))",
+        },
+        destructive: {
+          DEFAULT: "hsl(var(--destructive))",
+          foreground: "hsl(var(--destructive-foreground))",
+        },
+        muted: {
+          DEFAULT: "hsl(var(--muted))",
+          foreground: "hsl(var(--muted-foreground))",
+        },
+        accent: {
+          DEFAULT: "hsl(var(--accent))",
+          foreground: "hsl(var(--accent-foreground))",
+        },
+        popover: {
+          DEFAULT: "hsl(var(--popover))",
+          foreground: "hsl(var(--popover-foreground))",
+        },
+        card: {
+          DEFAULT: "hsl(var(--card))",
+          foreground: "hsl(var(--card-foreground))",
+        },
+      },
+      borderRadius: {
+        lg: `var(--radius)`,
+        md: `calc(var(--radius) - 2px)`,
+        sm: "calc(var(--radius) - 4px)",
+      },
+      fontFamily: {
+        sans: ["var(--font-sans)", ...fontFamily.sans],
+      },
+      fontSize: {
+        "5xs": "8px",
+        "4xs": "9px",
+        "3xs": "10px",
+        "2xs": "11px"
+      },
+      keyframes: {
+        "accordion-down": {
+          from: { height: "0" },
+          to: { height: "var(--radix-accordion-content-height)" },
+        },
+        "accordion-up": {
+          from: { height: "var(--radix-accordion-content-height)" },
+          to: { height: "0" },
+        },
+      },
+      animation: {
+        "accordion-down": "accordion-down 0.2s ease-out",
+        "accordion-up": "accordion-up 0.2s ease-out",
+      },
+    },
+  },
+  plugins: [require("tailwindcss-animate")],
+}

client/tsconfig.json

@@ -0,0 +1,32 @@
+{
+  "compilerOptions": {
+    // Enable latest features
+    "lib": ["ESNext", "DOM", "DOM.Iterable"],
+    "target": "ESNext",
+    "module": "ESNext",
+    "moduleDetection": "force",
+    "jsx": "react-jsx",
+    "allowJs": true,
+
+    // Bundler mode
+    "moduleResolution": "bundler",
+    "allowImportingTsExtensions": true,
+    "verbatimModuleSyntax": true,
+    "noEmit": true,
+
+    "baseUrl": ".",
+    "paths": {
+      "@/*": ["./src/*"]
+    },
+
+    // Best practices
+    "strict": true,
+    "skipLibCheck": true,
+    "noFallthroughCasesInSwitch": true,
+
+    // Some stricter flags (disabled by default)
+    "noUnusedLocals": false,
+    "noUnusedParameters": false,
+    "noPropertyAccessFromIndexSignature": false
+  }
+}

engine.py

@@ -0,0 +1,300 @@
+import time
+import logging
+import hashlib
+import uuid
+import os
+import io
+import shutil
+import asyncio
+import base64
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from typing import Dict, Any, List, Optional, AsyncGenerator, Tuple, Union
+from functools import lru_cache
+import av
+import numpy as np
+import cv2
+import torch
+import torch.nn.functional as F
+from PIL import Image
+
+from liveportrait.config.argument_config import ArgumentConfig
+from liveportrait.utils.camera import get_rotation_matrix
+from liveportrait.utils.io import load_image_rgb, load_driving_info, resize_to_limit
+from liveportrait.utils.crop import prepare_paste_back, paste_back
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Global constants
+DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
+MODELS_DIR = os.path.join(DATA_ROOT, "models")
+
+def base64_data_uri_to_PIL_Image(base64_string: str) -> Image.Image:
+    """
+    Convert a base64 data URI to a PIL Image.
+
+    Args:
+        base64_string (str): The base64 encoded image data.
+
+    Returns:
+        Image.Image: The decoded PIL Image.
+    """
+    if ',' in base64_string:
+        base64_string = base64_string.split(',')[1]
+    img_data = base64.b64decode(base64_string)
+    return Image.open(io.BytesIO(img_data))
+
+class Engine:
+    """
+    The main engine class for FacePoke
+    """
+
+    def __init__(self, live_portrait):
+        """
+        Initialize the FacePoke engine with necessary models and processors.
+
+        Args:
+            live_portrait (LivePortraitPipeline): The LivePortrait model for video generation.
+        """
+        self.live_portrait = live_portrait
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        # cache for the "modify image" workflow
+        self.image_cache = {}  # Stores the original images
+        self.processed_cache = {}  # Stores the processed image data
+
+        logger.info("✅ FacePoke Engine initialized successfully.")
+
+    def get_image_hash(self, image: Union[Image.Image, str, bytes]) -> str:
+        """
+        Compute or retrieve the hash for an image.
+
+        Args:
+            image (Union[Image.Image, str, bytes]): The input image, either as a PIL Image,
+                base64 string, or bytes.
+
+        Returns:
+            str: The computed hash of the image.
+        """
+        if isinstance(image, str):
+            # Assume it's already a hash if it's a string of the right length
+            if len(image) == 32:
+                return image
+            # Otherwise, assume it's a base64 string
+            image = base64_data_uri_to_PIL_Image(image)
+
+        if isinstance(image, Image.Image):
+            return hashlib.md5(image.tobytes()).hexdigest()
+        elif isinstance(image, bytes):
+            return hashlib.md5(image).hexdigest()
+        else:
+            raise ValueError("Unsupported image type")
+
+    @lru_cache(maxsize=128)
+    def _process_image(self, image_hash: str) -> Dict[str, Any]:
+        """
+        Process the input image and cache the results.
+
+        Args:
+            image_hash (str): Hash of the input image.
+
+        Returns:
+            Dict[str, Any]: Processed image data.
+        """
+        logger.info(f"Processing image with hash: {image_hash}")
+        if image_hash not in self.image_cache:
+            raise ValueError(f"Image with hash {image_hash} not found in cache")
+
+        image = self.image_cache[image_hash]
+        img_rgb = np.array(image)
+
+        inference_cfg = self.live_portrait.live_portrait_wrapper.cfg
+        img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n)
+        crop_info = self.live_portrait.cropper.crop_single_image(img_rgb)
+        img_crop_256x256 = crop_info['img_crop_256x256']
+
+        I_s = self.live_portrait.live_portrait_wrapper.prepare_source(img_crop_256x256)
+        x_s_info = self.live_portrait.live_portrait_wrapper.get_kp_info(I_s)
+        f_s = self.live_portrait.live_portrait_wrapper.extract_feature_3d(I_s)
+        x_s = self.live_portrait.live_portrait_wrapper.transform_keypoint(x_s_info)
+
+        processed_data = {
+            'img_rgb': img_rgb,
+            'crop_info': crop_info,
+            'x_s_info': x_s_info,
+            'f_s': f_s,
+            'x_s': x_s,
+            'inference_cfg': inference_cfg
+        }
+
+        self.processed_cache[image_hash] = processed_data
+
+        return processed_data
+
+    async def modify_image(self, image_or_hash: Union[Image.Image, str, bytes], params: Dict[str, float]) -> str:
+        """
+        Modify the input image based on the provided parameters, using caching for efficiency
+        and outputting the result as a WebP image.
+
+        Args:
+            image_or_hash (Union[Image.Image, str, bytes]): Input image as a PIL Image, base64-encoded string,
+                image bytes, or a hash string.
+            params (Dict[str, float]): Parameters for face transformation.
+
+        Returns:
+            str: Modified image as a base64-encoded WebP data URI.
+
+        Raises:
+            ValueError: If there's an error modifying the image or WebP is not supported.
+        """
+        logger.info("Starting image modification")
+        logger.debug(f"Modification parameters: {params}")
+
+        try:
+            image_hash = self.get_image_hash(image_or_hash)
+
+            # If we don't have the image in cache yet, add it
+            if image_hash not in self.image_cache:
+                if isinstance(image_or_hash, (Image.Image, bytes)):
+                    self.image_cache[image_hash] = image_or_hash
+                elif isinstance(image_or_hash, str) and len(image_or_hash) != 32:
+                    # It's a base64 string, not a hash
+                    self.image_cache[image_hash] = base64_data_uri_to_PIL_Image(image_or_hash)
+                else:
+                    raise ValueError("Image not found in cache and no valid image provided")
+
+            # Process the image (this will use the cache if available)
+            if image_hash not in self.processed_cache:
+                processed_data = await asyncio.to_thread(self._process_image, image_hash)
+            else:
+                processed_data = self.processed_cache[image_hash]
+
+            # Apply modifications based on params
+            x_d_new = processed_data['x_s_info']['kp'].clone()
+            await self._apply_facial_modifications(x_d_new, params)
+
+            # Apply rotation
+            R_new = get_rotation_matrix(
+                processed_data['x_s_info']['pitch'] + params.get('rotate_pitch', 0),
+                processed_data['x_s_info']['yaw'] + params.get('rotate_yaw', 0),
+                processed_data['x_s_info']['roll'] + params.get('rotate_roll', 0)
+            )
+            x_d_new = processed_data['x_s_info']['scale'] * (x_d_new @ R_new) + processed_data['x_s_info']['t']
+
+            # Apply stitching
+            x_d_new = await asyncio.to_thread(self.live_portrait.live_portrait_wrapper.stitching, processed_data['x_s'], x_d_new)
+
+            # Generate the output
+            out = await asyncio.to_thread(self.live_portrait.live_portrait_wrapper.warp_decode, processed_data['f_s'], processed_data['x_s'], x_d_new)
+            I_p = self.live_portrait.live_portrait_wrapper.parse_output(out['out'])[0]
+
+            # Paste back to full size
+            mask_ori = await asyncio.to_thread(
+                prepare_paste_back,
+                processed_data['inference_cfg'].mask_crop, processed_data['crop_info']['M_c2o'],
+                dsize=(processed_data['img_rgb'].shape[1], processed_data['img_rgb'].shape[0])
+            )
+            I_p_to_ori_blend = await asyncio.to_thread(
+                paste_back,
+                I_p, processed_data['crop_info']['M_c2o'], processed_data['img_rgb'], mask_ori
+            )
+
+            # Convert the result to a PIL Image
+            result_image = Image.fromarray(I_p_to_ori_blend)
+
+            # Save as WebP
+            buffered = io.BytesIO()
+            result_image.save(buffered, format="WebP", quality=85)  # Adjust quality as needed
+            modified_image_base64 = base64.b64encode(buffered.getvalue()).decode('utf-8')
+
+            logger.info("Image modification completed successfully")
+            return f"data:image/webp;base64,{modified_image_base64}"
+
+        except Exception as e:
+            logger.error(f"Error in modify_image: {str(e)}")
+            logger.exception("Full traceback:")
+            raise ValueError(f"Failed to modify image: {str(e)}")
+
+    async def _apply_facial_modifications(self, x_d_new: torch.Tensor, params: Dict[str, float]) -> None:
+        """
+        Apply facial modifications to the keypoints based on the provided parameters.
+
+        Args:
+            x_d_new (torch.Tensor): Tensor of facial keypoints to be modified.
+            params (Dict[str, float]): Parameters for face transformation.
+        """
+        modifications = [
+            ('smile', [
+                (0, 20, 1, -0.01), (0, 14, 1, -0.02), (0, 17, 1, 0.0065), (0, 17, 2, 0.003),
+                (0, 13, 1, -0.00275), (0, 16, 1, -0.00275), (0, 3, 1, -0.0035), (0, 7, 1, -0.0035)
+            ]),
+            ('aaa', [
+                (0, 19, 1, 0.001), (0, 19, 2, 0.0001), (0, 17, 1, -0.0001)
+            ]),
+            ('eee', [
+                (0, 20, 2, -0.001), (0, 20, 1, -0.001), (0, 14, 1, -0.001)
+            ]),
+            ('woo', [
+                (0, 14, 1, 0.001), (0, 3, 1, -0.0005), (0, 7, 1, -0.0005), (0, 17, 2, -0.0005)
+            ]),
+            ('wink', [
+                (0, 11, 1, 0.001), (0, 13, 1, -0.0003), (0, 17, 0, 0.0003),
+                (0, 17, 1, 0.0003), (0, 3, 1, -0.0003)
+            ]),
+            ('pupil_x', [
+                (0, 11, 0, 0.0007 if params.get('pupil_x', 0) > 0 else 0.001),
+                (0, 15, 0, 0.001 if params.get('pupil_x', 0) > 0 else 0.0007)
+            ]),
+            ('pupil_y', [
+                (0, 11, 1, -0.001), (0, 15, 1, -0.001)
+            ]),
+            ('eyes', [
+                (0, 11, 1, -0.001), (0, 13, 1, 0.0003), (0, 15, 1, -0.001), (0, 16, 1, 0.0003),
+                (0, 1, 1, -0.00025), (0, 2, 1, 0.00025)
+            ]),
+            ('eyebrow', [
+                (0, 1, 1, 0.001 if params.get('eyebrow', 0) > 0 else 0.0003),
+                (0, 2, 1, -0.001 if params.get('eyebrow', 0) > 0 else -0.0003),
+                (0, 1, 0, -0.001 if params.get('eyebrow', 0) <= 0 else 0),
+                (0, 2, 0, 0.001 if params.get('eyebrow', 0) <= 0 else 0)
+            ])
+        ]
+
+        for param_name, adjustments in modifications:
+            param_value = params.get(param_name, 0)
+            for i, j, k, factor in adjustments:
+                x_d_new[i, j, k] += param_value * factor
+
+        # Special case for pupil_y affecting eyes
+        x_d_new[0, 11, 1] -= params.get('pupil_y', 0) * 0.001
+        x_d_new[0, 15, 1] -= params.get('pupil_y', 0) * 0.001
+        params['eyes'] = params.get('eyes', 0) - params.get('pupil_y', 0) / 2.
+
+    async def cleanup(self):
+        """
+        Perform cleanup operations for the Engine.
+        This method should be called when shutting down the application.
+        """
+        logger.info("Starting Engine cleanup")
+        try:
+            # TODO: Add any additional cleanup operations here
+            logger.info("Engine cleanup completed successfully")
+        except Exception as e:
+            logger.error(f"Error during Engine cleanup: {str(e)}")
+            logger.exception("Full traceback:")
+
+def create_engine(models):
+    logger.info("Creating Engine instance...")
+
+    live_portrait = models
+
+    engine = Engine(
+        live_portrait=live_portrait,
+        # we might have more in the future
+    )
+
+    logger.info("Engine instance created successfully")
+    return engine

liveportrait/config/argument_config.py

@@ -0,0 +1,44 @@
+# coding: utf-8
+
+"""
+config for user
+"""
+
+import os.path as osp
+from dataclasses import dataclass
+import tyro
+from typing_extensions import Annotated
+from .base_config import PrintableConfig, make_abs_path
+
+
+@dataclass(repr=False)  # use repr from PrintableConfig
+class ArgumentConfig(PrintableConfig):
+    ########## input arguments ##########
+    source_image: Annotated[str, tyro.conf.arg(aliases=["-s"])] = make_abs_path('../../assets/examples/source/s6.jpg')  # path to the source portrait
+    driving_info:  Annotated[str, tyro.conf.arg(aliases=["-d"])] = make_abs_path('../../assets/examples/driving/d0.mp4')  # path to driving video or template (.pkl format)
+    output_dir: Annotated[str, tyro.conf.arg(aliases=["-o"])] = 'animations/'  # directory to save output video
+    #####################################
+
+    ########## inference arguments ##########
+    device_id: int = 0
+    flag_lip_zero : bool = True # whether let the lip to close state before animation, only take effect when flag_eye_retargeting and flag_lip_retargeting is False
+    flag_eye_retargeting: bool = False
+    flag_lip_retargeting: bool = False
+    flag_stitching: bool = True  # we recommend setting it to True!
+    flag_relative: bool = True  # whether to use relative motion
+    flag_pasteback: bool = True  # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
+    flag_do_crop: bool = True  # whether to crop the source portrait to the face-cropping space
+    flag_do_rot: bool = True  # whether to conduct the rotation when flag_do_crop is True
+    #########################################
+
+    ########## crop arguments ##########
+    dsize: int = 512
+    scale: float = 2.3
+    vx_ratio: float = 0  # vx ratio
+    vy_ratio: float = -0.125  # vy ratio +up, -down
+    ####################################
+
+    ########## gradio arguments ##########
+    server_port: Annotated[int, tyro.conf.arg(aliases=["-p"])]  = 8890
+    share: bool = True
+    server_name: str = "0.0.0.0"

liveportrait/config/base_config.py

@@ -0,0 +1,29 @@
+# coding: utf-8
+
+"""
+pretty printing class
+"""
+
+from __future__ import annotations
+import os.path as osp
+from typing import Tuple
+
+
+def make_abs_path(fn):
+    return osp.join(osp.dirname(osp.realpath(__file__)), fn)
+
+
+class PrintableConfig:  # pylint: disable=too-few-public-methods
+    """Printable Config defining str function"""
+
+    def __repr__(self):
+        lines = [self.__class__.__name__ + ":"]
+        for key, val in vars(self).items():
+            if isinstance(val, Tuple):
+                flattened_val = "["
+                for item in val:
+                    flattened_val += str(item) + "\n"
+                flattened_val = flattened_val.rstrip("\n")
+                val = flattened_val + "]"
+            lines += f"{key}: {str(val)}".split("\n")
+        return "\n    ".join(lines)

liveportrait/config/crop_config.py

@@ -0,0 +1,18 @@
+# coding: utf-8
+
+"""
+parameters used for crop faces
+"""
+
+import os.path as osp
+from dataclasses import dataclass
+from typing import Union, List
+from .base_config import PrintableConfig
+
+
+@dataclass(repr=False)  # use repr from PrintableConfig
+class CropConfig(PrintableConfig):
+    dsize: int = 512  # crop size
+    scale: float = 2.3  # scale factor
+    vx_ratio: float = 0  # vx ratio
+    vy_ratio: float = -0.125  # vy ratio +up, -down

liveportrait/config/inference_config.py

@@ -0,0 +1,53 @@
+# coding: utf-8
+
+"""
+config dataclass used for inference
+"""
+
+import os
+import os.path as osp
+from dataclasses import dataclass
+from typing import Literal, Tuple
+from .base_config import PrintableConfig, make_abs_path
+
+# Configuration
+DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
+MODELS_DIR = os.path.join(DATA_ROOT, "models")
+
+@dataclass(repr=False)  # use repr from PrintableConfig
+class InferenceConfig(PrintableConfig):
+    models_config: str = make_abs_path('./models.yaml')  # portrait animation config
+    checkpoint_F = os.path.join(MODELS_DIR, "liveportrait", "appearance_feature_extractor.pth")
+    checkpoint_M = os.path.join(MODELS_DIR, "liveportrait", "motion_extractor.pth")
+    checkpoint_W = os.path.join(MODELS_DIR, "liveportrait", "warping_module.pth")
+    checkpoint_G = os.path.join(MODELS_DIR, "liveportrait", "spade_generator.pth")
+    checkpoint_S = os.path.join(MODELS_DIR, "liveportrait", "stitching_retargeting_module.pth")
+
+    flag_use_half_precision: bool = True  # whether to use half precision
+
+    flag_lip_zero: bool = True  # whether let the lip to close state before animation, only take effect when flag_eye_retargeting and flag_lip_retargeting is False
+    lip_zero_threshold: float = 0.03
+
+    flag_eye_retargeting: bool = False
+    flag_lip_retargeting: bool = False
+    flag_stitching: bool = True  # we recommend setting it to True!
+
+    flag_relative: bool = True  # whether to use relative motion
+    anchor_frame: int = 0  # set this value if find_best_frame is True
+
+    input_shape: Tuple[int, int] = (256, 256)  # input shape
+    output_format: Literal['mp4', 'gif'] = 'mp4'  # output video format
+    output_fps: int = 25  # MuseTalk prefers 25 fps, so we use 25 as default fps for output video
+    crf: int = 15  # crf for output video
+
+    flag_write_result: bool = True  # whether to write output video
+    flag_pasteback: bool = True  # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
+    mask_crop = None
+    flag_write_gif: bool = False
+    size_gif: int = 256
+    ref_max_shape: int = 1280
+    ref_shape_n: int = 2
+
+    device_id: int = 0
+    flag_do_crop: bool = False  # whether to crop the source portrait to the face-cropping space
+    flag_do_rot: bool = True  # whether to conduct the rotation when flag_do_crop is True

liveportrait/config/models.yaml

@@ -0,0 +1,43 @@
+model_params:
+  appearance_feature_extractor_params: # the F in the paper
+    image_channel: 3
+    block_expansion: 64
+    num_down_blocks: 2
+    max_features: 512
+    reshape_channel: 32
+    reshape_depth: 16
+    num_resblocks: 6
+  motion_extractor_params: # the M in the paper
+    num_kp: 21
+    backbone: convnextv2_tiny
+  warping_module_params: # the W in the paper
+    num_kp: 21
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 2
+    reshape_channel: 32
+    estimate_occlusion_map: True
+    dense_motion_params:
+      block_expansion: 32
+      max_features: 1024
+      num_blocks: 5
+      reshape_depth: 16
+      compress: 4
+  spade_generator_params: # the G in the paper
+    upscale: 2 # represents upsample factor 256x256 -> 512x512
+    block_expansion: 64
+    max_features: 512
+    num_down_blocks: 2
+  stitching_retargeting_module_params: # the S in the paper
+    stitching:
+      input_size: 126 # (21*3)*2
+      hidden_sizes: [128, 128, 64]
+      output_size: 65 # (21*3)+2(tx,ty)
+    lip:
+      input_size: 65 # (21*3)+2
+      hidden_sizes: [128, 128, 64]
+      output_size: 63 # (21*3)
+    eye:
+      input_size: 66 # (21*3)+3
+      hidden_sizes: [256, 256, 128, 128, 64]
+      output_size: 63 # (21*3)

liveportrait/gradio_pipeline.py

@@ -0,0 +1,140 @@
+# coding: utf-8
+
+"""
+Pipeline for gradio
+"""
+import gradio as gr
+from .config.argument_config import ArgumentConfig
+from .live_portrait_pipeline import LivePortraitPipeline
+from .utils.io import load_img_online
+from .utils.rprint import rlog as log
+from .utils.crop import prepare_paste_back, paste_back
+from .utils.camera import get_rotation_matrix
+from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
+
+def update_args(args, user_args):
+    """update the args according to user inputs
+    """
+    for k, v in user_args.items():
+        if hasattr(args, k):
+            setattr(args, k, v)
+    return args
+
+class GradioPipeline(LivePortraitPipeline):
+
+    def __init__(self, inference_cfg, crop_cfg, args: ArgumentConfig):
+        super().__init__(inference_cfg, crop_cfg)
+        # self.live_portrait_wrapper = self.live_portrait_wrapper
+        self.args = args
+        # for single image retargeting
+        self.start_prepare = False
+        self.f_s_user = None
+        self.x_c_s_info_user = None
+        self.x_s_user = None
+        self.source_lmk_user = None
+        self.mask_ori = None
+        self.img_rgb = None
+        self.crop_M_c2o = None
+
+
+    def execute_video(
+        self,
+        input_image_path,
+        input_video_path,
+        flag_relative_input,
+        flag_do_crop_input,
+        flag_remap_input,
+        ):
+        """ for video driven potrait animation
+        """
+        if input_image_path is not None and input_video_path is not None:
+            args_user = {
+                'source_image': input_image_path,
+                'driving_info': input_video_path,
+                'flag_relative': flag_relative_input,
+                'flag_do_crop': flag_do_crop_input,
+                'flag_pasteback': flag_remap_input,
+            }
+            # update config from user input
+            self.args = update_args(self.args, args_user)
+            self.live_portrait_wrapper.update_config(self.args.__dict__)
+            self.cropper.update_config(self.args.__dict__)
+            # video driven animation
+            video_path, video_path_concat = self.execute(self.args)
+            gr.Info("Run successfully!", duration=2)
+            return video_path, video_path_concat,
+        else:
+            raise gr.Error("The input source portrait or driving video hasn't been prepared yet 💥!", duration=5)
+
+    def execute_image(self, input_eye_ratio: float, input_lip_ratio: float):
+        """ for single image retargeting
+        """
+        if input_eye_ratio is None or input_eye_ratio is None:
+            raise gr.Error("Invalid ratio input 💥!", duration=5)
+        elif self.f_s_user is None:
+            if self.start_prepare:
+                raise gr.Error(
+                    "The source portrait is under processing 💥! Please wait for a second.",
+                    duration=5
+                )
+            else:
+                raise gr.Error(
+                    "The source portrait hasn't been prepared yet 💥! Please scroll to the top of the page to upload.",
+                    duration=5
+                )
+        else:
+            # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
+            combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio([[input_eye_ratio]], self.source_lmk_user)
+            eyes_delta = self.live_portrait_wrapper.retarget_eye(self.x_s_user, combined_eye_ratio_tensor)
+            # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
+            combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio([[input_lip_ratio]], self.source_lmk_user)
+            lip_delta = self.live_portrait_wrapper.retarget_lip(self.x_s_user, combined_lip_ratio_tensor)
+            num_kp = self.x_s_user.shape[1]
+            # default: use x_s
+            x_d_new = self.x_s_user + eyes_delta.reshape(-1, num_kp, 3) + lip_delta.reshape(-1, num_kp, 3)
+            # D(W(f_s; x_s, x′_d))
+            out = self.live_portrait_wrapper.warp_decode(self.f_s_user, self.x_s_user, x_d_new)
+            out = self.live_portrait_wrapper.parse_output(out['out'])[0]
+            out_to_ori_blend = paste_back(out, self.crop_M_c2o, self.img_rgb, self.mask_ori)
+            gr.Info("Run successfully!", duration=2)
+            return out, out_to_ori_blend
+
+
+    def prepare_retargeting(self, input_image_path, flag_do_crop = True):
+        """ for single image retargeting
+        """
+        if input_image_path is not None:
+            gr.Info("Upload successfully!", duration=2)
+            self.start_prepare = True
+            inference_cfg = self.live_portrait_wrapper.cfg
+            ######## process source portrait ########
+            img_rgb = load_img_online(input_image_path, mode='rgb', max_dim=1280, n=16)
+            log(f"Load source image from {input_image_path}.")
+            crop_info = self.cropper.crop_single_image(img_rgb)
+            if flag_do_crop:
+                I_s = self.live_portrait_wrapper.prepare_source(crop_info['img_crop_256x256'])
+            else:
+                I_s = self.live_portrait_wrapper.prepare_source(img_rgb)
+            x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)
+            R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
+            ############################################
+
+            # record global info for next time use
+            self.f_s_user = self.live_portrait_wrapper.extract_feature_3d(I_s)
+            self.x_s_user = self.live_portrait_wrapper.transform_keypoint(x_s_info)
+            self.x_s_info_user = x_s_info
+            self.source_lmk_user = crop_info['lmk_crop']
+            self.img_rgb = img_rgb
+            self.crop_M_c2o = crop_info['M_c2o']
+            self.mask_ori = prepare_paste_back(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
+            # update slider
+            eye_close_ratio = calc_eye_close_ratio(self.source_lmk_user[None])
+            eye_close_ratio = float(eye_close_ratio.squeeze(0).mean())
+            lip_close_ratio = calc_lip_close_ratio(self.source_lmk_user[None])
+            lip_close_ratio = float(lip_close_ratio.squeeze(0).mean())
+            # for vis
+            self.I_s_vis = self.live_portrait_wrapper.parse_output(I_s)[0]
+            return eye_close_ratio, lip_close_ratio, self.I_s_vis
+        else:
+            # when press the clear button, go here
+            return 0.8, 0.8, self.I_s_vis

liveportrait/live_portrait_pipeline.py

@@ -0,0 +1,193 @@
+# coding: utf-8
+
+"""
+Pipeline of LivePortrait
+"""
+
+# TODO:
+# 1. 当前假定所有的模板都是已经裁好的，需要修改下
+# 2. pick样例图 source + driving
+
+import cv2
+import numpy as np
+import pickle
+import os.path as osp
+from rich.progress import track
+
+from .config.argument_config import ArgumentConfig
+from .config.inference_config import InferenceConfig
+from .config.crop_config import CropConfig
+from .utils.cropper import Cropper
+from .utils.camera import get_rotation_matrix
+from .utils.video import images2video, concat_frames
+from .utils.crop import _transform_img, prepare_paste_back, paste_back
+from .utils.retargeting_utils import calc_lip_close_ratio
+from .utils.io import load_image_rgb, load_driving_info, resize_to_limit
+from .utils.helper import mkdir, basename, dct2cuda, is_video, is_template
+from .utils.rprint import rlog as log
+from .live_portrait_wrapper import LivePortraitWrapper
+
+
+def make_abs_path(fn):
+    return osp.join(osp.dirname(osp.realpath(__file__)), fn)
+
+
+class LivePortraitPipeline(object):
+
+    def __init__(self, inference_cfg: InferenceConfig, crop_cfg: CropConfig):
+        self.live_portrait_wrapper: LivePortraitWrapper = LivePortraitWrapper(cfg=inference_cfg)
+        self.cropper = Cropper(crop_cfg=crop_cfg)
+
+    def execute(self, args: ArgumentConfig):
+        inference_cfg = self.live_portrait_wrapper.cfg # for convenience
+        ######## process source portrait ########
+        img_rgb = load_image_rgb(args.source_image)
+        img_rgb = resize_to_limit(img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n)
+        log(f"Load source image from {args.source_image}")
+        crop_info = self.cropper.crop_single_image(img_rgb)
+        source_lmk = crop_info['lmk_crop']
+        img_crop, img_crop_256x256 = crop_info['img_crop'], crop_info['img_crop_256x256']
+        if inference_cfg.flag_do_crop:
+            I_s = self.live_portrait_wrapper.prepare_source(img_crop_256x256)
+        else:
+            I_s = self.live_portrait_wrapper.prepare_source(img_rgb)
+        x_s_info = self.live_portrait_wrapper.get_kp_info(I_s)
+        x_c_s = x_s_info['kp']
+        R_s = get_rotation_matrix(x_s_info['pitch'], x_s_info['yaw'], x_s_info['roll'])
+        f_s = self.live_portrait_wrapper.extract_feature_3d(I_s)
+        x_s = self.live_portrait_wrapper.transform_keypoint(x_s_info)
+
+        if inference_cfg.flag_lip_zero:
+            # let lip-open scalar to be 0 at first
+            c_d_lip_before_animation = [0.]
+            combined_lip_ratio_tensor_before_animation = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_before_animation, source_lmk)
+            if combined_lip_ratio_tensor_before_animation[0][0] < inference_cfg.lip_zero_threshold:
+                inference_cfg.flag_lip_zero = False
+            else:
+                lip_delta_before_animation = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor_before_animation)
+        ############################################
+
+        ######## process driving info ########
+        if is_video(args.driving_info):
+            log(f"Load from video file (mp4 mov avi etc...): {args.driving_info}")
+            # TODO: 这里track一下驱动视频 -> 构建模板
+            driving_rgb_lst = load_driving_info(args.driving_info)
+            driving_rgb_lst_256 = [cv2.resize(_, (256, 256)) for _ in driving_rgb_lst]
+            I_d_lst = self.live_portrait_wrapper.prepare_driving_videos(driving_rgb_lst_256)
+            n_frames = I_d_lst.shape[0]
+            if inference_cfg.flag_eye_retargeting or inference_cfg.flag_lip_retargeting:
+                driving_lmk_lst = self.cropper.get_retargeting_lmk_info(driving_rgb_lst)
+                input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)
+        elif is_template(args.driving_info):
+            log(f"Load from video templates {args.driving_info}")
+            with open(args.driving_info, 'rb') as f:
+                template_lst, driving_lmk_lst = pickle.load(f)
+            n_frames = template_lst[0]['n_frames']
+            input_eye_ratio_lst, input_lip_ratio_lst = self.live_portrait_wrapper.calc_retargeting_ratio(source_lmk, driving_lmk_lst)
+        else:
+            raise Exception("Unsupported driving types!")
+        #########################################
+
+        ######## prepare for pasteback ########
+        if inference_cfg.flag_pasteback:
+            mask_ori = prepare_paste_back(inference_cfg.mask_crop, crop_info['M_c2o'], dsize=(img_rgb.shape[1], img_rgb.shape[0]))
+            I_p_paste_lst = []
+        #########################################
+
+        I_p_lst = []
+        R_d_0, x_d_0_info = None, None
+        for i in track(range(n_frames), description='Animating...', total=n_frames):
+            if is_video(args.driving_info):
+                # extract kp info by M
+                I_d_i = I_d_lst[i]
+                x_d_i_info = self.live_portrait_wrapper.get_kp_info(I_d_i)
+                R_d_i = get_rotation_matrix(x_d_i_info['pitch'], x_d_i_info['yaw'], x_d_i_info['roll'])
+            else:
+                # from template
+                x_d_i_info = template_lst[i]
+                x_d_i_info = dct2cuda(x_d_i_info, inference_cfg.device_id)
+                R_d_i = x_d_i_info['R_d']
+
+            if i == 0:
+                R_d_0 = R_d_i
+                x_d_0_info = x_d_i_info
+
+            if inference_cfg.flag_relative:
+                R_new = (R_d_i @ R_d_0.permute(0, 2, 1)) @ R_s
+                delta_new = x_s_info['exp'] + (x_d_i_info['exp'] - x_d_0_info['exp'])
+                scale_new = x_s_info['scale'] * (x_d_i_info['scale'] / x_d_0_info['scale'])
+                t_new = x_s_info['t'] + (x_d_i_info['t'] - x_d_0_info['t'])
+            else:
+                R_new = R_d_i
+                delta_new = x_d_i_info['exp']
+                scale_new = x_s_info['scale']
+                t_new = x_d_i_info['t']
+
+            t_new[..., 2].fill_(0) # zero tz
+            x_d_i_new = scale_new * (x_c_s @ R_new + delta_new) + t_new
+
+            # Algorithm 1:
+            if not inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:
+                # without stitching or retargeting
+                if inference_cfg.flag_lip_zero:
+                    x_d_i_new += lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
+                else:
+                    pass
+            elif inference_cfg.flag_stitching and not inference_cfg.flag_eye_retargeting and not inference_cfg.flag_lip_retargeting:
+                # with stitching and without retargeting
+                if inference_cfg.flag_lip_zero:
+                    x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new) + lip_delta_before_animation.reshape(-1, x_s.shape[1], 3)
+                else:
+                    x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
+            else:
+                eyes_delta, lip_delta = None, None
+                if inference_cfg.flag_eye_retargeting:
+                    c_d_eyes_i = input_eye_ratio_lst[i]
+                    combined_eye_ratio_tensor = self.live_portrait_wrapper.calc_combined_eye_ratio(c_d_eyes_i, source_lmk)
+                    # ∆_eyes,i = R_eyes(x_s; c_s,eyes, c_d,eyes,i)
+                    eyes_delta = self.live_portrait_wrapper.retarget_eye(x_s, combined_eye_ratio_tensor)
+                if inference_cfg.flag_lip_retargeting:
+                    c_d_lip_i = input_lip_ratio_lst[i]
+                    combined_lip_ratio_tensor = self.live_portrait_wrapper.calc_combined_lip_ratio(c_d_lip_i, source_lmk)
+                    # ∆_lip,i = R_lip(x_s; c_s,lip, c_d,lip,i)
+                    lip_delta = self.live_portrait_wrapper.retarget_lip(x_s, combined_lip_ratio_tensor)
+
+                if inference_cfg.flag_relative:  # use x_s
+                    x_d_i_new = x_s + \
+                        (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
+                        (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
+                else:  # use x_d,i
+                    x_d_i_new = x_d_i_new + \
+                        (eyes_delta.reshape(-1, x_s.shape[1], 3) if eyes_delta is not None else 0) + \
+                        (lip_delta.reshape(-1, x_s.shape[1], 3) if lip_delta is not None else 0)
+
+                if inference_cfg.flag_stitching:
+                    x_d_i_new = self.live_portrait_wrapper.stitching(x_s, x_d_i_new)
+
+            out = self.live_portrait_wrapper.warp_decode(f_s, x_s, x_d_i_new)
+            I_p_i = self.live_portrait_wrapper.parse_output(out['out'])[0]
+            I_p_lst.append(I_p_i)
+
+            if inference_cfg.flag_pasteback:
+                I_p_i_to_ori_blend = paste_back(I_p_i, crop_info['M_c2o'], img_rgb, mask_ori)
+                I_p_paste_lst.append(I_p_i_to_ori_blend)
+
+        mkdir(args.output_dir)
+        wfp_concat = None
+
+        # note by @jbilcke-hf:
+        # I have disabled this block, since we don't need to debug it
+        #if is_video(args.driving_info):
+        #    frames_concatenated = concat_frames(I_p_lst, driving_rgb_lst, img_crop_256x256)
+        #    # save (driving frames, source image, drived frames) result
+        #    wfp_concat = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}_concat.mp4')
+        #    images2video(frames_concatenated, wfp=wfp_concat)#
+
+        # save drived result
+        wfp = osp.join(args.output_dir, f'{basename(args.source_image)}--{basename(args.driving_info)}.mp4')
+        if inference_cfg.flag_pasteback:
+            images2video(I_p_paste_lst, wfp=wfp)
+        else:
+            images2video(I_p_lst, wfp=wfp)
+
+        return wfp, wfp_concat

liveportrait/live_portrait_wrapper.py

@@ -0,0 +1,307 @@
+# coding: utf-8
+
+"""
+Wrapper for LivePortrait core functions
+"""
+
+import os.path as osp
+import numpy as np
+import cv2
+import torch
+import yaml
+
+from .utils.timer import Timer
+from .utils.helper import load_model, concat_feat
+from .utils.camera import headpose_pred_to_degree, get_rotation_matrix
+from .utils.retargeting_utils import calc_eye_close_ratio, calc_lip_close_ratio
+from .config.inference_config import InferenceConfig
+from .utils.rprint import rlog as log
+
+
+class LivePortraitWrapper(object):
+
+    def __init__(self, cfg: InferenceConfig):
+
+        model_config = yaml.load(open(cfg.models_config, 'r'), Loader=yaml.SafeLoader)
+
+        # init F
+        self.appearance_feature_extractor = load_model(cfg.checkpoint_F, model_config, cfg.device_id, 'appearance_feature_extractor')
+        #log(f'Load appearance_feature_extractor done.')
+        # init M
+        self.motion_extractor = load_model(cfg.checkpoint_M, model_config, cfg.device_id, 'motion_extractor')
+        #log(f'Load motion_extractor done.')
+        # init W
+        self.warping_module = load_model(cfg.checkpoint_W, model_config, cfg.device_id, 'warping_module')
+        #log(f'Load warping_module done.')
+        # init G
+        self.spade_generator = load_model(cfg.checkpoint_G, model_config, cfg.device_id, 'spade_generator')
+        #log(f'Load spade_generator done.')
+        # init S and R
+        if cfg.checkpoint_S is not None and osp.exists(cfg.checkpoint_S):
+            self.stitching_retargeting_module = load_model(cfg.checkpoint_S, model_config, cfg.device_id, 'stitching_retargeting_module')
+            #log(f'Load stitching_retargeting_module done.')
+        else:
+            self.stitching_retargeting_module = None
+
+        self.cfg = cfg
+        self.device_id = cfg.device_id
+        self.timer = Timer()
+
+    def update_config(self, user_args):
+        for k, v in user_args.items():
+            if hasattr(self.cfg, k):
+                setattr(self.cfg, k, v)
+
+    def prepare_source(self, img: np.ndarray) -> torch.Tensor:
+        """ construct the input as standard
+        img: HxWx3, uint8, 256x256
+        """
+        h, w = img.shape[:2]
+        if h != self.cfg.input_shape[0] or w != self.cfg.input_shape[1]:
+            x = cv2.resize(img, (self.cfg.input_shape[0], self.cfg.input_shape[1]))
+        else:
+            x = img.copy()
+
+        if x.ndim == 3:
+            x = x[np.newaxis].astype(np.float32) / 255.  # HxWx3 -> 1xHxWx3, normalized to 0~1
+        elif x.ndim == 4:
+            x = x.astype(np.float32) / 255.  # BxHxWx3, normalized to 0~1
+        else:
+            raise ValueError(f'img ndim should be 3 or 4: {x.ndim}')
+        x = np.clip(x, 0, 1)  # clip to 0~1
+        x = torch.from_numpy(x).permute(0, 3, 1, 2)  # 1xHxWx3 -> 1x3xHxW
+        x = x.cuda(self.device_id)
+        return x
+
+    def prepare_driving_videos(self, imgs) -> torch.Tensor:
+        """ construct the input as standard
+        imgs: NxBxHxWx3, uint8
+        """
+        if isinstance(imgs, list):
+            _imgs = np.array(imgs)[..., np.newaxis]  # TxHxWx3x1
+        elif isinstance(imgs, np.ndarray):
+            _imgs = imgs
+        else:
+            raise ValueError(f'imgs type error: {type(imgs)}')
+
+        y = _imgs.astype(np.float32) / 255.
+        y = np.clip(y, 0, 1)  # clip to 0~1
+        y = torch.from_numpy(y).permute(0, 4, 3, 1, 2)  # TxHxWx3x1 -> Tx1x3xHxW
+        y = y.cuda(self.device_id)
+
+        return y
+
+    def extract_feature_3d(self, x: torch.Tensor) -> torch.Tensor:
+        """ get the appearance feature of the image by F
+        x: Bx3xHxW, normalized to 0~1
+        """
+        with torch.no_grad():
+            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):
+                feature_3d = self.appearance_feature_extractor(x)
+
+        return feature_3d.float()
+
+    def get_kp_info(self, x: torch.Tensor, **kwargs) -> dict:
+        """ get the implicit keypoint information
+        x: Bx3xHxW, normalized to 0~1
+        flag_refine_info: whether to trandform the pose to degrees and the dimention of the reshape
+        return: A dict contains keys: 'pitch', 'yaw', 'roll', 't', 'exp', 'scale', 'kp'
+        """
+        with torch.no_grad():
+            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):
+                kp_info = self.motion_extractor(x)
+
+            if self.cfg.flag_use_half_precision:
+                # float the dict
+                for k, v in kp_info.items():
+                    if isinstance(v, torch.Tensor):
+                        kp_info[k] = v.float()
+
+        flag_refine_info: bool = kwargs.get('flag_refine_info', True)
+        if flag_refine_info:
+            bs = kp_info['kp'].shape[0]
+            kp_info['pitch'] = headpose_pred_to_degree(kp_info['pitch'])[:, None]  # Bx1
+            kp_info['yaw'] = headpose_pred_to_degree(kp_info['yaw'])[:, None]  # Bx1
+            kp_info['roll'] = headpose_pred_to_degree(kp_info['roll'])[:, None]  # Bx1
+            kp_info['kp'] = kp_info['kp'].reshape(bs, -1, 3)  # BxNx3
+            kp_info['exp'] = kp_info['exp'].reshape(bs, -1, 3)  # BxNx3
+
+        return kp_info
+
+    def get_pose_dct(self, kp_info: dict) -> dict:
+        pose_dct = dict(
+            pitch=headpose_pred_to_degree(kp_info['pitch']).item(),
+            yaw=headpose_pred_to_degree(kp_info['yaw']).item(),
+            roll=headpose_pred_to_degree(kp_info['roll']).item(),
+        )
+        return pose_dct
+
+    def get_fs_and_kp_info(self, source_prepared, driving_first_frame):
+
+        # get the canonical keypoints of source image by M
+        source_kp_info = self.get_kp_info(source_prepared, flag_refine_info=True)
+        source_rotation = get_rotation_matrix(source_kp_info['pitch'], source_kp_info['yaw'], source_kp_info['roll'])
+
+        # get the canonical keypoints of first driving frame by M
+        driving_first_frame_kp_info = self.get_kp_info(driving_first_frame, flag_refine_info=True)
+        driving_first_frame_rotation = get_rotation_matrix(
+            driving_first_frame_kp_info['pitch'],
+            driving_first_frame_kp_info['yaw'],
+            driving_first_frame_kp_info['roll']
+        )
+
+        # get feature volume by F
+        source_feature_3d = self.extract_feature_3d(source_prepared)
+
+        return source_kp_info, source_rotation, source_feature_3d, driving_first_frame_kp_info, driving_first_frame_rotation
+
+    def transform_keypoint(self, kp_info: dict):
+        """
+        transform the implicit keypoints with the pose, shift, and expression deformation
+        kp: BxNx3
+        """
+        kp = kp_info['kp']    # (bs, k, 3)
+        pitch, yaw, roll = kp_info['pitch'], kp_info['yaw'], kp_info['roll']
+
+        t, exp = kp_info['t'], kp_info['exp']
+        scale = kp_info['scale']
+
+        pitch = headpose_pred_to_degree(pitch)
+        yaw = headpose_pred_to_degree(yaw)
+        roll = headpose_pred_to_degree(roll)
+
+        bs = kp.shape[0]
+        if kp.ndim == 2:
+            num_kp = kp.shape[1] // 3  # Bx(num_kpx3)
+        else:
+            num_kp = kp.shape[1]  # Bxnum_kpx3
+
+        rot_mat = get_rotation_matrix(pitch, yaw, roll)    # (bs, 3, 3)
+
+        # Eqn.2: s * (R * x_c,s + exp) + t
+        kp_transformed = kp.view(bs, num_kp, 3) @ rot_mat + exp.view(bs, num_kp, 3)
+        kp_transformed *= scale[..., None]  # (bs, k, 3) * (bs, 1, 1) = (bs, k, 3)
+        kp_transformed[:, :, 0:2] += t[:, None, 0:2]  # remove z, only apply tx ty
+
+        return kp_transformed
+
+    def retarget_eye(self, kp_source: torch.Tensor, eye_close_ratio: torch.Tensor) -> torch.Tensor:
+        """
+        kp_source: BxNx3
+        eye_close_ratio: Bx3
+        Return: Bx(3*num_kp+2)
+        """
+        feat_eye = concat_feat(kp_source, eye_close_ratio)
+
+        with torch.no_grad():
+            delta = self.stitching_retargeting_module['eye'](feat_eye)
+
+        return delta
+
+    def retarget_lip(self, kp_source: torch.Tensor, lip_close_ratio: torch.Tensor) -> torch.Tensor:
+        """
+        kp_source: BxNx3
+        lip_close_ratio: Bx2
+        """
+        feat_lip = concat_feat(kp_source, lip_close_ratio)
+
+        with torch.no_grad():
+            delta = self.stitching_retargeting_module['lip'](feat_lip)
+
+        return delta
+
+    def stitch(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:
+        """
+        kp_source: BxNx3
+        kp_driving: BxNx3
+        Return: Bx(3*num_kp+2)
+        """
+        feat_stiching = concat_feat(kp_source, kp_driving)
+
+        with torch.no_grad():
+            delta = self.stitching_retargeting_module['stitching'](feat_stiching)
+
+        return delta
+
+    def stitching(self, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:
+        """ conduct the stitching
+        kp_source: Bxnum_kpx3
+        kp_driving: Bxnum_kpx3
+        """
+
+        if self.stitching_retargeting_module is not None:
+
+            bs, num_kp = kp_source.shape[:2]
+
+            kp_driving_new = kp_driving.clone()
+            delta = self.stitch(kp_source, kp_driving_new)
+
+            delta_exp = delta[..., :3*num_kp].reshape(bs, num_kp, 3)  # 1x20x3
+            delta_tx_ty = delta[..., 3*num_kp:3*num_kp+2].reshape(bs, 1, 2)  # 1x1x2
+
+            kp_driving_new += delta_exp
+            kp_driving_new[..., :2] += delta_tx_ty
+
+            return kp_driving_new
+
+        return kp_driving
+
+    def warp_decode(self, feature_3d: torch.Tensor, kp_source: torch.Tensor, kp_driving: torch.Tensor) -> torch.Tensor:
+        """ get the image after the warping of the implicit keypoints
+        feature_3d: Bx32x16x64x64, feature volume
+        kp_source: BxNx3
+        kp_driving: BxNx3
+        """
+        # The line 18 in Algorithm 1: D(W(f_s; x_s, x′_d,i)）
+        with torch.no_grad():
+            with torch.autocast(device_type='cuda', dtype=torch.float16, enabled=self.cfg.flag_use_half_precision):
+                # get decoder input
+                ret_dct = self.warping_module(feature_3d, kp_source=kp_source, kp_driving=kp_driving)
+                # decode
+                ret_dct['out'] = self.spade_generator(feature=ret_dct['out'])
+
+            # float the dict
+            if self.cfg.flag_use_half_precision:
+                for k, v in ret_dct.items():
+                    if isinstance(v, torch.Tensor):
+                        ret_dct[k] = v.float()
+
+        return ret_dct
+
+    def parse_output(self, out: torch.Tensor) -> np.ndarray:
+        """ construct the output as standard
+        return: 1xHxWx3, uint8
+        """
+        out = np.transpose(out.data.cpu().numpy(), [0, 2, 3, 1])  # 1x3xHxW -> 1xHxWx3
+        out = np.clip(out, 0, 1)  # clip to 0~1
+        out = np.clip(out * 255, 0, 255).astype(np.uint8)  # 0~1 -> 0~255
+
+        return out
+
+    def calc_retargeting_ratio(self, source_lmk, driving_lmk_lst):
+        input_eye_ratio_lst = []
+        input_lip_ratio_lst = []
+        for lmk in driving_lmk_lst:
+            # for eyes retargeting
+            input_eye_ratio_lst.append(calc_eye_close_ratio(lmk[None]))
+            # for lip retargeting
+            input_lip_ratio_lst.append(calc_lip_close_ratio(lmk[None]))
+        return input_eye_ratio_lst, input_lip_ratio_lst
+
+    def calc_combined_eye_ratio(self, input_eye_ratio, source_lmk):
+        eye_close_ratio = calc_eye_close_ratio(source_lmk[None])
+        eye_close_ratio_tensor = torch.from_numpy(eye_close_ratio).float().cuda(self.device_id)
+        input_eye_ratio_tensor = torch.Tensor([input_eye_ratio[0][0]]).reshape(1, 1).cuda(self.device_id)
+        # [c_s,eyes, c_d,eyes,i]
+        combined_eye_ratio_tensor = torch.cat([eye_close_ratio_tensor, input_eye_ratio_tensor], dim=1)
+        return combined_eye_ratio_tensor
+
+    def calc_combined_lip_ratio(self, input_lip_ratio, source_lmk):
+        lip_close_ratio = calc_lip_close_ratio(source_lmk[None])
+        lip_close_ratio_tensor = torch.from_numpy(lip_close_ratio).float().cuda(self.device_id)
+        # [c_s,lip, c_d,lip,i]
+        input_lip_ratio_tensor = torch.Tensor([input_lip_ratio[0]]).cuda(self.device_id)
+        if input_lip_ratio_tensor.shape != [1, 1]:
+            input_lip_ratio_tensor = input_lip_ratio_tensor.reshape(1, 1)
+        combined_lip_ratio_tensor = torch.cat([lip_close_ratio_tensor, input_lip_ratio_tensor], dim=1)
+        return combined_lip_ratio_tensor

liveportrait/modules/appearance_feature_extractor.py

@@ -0,0 +1,48 @@
+# coding: utf-8
+
+"""
+Appearance extractor(F) defined in paper, which maps the source image s to a 3D appearance feature volume.
+"""
+
+import torch
+from torch import nn
+from .util import SameBlock2d, DownBlock2d, ResBlock3d
+
+
+class AppearanceFeatureExtractor(nn.Module):
+
+    def __init__(self, image_channel, block_expansion, num_down_blocks, max_features, reshape_channel, reshape_depth, num_resblocks):
+        super(AppearanceFeatureExtractor, self).__init__()
+        self.image_channel = image_channel
+        self.block_expansion = block_expansion
+        self.num_down_blocks = num_down_blocks
+        self.max_features = max_features
+        self.reshape_channel = reshape_channel
+        self.reshape_depth = reshape_depth
+
+        self.first = SameBlock2d(image_channel, block_expansion, kernel_size=(3, 3), padding=(1, 1))
+
+        down_blocks = []
+        for i in range(num_down_blocks):
+            in_features = min(max_features, block_expansion * (2 ** i))
+            out_features = min(max_features, block_expansion * (2 ** (i + 1)))
+            down_blocks.append(DownBlock2d(in_features, out_features, kernel_size=(3, 3), padding=(1, 1)))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+        self.second = nn.Conv2d(in_channels=out_features, out_channels=max_features, kernel_size=1, stride=1)
+
+        self.resblocks_3d = torch.nn.Sequential()
+        for i in range(num_resblocks):
+            self.resblocks_3d.add_module('3dr' + str(i), ResBlock3d(reshape_channel, kernel_size=3, padding=1))
+
+    def forward(self, source_image):
+        out = self.first(source_image)  # Bx3x256x256 -> Bx64x256x256
+
+        for i in range(len(self.down_blocks)):
+            out = self.down_blocks[i](out)
+        out = self.second(out)
+        bs, c, h, w = out.shape  # ->Bx512x64x64
+
+        f_s = out.view(bs, self.reshape_channel, self.reshape_depth, h, w)  # ->Bx32x16x64x64
+        f_s = self.resblocks_3d(f_s)  # ->Bx32x16x64x64
+        return f_s

liveportrait/modules/convnextv2.py

@@ -0,0 +1,149 @@
+# coding: utf-8
+
+"""
+This moudle is adapted to the ConvNeXtV2 version for the extraction of implicit keypoints, poses, and expression deformation.
+"""
+
+import torch
+import torch.nn as nn
+# from timm.models.layers import trunc_normal_, DropPath
+from .util import LayerNorm, DropPath, trunc_normal_, GRN
+
+__all__ = ['convnextv2_tiny']
+
+
+class Block(nn.Module):
+    """ ConvNeXtV2 Block.
+
+    Args:
+        dim (int): Number of input channels.
+        drop_path (float): Stochastic depth rate. Default: 0.0
+    """
+
+    def __init__(self, dim, drop_path=0.):
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)  # depthwise conv
+        self.norm = LayerNorm(dim, eps=1e-6)
+        self.pwconv1 = nn.Linear(dim, 4 * dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.act = nn.GELU()
+        self.grn = GRN(4 * dim)
+        self.pwconv2 = nn.Linear(4 * dim, dim)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+
+    def forward(self, x):
+        input = x
+        x = self.dwconv(x)
+        x = x.permute(0, 2, 3, 1)  # (N, C, H, W) -> (N, H, W, C)
+        x = self.norm(x)
+        x = self.pwconv1(x)
+        x = self.act(x)
+        x = self.grn(x)
+        x = self.pwconv2(x)
+        x = x.permute(0, 3, 1, 2)  # (N, H, W, C) -> (N, C, H, W)
+
+        x = input + self.drop_path(x)
+        return x
+
+
+class ConvNeXtV2(nn.Module):
+    """ ConvNeXt V2
+
+    Args:
+        in_chans (int): Number of input image channels. Default: 3
+        num_classes (int): Number of classes for classification head. Default: 1000
+        depths (tuple(int)): Number of blocks at each stage. Default: [3, 3, 9, 3]
+        dims (int): Feature dimension at each stage. Default: [96, 192, 384, 768]
+        drop_path_rate (float): Stochastic depth rate. Default: 0.
+        head_init_scale (float): Init scaling value for classifier weights and biases. Default: 1.
+    """
+
+    def __init__(
+        self,
+        in_chans=3,
+        depths=[3, 3, 9, 3],
+        dims=[96, 192, 384, 768],
+        drop_path_rate=0.,
+        **kwargs
+    ):
+        super().__init__()
+        self.depths = depths
+        self.downsample_layers = nn.ModuleList()  # stem and 3 intermediate downsampling conv layers
+        stem = nn.Sequential(
+            nn.Conv2d(in_chans, dims[0], kernel_size=4, stride=4),
+            LayerNorm(dims[0], eps=1e-6, data_format="channels_first")
+        )
+        self.downsample_layers.append(stem)
+        for i in range(3):
+            downsample_layer = nn.Sequential(
+                LayerNorm(dims[i], eps=1e-6, data_format="channels_first"),
+                nn.Conv2d(dims[i], dims[i+1], kernel_size=2, stride=2),
+            )
+            self.downsample_layers.append(downsample_layer)
+
+        self.stages = nn.ModuleList()  # 4 feature resolution stages, each consisting of multiple residual blocks
+        dp_rates = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(4):
+            stage = nn.Sequential(
+                *[Block(dim=dims[i], drop_path=dp_rates[cur + j]) for j in range(depths[i])]
+            )
+            self.stages.append(stage)
+            cur += depths[i]
+
+        self.norm = nn.LayerNorm(dims[-1], eps=1e-6)  # final norm layer
+
+        # NOTE: the output semantic items
+        num_bins = kwargs.get('num_bins', 66)
+        num_kp = kwargs.get('num_kp', 24)  # the number of implicit keypoints
+        self.fc_kp = nn.Linear(dims[-1], 3 * num_kp)  # implicit keypoints
+
+        # print('dims[-1]: ', dims[-1])
+        self.fc_scale = nn.Linear(dims[-1], 1)  # scale
+        self.fc_pitch = nn.Linear(dims[-1], num_bins)  # pitch bins
+        self.fc_yaw = nn.Linear(dims[-1], num_bins)  # yaw bins
+        self.fc_roll = nn.Linear(dims[-1], num_bins)  # roll bins
+        self.fc_t = nn.Linear(dims[-1], 3)  # translation
+        self.fc_exp = nn.Linear(dims[-1], 3 * num_kp)  # expression / delta
+
+    def _init_weights(self, m):
+        if isinstance(m, (nn.Conv2d, nn.Linear)):
+            trunc_normal_(m.weight, std=.02)
+            nn.init.constant_(m.bias, 0)
+
+    def forward_features(self, x):
+        for i in range(4):
+            x = self.downsample_layers[i](x)
+            x = self.stages[i](x)
+        return self.norm(x.mean([-2, -1]))  # global average pooling, (N, C, H, W) -> (N, C)
+
+    def forward(self, x):
+        x = self.forward_features(x)
+
+        # implicit keypoints
+        kp = self.fc_kp(x)
+
+        # pose and expression deformation
+        pitch = self.fc_pitch(x)
+        yaw = self.fc_yaw(x)
+        roll = self.fc_roll(x)
+        t = self.fc_t(x)
+        exp = self.fc_exp(x)
+        scale = self.fc_scale(x)
+
+        ret_dct = {
+            'pitch': pitch,
+            'yaw': yaw,
+            'roll': roll,
+            't': t,
+            'exp': exp,
+            'scale': scale,
+
+            'kp': kp,  # canonical keypoint
+        }
+
+        return ret_dct
+
+
+def convnextv2_tiny(**kwargs):
+    model = ConvNeXtV2(depths=[3, 3, 9, 3], dims=[96, 192, 384, 768], **kwargs)
+    return model

liveportrait/modules/dense_motion.py

@@ -0,0 +1,104 @@
+# coding: utf-8
+
+"""
+The module that predicting a dense motion from sparse motion representation given by kp_source and kp_driving
+"""
+
+from torch import nn
+import torch.nn.functional as F
+import torch
+from .util import Hourglass, make_coordinate_grid, kp2gaussian
+
+
+class DenseMotionNetwork(nn.Module):
+    def __init__(self, block_expansion, num_blocks, max_features, num_kp, feature_channel, reshape_depth, compress, estimate_occlusion_map=True):
+        super(DenseMotionNetwork, self).__init__()
+        self.hourglass = Hourglass(block_expansion=block_expansion, in_features=(num_kp+1)*(compress+1), max_features=max_features, num_blocks=num_blocks)  # ~60+G
+
+        self.mask = nn.Conv3d(self.hourglass.out_filters, num_kp + 1, kernel_size=7, padding=3)  # 65G! NOTE: computation cost is large
+        self.compress = nn.Conv3d(feature_channel, compress, kernel_size=1)  # 0.8G
+        self.norm = nn.BatchNorm3d(compress, affine=True)
+        self.num_kp = num_kp
+        self.flag_estimate_occlusion_map = estimate_occlusion_map
+
+        if self.flag_estimate_occlusion_map:
+            self.occlusion = nn.Conv2d(self.hourglass.out_filters*reshape_depth, 1, kernel_size=7, padding=3)
+        else:
+            self.occlusion = None
+
+    def create_sparse_motions(self, feature, kp_driving, kp_source):
+        bs, _, d, h, w = feature.shape  # (bs, 4, 16, 64, 64)
+        identity_grid = make_coordinate_grid((d, h, w), ref=kp_source)  # (16, 64, 64, 3)
+        identity_grid = identity_grid.view(1, 1, d, h, w, 3)  # (1, 1, d=16, h=64, w=64, 3)
+        coordinate_grid = identity_grid - kp_driving.view(bs, self.num_kp, 1, 1, 1, 3)
+
+        k = coordinate_grid.shape[1]
+
+        # NOTE: there lacks an one-order flow
+        driving_to_source = coordinate_grid + kp_source.view(bs, self.num_kp, 1, 1, 1, 3)    # (bs, num_kp, d, h, w, 3)
+
+        # adding background feature
+        identity_grid = identity_grid.repeat(bs, 1, 1, 1, 1, 1)
+        sparse_motions = torch.cat([identity_grid, driving_to_source], dim=1)  # (bs, 1+num_kp, d, h, w, 3)
+        return sparse_motions
+
+    def create_deformed_feature(self, feature, sparse_motions):
+        bs, _, d, h, w = feature.shape
+        feature_repeat = feature.unsqueeze(1).unsqueeze(1).repeat(1, self.num_kp+1, 1, 1, 1, 1, 1)      # (bs, num_kp+1, 1, c, d, h, w)
+        feature_repeat = feature_repeat.view(bs * (self.num_kp+1), -1, d, h, w)                         # (bs*(num_kp+1), c, d, h, w)
+        sparse_motions = sparse_motions.view((bs * (self.num_kp+1), d, h, w, -1))                       # (bs*(num_kp+1), d, h, w, 3)
+        sparse_deformed = F.grid_sample(feature_repeat, sparse_motions, align_corners=False)
+        sparse_deformed = sparse_deformed.view((bs, self.num_kp+1, -1, d, h, w))                        # (bs, num_kp+1, c, d, h, w)
+
+        return sparse_deformed
+
+    def create_heatmap_representations(self, feature, kp_driving, kp_source):
+        spatial_size = feature.shape[3:]  # (d=16, h=64, w=64)
+        gaussian_driving = kp2gaussian(kp_driving, spatial_size=spatial_size, kp_variance=0.01)  # (bs, num_kp, d, h, w)
+        gaussian_source = kp2gaussian(kp_source, spatial_size=spatial_size, kp_variance=0.01)  # (bs, num_kp, d, h, w)
+        heatmap = gaussian_driving - gaussian_source  # (bs, num_kp, d, h, w)
+
+        # adding background feature
+        zeros = torch.zeros(heatmap.shape[0], 1, spatial_size[0], spatial_size[1], spatial_size[2]).type(heatmap.type()).to(heatmap.device)
+        heatmap = torch.cat([zeros, heatmap], dim=1)
+        heatmap = heatmap.unsqueeze(2)         # (bs, 1+num_kp, 1, d, h, w)
+        return heatmap
+
+    def forward(self, feature, kp_driving, kp_source):
+        bs, _, d, h, w = feature.shape  # (bs, 32, 16, 64, 64)
+
+        feature = self.compress(feature)  # (bs, 4, 16, 64, 64)
+        feature = self.norm(feature)  # (bs, 4, 16, 64, 64)
+        feature = F.relu(feature)  # (bs, 4, 16, 64, 64)
+
+        out_dict = dict()
+
+        # 1. deform 3d feature
+        sparse_motion = self.create_sparse_motions(feature, kp_driving, kp_source)  # (bs, 1+num_kp, d, h, w, 3)
+        deformed_feature = self.create_deformed_feature(feature, sparse_motion)  # (bs, 1+num_kp, c=4, d=16, h=64, w=64)
+
+        # 2. (bs, 1+num_kp, d, h, w)
+        heatmap = self.create_heatmap_representations(deformed_feature, kp_driving, kp_source)  # (bs, 1+num_kp, 1, d, h, w)
+
+        input = torch.cat([heatmap, deformed_feature], dim=2)  # (bs, 1+num_kp, c=5, d=16, h=64, w=64)
+        input = input.view(bs, -1, d, h, w)  # (bs, (1+num_kp)*c=105, d=16, h=64, w=64)
+
+        prediction = self.hourglass(input)
+
+        mask = self.mask(prediction)
+        mask = F.softmax(mask, dim=1)  # (bs, 1+num_kp, d=16, h=64, w=64)
+        out_dict['mask'] = mask
+        mask = mask.unsqueeze(2)                                   # (bs, num_kp+1, 1, d, h, w)
+        sparse_motion = sparse_motion.permute(0, 1, 5, 2, 3, 4)    # (bs, num_kp+1, 3, d, h, w)
+        deformation = (sparse_motion * mask).sum(dim=1)            # (bs, 3, d, h, w)  mask take effect in this place
+        deformation = deformation.permute(0, 2, 3, 4, 1)           # (bs, d, h, w, 3)
+
+        out_dict['deformation'] = deformation
+
+        if self.flag_estimate_occlusion_map:
+            bs, _, d, h, w = prediction.shape
+            prediction_reshape = prediction.view(bs, -1, h, w)
+            occlusion_map = torch.sigmoid(self.occlusion(prediction_reshape))  # Bx1x64x64
+            out_dict['occlusion_map'] = occlusion_map
+
+        return out_dict

liveportrait/modules/motion_extractor.py

@@ -0,0 +1,35 @@
+# coding: utf-8
+
+"""
+Motion extractor(M), which directly predicts the canonical keypoints, head pose and expression deformation of the input image
+"""
+
+from torch import nn
+import torch
+
+from .convnextv2 import convnextv2_tiny
+from .util import filter_state_dict
+
+model_dict = {
+    'convnextv2_tiny': convnextv2_tiny,
+}
+
+
+class MotionExtractor(nn.Module):
+    def __init__(self, **kwargs):
+        super(MotionExtractor, self).__init__()
+
+        # default is convnextv2_base
+        backbone = kwargs.get('backbone', 'convnextv2_tiny')
+        self.detector = model_dict.get(backbone)(**kwargs)
+
+    def load_pretrained(self, init_path: str):
+        if init_path not in (None, ''):
+            state_dict = torch.load(init_path, map_location=lambda storage, loc: storage)['model']
+            state_dict = filter_state_dict(state_dict, remove_name='head')
+            ret = self.detector.load_state_dict(state_dict, strict=False)
+            print(f'Load pretrained model from {init_path}, ret: {ret}')
+
+    def forward(self, x):
+        out = self.detector(x)
+        return out

liveportrait/modules/spade_generator.py

@@ -0,0 +1,59 @@
+# coding: utf-8
+
+"""
+Spade decoder(G) defined in the paper, which input the warped feature to generate the animated image.
+"""
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+from .util import SPADEResnetBlock
+
+
+class SPADEDecoder(nn.Module):
+    def __init__(self, upscale=1, max_features=256, block_expansion=64, out_channels=64, num_down_blocks=2):
+        for i in range(num_down_blocks):
+            input_channels = min(max_features, block_expansion * (2 ** (i + 1)))
+        self.upscale = upscale
+        super().__init__()
+        norm_G = 'spadespectralinstance'
+        label_num_channels = input_channels  # 256
+
+        self.fc = nn.Conv2d(input_channels, 2 * input_channels, 3, padding=1)
+        self.G_middle_0 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
+        self.G_middle_1 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
+        self.G_middle_2 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
+        self.G_middle_3 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
+        self.G_middle_4 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
+        self.G_middle_5 = SPADEResnetBlock(2 * input_channels, 2 * input_channels, norm_G, label_num_channels)
+        self.up_0 = SPADEResnetBlock(2 * input_channels, input_channels, norm_G, label_num_channels)
+        self.up_1 = SPADEResnetBlock(input_channels, out_channels, norm_G, label_num_channels)
+        self.up = nn.Upsample(scale_factor=2)
+
+        if self.upscale is None or self.upscale <= 1:
+            self.conv_img = nn.Conv2d(out_channels, 3, 3, padding=1)
+        else:
+            self.conv_img = nn.Sequential(
+                nn.Conv2d(out_channels, 3 * (2 * 2), kernel_size=3, padding=1),
+                nn.PixelShuffle(upscale_factor=2)
+            )
+
+    def forward(self, feature):
+        seg = feature  # Bx256x64x64
+        x = self.fc(feature)  # Bx512x64x64
+        x = self.G_middle_0(x, seg)
+        x = self.G_middle_1(x, seg)
+        x = self.G_middle_2(x, seg)
+        x = self.G_middle_3(x, seg)
+        x = self.G_middle_4(x, seg)
+        x = self.G_middle_5(x, seg)
+
+        x = self.up(x)  # Bx512x64x64 -> Bx512x128x128
+        x = self.up_0(x, seg)  # Bx512x128x128 -> Bx256x128x128
+        x = self.up(x)  # Bx256x128x128 -> Bx256x256x256
+        x = self.up_1(x, seg)  # Bx256x256x256 -> Bx64x256x256
+
+        x = self.conv_img(F.leaky_relu(x, 2e-1))  # Bx64x256x256 -> Bx3xHxW
+        x = torch.sigmoid(x)  # Bx3xHxW
+
+        return x

liveportrait/modules/stitching_retargeting_network.py

@@ -0,0 +1,38 @@
+# coding: utf-8
+
+"""
+Stitching module(S) and two retargeting modules(R) defined in the paper.
+
+- The stitching module pastes the animated portrait back into the original image space without pixel misalignment, such as in
+the stitching region.
+
+- The eyes retargeting module is designed to address the issue of incomplete eye closure during cross-id reenactment, especially
+when a person with small eyes drives a person with larger eyes.
+
+- The lip retargeting module is designed similarly to the eye retargeting module, and can also normalize the input by ensuring that
+the lips are in a closed state, which facilitates better animation driving.
+"""
+from torch import nn
+
+
+class StitchingRetargetingNetwork(nn.Module):
+    def __init__(self, input_size, hidden_sizes, output_size):
+        super(StitchingRetargetingNetwork, self).__init__()
+        layers = []
+        for i in range(len(hidden_sizes)):
+            if i == 0:
+                layers.append(nn.Linear(input_size, hidden_sizes[i]))
+            else:
+                layers.append(nn.Linear(hidden_sizes[i - 1], hidden_sizes[i]))
+            layers.append(nn.ReLU(inplace=True))
+        layers.append(nn.Linear(hidden_sizes[-1], output_size))
+        self.mlp = nn.Sequential(*layers)
+
+    def initialize_weights_to_zero(self):
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.zeros_(m.weight)
+                nn.init.zeros_(m.bias)
+
+    def forward(self, x):
+        return self.mlp(x)

liveportrait/modules/util.py

@@ -0,0 +1,441 @@
+# coding: utf-8
+
+"""
+This file defines various neural network modules and utility functions, including convolutional and residual blocks,
+normalizations, and functions for spatial transformation and tensor manipulation.
+"""
+
+from torch import nn
+import torch.nn.functional as F
+import torch
+import torch.nn.utils.spectral_norm as spectral_norm
+import math
+import warnings
+
+
+def kp2gaussian(kp, spatial_size, kp_variance):
+    """
+    Transform a keypoint into gaussian like representation
+    """
+    mean = kp
+
+    coordinate_grid = make_coordinate_grid(spatial_size, mean)
+    number_of_leading_dimensions = len(mean.shape) - 1
+    shape = (1,) * number_of_leading_dimensions + coordinate_grid.shape
+    coordinate_grid = coordinate_grid.view(*shape)
+    repeats = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 1)
+    coordinate_grid = coordinate_grid.repeat(*repeats)
+
+    # Preprocess kp shape
+    shape = mean.shape[:number_of_leading_dimensions] + (1, 1, 1, 3)
+    mean = mean.view(*shape)
+
+    mean_sub = (coordinate_grid - mean)
+
+    out = torch.exp(-0.5 * (mean_sub ** 2).sum(-1) / kp_variance)
+
+    return out
+
+
+def make_coordinate_grid(spatial_size, ref, **kwargs):
+    d, h, w = spatial_size
+    x = torch.arange(w).type(ref.dtype).to(ref.device)
+    y = torch.arange(h).type(ref.dtype).to(ref.device)
+    z = torch.arange(d).type(ref.dtype).to(ref.device)
+
+    # NOTE: must be right-down-in
+    x = (2 * (x / (w - 1)) - 1)  # the x axis faces to the right
+    y = (2 * (y / (h - 1)) - 1)  # the y axis faces to the bottom
+    z = (2 * (z / (d - 1)) - 1)  # the z axis faces to the inner
+
+    yy = y.view(1, -1, 1).repeat(d, 1, w)
+    xx = x.view(1, 1, -1).repeat(d, h, 1)
+    zz = z.view(-1, 1, 1).repeat(1, h, w)
+
+    meshed = torch.cat([xx.unsqueeze_(3), yy.unsqueeze_(3), zz.unsqueeze_(3)], 3)
+
+    return meshed
+
+
+class ConvT2d(nn.Module):
+    """
+    Upsampling block for use in decoder.
+    """
+
+    def __init__(self, in_features, out_features, kernel_size=3, stride=2, padding=1, output_padding=1):
+        super(ConvT2d, self).__init__()
+
+        self.convT = nn.ConvTranspose2d(in_features, out_features, kernel_size=kernel_size, stride=stride,
+                                        padding=padding, output_padding=output_padding)
+        self.norm = nn.InstanceNorm2d(out_features)
+
+    def forward(self, x):
+        out = self.convT(x)
+        out = self.norm(out)
+        out = F.leaky_relu(out)
+        return out
+
+
+class ResBlock3d(nn.Module):
+    """
+    Res block, preserve spatial resolution.
+    """
+
+    def __init__(self, in_features, kernel_size, padding):
+        super(ResBlock3d, self).__init__()
+        self.conv1 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding)
+        self.conv2 = nn.Conv3d(in_channels=in_features, out_channels=in_features, kernel_size=kernel_size, padding=padding)
+        self.norm1 = nn.BatchNorm3d(in_features, affine=True)
+        self.norm2 = nn.BatchNorm3d(in_features, affine=True)
+
+    def forward(self, x):
+        out = self.norm1(x)
+        out = F.relu(out)
+        out = self.conv1(out)
+        out = self.norm2(out)
+        out = F.relu(out)
+        out = self.conv2(out)
+        out += x
+        return out
+
+
+class UpBlock3d(nn.Module):
+    """
+    Upsampling block for use in decoder.
+    """
+
+    def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
+        super(UpBlock3d, self).__init__()
+
+        self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
+                              padding=padding, groups=groups)
+        self.norm = nn.BatchNorm3d(out_features, affine=True)
+
+    def forward(self, x):
+        out = F.interpolate(x, scale_factor=(1, 2, 2))
+        out = self.conv(out)
+        out = self.norm(out)
+        out = F.relu(out)
+        return out
+
+
+class DownBlock2d(nn.Module):
+    """
+    Downsampling block for use in encoder.
+    """
+
+    def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
+        super(DownBlock2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups)
+        self.norm = nn.BatchNorm2d(out_features, affine=True)
+        self.pool = nn.AvgPool2d(kernel_size=(2, 2))
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.norm(out)
+        out = F.relu(out)
+        out = self.pool(out)
+        return out
+
+
+class DownBlock3d(nn.Module):
+    """
+    Downsampling block for use in encoder.
+    """
+
+    def __init__(self, in_features, out_features, kernel_size=3, padding=1, groups=1):
+        super(DownBlock3d, self).__init__()
+        '''
+        self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
+                                padding=padding, groups=groups, stride=(1, 2, 2))
+        '''
+        self.conv = nn.Conv3d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size,
+                              padding=padding, groups=groups)
+        self.norm = nn.BatchNorm3d(out_features, affine=True)
+        self.pool = nn.AvgPool3d(kernel_size=(1, 2, 2))
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.norm(out)
+        out = F.relu(out)
+        out = self.pool(out)
+        return out
+
+
+class SameBlock2d(nn.Module):
+    """
+    Simple block, preserve spatial resolution.
+    """
+
+    def __init__(self, in_features, out_features, groups=1, kernel_size=3, padding=1, lrelu=False):
+        super(SameBlock2d, self).__init__()
+        self.conv = nn.Conv2d(in_channels=in_features, out_channels=out_features, kernel_size=kernel_size, padding=padding, groups=groups)
+        self.norm = nn.BatchNorm2d(out_features, affine=True)
+        if lrelu:
+            self.ac = nn.LeakyReLU()
+        else:
+            self.ac = nn.ReLU()
+
+    def forward(self, x):
+        out = self.conv(x)
+        out = self.norm(out)
+        out = self.ac(out)
+        return out
+
+
+class Encoder(nn.Module):
+    """
+    Hourglass Encoder
+    """
+
+    def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
+        super(Encoder, self).__init__()
+
+        down_blocks = []
+        for i in range(num_blocks):
+            down_blocks.append(DownBlock3d(in_features if i == 0 else min(max_features, block_expansion * (2 ** i)), min(max_features, block_expansion * (2 ** (i + 1))), kernel_size=3, padding=1))
+        self.down_blocks = nn.ModuleList(down_blocks)
+
+    def forward(self, x):
+        outs = [x]
+        for down_block in self.down_blocks:
+            outs.append(down_block(outs[-1]))
+        return outs
+
+
+class Decoder(nn.Module):
+    """
+    Hourglass Decoder
+    """
+
+    def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
+        super(Decoder, self).__init__()
+
+        up_blocks = []
+
+        for i in range(num_blocks)[::-1]:
+            in_filters = (1 if i == num_blocks - 1 else 2) * min(max_features, block_expansion * (2 ** (i + 1)))
+            out_filters = min(max_features, block_expansion * (2 ** i))
+            up_blocks.append(UpBlock3d(in_filters, out_filters, kernel_size=3, padding=1))
+
+        self.up_blocks = nn.ModuleList(up_blocks)
+        self.out_filters = block_expansion + in_features
+
+        self.conv = nn.Conv3d(in_channels=self.out_filters, out_channels=self.out_filters, kernel_size=3, padding=1)
+        self.norm = nn.BatchNorm3d(self.out_filters, affine=True)
+
+    def forward(self, x):
+        out = x.pop()
+        for up_block in self.up_blocks:
+            out = up_block(out)
+            skip = x.pop()
+            out = torch.cat([out, skip], dim=1)
+        out = self.conv(out)
+        out = self.norm(out)
+        out = F.relu(out)
+        return out
+
+
+class Hourglass(nn.Module):
+    """
+    Hourglass architecture.
+    """
+
+    def __init__(self, block_expansion, in_features, num_blocks=3, max_features=256):
+        super(Hourglass, self).__init__()
+        self.encoder = Encoder(block_expansion, in_features, num_blocks, max_features)
+        self.decoder = Decoder(block_expansion, in_features, num_blocks, max_features)
+        self.out_filters = self.decoder.out_filters
+
+    def forward(self, x):
+        return self.decoder(self.encoder(x))
+
+
+class SPADE(nn.Module):
+    def __init__(self, norm_nc, label_nc):
+        super().__init__()
+
+        self.param_free_norm = nn.InstanceNorm2d(norm_nc, affine=False)
+        nhidden = 128
+
+        self.mlp_shared = nn.Sequential(
+            nn.Conv2d(label_nc, nhidden, kernel_size=3, padding=1),
+            nn.ReLU())
+        self.mlp_gamma = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+        self.mlp_beta = nn.Conv2d(nhidden, norm_nc, kernel_size=3, padding=1)
+
+    def forward(self, x, segmap):
+        normalized = self.param_free_norm(x)
+        segmap = F.interpolate(segmap, size=x.size()[2:], mode='nearest')
+        actv = self.mlp_shared(segmap)
+        gamma = self.mlp_gamma(actv)
+        beta = self.mlp_beta(actv)
+        out = normalized * (1 + gamma) + beta
+        return out
+
+
+class SPADEResnetBlock(nn.Module):
+    def __init__(self, fin, fout, norm_G, label_nc, use_se=False, dilation=1):
+        super().__init__()
+        # Attributes
+        self.learned_shortcut = (fin != fout)
+        fmiddle = min(fin, fout)
+        self.use_se = use_se
+        # create conv layers
+        self.conv_0 = nn.Conv2d(fin, fmiddle, kernel_size=3, padding=dilation, dilation=dilation)
+        self.conv_1 = nn.Conv2d(fmiddle, fout, kernel_size=3, padding=dilation, dilation=dilation)
+        if self.learned_shortcut:
+            self.conv_s = nn.Conv2d(fin, fout, kernel_size=1, bias=False)
+        # apply spectral norm if specified
+        if 'spectral' in norm_G:
+            self.conv_0 = spectral_norm(self.conv_0)
+            self.conv_1 = spectral_norm(self.conv_1)
+            if self.learned_shortcut:
+                self.conv_s = spectral_norm(self.conv_s)
+        # define normalization layers
+        self.norm_0 = SPADE(fin, label_nc)
+        self.norm_1 = SPADE(fmiddle, label_nc)
+        if self.learned_shortcut:
+            self.norm_s = SPADE(fin, label_nc)
+
+    def forward(self, x, seg1):
+        x_s = self.shortcut(x, seg1)
+        dx = self.conv_0(self.actvn(self.norm_0(x, seg1)))
+        dx = self.conv_1(self.actvn(self.norm_1(dx, seg1)))
+        out = x_s + dx
+        return out
+
+    def shortcut(self, x, seg1):
+        if self.learned_shortcut:
+            x_s = self.conv_s(self.norm_s(x, seg1))
+        else:
+            x_s = x
+        return x_s
+
+    def actvn(self, x):
+        return F.leaky_relu(x, 2e-1)
+
+
+def filter_state_dict(state_dict, remove_name='fc'):
+    new_state_dict = {}
+    for key in state_dict:
+        if remove_name in key:
+            continue
+        new_state_dict[key] = state_dict[key]
+    return new_state_dict
+
+
+class GRN(nn.Module):
+    """ GRN (Global Response Normalization) layer
+    """
+
+    def __init__(self, dim):
+        super().__init__()
+        self.gamma = nn.Parameter(torch.zeros(1, 1, 1, dim))
+        self.beta = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+    def forward(self, x):
+        Gx = torch.norm(x, p=2, dim=(1, 2), keepdim=True)
+        Nx = Gx / (Gx.mean(dim=-1, keepdim=True) + 1e-6)
+        return self.gamma * (x * Nx) + self.beta + x
+
+
+class LayerNorm(nn.Module):
+    r""" LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with
+    shape (batch_size, height, width, channels) while channels_first corresponds to inputs
+    with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, eps=1e-6, data_format="channels_last"):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(normalized_shape))
+        self.bias = nn.Parameter(torch.zeros(normalized_shape))
+        self.eps = eps
+        self.data_format = data_format
+        if self.data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError
+        self.normalized_shape = (normalized_shape, )
+
+    def forward(self, x):
+        if self.data_format == "channels_last":
+            return F.layer_norm(x, self.normalized_shape, self.weight, self.bias, self.eps)
+        elif self.data_format == "channels_first":
+            u = x.mean(1, keepdim=True)
+            s = (x - u).pow(2).mean(1, keepdim=True)
+            x = (x - u) / torch.sqrt(s + self.eps)
+            x = self.weight[:, None, None] * x + self.bias[:, None, None]
+            return x
+
+
+def _no_grad_trunc_normal_(tensor, mean, std, a, b):
+    # Cut & paste from PyTorch official master until it's in a few official releases - RW
+    # Method based on https://people.sc.fsu.edu/~jburkardt/presentations/truncated_normal.pdf
+    def norm_cdf(x):
+        # Computes standard normal cumulative distribution function
+        return (1. + math.erf(x / math.sqrt(2.))) / 2.
+
+    if (mean < a - 2 * std) or (mean > b + 2 * std):
+        warnings.warn("mean is more than 2 std from [a, b] in nn.init.trunc_normal_. "
+                      "The distribution of values may be incorrect.",
+                      stacklevel=2)
+
+    with torch.no_grad():
+        # Values are generated by using a truncated uniform distribution and
+        # then using the inverse CDF for the normal distribution.
+        # Get upper and lower cdf values
+        l = norm_cdf((a - mean) / std)
+        u = norm_cdf((b - mean) / std)
+
+        # Uniformly fill tensor with values from [l, u], then translate to
+        # [2l-1, 2u-1].
+        tensor.uniform_(2 * l - 1, 2 * u - 1)
+
+        # Use inverse cdf transform for normal distribution to get truncated
+        # standard normal
+        tensor.erfinv_()
+
+        # Transform to proper mean, std
+        tensor.mul_(std * math.sqrt(2.))
+        tensor.add_(mean)
+
+        # Clamp to ensure it's in the proper range
+        tensor.clamp_(min=a, max=b)
+        return tensor
+
+
+def drop_path(x, drop_prob=0., training=False, scale_by_keep=True):
+    """ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+    This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
+    the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
+    See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
+    changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
+    'survival rate' as the argument.
+
+    """
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = x.new_empty(shape).bernoulli_(keep_prob)
+    if keep_prob > 0.0 and scale_by_keep:
+        random_tensor.div_(keep_prob)
+    return x * random_tensor
+
+
+class DropPath(nn.Module):
+    """ Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).
+    """
+
+    def __init__(self, drop_prob=None, scale_by_keep=True):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+        self.scale_by_keep = scale_by_keep
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
+
+
+def trunc_normal_(tensor, mean=0., std=1., a=-2., b=2.):
+    return _no_grad_trunc_normal_(tensor, mean, std, a, b)