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 import gc import huggingface_hub import os import logging # 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' self.video_limit_cpu = 100 self.video_limit_gpu = 300 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[torch.Tensor, str]: 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[torch.Tensor, str]: 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 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 h, w, scale = 0, 0, 0 # 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] bbox = face.bbox.astype(int) landmarks = face.landmark_2d_106 # Align face based on landmarks aligned_face = self.align_face(frame, landmarks) if aligned_face is not None: 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: 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, h, w, top, bottom, left, right, scale return frame, instyle, message def align_face(self, image, landmarks): # Calculate auxiliary vectors for alignment 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 = PIL.Image.fromarray(image) img = img.transform((transform_size, transform_size), PIL.Image.QUAD, (quad + 0.5).flatten(), PIL.Image.BILINEAR) if output_size < transform_size: img = img.resize((output_size, output_size), PIL.Image.ANTIALIAS) return np.array(img) # Other methods remain unchanged def detect_and_align_image(self, frame_rgb: np.ndarray, top: int, bottom: int, left: int, right: int) -> tuple: if frame_rgb is None: return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load the image.' # Convert RGB to BGR frame_bgr = cv2.cvtColor(frame_rgb, cv2.COLOR_RGB2BGR) return self.detect_and_align(frame_bgr, top, bottom, left, right) def detect_and_align_video(self, video: str, top: int, bottom: int, left: int, right: int) -> tuple: if video is None: return np.zeros((256, 256, 3), np.uint8), None, 'Error: fail to load empty file.' video_cap = cv2.VideoCapture(video) if video_cap.get(7) == 0: video_cap.release() return np.zeros((256, 256, 3), np.uint8), torch.zeros(1, 18, 512).to(self.device), 'Error: fail to load the video.' success, frame = video_cap.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) video_cap.release() return self.detect_and_align(frame, top, bottom, left, right) def detect_and_align_full_video(self, video: str, top: int, bottom: int, left: int, right: int) -> tuple: message = 'Error: no face detected! Please retry or change the video.' instyle = None if video is None: return 'default.mp4', instyle, 'Error: fail to load empty file.' video_cap = cv2.VideoCapture(video) if video_cap.get(7) == 0: video_cap.release() return 'default.mp4', instyle, 'Error: fail to load the video.' num = min(self.video_limit_gpu, int(video_cap.get(7))) if self.device == 'cpu': num = min(self.video_limit_cpu, num) success, frame = video_cap.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frame, instyle, message, w, h, top, bottom, left, right, scale = self.detect_and_align(frame, top, bottom, left, right, True) if instyle is None: return 'default.mp4', instyle, message fourcc = cv2.VideoWriter_fourcc(*'mp4v') videoWriter = cv2.VideoWriter('input.mp4', fourcc, video_cap.get(5), (int(right-left), int(bottom-top))) videoWriter.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) kernel_1d = np.array([[0.125], [0.375], [0.375], [0.125]]) for i in range(num-1): success, frame = video_cap.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) if scale <= 0.75: frame = cv2.sepFilter2D(frame, -1, kernel_1d, kernel_1d) if scale <= 0.375: frame = cv2.sepFilter2D(frame, -1, kernel_1d, kernel_1d) frame = cv2.resize(frame, (w, h))[top:bottom, left:right] videoWriter.write(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)) videoWriter.release() video_cap.release() return 'input.mp4', instyle, 'Successfully rescale the video to (%d, %d)' % (bottom-top, right-left) 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: return np.zeros((256, 256, 3), np.uint8), 'Opps, 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: return np.zeros((256, 256, 3), np.uint8), 'Opps, something wrong with the style type. Please go to Step 1 and load model again.' with torch.no_grad(): if self.color_transfer: s_w = exstyle else: s_w = instyle.clone() s_w[:, :7] = exstyle[:, :7] x = self.transform(aligned_face).unsqueeze(dim=0).to(self.device) 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) 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) def video_toonify(self, aligned_video: str, instyle: torch.Tensor, exstyle: torch.Tensor, style_degree: float, style_type: str) -> tuple: if aligned_video is None: return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.' video_cap = cv2.VideoCapture(aligned_video) if instyle is None or aligned_video is None or video_cap.get(7) == 0: video_cap.release() return 'default.mp4', 'Opps, something wrong with the input. Please go to Step 2 and Rescale Video again.' if self.style_name != style_type: exstyle, _ = self.load_model(style_type) if exstyle is None: return 'default.mp4', 'Opps, something wrong with the style type. Please go to Step 1 and load model again.' num = min(self.video_limit_gpu, int(video_cap.get(7))) if self.device == 'cpu': num = min(self.video_limit_cpu, num) fourcc = cv2.VideoWriter_fourcc(*'mp4v') videoWriter = cv2.VideoWriter('output.mp4', fourcc, video_cap.get(5), (int(video_cap.get(3)*4), int(video_cap.get(4)*4))) batch_frames = [] if video_cap.get(3) != 0: if self.device == 'cpu': batch_size = max(1, int(4 * 256 * 256 / video_cap.get(3) / video_cap.get(4))) else: batch_size = min(max(1, int(4 * 400 * 360 / video_cap.get(3) / video_cap.get(4))), 4) else: batch_size = 1 print('*** Toonify using batch size of %d on %dx%d video of %d frames with style of %s' % (batch_size, int(video_cap.get(3)*4), int(video_cap.get(4)*4), num, style_type)) with torch.no_grad(): if self.color_transfer: s_w = exstyle else: s_w = instyle.clone() s_w[:, :7] = exstyle[:, :7] for i in range(num): success, frame = video_cap.read() frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) batch_frames += [self.transform(frame).unsqueeze(dim=0).to(self.device)] if len(batch_frames) == batch_size or (i+1) == num: x = torch.cat(batch_frames, dim=0) batch_frames = [] with torch.no_grad(): 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), style_degree) y_tilde = torch.clamp(y_tilde, -1, 1) for k in range(y_tilde.size(0)): videoWriter.write(tensor2cv2(y_tilde[k].cpu())) gc.collect() videoWriter.release() video_cap.release() return 'output.mp4', 'Successfully toonify video of %d frames with style of %s' % (num, self.style_name) def tensor2cv2(self, img): """Convert a tensor image to OpenCV format.""" tmp = ((img.cpu().numpy().transpose(1, 2, 0) + 1.0) * 127.5).astype(np.uint8).copy() logging.debug(f"Converted image shape: {tmp.shape}, strides: {tmp.strides}") return cv2.cvtColor(tmp, cv2.COLOR_RGB2BGR)