vtoonify_model.py

from __future__ import annotations
import gradio as gr
import pathlib
import sys
sys.path.insert(0, 'vtoonify')

from util import load_psp_standalone, get_video_crop_parameter, tensor2cv2
import torch
import torch.nn as nn
import numpy as np
import insightface
import cv2
from model.vtoonify import VToonify
from model.bisenet.model import BiSeNet
import torch.nn.functional as F
from torchvision import transforms
from model.encoder.align_all_parallel import align_face
import gc
import huggingface_hub
import os
import logging
from PIL import Image




# Configure logging
logging.basicConfig(level=logging.INFO)

MODEL_REPO = 'PKUWilliamYang/VToonify'

class Model():
    def __init__(self, device):
        super().__init__()

        self.device = device
        self.style_types = {
            'cartoon1': ['vtoonify_d_cartoon/vtoonify_s026_d0.5.pt', 26],
            'cartoon1-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 26],
            'cartoon2-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 64],
            'cartoon3-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 153],
            'cartoon4': ['vtoonify_d_cartoon/vtoonify_s299_d0.5.pt', 299],
            'cartoon4-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 299],
            'cartoon5-d': ['vtoonify_d_cartoon/vtoonify_s_d.pt', 8],
            'comic1-d': ['vtoonify_d_comic/vtoonify_s_d.pt', 28],
            'comic2-d': ['vtoonify_d_comic/vtoonify_s_d.pt', 18],
            'arcane1': ['vtoonify_d_arcane/vtoonify_s000_d0.5.pt', 0],
            'arcane1-d': ['vtoonify_d_arcane/vtoonify_s_d.pt', 0],
            'arcane2': ['vtoonify_d_arcane/vtoonify_s077_d0.5.pt', 77],
            'arcane2-d': ['vtoonify_d_arcane/vtoonify_s_d.pt', 77],
            'caricature1': ['vtoonify_d_caricature/vtoonify_s039_d0.5.pt', 39],
            'caricature2': ['vtoonify_d_caricature/vtoonify_s068_d0.5.pt', 68],
            'pixar': ['vtoonify_d_pixar/vtoonify_s052_d0.5.pt', 52],
            'pixar-d': ['vtoonify_d_pixar/vtoonify_s_d.pt', 52],
            'illustration1-d': ['vtoonify_d_illustration/vtoonify_s054_d_c.pt', 54],
            'illustration2-d': ['vtoonify_d_illustration/vtoonify_s004_d_c.pt', 4],
            'illustration3-d': ['vtoonify_d_illustration/vtoonify_s009_d_c.pt', 9],
            'illustration4-d': ['vtoonify_d_illustration/vtoonify_s043_d_c.pt', 43],
            'illustration5-d': ['vtoonify_d_illustration/vtoonify_s086_d_c.pt', 86],
        }

        self.face_detector = self._create_insightface_detector()
        self.parsingpredictor = self._create_parsing_model()
        self.pspencoder = self._load_encoder()
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
        ])

        self.vtoonify, self.exstyle = self._load_default_model()
        self.color_transfer = False
        self.style_name = 'cartoon1'

    def _create_insightface_detector(self):
        # Initialize InsightFace
        app = insightface.app.FaceAnalysis()
        app.prepare(ctx_id=0 if self.device == 'cuda' else -1, det_size=(640, 640))
        return app

    def _create_parsing_model(self):
        parsingpredictor = BiSeNet(n_classes=19)
        parsingpredictor.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/faceparsing.pth'),
                                                    map_location=lambda storage, loc: storage))
        parsingpredictor.to(self.device).eval()
        return parsingpredictor

    def _load_encoder(self) -> nn.Module:
        style_encoder_path = huggingface_hub.hf_hub_download(MODEL_REPO, 'models/encoder.pt')
        return load_psp_standalone(style_encoder_path, self.device)

    def _load_default_model(self) -> tuple:
        vtoonify = VToonify(backbone='dualstylegan')
        vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO,
                                            'models/vtoonify_d_cartoon/vtoonify_s026_d0.5.pt'), 
                                            map_location=lambda storage, loc: storage)['g_ema'])
        vtoonify.to(self.device)
        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/vtoonify_d_cartoon/exstyle_code.npy'), allow_pickle=True).item()
        exstyle = torch.tensor(tmp[list(tmp.keys())[26]]).to(self.device)
        with torch.no_grad():  
            exstyle = vtoonify.zplus2wplus(exstyle)
        return vtoonify, exstyle

    def load_model(self, style_type: str) -> tuple:
        if 'illustration' in style_type:
            self.color_transfer = True
        else:
            self.color_transfer = False
        if style_type not in self.style_types.keys():
            return None, 'Oops, wrong Style Type. Please select a valid model.'
        self.style_name = style_type
        model_path, ind = self.style_types[style_type]
        style_path = os.path.join('models', os.path.dirname(model_path), 'exstyle_code.npy')
        self.vtoonify.load_state_dict(torch.load(huggingface_hub.hf_hub_download(MODEL_REPO, 'models/' + model_path), 
                                            map_location=lambda storage, loc: storage)['g_ema'])
        tmp = np.load(huggingface_hub.hf_hub_download(MODEL_REPO, style_path), allow_pickle=True).item()
        exstyle = torch.tensor(tmp[list(tmp.keys())[ind]]).to(self.device)
        with torch.no_grad():  
            exstyle = self.vtoonify.zplus2wplus(exstyle)
        return exstyle, 'Model of %s loaded.' % (style_type)

    def convert_106_to_68(self, landmarks_106):
        # Mapping from 106 landmarks to 68 landmarks
        landmark106to68 = [
            1, 10, 12, 14, 16, 3, 5, 7, 0, 23, 21, 19, 32, 30, 28, 26, 17,  # Face outline
            43, 48, 49, 51, 50,  # Left eyebrow
            102, 103, 104, 105, 101,  # Right eyebrow
            72, 73, 74, 86, 78, 79, 80, 85, 84,  # Nose
            35, 41, 42, 39, 37, 36,  # Left eye
            89, 95, 96, 93, 91, 90,  # Right eye
            52, 64, 63, 71, 67, 68, 61, 58, 59, 53, 56, 55, 65, 66, 62, 70, 69, 57, 60, 54  # Mouth
        ]
        
        # Convert 106 landmarks to 68 landmarks
        landmarks_68 = [landmarks_106[index] for index in landmark106to68]
        
        return landmarks_68

    def detect_and_align_image(self, image: str, top: int, bottom: int, left: int, right: int
                              ) -> tuple[np.ndarray, torch.Tensor, str]:
        if image is None:
            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load empty file.'
        frame = cv2.imread(image)
        if frame is None:
            return np.zeros((256,256,3), np.uint8), None, 'Error: fail to load the image.'       
        frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
        return self.detect_and_align(frame, top, bottom, left, right)
    def detect_and_align(self, frame, top, bottom, left, right, return_para=False):
        message = 'Error: no face detected! Please retry or change the photo.'
        instyle = None

        # Use InsightFace for face detection
        faces = self.face_detector.get(frame)
        if len(faces) > 0:
            logging.info(f"Detected {len(faces)} face(s).")
            face = faces[0]
            landmarks_106 = face.landmark_2d_106
            landmarks_68 = self.convert_106_to_68(landmarks_106)

            # Align face based on mapped landmarks
            aligned_face = self.align_face(frame, landmarks_68)
            if aligned_face is not None:
                logging.info(f"Aligned face shape: {aligned_face.shape}")
                with torch.no_grad():
                    I = self.transform(aligned_face).unsqueeze(dim=0).to(self.device)
                    instyle = self.pspencoder(I)
                    instyle = self.vtoonify.zplus2wplus(instyle)
                    message = 'Successfully aligned the face.'
            else:
                logging.warning("Failed to align face.")
                frame = np.zeros((256, 256, 3), np.uint8)
        else:
            logging.warning("No face detected.")
            frame = np.zeros((256, 256, 3), np.uint8)

        if return_para:
            return frame, instyle, message
        return frame, instyle, message

    def align_face(self, image, landmarks):
        # Example alignment logic using 68 landmarks
        eye_left = np.mean(landmarks[36:42], axis=0)
        eye_right = np.mean(landmarks[42:48], axis=0)
        mouth_left = landmarks[48]
        mouth_right = landmarks[54]

        # Calculate transformation parameters
        eye_center = (eye_left + eye_right) / 2
        mouth_center = (mouth_left + mouth_right) / 2
        eye_to_eye = eye_right - eye_left
        eye_to_mouth = mouth_center - eye_center

        # Define the transformation matrix
        x = eye_to_eye - np.flipud(eye_to_mouth) * [-1, 1]
        x /= np.hypot(*x)
        x *= np.hypot(*eye_to_eye) * 2.0
        y = np.flipud(x) * [-1, 1]
        c = eye_center + eye_to_mouth * 0.1
        quad = np.stack([c - x - y, c - x + y, c + x + y, c + x - y])
        qsize = np.hypot(*x) * 2

        # Transform and crop the image
        transform_size = 256
        output_size = 256
        img = Image.fromarray(image)
        img = img.transform((transform_size, transform_size), Image.QUAD, (quad + 0.5).flatten(), Image.BILINEAR)
        if output_size < transform_size:
            img = img.resize((output_size, output_size), Image.ANTIALIAS)

        return np.array(img)

    
    def image_toonify(self, aligned_face: np.ndarray, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple:
        if instyle is None or aligned_face is None:
            logging.error("Invalid input: instyle or aligned_face is None.")
            return np.zeros((256, 256, 3), np.uint8), 'Oops, something wrong with the input. Please go to Step 2 and Rescale Image/First Frame again.'
    
        if self.style_name != style_type:
            exstyle, _ = self.load_model(style_type)
        if exstyle is None:
            logging.error("Failed to load style model.")
            return np.zeros((256, 256, 3), np.uint8), 'Oops, something wrong with the style type. Please go to Step 1 and load model again.'
    
        try:
            with torch.no_grad():
               if self.color_transfer:
                   s_w = exstyle
               else:
                   s_w = instyle.clone()
                   s_w[:, :7] = exstyle[:, :7]

            # Ensure the input is resized to 256x256
               aligned_face_resized = cv2.resize(aligned_face, (256, 256))
               x = self.transform(aligned_face_resized).unsqueeze(dim=0).to(self.device)
               logging.info(f"Input to VToonify shape: {x.shape}")
               x_p = F.interpolate(self.parsingpredictor(2 * (F.interpolate(x, scale_factor=2, mode='bilinear', align_corners=False)))[0], 
                                scale_factor=0.5, recompute_scale_factor=False).detach()
               inputs = torch.cat((x, x_p / 16.), dim=1)
               y_tilde = self.vtoonify(inputs, s_w.repeat(inputs.size(0), 1, 1), d_s=style_degree)        
               y_tilde = torch.clamp(y_tilde, -1, 1)
               logging.info(f"Output from VToonify shape: {y_tilde.shape}")
               print('*** Toonify %dx%d image with style of %s' % (y_tilde.shape[2], y_tilde.shape[3], style_type))
        
            return ((y_tilde[0].cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8), 'Successfully toonify the image with style of %s'%(self.style_name)
    
        except Exception as e:
            logging.error(f"Error during model execution: {e}")
            return np.zeros((256, 256, 3), np.uint8), f"Error during processing: {str(e)}"