From daaa8522a09e59ac93394ba1fea5dde09d71978f Mon Sep 17 00:00:00 2001 From: nick-haoran <33177696+nick-haoran@users.noreply.github.com> Date: Tue, 18 Nov 2025 22:23:38 +0800 Subject: [PATCH] =?UTF-8?q?clear=20and=20add=20bat(TODOFIX:=20tha4=20stude?= =?UTF-8?q?nt=20model=20vram=20cache=200.0%=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- 03.构建tha4的ONNX和TensorRT学生模型.bat | 74 ++++++++++ PerformanceTest.md | 123 ---------------- THA4_SETUP.md | 65 -------- THA4_USAGE.md | 127 ---------------- args.py | 2 + ezvtb_rt_interface.py | 96 ++++++++---- ezvtuber-rt | 2 +- launcher2.py | 102 ++++++++++++- main.py | 2 + onnx_export_tha4_student.py | 21 +-- test_tha4_student.py | 187 ++++++++++++++++++++++++ 11 files changed, 440 insertions(+), 361 deletions(-) create mode 100644 03.构建tha4的ONNX和TensorRT学生模型.bat delete mode 100644 PerformanceTest.md delete mode 100644 THA4_SETUP.md delete mode 100644 THA4_USAGE.md create mode 100644 test_tha4_student.py diff --git a/03.构建tha4的ONNX和TensorRT学生模型.bat b/03.构建tha4的ONNX和TensorRT学生模型.bat new file mode 100644 index 0000000..f7261fa --- /dev/null +++ b/03.构建tha4的ONNX和TensorRT学生模型.bat @@ -0,0 +1,74 @@ +@echo off +setlocal enabledelayedexpansion +cd /D "%~dp0" + +SET PATH=%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% + +echo ================================================================================ +echo THA4 ѧģ ONNX TensorRT ת +echo ================================================================================ +echo. + +REM ģĿ¼Ƿ +if not exist "data\models\custom_tha4_models" ( + echo [] ҲģĿ¼: data\models\custom_tha4_models + pause + exit /b 1 +) + +echo õ THA4 ѧģ: +echo -------------------------------------------------------------------------------- +set count=0 +for /d %%i in (data\models\custom_tha4_models\*) do ( + set /a count+=1 + echo [!count!] %%~nxi +) +echo -------------------------------------------------------------------------------- +echo. + +if %count%==0 ( + echo [] δҵκģļ + pause + exit /b 1 +) + +echo Ҫתģļ: +set /p MODEL_NAME=^> + +REM ǷΪ +if "%MODEL_NAME%"=="" ( + echo [] ģƲΪ + pause + exit /b 1 +) + +REM ģļǷ +if not exist "data\models\custom_tha4_models\%MODEL_NAME%" ( + echo [] ģļв: data\models\custom_tha4_models\%MODEL_NAME% + pause + exit /b 1 +) + +echo. +echo ================================================================================ +echo 1/2: ONNX ģ +echo ================================================================================ +echo ģ: %MODEL_NAME% +echo : python onnx_export_tha4_student.py %MODEL_NAME% --cpu +echo -------------------------------------------------------------------------------- +echo. + +python onnx_export_tha4_student.py %MODEL_NAME% --cpu + +echo ================================================================================ +echo 2/2: תΪ TensorRT ģ +echo ================================================================================ +echo ģ: %MODEL_NAME% +echo : python tha4_student_convert_onnx_to_trt.py %MODEL_NAME% fp32 +echo -------------------------------------------------------------------------------- +echo. + +python tha4_student_convert_onnx_to_trt.py %MODEL_NAME% fp32 + +pause + diff --git a/PerformanceTest.md b/PerformanceTest.md deleted file mode 100644 index 1d048ea..0000000 --- a/PerformanceTest.md +++ /dev/null @@ -1,123 +0,0 @@ -# 测试环境 -### 基于 -* Release版本v0.3 -* 测试平台为windows10 22H2 -* CPU Intel 10400f -* GPU Nvidia 3060 12G, Nvidia 1660 Super 6G, ARC A770 16G，RX580 8G. - -### 测试定量 -* Jpeg compression 和 Input Simplify 设为 medium -* 使用debug input 和 debug output -* 不使用eyebrow参数 -* 帧数限制60fps -* 英伟达显卡使用TensorRT测试 - -观察开始后约5-10秒的数值记录, 即不考虑长期缓存效果(没时间测) - -# 测试结果 -请注意测试结果仅限参考显卡大概性能，具体性能取决于用户周边配置和显卡驱动。 - -GPU | model | VRAM Cache | RAM Cache | Interpolation | SR | Frame Rate | GPU Usage -:---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: -A770 | sepe_half | N/A | No | No | No | 33fps | 97% -A770 | sepe_full | N/A | No | No | No | 48fps | 97% -A770 | stan_half | N/A | No | No | No | 27fps | 97% -A770 | stan_full | N/A | No | No | No | 29fps | 97% -A770 | sepe_half | N/A | 2GB | No | No | 30fps | 82% -A770 | sepe_full | N/A | 2GB | No | No | 43fps | 78% -A770 | stan_half | N/A | 2GB | No | No | 30fps | 82% -A770 | stan_full | N/A | 2GB | No | No | 23fps | 85% -A770 | sepe_full | N/A | No | No | x2_half | 17fps | 97% -A770 | sepe_full | N/A | No | No | x4_half | 29fps | 97% -A770 | sepe_full | N/A | No | No | x4_full | 25fps | 93% -A770 | sepe_full | N/A | No | x2_half | No | 60fps | 93% -A770 | sepe_full | N/A | No | x3_half | No | 60fps | 85% -A770 | sepe_full | N/A | No | x4_half | No | 60fps | 81% --- | -- | --| --|--|--|--|-- -RX580 | sepe_full | N/A | No | No | No | 12fps | 98% -RX580 | stan_full | N/A | No | No | No | 15fps | 98% -RX580 | sepe_full | N/A | 2GB | No | No | 11fps | 93% -RX580 | stan_full | N/A | 2GB | No | No | 13fps | 93% -RX580 | stan_full | N/A | No | No | x2_half | 10fps | 97% -RX580 | stan_full | N/A | No | No | x2_full | 10fps | 97% -RX580 | stan_full | N/A | No | No | x4_half | 11fps | 98% -RX580 | stan_full | N/A | No | No | x4_full | 11fps | 98% -RX580 | stan_full | N/A | No | x2_half | no | 22fps | 98% -RX580 | stan_full | N/A | No | x2_full | no | 22fps | 98% -RX580 | stan_full | N/A | No | x3_half | no | 28fps | 98% -RX580 | stan_full | N/A | No | x3_full | no | 28fps | 98% -RX580 | stan_full | N/A | No | x4_half | no | 33fps | 98% -RX580 | stan_full | N/A | No | x4_full | no | 33fps | 98% --- | -- | --| --|--|--|--|-- -1660s | sepe_half | No | No | No | No | 38fps | 95% -1660s | sepe_full | No | No | No | No | 28fps | 95% -1660s | stan_half | No | No | No | No | 21fps | 96% -1660s | stan_full | No | No | No | No | 12fps | 98% -1660s | sepe_half | 2GB | No | No | No | 39fps | 95% -1660s | sepe_full | 2GB | No | No | No | 28fps | 95% -1660s | stan_half | 2GB | No | No | No | 29fps | 96% -1660s | stan_full | 2GB | No | No | No | 12fps | 96% -1660s | sepe_half | 2GB | 2GB | No | No | 40fps | 75% -1660s | sepe_full | 2GB | 2GB | No | No | 24fps | 80% -1660s | stan_half | 2GB | 2GB | No | No | 26fps | 83% -1660s | stan_full | 2GB | 2GB | No | No | 12fps | 92% -1660s | sepe_half | No | No | No | x2_half | 23fps | 96% -1660s | sepe_half | No | No | No | x2_full | 18fps | 96% -1660s | sepe_half | No | No | No | x4_half | 23fps | 96% -1660s | sepe_half | No | No | No | x4_full | 26fps | 97% -1660s | sepe_half | No | No | x2_half | No | 63fps | 75% -1660s | sepe_half | No | No | x2_full | No | 63fps | 80% -1660s | sepe_half | No | No | x3_half | No | 63fps | 65% -1660s | sepe_half | No | No | x3_full | No | 64fps | 66% -1660s | sepe_half | No | No | x4_half | No | 64fps | 53% -1660s | sepe_half | No | No | x4_full | No | 64fps | 60% --- | -- | --| --|--|--|--|-- -3060 | sepe_half | No | No | No | No | 63fps | 62% -3060 | sepe_full | No | No | No | No | 59fps | 93% -3060 | stan_half | No | No | No | No | 63fps | 80% -3060 | stan_full | No | No | No | No | 31fps | 96% -3060 | sepe_half | 2GB | No | No | No | 63fps | 60% -3060 | sepe_full | 2GB | No | No | No | 60fps | 93% -3060 | stan_half | 2GB | No | No | No | 63fps | 76% -3060 | stan_full | 2GB | No | No | No | 33fps | 96% -3060 | sepe_half | 2GB | 2GB | No | No | 63fps | 58% -3060 | sepe_full | 2GB | 2GB | No | No | 46fps | 70% -3060 | stan_half | 2GB | 2GB | No | No | 43fps | 62% -3060 | stan_full | 2GB | 2GB | No | No | 27fps | 80% -3060 | sepe_half | No | No | No | x2_half | 53fps | 85% -3060 | sepe_half | No | No | No | x2_full | 30fps | 94% -3060 | sepe_half | No | No | No | x4_half | 49fps | 90% -3060 | sepe_half | No | No | No | x4_full | 45fps | 90% -3060 | sepe_half | No | No | x2_half | No | 64fps | 45% -3060 | sepe_half | No | No | x2_full | No | 63fps | 50% -3060 | sepe_half | No | No | x3_half | No | 64fps | 35% -3060 | sepe_half | No | No | x3_full | No | 64fps | 42% -3060 | sepe_half | No | No | x4_half | No | 64fps | 30% -3060 | sepe_half | No | No | x4_full | No | 64fps | 40% - -# 已知问题: -* A770: rife插帧全精度会出现奇怪结果, waifu2x 超分的全精度会输出全绿色 -* RX580 tha3 模型半精度会输出奇怪结果 -* 3060 有时出现奇怪运行时占用上升，猜测为插拔显卡导致的驱动问题。 - -# DirectML 思考 -DirectML 允许所有 Windows 显卡设备以及 AIPC 的 NPU 运行此项目，但由于不同设备的 Kernel 在不同精度下实现区别和驱动兼容问题， -导致 DirectML 的运行结果在某个模型的某些精度出现不可用的扭曲结果，且运行性能与对模型的常识理解不相符。 - -例如适用于本项目插帧的RIFE模型的调用封装经历过一轮迭代, -由原先的n次输入(N=1,C,H,W)独立图像调用RIFE模型生成n个独立结果， -改为输入(N=n,C,H,W)的叠加图像一次性调用RIFE模型得到叠加结果， -此改动下模型获得了稍小的显存占用且在TensorRT实现中获得一定性能提升（虽然改动本意并不在此）。 -接着在笔者环境中的因特尔 A770 显卡 DirectML 在RIFE模型迭代后全精度推理结果由原先的符合预期变为扭曲图像， -半精度迭代后模型不仅保持推理结果符合预期，且相较模型迭代前得到了一定性能提升。 - -再例如笔者测试某 AMD 核显运行时发现其运行THA3的半精度模型时会因为此设备无法运行其中的某个kernel而将kernel运算单独卸载到cpu进行， -但保留了其他大部分能在核显上运行的kernel，由此推理模型需要结果在两个设备间反复，带来了极大的性能倒退。 - -目前阶段，受限于Python OnnxRuntime api 的支持匮乏（相较C++），使用 python 运行 Directml 无法避免上述问题， -在代码层面无法在运行前做出是否适配的条件判断，也无法在出现kernel卸载到cpu性能剧降时raise错误将运行打断。 - -所以，DirectML运行本项目的任意模型出现奇怪结果，或不符合预期的性能倒退，都应该属于意料之中 -换个模型精度，观察调试输出是否有DirectML黄字警告 通常可以解决大部分问题。 -具体某个硬件在DirectML的运行异常修复，已远超本项目的目标范围，请自行验证失效kernel，向硬件厂商和微软求助并期待驱动更新带来解决 -（建议买张N卡吧别折腾自己了）。 diff --git a/THA4_SETUP.md b/THA4_SETUP.md deleted file mode 100644 index a569f99..0000000 --- a/THA4_SETUP.md +++ /dev/null @@ -1,65 +0,0 @@ -# THA4 Model Setup Guide - -## Required Files Structure - -Place your THA4 model files in `data/models/tha4/` with the following structure: - -``` -data/models/tha4/ -├── character_model.yaml # Configuration file (paths to other files) -├── character.png # Character source image (RGBA, 512x512) -├── face_morpher.pt # Face morphing model weights -└── body_morpher.pt # Body morphing model weights -``` - -## Example character_model.yaml - -Create a file named `character_model.yaml` with the following content: - -```yaml -character_image_file_name: character.png -face_morpher_file_name: face_morpher.pt -body_morpher_file_name: body_morpher.pt -``` - -**Note:** All paths in the YAML are relative to the directory containing the YAML file. - -## Alternative Structure - -You can also organize files in subdirectories: - -``` -data/models/tha4/ -├── my_character/ -│ ├── character_model.yaml -│ ├── image.png -│ ├── face_morpher.pt -│ └── body_morpher.pt -``` - -Then update your YAML accordingly: - -```yaml -character_image_file_name: image.png -face_morpher_file_name: face_morpher.pt -body_morpher_file_name: body_morpher.pt -``` - -## Usage in EasyVtuber - -1. Place your files as shown above -2. Launch the EasyVtuber launcher -3. Select "THA4" from the "Model Select" dropdown -4. The system will automatically: - - Load `character_model.yaml` - - Read the PNG character image - - Load both PT model files - - Use the PNG image as the source for animation - -## Differences from THA3 - -- **THA4**: Uses YAML to define character (image + models as a package) -- **THA3**: Uses separate image file selected by user + shared model files - -The THA4 approach bundles character and model together, making it easier to -switch between different characters with their trained models. diff --git a/THA4_USAGE.md b/THA4_USAGE.md deleted file mode 100644 index f3f42a2..0000000 --- a/THA4_USAGE.md +++ /dev/null @@ -1,127 +0,0 @@ -# THA4 模型使用指南 (新版本) - -## 文件结构 - -THA4 模型现在支持直接放在 `data/images/` 目录中，与 THA3 的 PNG 图片一起管理。 - -### 所需文件 - -对于名为 `mycharacter` 的角色，需要在 `data/images/` 中放置以下文件： - -``` -data/images/ -├── mycharacter.yaml # 配置文件 -├── mycharacter.png # 角色源图片 (512x512 RGBA) -├── mycharacter_face_morpher.pt # 面部变形模型权重 -└── mycharacter_body_morpher.pt # 身体变形模型权重 -``` - -### YAML 配置格式 - -YAML 文件 (`mycharacter.yaml`) 应包含相对路径指向其他文件： - -```yaml -character_image_file_name: mycharacter.png -face_morpher_file_name: mycharacter_face_morpher.pt -body_morpher_file_name: mycharacter_body_morpher.pt -``` - -**重要说明：** -- YAML 中的所有路径都是相对于 YAML 文件所在位置的相对路径 -- 由于所有文件都在 `data/images/` 中，只需使用文件名即可 -- 角色图片必须是 512x512 像素的 RGBA 格式 -- 模型文件 (`.pt`) 是从 THA4 训练得到的 PyTorch 权重 - -## 在 EasyVtuber 中使用 - -### 自动检测 - -1. 将所有 4 个文件放入 `data/images/` 目录 -2. 打开启动器 (`02A.启动器.bat`) -3. 在 Character 下拉菜单中选择 `.yaml` 文件（例如 `mycharacter`） -4. 程序将自动： - - 检测这是 THA4 模型 - - 切换到 THA4 推理模式 - - 从 YAML 配置加载角色图片和模型 - -### 模型选择锁定 - -当你选择 YAML 文件时： -- 模型**自动锁定到 THA4** -- 使用 YAML 中的角色图片（启动器中的 Character 选择仅用于选择 YAML） -- 手动的模型选择将被覆盖 - -## 示例设置 - -假设你为名为 "alice" 的角色训练了 THA4 模型： - -1. 将这些文件复制到 `data/images/`： - ``` - alice.yaml - alice.png - alice_face_morpher.pt - alice_body_morpher.pt - ``` - -2. 创建 `alice.yaml`，内容为： - ```yaml - character_image_file_name: alice.png - face_morpher_file_name: alice_face_morpher.pt - body_morpher_file_name: alice_body_morpher.pt - ``` - -3. 在启动器中，从 Character 下拉菜单选择 "alice" -4. 启动 - 程序将检测到 YAML 并自动使用 THA4！ - -## 与 THA3 的区别 - -| 特性 | THA3 | THA4 | -|------|------|------| -| 角色选择 | data/images/ 中的任意 PNG | 选择 YAML 文件 | -| 模型文件 | data/models/tha3/ 中的共享模型 | 每个角色独立的 .pt 文件 | -| 图片加载 | 运行时选择 | 与训练模型绑定 | -| 模型类型 | Standard/Separate, Half/Full | 单一类型 (float32) | -| 可互换性 | 可使用任意图片 | 图片与训练模型绑定 | - -## 文件命名建议 - -为了保持清晰，建议使用一致的命名规则： - -``` -<角色名>.yaml # 配置文件 -<角色名>.png # 源图片 -<角色名>_face_morpher.pt # 面部模型 -<角色名>_body_morpher.pt # 身体模型 -``` - -## 故障排除 - -**错误："THA4 character model YAML not found"** -- 检查 `.yaml` 文件是否存在于 `data/images/` 中 -- 验证文件名是否与你选择的一致 - -**错误："No such file or directory" 找不到 PNG 或 PT 文件** -- 检查 YAML 中的路径是否正确 -- 确保所有 4 个文件都在同一个 `data/images/` 目录中 -- 验证文件名完全匹配（区分大小写） - -**黑屏 / 无输出** -- 验证角色 PNG 是 512x512 RGBA 格式 -- 检查 .pt 模型文件是否来自正确训练的 THA4 模型 -- 确保模型文件与用于训练的角色图片匹配 - -## 旧文件迁移 - -如果你之前将文件放在 `data/models/tha4/` 中： - -1. 将所有文件移动到 `data/images/`： - ``` - mv data/models/tha4/character_model.yaml data/images/mycharacter.yaml - mv data/models/tha4/character.png data/images/mycharacter.png - mv data/models/tha4/face_morpher.pt data/images/mycharacter_face_morpher.pt - mv data/models/tha4/body_morpher.pt data/images/mycharacter_body_morpher.pt - ``` - -2. 更新 YAML 文件中的路径为新的文件名 - -3. 在启动器中选择新的 YAML 文件名 diff --git a/args.py b/args.py index 28032b1..ab6ae1e 100644 --- a/args.py +++ b/args.py @@ -38,6 +38,8 @@ 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('--use_tha4_student', action='store_true') +parser.add_argument('--tha4_student_model', type=str, default=None) parser.add_argument('--model_seperable', action='store_true') parser.add_argument('--model_half', action='store_true') diff --git a/ezvtb_rt_interface.py b/ezvtb_rt_interface.py index cf7dc0b..25d2aa8 100644 --- a/ezvtb_rt_interface.py +++ b/ezvtb_rt_interface.py @@ -44,7 +44,9 @@ def get_core( sr_half:bool = True, sr_noise:int = 1, #THA4 model setting (NEW) - use_tha4:bool = False # Use THA4 instead of THA3 + use_tha4: bool = False, # Use THA4 instead of THA3 + use_tha4_student: bool = False, # Use THA4 Student model + tha4_student_model: str = None # Student model name ): support_trt = False if use_tensorrt: @@ -65,7 +67,14 @@ def get_core( print(f'Using device {pycuda.autoinit.device.name()}') # Determine model directory based on THA version - if use_tha4: + if use_tha4_student and tha4_student_model: + # THA4 Student: Use data/models/custom_tha4_models/{model_name} + tha_model_dir = os.path.join( + '.', 'data', 'models', 'custom_tha4_models', + tha4_student_model) + print(f'THA4 Student Mode - Model: {tha4_student_model}') + print(f'THA4 Student Mode - Model Path: {tha_model_dir}') + elif use_tha4: # THA4: Use data/models/tha4 directory with precision selection tha_model_dir = os.path.join( '.', 'data', 'models', 'tha4', @@ -115,35 +124,64 @@ def get_core( core = None if use_tensorrt: from ezvtb_rt.core import CoreTRT - core = CoreTRT( - tha_model_version='v4' if use_tha4 else 'v3', - tha_model_seperable=model_seperable, - tha_model_fp16=model_half, - rife_model_enable=use_interpolation, - rife_model_scale=interpolation_scale, - rife_model_fp16=interpolation_half, - sr_model_enable=use_sr, - sr_model_scale=4 if sr_x4 else 2, - sr_model_noise=sr_noise, - sr_model_fp16=sr_half, - vram_cache_size=model_cache_size, - cache_max_giga=cacher_ram_size, - use_eyebrow=model_use_eyebrow) + if use_tha4_student and tha4_student_model: + core = CoreTRT( + tha_model_version='v4_student', + tha_model_name=tha4_student_model, + rife_model_enable=use_interpolation, + rife_model_scale=interpolation_scale, + rife_model_fp16=interpolation_half, + sr_model_enable=use_sr, + sr_model_scale=4 if sr_x4 else 2, + sr_model_noise=sr_noise, + sr_model_fp16=sr_half, + vram_cache_size=model_cache_size, + cache_max_giga=cacher_ram_size, + use_eyebrow=model_use_eyebrow) + else: + core = CoreTRT( + tha_model_version='v4' if use_tha4 else 'v3', + tha_model_seperable=model_seperable, + tha_model_fp16=model_half, + rife_model_enable=use_interpolation, + rife_model_scale=interpolation_scale, + rife_model_fp16=interpolation_half, + sr_model_enable=use_sr, + sr_model_scale=4 if sr_x4 else 2, + sr_model_noise=sr_noise, + sr_model_fp16=sr_half, + vram_cache_size=model_cache_size, + cache_max_giga=cacher_ram_size, + use_eyebrow=model_use_eyebrow) else: from ezvtb_rt.core_ort import CoreORT - core = CoreORT( - tha_model_version='v4' if use_tha4 else 'v3', - tha_model_seperable=model_seperable, - tha_model_fp16=model_half, - rife_model_enable=use_interpolation, - rife_model_scale=interpolation_scale, - rife_model_fp16=interpolation_half, - sr_model_enable=use_sr, - sr_model_scale=4 if sr_x4 else 2, - sr_model_noise=sr_noise, - sr_model_fp16=sr_half, - cache_max_giga=cacher_ram_size, - use_eyebrow=model_use_eyebrow) + if use_tha4_student and tha4_student_model: + core = CoreORT( + tha_model_version='v4_student', + tha_model_name=tha4_student_model, + rife_model_enable=use_interpolation, + rife_model_scale=interpolation_scale, + rife_model_fp16=interpolation_half, + sr_model_enable=use_sr, + sr_model_scale=4 if sr_x4 else 2, + sr_model_noise=sr_noise, + sr_model_fp16=sr_half, + cache_max_giga=cacher_ram_size, + use_eyebrow=model_use_eyebrow) + else: + core = CoreORT( + tha_model_version='v4' if use_tha4 else 'v3', + tha_model_seperable=model_seperable, + tha_model_fp16=model_half, + rife_model_enable=use_interpolation, + rife_model_scale=interpolation_scale, + rife_model_fp16=interpolation_half, + sr_model_enable=use_sr, + sr_model_scale=4 if sr_x4 else 2, + sr_model_noise=sr_noise, + sr_model_fp16=sr_half, + cache_max_giga=cacher_ram_size, + use_eyebrow=model_use_eyebrow) return core diff --git a/ezvtuber-rt b/ezvtuber-rt index d23cfb3..b59e1ba 160000 --- a/ezvtuber-rt +++ b/ezvtuber-rt @@ -1 +1 @@ -Subproject commit d23cfb3a3649193e88042df1079ecb33a4cf805a +Subproject commit b59e1bab05b5c007a09690f83bf656d9c77a420d diff --git a/launcher2.py b/launcher2.py index 7f1505a..5715feb 100644 --- a/launcher2.py +++ b/launcher2.py @@ -64,6 +64,8 @@ finally: p = None dirPath = 'data/images' characterList = [] +studentModelList = [] +studentModelCharacterMap = {} hasRTModel = False try: @@ -82,7 +84,39 @@ def refreshList(): characterList.append(item[:-4]) +def scanStudentModels(): + """Scan custom_tha4_models directory for student models""" + global studentModelList, studentModelCharacterMap + studentModelList = [] + studentModelCharacterMap = {} + + custom_models_path = 'data/models/custom_tha4_models' + if os.path.exists(custom_models_path): + try: + for model_name in os.listdir(custom_models_path): + model_path = os.path.join(custom_models_path, model_name) + if os.path.isdir(model_path): + # Check if it's a valid student model + face_trt = os.path.join(model_path, 'face_morpher.trt') + body_trt = os.path.join(model_path, 'body_morpher.trt') + character_png = os.path.join(model_path, + 'character.png') + has_trt = (os.path.exists(face_trt) and + os.path.exists(body_trt)) + has_character = os.path.exists(character_png) + + if has_trt and has_character: + studentModelList.append(model_name) + studentModelCharacterMap[model_name] = model_name + except Exception: + pass + + # Sort alphabetically + studentModelList.sort() + + refreshList() +scanStudentModels() class OptionPanel(wx.Panel): @@ -227,9 +261,23 @@ class LauncherPanel(wx.Panel): choices=['10', '15', '20', '30', '60']) addOption('preset', title='Performance Preset', desc='性能预设，注意修改后会覆盖后续配置', choices=['Low', 'Medium', 'High', 'Ultra', 'Custom']) - addOption('model_select', title='Model Select', desc='选择使用的模型\nStandard Full精度较高性能较低', - choices=['Seperable Half', 'Seperable Full', 'Standard Half', 'Standard Full', 'THA4 Half', 'THA4 Full'], - mapping=['seperable_half', 'seperable_full', 'standard_half', 'standard_full', 'tha4_half', 'tha4_full']) + + # Build model_select choices + model_choices = ['Seperable Half', 'Seperable Full', 'Standard Half', + 'Standard Full', 'THA4 Half', 'THA4 Full'] + model_mapping = ['seperable_half', 'seperable_full', + 'standard_half', 'standard_full', 'tha4_half', + 'tha4_full'] + + # Add student models if available + for student_model in studentModelList: + model_choices.append(f'THA4 Student ({student_model})') + model_mapping.append(f'tha4_student_{student_model}') + + addOption('model_select', title='Model Select', + desc='选择使用的模型\nStandard Full精度较高性能较低', + choices=model_choices, + mapping=model_mapping) addOption('ram_cache_size', title='RAM Cache Size', desc='分配内存缓存大小\n用于存储最终运算结果', choices=['Off', '1GB', '2GB', '4GB', '8GB', '16GB'], mapping=['0b', '1gb', '2gb', '4gb', '8gb', '16gb']) @@ -307,17 +355,50 @@ class LauncherPanel(wx.Panel): self.optionDict['preset'].Bind(wx.EVT_CHOICE, presetChoice) presetChoice() + def onModelSelect(e=None): + """Handle model selection change""" + model_value = self.optionDict['model_select'].GetValue() + is_student_model = 'tha4_student_' in model_value + + char_ctrl = self.optionDict['character'] + + if is_student_model: + # Disable character selection for student models + # Student models have their own built-in character + char_ctrl.control.Enable(False) + char_ctrl.control.SetToolTip( + 'Locked: Student model includes built-in character') + else: + # Re-enable character selection for non-student models + char_ctrl.control.Enable(True) + char_ctrl.control.SetToolTip( + 'Select a character from data/images') + + self.optionDict['model_select'].Bind( + wx.EVT_CHOICE, onModelSelect) + + # Check initial model selection and lock character if needed + onModelSelect() + def onActivate(e): global characterList tName = self.optionDict['character'].GetValue() refreshList() - self.optionDict['character'].control.SetItems(characterList) - self.optionDict['character'].control.SetSelection(characterList.index(tName)) + self.optionDict['character'].control.SetItems( + characterList) + try: + idx = characterList.index(tName) + self.optionDict['character'].control.SetSelection(idx) + except ValueError: + # Character not found, select first available + if characterList: + self.optionDict['character'].control.SetSelection(0) if not hasRTModel: self.optionDict['use_tensorrt'].control.SetValue(False) self.optionDict['use_tensorrt'].control.Enable(False) - self.optionDict['use_tensorrt'].control.SetToolTip('需要先构建TensorRT模型') + self.optionDict['use_tensorrt'].control.SetToolTip( + '需要先构建TensorRT模型') self.frame.Bind(wx.EVT_ACTIVATE, onActivate) @@ -411,7 +492,14 @@ class LauncherPanel(wx.Panel): run_args.append('4') if args['model_select'] is not None: - if 'tha4' in args['model_select']: + if 'tha4_student_' in args['model_select']: + # Student model: tha4_student_{model_name} + model_name = args['model_select'].replace( + 'tha4_student_', '') + run_args.append('--use_tha4_student') + run_args.append('--tha4_student_model') + run_args.append(model_name) + elif 'tha4' in args['model_select']: run_args.append('--use_tha4') if 'half' in args['model_select']: run_args.append('--model_half') diff --git a/main.py b/main.py index 51b84ca..f4d0f9b 100644 --- a/main.py +++ b/main.py @@ -350,6 +350,8 @@ class ModelClientProcess(Process): self.model = get_core(device_id=args.device_id, use_tensorrt=args.use_tensorrt, use_tha4=args.use_tha4, + use_tha4_student=args.use_tha4_student, + tha4_student_model=args.tha4_student_model, model_seperable = args.model_seperable, model_half=args.model_half, diff --git a/onnx_export_tha4_student.py b/onnx_export_tha4_student.py index 5d3366b..1fe4436 100644 --- a/onnx_export_tha4_student.py +++ b/onnx_export_tha4_student.py @@ -73,13 +73,17 @@ elif not torch.cuda.is_available(): else: DEVICE_NAME = 'cuda:0' -IMAGE_INPUT = os.path.join('.', 'data', 'images', 'character.png') +# Use image from model folder +MODEL_BASE_PATH = join('.', 'data', 'models', 'custom_tha4_models', MODEL_NAME) +IMAGE_INPUT = os.path.join(MODEL_BASE_PATH, 'character.png') -# Create output directories -TMP_DIR = join('.', 'onnx_model_tha4_student', MODEL_NAME, 'tmp') +if not os.path.exists(IMAGE_INPUT): + raise FileNotFoundError(f"Image not found: {IMAGE_INPUT}") + +# Create output directories (save models in the model folder) +TMP_DIR = join(MODEL_BASE_PATH, 'tmp') Path(TMP_DIR).mkdir(parents=True, exist_ok=True) -MODEL_DIR = join('.', 'onnx_model_tha4_student', MODEL_NAME, 'fp32' if not HALF else 'fp16') -Path(MODEL_DIR).mkdir(parents=True, exist_ok=True) +MODEL_DIR = MODEL_BASE_PATH # Save directly in model folder TMP_FILE_WRITE = join(TMP_DIR, 'tmp.onnx') validation_device = torch.device(DEVICE_NAME) @@ -107,11 +111,10 @@ def numpy_linear_to_srgb(x): return np.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) # Load THA4 Student Model (mode_14: 2-step architecture) -def load_tha4_student_poser(device: torch.device, use_half: bool, model_name: str): +def load_tha4_student_poser(device: torch.device, use_half: bool, model_base_path: str): from tha4.poser.modes.mode_14 import create_poser # Model file paths - model_base_path = join('.', 'data', 'models', 'custom_tha4_models', model_name) m_files = { 'face_morpher': join(model_base_path, 'face_morpher.pt'), 'body_morpher': join(model_base_path, 'body_morpher.pt') @@ -130,14 +133,14 @@ def load_tha4_student_poser(device: torch.device, use_half: bool, model_name: st return poser # Load posers for validation and export -validation_poser = load_tha4_student_poser(validation_device, HALF and DEVICE_NAME != 'cpu', MODEL_NAME) +validation_poser = load_tha4_student_poser(validation_device, HALF and DEVICE_NAME != 'cpu', MODEL_BASE_PATH) if DEVICE_NAME == 'cpu': # If using CPU for validation, reuse the same poser for export export_poser = validation_poser print("Using same poser for validation and export (CPU mode)") else: # If using GPU for validation, load separate CPU poser for export - export_poser = load_tha4_student_poser(export_device, False, MODEL_NAME) + export_poser = load_tha4_student_poser(export_device, False, MODEL_BASE_PATH) pose_size = validation_poser.get_num_parameters() print(f"Pose parameters: {pose_size}") diff --git a/test_tha4_student.py b/test_tha4_student.py new file mode 100644 index 0000000..f353120 --- /dev/null +++ b/test_tha4_student.py @@ -0,0 +1,187 @@ +#!/usr/bin/env python3 +"""Test script for THA4 Student model inference +Tests both TensorRT and ONNX Runtime backends +""" +import sys +import os +import numpy as np + +# Add ezvtuber-rt to path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'ezvtuber-rt')) + + +def test_tensorrt_backend(): + """Test TensorRT backend""" + print("\n=== Testing TensorRT Backend ===") + try: + from ezvtb_rt import CoreTRT + + # Initialize with student model (kanori_tha4) + core = CoreTRT( + tha_model_version='v4_student', + tha_model_name='kanori_tha4' + ) + + # Create dummy image (512x512 RGBA) + test_image = np.random.randint(0, 256, (512, 512, 4), dtype=np.uint8) + + # Create dummy pose (1x45) + test_pose = np.random.randn(1, 45).astype(np.float32) + + # Run inference + core.setImage(test_image) + result = core.inference(test_pose) + + print("PASS: TensorRT Backend: SUCCESS") + print(f" Input image shape: {test_image.shape}") + print(f" Input pose shape: {test_pose.shape}") + out_shape = (result[0].shape if isinstance(result, list) + else result.shape) + print(f" Output shape: {out_shape}") + return True + + except Exception as e: + print("FAIL: TensorRT Backend: FAILED") + print(f" Error: {str(e)}") + import traceback + traceback.print_exc() + return False + + +def test_onnxruntime_backend(): + """Test ONNX Runtime backend""" + print("\n=== Testing ONNX Runtime Backend ===") + try: + from ezvtb_rt import CoreORT + + # Initialize with student model (kanori_tha4) + core = CoreORT( + tha_model_version='v4_student', + tha_model_name='kanori_tha4' + ) + + # Create dummy image (512x512 RGBA) + test_image = np.random.randint(0, 256, (512, 512, 4), dtype=np.uint8) + + # Create dummy pose (1x45) + test_pose = np.random.randn(1, 45).astype(np.float32) + + # Run inference + core.setImage(test_image) + result = core.inference(test_pose) + + print("PASS: ONNX Runtime Backend: SUCCESS") + print(f" Input image shape: {test_image.shape}") + print(f" Input pose shape: {test_pose.shape}") + out_shape = (result[0].shape if isinstance(result, list) + else result.shape) + print(f" Output shape: {out_shape}") + return True + + except Exception as e: + print("FAIL: ONNX Runtime Backend: FAILED") + print(f" Error: {str(e)}") + import traceback + traceback.print_exc() + return False + + +def test_direct_models(): + """Test direct model classes""" + print("\n=== Testing Direct Model Classes ===") + + success = True + + # Test THA4Student (TensorRT) + try: + from ezvtb_rt.tha4_student import THA4Student + model_dir = os.path.join( + os.path.dirname(__file__), + 'data', 'models', 'custom_tha4_models', 'kanori_tha4' + ) + + if os.path.exists(model_dir): + model = THA4Student(model_dir) + test_image = np.random.randint( + 0, 256, (512, 512, 4), dtype=np.uint8 + ) + test_pose = np.random.randn(1, 45).astype(np.float32) + + model.setImage(test_image) + result = model.inference(test_pose) + + print("PASS: THA4Student (TensorRT): SUCCESS") + out_shape = result[0].shape if isinstance(result, list) else result.shape + print(f" Output shape: {out_shape}") + else: + print(f"SKIP: THA4Student: Model not found at {model_dir}") + + except Exception as e: + print("FAIL: THA4Student (TensorRT): FAILED") + print(f" Error: {str(e)}") + success = False + + # Test THA4StudentORT (ONNX Runtime) + try: + from ezvtb_rt.tha4_student_ort import THA4StudentORT + model_dir = os.path.join( + os.path.dirname(__file__), + 'data', 'models', 'custom_tha4_models', 'kanori_tha4' + ) + + if os.path.exists(model_dir): + model = THA4StudentORT(model_dir, device_id=0) + test_image = np.random.randint( + 0, 256, (512, 512, 4), dtype=np.uint8 + ) + test_pose = np.random.randn(1, 45).astype(np.float32) + + model.setImage(test_image) + result = model.inference(test_pose) + + print("PASS: THA4StudentORT (ONNX Runtime): SUCCESS") + out_shape = result[0].shape if isinstance(result, list) else result.shape + print(f" Output shape: {out_shape}") + else: + print(f"SKIP: THA4StudentORT: Model not found at {model_dir}") + + except Exception as e: + print("FAIL: THA4StudentORT (ONNX Runtime): FAILED") + print(f" Error: {str(e)}") + success = False + + return success + + +def main(): + print("THA4 Student Model Integration Tests") + print("=" * 50) + + results = [] + + # Test direct models first + print("\n--- Direct Model Classes ---") + results.append(("Direct Models", test_direct_models())) + + # Test backends + print("\n--- Framework Backends ---") + results.append(("TensorRT Backend", test_tensorrt_backend())) + results.append(("ONNX Runtime Backend", test_onnxruntime_backend())) + + # Summary + print("\n" + "=" * 50) + print("TEST SUMMARY") + print("=" * 50) + for name, passed in results: + status = "PASS" if passed else "FAIL" + print(f" {name:<30} [{status}]") + + all_passed = all(r[1] for r in results) + overall = "ALL TESTS PASSED" if all_passed else "SOME TESTS FAILED" + print("\nOverall: " + overall) + + return 0 if all_passed else 1 + + +if __name__ == '__main__': + sys.exit(main())