diff --git a/.gitignore b/.gitignore index 0241e25..25a244d 100644 --- a/.gitignore +++ b/.gitignore @@ -142,4 +142,6 @@ data/ !data/models/placeholder.txt *.iml .idea/ -*.pt \ No newline at end of file +*.pt + +.serena \ No newline at end of file diff --git a/THA4_SETUP.md b/THA4_SETUP.md new file mode 100644 index 0000000..a569f99 --- /dev/null +++ b/THA4_SETUP.md @@ -0,0 +1,65 @@ +# 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 new file mode 100644 index 0000000..f3f42a2 --- /dev/null +++ b/THA4_USAGE.md @@ -0,0 +1,127 @@ +# 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/launcher2.py b/launcher2.py index fdd90aa..6abcfdc 100644 --- a/launcher2.py +++ b/launcher2.py @@ -80,6 +80,9 @@ def refreshList(): 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]) + elif '.yaml' == item[-5:] or '.yml' == item[-4:]: + # Support THA4 YAML files + characterList.append(item[:-5] if item.endswith('.yaml') else item[:-4]) refreshList() @@ -227,9 +230,32 @@ 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'], - mapping=['seperable_half', 'seperable_full', 'standard_half', 'standard_full']) + model_select_option = addOption('model_select', title='Model Select', desc='选择使用的模型\nStandard Full精度较高性能较低', + choices=['Seperable Half', 'Seperable Full', 'Standard Half', 'Standard Full', 'THA4'], + mapping=['seperable_half', 'seperable_full', 'standard_half', 'standard_full', 'tha4']) + + # Add event handler for character selection to auto-lock THA4 when YAML is selected + def on_character_change(event): + character_name = self.optionDict['character'].GetValue() + # Check if selected character has a YAML file + yaml_exists = False + for ext in ['.yaml', '.yml']: + if os.path.exists(os.path.join(dirPath, character_name + ext)): + yaml_exists = True + break + + if yaml_exists: + # Auto-select THA4 model and disable the control + model_select_option.control.SetSelection(4) # THA4 is index 4 + model_select_option.control.Enable(False) + else: + # Re-enable model selection for non-YAML files + model_select_option.control.Enable(True) + + self.optionDict['character'].control.Bind(wx.EVT_CHOICE, on_character_change) + + # Trigger initial check + on_character_change(None) addOption('ram_cache_size', title='RAM Cache Size', desc='分配内存缓存大小\n用于存储最终运算结果', choices=['Off', '1GB', '2GB', '4GB', '8GB', '16GB'], mapping=['0b', '1gb', '2gb', '4gb', '8gb', '16gb']) @@ -411,10 +437,14 @@ class LauncherPanel(wx.Panel): 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 'tha4' in args['model_select']: + run_args.append('--model') + run_args.append('tha4') + else: + 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') diff --git a/main.py b/main.py index ca072b0..f250514 100644 --- a/main.py +++ b/main.py @@ -330,12 +330,13 @@ class MouseClientProcess(Process): class ModelClientProcess(Process): - def __init__(self, input_image): + def __init__(self, input_image, yaml_path=None): super().__init__() self.should_terminate = Value('b', False) self.updated = Value('b', False) self.data = None self.input_image = input_image + self.yaml_path = yaml_path self.output_queue = Queue() self.input_queue = Queue() self.shms = [shared_memory.SharedMemory(create=True, size=args.model_output_size * args.model_output_size * 4) @@ -346,25 +347,36 @@ class ModelClientProcess(Process): self.gpu_cache_hit_ratio = Value('f', 0.0) def run(self): - self.model = get_core(device_id=args.device_id, - use_tensorrt=args.use_tensorrt, + # Check if using THA4 model + if args.model == 'tha4': + from tha4_core import get_tha4_core + self.model = get_tha4_core( + device_id=args.device_id, + use_eyebrow=args.eyebrow, + interpolation_scale=args.interpolation_scale if args.use_interpolation else 1, + yaml_path=self.yaml_path + ) + else: + # Use existing ezvtb_rt interface for THA3 models + self.model = get_core(device_id=args.device_id, + use_tensorrt=args.use_tensorrt, - model_seperable = args.model_seperable, - model_half=args.model_half, - model_cache_size=args.max_gpu_cache_len, - model_use_eyebrow=args.eyebrow, + model_seperable = args.model_seperable, + model_half=args.model_half, + model_cache_size=args.max_gpu_cache_len, + model_use_eyebrow=args.eyebrow, - use_interpolation=args.use_interpolation, - interpolation_scale=args.interpolation_scale, - interpolation_half=args.interpolation_half, + use_interpolation=args.use_interpolation, + interpolation_scale=args.interpolation_scale, + interpolation_half=args.interpolation_half, - cacher_quality=args.cacher_quality, - cacher_ram_size=args.max_cache_len, + cacher_quality=args.cacher_quality, + cacher_ram_size=args.max_cache_len, - use_sr=args.use_sr, - sr_half=args.sr_half, - sr_x4=args.sr_x4, - sr_noise=args.sr_noise) + use_sr=args.use_sr, + sr_half=args.sr_half, + sr_x4=args.sr_x4, + sr_noise=args.sr_noise) self.model.setImage(self.input_image) input_pose = np.zeros((1, 45), dtype=np.float32) @@ -529,26 +541,47 @@ class ModelClientProcess(Process): @torch.no_grad() def main(): - img = Image.open(f"data/images/{args.character}.png") - img = img.convert('RGBA') - ow, oh = img.size - for i, px in enumerate(img.getdata()): - if px[3] <= 0: - y = i // ow - x = i % ow - img.putpixel((x, y), (0, 0, 0, 0)) - if ow != 512 or oh != 512: - img = resize_to_512_center(img) - if args.alpha_clean: - curves = { - 'a': [ - (60, 0), - (200, 255) - ] - } - img = apply_color_curves(img, curves) - input_image = np.array(img) - input_image = cv2.cvtColor(input_image, cv2.COLOR_RGBA2BGRA) + # Check if character is a YAML file (THA4 model) + yaml_path = None + is_tha4 = False + + # Try to find YAML file + for ext in ['.yaml', '.yml']: + yaml_candidate = f"data/images/{args.character}{ext}" + if os.path.exists(yaml_candidate): + yaml_path = yaml_candidate + is_tha4 = True + print(f"Detected THA4 YAML config: {yaml_path}") + # Override model selection to THA4 + args.model = 'tha4' + break + + # Load character image (only for non-THA4 models) + if not is_tha4: + img = Image.open(f"data/images/{args.character}.png") + img = img.convert('RGBA') + ow, oh = img.size + for i, px in enumerate(img.getdata()): + if px[3] <= 0: + y = i // ow + x = i % ow + img.putpixel((x, y), (0, 0, 0, 0)) + if ow != 512 or oh != 512: + img = resize_to_512_center(img) + if args.alpha_clean: + curves = { + 'a': [ + (60, 0), + (200, 255) + ] + } + img = apply_color_curves(img, curves) + input_image = np.array(img) + input_image = cv2.cvtColor(input_image, cv2.COLOR_RGBA2BGRA) + else: + # For THA4, create a dummy image (will be ignored) + input_image = np.zeros((512, 512, 4), dtype=np.uint8) + print("THA4 mode: Character image will be loaded from YAML") print("Character Image Loaded:", args.character) cap = None @@ -652,7 +685,7 @@ def main(): } model_output = None - model_process = ModelClientProcess(input_image) + model_process = ModelClientProcess(input_image, yaml_path=yaml_path) model_process.daemon = True model_process.start() model_output_nps = [np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, diff --git a/models.py b/models.py index e91620a..34abb08 100644 --- a/models.py +++ b/models.py @@ -1,5 +1,6 @@ import time +import torch import torch.nn as nn import tha2.poser.modes.mode_20 @@ -13,6 +14,14 @@ 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): @@ -171,3 +180,58 @@ class TalkingAnime(nn.Module): 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 2704cef..747e7df 100644 --- a/pack_release.bat +++ b/pack_release.bat @@ -7,57 +7,57 @@ REM ======================================================================== cd /D "%~dp0" @REM git clean -fdx -git submodule update --init --recursive +@REM git submodule update --init --recursive -if not exist "envs" mkdir envs -if not exist "envs\miniconda3" mkdir envs\miniconda3 +@REM if not exist "envs" mkdir envs +@REM if not exist "envs\miniconda3" mkdir envs\miniconda3 -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'}" +@REM IF not EXIST %~dp0envs\miniconda3\Scripts ( +@REM @RD /S /Q %~dp0envs\miniconda3 +@REM mkdir %~dp0envs\miniconda3 +@REM echo "Downloading miniconda..." +@REM 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'}" - echo "Installling minconda..." - start /wait "" %~dp0envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=%~dp0envs\miniconda3 - echo "Successfully install minconda" - del envs\miniconda3.exe -) +@REM echo "Installling minconda..." +@REM start /wait "" %~dp0envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=%~dp0envs\miniconda3 +@REM echo "Successfully install minconda" +@REM del envs\miniconda3.exe +@REM ) -SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% +@REM SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% -call activate -call conda env list -call conda update -y --all -call conda create -n ezvtb_rt_venv_release -c conda-forge conda-pack python=3.10 -y -call conda activate ezvtb_rt_venv_release -call conda env list +@REM call activate +@REM call conda env list +@REM call conda update -y --all +@REM call conda create -n ezvtb_rt_venv_release -c conda-forge conda-pack python=3.10 -y +@REM call conda activate ezvtb_rt_venv_release +@REM call conda env list -call conda install -y -c conda-forge cudatoolkit=12.9.1 cudnn=12.9.1 -call conda install -y nvidia/label/cuda-12.9.1::cuda-nvcc-dev_win-64 -call conda install -y pycuda -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ +@REM call conda install -y -c conda-forge cudatoolkit=12.9.1 cudnn=12.9.1 +@REM call conda install -y nvidia/label/cuda-12.9.1::cuda-nvcc-dev_win-64 +@REM call conda install -y pycuda -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ -call conda-pack -n ezvtb_rt_venv_release -o %~dp0envs\python_embedded --format no-archive +@REM 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% +@REM SET PATH=%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% -call python -m pip install --upgrade pip wheel -i https://mirrors.aliyun.com/pypi/simple/ -echo yes|python -m pip install nvidia-cudnn-cu12 -i https://mirrors.aliyun.com/pypi/simple/ +@REM call python -m pip install --upgrade pip wheel -i https://mirrors.aliyun.com/pypi/simple/ +@REM echo yes|python -m pip install nvidia-cudnn-cu12 -i https://mirrors.aliyun.com/pypi/simple/ -echo yes|pip install tensorrt_cu12_libs==10.11.0.33 tensorrt_cu12_bindings==10.11.0.33 tensorrt==10.11.0.33 --extra-index-url https://pypi.nvidia.com +@REM echo yes|pip install tensorrt_cu12_libs==10.11.0.33 tensorrt_cu12_bindings==10.11.0.33 tensorrt==10.11.0.33 --extra-index-url https://pypi.nvidia.com -call python -m pip install -r requirements.txt --no-warn-script-location -i https://mirrors.aliyun.com/pypi/simple/ +@REM call python -m pip install -r requirements.txt --no-warn-script-location -i https://mirrors.aliyun.com/pypi/simple/ -if not exist "data\models" mkdir data\models -echo "Downloading model zip file..." -powershell -Command "(New-Object Net.WebClient).DownloadFile('https://github.com/zpeng11/ezvtuber-rt/releases/download/0.0.1/20241220.zip', '.\data\models\model.zip')" -echo "Unzipping model zip file..." -powershell -Command "Expand-Archive -Path '.\data\models\model.zip' -DestinationPath '.\data\models\temp' -Force; Move-Item -Path '.\data\models\temp\20241220\*' -Destination '.\data\models\' -Force; Remove-Item -Path '.\data\models\temp' -Recurse -Force" -del .\data\models\model.zip +@REM if not exist "data\models" mkdir data\models +@REM echo "Downloading model zip file..." +@REM powershell -Command "(New-Object Net.WebClient).DownloadFile('https://github.com/zpeng11/ezvtuber-rt/releases/download/0.0.1/20241220.zip', '.\data\models\model.zip')" +@REM echo "Unzipping model zip file..." +@REM powershell -Command "Expand-Archive -Path '.\data\models\model.zip' -DestinationPath '.\data\models\temp' -Force; Move-Item -Path '.\data\models\temp\20241220\*' -Destination '.\data\models\' -Force; Remove-Item -Path '.\data\models\temp' -Recurse -Force" +@REM del .\data\models\model.zip -@REM To verify if environment works -call python ezvtb_rt_interface.py +@REM @REM To verify if environment works +@REM call python ezvtb_rt_interface.py @REM Clean up step echo "Cleaning up .trt files in data\models directory..." diff --git a/tha4 b/tha4 new file mode 160000 index 0000000..3206401 --- /dev/null +++ b/tha4 @@ -0,0 +1 @@ +Subproject commit 320640116abd604de22d7b56faec85b4d95770e0 diff --git a/tha4_adapter.py b/tha4_adapter.py new file mode 100644 index 0000000..73c6432 --- /dev/null +++ b/tha4_adapter.py @@ -0,0 +1,128 @@ +""" +THA4 Adapter Module - Simplified Version +PyTorch only, no TensorRT/ONNX support +Adapts THA4 models to existing EasyVtuber architecture with THA3-compatible interface +Uses CharacterModel to properly load YAML + PNG + PT files +""" +import sys +import os +import torch + +# Add tha4 source code path +sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'tha4', 'src')) + +from tha4.charmodel.character_model import CharacterModel + + +def convert_tha3_pose_to_tha4(eyebrow_vector, mouth_eye_vector, pose_vector): + """ + Convert THA3 pose format to THA4 format + + THA3 inputs: + - eyebrow_vector: [12] eyebrow parameters + - mouth_eye_vector: [27] facial expression parameters + - pose_vector: [6] head pose parameters + + THA4 input: + - pose: [45] unified pose vector + - [0:39] face pose (12 eyebrow + 27 face) + - [39:45] body pose (6 head rotation) + + Returns: + torch.Tensor [batch, 45] + """ + # Concatenate: [eyebrow(12), mouth_eye(27), pose(6)] = [45] + face_pose = torch.cat([eyebrow_vector, mouth_eye_vector], dim=1) + tha4_pose = torch.cat([face_pose, pose_vector], dim=1) + return tha4_pose + + +class THA4Wrapper: + """ + THA4 Poser wrapper providing THA3 TalkingAnime3-compatible interface + Uses CharacterModel to load YAML config + PNG image + PT models + """ + + def __init__(self, device: torch.device, yaml_path: str = None): + """ + Initialize THA4 wrapper + + Args: + device: torch device (cuda/cpu) + yaml_path: path to character model YAML (required) + """ + self.device = device + + # YAML path is required + if yaml_path is None: + raise ValueError( + "THA4 requires a YAML config path. " + "Please select a .yaml file in the Character selection." + ) + + # Check if YAML exists + if not os.path.exists(yaml_path): + raise FileNotFoundError( + f"THA4 character model YAML not found: {yaml_path}\n" + f"Expected files in data/images/:\n" + f" - .yaml (config with relative paths)\n" + f" - .png (source image)\n" + f" - _face_morpher.pt (face model)\n" + f" - _body_morpher.pt (body model)" + ) + + # Load CharacterModel from YAML + print(f"Loading THA4 CharacterModel from: {yaml_path}") + self.character_model = CharacterModel.load(yaml_path) + + # Get poser (lazy loads .pt files) + self.poser = self.character_model.get_poser(device) + + # Get character image tensor (loads .png file) + char_img = self.character_model.get_character_image(device) + + # Add batch dimension if needed: [4, 512, 512] -> [1, 4, 512, 512] + if len(char_img.shape) == 3: + self.character_image = char_img.unsqueeze(0) + else: + self.character_image = char_img + + print(f"THA4 model loaded successfully:") + print(f" Device: {device}") + print(f" Image: {self.character_model.character_image_file_name}") + print(f" Face: {self.character_model.face_morpher_file_name}") + print(f" Body: {self.character_model.body_morpher_file_name}") + + 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 TalkingAnime3.forward() + + Args: + image: [batch, 4, 512, 512] + mouth_eye_vector: [batch, 27] + pose_vector: [batch, 6] + eyebrow_vector: [batch, 12] + mouth_eye_vector_c: compressed version (not used by THA4) + eyebrow_vector_c: compressed version (not used by THA4) + ratio: GPU cache hit ratio (not used by THA4) + + Returns: + output image [batch, 4, 512, 512] + """ + # Convert pose format + tha4_pose = convert_tha3_pose_to_tha4( + eyebrow_vector, mouth_eye_vector, pose_vector + ) + + # Call THA4 poser + with torch.no_grad(): + output = self.poser.pose(image, tha4_pose, output_index=0) + + return output + + def to(self, device: torch.device): + """Move model to specified device""" + self.device = device + self.poser.to(device) + return self diff --git a/tha4_core.py b/tha4_core.py new file mode 100644 index 0000000..d57ad70 --- /dev/null +++ b/tha4_core.py @@ -0,0 +1,164 @@ +""" +THA4 Core - Pure PyTorch implementation +Provides interface compatible with CoreTRT/CoreORT for THA4 models +""" +import torch +import numpy as np +import cv2 +from tha4_adapter import THA4Wrapper, convert_tha3_pose_to_tha4 + +# Import color space conversion from THA4 +import sys +import os +sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'tha4', 'src')) +from tha4.image_util import convert_linear_to_srgb + + +class THA4Core: + """ + THA4 Core class with interface compatible with CoreTRT/CoreORT + + Uses CharacterModel to load YAML + PNG + PT files properly. + The PNG image from YAML is used as the source image. + + Expected interface: + - setImage(img: np.ndarray) - NOT USED, image comes from YAML + - inference(pose: np.ndarray) -> List[np.ndarray] - run inference + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1): + """ + Initialize THA4 Core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + Default: 'data/models/tha4/character_model.yaml' + interpolation_scale: number of output frames to generate + (THA4 doesn't support real interpolation, will duplicate frames) + """ + device_str = f'cuda:{device_id}' if torch.cuda.is_available() else 'cpu' + self.device = torch.device(device_str) + self.use_eyebrow = use_eyebrow + self.interpolation_scale = interpolation_scale + + # Create THA4 wrapper (loads YAML + PNG + PT) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + + # Use character image from YAML (already loaded by wrapper) + self.input_image_tensor = self.wrapper.character_image + + # For compatibility (no caching in THA4 core) + self.cacher = None + + print(f"THA4 Core initialized on {self.device}") + print(f"Using character image from model (ignoring setImage calls)") + if interpolation_scale > 1: + print(f"Note: THA4 doesn't support interpolation, " + f"will duplicate frames {interpolation_scale}x") + + def setImage(self, img: np.ndarray): + """ + Set input character image - IGNORED for THA4 + + THA4 uses the character image from the YAML/PNG files loaded + during initialization. This method exists only for interface + compatibility with CoreTRT/CoreORT. + + Args: + img: numpy array in BGRA format (ignored) + """ + # THA4 uses character image from YAML, ignore external image + print("THA4: setImage() called but ignored (using image from YAML)") + + def inference(self, pose: np.ndarray) -> list: + """ + Run inference + + Args: + pose: numpy array of shape (1, 45) + [0:12] - eyebrow parameters + [12:39] - face parameters + [39:45] - head pose parameters + + Returns: + List of output images as numpy arrays [H, W, 4] in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set. Call setImage() first.") + + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] # [1, 12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] # [1, 27] + head_pose = pose[:, 39:45] # [1, 6] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + # Call wrapper (it handles the pose conversion internally) + # We need to create dummy compressed versions for interface compatibility + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # THA4 outputs in linear color space, convert to sRGB + # Output is [batch, 4, H, W] in range [-1, 1] + output_image = output_tensor[0].float() # [4, 512, 512] + + # Clip to [-1, 1] and convert to [0, 1] + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + + # Convert linear RGB to sRGB (keep alpha as-is) + output_image = convert_linear_to_srgb(output_image) + + # Convert to [H, W, C] numpy + output_np = output_image.permute(1, 2, 0).cpu().numpy() + + # Convert to uint8 + output_np = (output_np * 255.0).astype(np.uint8) + + # Convert RGBA to BGRA for output + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Return as list with proper number of frames + # THA4 doesn't support real interpolation, so duplicate frames if needed + return [output_bgra] * self.interpolation_scale + + +def get_tha4_core(device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1): + """ + Factory function to create THA4 core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow control + yaml_path: path to character_model.yaml + Default: 'data/models/tha4/character_model.yaml' + interpolation_scale: number of frames (for compatibility with interpolation) + + Returns: + THA4Core instance + """ + return THA4Core(device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + interpolation_scale=interpolation_scale)