clear and add bat(TODOFIX: tha4 student model vram cache 0.0%）
This commit is contained in:
nick-haoran committed
1 parent
0b6cb9c8d3
commit
daaa8522a0
11 files changed
+440
-361
No files matched your search
@@ -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
|
||||
|
||||
@@ -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卡吧别折腾自己了）。
|
||||
@@ -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.
|
||||
-127
@@ -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 文件名
|
||||
@@ -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')
|
||||
|
||||
+67
-29
@@ -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
|
||||
|
||||
|
||||
+1
-1
Submodule ezvtuber-rt updated: d23cfb3a36...b59e1bab05.
+95
-7
@@ -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')
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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())
|
||||
Reference in new issue
Block a user