diff --git a/.gitignore b/.gitignore index 25a244d..64d75dc 100644 --- a/.gitignore +++ b/.gitignore @@ -144,4 +144,7 @@ data/ .idea/ *.pt -.serena \ No newline at end of file +.serena +onnx_model_tha4/fp16 +onnx_model_tha4/fp32 +onnx_model_tha4/tmp \ No newline at end of file diff --git a/args.py b/args.py index 2c2d6eb..28032b1 100644 --- a/args.py +++ b/args.py @@ -37,6 +37,7 @@ parser.add_argument('--simplify', type=int, default=1) parser.add_argument('--device_id', type=int, default=0) parser.add_argument('--use_tensorrt', action='store_true') +parser.add_argument('--use_tha4', action='store_true') parser.add_argument('--model_seperable', action='store_true') parser.add_argument('--model_half', action='store_true') diff --git a/convert_onnx_to_trt_fp16.py b/convert_onnx_to_trt_fp16.py new file mode 100644 index 0000000..3825543 --- /dev/null +++ b/convert_onnx_to_trt_fp16.py @@ -0,0 +1,51 @@ +""" +将 onnx_model_tha4/fp32 目录中的 ONNX 模型转换为 fp16 精度的 TensorRT 模型 +并保存到 onnx_model_tha4/fp16 目录 +""" +import os +import sys +from pathlib import Path + +# 添加 ezvtuber-rt 到路径 +sys.path.insert(0, str(Path(__file__).parent / 'ezvtuber-rt')) + +from ezvtb_rt.trt_utils import build_engine, save_engine + +def convert_onnx_to_trt(): + # 输入和输出目录 + input_dir = Path('onnx_model_tha4/fp32') + output_dir = Path('onnx_model_tha4/fp16') + + # 创建输出目录 + output_dir.mkdir(parents=True, exist_ok=True) + + # 获取所有 ONNX 文件 + onnx_files = list(input_dir.glob('*.onnx')) + + if not onnx_files: + print(f"在 {input_dir} 中未找到 ONNX 模型文件") + return + + print(f"找到 {len(onnx_files)} 个 ONNX 模型文件") + print("开始转换为 fp16 精度的 TensorRT 模型...\n") + + for onnx_file in onnx_files: + print(f"正在转换: {onnx_file.name}") + + # 构建输出文件路径 + trt_filename = onnx_file.stem + '.trt' + trt_path = output_dir / trt_filename + + try: + # 构建并保存 TensorRT 引擎 + engine = build_engine(str(onnx_file), 'fp16') + save_engine(engine, str(trt_path)) + print(f"✓ 成功转换: {onnx_file.name} -> {trt_path}\n") + except Exception as e: + print(f"✗ 转换失败: {onnx_file.name}") + print(f" 错误信息: {str(e)}\n") + + print("转换完成！") + +if __name__ == '__main__': + convert_onnx_to_trt() diff --git a/ezvtb_rt_interface.py b/ezvtb_rt_interface.py index 5272a69..4c0b537 100644 --- a/ezvtb_rt_interface.py +++ b/ezvtb_rt_interface.py @@ -37,72 +37,120 @@ def get_core( use_sr:bool = False, sr_x4:bool = True, sr_half:bool = True, - sr_noise:int = 1 + sr_noise:int = 1, + #THA4 model setting (NEW) + use_tha4:bool = False # Use THA4 instead of THA3 ): support_trt = False - if use_tensorrt: #Verify tensorrt + if use_tensorrt: print('Verifying TensorRT') try: - from ezvtb_rt.trt_utils import cudaSetDevice, check_build_all_models + from ezvtb_rt.trt_utils import cudaSetDevice cudaSetDevice(device_id) - # check_build_all_models() support_trt = True - except: + except Exception: support_trt = False - if support_trt == False and use_tensorrt == True: + if not support_trt and use_tensorrt: print('TensorRT option selected but not supported') use_tensorrt = False if use_tensorrt: os.environ['CUDA_DEVICE'] = str(device_id) import pycuda.autoinit - print(f'Using devcie {pycuda.autoinit.device.name()}') + print(f'Using device {pycuda.autoinit.device.name()}') + + # Determine model directory based on THA version + if use_tha4: + # THA4: Use onnx_model_tha4 directory + if use_tensorrt: + tha_model_dir = os.path.join('.', 'onnx_model_tha4', 'fp16') + else: + tha_model_dir = os.path.join('.', 'onnx_model_tha4', 'fp32') + print(f'THA4 Mode - Model Path: {tha_model_dir}') + else: + # THA3: Use data/models/tha3 directory + tha_model_dir = os.path.join( + '.', 'data', 'models', 'tha3', + 'seperable' if model_seperable else 'standard', + 'fp16' if model_half else 'fp32') + print(f'THA3 Mode - Model Path: {tha_model_dir}') - tha_model_dir = os.path.join('.','data', 'models', 'tha3', - 'seperable' if model_seperable else 'standard', - 'fp16' if model_half else 'fp32') rife_model_path = '' if use_interpolation: - rife_model_path = os.path.join('.','data', 'models','rife_512', - f'x{interpolation_scale}', 'fp16' if interpolation_half else 'fp32') - + rife_model_path = os.path.join( + '.', 'data', 'models', 'rife_512', + f'x{interpolation_scale}', + 'fp16' if interpolation_half else 'fp32') + sr_model_path = '' if use_sr: if sr_x4: if sr_half: - sr_model_path = os.path.join('.','data', 'models', 'Real-ESRGAN', 'exported_256_fp16') + sr_model_path = os.path.join( + '.', 'data', 'models', 'Real-ESRGAN', + 'exported_256_fp16') else: - sr_model_path = os.path.join('.','data', 'models', 'Real-ESRGAN', 'exported_256') - else: #x2 + sr_model_path = os.path.join( + '.', 'data', 'models', 'Real-ESRGAN', + 'exported_256') + else: if sr_half: - sr_model_path = os.path.join('.','data', 'models', 'waifu2x_upconv', 'fp16', 'upconv_7', 'art', f'noise{sr_noise}_scale2x') + sr_model_path = os.path.join( + '.', 'data', 'models', 'waifu2x_upconv', + 'fp16', 'upconv_7', 'art', + f'noise{sr_noise}_scale2x') else: - sr_model_path = os.path.join('.','data', 'models', 'waifu2x_upconv', 'fp32', 'upconv_7', 'art', f'noise{sr_noise}_scale2x') + sr_model_path = os.path.join( + '.', 'data', 'models', 'waifu2x_upconv', + 'fp32', 'upconv_7', 'art', + f'noise{sr_noise}_scale2x') - print(f'THA3 Path:{tha_model_dir}') - print(f'RIFE Path:{rife_model_path}') - print(f'SR Path:{sr_model_path}') + print(f'RIFE Path: {rife_model_path}') + print(f'SR Path: {sr_model_path}') core = None if use_tensorrt: - from ezvtb_rt.core import CoreTRT - core = CoreTRT(tha_model_dir, - vram_cache_size = model_cache_size, - use_eyebrow = model_use_eyebrow, - rife_dir = rife_model_path if len(rife_model_path) > 0 else None, - sr_dir = sr_model_path if len(sr_model_path) > 0 else None, - cache_max_volume = cacher_ram_size, - cache_quality = cacher_quality, - ) - else: #Use directml - from ezvtb_rt.core_ort import CoreORT - core = CoreORT(tha_model_dir, - rife_path = rife_model_path if len(rife_model_path) > 0 else None, - sr_path = sr_model_path if len(sr_model_path) > 0 else None, - device_id=device_id, - cache_max_volume = cacher_ram_size, - cache_quality = cacher_quality, - use_eyebrow = model_use_eyebrow) + if use_tha4: + from ezvtb_rt.core_tha4 import CoreTHA4TRT + core = CoreTHA4TRT( + tha_model_dir, + vram_cache_size=model_cache_size, + use_eyebrow=model_use_eyebrow, + rife_dir=rife_model_path if rife_model_path else None, + sr_dir=sr_model_path if sr_model_path else None, + cache_max_volume=cacher_ram_size, + cache_quality=cacher_quality) + else: + from ezvtb_rt.core import CoreTRT + core = CoreTRT( + tha_model_dir, + vram_cache_size=model_cache_size, + use_eyebrow=model_use_eyebrow, + rife_dir=rife_model_path if rife_model_path else None, + sr_dir=sr_model_path if sr_model_path else None, + cache_max_volume=cacher_ram_size, + cache_quality=cacher_quality) + else: + if use_tha4: + from ezvtb_rt.core_tha4_ort import CoreTHA4ORT + core = CoreTHA4ORT( + tha_model_dir, + rife_path=rife_model_path if rife_model_path else None, + sr_path=sr_model_path if sr_model_path else None, + device_id=device_id, + cache_max_volume=cacher_ram_size, + cache_quality=cacher_quality, + use_eyebrow=model_use_eyebrow) + else: + from ezvtb_rt.core_ort import CoreORT + core = CoreORT( + tha_model_dir, + rife_path=rife_model_path if rife_model_path else None, + sr_path=sr_model_path if sr_model_path else None, + device_id=device_id, + cache_max_volume=cacher_ram_size, + cache_quality=cacher_quality, + use_eyebrow=model_use_eyebrow) return core diff --git a/ezvtuber-rt b/ezvtuber-rt index d602899..7ce4ce8 160000 --- a/ezvtuber-rt +++ b/ezvtuber-rt @@ -1 +1 @@ -Subproject commit d602899f2dde6561d6d5dfbbfa908117c10900b3 +Subproject commit 7ce4ce8d50ca17409b01644cbd714b3f8a5c1412 diff --git a/launcher2.py b/launcher2.py index 6abcfdc..8e1fcab 100644 --- a/launcher2.py +++ b/launcher2.py @@ -438,8 +438,7 @@ class LauncherPanel(wx.Panel): if args['model_select'] is not None: if 'tha4' in args['model_select']: - run_args.append('--model') - run_args.append('tha4') + run_args.append('--use_tha4') else: if 'seperable' in args['model_select']: run_args.append('--model_seperable') diff --git a/main.py b/main.py index f250514..4bcfd0b 100644 --- a/main.py +++ b/main.py @@ -347,36 +347,27 @@ class ModelClientProcess(Process): self.gpu_cache_hit_ratio = Value('f', 0.0) def run(self): - # Check if using THA4 model - if args.model == 'tha4': - from tha4_core import get_tha4_core - self.model = get_tha4_core( - device_id=args.device_id, - use_eyebrow=args.eyebrow, - interpolation_scale=args.interpolation_scale if args.use_interpolation else 1, - yaml_path=self.yaml_path - ) - else: - # Use existing ezvtb_rt interface for THA3 models - self.model = get_core(device_id=args.device_id, - use_tensorrt=args.use_tensorrt, + # Use unified ezvtb_rt interface for both THA3 and THA4 + self.model = get_core(device_id=args.device_id, + use_tensorrt=args.use_tensorrt, + use_tha4=args.use_tha4, - model_seperable = args.model_seperable, - model_half=args.model_half, - model_cache_size=args.max_gpu_cache_len, - model_use_eyebrow=args.eyebrow, + model_seperable = args.model_seperable, + model_half=args.model_half, + model_cache_size=args.max_gpu_cache_len, + model_use_eyebrow=args.eyebrow, - use_interpolation=args.use_interpolation, - interpolation_scale=args.interpolation_scale, - interpolation_half=args.interpolation_half, + use_interpolation=args.use_interpolation, + interpolation_scale=args.interpolation_scale, + interpolation_half=args.interpolation_half, - cacher_quality=args.cacher_quality, - cacher_ram_size=args.max_cache_len, + cacher_quality=args.cacher_quality, + cacher_ram_size=args.max_cache_len, - use_sr=args.use_sr, - sr_half=args.sr_half, - sr_x4=args.sr_x4, - sr_noise=args.sr_noise) + use_sr=args.use_sr, + sr_half=args.sr_half, + sr_x4=args.sr_x4, + sr_noise=args.sr_noise) self.model.setImage(self.input_image) input_pose = np.zeros((1, 45), dtype=np.float32) diff --git a/onnx_export_tha4.py b/onnx_export_tha4.py new file mode 100644 index 0000000..b346189 --- /dev/null +++ b/onnx_export_tha4.py @@ -0,0 +1,232 @@ +import torch +from os.path import join +from pathlib import Path +import sys +import os +sys.path.append(os.getcwd()) +sys.path.append(join(os.getcwd(), 'tha4', 'src')) +import numpy as np +import PIL.Image +import cv2 +import onnx +from onnxsim import simplify +import onnxruntime as ort +from torch import Tensor +from torch.nn import Module +from typing import List, Optional +from torch.nn.functional import interpolate +import matplotlib.pyplot as plt + +# ================================================================================= +# START: Monkey Patch for torch.nn.functional.affine_grid (CORRECTED, DEVICE-AWARE) +# ================================================================================= +import torch.nn.functional + +def linspace_from_neg_one(num_steps, dtype=torch.float32, device='cpu'): + r = torch.linspace(-1, 1, num_steps, dtype=torch.float32, device=device) + if num_steps > 1: r = r * (num_steps - 1) / num_steps + return r +def create_grid(N, C, H, W, device='cpu'): + grid = torch.empty((N, H, W, 3), dtype=torch.float32, device=device) + grid.select(-1, 0).copy_(linspace_from_neg_one(W, device=device)) + grid.select(-1, 1).copy_(linspace_from_neg_one(H, device=device).unsqueeze_(-1)) + grid.select(-1, 2).fill_(1); return grid +def custom_affine_grid(theta, size, align_corners=False): + N, C, H, W = size; device = theta.device + grid = create_grid(N, C, H, W, device=device) + grid = grid.view(N, H * W, 3).bmm(theta.transpose(1, 2)) + grid = grid.view(N, H, W, 2); return grid +def apply_affine_grid_patch(): + torch.nn.functional.affine_grid = custom_affine_grid + print("✅ Applied monkey patch to torch.nn.functional.affine_grid (Device-aware version).") +# ================================================================================= +# END: Monkey Patch +# ================================================================================= + +apply_affine_grid_patch() + +if len(sys.argv) != 2: + print("Usage: python onnx_export_tha4.py "); print("Available models: float, half") + raise ValueError('Do not get model name') + +MODEL_NAME = sys.argv[1] +HALF = (MODEL_NAME == 'half') +DEVICE_NAME = 'cuda:0' +IMAGE_INPUT = os.path.join('.', 'data', 'images', 'character.png') + +TMP_DIR = join('.', 'onnx_model_tha4', 'tmp'); Path(TMP_DIR).mkdir(parents=True, exist_ok=True) +MODEL_DIR = join('.', 'onnx_model_tha4', 'fp32' if not HALF else 'fp16'); Path(MODEL_DIR).mkdir(parents=True, exist_ok=True) +TMP_FILE_WRITE = join(TMP_DIR, 'tmp.onnx') + +gpu_device = torch.device(DEVICE_NAME); cpu_device = torch.device('cpu') +gpu_dtype = torch.float16 if HALF else torch.float32 +cpu_dtype = torch.float32 + +print(f"Using GPU device for validation: {gpu_device}, dtype: {gpu_dtype}") +print(f"Using CPU device for ONNX export, dtype: {cpu_dtype}") + +def torch_srgb_to_linear(x: torch.Tensor): + x = torch.clip(x, 0.0, 1.0); return torch.where(torch.le(x, 0.04045), x / 12.92, ((x + 0.055) / 1.055) ** 2.4) +def torch_linear_to_srgb(x): + x = torch.clip(x, 0.0, 1.0); return torch.where(torch.le(x, 0.003130804953560372), x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) +def numpy_linear_to_srgb(x): + x = np.clip(x, 0.0, 1.0); return np.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) +def load_tha4_poser(device: torch.device, use_half: bool): + from tha4.poser.modes.mode_07 import create_poser + m_files = { 'eyebrow_decomposer': 'data/models/tha4/eyebrow_decomposer.pt', 'eyebrow_morphing_combiner': 'data/models/tha4/eyebrow_morphing_combiner.pt', 'face_morpher': 'data/models/tha4/face_morpher.pt', 'body_morpher': 'data/models/tha4/body_morpher.pt', 'upscaler': 'data/models/tha4/upscaler.pt' } + print(f"Loading THA4 poser onto {device}..."); poser = create_poser(device=device, module_file_names=m_files) + if use_half and device.type == 'cuda': print("Converting poser modules to half precision (FP16)..."); poser.half() + return poser + +gpu_poser = load_tha4_poser(gpu_device, HALF) +cpu_poser = load_tha4_poser(cpu_device, False) +pose_size = gpu_poser.get_num_parameters(); print(f"Pose parameters: {pose_size}") + +def resize_PIL_image(img, size=(512, 512)): return img.resize(size, PIL.Image.Resampling.LANCZOS) +def extract_PIL_image_from_filelike(path): return PIL.Image.open(path) +def extract_pytorch_image_from_PIL_image(img): + arr = np.array(img) + if arr.shape[2] != 4: raise ValueError("Input image must be RGBA") + t = torch.from_numpy(arr).float() / 255.0 + t[:, :, :3] = torch_srgb_to_linear(t[:, :, :3]) + t = t.permute(2, 0, 1) * 2.0 - 1.0; return t + +pil_image = resize_PIL_image(extract_PIL_image_from_filelike(IMAGE_INPUT)) +pt_img_gpu = extract_pytorch_image_from_PIL_image(pil_image).to(gpu_dtype).unsqueeze(0).to(gpu_device) +zero_pose_gpu = torch.zeros(1, pose_size, dtype=gpu_dtype, device=gpu_device) +with torch.no_grad(): poser_torch_res = gpu_poser.pose(pt_img_gpu, zero_pose_gpu) +print(f"Reference GPU output shape: {poser_torch_res.shape}") + +print("\n" + "="*80); print("Starting ONNX export process on CPU..."); print("="*80) +cpu_modules = cpu_poser.get_modules() +cv_img_numpy = cv2.imread(IMAGE_INPUT, cv2.IMREAD_UNCHANGED) +cv_img_tensor_cpu = torch.from_numpy(cv_img_numpy) + +class EyebrowDecomposerWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor) -> List[Tensor]: + img_f = img.to(cpu_dtype) / 255.0; img_f[:, :, :3] = torch_srgb_to_linear(img_f[:, :, :3]) + # --- THE FIX: Swap BGR to RGB --- + img_f = img_f[:, :, [2, 1, 0, 3]] + img_c = img_f.permute(2, 0, 1).unsqueeze(0) * 2.0 - 1.0 + crop = img_c[:, :, 64:192, 192:320]; res = self.m(crop) + return [res[3], res[0], img_c] +eyebrow_decomposer_wrapper = EyebrowDecomposerWrapper(cpu_modules['eyebrow_decomposer']).eval() +with torch.no_grad(): decomposer_res_cpu = eyebrow_decomposer_wrapper(cv_img_tensor_cpu) +torch.onnx.export(eyebrow_decomposer_wrapper, cv_img_tensor_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['input_image'], output_names=["background_layer", "eyebrow_layer", "image_prepared"]) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "decomposer.onnx")); print("✓ Decomposer exported") + +eyebrow_pose_cpu = torch.zeros(1, 12, dtype=cpu_dtype) +class EyebrowCombinerWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor, bg: Tensor, layer: Tensor, pose: Tensor) -> Tensor: + crop = img[:, :, 32:224, 160:352].clone() + morphed = self.m(bg, layer, pose)[2] + crop[:, :, 32:160, 32:160] = morphed; return crop +eyebrow_combiner_wrapper = EyebrowCombinerWrapper(cpu_modules['eyebrow_morphing_combiner']).eval() +combiner_inputs_cpu = (decomposer_res_cpu[2], decomposer_res_cpu[0], decomposer_res_cpu[1], eyebrow_pose_cpu) +with torch.no_grad(): combiner_res_cpu = eyebrow_combiner_wrapper(*combiner_inputs_cpu) +torch.onnx.export(eyebrow_combiner_wrapper, combiner_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['image_prepared', 'eyebrow_background_layer', "eyebrow_layer", 'eyebrow_pose'], output_names=['eyebrow_image']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "combiner.onnx")); print("✓ Combiner exported") + +face_pose_cpu = torch.zeros(1, 27, dtype=cpu_dtype) +class FaceMorpherWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor, crop: Tensor, pose: Tensor) -> List[Tensor]: + morphed = self.m(crop, pose)[0]; full = img.clone() + full[:, :, 32:224, 160:352] = morphed + half = interpolate(full, size=(256, 256), mode='bilinear', align_corners=False); return [full, half] +face_morpher_wrapper = FaceMorpherWrapper(cpu_modules['face_morpher']).eval() +morpher_inputs_cpu = (decomposer_res_cpu[2], combiner_res_cpu, face_pose_cpu) +with torch.no_grad(): morpher_res_cpu = face_morpher_wrapper(*morpher_inputs_cpu) +torch.onnx.export(face_morpher_wrapper, morpher_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['image_prepared', 'im_morpher_crop', 'face_pose'], output_names=['face_morphed_full', 'face_morphed_half']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "morpher.onnx")); print("✓ Face Morpher exported") + +rotation_pose_cpu = torch.zeros(1, 6, dtype=cpu_dtype) +class BodyMorpherWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor, pose: Tensor) -> List[Tensor]: + res = self.m(img, pose); return [res[0], res[3]] +body_morpher_wrapper = BodyMorpherWrapper(cpu_modules['body_morpher']).eval() +body_morpher_inputs_cpu = (morpher_res_cpu[1], rotation_pose_cpu) +with torch.no_grad(): body_morpher_res_cpu = body_morpher_wrapper(*body_morpher_inputs_cpu) +torch.onnx.export(body_morpher_wrapper, body_morpher_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['face_morphed_half', 'rotation_pose'], output_names=['half_res_posed_image', 'half_res_grid_change']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "body_morpher.onnx")); print("✓ Body Morpher exported") + +class UpscalerWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, rest_img: Tensor, posed_half: Tensor, grid_half: Tensor, pose: Tensor) -> List[Tensor]: + posed_full = interpolate(posed_half, size=(512, 512), mode='bilinear', align_corners=False) + grid_full = interpolate(grid_half, size=(512, 512), mode='bilinear', align_corners=False) + res = self.m(rest_img, posed_full, grid_full, pose)[0] + output = res / 2.0 + 0.5; output[:, :3, :, :] = torch_linear_to_srgb(output[:, :3, :, :]) + cv_res = (output.squeeze(0).permute(1, 2, 0)[:, :, [2, 1, 0, 3]] * 255).clamp(0, 255).to(torch.uint8) + return [output, cv_res] +upscaler_wrapper = UpscalerWrapper(cpu_modules['upscaler']).eval() +upscaler_inputs_cpu = (morpher_res_cpu[0], body_morpher_res_cpu[0], body_morpher_res_cpu[1], rotation_pose_cpu) +with torch.no_grad(): upscaler_res_cpu = upscaler_wrapper(*upscaler_inputs_cpu) +torch.onnx.export(upscaler_wrapper, upscaler_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['rest_image', 'half_res_posed_image', 'half_res_grid_change', 'rotation_pose'], output_names=['result', 'cv_result']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "upscaler.onnx")); print("✓ Upscaler exported") + +print("\n" + "="*80); print("Validation: Running ONNX inference..."); print("="*80) +class THA4ONNXRunner: + def __init__(self, model_dir, use_gpu=True): + providers = [("CUDAExecutionProvider", {"device_id": 0})] if use_gpu else ['CPUExecutionProvider'] + self.decomposer_sess = self._create_session(join(model_dir, 'decomposer.onnx'), providers) + self.combiner_sess = self._create_session(join(model_dir, "combiner.onnx"), providers) + self.morpher_sess = self._create_session(join(model_dir, "morpher.onnx"), providers) + self.body_morpher_sess = self._create_session(join(model_dir, "body_morpher.onnx"), providers) + self.upscaler_sess = self._create_session(join(model_dir, "upscaler.onnx"), providers) + print("✓ All ONNX models loaded") + def _create_session(self, path, providers): + try: return ort.InferenceSession(path, providers=providers) + except Exception: print(f"Failed with {providers}, falling back to CPU for {os.path.basename(path)}"); return ort.InferenceSession(path, providers=['CPUExecutionProvider']) + def run(self, img_path, p_eye, p_face, p_rot): + img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) + d_res = self.decomposer_sess.run(None, {'input_image': img}) + c_res = self.combiner_sess.run(None, {'image_prepared': d_res[2], 'eyebrow_background_layer': d_res[0], "eyebrow_layer": d_res[1], 'eyebrow_pose': p_eye}) + m_res = self.morpher_sess.run(None, {'image_prepared': d_res[2], 'im_morpher_crop': c_res[0], 'face_pose': p_face}) + b_res = self.body_morpher_sess.run(None, {'face_morphed_half': m_res[1], 'rotation_pose': p_rot}) + u_res = self.upscaler_sess.run(None, {'rest_image': m_res[0], 'half_res_posed_image': b_res[0], 'half_res_grid_change': b_res[1], 'rotation_pose': p_rot}) + return u_res[0] + +def save_comparison_images(onnx_img_np, torch_img_np, prefix=""): + onnx_hwc = np.transpose(onnx_img_np.squeeze(0), (1, 2, 0))[:, :, :3] + torch_hwc = np.transpose(torch_img_np.squeeze(0), (1, 2, 0))[:, :, :3] + onnx_hwc = np.clip(onnx_hwc, 0, 1) + torch_hwc = np.clip(torch_hwc, 0, 1) + onnx_path = f"{prefix}onnx_output.png"; torch_path = f"{prefix}pytorch_output.png" + plt.imsave(onnx_path, onnx_hwc); plt.imsave(torch_path, torch_hwc) + print(f"✅ Saved comparison images: {onnx_path} and {torch_path}") + +runner = THA4ONNXRunner(MODEL_DIR) +p_eye_z, p_face_z, p_rot_z = np.zeros((1, 12), dtype=np.float32), np.zeros((1, 27), dtype=np.float32), np.zeros((1, 6), dtype=np.float32) + +print("\nRunning ONNX inference with zero pose...") +onnx_output = runner.run(IMAGE_INPUT, p_eye_z, p_face_z, p_rot_z) +ref_np = poser_torch_res.cpu().numpy() / 2.0 + 0.5; ref_np[0, :3, :, :] = numpy_linear_to_srgb(ref_np[0, :3, :, :]) +save_comparison_images(onnx_output, ref_np, "zero_pose_") +mse = ((onnx_output - ref_np) ** 2).mean(); print(f"\n{'='*80}\nMSE (zero pose): {mse:.10f}\n{'='*80}") +ACCEPTABLE_MSE_THRESHOLD = 1e-2 +if mse < ACCEPTABLE_MSE_THRESHOLD: print(f"✅ Validation PASSED! MSE is within the acceptable range of {ACCEPTABLE_MSE_THRESHOLD}.") +else: print(f"⚠️ Warning: MSE is higher than expected! But this can be normal due to precision differences.") + +print("\n" + "="*80); print("Testing with random poses..."); print("="*80) +p_eye_r = np.random.randn(1, 12).astype(np.float32) * 0.5 +p_face_r = np.random.randn(1, 27).astype(np.float32) * 0.5 +p_rot_r = np.random.randn(1, 6).astype(np.float32) * 0.3 +onnx_output_r = runner.run(IMAGE_INPUT, p_eye_r, p_face_r, p_rot_r) +pose_r_gpu = torch.from_numpy(np.concatenate([p_eye_r, p_face_r, p_rot_r], axis=1)).to(gpu_dtype).to(gpu_device) +with torch.no_grad(): ref_output_r = gpu_poser.pose(pt_img_gpu, pose_r_gpu) +ref_np_r = ref_output_r.cpu().numpy() / 2.0 + 0.5; ref_np_r[0, :3, :, :] = numpy_linear_to_srgb(ref_np_r[0, :3, :, :]) +save_comparison_images(onnx_output_r, ref_np_r, "random_pose_") +mse_r = ((onnx_output_r - ref_np_r) ** 2).mean(); print(f"MSE (random pose): {mse_r:.10f}") +if mse_r < ACCEPTABLE_MSE_THRESHOLD: print(f"✅ Random pose test PASSED! (Threshold: {ACCEPTABLE_MSE_THRESHOLD})") +else: print(f"⚠️ Warning: High MSE with random pose! But this can be normal.") + +print(f"\n{'='*80}\nExport and validation completed!\nModels saved to: {MODEL_DIR}\n{'='*80}") \ No newline at end of file diff --git a/tha4_core.py b/tha4_core.py index d57ad70..6d2d33d 100644 --- a/tha4_core.py +++ b/tha4_core.py @@ -1,5 +1,5 @@ """ -THA4 Core - Pure PyTorch implementation +THA4 Core - PyTorch THA4 model with optional TensorRT RIFE/SR acceleration Provides interface compatible with CoreTRT/CoreORT for THA4 models """ import torch @@ -13,10 +13,22 @@ import os sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'tha4', 'src')) from tha4.image_util import convert_linear_to_srgb +# Add ezvtuber-rt path for TensorRT modules +dir_path = os.path.dirname(os.path.realpath(__file__)) +ezvtb_path = os.path.join(dir_path, 'ezvtuber-rt') +if os.path.exists(ezvtb_path) and ezvtb_path not in sys.path: + sys.path.append(ezvtb_path) + class THA4Core: """ - THA4 Core class with interface compatible with CoreTRT/CoreORT + THA4 Core class - Pure PyTorch implementation (No TensorRT) + + This class only uses PyTorch for THA4 model inference. + It does NOT support real frame interpolation or super resolution. + If interpolation_scale > 1, frames will be simply duplicated. + + For real RIFE interpolation and SR with TensorRT, use THA4CoreWithAccel. Uses CharacterModel to load YAML + PNG + PT files properly. The PNG image from YAML is used as the source image. @@ -26,7 +38,7 @@ class THA4Core: - inference(pose: np.ndarray) -> List[np.ndarray] - run inference """ - def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, interpolation_scale=1): """ Initialize THA4 Core @@ -35,11 +47,12 @@ class THA4Core: device_id: GPU device ID use_eyebrow: whether to use eyebrow parameters yaml_path: path to character_model.yaml - Default: 'data/models/tha4/character_model.yaml' interpolation_scale: number of output frames to generate - (THA4 doesn't support real interpolation, will duplicate frames) + (THA4 doesn't support real interpolation, + will duplicate frames) """ - device_str = f'cuda:{device_id}' if torch.cuda.is_available() else 'cpu' + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' self.device = torch.device(device_str) self.use_eyebrow = use_eyebrow self.interpolation_scale = interpolation_scale @@ -54,10 +67,12 @@ class THA4Core: self.cacher = None print(f"THA4 Core initialized on {self.device}") - print(f"Using character image from model (ignoring setImage calls)") + print("Using character image from model (ignoring setImage calls)") if interpolation_scale > 1: - print(f"Note: THA4 doesn't support interpolation, " - f"will duplicate frames {interpolation_scale}x") + print(f"Note: Pure PyTorch mode doesn't support real " + f"interpolation, will duplicate frames {interpolation_scale}x") + print(" Use THA4CoreWithAccel with TensorRT for " + "real RIFE interpolation") def setImage(self, img: np.ndarray): """ @@ -80,11 +95,12 @@ class THA4Core: Args: pose: numpy array of shape (1, 45) [0:12] - eyebrow parameters - [12:39] - face parameters + [12:39] - face parameters [39:45] - head pose parameters Returns: - List of output images as numpy arrays [H, W, 4] in BGRA format, uint8 + List of output images as numpy arrays [H, W, 4] + in BGRA format, uint8 """ if self.input_image_tensor is None: raise RuntimeError("Image not set. Call setImage() first.") @@ -104,7 +120,8 @@ class THA4Core: head_tensor = torch.from_numpy(head_pose).to(self.device) # Call wrapper (it handles the pose conversion internally) - # We need to create dummy compressed versions for interface compatibility + # We need to create dummy compressed versions for + # interface compatibility eyebrow_c = eyebrow_pose[0].tolist() face_c = face_pose[0].tolist() @@ -139,12 +156,305 @@ class THA4Core: output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) # Return as list with proper number of frames - # THA4 doesn't support real interpolation, so duplicate frames if needed + # Pure PyTorch mode: duplicate frames (no real interpolation) + # Use THA4CoreWithAccel with TensorRT for real RIFE interpolation return [output_bgra] * self.interpolation_scale +class THA4CoreWithAccel: + """ + THA4 Core with TensorRT acceleration for RIFE and SR + + Architecture: + - THA4 main model: PyTorch (PT files) - always PyTorch + - RIFE interpolation: TensorRT (optional) - REAL frame interpolation + - SR super-resolution: TensorRT (optional) - REAL upscaling + - Cacher: RAM cache (optional) + + This class provides REAL interpolation and super resolution using + TensorRT, unlike THA4Core which only duplicates frames. + + Follows the same design as CoreTRT for THA3 models. + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + use_tensorrt=False, use_interpolation=False, + interpolation_scale=2, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Initialize THA4 Core with optional TensorRT acceleration + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + use_tensorrt: enable TensorRT for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_scale: interpolation scale (2/3/4) + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x (1/2/3) + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + """ + self.device_id = device_id + self.use_eyebrow = use_eyebrow + self.use_tensorrt = use_tensorrt + self.use_interpolation = use_interpolation + self.use_sr = use_sr + + # Initialize PyTorch device + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + + # Initialize THA4 wrapper (PyTorch) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + self.input_image_tensor = self.wrapper.character_image + + # Initialize TensorRT components if requested + self.rife = None + self.sr = None + self.cacher = None + self.scale = 1 + self.instream = None + self.output_memory = None + + if use_tensorrt and (use_interpolation or use_sr): + # Initialize CUDA context for TensorRT + self._init_tensorrt(device_id) + + # Build model paths + rife_model_path = '' + if use_interpolation: + rife_model_path = os.path.join( + '.', 'data', 'models', 'rife_512', + f'x{interpolation_scale}', + 'fp16' if interpolation_half else 'fp32') + + sr_model_path = '' + if use_sr: + if sr_x4: + base_path = os.path.join('.', 'data', 'models', + 'Real-ESRGAN') + sr_model_path = os.path.join( + base_path, + 'exported_256_fp16' if sr_half else 'exported_256') + else: # x2 + base_path = os.path.join('.', 'data', 'models', + 'waifu2x_upconv', + 'fp16' if sr_half else 'fp32', + 'upconv_7', 'art') + sr_model_path = os.path.join( + base_path, f'noise{sr_noise}_scale2x') + + # Initialize TensorRT modules + self._init_trt_modules(rife_model_path, sr_model_path) + + # Initialize cacher + if cacher_ram_size > 0.0: + from ezvtb_rt.cache import Cacher + self.cacher = Cacher(cacher_ram_size, cacher_quality) + + print("THA4 Core initialized:") + print(f" Device: {self.device}") + print(f" TensorRT: {use_tensorrt}") + interp_info = f"scale: {interpolation_scale}" \ + if use_interpolation else "N/A" + print(f" Interpolation: {use_interpolation} ({interp_info})") + if use_interpolation: + print(" (Real RIFE interpolation via TensorRT)") + print(f" Super Resolution: {use_sr}") + + def _init_tensorrt(self, device_id): + """Initialize TensorRT and CUDA context""" + try: + from ezvtb_rt.trt_utils import cudaSetDevice + cudaSetDevice(device_id) + os.environ['CUDA_DEVICE'] = str(device_id) + import pycuda.autoinit + device_name = pycuda.autoinit.device.name() + print(f'TensorRT initialized on device: {device_name}') + except Exception as e: + print(f'Warning: Failed to initialize TensorRT: {e}') + self.use_tensorrt = False + + def _init_trt_modules(self, rife_model_path, sr_model_path): + """Initialize RIFE and SR TensorRT modules""" + try: + import pycuda.driver as cuda + + # Create CUDA stream + self.instream = cuda.Stream() + + # Create output memory compatible with RIFE input + # [1, 512, 512, 4] + output_shape = (1, 512, 512, 4) + self.output_memory = self._create_host_device_mem( + output_shape, np.uint8) + + # Initialize RIFE + if self.use_interpolation and rife_model_path: + from ezvtb_rt.rife import RIFE + + self.rife = RIFE(rife_model_path, self.instream, + self.output_memory) + self.scale = self.rife.scale + print(f"RIFE TensorRT module loaded: {rife_model_path}") + + # Initialize SR + if self.use_sr and sr_model_path: + from ezvtb_rt.sr import SR + + if self.rife is not None: + instream = self.rife.instream + mems = [self.rife.memories['framegen_'+str(i)] + for i in range(self.rife.scale)] + else: + instream = self.instream + mems = [self.output_memory] + + self.sr = SR(sr_model_path, instream, mems) + print(f"SR TensorRT module loaded: {sr_model_path}") + + except Exception as e: + print(f'Warning: Failed to initialize TensorRT modules: {e}') + self.rife = None + self.sr = None + + def _create_host_device_mem(self, shape, dtype): + """Create HostDeviceMem for TensorRT modules""" + import pycuda.driver as cuda + from ezvtb_rt.trt_utils import HostDeviceMem + + host_mem = cuda.pagelocked_empty(shape, dtype) + device_mem = cuda.mem_alloc(host_mem.nbytes) + return HostDeviceMem(host_mem, device_mem) + + def setImage(self, img: np.ndarray): + """Set input character image - IGNORED for THA4""" + # THA4 uses character image from YAML, ignore external image + pass + + def inference(self, pose: np.ndarray) -> list: + """ + Run full inference pipeline: THA4 -> RIFE (TRT) -> SR (TRT) + + Args: + pose: numpy array of shape (1, 45) + + Returns: + List of output images [H, W, 4] in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set.") + + pose = pose.astype(np.float32) + + # Cache management + need_cache_write = 0 + res_carrier = None + + # Step 1: THA4 inference (PyTorch) with optional caching + if self.cacher is None: + output_bgra = self._run_tha4_inference(pose) + res_carrier = [output_bgra] + else: + hs = hash(str(pose)) + cached = self.cacher.read(hs) + + if cached is not None: # Cache hit + output_bgra = cached + res_carrier = [cached] + + # If using TensorRT, upload to GPU + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + else: # Cache miss + output_bgra = self._run_tha4_inference(pose) + need_cache_write = hs + res_carrier = [output_bgra] + + # Step 2: RIFE interpolation (TensorRT) + if self.rife is not None: + self.rife.inference() + res_carrier = self.rife + + # Step 3: SR super resolution (TensorRT) + if self.sr is not None: + self.sr.inference() + res_carrier = self.sr + + # Step 4: Fetch results + if need_cache_write != 0: + self.cacher.write(need_cache_write, output_bgra) + + # Get final results + if type(res_carrier) is not list: + res_carrier = res_carrier.fetchRes() + + return res_carrier + + def _run_tha4_inference(self, pose: np.ndarray) -> np.ndarray: + """Run THA4 PyTorch inference and return BGRA numpy array""" + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] + head_pose = pose[:, 39:45] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + # Run THA4 inference + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # Post-process output + output_image = output_tensor[0].float() + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + output_image = convert_linear_to_srgb(output_image) + output_np = output_image.permute(1, 2, 0).cpu().numpy() + output_np = (output_np * 255.0).astype(np.uint8) + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Upload to GPU memory if TensorRT is active + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + + return output_bgra + + def get_tha4_core(device_id=0, use_eyebrow=True, yaml_path=None, - interpolation_scale=1): + interpolation_scale=1, use_tensorrt=False, + use_interpolation=False, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): """ Factory function to create THA4 core @@ -152,13 +462,43 @@ def get_tha4_core(device_id=0, use_eyebrow=True, yaml_path=None, device_id: GPU device ID use_eyebrow: whether to use eyebrow control yaml_path: path to character_model.yaml - Default: 'data/models/tha4/character_model.yaml' - interpolation_scale: number of frames (for compatibility with interpolation) + interpolation_scale: number of frames for interpolation + use_tensorrt: enable TensorRT acceleration for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB Returns: - THA4Core instance + THA4Core or THA4CoreWithAccel instance """ - return THA4Core(device_id=device_id, - use_eyebrow=use_eyebrow, - yaml_path=yaml_path, - interpolation_scale=interpolation_scale) + # Use accelerated version if TensorRT is requested + if use_tensorrt and (use_interpolation or use_sr): + return THA4CoreWithAccel( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + use_tensorrt=use_tensorrt, + use_interpolation=use_interpolation, + interpolation_scale=interpolation_scale, + interpolation_half=interpolation_half, + use_sr=use_sr, + sr_x4=sr_x4, + sr_half=sr_half, + sr_noise=sr_noise, + cacher_quality=cacher_quality, + cacher_ram_size=cacher_ram_size + ) + else: + # Use simple PyTorch-only version + return THA4Core( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + interpolation_scale=interpolation_scale if use_interpolation + else 1 + ) diff --git a/tha4_core_old.py b/tha4_core_old.py new file mode 100644 index 0000000..5add4d5 --- /dev/null +++ b/tha4_core_old.py @@ -0,0 +1,498 @@ +""" +THA4 Core - ONNX-based THA4 model with optional TensorRT/RIFE/SR acceleration +Provides interface compatible with CoreTRT/CoreORT for THA4 models +Uses teacher models (5 models) from onnx_model_tha4 folder +""" +import numpy as np +import sys +import os + +# Add ezvtuber-rt path for TensorRT/RIFE/SR modules +dir_path = os.path.dirname(os.path.realpath(__file__)) +ezvtb_path = os.path.join(dir_path, 'ezvtuber-rt') +if os.path.exists(ezvtb_path) and ezvtb_path not in sys.path: + sys.path.append(ezvtb_path) + + +class THA4CoreTRT: + """ + THA4 Core class - Pure PyTorch implementation (No TensorRT) + + This class only uses PyTorch for THA4 model inference. + It does NOT support real frame interpolation or super resolution. + If interpolation_scale > 1, frames will be simply duplicated. + + For real RIFE interpolation and SR with TensorRT, use THA4CoreWithAccel. + + Uses CharacterModel to load YAML + PNG + PT files properly. + The PNG image from YAML is used as the source image. + + Expected interface: + - setImage(img: np.ndarray) - NOT USED, image comes from YAML + - inference(pose: np.ndarray) -> List[np.ndarray] - run inference + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1): + """ + Initialize THA4 Core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + interpolation_scale: number of output frames to generate + (THA4 doesn't support real interpolation, + will duplicate frames) + """ + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + self.use_eyebrow = use_eyebrow + self.interpolation_scale = interpolation_scale + + # Create THA4 wrapper (loads YAML + PNG + PT) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + + # Use character image from YAML (already loaded by wrapper) + self.input_image_tensor = self.wrapper.character_image + + # For compatibility (no caching in THA4 core) + self.cacher = None + + print(f"THA4 Core initialized on {self.device}") + print("Using character image from model (ignoring setImage calls)") + if interpolation_scale > 1: + print(f"Note: Pure PyTorch mode doesn't support real " + f"interpolation, will duplicate frames {interpolation_scale}x") + print(" Use THA4CoreWithAccel with TensorRT for " + "real RIFE interpolation") + + def setImage(self, img: np.ndarray): + """ + Set input character image - IGNORED for THA4 + + THA4 uses the character image from the YAML/PNG files loaded + during initialization. This method exists only for interface + compatibility with CoreTRT/CoreORT. + + Args: + img: numpy array in BGRA format (ignored) + """ + # THA4 uses character image from YAML, ignore external image + print("THA4: setImage() called but ignored (using image from YAML)") + + def inference(self, pose: np.ndarray) -> list: + """ + Run inference + + Args: + pose: numpy array of shape (1, 45) + [0:12] - eyebrow parameters + [12:39] - face parameters + [39:45] - head pose parameters + + Returns: + List of output images as numpy arrays [H, W, 4] + in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set. Call setImage() first.") + + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] # [1, 12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] # [1, 27] + head_pose = pose[:, 39:45] # [1, 6] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + # Call wrapper (it handles the pose conversion internally) + # We need to create dummy compressed versions for + # interface compatibility + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # THA4 outputs in linear color space, convert to sRGB + # Output is [batch, 4, H, W] in range [-1, 1] + output_image = output_tensor[0].float() # [4, 512, 512] + + # Clip to [-1, 1] and convert to [0, 1] + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + + # Convert linear RGB to sRGB (keep alpha as-is) + output_image = convert_linear_to_srgb(output_image) + + # Convert to [H, W, C] numpy + output_np = output_image.permute(1, 2, 0).cpu().numpy() + + # Convert to uint8 + output_np = (output_np * 255.0).astype(np.uint8) + + # Convert RGBA to BGRA for output + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Return as list with proper number of frames + # Pure PyTorch mode: duplicate frames (no real interpolation) + # Use THA4CoreWithAccel with TensorRT for real RIFE interpolation + return [output_bgra] * self.interpolation_scale + + +class THA4CoreWithAccel: + """ + THA4 Core with TensorRT acceleration for RIFE and SR + + Architecture: + - THA4 main model: PyTorch (PT files) - always PyTorch + - RIFE interpolation: TensorRT (optional) - REAL frame interpolation + - SR super-resolution: TensorRT (optional) - REAL upscaling + - Cacher: RAM cache (optional) + + This class provides REAL interpolation and super resolution using + TensorRT, unlike THA4Core which only duplicates frames. + + Follows the same design as CoreTRT for THA3 models. + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + use_tensorrt=False, use_interpolation=False, + interpolation_scale=2, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Initialize THA4 Core with optional TensorRT acceleration + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + use_tensorrt: enable TensorRT for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_scale: interpolation scale (2/3/4) + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x (1/2/3) + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + """ + self.device_id = device_id + self.use_eyebrow = use_eyebrow + self.use_tensorrt = use_tensorrt + self.use_interpolation = use_interpolation + self.use_sr = use_sr + + # Initialize PyTorch device + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + + # Initialize THA4 wrapper (PyTorch) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + self.input_image_tensor = self.wrapper.character_image + + # Initialize TensorRT components if requested + self.rife = None + self.sr = None + self.cacher = None + self.scale = 1 + self.instream = None + self.output_memory = None + + if use_tensorrt and (use_interpolation or use_sr): + # Initialize CUDA context for TensorRT + self._init_tensorrt(device_id) + + # Build model paths + rife_model_path = '' + if use_interpolation: + rife_model_path = os.path.join( + '.', 'data', 'models', 'rife_512', + f'x{interpolation_scale}', + 'fp16' if interpolation_half else 'fp32') + + sr_model_path = '' + if use_sr: + if sr_x4: + base_path = os.path.join('.', 'data', 'models', + 'Real-ESRGAN') + sr_model_path = os.path.join( + base_path, + 'exported_256_fp16' if sr_half else 'exported_256') + else: # x2 + base_path = os.path.join('.', 'data', 'models', + 'waifu2x_upconv', + 'fp16' if sr_half else 'fp32', + 'upconv_7', 'art') + sr_model_path = os.path.join( + base_path, f'noise{sr_noise}_scale2x') + + # Initialize TensorRT modules + self._init_trt_modules(rife_model_path, sr_model_path) + + # Initialize cacher + if cacher_ram_size > 0.0: + from ezvtb_rt.cache import Cacher + self.cacher = Cacher(cacher_ram_size, cacher_quality) + + print("THA4 Core initialized:") + print(f" Device: {self.device}") + print(f" TensorRT: {use_tensorrt}") + interp_info = f"scale: {interpolation_scale}" \ + if use_interpolation else "N/A" + print(f" Interpolation: {use_interpolation} ({interp_info})") + if use_interpolation: + print(" (Real RIFE interpolation via TensorRT)") + print(f" Super Resolution: {use_sr}") + + def _init_tensorrt(self, device_id): + """Initialize TensorRT and CUDA context""" + try: + from ezvtb_rt.trt_utils import cudaSetDevice + cudaSetDevice(device_id) + os.environ['CUDA_DEVICE'] = str(device_id) + import pycuda.autoinit + device_name = pycuda.autoinit.device.name() + print(f'TensorRT initialized on device: {device_name}') + except Exception as e: + print(f'Warning: Failed to initialize TensorRT: {e}') + self.use_tensorrt = False + + def _init_trt_modules(self, rife_model_path, sr_model_path): + """Initialize RIFE and SR TensorRT modules""" + try: + import pycuda.driver as cuda + + # Create CUDA stream + self.instream = cuda.Stream() + + # Create output memory compatible with RIFE input + # [1, 512, 512, 4] + output_shape = (1, 512, 512, 4) + self.output_memory = self._create_host_device_mem( + output_shape, np.uint8) + + # Initialize RIFE + if self.use_interpolation and rife_model_path: + from ezvtb_rt.rife import RIFE + + self.rife = RIFE(rife_model_path, self.instream, + self.output_memory) + self.scale = self.rife.scale + print(f"RIFE TensorRT module loaded: {rife_model_path}") + + # Initialize SR + if self.use_sr and sr_model_path: + from ezvtb_rt.sr import SR + + if self.rife is not None: + instream = self.rife.instream + mems = [self.rife.memories['framegen_'+str(i)] + for i in range(self.rife.scale)] + else: + instream = self.instream + mems = [self.output_memory] + + self.sr = SR(sr_model_path, instream, mems) + print(f"SR TensorRT module loaded: {sr_model_path}") + + except Exception as e: + print(f'Warning: Failed to initialize TensorRT modules: {e}') + self.rife = None + self.sr = None + + def _create_host_device_mem(self, shape, dtype): + """Create HostDeviceMem for TensorRT modules""" + import pycuda.driver as cuda + from ezvtb_rt.trt_utils import HostDeviceMem + + host_mem = cuda.pagelocked_empty(shape, dtype) + device_mem = cuda.mem_alloc(host_mem.nbytes) + return HostDeviceMem(host_mem, device_mem) + + def setImage(self, img: np.ndarray): + """Set input character image - IGNORED for THA4""" + # THA4 uses character image from YAML, ignore external image + pass + + def inference(self, pose: np.ndarray) -> list: + """ + Run full inference pipeline: THA4 -> RIFE (TRT) -> SR (TRT) + + Args: + pose: numpy array of shape (1, 45) + + Returns: + List of output images [H, W, 4] in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set.") + + pose = pose.astype(np.float32) + + # Cache management + need_cache_write = 0 + res_carrier = None + + # Step 1: THA4 inference (PyTorch) with optional caching + if self.cacher is None: + output_bgra = self._run_tha4_inference(pose) + res_carrier = [output_bgra] + else: + hs = hash(str(pose)) + cached = self.cacher.read(hs) + + if cached is not None: # Cache hit + output_bgra = cached + res_carrier = [cached] + + # If using TensorRT, upload to GPU + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + else: # Cache miss + output_bgra = self._run_tha4_inference(pose) + need_cache_write = hs + res_carrier = [output_bgra] + + # Step 2: RIFE interpolation (TensorRT) + if self.rife is not None: + self.rife.inference() + res_carrier = self.rife + + # Step 3: SR super resolution (TensorRT) + if self.sr is not None: + self.sr.inference() + res_carrier = self.sr + + # Step 4: Fetch results + if need_cache_write != 0: + self.cacher.write(need_cache_write, output_bgra) + + # Get final results + if type(res_carrier) is not list: + res_carrier = res_carrier.fetchRes() + + return res_carrier + + def _run_tha4_inference(self, pose: np.ndarray) -> np.ndarray: + """Run THA4 PyTorch inference and return BGRA numpy array""" + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] + head_pose = pose[:, 39:45] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + # Run THA4 inference + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # Post-process output + output_image = output_tensor[0].float() + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + output_image = convert_linear_to_srgb(output_image) + output_np = output_image.permute(1, 2, 0).cpu().numpy() + output_np = (output_np * 255.0).astype(np.uint8) + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Upload to GPU memory if TensorRT is active + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + + return output_bgra + + +def get_tha4_core(device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1, use_tensorrt=False, + use_interpolation=False, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Factory function to create THA4 core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow control + yaml_path: path to character_model.yaml + interpolation_scale: number of frames for interpolation + use_tensorrt: enable TensorRT acceleration for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + + Returns: + THA4Core or THA4CoreWithAccel instance + """ + # Use accelerated version if TensorRT is requested + if use_tensorrt and (use_interpolation or use_sr): + return THA4CoreWithAccel( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + use_tensorrt=use_tensorrt, + use_interpolation=use_interpolation, + interpolation_scale=interpolation_scale, + interpolation_half=interpolation_half, + use_sr=use_sr, + sr_x4=sr_x4, + sr_half=sr_half, + sr_noise=sr_noise, + cacher_quality=cacher_quality, + cacher_ram_size=cacher_ram_size + ) + else: + # Use simple PyTorch-only version + return THA4Core( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + interpolation_scale=interpolation_scale if use_interpolation + else 1 + )