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
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))
# 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()
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.
# 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)}")