From b406da10daa043cc539ad4509ff517be602860ba Mon Sep 17 00:00:00 2001 From: zpeng11 Date: Sun, 29 Mar 2026 16:41:33 -0400 Subject: [PATCH] feat:maintainance for better startup speed --- .gitignore | 5 +- ....\345\220\257\345\212\250\345\231\250.bat" | 2 +- ...5\350\276\223\345\207\272\357\274\211.bat" | 2 +- README.md | 19 +- ezvtuber-rt | 2 +- launcher.py | 360 ------------------ launcher2.py | 8 + models.py | 237 ------------ pack_release.bat | 43 ++- src/main.py | 2 - 10 files changed, 55 insertions(+), 625 deletions(-) delete mode 100644 launcher.py delete mode 100644 models.py diff --git a/.gitignore b/.gitignore index 3f5ecf6c..1b2b5f49 100644 --- a/.gitignore +++ b/.gitignore @@ -151,4 +151,7 @@ data/ .serena onnx_model_tha4/fp16 onnx_model_tha4/fp32 -onnx_model_tha4/ \ No newline at end of file +onnx_model_tha4/ + +# Tensorrt temporary files +*.pch \ No newline at end of file diff --git "a/01A.\345\220\257\345\212\250\345\231\250.bat" "b/01A.\345\220\257\345\212\250\345\231\250.bat" index 08cca032..785cb620 100644 --- "a/01A.\345\220\257\345\212\250\345\231\250.bat" +++ "b/01A.\345\220\257\345\212\250\345\231\250.bat" @@ -1,7 +1,7 @@ @echo on cd /D "%~dp0" -SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% +SET PATH=%~dp0envs\TensorRT-RTX\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% start "" pythonw launcher2.py exit \ No newline at end of file diff --git "a/01B.\345\220\257\345\212\250\345\231\250\357\274\210\350\260\203\350\257\225\350\276\223\345\207\272\357\274\211.bat" "b/01B.\345\220\257\345\212\250\345\231\250\357\274\210\350\260\203\350\257\225\350\276\223\345\207\272\357\274\211.bat" index 453f60c0..23309b22 100644 --- "a/01B.\345\220\257\345\212\250\345\231\250\357\274\210\350\260\203\350\257\225\350\276\223\345\207\272\357\274\211.bat" +++ "b/01B.\345\220\257\345\212\250\345\231\250\357\274\210\350\260\203\350\257\225\350\276\223\345\207\272\357\274\211.bat" @@ -1,7 +1,7 @@ @echo on cd /D "%~dp0" -SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% +SET PATH=%~dp0envs\TensorRT-RTX\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% python launcher2.py diff --git a/README.md b/README.md index 989ca5df..ac293e1f 100644 --- a/README.md +++ b/README.md @@ -165,14 +165,6 @@ conda 24.11.3 ``` 以上方式可以保证conda环境被你选定的Ananconda统一管理。若没有此需求可以忽略直接运行下一步。 -### 下载 TensorRT-RTX （Nvidia开发必选） -``` -curl -L -o trt_rtx.zip https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.3/TensorRT-RTX-1.3.0.35-win10-amd64-cuda-12.9-Release-external.zip -tar -xf trt_rtx.zip "TensorRT-RTX-1.3.0.35" && ren "TensorRT-RTX-1.3.0.35" "TensorRT-RTX-1.3.0.35_cu129" && del trt_rtx.zip -# Adding the bin folder to PATH -cd TensorRT-RTX-1.3.0.35_cu129\bin && set PATH=%CD%;%PATH% -``` -也可以在Window 环境变量管理器 GUI中将`TensorRT-RTX-1.3.0.35_cu129\bin`在`PATH`环境变量中添加来持久化 ### 准备Conda 环境 ``` @@ -182,10 +174,15 @@ conda install conda-forge::pycuda conda install -c nvidia/label/cuda-12.9.1 cuda-nvcc-dev_win-64 cudnn cuda-runtime ``` -### 安装 TensorRT-RTX 的 Python Binding -``` -pip install TensorRT-RTX-1.3.0.35_cu129\python\tensorrt_rtx-1.3.0.35-cp310-none-win_amd64.whl +### 下载 TensorRT-RTX （Nvidia开发必选） +1. Go to https://developer.nvidia.com/tensorrt-rtx download for windows cuda129 +2. Unzip downloaded folder +3. Add bin folder to environment PATH +4. Activate your python environment and install like below: +```bash +pip install D:\TensorRT-RTX-1.4.0.76_cu129\python\tensorrt_rtx-1.4.0.76-cp310-none-win_amd64.whl ``` +也可以在Window 环境变量管理器 GUI中将`TensorRT-RTX-1.4.0.76_cu129\bin`在`PATH`环境变量中添加来持久化 ### 克隆项目和子项目 ``` diff --git a/ezvtuber-rt b/ezvtuber-rt index 07c30292..ecd0eeb8 160000 --- a/ezvtuber-rt +++ b/ezvtuber-rt @@ -1 +1 @@ -Subproject commit 07c302921e67b483b023a7db5c0fa96ac5fd8e6d +Subproject commit ecd0eeb8cdcbf5acfcd4564bd56abcd5ab09d072 diff --git a/launcher.py b/launcher.py deleted file mode 100644 index 984cdbca..00000000 --- a/launcher.py +++ /dev/null @@ -1,360 +0,0 @@ -import os -import subprocess -import tkinter as tk -import tkinter.messagebox -from tkinter import ttk -import json -import sys - -cache_simplify_map = { - 'Off':0, - 'Low':1, - 'Medium':2, - 'High':3, - 'Higher':4, - 'Highest':6, - 'Gaming':8 -} - -cache_simplify_quality_map = { - 'Off':100, - 'Low':99, - 'Medium':95, - 'High':90, - 'Higher':85, - 'Highest':80, - 'Gaming':75 -} - -default_arg = { - 'character': 'lambda_00', - 'input': 2, - 'output': 2, - 'ifm': None, - 'osf': '127.0.0.1:11573', - 'is_extend_movement': False, - 'is_alpha_split': False, - 'is_bongo': False, - 'is_eyebrow': False, - 'cache_simplify': 'High', - 'ram_cache_size': '2gb', - 'vram_cache_size': '2gb', - 'model_select':'seperable_half', - 'interpolation':'x2_half', - 'frame_rate_limit':'30', - 'sr':'Off', - 'device_id':'0', - 'use_tensorrt':True -} - -try: - f = open('launcher.json') - args = json.load(f) - default_arg.update(args) - f.close() -except: - pass -finally: - args = default_arg - -p = None -dirPath = 'data/images' -characterList = [] -for item in sorted(os.listdir(dirPath), key=lambda x: -os.path.getmtime(os.path.join(dirPath, x))): - if '.png' == item[-4:]: - characterList.append(item[:-4]) - -root = tk.Tk() -root.resizable(False, False) -root.title('EasyVtuber Launcher') - -launcher = ttk.Frame(root) -launcher.pack(fill='x', expand=True) - - -def launch(): - global p - global launch_btn - args = { - 'character': character.get(), - 'input': input.get(), - 'output': output.get(), - 'ifm': ifm.get(), - 'osf': osf.get(), - 'is_extend_movement': is_extend_movement.get(), - 'is_alpha_split': is_alpha_split.get(), - 'is_bongo': is_bongo.get(), - 'is_eyebrow': is_eyebrow.get(), - 'cache_simplify': cache_simplify.get(), - 'ram_cache_size': ram_cache_size.get(), - 'vram_cache_size': vram_cache_size.get(), - 'model_select': model_select.get(), - 'interpolation':interpolation.get(), - 'frame_rate_limit':frame_rate_limit.get(), - 'sr':sr.get(), - 'device_id':device_id.get(), - 'use_tensorrt':use_tensorrt.get() - } - - if args['input'] == 0: - if len(args['ifm']) == 0: - tkinter.messagebox.showinfo('EasyVtuber Launcher', 'Please Input iFacialMocap IP:Port') - return - if args['input'] == 4: - if len(args['osf']) == 0: - tkinter.messagebox.showinfo('EasyVtuber Launcher', 'Please Input OpenSeeFace IP:Port') - return - - f = open('launcher.json', mode='w') - json.dump(args, f) - f.close() - if p is not None: - subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL) - p = None - launch_btn.config(text="Save & Launch") - else: - run_args = [sys.executable, 'main.py'] - if len(args['character']): - run_args.append('--character') - run_args.append(args['character']) - - if args['input'] == 0: - if len(args['ifm']): - run_args.append('--ifm') - run_args.append(args['ifm']) - elif args['input'] == 1: - run_args.append('--input') - run_args.append('cam') - elif args['input'] == 2: - run_args.append('--debug_input') - elif args['input'] == 3: - run_args.append('--mouse_input') - run_args.append('0,0,' + str(root.winfo_screenwidth()) + ',' + str(root.winfo_screenheight())) - elif args['input'] == 4: - if len(args['osf']): - run_args.append('--osf') - run_args.append(args['osf']) - - if args['output'] == 0: - run_args.append('--output_webcam') - run_args.append('unitycapture') - elif args['output'] == 1: - run_args.append('--output_webcam') - run_args.append('obs') - elif args['output'] == 3: - run_args.append('--output_webcam') - run_args.append('spout') - elif args['output'] == 2: - run_args.append('--debug') - if args['is_alpha_split']: - run_args.append('--alpha_split') - if args['is_extend_movement']: - run_args.append('--extend_movement') - run_args.append('1') - if args['is_bongo']: - run_args.append('--bongo') - if args['is_eyebrow']: - run_args.append('--eyebrow') - if args['cache_simplify'] is not None: - run_args.append('--simplify') - run_args.append(str(cache_simplify_map[args['cache_simplify']])) - run_args.append('--cacher_quality') - run_args.append(str(cache_simplify_quality_map[args['cache_simplify']])) - if args['cache_simplify'] != 'Off': - run_args.append('--use_cacher') - if args['ram_cache_size'] is not None: - run_args.append('--cache') - run_args.append(args['ram_cache_size']) - run_args.append('--gpu_cache') - run_args.append(args['vram_cache_size']) - - if args['interpolation'] is not None: - if not 'off' == args['interpolation']: - run_args.append('--use_interpolation') - if 'half' in args['interpolation']: - run_args.append('--interpolation_half') - - if 'x2' in args['interpolation']: - run_args.append('--interpolation_scale') - run_args.append('2') - elif 'x3' in args['interpolation']: - run_args.append('--interpolation_scale') - run_args.append('3') - elif 'x4' in args['interpolation']: - run_args.append('--interpolation_scale') - run_args.append('4') - - if args['model_select'] is not None: - if 'seperable' in args['model_select']: - run_args.append('--model_seperable') - if 'half' in args['model_select']: - run_args.append('--model_half') - - if args['frame_rate_limit'] is not None: - run_args.append('--frame_rate_limit') - run_args.append(args['frame_rate_limit']) - - if args['sr'] is not None and args['sr'] != 'Off': - if 'anime4k' in args['sr']: - run_args.append('--anime4k') - else: - run_args.append('--use_sr') - if 'x4' in args['sr']: - run_args.append('--sr_x4') - if 'half' in args['sr']: - run_args.append('--sr_half') - - if args['device_id'] is not None: - run_args.append('--device_id') - run_args.append(args['device_id']) - - if args['use_tensorrt'] is not None and args['use_tensorrt']: - run_args.append('--use_tensorrt') - run_args.append('--model_cache') - run_args.append('--model_vram_cache') - - run_args.append('--output_size') - run_args.append('512x512') - print('Launched: ' + ' '.join(run_args)) - p = subprocess.Popen(run_args) - launch_btn.config(text='Stop') - - -launch_btn = ttk.Button(launcher, text="Save & Launch", command=launch) -launch_btn.pack(side='bottom', fill='x', expand=True, pady=10, padx=10) - -frameL = ttk.Frame(launcher) -frameL.pack(padx=10, pady=10, fill='both', side='left', expand=True) -frameR = ttk.Frame(launcher) -frameR.pack(padx=10, pady=10, fill='both', side='left', expand=True) - -character = tk.StringVar(value=args['character']) -ttk.Label(frameL, text="Character").pack(fill='x', expand=True) - -# ttk.Entry(frameL, textvariable=character).pack(fill='x', expand=True) -char_combo = ttk.Combobox(frameL, textvariable=character, value=characterList) -char_combo.pack(fill='x', expand=True) - - - -def inputChange(): - i=input.get() - if i==0: - ifmLbl.pack(fill='x', expand=True) - ifmEnt.pack(fill='x', expand=True) - osfLbl.pack_forget() - osfEnt.pack_forget() - elif i==4: - ifmLbl.pack_forget() - ifmEnt.pack_forget() - osfLbl.pack(fill='x', expand=True) - osfEnt.pack(fill='x', expand=True) - else: - ifmLbl.pack_forget() - ifmEnt.pack_forget() - osfLbl.pack_forget() - osfEnt.pack_forget() - -input = tk.IntVar(value=args['input']) -ttk.Label(frameL, text="Face Data Source").pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='iFacialMocap', value=0, variable=input, command=inputChange).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='OpenSeeFace', value=4, variable=input, command=inputChange).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='Webcam(opencv)', value=1, variable=input, command=inputChange).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='Mouse Input', value=3, variable=input, command=inputChange).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='Initial Debug Input', value=2, variable=input, command=inputChange).pack(fill='x', - expand=True) -frameLTxt = ttk.Frame(frameL) -frameLTxt.pack(fill='x', expand=True) -ifmLbl = ttk.Label(frameLTxt, text="iFacialMocap IP:Port") -ifmLbl.pack(fill='x', expand=True) - -ifm = tk.StringVar(value=args['ifm']) -ifmEnt = ttk.Entry(frameLTxt, textvariable=ifm, state=False) -ifmEnt.pack(fill='x', expand=True) - -osfLbl = ttk.Label(frameLTxt, text="OpenSeeFace IP:Port") -osfLbl.pack(fill='x', expand=True) -osf = tk.StringVar(value=args['osf']) -osfEnt = ttk.Entry(frameLTxt, textvariable=osf, state=False) -osfEnt.pack(fill='x', expand=True) -inputChange() - -ttk.Label(frameR, text="GPU Device ID").pack(fill='x', expand=True) -device_id = tk.StringVar(value=args['device_id']) -device_id_combo = ttk.Combobox(frameR, textvariable=device_id, value=['0', '1', '2', '3', '4'], state='readonly').pack(fill='x', expand=True) - -use_tensorrt = tk.BooleanVar(value=args['use_tensorrt']) -ttk.Checkbutton(frameR, text='Use TensorRT', variable=use_tensorrt).pack(fill='x', expand=True) - -ttk.Label(frameR, text="Model Select").pack(fill='x', expand=True) -model_select = tk.StringVar(value=args['model_select']) -model_select_combo = ttk.Combobox(frameR, textvariable=model_select, value=['seperable_half', 'seperable_full','standard_half','standard_full'], state='readonly').pack(fill='x', expand=True) - -ttk.Label(frameR, text="VRAM Cache Size (TensorRT only)").pack(fill='x', expand=True) -vram_cache_size = tk.StringVar(value=args['vram_cache_size']) -vram_cache_size_combo = ttk.Combobox(frameR, textvariable=vram_cache_size, value=['0b', '128mb', '256mb','512mb', '1gb', '2gb', '4gb'], state='readonly').pack(fill='x', expand=True) - -ttk.Label(frameR, text="Facial Input Simplify & Cache Quality").pack(fill='x', expand=True) -cache_simplify = tk.StringVar(value=args['cache_simplify']) -cache_simplify_combo = ttk.Combobox(frameR, textvariable=cache_simplify, value=['Off', 'Low','Medium','High', 'Higher', 'Highest', 'Gaming'], state='readonly').pack(fill='x', expand=True) - -ttk.Label(frameR, text="RAM Cache Size").pack(fill='x', expand=True) -ram_cache_size = tk.StringVar(value=args['ram_cache_size']) -ram_cache_size_combo = ttk.Combobox(frameR, textvariable=ram_cache_size, value=['0b', '1gb', '2gb', '4gb', '8gb', '16gb'], state='readonly').pack(fill='x', expand=True) - -ttk.Label(frameR, text="Frame Interpolation").pack(fill='x', expand=True) -interpolation = tk.StringVar(value=args['interpolation']) -interpolation_combo = ttk.Combobox(frameR, textvariable=interpolation, value=['off', 'x2_half', 'x3_half','x4_half', 'x2_full', 'x3_full', 'x4_full'], state='readonly').pack(fill='x', expand=True) - -ttk.Label(frameR, text="Frame Rate Limitation").pack(fill='x', expand=True) -frame_rate_limit = tk.StringVar(value=args['frame_rate_limit']) -frame_rate_limit_combo = ttk.Combobox(frameR, textvariable=frame_rate_limit, value=['20', '25', '30','40', '50', '60'], state='readonly').pack(fill='x', expand=True) - -ttk.Label(frameR, text="Super Resolution").pack(fill='x', expand=True) -sr = tk.StringVar(value=args['sr']) -sr_combo = ttk.Combobox(frameR, textvariable=sr, value=['Off', 'anime4k_x2', 'waifu2x_x2_half', 'real-esrgan_x4_half', 'waifu2x_x2_full', 'real-esrgan_x4_full'], state='readonly').pack(fill='x', expand=True) - - - -ttk.Label(frameL, text="Extra Options").pack(fill='x', expand=True) -is_eyebrow = tk.BooleanVar(value=args['is_eyebrow']) -ttk.Checkbutton(frameL, text='Eyebrow (iFM Only)', variable=is_eyebrow).pack(fill='x', expand=True) - -is_extend_movement = tk.BooleanVar(value=args['is_extend_movement']) -ttk.Checkbutton(frameL, text='Extend Movement', variable=is_extend_movement).pack(fill='x', expand=True) - - -is_alpha_split = tk.BooleanVar(value=args['is_alpha_split']) -ttk.Checkbutton(frameL, text='Alpha Split', variable=is_alpha_split).pack(fill='x', expand=True) - -is_bongo = tk.BooleanVar(value=args['is_bongo']) -ttk.Checkbutton(frameL, text='Bongocat Mode', variable=is_bongo).pack(fill='x', expand=True) - -output = tk.IntVar(value=args['output']) -ttk.Label(frameL, text="Output").pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='Unity Capture', value=0, variable=output).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='OBS Virtual Camera', value=1, variable=output).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='Spout2', value=3, variable=output).pack(fill='x', expand=True) -ttk.Radiobutton(frameL, text='Initial Debug Output', value=2, variable=output).pack(fill='x', expand=True) - - -def closeWindow(): - if p is not None: - subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL) - root.destroy() - - -def handle_focus(event): - characterList = [] - if event.widget == root: - for item in sorted(os.listdir(dirPath), key=lambda x: -os.path.getmtime(os.path.join(dirPath, x))): - if '.png' == item[-4:]: - characterList.append(item[:-4]) - char_combo.config(value=characterList) - - -root.bind("", handle_focus) -root.protocol('WM_DELETE_WINDOW', closeWindow) -root.mainloop() diff --git a/launcher2.py b/launcher2.py index c30d4cae..28d540a8 100644 --- a/launcher2.py +++ b/launcher2.py @@ -269,6 +269,14 @@ def _important_log_line(line): name = os.path.basename(path) if name: return f'Building: {name}' + # TRT: Loading engine from file ...\filename.trt + if '[TRT]' in line and 'Loading engine from file' in line: + idx = line.find('Loading engine from file') + if idx != -1: + path = line[idx + len('Loading engine from file'):].strip().strip('.').strip().rstrip('\r\n') + name = os.path.basename(path) + if name: + return f'Loading: {name}' # TRT: Loading ONNX file from path ...\filename.onnx if '[TRT]' in line and 'Loading ONNX file from path' in line: idx = line.find('Loading ONNX file from path') diff --git a/models.py b/models.py deleted file mode 100644 index 34abb088..00000000 --- a/models.py +++ /dev/null @@ -1,237 +0,0 @@ -import time - -import torch -import torch.nn as nn - -import tha2.poser.modes.mode_20 -import tha3.poser.modes.standard_float -import tha3.poser.modes.separable_float -import tha3.poser.modes.standard_half -import tha3.poser.modes.separable_half -from torch.nn.functional import interpolate - -from args import args - -from collections import OrderedDict - -# THA4 适配器导入 -try: - from tha4_adapter import THA4Wrapper - THA4_AVAILABLE = True -except ImportError: - THA4_AVAILABLE = False - print("Warning: THA4 adapter not available") - - -class TalkingAnimeLight(nn.Module): - def __init__(self): - super(TalkingAnimeLight, self).__init__() - self.face_morpher = tha2.poser.modes.mode_20.load_face_morpher('pretrained/face_morpher.pt') - self.two_algo_face_rotator = tha2.poser.modes.mode_20.load_face_rotater('pretrained/two_algo_face_rotator.pt') - self.combiner = tha2.poser.modes.mode_20.load_combiner('pretrained/combiner.pt') - self.face_cache = OrderedDict() - self.tot = 0 - self.hit = 0 - - def forward(self, image, mouth_eye_vector, pose_vector, mouth_eye_vector_c, ratio=None): - x = image.clone() - if args.perf == 'model': - tic = time.perf_counter() - input_hash = hash(tuple(mouth_eye_vector_c)) - cached = self.face_cache.get(input_hash) - self.tot += 1 - if cached is None: - mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) - self.face_cache[input_hash] = mouth_eye_morp_image.detach() - if len(self.face_cache) > args.max_gpu_cache_len: - self.face_cache.popitem(last=False) - else: - self.hit += 1 - mouth_eye_morp_image = cached - self.face_cache.move_to_end(input_hash) - if args.debug and ratio is not None: - ratio.value = self.hit / self.tot - if args.perf == 'model': - print(" - face_morpher", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - x[:, :, 32:224, 32:224] = mouth_eye_morp_image - rotate_image = self.two_algo_face_rotator(x, pose_vector)[:2] - if args.perf == 'model': - print(" - rotator", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - output_image = self.combiner(rotate_image[0], rotate_image[1], pose_vector) - if args.perf == 'model': - print(" - combiner", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - return output_image - - -class TalkingAnime3(nn.Module): - def __init__(self): - super(TalkingAnime3, self).__init__() - if args.model == "standard_float": - if args.eyebrow: - self.eyebrow_decomposer = tha3.poser.modes.standard_float.load_eyebrow_decomposer( - 'data/models/standard_float/eyebrow_decomposer.pt') - self.eyebrow_morphing_combiner = tha3.poser.modes.standard_float.load_eyebrow_morphing_combiner( - 'data/models/standard_float/eyebrow_morphing_combiner.pt') - self.face_morpher = tha3.poser.modes.standard_float.load_face_morpher( - 'data/models/standard_float/face_morpher.pt') - self.two_algo_face_body_rotator = tha3.poser.modes.standard_float.load_two_algo_generator( - 'data/models/standard_float/two_algo_face_body_rotator.pt') - self.editor = tha3.poser.modes.standard_float.load_editor('data/models/standard_float/editor.pt') - elif args.model == "standard_half": - if args.eyebrow: - self.eyebrow_decomposer = tha3.poser.modes.standard_half.load_eyebrow_decomposer( - 'data/models/standard_half/eyebrow_decomposer.pt') - self.eyebrow_morphing_combiner = tha3.poser.modes.standard_half.load_eyebrow_morphing_combiner( - 'data/models/standard_half/eyebrow_morphing_combiner.pt') - - self.face_morpher = tha3.poser.modes.standard_half.load_face_morpher( - 'data/models/standard_half/face_morpher.pt') - self.two_algo_face_body_rotator = tha3.poser.modes.standard_half.load_two_algo_generator( - 'data/models/standard_half/two_algo_face_body_rotator.pt') - self.editor = tha3.poser.modes.standard_half.load_editor('data/models/standard_half/editor.pt') - elif args.model == "separable_float": - if args.eyebrow: - self.eyebrow_decomposer = tha3.poser.modes.separable_float.load_eyebrow_decomposer( - 'data/models/separable_float/eyebrow_decomposer.pt') - self.eyebrow_morphing_combiner = tha3.poser.modes.separable_float.load_eyebrow_morphing_combiner( - 'data/models/separable_float/eyebrow_morphing_combiner.pt') - - self.face_morpher = tha3.poser.modes.separable_float.load_face_morpher( - 'data/models/separable_float/face_morpher.pt') - self.two_algo_face_body_rotator = tha3.poser.modes.separable_float.load_two_algo_generator( - 'data/models/separable_float/two_algo_face_body_rotator.pt') - self.editor = tha3.poser.modes.separable_float.load_editor('data/models/separable_float/editor.pt') - elif args.model == "separable_half": - if args.eyebrow: - self.eyebrow_decomposer = tha3.poser.modes.separable_half.load_eyebrow_decomposer( - 'data/models/separable_half/eyebrow_decomposer.pt') - self.eyebrow_morphing_combiner = tha3.poser.modes.separable_half.load_eyebrow_morphing_combiner( - 'data/models/separable_half/eyebrow_morphing_combiner.pt') - - self.face_morpher = tha3.poser.modes.separable_half.load_face_morpher( - 'data/models/separable_half/face_morpher.pt') - self.two_algo_face_body_rotator = tha3.poser.modes.separable_half.load_two_algo_generator( - 'data/models/separable_half/two_algo_face_body_rotator.pt') - self.editor = tha3.poser.modes.separable_half.load_editor('data/models/separable_half/editor.pt') - else: - raise RuntimeError("Invalid model: '%s'" % args.model) - self.face_cache = OrderedDict() - self.tot = 0 - self.hit = 0 - - def forward(self, image, mouth_eye_vector, pose_vector, eyebrow_vector, mouth_eye_vector_c, eyebrow_vector_c, - ratio=None): - if args.perf == 'model': - tic = time.perf_counter() - x = image.clone() - if args.eyebrow: - input_hash = hash(tuple(eyebrow_vector_c + mouth_eye_vector_c)) - else: - input_hash = hash(tuple(mouth_eye_vector_c)) - cached = self.face_cache.get(input_hash) - self.tot += 1 - if cached is None: - face_image = x[:, :, 32:32 + 192, (32 + 128):(32 + 192 + 128)].clone() - if args.eyebrow: - eyebrow_morp_image = self.eyebrow_decomposer(x[:, :, 64:192, 64 + 128:192 + 128].clone()) - eyebrow_morp_image = \ - self.eyebrow_morphing_combiner(eyebrow_morp_image[3], eyebrow_morp_image[0], eyebrow_vector)[2] - face_image[:, :, 32:32 + 128, 32:32 + 128] = eyebrow_morp_image - mouth_eye_morp_image = self.face_morpher(face_image, mouth_eye_vector)[0] - self.face_cache[input_hash] = mouth_eye_morp_image.detach() - if len(self.face_cache) > args.max_gpu_cache_len: - self.face_cache.popitem(last=False) - else: - self.hit += 1 - mouth_eye_morp_image = cached - self.face_cache.move_to_end(input_hash) - if args.debug and ratio is not None: - ratio.value = self.hit / self.tot - if args.perf == 'model': - print(" - face_morpher", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - x[:, :, 32:32 + 192, (32 + 128):(32 + 192 + 128)] = mouth_eye_morp_image - x_half = interpolate(x, size=(256, 256), mode='bilinear', align_corners=False) - rotate_image = self.two_algo_face_body_rotator(x_half, pose_vector) - if args.perf == 'model': - print(" - rotator", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - output_image = self.editor(x, - interpolate(rotate_image[1], size=(512, 512), mode='bilinear', align_corners=False), - interpolate(rotate_image[2], size=(512, 512), mode='bilinear', align_corners=False), - pose_vector)[0] - if args.perf == 'model': - print(" - editor", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - return output_image - - -class TalkingAnime(nn.Module): - def __init__(self): - super(TalkingAnime, self).__init__() - - def forward(self, image, mouth_eye_vector, pose_vector): - x = image.clone() - mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) - x[:, :, 32:224, 32:224] = mouth_eye_morp_image - rotate_image = self.two_algo_face_rotator(x, pose_vector)[:2] - output_image = self.combiner(rotate_image[0], rotate_image[1], pose_vector) - return output_image - - -class TalkingAnime4(nn.Module): - """ - THA4 model class, compatible with THA3 TalkingAnime3 interface - Uses THA4's Siren architecture for inference - PyTorch only, float32 precision - """ - def __init__(self, device): - super(TalkingAnime4, self).__init__() - if not THA4_AVAILABLE: - raise RuntimeError( - "THA4 adapter not available. Check tha4_adapter.py" - ) - - # Create THA4 wrapper - self.wrapper = THA4Wrapper(device=device) - - print("THA4 model initialized (float32)") - - def forward(self, image, mouth_eye_vector, pose_vector, eyebrow_vector, - mouth_eye_vector_c, eyebrow_vector_c, ratio=None): - """ - Forward inference, compatible with THA3 interface - - Args: - image: [batch, 4, 512, 512] - mouth_eye_vector: [batch, 27] - pose_vector: [batch, 6] - eyebrow_vector: [batch, 12] - mouth_eye_vector_c: compressed (for caching) - eyebrow_vector_c: compressed (for caching) - ratio: GPU cache hit ratio - - Returns: - output_image: [batch, 4, 512, 512] - """ - if args.perf == 'model': - tic = time.perf_counter() - - # Call THA4 wrapper - output_image = self.wrapper.forward( - image, mouth_eye_vector, pose_vector, eyebrow_vector, - mouth_eye_vector_c, eyebrow_vector_c, ratio - ) - - if args.perf == 'model': - print(" - tha4_inference", (time.perf_counter() - tic) * 1000) - - return output_image - - def to(self, device): - """Move model to specified device""" - self.wrapper.to(device) - return super().to(device) diff --git a/pack_release.bat b/pack_release.bat index 922dd7b6..a4cd9352 100644 --- a/pack_release.bat +++ b/pack_release.bat @@ -5,16 +5,37 @@ REM 1. Checkout to the release_package branch: git checkout release_package REM 2. Rebase with the latest changes: git rebase main REM ======================================================================== +REM ======================================================================== +REM Version Configuration +REM ======================================================================== +SET TENSORRT_VERSION=1.4.0.76 +SET CUDA_VERSION=12.9 +SET CUDA_VERSION_FULL=12.9.1 +SET PYTHON_VERSION=3.10 +SET MINICONDA_VERSION=Miniconda3-py310_24.9.2-0-Windows-x86_64 + +REM ======================================================================== +REM Download URLs +REM ======================================================================== +SET TENSORRT_DOWNLOAD_URL=https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.4/TensorRT-RTX-%TENSORRT_VERSION%-Windows-amd64-cuda-%CUDA_VERSION%-Release-external.zip +SET MINICONDA_DOWNLOAD_URL=https://repo.anaconda.com/miniconda/%MINICONDA_VERSION%.exe + +REM ======================================================================== +REM Derived Paths (auto-generated from versions) +REM ======================================================================== +SET TENSORRT_DIR_NAME=TensorRT-RTX-%TENSORRT_VERSION%_cu%CUDA_VERSION% +SET TENSORRT_WHL_NAME=tensorrt_rtx-%TENSORRT_VERSION%-cp310-none-win_amd64.whl + cd /D "%~dp0" if not exist "envs" mkdir envs if not exist "envs\miniconda3" mkdir envs\miniconda3 -IF not EXIST %~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin ( - @RD /S /Q %~dp0envs\TensorRT-RTX-1.3.0.35_cu129 +IF not EXIST %~dp0envs\%TENSORRT_DIR_NAME%\bin ( + @RD /S /Q %~dp0envs\%TENSORRT_DIR_NAME% cd /D "%~dp0\envs" - curl -L -o trt_rtx.zip https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.3/TensorRT-RTX-1.3.0.35-win10-amd64-cuda-12.9-Release-external.zip - tar -xf trt_rtx.zip "TensorRT-RTX-1.3.0.35" && ren "TensorRT-RTX-1.3.0.35" "TensorRT-RTX-1.3.0.35_cu129" && del trt_rtx.zip + curl -L -f -o trt_rtx.zip %TENSORRT_DOWNLOAD_URL% + tar -xf trt_rtx.zip "TensorRT-RTX-%TENSORRT_VERSION%" && ren "TensorRT-RTX-%TENSORRT_VERSION%" "TensorRT-RTX" && del trt_rtx.zip cd /D "%~dp0" ) @@ -22,32 +43,32 @@ IF not EXIST %~dp0envs\miniconda3\Scripts ( @RD /S /Q %~dp0envs\miniconda3 mkdir %~dp0envs\miniconda3 echo "Downloading miniconda..." - powershell -Command "Invoke-WebRequest -Uri 'https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe' -OutFile '.\envs\miniconda3.exe' -UseBasicParsing -Headers @{'User-Agent'='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36'}" + powershell -Command "Invoke-WebRequest -Uri '%MINICONDA_DOWNLOAD_URL%' -OutFile '.\envs\miniconda3.exe' -UseBasicParsing -Headers @{'User-Agent'='Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/123.0.0.0 Safari/537.36'}" - echo "Installling minconda..." + echo "Installing miniconda..." start /wait "" %~dp0envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=%~dp0envs\miniconda3 - echo "Successfully install minconda" + echo "Successfully installed miniconda, now cleaning up the installer..." del envs\miniconda3.exe ) -SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\miniconda3\Scripts;%PATH% +SET PATH=%~dp0envs\%TENSORRT_DIR_NAME%\bin;%~dp0envs\miniconda3\Scripts;%PATH% call activate call conda env list call conda update -y --all -call conda create -n ezvtb_rt_venv_release python=3.10 -y +call conda create -n ezvtb_rt_venv_release python=%PYTHON_VERSION% -y call conda activate ezvtb_rt_venv_release call conda env list call conda install -y conda-pack call conda install -y conda-forge::pycuda -call conda install -y -c nvidia/label/cuda-12.9.1 cuda-nvcc-dev_win-64 cudnn cuda-runtime +call conda install -y -c nvidia/label/cuda-%CUDA_VERSION_FULL% cuda-nvcc-dev_win-64 cudnn cuda-runtime call conda-pack -n ezvtb_rt_venv_release -o %~dp0envs\python_embedded --format no-archive SET PATH=%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% -call python -m pip install %~dp0envs\TensorRT-RTX-1.3.0.35_cu129\python\tensorrt_rtx-1.3.0.35-cp310-none-win_amd64.whl +call python -m pip install %~dp0envs\TensorRT-RTX\python\%TENSORRT_WHL_NAME% call python -m pip install -r requirements.txt --no-warn-script-location @RD /S /Q %~dp0envs\miniconda3 diff --git a/src/main.py b/src/main.py index 2545b5ec..6f36e74f 100644 --- a/src/main.py +++ b/src/main.py @@ -98,7 +98,6 @@ def main(): print("Using OpenCV windows for output display.") pipeline_fps = FPS() - last_frame_time = None # 上一帧输出时间，用于打印帧时间差 last_batch_start_time = None # 上一批就绪时间，用于周期估计 n_frames = args.interpolation_scale min_period = n_frames * interval if interval > 0 else n_frames / 60.0 # 60fps 下本批最少占用时间 @@ -135,7 +134,6 @@ def main(): cv2.imshow("EasyVtuber Debug Frame", np_ret_shms[i]) cv2.waitKey(1) now_send = time.perf_counter() - last_frame_time = now_send # 限速：下一帧最早在 last_time + interval，若已落后于当前时间则对齐到 now if interval > 0: last_time += interval