tha4 teachermodel support

.gitignore

@@ -144,4 +144,7 @@ data/
 .idea/
 *.pt
 
-.serena
+.serena
+onnx_model_tha4/fp16
+onnx_model_tha4/fp32
+onnx_model_tha4/tmp

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')

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()

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
     

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')

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)
 

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 <model_name>"); 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}")

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
+        )

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
+        )