Upload 11 files

.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

@@ -1,35 +1,13 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
+*.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
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* 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,33 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+**/__pycache__/
+*.py[cod]
+**/*.py[cod]
+*$py.class
+
+miniserver.py
+
+# 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,60 @@
+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/
+
+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,161 @@
+---
+title: FacePoke
+emoji: 🙂‍↔️👈
+colorFrom: yellow
+colorTo: red
+sdk: docker
+pinned: true
+license: mit
+header: mini
+app_file: app.py
+app_port: 8080
+disable_embedding: true
+short_description: Import a portrait, click to move the head!
+---
+
+# FacePoke
+
+## 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)
+
+# Funding
+
+FacePoke is just a tiny project!
+
+There are no plans to create a proprietary/expensive cloud-only blackbox SaaS for FacePoke, so if you like it, you can buy me a coffee 🫶
+
+<a href="https://www.buymeacoffee.com/flngr" target="_blank"><img src="https://www.buymeacoffee.com/assets/img/custom_images/orange_img.png" alt="Buy Me A Coffee" style="height: 41px !important;width: 174px !important;box-shadow: 0px 3px 2px 0px rgba(190, 190, 190, 0.5) !important;-webkit-box-shadow: 0px 3px 2px 0px rgba(190, 190, 190, 0.5) !important;" ></a>
+
+## 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
+
+### Before you install
+
+FacePoke has only been tested in a Linux environment, using `Python 3.10` and `CUDA 12.4` (so a NVIDIA GPU).
+
+Contributions are welcome to help supporting other platforms!
+
+### Local Setup
+
+1. Make sure you have Git and Git LFS installed globally (https://git-lfs.com):
+
+   ```bash
+   git lfs install
+   ```
+
+2. Clone the repository:
+   ```bash
+   git clone https://github.com/jbilcke-hf/FacePoke.git
+   cd FacePoke
+   ```
+
+3. Install Python dependencies:
+
+   Using a virtual environment (Python venv) is strongly recommended.
+
+   FacePoke has been tested with `Python 3.10`.
+
+   ```bash
+   pip3 install --upgrade -r requirements.txt
+   ```
+
+4. Install frontend dependencies:
+   ```bash
+   cd client
+   bun install
+   ```
+
+5. Build the frontend:
+   ```bash
+   bun build ./src/index.tsx --outdir ../public/
+   ```
+
+6. Start the backend server:
+   ```bash
+   python app.py
+   ```
+
+7. 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.
+
+
+Note: by default Hugging Face runs the `main` branch, so if you want to push a feature branch you need to do this:
+
+```bash
+git push <space_repo> <feature_branch>:main -f
+```
+
+## 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.
+
+### Increasing the framerate
+
+I am testing various things to increase the framerate.
+
+One project is to only transmit the modified head, instead of the whole image.
+
+Another one is to automatically adapt to the server and network speed.
+
+## 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.
+
+Please note that while the code of LivePortrait and Insightface are open-source with "no limitation for both academic and commercial usage", the model weights trained from Insightface data are available for [non-commercial research purposes only](https://github.com/deepinsight/insightface?tab=readme-ov-file#license).
+
+---
+
+Developed with ❤️ by Julian Bilcke at Hugging Face

app.py

@@ -0,0 +1,145 @@
+"""
+FacePoke API
+
+Author: Julian Bilcke
+Date: September 30, 2024
+"""
+
+import sys
+import asyncio
+from aiohttp import web, WSMsgType
+import json
+from json import JSONEncoder
+import numpy as np
+import uuid
+import logging
+import os
+import signal
+from typing import Dict, Any, List, Optional
+import base64
+import io
+
+from PIL import Image
+
+# by popular demand, let's add support for avif
+import pillow_avif
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Set asyncio logger to DEBUG level
+#logging.getLogger("asyncio").setLevel(logging.INFO)
+
+#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
+
+# Global constants
+DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
+MODELS_DIR = os.path.join(DATA_ROOT, "models")
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.integer):
+            return int(obj)
+        elif isinstance(obj, np.floating):
+            return float(obj)
+        elif isinstance(obj, np.ndarray):
+            return obj.tolist()
+        else:
+            return super(NumpyEncoder, self).default(obj)
+
+async def websocket_handler(request: web.Request) -> web.WebSocketResponse:
+    ws = web.WebSocketResponse()
+    await ws.prepare(request)
+    engine = request.app['engine']
+    try:
+        #logger.info("New WebSocket connection established")
+        while True:
+            msg = await ws.receive()
+
+            if msg.type in (WSMsgType.CLOSE, WSMsgType.ERROR):
+                #logger.warning(f"WebSocket connection closed: {msg.type}")
+                break
+
+            try:
+                if msg.type == WSMsgType.BINARY:
+                    res = await engine.load_image(msg.data)
+                    json_res = json.dumps(res, cls=NumpyEncoder)
+                    await ws.send_str(json_res)
+
+                elif msg.type == WSMsgType.TEXT:
+                    data = json.loads(msg.data)
+                    webp_bytes = await engine.transform_image(data.get('uuid'), data.get('params'))
+                    await ws.send_bytes(webp_bytes)
+
+            except Exception as e:
+                logger.error(f"Error in engine: {str(e)}")
+                logger.exception("Full traceback:")
+                await ws.send_json({"error": str(e)})
+
+    except Exception as e:
+        logger.error(f"Error in websocket_handler: {str(e)}")
+        logger.exception("Full traceback:")
+    return ws
+
+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 initialize_app() -> web.Application:
+    """Initialize and configure the web application."""
+    try:
+        logger.info("Initializing application...")
+        live_portrait = await initialize_models()
+
+        logger.info("🚀 Creating Engine instance...")
+        engine = Engine(live_portrait=live_portrait)
+        logger.info("✅ Engine instance created.")
+
+        app = web.Application()
+        app['engine'] = engine
+
+        # 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
+
+if __name__ == "__main__":
+    try:
+        logger.info("Starting FacePoke application")
+        app = asyncio.run(initialize_app())
+        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:")

build.sh

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

engine.py

@@ -0,0 +1,228 @@
+import uuid
+import logging
+import hashlib
+import os
+import io
+import asyncio
+from async_lru import alru_cache
+import base64
+from queue import Queue
+from typing import Dict, Any, List, Optional, Union
+from functools import lru_cache
+import numpy as np
+import torch
+import torch.nn.functional as F
+from PIL import Image, ImageOps
+
+from liveportrait.config.argument_config import ArgumentConfig
+from liveportrait.utils.camera import get_rotation_matrix
+from liveportrait.utils.io import resize_to_limit
+from liveportrait.utils.crop import prepare_paste_back, paste_back, parse_bbox_from_landmark
+
+# 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")
+
+        self.processed_cache = {}  # Stores the processed image data
+
+        logger.info("✅ FacePoke Engine initialized successfully.")
+
+    @alru_cache(maxsize=512)
+    async def load_image(self, data):
+        image = Image.open(io.BytesIO(data))
+
+        # keep the exif orientation (fix the selfie issue on iphone)
+        image = ImageOps.exif_transpose(image)
+
+        # Convert the image to RGB mode (removes alpha channel if present)
+        image = image.convert('RGB')
+
+        uid = str(uuid.uuid4())
+        img_rgb = np.array(image)
+
+        inference_cfg = self.live_portrait.live_portrait_wrapper.cfg
+        img_rgb = await asyncio.to_thread(resize_to_limit, img_rgb, inference_cfg.ref_max_shape, inference_cfg.ref_shape_n)
+        crop_info = await asyncio.to_thread(self.live_portrait.cropper.crop_single_image, img_rgb)
+        img_crop_256x256 = crop_info['img_crop_256x256']
+
+        I_s = await asyncio.to_thread(self.live_portrait.live_portrait_wrapper.prepare_source, img_crop_256x256)
+        x_s_info = await asyncio.to_thread(self.live_portrait.live_portrait_wrapper.get_kp_info, I_s)
+        f_s = await asyncio.to_thread(self.live_portrait.live_portrait_wrapper.extract_feature_3d, I_s)
+        x_s = await asyncio.to_thread(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[uid] = processed_data
+
+        # Calculate the bounding box
+        bbox_info = parse_bbox_from_landmark(processed_data['crop_info']['lmk_crop'], scale=1.0)
+
+        return {
+            'u': uid,
+
+            # those aren't easy to serialize
+            'c': bbox_info['center'], # 2x1
+            's': bbox_info['size'], # scalar
+            'b': bbox_info['bbox'],  # 4x2
+            'a': bbox_info['angle'],  # rad, counterclockwise
+            # 'bbox_rot': bbox_info['bbox_rot'].toList(),  # 4x2
+        }
+
+    async def transform_image(self, uid: str, params: Dict[str, float]) -> bytes:
+        # If we don't have the image in cache yet, add it
+        if uid not in self.processed_cache:
+            raise ValueError("cache miss")
+
+        processed_data = self.processed_cache[uid]
+
+        try:
+            # Apply modifications based on params
+            x_d_new = processed_data['x_s_info']['kp'].clone()
+
+            # Adapted from https://github.com/PowerHouseMan/ComfyUI-AdvancedLivePortrait/blob/main/nodes.py#L408-L472
+            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)
+                ]),
+                # Some other ones: https://github.com/jbilcke-hf/FacePoke/issues/22#issuecomment-2408708028
+                # Still need to check how exactly we would control those in the UI,
+                # as we don't have yet segmentation in the frontend UI for those body parts
+                #('lower_lip', [
+                #    (0, 19, 1, 0.02)
+                #]),
+                #('upper_lip', [
+                #    (0, 20, 1, -0.01)
+                #]),
+                #('neck', [(0, 5, 1, 0.01)]),
+            ]
+
+            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.
+
+
+            # 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 = await asyncio.to_thread(self.live_portrait.live_portrait_wrapper.parse_output, out['out'])
+
+            buffered = io.BytesIO()
+
+            ####################################################
+            # this part is about stitching the image back into the original.
+            #
+            # this is an expensive operation, not just because of the compute
+            # but because the payload will also be bigger (we send back the whole pic)
+            #
+            # I'm currently running some experiments to do it in the frontend
+            #
+            #  --- old way: we do it in the server-side: ---
+            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[0], processed_data['crop_info']['M_c2o'], processed_data['img_rgb'], mask_ori
+            )
+            result_image = Image.fromarray(I_p_to_ori_blend)
+
+            # --- maybe future way: do it in the frontend: ---
+            #result_image = Image.fromarray(I_p[0])
+            ####################################################
+
+            # write it into a webp
+            result_image.save(buffered, format="WebP", quality=82, lossless=False, method=6)
+
+            return buffered.getvalue()
+
+        except Exception as e:
+            raise ValueError(f"Failed to modify image: {str(e)}")

loader.py

@@ -0,0 +1,165 @@
+import os
+import logging
+import torch
+import asyncio
+import aiohttp
+import requests
+from huggingface_hub import hf_hub_download
+
+# Configure logging
+logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
+# Configuration
+DATA_ROOT = os.environ.get('DATA_ROOT', '/tmp/data')
+MODELS_DIR = os.path.join(DATA_ROOT, "models")
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+# Hugging Face repository information
+HF_REPO_ID = "jbilcke-hf/model-cocktail"
+
+# Model files to download
+MODEL_FILES = [
+    "dwpose/dw-ll_ucoco_384.pth",
+    "face-detector/s3fd-619a316812.pth",
+
+    "liveportrait/spade_generator.pth",
+    "liveportrait/warping_module.pth",
+    "liveportrait/motion_extractor.pth",
+    "liveportrait/stitching_retargeting_module.pth",
+    "liveportrait/appearance_feature_extractor.pth",
+    "liveportrait/landmark.onnx",
+
+    # For animal mode 🐶🐱
+    # however they say animal mode doesn't support stitching yet?
+    # https://github.com/KwaiVGI/LivePortrait/blob/main/assets/docs/changelog/2024-08-02.md#updates-on-animals-mode
+    #"liveportrait-animals/warping_module.pth",
+    #"liveportrait-animals/spade_generator.pth",
+    #"liveportrait-animals/motion_extractor.pth",
+    #"liveportrait-animals/appearance_feature_extractor.pth",
+    #"liveportrait-animals/stitching_retargeting_module.pth",
+    #"liveportrait-animals/xpose.pth",
+
+    # this is a hack, instead we should probably try to
+    # fix liveportrait/utils/dependencies/insightface/utils/storage.py
+    "insightface/models/buffalo_l.zip",
+
+    "insightface/buffalo_l/det_10g.onnx",
+    "insightface/buffalo_l/2d106det.onnx",
+    "sd-vae-ft-mse/diffusion_pytorch_model.bin",
+    "sd-vae-ft-mse/diffusion_pytorch_model.safetensors",
+    "sd-vae-ft-mse/config.json",
+
+    # we don't use those yet
+    #"flux-dev/flux-dev-fp8.safetensors",
+    #"flux-dev/flux_dev_quantization_map.json",
+    #"pulid-flux/pulid_flux_v0.9.0.safetensors",
+    #"pulid-flux/pulid_v1.bin"
+]
+
+def create_directory(directory):
+    """Create a directory if it doesn't exist and log its status."""
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+        logger.info(f"  Directory created: {directory}")
+    else:
+        logger.info(f"  Directory already exists: {directory}")
+
+def print_directory_structure(startpath):
+    """Print the directory structure starting from the given path."""
+    for root, dirs, files in os.walk(startpath):
+        level = root.replace(startpath, '').count(os.sep)
+        indent = ' ' * 4 * level
+        logger.info(f"{indent}{os.path.basename(root)}/")
+        subindent = ' ' * 4 * (level + 1)
+        for f in files:
+            logger.info(f"{subindent}{f}")
+
+async def download_hf_file(filename: str) -> None:
+    """Download a file from Hugging Face to the models directory."""
+    dest = os.path.join(MODELS_DIR, filename)
+    os.makedirs(os.path.dirname(dest), exist_ok=True)
+    if os.path.exists(dest):
+        # this is really for debugging purposes only
+        logger.debug(f"    ✅ {filename}")
+        return
+
+    logger.info(f"    ⏳ Downloading {HF_REPO_ID}/{filename}")
+
+    try:
+        await asyncio.get_event_loop().run_in_executor(
+            None,
+            lambda: hf_hub_download(
+                repo_id=HF_REPO_ID,
+                filename=filename,
+                local_dir=MODELS_DIR
+            )
+        )
+        logger.info(f"    ✅ Downloaded {filename}")
+    except Exception as e:
+        logger.error(f"🚨 Error downloading file from Hugging Face: {e}")
+        if os.path.exists(dest):
+            os.remove(dest)
+        raise
+
+async def download_all_models():
+    """Download all required models from the Hugging Face repository."""
+    logger.info("  🔎 Looking for models...")
+    tasks = [download_hf_file(filename) for filename in MODEL_FILES]
+    await asyncio.gather(*tasks)
+    logger.info("  ✅ All models are available")
+
+    # are you looking to debug the app and verify that models are downloaded properly?
+    # then un-comment the two following lines:
+    #logger.info("💡 Printing directory structure of models:")
+    #print_directory_structure(MODELS_DIR)
+
+class ModelLoader:
+    """A class responsible for loading and initializing all required models."""
+
+    def __init__(self):
+        self.device = DEVICE
+        self.models_dir = MODELS_DIR
+
+    async def load_live_portrait(self):
+        """Load LivePortrait models."""
+        from liveportrait.config.inference_config import InferenceConfig
+        from liveportrait.config.crop_config import CropConfig
+        from liveportrait.live_portrait_pipeline import LivePortraitPipeline
+
+        logger.info("  ⏳ Loading LivePortrait models...")
+        live_portrait_pipeline = await asyncio.to_thread(
+            LivePortraitPipeline,
+            inference_cfg=InferenceConfig(
+                # default values
+                flag_stitching=True,  # we recommend setting it to True!
+                flag_relative=True,  # whether to use relative motion
+                flag_pasteback=True,  # whether to paste-back/stitch the animated face cropping from the face-cropping space to the original image space
+                flag_do_crop= True,  # whether to crop the source portrait to the face-cropping space
+                flag_do_rot=True,  # whether to conduct the rotation when flag_do_crop is True
+            ),
+            crop_cfg=CropConfig()
+        )
+        logger.info("  ✅ LivePortrait models loaded successfully.")
+        return live_portrait_pipeline
+
+async def initialize_models():
+    """Initialize and load all required models."""
+    logger.info("🚀 Starting model initialization...")
+
+    # Ensure all required models are downloaded
+    await download_all_models()
+
+    # Initialize the ModelLoader
+    loader = ModelLoader()
+
+    # Load LivePortrait models
+    live_portrait = await loader.load_live_portrait()
+
+    logger.info("✅ Model initialization completed.")
+    return live_portrait
+
+# Initial setup
+logger.info("🚀 Setting up storage directories...")
+create_directory(MODELS_DIR)
+logger.info("✅ Storage directories setup completed.")

requirements.txt

@@ -0,0 +1,83 @@
+
+# --------------------------------------------------------------------
+# Cuda 12.4
+# --------------------------------------------------------------------
+
+--extra-index-url https://download.pytorch.org/whl/cu124
+torch
+torchvision
+torchaudio
+torchgeometry
+
+# --------------------------------------------------------------------
+# Common libraries for LivePortrait and all
+# --------------------------------------------------------------------
+
+# LRU cache compatible with asyncio
+async-lru==2.0.4
+
+# note: gradio is only used for the cropping utility
+gradio==5.6.0
+
+pyyaml==6.0.1
+numpy==1.22.0 # Updated to resolve conflicts
+opencv-python==4.8.1.78 # Downgraded to be compatible with numpy
+scipy==1.10.1
+imageio==2.31.1
+imageio-ffmpeg==0.5.1
+lmdb==1.4.1
+tqdm==4.66.4
+rich==13.7.1
+
+# this one is from 5 years ago, we should probably
+# use python-ffmpeg or typed-ffmpeg instead
+# https://pypi.org/search/?q=ffmpeg-python
+ffmpeg-python==0.2.0
+
+# versions of onnx gpu <= 1.17 did not support Cuda 12
+onnxruntime-gpu==1.19.2
+
+# onnx 1.16.2 has some issues it seems:
+# https://github.com/onnx/onnx/issues/6267
+#onnx==1.16.2
+# update: I've rolled it back to 1.16.1
+# see: https://github.com/jbilcke-hf/FacePoke/issues/23#issuecomment-2414714490
+onnx==1.16.1
+
+scikit-image==0.20.0
+albumentations==1.3.1
+matplotlib==3.7.2
+tyro==0.8.5
+chumpy==0.70
+
+diffusers==0.30.3
+accelerate==0.34.2
+tensorflow==2.12.0
+tensorboard==2.12.0
+transformers==4.39.2
+
+gdown==5.2.0
+requests==2.32.3
+omegaconf==2.3.0
+
+pydantic==2.9.2
+
+# --------------------------------------------------------------------
+# RESERVED FOR FUTURE USAGE
+#
+# (it adds bloat, so you can remove them if you want)
+# --------------------------------------------------------------------
+aiohttp==3.10.5
+av==12.3.0
+einops==0.7.0
+safetensors==0.4.5
+huggingface-hub==0.25.1
+optimum-quanto==0.2.4
+
+# --------------------------------------------------------------------
+# Used for advanced LivePortrait features
+# --------------------------------------------------------------------
+pillow==10.4.0
+
+# by popular demand, let's add AVIF
+pillow-avif-plugin==1.4.6