diff --git a/.gitignore b/.gitignore index 0241e25..3f5ecf6 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,11 @@ pretrained/** !pretrained/PUT_MODEL_HERE character -src +!src +!tha4/src +tha4/distiller-ui-doc +tha4/docs +tha4/poetry dst .idea bin @@ -142,4 +146,9 @@ data/ !data/models/placeholder.txt *.iml .idea/ -*.pt \ No newline at end of file +*.pt + +.serena +onnx_model_tha4/fp16 +onnx_model_tha4/fp32 +onnx_model_tha4/ \ No newline at end of file diff --git a/.gitmodules b/.gitmodules index ebfeccc..2abdf60 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,4 +1,4 @@ [submodule "ezvtuber-rt"] path = ezvtuber-rt - url = https://github.com/zpeng11/ezvtuber-rt.git - branch = main \ No newline at end of file + url = git@github.com:yuyuyzl/ezvtuber-rt.git + branch = tha4 \ No newline at end of file diff --git a/01A.启动器.bat b/01A.启动器.bat new file mode 100644 index 0000000..08cca03 --- /dev/null +++ b/01A.启动器.bat @@ -0,0 +1,7 @@ +@echo on +cd /D "%~dp0" + +SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% + +start "" pythonw launcher2.py +exit \ No newline at end of file diff --git a/01A.构建运行环境（默认源）.bat b/01A.构建运行环境（默认源）.bat deleted file mode 100644 index dfba141..0000000 --- a/01A.构建运行环境（默认源）.bat +++ /dev/null @@ -1,41 +0,0 @@ -@echo on -cd /D "%~dp0" - -WHERE conda -IF %ERRORLEVEL% NEQ 0 ( - if not exist "envs" mkdir envs - 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 "(New-Object Net.WebClient).DownloadFile('https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe', '.\envs\miniconda3.exe')" - - 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 - ) -) - -IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% - -call activate -call conda env list -call conda update -y --all -call conda create -y -n ezvtb_rt_venv python=3.10 -call conda activate ezvtb_rt_venv -call conda env list - -call conda install -y nvidia/label/cuda-12.9.0::cuda-nvcc-dev_win-64 -call conda install -y conda-forge::pycuda - -call python -m pip install --upgrade pip wheel -echo yes|python -m pip install nvidia-cudnn-cu12 - -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 - -pause \ No newline at end of file diff --git a/01B.启动器（调试输出）.bat b/01B.启动器（调试输出）.bat new file mode 100644 index 0000000..453f60c --- /dev/null +++ b/01B.启动器（调试输出）.bat @@ -0,0 +1,8 @@ +@echo on +cd /D "%~dp0" + +SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% + +python launcher2.py + +pause \ No newline at end of file diff --git a/01B.构建运行环境（国内源）.bat b/01B.构建运行环境（国内源）.bat deleted file mode 100644 index 989ca17..0000000 --- a/01B.构建运行环境（国内源）.bat +++ /dev/null @@ -1,41 +0,0 @@ -@echo on -cd /D "%~dp0" - -WHERE conda -IF %ERRORLEVEL% NEQ 0 ( - if not exist "envs" mkdir envs - 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 "(New-Object Net.WebClient).DownloadFile('https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe', '.\envs\miniconda3.exe')" - - 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 - ) -) - -IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% - -call activate -call conda env list -call conda update -y --all -call conda create -y -n ezvtb_rt_venv python=3.10 -call conda activate ezvtb_rt_venv -call conda env list - -call conda install -y nvidia/label/cuda-12.9.0::cuda-nvcc-dev_win-64 -call conda install -y pycuda -c https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ - -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/ - -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/ - -pause \ No newline at end of file diff --git a/02.构建TensorRT模型（N卡开启TensorRT使用）.bat b/02.构建TensorRT模型（N卡开启TensorRT使用）.bat deleted file mode 100644 index 0315d6f..0000000 --- a/02.构建TensorRT模型（N卡开启TensorRT使用）.bat +++ /dev/null @@ -1,11 +0,0 @@ -@echo off -cd /D "%~dp0" -IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% -call activate -call conda env list -call conda activate ezvtb_rt_venv -call conda env list - -python ezvtb_rt_interface.py - -pause \ No newline at end of file diff --git a/02A.启动器.bat b/02A.启动器.bat deleted file mode 100644 index 2a714e4..0000000 --- a/02A.启动器.bat +++ /dev/null @@ -1,10 +0,0 @@ -@echo off -cd /D "%~dp0" -IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% -call activate -call conda env list -call conda activate ezvtb_rt_venv -call conda env list - - -pythonw launcher2.py \ No newline at end of file diff --git a/02B.启动器（调试输出）.bat b/02B.启动器（调试输出）.bat deleted file mode 100644 index eee185a..0000000 --- a/02B.启动器（调试输出）.bat +++ /dev/null @@ -1,12 +0,0 @@ -@echo on -cd /D "%~dp0" -IF EXIST %~dp0envs\miniconda3\Scripts SET PATH=%~dp0envs\miniconda3\Scripts;%PATH% -call activate -call conda env list -call conda activate ezvtb_rt_venv -call conda env list - - -python launcher2.py - -pause \ No newline at end of file diff --git a/PerformanceTest.md b/PerformanceTest.md index 1d048ea..e2eb39d 100644 --- a/PerformanceTest.md +++ b/PerformanceTest.md @@ -120,4 +120,4 @@ DirectML 允许所有 Windows 显卡设备以及 AIPC 的 NPU 运行此项目， 所以，DirectML运行本项目的任意模型出现奇怪结果，或不符合预期的性能倒退，都应该属于意料之中 换个模型精度，观察调试输出是否有DirectML黄字警告 通常可以解决大部分问题。 具体某个硬件在DirectML的运行异常修复，已远超本项目的目标范围，请自行验证失效kernel，向硬件厂商和微软求助并期待驱动更新带来解决 -（建议买张N卡吧别折腾自己了）。 +（建议买张N卡吧别折腾自己了）。 \ No newline at end of file diff --git a/README.md b/README.md index 35d5107..989ca5d 100644 --- a/README.md +++ b/README.md @@ -15,21 +15,22 @@ Updates: * 使用 https://github.com/zpeng11/ezvtuber-rt/tree/main 项目转换的ONNX模型，实现TensorRT加速以及非N卡支持，在半精度和全精度下均有加速（具体取决于显卡） * 为A卡和I卡提供DirectML支持，人人都能玩。 * 引入[RIFE](https://github.com/hzwer/ECCV2022-RIFE)模型进行插帧，极限帧数提升达到50%-100%，配合锁帧可以有效降低使用占用。 -* 使用 [TurboJPEG](https://github.com/Dobatymo/turbojpeg-python) 获得成倍缓存命中率提升，在长时间使用上显卡减负效果相当出色。 +* 使用 `brotli` 压缩器获得成倍缓存命中率提升，在长时间使用上显卡减负效果相当出色。 * 使用 [waifu2x](https://github.com/nagadomi/waifu2x) 和 [real-esrgan](https://github.com/xinntao/Real-ESRGAN) 项目带来的输出超分辨率，对比anime4k效果提升显著（也使用更多gpu占用） * 更漂亮的 WxPython 新启动器界面，带中文介绍更加简单易用 * 添加Spout2 支持输出原生透明通道给OBS +* 添加 [THA V4](https://github.com/pkhungurn/talking-head-anime-4-demo) 支持，拥有更好的图形输出 ## Requirements ### 硬件 -- 支持FaceID的iPhone（使用ifacialmocap软件，需购买，需要稳定的WIFI连接）或网络摄像头（使用OpenCV） +- 支持无摄像头鼠标捕获/USB网络摄像头/iPhone结构光面捕（ifacialmocap软件，需购买） - 任意5年内的游戏级显卡，N卡/A卡/I卡均可使用，某些性能强的核显也可尝试，请参考[性能测试结果](PerformanceTest.md) ### 软件 -- 本方案在Windows 10上测试可用 -- OBS或Unity Capture（虚拟摄像头方案） +- 本方案在Windows 10/11上测试可用 +- OBS 使用经验 - Photoshop或其他图片处理软件 - 科学上网方案，看懂英文网站和报错的能力 @@ -40,100 +41,13 @@ Updates: * [谷歌网盘](https://drive.google.com/drive/folders/1cYj18EfVQ2Cl348_rkCu_fgaasHTI_io?usp=drive_link) ### 安装Spout2 OBS插件(可选) -安装整合包附带的 `OBS_Spout2_Plugin_Install_v1.8.exe`， 安装包来自 https://github.com/Off-World-Live/obs-spout2-plugin 项目，可以提供透明通道给OBS。 - -### 构建TensorRT模型（N卡专用，可选） -双击`02.构建TensorRT模型（N卡使用）.bat`，若是英伟达用户将会开始TensorRT构建, 每台配置仅需构建一次，更新显卡驱动可能需要重新构建 -额外需要>20min时间(取决于显卡)，结束后可以大幅提升运行速度（可简单理解为玩游戏之前编译着色器，和每个人显卡有关） -构建的开头和结束输出如下，请耐心等进度条百分百或完成后自动退出。 -若不慎关掉窗口也可以双击自动从断点重来，可以先放后台运行构建并尝试下一步运行暂时不勾选TensorRT： -``` -Start testing if TensorRT works on this machine - 0%| | 0/38 [00:00conda --version -conda 24.11.3 -``` -以上方式可以保证conda环境被你选定的Ananconda统一管理。若没有此需求可以忽略直接运行下一步，`01X.构建运行环境（XX源）.bat`脚本会自动安装miniconda。 - -### 克隆项目和子项目 -``` -git clone https://github.com/zpeng11/EasyVtuber.git -cd EasyVtuber -git submodule init -git submodule update --recursive --remote -``` - -### 安装环境 -根据网络环境双击运行`01A.构建运行环境（默认源）.bat`或者`01B.构建运行环境（国内源）.bat` -这两个脚本会构建名为`ezvtb_rt_venv`的conda环境 -或可以用如下命令手动创建： -``` -conda create -y -n ezvtb_rt_venv python=3.10 -conda activate ezvtb_rt_venv -conda install -y nvidia/label/cuda-12.6.3::cuda-nvcc-dev_win-64 -conda install -y conda-forge::pycuda -python -m pip install --upgrade pip wheel -python -m pip install nvidia-cudnn-cu12 -pip install tensorrt_cu12_libs==10.6.0 tensorrt_cu12_bindings==10.6.0 tensorrt==10.6.0 --extra-index-url https://pypi.nvidia.com -python -m pip install -r requirements.txt --no-warn-script-location -``` - -### 下载模型 -前往模型 release 地址 https://github.com/zpeng11/ezvtuber-rt/releases/tag/0.0.1 -下载 https://github.com/zpeng11/ezvtuber-rt/releases/download/0.0.1/20241220.zip 并解压到`data/models`文件夹下 - -### Pycharm 配置 -安装完成后，在Pycharm内打开本项目，右下角解释器菜单点开，`Add Interpreter...`->`Conda Environment`->`Existing environment` -选好自己电脑上的`conda.exe`和刚才创建好的`ezvtb_rt_venv`环境内的`python.exe` -点击OK，依赖全亮即可 - -### 检查环境模型，并构建TensorRT -在Conda环境中执行以下命令 -`python ezvtb_rt_interface.py` - -### 运行启动器 -在Conda环境中执行以下命令 -`python launcher2.py` - -### 环境错误排查 -当运行TensorRT构建或启动器出现错误时，请参考如下可能性： -1. 环境安装有错误（一般运行时缺少库99%的原因都是这个，pip下载并不稳定，各种网络问题都可能导致安装失败），检查源，手动删除`envs`文件夹再试, 对照 [此log](assets/complete_building_log.txt) 为成功案例排查原因。 实在不懂排查请使用整合包。 -2. 英伟达显卡但计算架构低于7.5（10系以及之前）不支持TensorRT - ## 输入输出设备 #### SPOUT2 OBS插件输出 @@ -143,7 +57,7 @@ python -m pip install -r requirements.txt --no-warn-script-location #### OBS Virtual Camera 无透明通道 -如果你选择自己进行抠像你可以直接仅输出RGB到OBS，使用补帧时可能也需要此方法 +次选方案，如果你选择自己进行抠像你可以直接仅输出RGB到OBS，使用补帧时可能也需要此方法 使用`视频采集设备`源，设置`滤镜`, 在`效果滤镜`添加`色值`并自定义颜色为纯黑色，多尝试调整颜色参数直到抠图合理 @@ -158,29 +72,21 @@ python -m pip install -r requirements.txt --no-warn-script-location 之后对OBS中的视频采集设备添加滤镜-着色器-选择你下载的`merge alpha2.hlsl`-关闭 这样透明通道就应用回左边的图像了 你可能需要手动调整一下裁剪把右侧的无用画面切掉 -(看不到着色器滤镜的话就是StreamFX没装好或者OBS不是最新版) +(看不到着色器滤镜的话就是StreamFX没装好或者OBS不是最新版)d -#### UnityCapture（Outdated，性能问题） - -存在未查明的性能瓶颈 -如果需要使用透明通道输出，参考 https://github.com/schellingb/UnityCapture#installation 安装好UnityCapture -只需要正常走完Install.bat，在OBS里能看到对应的设备（Unity Video Capture）就行 - -在OBS添加完摄像头以后，还需要手动配置一次摄像头属性才能支持ARGB -右键属性-取消激活-分辨率类型自定义-分辨率512x512(与`--output_size`参数一致)-视频格式ARGB-激活 - -#### iFacialMocap +#### iFacialMocap 面捕 目前最推荐的方案，使用苹果结构光硬件，捕捉效果追踪准确度都最好，demo中使用了此输入。 需要输入`ip:port`,其中`ip`会显示在手机上，`port`使用默认`49983` https://www.ifacialmocap.com/download/ 你大概率需要购买正式版（非广告，只是试用版不太够时长） -购买之前确认好自己的设备支持 -**不需要下载PC软件**，装好iOS端的软件即可，连接信息通过参数传入Python -游戏加速器，梯子等网络工具可能会导致ip失效，请自行解决。 +购买之前确认好自己的设备支持，装好iOS端的软件即可，连接信息通过参数传入Python +一般需要iphone与window在同一局域网下 +且游戏加速器，梯子等网络工具可能会导致ip失效，请自行解决。 -#### OpenSeeFace +#### OpenSeeFace 高精度摄像头面捕 +在没有苹果硬件时最推荐使用的捕捉方案，需要更多配置 https://github.com/emilianavt/OpenSeeFace/releases 直接下载最新版本的Release包并解压，或使用整合包附带的压缩包 之后进入解压目录的Binary文件夹 @@ -189,28 +95,34 @@ https://github.com/emilianavt/OpenSeeFace/releases 之后保存并双击`run.bat`运行，按照提示选择摄像头、分辨率、帧率，捕获正常的话可以看到输出画面 最后在启动器中选择OpenSeeFace输入，或添加启动参数`--osf 127.0.0.1:11573`即可接入OpenSeeFace +#### OpenCV 低精度摄像头面捕 + +此输入不需要配置，会尝试开启默认的USB或笔记本电脑摄像头开始面捕，稳定性和精确度略欠佳。 + +#### 鼠标音频输入 + +隐私考虑不想使用任何摄像头可使用这个输入，定时循环眨眼和呼吸，配合音频音量控制口型，可以在无摄像头时得到不错的动态。 + ## 性能配置(调参必读） ### 概念须知 #### 模型&精度 本项目完全基于AI算法生成，由3类AI模型提供服务（缝合怪）： -* THA模型：指本项目的基础模型，输入静态图像和姿态信息计算出新图像 +* THA模型：指本项目的基础模型，输入静态图像和姿态信息计算出新图像，有v3_seperable, v3_standard, v4, v4_student 四种版本。 * 插帧/补帧模型：指RIFE插帧模型，输入第1帧和第2帧，输出第1.5帧，类似DLSS3或小黄鸭，此模型是整体性能提升的关键 -* 超分模型：指waifu2x或realesrgan模型，输入一帧图像，解计算出放大且更清晰的同一帧，显著提升画面效果 +* 超分模型：指Anime4k, waifu2x或realesrgan模型，输入一帧图像，解算出放大且更清晰的同一帧，显著提升画面效果 -每一种模型提供了半精度(标记为half)和全精度(标记为full)选项，半精度通常运行速度比全精度快，但生图效果稍微逊色，需要根据性能进行取舍。 -THA模型有另外的`seperable`和`standard`两个版本，可以简单理解为`seperable`是性能更强但生图效果稍差的版本，所以加上两种精度THA模型一共有4种组合。 +每一种模型提供了半精度(标记为half)和全精度(标记为full)选项，半精度通常运行速度比全精度快，但生图效果稍微逊色，且兼容性可能不好。 #### 缓存 当我们当前的姿态与一秒前的姿态完全相同时，输入给模型的参数也完全相同，没有必要再做一次完全相同的运算得出完全相同的结果 -缓存的原理是将过去一段时间内产生过的结果储存起来。 -倘若发现姿态相同，就使用储存的结果输出而绕过模型运行（又叫缓存命中），从而减轻显卡负担。 -本项目同时使用显存VRAM和内存RAM用作缓存目的。在内存中还使用了JPEG格式对缓存的图片进行压缩。 +倘若发现姿态相同，就使用曾经的结果输出而绕过模型运行（又叫缓存命中），从而减轻显卡负担。 +本项目同时使用显存VRAM和内存RAM用作缓存目的。在内存中还对缓存的图片进行压缩。 缓存大小（单位GB）决定了缓存最多能使用的空间大小。更大的缓存能让缓存命中率更高，长时间使用时显卡占用更小。但也会挤占其他应用或游戏的空间。 需要自行理解并取舍。 -#### 输入简化 +#### 输入简化（量化） 输入简化与上文的缓存相辅相成，我们将姿态参数分入一些区间，这样相近的参数就可以命中同一个缓存。 举个非常简化的例子：缓存中的一个结果对应姿态参数`0.5`, 在不进行任何输入简化时，需要新输入参数完全相同于`0.5`即 `x = 0.5`才能判定命中并绕过GPU运算。 而开启了输入简化后，只需要新输入参数 `x` $\approx$ `0.5` 就可以命中。如`x = 0.49`或 `x = 0.51`. @@ -226,9 +138,9 @@ THA模型有另外的`seperable`和`standard`两个版本，可以简单理解 #### 超分 ![image](https://github.com/user-attachments/assets/afa78440-5122-4b57-b586-4487e125eb2f) -超分由waifu2x和realesrgan模型的两种精度一共4种组合，外加一种Anime4K纯cpu实现（效果比较差）。 +超分由waifu2x和realesrgan模型的两种精度一共4种组合，外加一种Anime4K实现。 通常选用半精度，根据使用半身或大头选择`waifu2x_x2_half`和`realesrgan_x4_half` -超分启动时会显著增加显卡占用，通常建议在低负载游戏或者杂谈状态下使用。 +超分启动时会显著增加显卡占用。 #### 插帧/补帧 ![image](https://github.com/user-attachments/assets/de392b93-139b-455b-b561-7713d26fca20) @@ -240,17 +152,70 @@ RIFE插帧的计算消耗比起一次THA计算要少约50%且画面观感差异 ## 脱离启动器运行本体 -`main.py` 的完全体运行命令请参考 `02B.启动器（调试输出）.bat` 的输出以及 `args.py`。 以下为部分基础参数 +`main.py` 的完全体运行命令请参考 `02B.启动器（调试输出）.bat` 的输出以及 `args.py`。 -参数名 | 值类型 | 说明 -:---: | :---: | :---: ---character|字符串|`character`目录下的输入图像文件名，不需要带扩展名 ---debug|无|打开OpenCV预览窗口输出渲染结果，如果没有任何输出配置，该参数默认生效 ---input|字符串|不使用iOS面捕时，传入要使用的摄像头设备名称，默认为设备0，有ifm参数时无效 ---ifm|字符串|使用iOS面捕时，传入设备的`IP:端口号`，如`192.168.31.182:49983` ---output_webcam|字符串|可用值为`obs` `unitycapture`，选择对应的输出种类，不传不输出到摄像头 ---extend_movement|浮点数|使用iOS面捕返回的头部位置，对模型输出图像进一步进行移动和旋转使得上半身可动
传入的数值表示移动倍率（建议值为1） ---output_size|字符串|格式为`256x256`，必须是4的倍数。
增大它并不会让图像更清晰，但配合extend_movement会增大可动范围 +## Installation(使用整合包请忽略) +可使用此安装方法对本项目二次开发 + +### 安装Anaconda（可选） +前往https://www.anaconda.com 安装Anaconda 并保证加入环境变量命令行可以找到，如： +``` +C:\Users\Eleven>conda --version +conda 24.11.3 +``` +以上方式可以保证conda环境被你选定的Ananconda统一管理。若没有此需求可以忽略直接运行下一步。 + +### 下载 TensorRT-RTX （Nvidia开发必选） +``` +curl -L -o trt_rtx.zip https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.3/TensorRT-RTX-1.3.0.35-win10-amd64-cuda-12.9-Release-external.zip +tar -xf trt_rtx.zip "TensorRT-RTX-1.3.0.35" && ren "TensorRT-RTX-1.3.0.35" "TensorRT-RTX-1.3.0.35_cu129" && del trt_rtx.zip +# Adding the bin folder to PATH +cd TensorRT-RTX-1.3.0.35_cu129\bin && set PATH=%CD%;%PATH% +``` +也可以在Window 环境变量管理器 GUI中将`TensorRT-RTX-1.3.0.35_cu129\bin`在`PATH`环境变量中添加来持久化 + +### 准备Conda 环境 +``` +conda create -y -n easyvtuber python=3.10 +conda activate easyvtuber +conda install conda-forge::pycuda +conda install -c nvidia/label/cuda-12.9.1 cuda-nvcc-dev_win-64 cudnn cuda-runtime +``` + +### 安装 TensorRT-RTX 的 Python Binding +``` +pip install TensorRT-RTX-1.3.0.35_cu129\python\tensorrt_rtx-1.3.0.35-cp310-none-win_amd64.whl +``` + +### 克隆项目和子项目 +``` +git clone https://github.com/yuyuyzl/EasyVtuber.git +cd EasyVtuber +git submodule init +git submodule update --recursive --remote +``` + +### 安装环境 +``` +pip install -r requirements.txt --no-warn-script-location +``` + +### 下载模型 +前往模型 release 地址 https://drive.google.com/file/d/1pWKIpjWeqfpa3Rub185FVvxDr5H09pOi/view?usp=drive_link 并解压到`data/models`文件夹下 + +### Pycharm 配置（可选） +安装完成后，在Pycharm内打开本项目，右下角解释器菜单点开，`Add Interpreter...`->`Conda Environment`->`Existing environment` +选好自己电脑上的`conda.exe`和刚才创建好的`easyvtuber`环境内的`python.exe` +点击OK，依赖全亮即可 + +### 运行启动器 +在Conda环境中执行以下命令 +`python launcher2.py` + +### 环境错误排查 +当运行TensorRT构建或启动器出现错误时，请参考如下可能性： +1. 环境安装有错误（一般运行时缺少库99%的原因都是这个，pip下载并不稳定，各种网络问题都可能导致安装失败），请检查源。 +2. 英伟达显卡但计算架构低于7.5（10系以及之前）不支持TensorRT ## FAQ @@ -258,9 +223,12 @@ RIFE插帧的计算消耗比起一次THA计算要少约50%且画面观感差异 A卡I卡使用 DirectML，这个框架与显卡厂商的显卡驱动及 Direct12 实现相关，很多算子实现存在数值错误或未实现并行优化被迫降级到cpu运行。方案：1. 将显卡驱动更新至最新，2. 打开调试启动器运行并检查模型或算子是否被降级到cpu实现，3.阅读上文的模型调参环节，切换不同模型和精度多尝试。 如有兴趣了解更多请参考[性能测试结果](PerformanceTest.md)末尾对DirectML的思考。 -> Q2: 我开启超分辨率后，为什么补帧没效果了？ +> Q2: 为什么无法找到我自己准备的Vtuber图像/使用自己的图像会报错？ -v0.5的实现中，补帧和超分只能二选一，开启超分后补帧会自动失效。欢迎重新在上文中更新下载v0.7之后的版本。 +1. 确认图像改为512x512的png，带Alpha透明通道即32比特位宽。 +2. 确认放置位置为`data/images`，在选项中应当能看到文件名。 +3. 仍然无法正确启动请尝试删除`launch.json`文件并重新启动软件。 +4. 99%的错误情况都来自于准备的图像不满足模型需求 > Q3: 我开启补帧后为什么在 Spout2 输出中边缘会有些抖动？ @@ -270,16 +238,30 @@ v0.5的实现中，补帧和超分只能二选一，开启超分后补帧会自 请在GPU页面中点击`Video Decode`或`Video Processing`按钮可以看到下拉，请对这里的选项挨个尝试找到正确的程序负载，一般在 `cuda` 或 `3D` 或 `Compute_X` 或 `Graphics_X`。 -> Q5: 我的配置是 CPU带核显加独立显卡，这个核显性能挺不错，打游戏时还用不上，能否把项目跑到这上面？ +> Q6: 本机有两张显卡，如何使用第二张副卡使用这个项目？ -完全可行，其实界面中留显卡选项有部分原因就是这方面需求。具体操作如下：去除TensorRT勾选，开启设备管理器观察核显占用，从`1`开始挨个尝试`GPU Device`选项并启动，直到观察到核显开始占用上升即找到了核显对应的编号。（os：DirectML的API就是这么难用没办法）接下来有可能出现Q1中提及的问题，请参照上文解决。 -因为性能问题很有可能需要使用插帧来提高输出帧数。目前核显相关测试样本较少，欢迎分享运行效果。 +在运行前配置环境变量`EZVTB_DEVICE_ID`为你想要运行的GPU ID,此变量缺省为`0` -> Q6: 我主机上插着两张英伟达显卡，如何使用第二张副卡使用这个项目？ -1. 对第二张卡构建TenosrRT：更改`02.构建TensorRT模型（N卡开启TensorRT使用）.bat` 中的 `python ezvtb_rt_interface.py` 一行，加上空格接参数cuda设备号，例如我的副卡设备号是`1`所以改成 `python ezvtb_rt_interface.py 1` 然后双击运行。 -2. 在启动器界面中选择相同的设备号。 +## References +``` +@inproceedings{huang2022rife, + title={Real-Time Intermediate Flow Estimation for Video Frame Interpolation}, + author={Huang, Zhewei and Zhang, Tianyuan and Heng, Wen and Shi, Boxin and Zhou, Shuchang}, + booktitle={Proceedings of the European Conference on Computer Vision (ECCV)}, + year={2022} +} +@misc{Khungurn:2022, + author = {Pramook Khungurn}, + title = {Talking Head(?) Anime from a Single Image 3: Now the Body Too}, + howpublished = {\url{http://pkhungurn.github.io/talking-head-anime-3/}}, + year = 2022, + note = {Accessed: YYYY-MM-DD}, +} -> Q7: 为什么无法使用我自己准备的Vtuber图像/使用自己的图像会报错？ -1. 确认图像改为512x512的png，带Alpha透明通道即32比特位宽。 -2. 确认放置位置为`data/images`，在选项中应当能看到文件名。 -3. 仍然无法正确启动请尝试删除`launch.json`文件并重新启动软件。 +@InProceedings{wang2021realesrgan, + author = {Xintao Wang and Liangbin Xie and Chao Dong and Ying Shan}, + title = {Real-ESRGAN: Training Real-World Blind Super-Resolution with Pure Synthetic Data}, + booktitle = {International Conference on Computer Vision Workshops (ICCVW)}, + date = {2021} +} +``` \ No newline at end of file diff --git a/args.py b/args.py deleted file mode 100644 index 2c2d6eb..0000000 --- a/args.py +++ /dev/null @@ -1,82 +0,0 @@ -import argparse -import re - - -def convert_to_byte(size): - result = re.search('(\d+\.?\d*)(b|kb|mb|gb|tb)', size.lower()) - if (result and result.groups()): - unit = result.groups()[1] - amount = float(result.groups()[0]) - index = ['b', 'kb', 'mb', 'gb', 'tb'].index(unit) - return amount * pow(1024, index) - raise ValueError("Invalid size provided, value is " + size) - - -parser = argparse.ArgumentParser() -parser.add_argument('--debug', action='store_true') -parser.add_argument('--eyebrow', action='store_true') -parser.add_argument('--extend_movement', type=float) -parser.add_argument('--input', type=str, default='cam') -parser.add_argument('--character', type=str, default='y') -parser.add_argument('--output_dir', type=str) -parser.add_argument('--output_webcam', type=str) -parser.add_argument('--output_size', type=str, default='512x512') -parser.add_argument('--model', type=str, default='standard_float') -parser.add_argument('--debug_input', action='store_true') -parser.add_argument('--mouse_input', type=str) -parser.add_argument('--perf', type=str) -parser.add_argument('--skip_model', action='store_true') -parser.add_argument('--ifm', type=str) -parser.add_argument('--osf', type=str) -parser.add_argument('--anime4k', action='store_true') -parser.add_argument('--alpha_split', action='store_true') -parser.add_argument('--bongo', action='store_true') -parser.add_argument('--cache', type=str, default='256mb') -parser.add_argument('--gpu_cache', type=str, default='512mb') -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('--model_seperable', action='store_true') -parser.add_argument('--model_half', action='store_true') -parser.add_argument('--model_cache', action='store_true') -parser.add_argument('--model_vram_cache', action='store_true') -# parser.add_argument('--model_cache_size', type=float, default=1.0) - -parser.add_argument('--use_interpolation', action='store_true') -parser.add_argument('--interpolation_scale', type=int, default=3) -parser.add_argument('--interpolation_half', action='store_true') - -parser.add_argument('--use_cacher', action='store_true') -parser.add_argument('--cacher_quality', type=int, default=85) -# parser.add_argument('--cacher_ram_size', type=float, default=1.0) - -parser.add_argument('--use_sr', action='store_true') -parser.add_argument('--sr_x4', action='store_true') -parser.add_argument('--sr_half', action='store_true') -parser.add_argument('--sr_noise', type=int, default=1) - -parser.add_argument('--frame_rate_limit', type=int, default=30) - -parser.add_argument('--alpha_clean', action='store_true') - -args = parser.parse_args() -# args.output_w = int(args.output_size.split('x')[0]) -# args.output_h = int(args.output_size.split('x')[1]) -if args.cache is not None: - args.max_cache_len = convert_to_byte(args.cache) / pow(1024, 3) # In gigabytes -else: - args.max_cache_len = 0 -if args.max_cache_len == 0: - args.use_cacher = False -if args.gpu_cache is not None: - args.max_gpu_cache_len = convert_to_byte(args.gpu_cache) / pow(1024, 3) # In gigabytes -else: - args.max_gpu_cache_len = 0 -if args.output_webcam is None and args.output_dir is None: args.debug = True - -args.model_output_size = 1024 if args.use_sr else 512 - -if args.use_sr and args.use_interpolation: - args.use_interpolation = False diff --git a/assets/complete_building_log.txt b/assets/complete_building_log.txt deleted file mode 100644 index dcc6465..0000000 --- a/assets/complete_building_log.txt +++ /dev/null @@ -1,548 +0,0 @@ -C:\Users\Eleven\Desktop\EasyVtuber-main>cd /D "C:\Users\Eleven\Desktop\EasyVtuber-main\" - -C:\Users\Eleven\Desktop\EasyVtuber-main>WHERE conda -INFO: Could not find files for the given pattern(s). - -C:\Users\Eleven\Desktop\EasyVtuber-main>IF 1 NEQ 0 ( -if not exist "envs" mkdir envs - if not exist "envs\miniconda3" mkdir envs\miniconda3 - IF not EXIST C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\Scripts ( - - mkdir C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3 - echo "Downloading miniconda..." - powershell -Command "(New-Object Net.WebClient).DownloadFile('https://repo.anaconda.com/miniconda/Miniconda3-py310_24.9.2-0-Windows-x86_64.exe', '.\envs\miniconda3.exe')" - echo "Installling minconda..." - start /wait "" C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3.exe /S /AddToPath=0 /RegisterPython=0 /InstallationType=JustMe /D=C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3 - echo "Successfully install minconda" - del envs\miniconda3.exe -) -) -"Downloading miniconda..." -"Installling minconda..." -"Successfully install minconda" - -C:\Users\Eleven\Desktop\EasyVtuber-main>IF EXIST C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\Scripts SET PATH=C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\Scripts;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6\bin;C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v12.6\libnvvp;;C:\Windows\system32;C:\Windows;C:\Windows\System32\Wbem;C:\Windows\System32\WindowsPowerShell\v1.0\;C:\Windows\System32\OpenSSH\;C:\Program Files (x86)\NVIDIA Corporation\PhysX\Common;C:\Program Files\Git\cmd;C:\Program Files\Tailscale\;C:\Program Files\NVIDIA Corporation\Nsight Compute 2024.3.2\;C:\Program Files\NVIDIA Corporation\NVIDIA NvDLISR;C:\Program Files\dotnet\;C:\Users\Eleven\AppData\Local\Programs\Python\Python312\Scripts\;C:\Users\Eleven\AppData\Local\Programs\Python\Python312\;C:\Users\Eleven\.cargo\bin;C:\Users\Eleven\AppData\Local\Microsoft\WindowsApps;C:\Users\Eleven\AppData\Local\Programs\Microsoft VS Code\bin;C:\Program Files\Git\bin;D:\ffmpeg\bin; - -C:\Users\Eleven\Desktop\EasyVtuber-main>call activate -# conda environments: -# - C:\EasyVtuber\envs\miniconda3 - C:\EasyVtuber\envs\miniconda3\envs\ezvtb_rt_venv -base * C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3 - -Retrieving notices: ...working... done -Channels: - - defaults -Platform: win-64 -Collecting package metadata (repodata.json): done -Solving environment: done - -## Package Plan ## - - environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3 - - -The following packages will be downloaded: - - package | build - ---------------------------|----------------- - anaconda-anon-usage-0.5.0 |py310hfc23b7f_100 28 KB - brotli-python-1.0.9 | py310h5da7b33_9 346 KB - ca-certificates-2024.11.26 | haa95532_0 132 KB - certifi-2024.12.14 | py310haa95532_0 162 KB - cffi-1.17.1 | py310h827c3e9_1 254 KB - conda-24.11.2 | py310haa95532_0 948 KB - conda-libmamba-solver-24.11.1| pyhd3eb1b0_0 39 KB - conda-package-handling-2.4.0| py310haa95532_0 297 KB - conda-package-streaming-0.11.0| py310haa95532_0 31 KB - cpp-expected-1.1.0 | h214f63a_0 131 KB - cryptography-43.0.3 | py310hbd6ee87_1 1.3 MB - libarchive-3.7.7 | h9243413_0 1.8 MB - libcurl-8.11.1 | haff574d_0 382 KB - libmamba-2.0.5 | hcd6fe79_1 4.7 MB - libmambapy-2.0.5 | py310h214f63a_1 459 KB - libsolv-0.7.30 | hf2fb9eb_1 473 KB - libssh2-1.11.1 | h2addb87_0 312 KB - libxml2-2.13.5 | h24da03e_0 2.9 MB - menuinst-2.2.0 | py310h5da7b33_0 200 KB - nlohmann_json-3.11.2 | h6c2663c_0 126 KB - packaging-24.2 | py310haa95532_0 175 KB - pluggy-1.5.0 | py310haa95532_0 39 KB - requests-2.32.3 | py310haa95532_1 101 KB - simdjson-3.10.1 | h214f63a_0 275 KB - spdlog-1.11.0 | h59b6b97_0 195 KB - vc-14.40 | haa95532_2 10 KB - vs2015_runtime-14.42.34433 | h9531ae6_2 1.2 MB - zstandard-0.23.0 | py310h4fc1ca9_1 342 KB - ------------------------------------------------------------ - Total: 17.3 MB - -The following NEW packages will be INSTALLED: - - cpp-expected pkgs/main/win-64::cpp-expected-1.1.0-h214f63a_0 - nlohmann_json pkgs/main/win-64::nlohmann_json-3.11.2-h6c2663c_0 - simdjson pkgs/main/win-64::simdjson-3.10.1-h214f63a_0 - spdlog pkgs/main/win-64::spdlog-1.11.0-h59b6b97_0 - -The following packages will be UPDATED: - - anaconda-anon-usa~ 0.4.4-py310hfc23b7f_100 --> 0.5.0-py310hfc23b7f_100 - brotli-python 1.0.9-py310hd77b12b_8 --> 1.0.9-py310h5da7b33_9 - ca-certificates 2024.9.24-haa95532_0 --> 2024.11.26-haa95532_0 - certifi 2024.8.30-py310haa95532_0 --> 2024.12.14-py310haa95532_0 - cffi 1.17.1-py310h827c3e9_0 --> 1.17.1-py310h827c3e9_1 - conda 24.9.2-py310haa95532_0 --> 24.11.2-py310haa95532_0 - conda-libmamba-so~ 24.9.0-pyhd3eb1b0_0 --> 24.11.1-pyhd3eb1b0_0 - conda-package-han~ 2.3.0-py310haa95532_0 --> 2.4.0-py310haa95532_0 - conda-package-str~ 0.10.0-py310haa95532_0 --> 0.11.0-py310haa95532_0 - cryptography 43.0.0-py310h89fc84f_0 --> 43.0.3-py310hbd6ee87_1 - libarchive 3.7.4-h9243413_0 --> 3.7.7-h9243413_0 - libcurl 8.9.1-h0416ee5_0 --> 8.11.1-haff574d_0 - libmamba 1.5.8-h99b1521_3 --> 2.0.5-hcd6fe79_1 - libmambapy 1.5.8-py310h77c03ed_3 --> 2.0.5-py310h214f63a_1 - libsolv 0.7.24-h23ce68f_1 --> 0.7.30-hf2fb9eb_1 - libssh2 1.11.0-h291bd65_0 --> 1.11.1-h2addb87_0 - libxml2 2.13.1-h24da03e_2 --> 2.13.5-h24da03e_0 - menuinst 2.1.2-py310h5da7b33_0 --> 2.2.0-py310h5da7b33_0 - packaging 24.1-py310haa95532_0 --> 24.2-py310haa95532_0 - pluggy 1.0.0-py310haa95532_1 --> 1.5.0-py310haa95532_0 - requests 2.32.3-py310haa95532_0 --> 2.32.3-py310haa95532_1 - vc 14.40-h2eaa2aa_1 --> 14.40-haa95532_2 - vs2015_runtime 14.40.33807-h98bb1dd_1 --> 14.42.34433-h9531ae6_2 - zstandard 0.23.0-py310h4fc1ca9_0 --> 0.23.0-py310h4fc1ca9_1 - - - -Downloading and Extracting Packages: - -Preparing transaction: done -Verifying transaction: done -Executing transaction: done -Channels: - - defaults -Platform: win-64 -Collecting package metadata (repodata.json): done -Solving environment: done - -## Package Plan ## - - environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv - - added / updated specs: - - python=3.10 - - -The following packages will be downloaded: - - package | build - ---------------------------|----------------- - python-3.10.16 | h4607a30_1 16.3 MB - ------------------------------------------------------------ - Total: 16.3 MB - -The following NEW packages will be INSTALLED: - - bzip2 pkgs/main/win-64::bzip2-1.0.8-h2bbff1b_6 - ca-certificates pkgs/main/win-64::ca-certificates-2024.11.26-haa95532_0 - libffi pkgs/main/win-64::libffi-3.4.4-hd77b12b_1 - openssl pkgs/main/win-64::openssl-3.0.15-h827c3e9_0 - pip pkgs/main/win-64::pip-24.2-py310haa95532_0 - python pkgs/main/win-64::python-3.10.16-h4607a30_1 - setuptools pkgs/main/win-64::setuptools-75.1.0-py310haa95532_0 - sqlite pkgs/main/win-64::sqlite-3.45.3-h2bbff1b_0 - tk pkgs/main/win-64::tk-8.6.14-h0416ee5_0 - tzdata pkgs/main/noarch::tzdata-2024b-h04d1e81_0 - vc pkgs/main/win-64::vc-14.40-haa95532_2 - vs2015_runtime pkgs/main/win-64::vs2015_runtime-14.42.34433-h9531ae6_2 - wheel pkgs/main/win-64::wheel-0.44.0-py310haa95532_0 - xz pkgs/main/win-64::xz-5.4.6-h8cc25b3_1 - zlib pkgs/main/win-64::zlib-1.2.13-h8cc25b3_1 - - - -Downloading and Extracting Packages: - -Preparing transaction: done -Verifying transaction: done -Executing transaction: done -# -# To activate this environment, use -# -# $ conda activate ezvtb_rt_venv -# -# To deactivate an active environment, use -# -# $ conda deactivate - - -# conda environments: -# - C:\EasyVtuber\envs\miniconda3 - C:\EasyVtuber\envs\miniconda3\envs\ezvtb_rt_venv -base C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3 -ezvtb_rt_venv * C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv - -Channels: - - https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - - defaults -Platform: win-64 -Collecting package metadata (repodata.json): done -Solving environment: done - -## Package Plan ## - - environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv - - added / updated specs: - - pycuda - - -The following packages will be downloaded: - - package | build - ---------------------------|----------------- - appdirs-1.4.4 | pyhd8ed1ab_1 14 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - boost-1.78.0 | py310h220cb41_4 555 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - boost-cpp-1.78.0 | h5b4e17d_0 17.0 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - ca-certificates-2024.12.14 | h56e8100_0 154 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - cudatoolkit-11.8.0 | h09e9e62_11 638.9 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - importlib-metadata-8.5.0 | pyha770c72_1 28 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - intel-openmp-2024.2.1 | h57928b3_1083 1.8 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - libblas-3.9.0 | 26_win64_mkl 3.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - libcblas-3.9.0 | 26_win64_mkl 3.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - libhwloc-2.9.1 | h51c2c0f_0 2.4 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - libiconv-1.17 | h8ffe710_0 698 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - liblapack-3.9.0 | 26_win64_mkl 3.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - lz4-c-1.9.4 | hcfcfb64_0 131 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - mako-1.3.8 | pyhd8ed1ab_0 66 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - markupsafe-2.1.2 | py310h8d17308_0 25 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - mkl-2024.2.2 | h66d3029_15 98.3 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - numpy-1.24.3 | py310hd02465a_0 5.7 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - openssl-3.1.0 | hcfcfb64_3 7.1 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - platformdirs-4.3.6 | pyhd8ed1ab_1 20 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - pthreads-win32-2.9.1 | hfa6e2cd_3 141 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - pycuda-2022.2.2 | py310ha2c4f5d_0 378 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - python_abi-3.10 | 2_cp310 4 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - pytools-2024.1.6 | pyhd8ed1ab_0 72 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - tbb-2021.9.0 | h91493d7_0 151 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - typing_extensions-4.12.2 | pyha770c72_1 39 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - ucrt-10.0.22621.0 | h57928b3_1 547 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - zipp-3.21.0 | pyhd8ed1ab_1 21 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge - ------------------------------------------------------------ - Total: 784.9 MB - -The following NEW packages will be INSTALLED: - - appdirs anaconda/cloud/conda-forge/noarch::appdirs-1.4.4-pyhd8ed1ab_1 - boost anaconda/cloud/conda-forge/win-64::boost-1.78.0-py310h220cb41_4 - boost-cpp anaconda/cloud/conda-forge/win-64::boost-cpp-1.78.0-h5b4e17d_0 - cudatoolkit anaconda/cloud/conda-forge/win-64::cudatoolkit-11.8.0-h09e9e62_11 - importlib-metadata anaconda/cloud/conda-forge/noarch::importlib-metadata-8.5.0-pyha770c72_1 - intel-openmp anaconda/cloud/conda-forge/win-64::intel-openmp-2024.2.1-h57928b3_1083 - libblas anaconda/cloud/conda-forge/win-64::libblas-3.9.0-26_win64_mkl - libcblas anaconda/cloud/conda-forge/win-64::libcblas-3.9.0-26_win64_mkl - libhwloc anaconda/cloud/conda-forge/win-64::libhwloc-2.9.1-h51c2c0f_0 - libiconv anaconda/cloud/conda-forge/win-64::libiconv-1.17-h8ffe710_0 - liblapack anaconda/cloud/conda-forge/win-64::liblapack-3.9.0-26_win64_mkl - libxml2 pkgs/main/win-64::libxml2-2.13.5-h24da03e_0 - lz4-c anaconda/cloud/conda-forge/win-64::lz4-c-1.9.4-hcfcfb64_0 - mako anaconda/cloud/conda-forge/noarch::mako-1.3.8-pyhd8ed1ab_0 - markupsafe anaconda/cloud/conda-forge/win-64::markupsafe-2.1.2-py310h8d17308_0 - mkl anaconda/cloud/conda-forge/win-64::mkl-2024.2.2-h66d3029_15 - numpy anaconda/cloud/conda-forge/win-64::numpy-1.24.3-py310hd02465a_0 - platformdirs anaconda/cloud/conda-forge/noarch::platformdirs-4.3.6-pyhd8ed1ab_1 - pthreads-win32 anaconda/cloud/conda-forge/win-64::pthreads-win32-2.9.1-hfa6e2cd_3 - pycuda anaconda/cloud/conda-forge/win-64::pycuda-2022.2.2-py310ha2c4f5d_0 - python_abi anaconda/cloud/conda-forge/win-64::python_abi-3.10-2_cp310 - pytools anaconda/cloud/conda-forge/noarch::pytools-2024.1.6-pyhd8ed1ab_0 - tbb anaconda/cloud/conda-forge/win-64::tbb-2021.9.0-h91493d7_0 - typing_extensions anaconda/cloud/conda-forge/noarch::typing_extensions-4.12.2-pyha770c72_1 - ucrt anaconda/cloud/conda-forge/win-64::ucrt-10.0.22621.0-h57928b3_1 - zipp anaconda/cloud/conda-forge/noarch::zipp-3.21.0-pyhd8ed1ab_1 - zstd pkgs/main/win-64::zstd-1.5.6-h8880b57_0 - -The following packages will be UPDATED: - - ca-certificates pkgs/main::ca-certificates-2024.11.26~ --> anaconda/cloud/conda-forge::ca-certificates-2024.12.14-h56e8100_0 - openssl pkgs/main::openssl-3.0.15-h827c3e9_0 --> anaconda/cloud/conda-forge::openssl-3.1.0-hcfcfb64_3 - - - -Downloading and Extracting Packages: - -Preparing transaction: done -Verifying transaction: done -Executing transaction: | "By downloading and using the CUDA Toolkit conda packages, you accept the terms and conditions of the CUDA End User License Agreement (EULA): https://docs.nvidia.com/cuda/eula/index.html" - -done -Channels: - - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - - defaults -Platform: win-64 -Collecting package metadata (repodata.json): done -Solving environment: done - -## Package Plan ## - - environment location: C:\Users\Eleven\Desktop\EasyVtuber-main\envs\miniconda3\envs\ezvtb_rt_venv - - added / updated specs: - - pytorch-cpu - - -The following packages will be downloaded: - - package | build - ---------------------------|----------------- - blas-1.0 | mkl 6 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - filelock-3.13.1 | py310haa95532_0 22 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - fsspec-2024.6.1 | py310haa95532_0 286 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - gmpy2-2.1.2 | py310h7f96b67_0 160 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - intel-openmp-2023.1.0 | h59b6b97_46320 2.7 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - jinja2-3.1.4 | py310haa95532_1 281 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - libuv-1.48.0 | h827c3e9_0 322 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mkl-2023.1.0 | h6b88ed4_46358 155.9 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mkl-service-2.4.0 | py310h2bbff1b_1 44 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mkl_fft-1.3.11 | py310h827c3e9_0 168 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mkl_random-1.2.8 | py310hc64d2fc_0 257 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mpc-1.1.0 | h7edee0f_1 260 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mpfr-4.0.2 | h62dcd97_1 1.5 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mpir-3.0.0 | hec2e145_1 1.3 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - mpmath-1.3.0 | py310haa95532_0 834 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - networkx-3.3 | py310haa95532_0 2.5 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - numpy-1.26.4 | py310h055cbcc_0 11 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - numpy-base-1.26.4 | py310h65a83cf_0 8.6 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - pytorch-2.3.1 |cpu_py310h0ce1571_0 111.2 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - pytorch-cpu-2.3.1 | haa95532_0 21 KB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - sympy-1.13.3 | py310haa95532_0 11.3 MB https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main - ------------------------------------------------------------ - Total: 297.6 MB - -The following NEW packages will be INSTALLED: - - blas anaconda/pkgs/main/win-64::blas-1.0-mkl - filelock anaconda/pkgs/main/win-64::filelock-3.13.1-py310haa95532_0 - fsspec anaconda/pkgs/main/win-64::fsspec-2024.6.1-py310haa95532_0 - gmpy2 anaconda/pkgs/main/win-64::gmpy2-2.1.2-py310h7f96b67_0 - jinja2 anaconda/pkgs/main/win-64::jinja2-3.1.4-py310haa95532_1 - libuv anaconda/pkgs/main/win-64::libuv-1.48.0-h827c3e9_0 - mkl-service anaconda/pkgs/main/win-64::mkl-service-2.4.0-py310h2bbff1b_1 - mkl_fft anaconda/pkgs/main/win-64::mkl_fft-1.3.11-py310h827c3e9_0 - mkl_random anaconda/pkgs/main/win-64::mkl_random-1.2.8-py310hc64d2fc_0 - mpc anaconda/pkgs/main/win-64::mpc-1.1.0-h7edee0f_1 - mpfr anaconda/pkgs/main/win-64::mpfr-4.0.2-h62dcd97_1 - mpir anaconda/pkgs/main/win-64::mpir-3.0.0-hec2e145_1 - mpmath anaconda/pkgs/main/win-64::mpmath-1.3.0-py310haa95532_0 - networkx anaconda/pkgs/main/win-64::networkx-3.3-py310haa95532_0 - numpy-base anaconda/pkgs/main/win-64::numpy-base-1.26.4-py310h65a83cf_0 - pytorch anaconda/pkgs/main/win-64::pytorch-2.3.1-cpu_py310h0ce1571_0 - pytorch-cpu anaconda/pkgs/main/win-64::pytorch-cpu-2.3.1-haa95532_0 - sympy anaconda/pkgs/main/win-64::sympy-1.13.3-py310haa95532_0 - -The following packages will be REMOVED: - - libblas-3.9.0-26_win64_mkl - libcblas-3.9.0-26_win64_mkl - liblapack-3.9.0-26_win64_mkl - -The following packages will be UPDATED: - - numpy anaconda/cloud/conda-forge::numpy-1.2~ --> anaconda/pkgs/main::numpy-1.26.4-py310h055cbcc_0 - -The following packages will be SUPERSEDED by a higher-priority channel: - - intel-openmp anaconda/cloud/conda-forge::intel-ope~ --> anaconda/pkgs/main::intel-openmp-2023.1.0-h59b6b97_46320 - mkl anaconda/cloud/conda-forge::mkl-2024.~ --> anaconda/pkgs/main::mkl-2023.1.0-h6b88ed4_46358 - - - -Downloading and Extracting Packages: - -Preparing transaction: done -Verifying transaction: done -Executing transaction: done -Looking in indexes: https://mirrors.aliyun.com/pypi/simple/ -Requirement already satisfied: pip in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (24.2) -Collecting pip - Downloading https://mirrors.aliyun.com/pypi/packages/ef/7d/500c9ad20238fcfcb4cb9243eede163594d7020ce87bd9610c9e02771876/pip-24.3.1-py3-none-any.whl (1.8 MB) - ---------------------------------------- 1.8/1.8 MB 1.4 MB/s eta 0:00:00 -Requirement already satisfied: wheel in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (0.44.0) -Collecting wheel - Downloading https://mirrors.aliyun.com/pypi/packages/0b/2c/87f3254fd8ffd29e4c02732eee68a83a1d3c346ae39bc6822dcbcb697f2b/wheel-0.45.1-py3-none-any.whl (72 kB) -Installing collected packages: wheel, pip - Attempting uninstall: wheel - Found existing installation: wheel 0.44.0 - Uninstalling wheel-0.44.0: - Successfully uninstalled wheel-0.44.0 - Attempting uninstall: pip - Found existing installation: pip 24.2 - Uninstalling pip-24.2: - Successfully uninstalled pip-24.2 -Successfully installed pip-24.3.1 wheel-0.45.1 -Looking in indexes: https://mirrors.aliyun.com/pypi/simple/ -Collecting nvidia-cudnn-cu12 - Downloading https://mirrors.aliyun.com/pypi/packages/07/1d/3c7055e1d189f55ed25183f50e1d82f477024289f925df22ac89c0f4d387/nvidia_cudnn_cu12-9.6.0.74-py3-none-win_amd64.whl (502.8 MB) - ---------------------------------------- 502.8/502.8 MB 6.3 MB/s eta 0:00:00 -Collecting nvidia-cublas-cu12 (from nvidia-cudnn-cu12) - Downloading https://mirrors.aliyun.com/pypi/packages/84/f7/985e9bdbe3e0ac9298fcc8cfa51a392862a46a0ffaccbbd56939b62a9c83/nvidia_cublas_cu12-12.6.4.1-py3-none-win_amd64.whl (434.5 MB) - ---------------------------------------- 434.5/434.5 MB 6.2 MB/s eta 0:00:00 -Installing collected packages: nvidia-cublas-cu12, nvidia-cudnn-cu12 -Successfully installed nvidia-cublas-cu12-12.6.4.1 nvidia-cudnn-cu12-9.6.0.74 -Looking in indexes: https://pypi.org/simple, https://pypi.nvidia.com -Collecting tensorrt_cu12_libs==10.6.0 - Downloading https://pypi.nvidia.com/tensorrt-cu12-libs/tensorrt_cu12_libs-10.6.0-py2.py3-none-win_amd64.whl (1075.0 MB) - ---------------------------------------- 1.1/1.1 GB 8.2 MB/s eta 0:00:00 -Collecting tensorrt_cu12_bindings==10.6.0 - Downloading https://pypi.nvidia.com/tensorrt-cu12-bindings/tensorrt_cu12_bindings-10.6.0-cp310-none-win_amd64.whl (814 kB) - ---------------------------------------- 814.1/814.1 kB 12.0 MB/s eta 0:00:00 -Collecting tensorrt==10.6.0 - Downloading https://pypi.nvidia.com/tensorrt/tensorrt-10.6.0.tar.gz (16 kB) - Preparing metadata (setup.py) ... done -Collecting nvidia-cuda-runtime-cu12 (from tensorrt_cu12_libs==10.6.0) - Downloading https://pypi.nvidia.com/nvidia-cuda-runtime-cu12/nvidia_cuda_runtime_cu12-12.6.77-py3-none-win_amd64.whl (891 kB) - ---------------------------------------- 891.8/891.8 kB 10.2 MB/s eta 0:00:00 -Collecting tensorrt-cu12==10.6.0 (from tensorrt==10.6.0) - Downloading https://pypi.nvidia.com/tensorrt-cu12/tensorrt-cu12-10.6.0.tar.gz (18 kB) - Preparing metadata (setup.py) ... done -Building wheels for collected packages: tensorrt, tensorrt-cu12 - Building wheel for tensorrt (setup.py) ... done - Created wheel for tensorrt: filename=tensorrt-10.6.0-py2.py3-none-any.whl size=16345 sha256=3fc94cd935389ea9535e279dc81bdf07b97da751150f64a1f45c39ba8b5c2663 - Stored in directory: c:\users\eleven\appdata\local\pip\cache\wheels\42\bd\86\598d26504705081ad82ae23a08a869322ee5a8bdb410d5d298 - Building wheel for tensorrt-cu12 (setup.py) ... done - Created wheel for tensorrt-cu12: filename=tensorrt_cu12-10.6.0-py2.py3-none-any.whl size=17568 sha256=1856795aec69a0f6767ab62c0962d586ea16c3301c0f7040cde242c5f92f9abc - Stored in directory: c:\users\eleven\appdata\local\pip\cache\wheels\13\99\06\326119ed088736cbd6ded866ada42bb29bbb67bed34903e57b -Successfully built tensorrt tensorrt-cu12 -Installing collected packages: tensorrt_cu12_bindings, tensorrt-cu12, nvidia-cuda-runtime-cu12, tensorrt_cu12_libs, tensorrt -Successfully installed nvidia-cuda-runtime-cu12-12.6.77 tensorrt-10.6.0 tensorrt-cu12-10.6.0 tensorrt_cu12_bindings-10.6.0 tensorrt_cu12_libs-10.6.0 -Looking in indexes: https://mirrors.aliyun.com/pypi/simple/ -Requirement already satisfied: numpy==1.26.4 in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from -r requirements.txt (line 1)) (1.26.4) -Collecting scipy (from -r requirements.txt (line 2)) - Downloading https://mirrors.aliyun.com/pypi/packages/51/8c/c2d371111961f737ae08881f654cf54eca796c42ec0429add2a06df97049/scipy-1.15.0-cp310-cp310-win_amd64.whl (43.9 MB) - ---------------------------------------- 43.9/43.9 MB 6.5 MB/s eta 0:00:00 -Collecting matplotlib (from -r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/5a/85/b9a54d64585a6b8737a78a61897450403c30f39e0bd3214270bb0b96f002/matplotlib-3.10.0-cp310-cp310-win_amd64.whl (8.0 MB) - ---------------------------------------- 8.0/8.0 MB 6.7 MB/s eta 0:00:00 -Collecting protobuf==3.20.1 (from -r requirements.txt (line 4)) - Downloading https://mirrors.aliyun.com/pypi/packages/c0/9c/bb88091287418ae1cf8af2bb9ed9710748a562b9abc227e4884d687a8650/protobuf-3.20.1-cp310-cp310-win_amd64.whl (903 kB) - ---------------------------------------- 903.8/903.8 kB 6.8 MB/s eta 0:00:00 -Collecting pynput==1.7.6 (from -r requirements.txt (line 5)) - Downloading https://mirrors.aliyun.com/pypi/packages/02/27/4de87850ff87c8dcecaaf8d27f28cec89ef17eeb6938f250449cb2635e06/pynput-1.7.6-py2.py3-none-any.whl (89 kB) -Collecting mediapipe (from -r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/f6/7a/448885fc960a972d8bf1939602919fdaa8de17f505b6c50782a074e7b70e/mediapipe-0.10.20-cp310-cp310-win_amd64.whl (51.0 MB) - ---------------------------------------- 51.0/51.0 MB 6.2 MB/s eta 0:00:00 -Collecting opencv_python==4.5.5.64 (from -r requirements.txt (line 7)) - Downloading https://mirrors.aliyun.com/pypi/packages/48/c3/798bd7b8f78430f82ec0660b753106717e4e4bb8032ce56f77d8577cb24b/opencv_python-4.5.5.64-cp36-abi3-win_amd64.whl (35.4 MB) - ---------------------------------------- 35.4/35.4 MB 6.3 MB/s eta 0:00:00 -Collecting Pillow==9.1.0 (from -r requirements.txt (line 8)) - Downloading https://mirrors.aliyun.com/pypi/packages/a8/df/1177786a2d1c0bf732ba6d5f05a2fa40f016e81e1c16d62f1101e35d271e/Pillow-9.1.0-cp310-cp310-win_amd64.whl (3.3 MB) - ---------------------------------------- 3.3/3.3 MB 6.4 MB/s eta 0:00:00 -Collecting pyanime4k==2.5.2 (from -r requirements.txt (line 9)) - Downloading https://mirrors.aliyun.com/pypi/packages/81/af/25a4a2e9a3d46820c016f6bb2cc235d953f104546a079c8c67bf09548734/pyanime4k-2.5.2-py3-none-win_amd64.whl (35.4 MB) - ---------------------------------------- 35.4/35.4 MB 6.3 MB/s eta 0:00:00 -Collecting pyvirtualcam==0.9.1 (from -r requirements.txt (line 10)) - Downloading https://mirrors.aliyun.com/pypi/packages/ca/59/e245e2d9832c6cefd212133d5c1ef988e66e7ccd93f4381017e844949334/pyvirtualcam-0.9.1-cp310-cp310-win_amd64.whl (158 kB) -Collecting pyopengl (from -r requirements.txt (line 11)) - Downloading https://mirrors.aliyun.com/pypi/packages/99/48/00e31747821d3fc56faddd00a4725454d1e694a8b67d715cf20f531506a5/PyOpenGL-3.1.7-py3-none-any.whl (2.4 MB) - ---------------------------------------- 2.4/2.4 MB 6.3 MB/s eta 0:00:00 -Collecting onnx (from -r requirements.txt (line 12)) - Downloading https://mirrors.aliyun.com/pypi/packages/0e/d3/d26ebf590a65686dde6b27fef32493026c5be9e42083340d947395f93405/onnx-1.17.0-cp310-cp310-win_amd64.whl (14.5 MB) - ---------------------------------------- 14.5/14.5 MB 6.0 MB/s eta 0:00:00 -Collecting onnxruntime-directml (from -r requirements.txt (line 13)) - Downloading https://mirrors.aliyun.com/pypi/packages/3c/4f/f433239b05304aa9af0217da20508abbbcec1dcd58ee821e3dab8939ecfe/onnxruntime_directml-1.20.1-cp310-cp310-win_amd64.whl (22.8 MB) - ---------------------------------------- 22.8/22.8 MB 6.3 MB/s eta 0:00:00 -Collecting turbojpeg (from -r requirements.txt (line 14)) - Downloading https://mirrors.aliyun.com/pypi/packages/a0/06/1a6ca4e7eb9c8342e5991093a36dcde37033a8f8da58dcf680282bb2b733/turbojpeg-0.0.2-cp310-cp310-win_amd64.whl (439 kB) -Collecting tqdm (from -r requirements.txt (line 15)) - Downloading https://mirrors.aliyun.com/pypi/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl (78 kB) -Collecting six (from pynput==1.7.6->-r requirements.txt (line 5)) - Downloading https://mirrors.aliyun.com/pypi/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl (11 kB) -Collecting ffmpeg-python>=0.2.0 (from pyanime4k==2.5.2->-r requirements.txt (line 9)) - Downloading https://mirrors.aliyun.com/pypi/packages/d7/0c/56be52741f75bad4dc6555991fabd2e07b432d333da82c11ad701123888a/ffmpeg_python-0.2.0-py3-none-any.whl (25 kB) -Collecting contourpy>=1.0.1 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/e0/82/c372be3fc000a3b2005061ca623a0d1ecd2eaafb10d9e883a2fc8566e951/contourpy-1.3.1-cp310-cp310-win_amd64.whl (218 kB) -Collecting cycler>=0.10 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl (8.3 kB) -Collecting fonttools>=4.22.0 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/19/d1/4dcd865360fb2c499749a913fe80e41c26e8ae18629d87dfffa3de27e831/fonttools-4.55.3-cp310-cp310-win_amd64.whl (2.2 MB) - ---------------------------------------- 2.2/2.2 MB 6.3 MB/s eta 0:00:00 -Collecting kiwisolver>=1.3.1 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/98/99/0dd05071654aa44fe5d5e350729961e7bb535372935a45ac89a8924316e6/kiwisolver-1.4.8-cp310-cp310-win_amd64.whl (71 kB) -Collecting packaging>=20.0 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/88/ef/eb23f262cca3c0c4eb7ab1933c3b1f03d021f2c48f54763065b6f0e321be/packaging-24.2-py3-none-any.whl (65 kB) -Collecting pyparsing>=2.3.1 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/1c/a7/c8a2d361bf89c0d9577c934ebb7421b25dc84bf3a8e3ac0a40aed9acc547/pyparsing-3.2.1-py3-none-any.whl (107 kB) -Collecting python-dateutil>=2.7 (from matplotlib->-r requirements.txt (line 3)) - Downloading https://mirrors.aliyun.com/pypi/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl (229 kB) -Collecting absl-py (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/a2/ad/e0d3c824784ff121c03cc031f944bc7e139a8f1870ffd2845cc2dd76f6c4/absl_py-2.1.0-py3-none-any.whl (133 kB) -Collecting attrs>=19.1.0 (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/89/aa/ab0f7891a01eeb2d2e338ae8fecbe57fcebea1a24dbb64d45801bfab481d/attrs-24.3.0-py3-none-any.whl (63 kB) -Collecting flatbuffers>=2.0 (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/fb/b4/31c461eef98b96b8ab736d97274548eaf2b2e349bf09e4de3902f7d53084/flatbuffers-24.12.23-py2.py3-none-any.whl (30 kB) -Collecting jax (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/22/49/b4418a7a892c0dd64442bbbeef54e1cdfe722dfc5a7bf0d611d3f5f90e99/jax-0.4.38-py3-none-any.whl (2.2 MB) - ---------------------------------------- 2.2/2.2 MB 6.3 MB/s eta 0:00:00 -Collecting jaxlib (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/0d/04/120c4caac6151f7297fedf9dd776362aa2d417d3f87bda826050b4da45e8/jaxlib-0.4.38-cp310-cp310-win_amd64.whl (64.2 MB) - ---------------------------------------- 64.2/64.2 MB 6.2 MB/s eta 0:00:00 -WARNING: Retrying (Retry(total=4, connect=None, read=None, redirect=None, status=None)) after connection broken by 'ReadTimeoutError("HTTPSConnectionPool(host='mirrors.aliyun.com', port=443): Read timed out. (read timeout=15)")': /pypi/simple/opencv-contrib-python/ -Collecting opencv-contrib-python (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/a7/9e/7110d2c5d543ab03b9581dbb1f8e2429863e44e0c9b4960b766f230c1279/opencv_contrib_python-4.10.0.84-cp37-abi3-win_amd64.whl (45.5 MB) - ---------------------------------------- 45.5/45.5 MB 6.3 MB/s eta 0:00:00 -INFO: pip is looking at multiple versions of mediapipe to determine which version is compatible with other requirements. This could take a while. -Collecting mediapipe (from -r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/d0/ca/4a32e95f1767467f7c461d7d5b4b506faebb1e30e35c1031c1b31cf82846/mediapipe-0.10.18-cp310-cp310-win_amd64.whl (50.9 MB) - ---------------------------------------- 50.9/50.9 MB 6.3 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/ea/1b/79e1ddc041012f5e5a7d60cbfd8cfe6037c5c87f4e77b4fdc0f9dfa2d44e/mediapipe-0.10.14-cp310-cp310-win_amd64.whl (50.8 MB) - ---------------------------------------- 50.8/50.8 MB 6.3 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/bb/b7/ffcdd19de3bf443864d0b0fed2b102275b96cbfbc1ddcf96ad811bbc1af0/mediapipe-0.10.13-cp310-cp310-win_amd64.whl (50.8 MB) - ---------------------------------------- 50.8/50.8 MB 6.4 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/12/50/9c24e158350d3f93be669db291fb452f21a25d874c94c5758374be82fff1/mediapipe-0.10.11-cp310-cp310-win_amd64.whl (50.8 MB) - ---------------------------------------- 50.8/50.8 MB 6.4 MB/s eta 0:00:00 -Collecting sounddevice>=0.4.4 (from mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/af/9b/15217b04f3b36d30de55fef542389d722de63f1ad81f9c72d8afc98cb6ab/sounddevice-0.5.1-py3-none-win_amd64.whl (363 kB) -INFO: pip is looking at multiple versions of onnx to determine which version is compatible with other requirements. This could take a while. -Collecting onnx (from -r requirements.txt (line 12)) - Downloading https://mirrors.aliyun.com/pypi/packages/0c/e4/1bc3ae56e6581587926a50a5c9dce3cfacf510e592e2512c7f5f2a9a4859/onnx-1.16.2-cp310-cp310-win_amd64.whl (14.4 MB) - ---------------------------------------- 14.4/14.4 MB 6.4 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/9c/7c/40fbebcb30f0fb6a773ca36e3a9dd4bdccd4b15455ef6c21a335781db78f/onnx-1.16.1-cp310-cp310-win_amd64.whl (14.4 MB) - ---------------------------------------- 14.4/14.4 MB 6.3 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/3b/e2/471ff83b3862967791d67f630000afce038756afbdf0665a3d767677c851/onnx-1.16.0-cp310-cp310-win_amd64.whl (14.4 MB) - ---------------------------------------- 14.4/14.4 MB 6.3 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/1b/43/6e84bf35a9201fb344b1a98edf7095c4fa1cf1478dfb6638d7b77f5475e6/onnx-1.15.0-cp310-cp310-win_amd64.whl (14.3 MB) - ---------------------------------------- 14.3/14.3 MB 6.3 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/00/36/e7a7e7a85564e7d409e4e8addfa11d41015d2190bfff30064771e7c21ca0/onnx-1.14.1-cp310-cp310-win_amd64.whl (13.3 MB) - ---------------------------------------- 13.3/13.3 MB 6.2 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/22/5c/46298252ea9f92b6b94184e8f001e575f2c346a22011498110fd032fc921/onnx-1.14.0-cp310-cp310-win_amd64.whl (13.3 MB) - ---------------------------------------- 13.3/13.3 MB 6.2 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/bb/cd/e2e63b1f0922dd45f6941a3422ccdb400360955064654a1a3a20e6a613fb/onnx-1.13.1-cp310-cp310-win_amd64.whl (12.2 MB) - ---------------------------------------- 12.2/12.2 MB 6.2 MB/s eta 0:00:00 -INFO: pip is still looking at multiple versions of onnx to determine which version is compatible with other requirements. This could take a while. - Downloading https://mirrors.aliyun.com/pypi/packages/97/e1/bdb42dde046507250083862520428901af3f5182faac50d3df5f13a4ff3d/onnx-1.13.0-cp310-cp310-win_amd64.whl (12.2 MB) - ---------------------------------------- 12.2/12.2 MB 6.5 MB/s eta 0:00:00 - Downloading https://mirrors.aliyun.com/pypi/packages/b0/f0/5e383f64293857e86c300cb70010cce1ca6f038f1563b85501d30404ad2a/onnx-1.12.0-cp310-cp310-win_amd64.whl (11.5 MB) - ---------------------------------------- 11.5/11.5 MB 6.2 MB/s eta 0:00:00 -Requirement already satisfied: typing-extensions>=3.6.2.1 in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from onnx->-r requirements.txt (line 12)) (4.12.2) -Collecting coloredlogs (from onnxruntime-directml->-r requirements.txt (line 13)) - Downloading https://mirrors.aliyun.com/pypi/packages/a7/06/3d6badcf13db419e25b07041d9c7b4a2c331d3f4e7134445ec5df57714cd/coloredlogs-15.0.1-py2.py3-none-any.whl (46 kB) -Requirement already satisfied: sympy in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from onnxruntime-directml->-r requirements.txt (line 13)) (1.13.3) -Collecting colorama (from tqdm->-r requirements.txt (line 15)) - Downloading https://mirrors.aliyun.com/pypi/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl (25 kB) -Collecting future (from ffmpeg-python>=0.2.0->pyanime4k==2.5.2->-r requirements.txt (line 9)) - Downloading https://mirrors.aliyun.com/pypi/packages/da/71/ae30dadffc90b9006d77af76b393cb9dfbfc9629f339fc1574a1c52e6806/future-1.0.0-py3-none-any.whl (491 kB) -Collecting CFFI>=1.0 (from sounddevice>=0.4.4->mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/d1/b6/0b0f5ab93b0df4acc49cae758c81fe4e5ef26c3ae2e10cc69249dfd8b3ab/cffi-1.17.1-cp310-cp310-win_amd64.whl (181 kB) -Collecting humanfriendly>=9.1 (from coloredlogs->onnxruntime-directml->-r requirements.txt (line 13)) - Downloading https://mirrors.aliyun.com/pypi/packages/f0/0f/310fb31e39e2d734ccaa2c0fb981ee41f7bd5056ce9bc29b2248bd569169/humanfriendly-10.0-py2.py3-none-any.whl (86 kB) -Collecting ml_dtypes>=0.4.0 (from jax->mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/6b/7e/bc54ae885e4d702e60a4bf50aa9066ff35e9c66b5213d11091f6bffb3036/ml_dtypes-0.5.1-cp310-cp310-win_amd64.whl (209 kB) -Collecting opt_einsum (from jax->mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/23/cd/066e86230ae37ed0be70aae89aabf03ca8d9f39c8aea0dec8029455b5540/opt_einsum-3.4.0-py3-none-any.whl (71 kB) -Requirement already satisfied: mpmath<1.4,>=1.1.0 in c:\users\eleven\desktop\easyvtuber-main\envs\miniconda3\envs\ezvtb_rt_venv\lib\site-packages (from sympy->onnxruntime-directml->-r requirements.txt (line 13)) (1.3.0) -Collecting pycparser (from CFFI>=1.0->sounddevice>=0.4.4->mediapipe->-r requirements.txt (line 6)) - Downloading https://mirrors.aliyun.com/pypi/packages/13/a3/a812df4e2dd5696d1f351d58b8fe16a405b234ad2886a0dab9183fb78109/pycparser-2.22-py3-none-any.whl (117 kB) -Collecting pyreadline3 (from humanfriendly>=9.1->coloredlogs->onnxruntime-directml->-r requirements.txt (line 13)) - Downloading https://mirrors.aliyun.com/pypi/packages/5a/dc/491b7661614ab97483abf2056be1deee4dc2490ecbf7bff9ab5cdbac86e1/pyreadline3-3.5.4-py3-none-any.whl (83 kB) -Installing collected packages: pyopengl, flatbuffers, turbojpeg, six, scipy, pyvirtualcam, pyreadline3, pyparsing, pycparser, protobuf, Pillow, packaging, opt_einsum, opencv_python, opencv-contrib-python, ml_dtypes, kiwisolver, future, fonttools, cycler, contourpy, colorama, attrs, absl-py, tqdm, python-dateutil, pynput, onnx, jaxlib, humanfriendly, ffmpeg-python, CFFI, sounddevice, pyanime4k, matplotlib, jax, coloredlogs, onnxruntime-directml, mediapipe -Successfully installed CFFI-1.17.1 Pillow-9.1.0 absl-py-2.1.0 attrs-24.3.0 colorama-0.4.6 coloredlogs-15.0.1 contourpy-1.3.1 cycler-0.12.1 ffmpeg-python-0.2.0 flatbuffers-24.12.23 fonttools-4.55.3 future-1.0.0 humanfriendly-10.0 jax-0.4.38 jaxlib-0.4.38 kiwisolver-1.4.8 matplotlib-3.10.0 mediapipe-0.10.11 ml_dtypes-0.5.1 onnx-1.12.0 onnxruntime-directml-1.20.1 opencv-contrib-python-4.10.0.84 opencv_python-4.5.5.64 opt_einsum-3.4.0 packaging-24.2 protobuf-3.20.1 pyanime4k-2.5.2 pycparser-2.22 pynput-1.7.6 pyopengl-3.1.7 pyparsing-3.2.1 pyreadline3-3.5.4 python-dateutil-2.9.0.post0 pyvirtualcam-0.9.1 scipy-1.15.0 six-1.17.0 sounddevice-0.5.1 tqdm-4.67.1 turbojpeg-0.0.2 -Press any key to continue . . . \ No newline at end of file diff --git a/character/amelria.png b/character/amelria.png deleted file mode 100644 index e5431e8..0000000 Binary files a/character/amelria.png and /dev/null differ diff --git a/character/ddolsun.png b/character/ddolsun.png deleted file mode 100644 index 42a8d91..0000000 Binary files a/character/ddolsun.png and /dev/null differ diff --git a/character/lostark.png b/character/lostark.png deleted file mode 100644 index 6b9ebf5..0000000 Binary files a/character/lostark.png and /dev/null differ diff --git a/character/luda.png b/character/luda.png deleted file mode 100644 index 873b89c..0000000 Binary files a/character/luda.png and /dev/null differ diff --git a/character/lupy.png b/character/lupy.png deleted file mode 100644 index b153bfc..0000000 Binary files a/character/lupy.png and /dev/null differ diff --git a/character/overwatch.png b/character/overwatch.png deleted file mode 100644 index e1f9bbb..0000000 Binary files a/character/overwatch.png and /dev/null differ diff --git a/character/y.png b/character/y.png deleted file mode 100644 index 6f1efa2..0000000 Binary files a/character/y.png and /dev/null differ diff --git a/ezvtb_rt_interface.py b/ezvtb_rt_interface.py deleted file mode 100644 index 0801eb4..0000000 --- a/ezvtb_rt_interface.py +++ /dev/null @@ -1,127 +0,0 @@ -import os -import sys -import subprocess -import urllib.request -import shutil - -def init_ezvtb_rt(search_dir = '.'): - if not os.path.exists(ezvtb_path) and not os.path.exists(ezvtb_main_path): #Check if rt not exist - zip_path = 'https://github.com/zpeng11/ezvtuber-rt/archive/refs/heads/main.zip' - filehandle, _ = urllib.request.urlretrieve(zip_path) - os.rename(filehandle, filehandle + '.zip') - shutil.unpack_archive(filehandle + '.zip' , '.') - -dir_path = os.path.dirname(os.path.realpath(__file__)) -ezvtb_path = os.path.join(dir_path, 'ezvtuber-rt') -ezvtb_main_path = os.path.join(dir_path, 'ezvtuber-rt-main') - -project_path = '' -if os.path.exists(ezvtb_path): - project_path = ezvtb_path -else: - project_path = ezvtb_main_path - -if project_path not in sys.path: - sys.path.append(project_path) - -def get_core( - #Device setting - device_id:int = 0, - use_tensorrt:bool = True, - #THA3 model setting - model_seperable:bool = True, - model_half:bool = True, #If using directml+half, there is small numerical error on Nvidia, and huge numerical error on AMD - model_cache_size:float = 1.0, #unit of GigaBytes, only works for tensorrt - model_use_eyebrow:bool = True, - #RIFE interpolation setting - use_interpolation:bool = True, - interpolation_scale:int = 2, - interpolation_half:bool = True, #If using directml+half, there is small numerical error on Nvidia, and huge numerical error on AMD - #Cacher setting - cacher_quality:int = 85, - cacher_ram_size:float = 2.0,#unit of GigaBytes - #SR setting - use_sr:bool = False, - sr_x4:bool = True, - sr_half:bool = True, - sr_noise:int = 1 - ): - support_trt = False - if use_tensorrt: #Verify tensorrt - print('Verifying TensorRT') - try: - from ezvtb_rt.trt_utils import cudaSetDevice, check_build_all_models - cudaSetDevice(device_id) - # check_build_all_models() - support_trt = True - except: - support_trt = False - if support_trt == False and use_tensorrt == True: - print('TensorRT option selected but not supported') - use_tensorrt = False - - if use_tensorrt: - os.environ['CUDA_DEVICE'] = str(device_id) - import pycuda.autoinit - print(f'Using devcie {pycuda.autoinit.device.name()}') - - tha_model_dir = os.path.join('.','data', 'models', 'tha3', - 'seperable' if model_seperable else 'standard', - 'fp16' if model_half else 'fp32') - rife_model_path = '' - if use_interpolation: - rife_model_path = os.path.join('.','data', 'models','rife_512', - f'x{interpolation_scale}', 'fp16' if interpolation_half else 'fp32') - - sr_model_path = '' - if use_sr: - if sr_x4: - if sr_half: - sr_model_path = os.path.join('.','data', 'models', 'Real-ESRGAN', 'exported_256_fp16') - else: - sr_model_path = os.path.join('.','data', 'models', 'Real-ESRGAN', 'exported_256') - else: #x2 - if sr_half: - sr_model_path = os.path.join('.','data', 'models', 'waifu2x_upconv', 'fp16', 'upconv_7', 'art', f'noise{sr_noise}_scale2x') - else: - sr_model_path = os.path.join('.','data', 'models', 'waifu2x_upconv', 'fp32', 'upconv_7', 'art', f'noise{sr_noise}_scale2x') - - print(f'THA3 Path:{tha_model_dir}') - print(f'RIFE Path:{rife_model_path}') - print(f'SR Path:{sr_model_path}') - - core = None - if use_tensorrt: - from ezvtb_rt.core import CoreTRT - core = CoreTRT(tha_model_dir, - vram_cache_size = model_cache_size, - use_eyebrow = model_use_eyebrow, - rife_dir = rife_model_path if len(rife_model_path) > 0 else None, - sr_dir = sr_model_path if len(sr_model_path) > 0 else None, - cache_max_volume = cacher_ram_size, - cache_quality = cacher_quality, - ) - else: #Use directml - from ezvtb_rt.core_ort import CoreORT - core = CoreORT(tha_model_dir, - rife_path = rife_model_path if len(rife_model_path) > 0 else None, - sr_path = sr_model_path if len(sr_model_path) > 0 else None, - device_id=device_id, - cache_max_volume = cacher_ram_size, - cache_quality = cacher_quality, - use_eyebrow = model_use_eyebrow) - - return core - -if __name__ == '__main__': - init_ezvtb_rt() - from ezvtb_rt.trt_utils import check_build_all_models, cudaSetDevice - device_id = 0 - import sys - if len(sys.argv) > 1: - device_id = int(sys.argv[1]) - try: - cudaSetDevice(device_id) - check_build_all_models('./data/models') - except: - pass diff --git a/ezvtuber-rt b/ezvtuber-rt index d602899..07c3029 160000 --- a/ezvtuber-rt +++ b/ezvtuber-rt @@ -1 +1 @@ -Subproject commit d602899f2dde6561d6d5dfbbfa908117c10900b3 +Subproject commit 07c302921e67b483b023a7db5c0fa96ac5fd8e6d diff --git a/launcher2.py b/launcher2.py index fdd90aa..c30d4ca 100644 --- a/launcher2.py +++ b/launcher2.py @@ -1,6 +1,7 @@ import ctypes import os import subprocess +import threading import wx import json @@ -29,26 +30,31 @@ cache_simplify_quality_map = { } default_arg = { 'character': 'lambda_00', - 'input': 2, + 'input': 3, 'output': 2, 'ifm': None, 'osf': '127.0.0.1:11573', + 'min_cutoff': 50, + 'beta': 80, 'is_extend_movement': False, 'is_alpha_split': False, 'is_bongo': False, 'is_alpha_clean': False, 'is_eyebrow': False, 'cache_simplify': 'High', - 'cache_compression': 'High', 'ram_cache_size': '2gb', 'vram_cache_size': '2gb', 'model_select': 'seperable_half', - 'interpolation': 'x2_half', + 'interpolation': "Off", 'frame_rate_limit': '30', - 'sr': 'Off', - 'device_id': '0', + 'sr': "Off", 'use_tensorrt': False, - 'preset': 'Medium' + 'preset': 'Low', + 'mouse_audio_input': False, + 'audio_sensitivity': '0.02', + 'audio_threshold': '10.0', + 'blink_interval': '5.0', + 'breath_cycle': 'inf' } try: @@ -64,15 +70,17 @@ finally: p = None dirPath = 'data/images' characterList = [] +studentModelList = [] +studentModelCharacterMap = {} -hasRTModel = False -try: - f = open('data/models/tha3/standard/fp16/decomposer.trt') - f.close() - hasRTModel = True -except: - pass - +def is_nvidia_gpu(): + try: + # 获取显卡名称列表 + output = subprocess.check_output("wmic path Win32_VideoController get Name", shell=True).decode('gbk') + return "NVIDIA" in output.upper() + except Exception: + return False +hasTRTSupport = is_nvidia_gpu() def refreshList(): global characterList @@ -82,11 +90,66 @@ 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() + +def min_cutoff_mapper(value, revert=False): + """ + 非线性映射函数：0-100整数 <-> 0-100浮点数 + 使用平方函数，使得越接近0数字越密集 + """ + if revert: + # 浮点 -> 整数: 使用平方根反向映射 + return int((value / 100.0) ** 0.5 * 100) + # 整数 -> 浮点: 使用平方映射 + return (value / 100.0) ** 2 * 100.0 + + +def beta_mapper(value, revert=False): + """ + 非线性映射函数：0-100整数 <-> 0-1浮点数 + 使用平方函数，使得越接近0数字越密集 + """ + if revert: + # 浮点 -> 整数: 使用平方根反向映射 + return int((value ** 0.5) * 100) + # 整数 -> 浮点: 使用平方映射 + return (value / 100.0) ** 2 class OptionPanel(wx.Panel): - def __init__(self, parent, title='', desc='', choices=None, mapping=None, type=0, default=None, disabled=False): + def __init__(self, parent, title='', desc='', choices=None, mapping=None, type=0, default=None, disabled=False, mapper=min_cutoff_mapper): wx.Panel.__init__(self, parent) self.type = type if mapping is not None: @@ -140,8 +203,38 @@ class OptionPanel(wx.Panel): self.control.SetValue(default) except: pass + elif self.type == 3: + # Slider type for float values 0.0 to 1.0 + sliderSizer = wx.BoxSizer(wx.HORIZONTAL) + self.control = wx.Slider(self, wx.ID_ANY, value=50, minValue=0, maxValue=100, + style=wx.SL_HORIZONTAL) + self.control.SetMinSize(wx.Size(250, -1)) + + # Add a label to show the float value + self.valueLabel = wx.StaticText(self, wx.ID_ANY, "0.50") + self.valueLabel.SetMinSize(wx.Size(50, -1)) + + try: + if default is not None: + self.control.SetValue(default) + self.valueLabel.SetLabelText(f"{mapper(default):.4f}") + except: + pass + + # Update label when slider changes + def onSliderChange(event): + val = mapper(self.control.GetValue()) + self.valueLabel.SetLabelText(f"{val:.4f}") + self.control.Bind(wx.EVT_SLIDER, onSliderChange) + + sliderSizer.Add(self.control, 1, wx.ALIGN_CENTER_VERTICAL) + sliderSizer.Add(self.valueLabel, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 10) + mainSizer.Add(sliderSizer, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 20) + # Skip the normal control addition below + self.control._slider_added = True - mainSizer.Add(self.control, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 20) + if not (self.type == 3 and hasattr(self.control, '_slider_added')): + mainSizer.Add(self.control, 0, wx.ALIGN_CENTER_VERTICAL | wx.LEFT, 20) # self.SetBackgroundColour('#000000') def GetValue(self): @@ -151,18 +244,103 @@ class OptionPanel(wx.Panel): ret = self.control.GetValue() elif self.type == 2: ret = self.control.GetValue() + elif self.type == 3: + ret = self.control.GetValue() if self.mapping is not None: return self.mapping[ret] else: return ret +def _important_log_line(line): + """从 main 的一行日志中提取“重要”的简短描述，用于状态栏；无关行返回 None。""" + line = line.strip() + if not line: + return None + if line.startswith('Launched:'): + return 'Launched' + if 'Model Inference Ready' in line: + return 'Model Inference Ready' + # TRT: Building engine from ONNX: ...\filename.onnx + if '[TRT]' in line and 'Building engine from ONNX:' in line: + idx = line.find('Building engine from ONNX:') + if idx != -1: + path = line[idx + len('Building engine from ONNX:'):].strip().rstrip('\r\n') + name = os.path.basename(path) + if name: + return f'Building: {name}' + # TRT: Loading ONNX file from path ...\filename.onnx + if '[TRT]' in line and 'Loading ONNX file from path' in line: + idx = line.find('Loading ONNX file from path') + if idx != -1: + path = line[idx + len('Loading ONNX file from path'):].strip().strip('.').strip().rstrip('\r\n') + name = os.path.basename(path) + if name: + return f'Loading: {name}' + # ORT: Loading ONNX model from path ...\filename.onnx + if '[ORT]' in line and 'Loading ONNX model from path' in line: + idx = line.find('Loading ONNX model from path') + if idx != -1: + path = line[idx + len('Loading ONNX model from path'):].strip().strip('.').strip().rstrip('\r\n') + name = os.path.basename(path) + if name: + return f'Loading: {name}' + # ORT: Completed loading session: xxx.onnx + if '[ORT]' in line and 'Completed loading session:' in line: + idx = line.find('Completed loading session:') + if idx != -1: + name = line[idx + len('Completed loading session:'):].strip().rstrip('\r\n') + if name: + return f'Loaded: {name}' + return None + + +def _on_main_log_line(panel, line): + """在子线程中调用：若该行是重要日志，则用 wx.CallAfter 更新 panel 的状态框。""" + display = _important_log_line(line) + if display is not None: + wx.CallAfter(panel.statusCtrl.SetValue, display) + + +def _read_pipe_to_stream(pipe, dest_stream, out_lines=None, on_line_callback=None): + """从 pipe 读行，写回 dest_stream，可选追加到 out_lines，并对每行调用 on_line_callback(line)。""" + if pipe is None: + return + try: + for raw in iter(pipe.readline, b''): + try: + text = raw.decode('utf-8', errors='replace') + except Exception: + text = raw.decode('gbk', errors='replace') + if out_lines is not None: + out_lines.append(text) + if on_line_callback is not None: + on_line_callback(text) + # 检查dest_stream是否可用（pythonw环境下可能不可用） + if dest_stream is not None: + try: + dest_stream.write(text) + dest_stream.flush() + except (AttributeError, OSError, ValueError): + # pythonw环境下sys.stdout/sys.stderr可能不可用，忽略错误 + pass + except Exception: + pass + finally: + try: + pipe.close() + except Exception: + pass + + class LauncherPanel(wx.Panel): def __init__(self, parent): wx.Panel.__init__(self, parent) self.number_of_buttons = 0 self.frame = parent self.optionDict = {} + self.main_output_lines = [] # main 的 stdout 副本 + self.main_stderr_lines = [] # main 的 stderr 副本 self.mainSizer = wx.BoxSizer(wx.VERTICAL) controlSizer = wx.BoxSizer(wx.HORIZONTAL) self.widgetSizer = wx.BoxSizer(wx.VERTICAL) @@ -178,7 +356,16 @@ class LauncherPanel(wx.Panel): f = f.MakeLarger() stVtuber.SetFont(f) controlSizer.Add(stEasy, 0, wx.ALL | wx.CENTER, 0) - controlSizer.Add(stVtuber, 1, wx.RIGHT | wx.CENTER, 30) + controlSizer.Add(stVtuber, 0, wx.RIGHT | wx.CENTER, 30) + + self.statusCtrl = wx.TextCtrl( + self, wx.ID_ANY, '', + style=wx.TE_READONLY | wx.BORDER_NONE | wx.TE_RIGHT, + ) + self.statusCtrl.SetHint('当前操作') + f = self.statusCtrl.GetFont() + self.statusCtrl.SetFont(f.Smaller()) + controlSizer.Add(self.statusCtrl, 1, wx.ALIGN_CENTER_VERTICAL | wx.LEFT | wx.RIGHT, 8) self.btnLaunch = wx.Button(self, label="Save && Launch") self.btnLaunch.Bind(wx.EVT_BUTTON, self.OnLaunch) @@ -209,27 +396,54 @@ class LauncherPanel(wx.Panel): addOption('input', title='Input Device', desc='选择希望使用的面捕数据源', choices=['iFacialMocap', 'OpenSeeFace', 'OpenCV(Webcam)', 'Mouse Input', 'Debug Input'], mapping=[0, 4, 1, 3, 2]) - addOption('ifm', title='iFacialMocap IP:Port', desc='输入iFacialMocap连接使用的IP:端口号', type=2) - addOption('is_eyebrow', title='Eyebrow', desc='使用眉毛输入，对性能有一定影响\n仅支持iFacialMocap', type=1, + addOption('ifm', title='iFacialMocap IP', desc='输入iFacialMocap连接使用的IP地址，默认连接 49983 端口', type=2) + addOption('is_eyebrow', title='Eyebrow', desc='使用眉毛输入，对性能有一定影响', type=1, default=True) addOption('osf', title='OpenSeeFace IP:Port', desc='输入OpenSeeFace连接使用的IP:端口号', type=2) + addOption('mouse_audio_input', title='Audio Input', desc='启用WASAPI音频输入控制嘴部动作', type=1) + addOption('audio_sensitivity', title='Audio Sensitivity', desc='音频灵敏度，控制音频对嘴部动作的影响程度', type=2) + addOption('audio_threshold', title='Audio Threshold', desc='音频阈值，低于此值的音频将被忽略', type=2) + addOption('blink_interval', title='Blink Interval', desc='设置眨眼间隔时间', + choices=['No Blink', '3 seconds', '5 seconds', '7 seconds'], + mapping=['inf', '3.0', '5.0', '7.0']) + addOption('min_cutoff', title='Min CutOff', desc='输入滤波频率截断，\n越小越平滑，越大静止时越灵敏', + type=3, mapper=min_cutoff_mapper) + addOption('beta', title='Beta', desc='输入滤波速度补偿，\n越小越平滑，越大运动时越灵敏', + type=3, mapper=beta_mapper) + + addOption('breath_cycle', title='Breath Cycle', desc='设置呼吸循环时间(会增加占用)', + choices=['No Breath', '3 seconds', '5 seconds', '7 seconds'], + mapping=['inf', '3.0', '5.0', '7.0']) addOption('output', title='Output', desc='选择输出目标', - choices=['Spout2', 'OBS VirtualCam', 'UnityCapture', 'Debug Output'], - mapping=[3, 1, 0, 2]) + choices=['Spout2', 'OBS VirtualCam', 'Debug Output'], + mapping=[0, 1, 2]) - addOption('device_id', title='GPU Device', desc='选择需要使用的计算设备(默认0)', choices=['0', '1', '2', '3']) - addOption('use_tensorrt', title='TensorRT', - desc='运行过预构建脚本后开启TensorRT可进一步提升性能\n（未构建成功时不可用，仅NVIDIA显卡支持）', + addOption('use_tensorrt', title='TensorRT加速', + desc='需要更长启动和预热时间（仅NVIDIA显卡支持）', type=1) addOption('frame_rate_limit', title='FPS Limit', desc='选择帧率限制目标', 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']) + + # 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']) @@ -239,9 +453,6 @@ class LauncherPanel(wx.Panel): addOption('cache_simplify', title='Input Simplify', desc='设置输入简化级别\n输入越简化，缓存命中率越高，动作越不平滑', choices=['Off', 'Low', 'Medium', 'High', 'Higher', 'Highest', 'Gaming']) - addOption('cache_compression', title='JPEG Compression', - desc='设置内存缓存压缩等级\n压缩等级越高，缓存命中率越高，输出质量越差', - choices=['Off', 'Low', 'Medium', 'High']) addOption('sr', title='SuperResolution', desc='选择使用的超分模型\n由于性能原因，real-esrgan会进行裁切', choices=['Off', 'anime4k_x2', 'waifu2x_x2_half', 'real-esrgan_x4_half', 'waifu2x_x2_full', @@ -270,11 +481,44 @@ class LauncherPanel(wx.Panel): self.optionSizer.Hide(self.optionDict['osf']) else: self.optionSizer.Show(self.optionDict['osf']) + self.optionSizer.Show(self.optionDict['is_eyebrow']) + if s != 1 and s != 4: + self.optionSizer.Hide(self.optionDict['min_cutoff']) + self.optionSizer.Hide(self.optionDict['beta']) + else: + self.optionSizer.Show(self.optionDict['min_cutoff']) + self.optionSizer.Show(self.optionDict['beta']) + # Show/hide audio input options for Mouse Input (s == 3) + if s != 3: + self.optionSizer.Hide(self.optionDict['mouse_audio_input']) + self.optionSizer.Hide(self.optionDict['audio_sensitivity']) + self.optionSizer.Hide(self.optionDict['audio_threshold']) + self.optionSizer.Hide(self.optionDict['blink_interval']) + else: + self.optionSizer.Show(self.optionDict['mouse_audio_input']) + self.optionSizer.Show(self.optionDict['blink_interval']) + # Update audio fields based on checkbox state + audioInputChoice() self.frame.fSizer.Layout() self.frame.Fit() + def audioInputChoice(e=None): + """Handle mouse_audio_input checkbox changes""" + enabled = self.optionDict['mouse_audio_input'].GetValue() + if enabled: + self.optionSizer.Show(self.optionDict['audio_sensitivity']) + self.optionSizer.Show(self.optionDict['audio_threshold']) + self.optionDict['audio_sensitivity'].control.Enable(True) + self.optionDict['audio_threshold'].control.Enable(True) + else: + self.optionSizer.Hide(self.optionDict['audio_sensitivity']) + self.optionSizer.Hide(self.optionDict['audio_threshold']) + self.frame.fSizer.Layout() + self.frame.Fit() + self.optionDict['input'].Bind(wx.EVT_CHOICE, inputChoice) + self.optionDict['mouse_audio_input'].Bind(wx.EVT_CHECKBOX, audioInputChoice) inputChoice() def presetChoice(e=None): @@ -284,13 +528,12 @@ class LauncherPanel(wx.Panel): self.optionDict['ram_cache_size'], self.optionDict['vram_cache_size'], self.optionDict['cache_simplify'], - self.optionDict['cache_compression'], ] presets = { - 'Low': [0, 1, 1, 5, 3], - 'Medium': [1, 1, 1, 4, 2], - 'High': [1, 2, 2, 2, 2], - 'Ultra': [3, 3, 3, 1, 1] + 'Low': [0, 1, 1, 5], + 'Medium': [1, 1, 1, 4], + 'High': [1, 2, 2, 2], + 'Ultra': [3, 3, 3, 1] } if s == 'Custom': @@ -299,7 +542,7 @@ class LauncherPanel(wx.Panel): for c in presetControls: self.optionSizer.Hide(c) if s in presets: opt = presets[s] - for i in range(5): presetControls[i].control.SetSelection(opt[i]) + for i in range(4): presetControls[i].control.SetSelection(opt[i]) self.frame.fSizer.Layout() self.frame.Fit() @@ -307,17 +550,54 @@ 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() + # 用控件当前选中字符串，避免 GetValue() 用旧 mapping 按下标取导致 IndexError + char_ctrl = self.optionDict['character'].control + tName = char_ctrl.GetStringSelection() if char_ctrl.GetSelection() >= 0 else '' refreshList() - self.optionDict['character'].control.SetItems(characterList) - self.optionDict['character'].control.SetSelection(characterList.index(tName)) + scanStudentModels() + self.optionDict['character'].mapping = characterList # 同步 mapping，后续 GetValue() 才正确 + char_ctrl.SetItems(characterList) + try: + idx = characterList.index(tName) + char_ctrl.SetSelection(idx) + except (ValueError, TypeError): + if characterList: + char_ctrl.SetSelection(0) + # 重新应用“当前模型 → 是否锁定 character”的规则 + onModelSelect() - if not hasRTModel: + if not hasTRTSupport: 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( + '需要NVIDIA显卡支持才能使用TensorRT') self.frame.Bind(wx.EVT_ACTIVATE, onActivate) @@ -332,62 +612,85 @@ class LauncherPanel(wx.Panel): self.btnLaunch.SetLabelText('Working...') if p is not None: - subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL) + creation_flags = 0 + if sys.platform == 'win32': + # CREATE_NO_WINDOW = 0x08000000 + creation_flags = 0x08000000 + subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=creation_flags) p = None + self.statusCtrl.Clear() self.btnLaunch.SetLabelText("Save & Launch") else: - run_args = [sys.executable, 'main.py'] + # 如果启动器是用pythonw启动的，使用python.exe来启动main以便捕获控制台输出 + python_exe = sys.executable + if 'pythonw' in python_exe.lower(): + python_exe = python_exe.replace('pythonw.exe', 'python.exe').replace('pythonw', 'python') + run_args = [python_exe, '-m', 'src.main'] if len(args['character']): run_args.append('--character') run_args.append(args['character']) if args['input'] == 0: if len(args['ifm']): - run_args.append('--ifm') - run_args.append(args['ifm']) + run_args.append('--ifm_input') + if ':' in args['ifm']: + run_args.append(args['ifm']) + else: + run_args.append(args['ifm'] + ':49983') elif args['input'] == 1: - run_args.append('--input') - run_args.append('cam') + run_args.append('--cam_input') elif args['input'] == 2: run_args.append('--debug_input') elif args['input'] == 3: run_args.append('--mouse_input') run_args.append('0,0,' + str(wx.GetDisplaySize().width) + ',' + str(wx.GetDisplaySize().height)) + # Add audio input options for mouse input + if args['mouse_audio_input']: + run_args.append('--mouse_audio_input') + if args['audio_sensitivity']: + run_args.append('--audio_sensitivity') + run_args.append(str(args['audio_sensitivity'])) + if args['audio_threshold']: + run_args.append('--audio_threshold') + run_args.append(str(args['audio_threshold'])) + # Add blink interval for mouse input + if args['blink_interval']: + run_args.append('--blink_interval') + run_args.append(str(args['blink_interval'])) elif args['input'] == 4: if len(args['osf']): - run_args.append('--osf') + run_args.append('--osf_input') run_args.append(args['osf']) + # Add breath cycle option + if args['breath_cycle']: + run_args.append('--breath_cycle') + run_args.append(str(args['breath_cycle'])) + if args['output'] == 0: - run_args.append('--output_webcam') - run_args.append('unitycapture') + run_args.append('--output_spout2') elif args['output'] == 1: - run_args.append('--output_webcam') - run_args.append('obs') - elif args['output'] == 3: - run_args.append('--output_webcam') - run_args.append('spout') + run_args.append('--output_virtual_cam') elif args['output'] == 2: - run_args.append('--debug') + run_args.append('--output_debug') + if args['is_alpha_split']: run_args.append('--alpha_split') if args['is_extend_movement']: run_args.append('--extend_movement') - run_args.append('1') if args['is_bongo']: run_args.append('--bongo') if args['is_alpha_clean']: run_args.append('--alpha_clean') if args['is_eyebrow']: run_args.append('--eyebrow') + if args['cache_simplify'] is not None: run_args.append('--simplify') run_args.append(str(cache_simplify_map[args['cache_simplify']])) - run_args.append('--cacher_quality') - run_args.append(str(cache_simplify_quality_map[args['cache_compression']])) - if args['cache_simplify'] != 'Off': - run_args.append('--use_cacher') if args['ram_cache_size'] is not None: run_args.append('--cache') run_args.append(args['ram_cache_size']) @@ -411,6 +714,20 @@ class LauncherPanel(wx.Panel): run_args.append('4') if args['model_select'] is not None: + if 'tha4_student_' in args['model_select']: + # Student model: tha4_student_{model_name} + model_name = args['model_select'].replace( + 'tha4_student_', '') + run_args.append('--model_version') + run_args.append('v4_student') + run_args.append('--model_name') + run_args.append(model_name) + elif 'tha4' in args['model_select']: + run_args.append('--model_version') + run_args.append('v4') + else: + run_args.append('--model_version') + run_args.append('v3') if 'seperable' in args['model_select']: run_args.append('--model_seperable') if 'half' in args['model_select']: @@ -421,28 +738,49 @@ class LauncherPanel(wx.Panel): run_args.append(args['frame_rate_limit']) if args['sr'] is not None and args['sr'] != 'Off': + run_args.append('--use_sr') if 'anime4k' in args['sr']: - run_args.append('--anime4k') - else: - run_args.append('--use_sr') - if 'x4' in args['sr']: - run_args.append('--sr_x4') - if 'half' in args['sr']: - run_args.append('--sr_half') - - if args['device_id'] is not None: - run_args.append('--device_id') - run_args.append(args['device_id']) + run_args.append('--sr_a4k') + if 'x4' in args['sr']: + run_args.append('--sr_x4') + if 'half' in args['sr']: + run_args.append('--sr_half') if args['use_tensorrt'] is not None and args['use_tensorrt']: run_args.append('--use_tensorrt') - run_args.append('--model_cache') - run_args.append('--model_vram_cache') - run_args.append('--output_size') - run_args.append('512x512') + run_args.append('--filter_min_cutoff') + run_args.append(str(min_cutoff_mapper(args['min_cutoff']))) + + run_args.append('--filter_beta') + run_args.append(str(beta_mapper(args['beta']))) + print('Launched: ' + ' '.join(run_args)) - p = subprocess.Popen(run_args) + self.main_output_lines.clear() + self.main_stderr_lines.clear() + self.statusCtrl.SetValue('Launched') + on_line = lambda line: _on_main_log_line(self, line) + # 使用CREATE_NO_WINDOW标志隐藏控制台窗口，但仍可捕获输出 + creation_flags = 0 + if sys.platform == 'win32': + # CREATE_NO_WINDOW = 0x08000000 + creation_flags = 0x08000000 + p = subprocess.Popen( + run_args, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + creationflags=creation_flags, + ) + threading.Thread( + target=_read_pipe_to_stream, + args=(p.stdout, sys.stdout, self.main_output_lines, on_line), + daemon=True, + ).start() + threading.Thread( + target=_read_pipe_to_stream, + args=(p.stderr, sys.stderr, self.main_stderr_lines, on_line), + daemon=True, + ).start() self.btnLaunch.SetLabelText('Stop') @@ -456,8 +794,14 @@ class MainFrame(wx.Frame): def OnClose(self, e): global p if p is not None: - subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL) + creation_flags = 0 + if sys.platform == 'win32': + # CREATE_NO_WINDOW = 0x08000000 + creation_flags = 0x08000000 + subprocess.run(['taskkill', '/F', '/PID', str(p.pid), '/T'], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + creationflags=creation_flags) e.Skip() def InitUi(self): diff --git a/main.py b/main.py deleted file mode 100644 index ca072b0..0000000 --- a/main.py +++ /dev/null @@ -1,1014 +0,0 @@ -import struct - -import cv2 -import torch -import pyvirtualcam -import numpy as np -import mediapipe as mp -from PIL import Image - -import tha2.poser.modes.mode_20_wx -from pose import get_pose -from preprocess import resize_to_512_center, apply_color_curves -from utils import preprocessing_image, postprocessing_image -from ezvtb_rt_interface import get_core - -import errno -import json -import os -import queue -import socket -import time -import math -from pynput.mouse import Button, Controller -import re -from collections import OrderedDict -from multiprocessing import Value, Process, Queue, shared_memory - -from pyanime4k import ac - -from tha2.mocap.ifacialmocap_constants import * - -from args import args -import copy - -from tha3.util import torch_linear_to_srgb, resize_PIL_image, extract_PIL_image_from_filelike, \ - extract_pytorch_image_from_PIL_image - -import collections -from PySpout import SpoutSender -from OpenGL.GL import GL_RGBA - - -def convert_linear_to_srgb(image: torch.Tensor) -> torch.Tensor: - rgb_image = torch_linear_to_srgb(image[0:3, :, :]) - return torch.cat([rgb_image, image[3:4, :, :]], dim=0) - - -class FPS: - def __init__(self, avarageof=50): - self.frametimestamps = collections.deque(maxlen=avarageof) - - def __call__(self): - self.frametimestamps.append(time.time()) - return self.view() - - def view(self): - if len(self.frametimestamps) > 1: - return len(self.frametimestamps) / (self.frametimestamps[-1] - self.frametimestamps[0] + 1e-10) - else: - return 0.0 - - def last(self): - if len(self.frametimestamps): - return self.frametimestamps[-1] - else: - return time.time() - - -device = torch.device('cpu') - - -def create_default_blender_data(): - data = {} - - for blendshape_name in BLENDSHAPE_NAMES: - data[blendshape_name] = 0.0 - - data[HEAD_BONE_X] = 0.0 - data[HEAD_BONE_Y] = 0.0 - data[HEAD_BONE_Z] = 0.0 - data[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] - - data[LEFT_EYE_BONE_X] = 0.0 - data[LEFT_EYE_BONE_Y] = 0.0 - data[LEFT_EYE_BONE_Z] = 0.0 - data[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] - - data[RIGHT_EYE_BONE_X] = 0.0 - data[RIGHT_EYE_BONE_Y] = 0.0 - data[RIGHT_EYE_BONE_Z] = 0.0 - data[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] - - return data - - -class OSFClientProcess(Process): - def __init__(self): - super().__init__() - self.queue = Queue() - self.should_terminate = Value('b', False) - self.address = args.osf.split(':')[0] - self.port = int(args.osf.split(':')[1]) - self.ifm_fps_number = Value('f', 0.0) - self.perf_time = 0 - - def run(self): - self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.socket.setblocking(False) - self.socket.bind(("", self.port)) - self.socket.settimeout(0.1) - ifm_fps = FPS() - while True: - if self.should_terminate.value: - break - try: - socket_bytes = self.socket.recv(8192) - except socket.error as e: - err = e.args[0] - if err == errno.EAGAIN or err == errno.EWOULDBLOCK or err == 'timed out': - continue - else: - raise e - - # socket_string = socket_bytes.decode("utf-8") - osf_raw = (struct.unpack('=di2f2fB1f4f3f3f68f136f210f14f', socket_bytes)) - # print(osf_raw[432:]) - data = {} - OpenSeeDataIndex = [ - 'time', - 'id', - 'cameraResolutionW', - 'cameraResolutionH', - 'rightEyeOpen', - 'leftEyeOpen', - 'got3DPoints', - 'fit3DError', - 'rawQuaternionX', - 'rawQuaternionY', - 'rawQuaternionZ', - 'rawQuaternionW', - 'rawEulerX', - 'rawEulerY', - 'rawEulerZ', - 'translationY', - 'translationX', - 'translationZ', - ] - for i in range(len(OpenSeeDataIndex)): - data[OpenSeeDataIndex[i]] = osf_raw[i] - data['translationY'] *= -1 - data['translationZ'] *= -1 - data['rotationY'] = data['rawEulerY'] - 10 - data['rotationX'] = (-data['rawEulerX'] + 360) % 360 - 180 - data['rotationZ'] = (data['rawEulerZ'] - 90) - OpenSeeFeatureIndex = [ - 'EyeLeft', - 'EyeRight', - 'EyebrowSteepnessLeft', - 'EyebrowUpDownLeft', - 'EyebrowQuirkLeft', - 'EyebrowSteepnessRight', - 'EyebrowUpDownRight', - 'EyebrowQuirkRight', - 'MouthCornerUpDownLeft', - 'MouthCornerInOutLeft', - 'MouthCornerUpDownRight', - 'MouthCornerInOutRight', - 'MouthOpen', - 'MouthWide' - ] - - for i in range(68): - data['confidence' + str(i)] = osf_raw[i + 18] - for i in range(68): - data['pointsX' + str(i)] = osf_raw[i * 2 + 18 + 68] - data['pointsY' + str(i)] = osf_raw[i * 2 + 18 + 68 + 1] - for i in range(70): - data['points3DX' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2] - data['points3DY' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2 + 1] - data['points3DZ' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2 + 2] - - for i in range(len(OpenSeeFeatureIndex)): - data[OpenSeeFeatureIndex[i]] = osf_raw[i + 432] - # print(data['rotationX'],data['rotationY'],data['rotationZ']) - - a = np.array([ - data['points3DX66'] - data['points3DX68'] + data['points3DX67'] - data['points3DX69'], - data['points3DY66'] - data['points3DY68'] + data['points3DY67'] - data['points3DY69'], - data['points3DZ66'] - data['points3DZ68'] + data['points3DZ67'] - data['points3DZ69'] - ]) - a = (a / np.linalg.norm(a)) - data['eyeRotationX'] = a[0] - data['eyeRotationY'] = a[1] - try: - self.queue.put_nowait(data) - except queue.Full: - pass - self.queue.close() - self.socket.close() - - -ifm_converter = tha2.poser.modes.mode_20_wx.IFacialMocapPoseConverter20() - - -class IFMClientProcess(Process): - def __init__(self): - super().__init__() - self.queue = Queue() - self.should_terminate = Value('b', False) - self.address = args.ifm.split(':')[0] - self.port = int(args.ifm.split(':')[1]) - self.ifm_fps_number = Value('f', 0.0) - self.perf_time = 0 - - def run(self): - - udpClntSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - - data = "iFacialMocap_sahuasouryya9218sauhuiayeta91555dy3719" - - data = data.encode('utf-8') - - udpClntSock.sendto(data, (self.address, self.port)) - - self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.socket.setblocking(False) - self.socket.bind(("", self.port)) - self.socket.settimeout(0.1) - ifm_fps = FPS() - pre_socket_string = '' - while True: - if self.should_terminate.value: - break - try: - socket_bytes = self.socket.recv(8192) - except socket.error as e: - err = e.args[0] - if err == errno.EAGAIN or err == errno.EWOULDBLOCK or err == 'timed out': - continue - else: - raise e - socket_string = socket_bytes.decode("utf-8") - if args.debug and pre_socket_string != socket_string: - self.ifm_fps_number.value = ifm_fps() - pre_socket_string = socket_string - # print(socket_string) - # blender_data = json.loads(socket_string) - data = self.convert_from_blender_data(socket_string) - - try: - self.queue.put_nowait(data) - except queue.Full: - pass - self.queue.close() - self.socket.close() - - @staticmethod - def convert_from_blender_data(blender_data): - data = {} - - for item in blender_data.split('|'): - if item.find('#') != -1: - k, arr = item.split('#') - arr = [float(n) for n in arr.split(',')] - data[k.replace("_L", "Left").replace("_R", "Right")] = arr - elif item.find('-') != -1: - k, v = item.split("-") - data[k.replace("_L", "Left").replace("_R", "Right")] = float(v) / 100 - - to_rad = 57.3 - data[HEAD_BONE_X] = data["=head"][0] / to_rad - data[HEAD_BONE_Y] = data["=head"][1] / to_rad - data[HEAD_BONE_Z] = data["=head"][2] / to_rad - data[HEAD_BONE_QUAT] = [data["=head"][3], data["=head"][4], data["=head"][5], 1] - # print(data[HEAD_BONE_QUAT][2],min(data[EYE_BLINK_LEFT],data[EYE_BLINK_RIGHT])) - data[RIGHT_EYE_BONE_X] = data["rightEye"][0] / to_rad - data[RIGHT_EYE_BONE_Y] = data["rightEye"][1] / to_rad - data[RIGHT_EYE_BONE_Z] = data["rightEye"][2] / to_rad - data[LEFT_EYE_BONE_X] = data["leftEye"][0] / to_rad - data[LEFT_EYE_BONE_Y] = data["leftEye"][1] / to_rad - data[LEFT_EYE_BONE_Z] = data["leftEye"][2] / to_rad - - return data - - -class MouseClientProcess(Process): - def __init__(self): - super().__init__() - self.queue = Queue() - - def run(self): - mouse = Controller() - posLimit = [int(x) for x in args.mouse_input.split(',')] - prev = { - 'eye_l_h_temp': 0, - 'eye_r_h_temp': 0, - 'mouth_ratio': 0, - 'eye_y_ratio': 0, - 'eye_x_ratio': 0, - 'x_angle': 0, - 'y_angle': 0, - 'z_angle': 0, - } - while True: - pos = mouse.position - # print(pos) - eye_limit = [0.8, 0.5] - head_eye_reduce = 0.6 - head_slowness = 0.2 - mouse_data = { - 'eye_l_h_temp': 0, - 'eye_r_h_temp': 0, - 'mouth_ratio': 0, - 'eye_y_ratio': np.interp(pos[1], [posLimit[1], posLimit[3]], [1, -1]) * eye_limit[1], - 'eye_x_ratio': np.interp(pos[0], [posLimit[0], posLimit[2]], [1, -1]) * eye_limit[0], - 'x_angle': np.interp(pos[1], [posLimit[1], posLimit[3]], [1, -1]), - 'y_angle': np.interp(pos[0], [posLimit[0], posLimit[2]], [1, -1]), - 'z_angle': 0, - } - mouse_data['x_angle'] = np.interp(head_slowness, [0, 1], [prev['x_angle'], mouse_data['x_angle']]) - mouse_data['y_angle'] = np.interp(head_slowness, [0, 1], [prev['y_angle'], mouse_data['y_angle']]) - mouse_data['eye_y_ratio'] -= mouse_data['x_angle'] * eye_limit[1] * head_eye_reduce - mouse_data['eye_x_ratio'] -= mouse_data['y_angle'] * eye_limit[0] * head_eye_reduce - if args.bongo: - mouse_data['y_angle'] += 0.05 - mouse_data['x_angle'] += 0.05 - prev = mouse_data - self.queue.put_nowait(mouse_data) - time.sleep(1 / 60) - - -class ModelClientProcess(Process): - def __init__(self, input_image): - super().__init__() - self.should_terminate = Value('b', False) - self.updated = Value('b', False) - self.data = None - self.input_image = input_image - 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) - for _ in range(args.interpolation_scale if args.use_interpolation else 1)] - self.model_fps_number = Value('f', 0.0) - self.gpu_fps_number = Value('f', 0.0) - self.cache_hit_ratio = Value('f', 0.0) - 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, - - 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, - - 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) - self.model.setImage(self.input_image) - input_pose = np.zeros((1, 45), dtype=np.float32) - - self.shared_nps = [ - np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, buffer=self.shms[i].buf) for - i in range(len(self.shms))] - - model_fps = FPS() - - last_process_time = time.time() - while True: - model_input = None - try: - model_input = self.input_queue.get(block=True, timeout=1.0) # Save cpu from infinite loop - while not self.input_queue.empty(): - model_input = self.input_queue.get_nowait() - except queue.Empty: - continue - if model_input is None: continue - frame_interval = (1 / args.frame_rate_limit) if not args.use_interpolation else ( - 1 / args.frame_rate_limit * args.interpolation_scale) - now = time.time() - if now < frame_interval + last_process_time - 0.005: continue # - last_process_time = now - simplify_arr = [1000] * ifm_converter.pose_size - if args.simplify >= 1: - simplify_arr = [200] * ifm_converter.pose_size - simplify_arr[ifm_converter.eye_wink_left_index] = 50 - simplify_arr[ifm_converter.eye_wink_right_index] = 50 - simplify_arr[ifm_converter.eye_happy_wink_left_index] = 50 - simplify_arr[ifm_converter.eye_happy_wink_right_index] = 50 - simplify_arr[ifm_converter.eye_surprised_left_index] = 30 - simplify_arr[ifm_converter.eye_surprised_right_index] = 30 - simplify_arr[ifm_converter.iris_rotation_x_index] = 25 - simplify_arr[ifm_converter.iris_rotation_y_index] = 25 - simplify_arr[ifm_converter.eye_raised_lower_eyelid_left_index] = 10 - simplify_arr[ifm_converter.eye_raised_lower_eyelid_right_index] = 10 - simplify_arr[ifm_converter.mouth_lowered_corner_left_index] = 5 - simplify_arr[ifm_converter.mouth_lowered_corner_right_index] = 5 - simplify_arr[ifm_converter.mouth_raised_corner_left_index] = 5 - simplify_arr[ifm_converter.mouth_raised_corner_right_index] = 5 - if args.simplify >= 2: - simplify_arr[ifm_converter.head_x_index] = 100 - simplify_arr[ifm_converter.head_y_index] = 100 - simplify_arr[ifm_converter.eye_surprised_left_index] = 10 - simplify_arr[ifm_converter.eye_surprised_right_index] = 10 - model_input[ifm_converter.eye_wink_left_index] += model_input[ - ifm_converter.eye_happy_wink_left_index] - model_input[ifm_converter.eye_happy_wink_left_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - model_input[ifm_converter.eye_wink_left_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - model_input[ifm_converter.eye_wink_right_index] += model_input[ - ifm_converter.eye_happy_wink_right_index] - model_input[ifm_converter.eye_happy_wink_right_index] = model_input[ - ifm_converter.eye_wink_right_index] / 2 - model_input[ifm_converter.eye_wink_right_index] = model_input[ - ifm_converter.eye_wink_right_index] / 2 - - uosum = model_input[ifm_converter.mouth_uuu_index] + \ - model_input[ifm_converter.mouth_ooo_index] - model_input[ifm_converter.mouth_ooo_index] = uosum - model_input[ifm_converter.mouth_uuu_index] = 0 - is_open = (model_input[ifm_converter.mouth_aaa_index] + model_input[ - ifm_converter.mouth_iii_index] + uosum) > 0 - model_input[ifm_converter.mouth_lowered_corner_left_index] = 0 - model_input[ifm_converter.mouth_lowered_corner_right_index] = 0 - model_input[ifm_converter.mouth_raised_corner_left_index] = 0.5 if is_open else 0 - model_input[ifm_converter.mouth_raised_corner_right_index] = 0.5 if is_open else 0 - simplify_arr[ifm_converter.mouth_lowered_corner_left_index] = 0 - simplify_arr[ifm_converter.mouth_lowered_corner_right_index] = 0 - simplify_arr[ifm_converter.mouth_raised_corner_left_index] = 0 - simplify_arr[ifm_converter.mouth_raised_corner_right_index] = 0 - if args.simplify >= 3: - simplify_arr[ifm_converter.iris_rotation_x_index] = 20 - simplify_arr[ifm_converter.iris_rotation_y_index] = 20 - simplify_arr[ifm_converter.eye_wink_left_index] = 32 - simplify_arr[ifm_converter.eye_wink_right_index] = 32 - simplify_arr[ifm_converter.eye_happy_wink_left_index] = 32 - simplify_arr[ifm_converter.eye_happy_wink_right_index] = 32 - if args.simplify >= 4: - simplify_arr[ifm_converter.head_x_index] = 50 - simplify_arr[ifm_converter.head_y_index] = 50 - simplify_arr[ifm_converter.neck_z_index] = 100 - model_input[ifm_converter.eye_raised_lower_eyelid_left_index] = 0 - model_input[ifm_converter.eye_raised_lower_eyelid_right_index] = 0 - simplify_arr[ifm_converter.iris_rotation_x_index] = 10 - simplify_arr[ifm_converter.iris_rotation_y_index] = 10 - simplify_arr[ifm_converter.eye_wink_left_index] = 24 - simplify_arr[ifm_converter.eye_wink_right_index] = 24 - simplify_arr[ifm_converter.eye_happy_wink_left_index] = 24 - simplify_arr[ifm_converter.eye_happy_wink_right_index] = 24 - simplify_arr[ifm_converter.eye_surprised_left_index] = 8 - simplify_arr[ifm_converter.eye_surprised_right_index] = 8 - model_input[ifm_converter.eye_wink_left_index] += model_input[ - ifm_converter.eye_wink_right_index] - model_input[ifm_converter.eye_wink_right_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - model_input[ifm_converter.eye_wink_left_index] = model_input[ - ifm_converter.eye_wink_left_index] / 2 - - model_input[ifm_converter.eye_surprised_left_index] += model_input[ - ifm_converter.eye_surprised_right_index] - model_input[ifm_converter.eye_surprised_right_index] = model_input[ - ifm_converter.eye_surprised_left_index] / 2 - model_input[ifm_converter.eye_surprised_left_index] = model_input[ - ifm_converter.eye_surprised_left_index] / 2 - - model_input[ifm_converter.eye_happy_wink_left_index] += model_input[ - ifm_converter.eye_happy_wink_right_index] - model_input[ifm_converter.eye_happy_wink_right_index] = model_input[ - ifm_converter.eye_happy_wink_left_index] / 2 - model_input[ifm_converter.eye_happy_wink_left_index] = model_input[ - ifm_converter.eye_happy_wink_left_index] / 2 - model_input[ifm_converter.mouth_aaa_index] = min( - model_input[ifm_converter.mouth_aaa_index] + - model_input[ifm_converter.mouth_ooo_index] / 2 + - model_input[ifm_converter.mouth_iii_index] / 2 + - model_input[ifm_converter.mouth_uuu_index] / 2, 1 - ) - model_input[ifm_converter.mouth_ooo_index] = 0 - model_input[ifm_converter.mouth_iii_index] = 0 - model_input[ifm_converter.mouth_uuu_index] = 0 - for i in range(4, args.simplify): - simplify_arr = [max(math.ceil(x * 0.8), 5) for x in simplify_arr] - for i in range(0, len(simplify_arr)): - if simplify_arr[i] > 0: - model_input[i] = round(model_input[i] * simplify_arr[i]) / simplify_arr[i] - - if args.perf == 'model': - tic = time.perf_counter() - if args.eyebrow: - for i in range(12): - input_pose[0, i] = model_input[i] - for i in range(27): - input_pose[0, i + 12] = model_input[i + 12] - for i in range(6): - input_pose[0, i + 12 + 27] = model_input[i + 27 + 12] - - output_images = self.model.inference(input_pose) - - if args.perf == 'model': - print("postprocess", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - - for i in range(len(self.shared_nps)): - np.copyto(self.shared_nps[i], output_images[i]) - self.output_queue.put_nowait(True) - - if args.debug: - self.model_fps_number.value = model_fps() - if self.model.cacher is not None: - self.cache_hit_ratio.value = self.model.cacher.hits / ( - self.model.cacher.hits + self.model.cacher.miss + 1) - try: - if args.use_tensorrt and args.model_cache and args.model_vram_cache: - self.gpu_cache_hit_ratio.value = self.model.tha.morpher_cacher.hits / ( - self.model.tha.morpher_cacher.hits + self.model.tha.morpher_cacher.miss) - except: - pass - - -@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) - - print("Character Image Loaded:", args.character) - cap = None - - output_fps = FPS() - - if not args.debug_input: - - if args.ifm is not None: - client_process = IFMClientProcess() - client_process.daemon = True - client_process.start() - print("iFacialMocap Service Running:", args.ifm) - - elif args.osf is not None: - client_process = OSFClientProcess() - client_process.daemon = True - client_process.start() - print("OpenSeeFace Service Running:", args.osf) - - elif args.mouse_input is not None: - client_process = MouseClientProcess() - client_process.daemon = True - client_process.start() - print("Mouse Input Running") - - else: - - if args.input == 'cam': - cap = cv2.VideoCapture(0) - ret, frame = cap.read() - if ret is None: - raise Exception("Can't find Camera") - else: - cap = cv2.VideoCapture(args.input) - frame_count = 0 - os.makedirs(os.path.join('dst', args.character, args.output_dir), exist_ok=True) - print("Webcam Input Running") - - facemesh = mp.solutions.face_mesh.FaceMesh(refine_landmarks=True) - - if args.output_webcam: - cam_scale = 1 - cam_width_scale = 1 - if args.anime4k: - cam_scale = 2 - if args.alpha_split: - cam_width_scale = 2 - - if args.output_webcam == 'spout': - # Create the sender - sender = SpoutSender("EasyVtuber", args.model_output_size * cam_scale * cam_width_scale, - args.model_output_size * cam_scale, GL_RGBA) - else: - cam = pyvirtualcam.Camera(width=args.model_output_size * cam_scale * cam_width_scale, - height=args.model_output_size * cam_scale, - fps=60, - backend=args.output_webcam, - fmt= - {'unitycapture': pyvirtualcam.PixelFormat.RGBA, - 'obs': pyvirtualcam.PixelFormat.RGB}[ - args.output_webcam]) - print(f'Using virtual camera: {cam.device}') - - a = None - - if args.anime4k: - parameters = ac.Parameters() - # enable HDN for ACNet - parameters.HDN = True - - # a = ac.AC( - # managerList=ac.ManagerList([ac.CUDAManager(dID=0)]), - # type=ac.ProcessorType.Cuda_ACNet, - # ) - - a = ac.AC( - managerList=ac.ManagerList([ac.OpenCLACNetManager(pID=0, dID=0)]), - type=ac.ProcessorType.OpenCL_ACNet, - ) - a.set_arguments(parameters) - print("Anime4K Loaded") - - position_vector = [0, 0, 0, 1] - position_vector_0 = None - pose_vector_0 = None - - pose_queue = [] - blender_data = {} - if (args.ifm): - blender_data = create_default_blender_data() - mouse_data = { - 'eye_l_h_temp': 0, - 'eye_r_h_temp': 0, - 'mouth_ratio': 0, - 'eye_y_ratio': 0, - 'eye_x_ratio': 0, - 'x_angle': 0, - 'y_angle': 0, - 'z_angle': 0, - } - - model_output = None - model_process = ModelClientProcess(input_image) - model_process.daemon = True - model_process.start() - model_output_nps = [np.ndarray((args.model_output_size, args.model_output_size, 4), dtype=np.uint8, - buffer=model_process.shms[i].buf) for i in range(len(model_process.shms))] - model_result_read_ptr = 0 - model_return_fps = FPS(5) - need_a_copy = True - - print("Ready. Close this console to exit.") - - mouth_eye_vector_c_hist = [] - - while True: - # ret, frame = cap.read() - # input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - # results = facemesh.process(input_frame) - - if args.perf == 'main': - tic = time.perf_counter() - if args.debug_input: - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - mouth_eye_vector_c[2] = math.sin(time.perf_counter() * 3) - mouth_eye_vector_c[3] = math.sin(time.perf_counter() * 3) - - mouth_eye_vector_c[14] = 0 - - mouth_eye_vector_c[25] = math.sin(time.perf_counter() * 2.2) * 0.2 - mouth_eye_vector_c[26] = math.sin(time.perf_counter() * 3.5) * 0.8 - - pose_vector_c[0] = math.sin(time.perf_counter() * 1.1) - pose_vector_c[1] = math.sin(time.perf_counter() * 1.2) - pose_vector_c[2] = math.sin(time.perf_counter() * 1.5) - - eyebrow_vector_c[6] = math.sin(time.perf_counter() * 1.1) - eyebrow_vector_c[7] = math.sin(time.perf_counter() * 1.1) - - elif args.osf is not None: - try: - new_blender_data = blender_data - while not client_process.should_terminate.value and not client_process.queue.empty(): - new_blender_data = client_process.queue.get_nowait() - blender_data = new_blender_data - except queue.Empty: - pass - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - if len(blender_data) != 0: - mouth_eye_vector_c[2] = 1 - blender_data['leftEyeOpen'] - mouth_eye_vector_c[3] = 1 - blender_data['rightEyeOpen'] - - mouth_eye_vector_c[14] = max(blender_data['MouthOpen'], 0) - # print(mouth_eye_vector_c[14]) - - mouth_eye_vector_c[25] = -blender_data['eyeRotationY'] * 3 - (blender_data['rotationX']) / 57.3 * 1.5 - mouth_eye_vector_c[26] = blender_data['eyeRotationX'] * 3 + (blender_data['rotationY']) / 57.3 - # print(mouth_eye_vector_c[25:27]) - eyebrow_vector_c[6] = blender_data['EyebrowUpDownLeft'] - eyebrow_vector_c[7] = blender_data['EyebrowUpDownRight'] - # print(blender_data['EyebrowUpDownLeft'],blender_data['EyebrowUpDownRight']) - - # if pose_vector_0==None: - # pose_vector_0=[0,0,0] - # pose_vector_0[0] = blender_data['rotationX'] - # pose_vector_0[1] = blender_data['rotationY'] - # pose_vector_0[2] = blender_data['rotationZ'] - # pose_vector_c[0] = (blender_data['rotationX']-pose_vector_0[0])/57.3*3 - # pose_vector_c[1] = -(blender_data['rotationY']-pose_vector_0[1])/57.3*3 - # pose_vector_c[2] = (blender_data['rotationZ']-pose_vector_0[2])/57.3 - pose_vector_c[0] = (blender_data['rotationX']) / 57.3 * 3 - pose_vector_c[1] = -(blender_data['rotationY']) / 57.3 * 3 - pose_vector_c[2] = (blender_data['rotationZ']) / 57.3 * 2 - # print(pose_vector_c) - - if position_vector_0 == None: - position_vector_0 = [0, 0, 0, 1] - position_vector_0[0] = blender_data['translationX'] - position_vector_0[1] = blender_data['translationY'] - position_vector_0[2] = blender_data['translationZ'] - position_vector[0] = -(blender_data['translationX'] - position_vector_0[0]) * 0.1 - position_vector[1] = -(blender_data['translationY'] - position_vector_0[1]) * 0.1 - position_vector[2] = -(blender_data['translationZ'] - position_vector_0[2]) * 0.1 - - # Apply filter - if len(mouth_eye_vector_c_hist) == 0: - for _ in range(3): - mouth_eye_vector_c_hist.append(copy.deepcopy(mouth_eye_vector_c)) - else: - mouth_eye_vector_c_hist.pop(0) - mouth_eye_vector_c_hist.append(mouth_eye_vector_c) - mouth_eye_vector_c[25] = mouth_eye_vector_c_hist[0][25] * 0.2 + mouth_eye_vector_c_hist[1][25] * 0.5 + mouth_eye_vector_c_hist[2][25] * 0.3 - mouth_eye_vector_c[26] = mouth_eye_vector_c_hist[0][26] * 0.2 + mouth_eye_vector_c_hist[1][26] * 0.5 + mouth_eye_vector_c_hist[2][26] * 0.3 - - elif args.ifm is not None: - # get pose from ifm - try: - new_blender_data = blender_data - while not client_process.should_terminate.value and not client_process.queue.empty(): - new_blender_data = client_process.queue.get_nowait() - blender_data = new_blender_data - except queue.Empty: - pass - - ifacialmocap_pose_converted = ifm_converter.convert(blender_data) - - # ifacialmocap_pose = blender_data - # - # eye_l_h_temp = ifacialmocap_pose[EYE_BLINK_LEFT] - # eye_r_h_temp = ifacialmocap_pose[EYE_BLINK_RIGHT] - # mouth_ratio = (ifacialmocap_pose[JAW_OPEN] - 0.10)*1.3 - # x_angle = -ifacialmocap_pose[HEAD_BONE_X] * 1.5 + 1.57 - # y_angle = -ifacialmocap_pose[HEAD_BONE_Y] - # z_angle = ifacialmocap_pose[HEAD_BONE_Z] - 1.57 - # - # eye_x_ratio = (ifacialmocap_pose[EYE_LOOK_IN_LEFT] - - # ifacialmocap_pose[EYE_LOOK_OUT_LEFT] - - # ifacialmocap_pose[EYE_LOOK_IN_RIGHT] + - # ifacialmocap_pose[EYE_LOOK_OUT_RIGHT]) / 2.0 / 0.75 - # - # eye_y_ratio = (ifacialmocap_pose[EYE_LOOK_UP_LEFT] - # + ifacialmocap_pose[EYE_LOOK_UP_RIGHT] - # - ifacialmocap_pose[EYE_LOOK_DOWN_RIGHT] - # + ifacialmocap_pose[EYE_LOOK_DOWN_LEFT]) / 2.0 / 0.75 - - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - for i in range(0, 12): - eyebrow_vector_c[i] = ifacialmocap_pose_converted[i] - for i in range(12, 39): - mouth_eye_vector_c[i - 12] = ifacialmocap_pose_converted[i] - for i in range(39, 42): - pose_vector_c[i - 39] = ifacialmocap_pose_converted[i] - - position_vector = blender_data[HEAD_BONE_QUAT] - - elif args.mouse_input is not None: - - try: - new_blender_data = mouse_data - while not client_process.queue.empty(): - new_blender_data = client_process.queue.get_nowait() - mouse_data = new_blender_data - except queue.Empty: - pass - - eye_l_h_temp = mouse_data['eye_l_h_temp'] - eye_r_h_temp = mouse_data['eye_r_h_temp'] - mouth_ratio = mouse_data['mouth_ratio'] - eye_y_ratio = mouse_data['eye_y_ratio'] - eye_x_ratio = mouse_data['eye_x_ratio'] - x_angle = mouse_data['x_angle'] - y_angle = mouse_data['y_angle'] - z_angle = mouse_data['z_angle'] - - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - mouth_eye_vector_c[2] = eye_l_h_temp - mouth_eye_vector_c[3] = eye_r_h_temp - - mouth_eye_vector_c[14] = mouth_ratio * 1.5 - - mouth_eye_vector_c[25] = eye_y_ratio - mouth_eye_vector_c[26] = eye_x_ratio - - pose_vector_c[0] = x_angle - pose_vector_c[1] = y_angle - pose_vector_c[2] = z_angle - - else: - ret, frame = cap.read() - input_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - results = facemesh.process(input_frame) - - if results.multi_face_landmarks is None: - continue - - facial_landmarks = results.multi_face_landmarks[0].landmark - - if args.debug: - pose, debug_image = get_pose(facial_landmarks, frame) - else: - pose = get_pose(facial_landmarks) - - if len(pose_queue) < 3: - pose_queue.append(pose) - pose_queue.append(pose) - pose_queue.append(pose) - else: - pose_queue.pop(0) - pose_queue.append(pose) - - np_pose = np.average(np.array(pose_queue), axis=0, weights=[0.6, 0.3, 0.1]) - - eye_l_h_temp = np_pose[0] - eye_r_h_temp = np_pose[1] - mouth_ratio = np_pose[2] - eye_y_ratio = np_pose[3] - eye_x_ratio = np_pose[4] - x_angle = np_pose[5] - y_angle = np_pose[6] - z_angle = np_pose[7] - - eyebrow_vector_c = [0.0] * 12 - mouth_eye_vector_c = [0.0] * 27 - pose_vector_c = [0.0] * 6 - - mouth_eye_vector_c[2] = eye_l_h_temp - mouth_eye_vector_c[3] = eye_r_h_temp - - mouth_eye_vector_c[14] = mouth_ratio * 1.5 - - mouth_eye_vector_c[25] = eye_y_ratio - mouth_eye_vector_c[26] = eye_x_ratio - - pose_vector_c[0] = (x_angle - 1.5) * 1.6 - pose_vector_c[1] = y_angle * 2.0 # temp weight - pose_vector_c[2] = (z_angle + 1.5) * 2 # temp weight - - pose_vector_c[3] = pose_vector_c[1] - pose_vector_c[4] = pose_vector_c[2] - - model_input_arr = eyebrow_vector_c - model_input_arr.extend(mouth_eye_vector_c) - model_input_arr.extend(pose_vector_c) - - if args.perf == 'main': - print('===') - print("input", time.perf_counter() - tic) - tic = time.perf_counter() - model_process.input_queue.put_nowait(model_input_arr) - - try: - while not model_process.output_queue.empty(): - model_process.output_queue.get_nowait() - model_result_read_ptr = 0 - model_return_fps() - need_a_copy = True - except queue.Empty: - pass - - if need_a_copy: - model_output = model_output_nps[model_result_read_ptr].copy() - need_a_copy = False - interval = 1 / (model_return_fps.view() + 1e-10) / len(model_output_nps) * 0.8 - if model_result_read_ptr + 1 < len(model_output_nps) and time.time() >= ( - (model_result_read_ptr + 1) * interval + model_return_fps.last()): - model_result_read_ptr += 1 - need_a_copy = True - - postprocessed_image = model_output - if args.perf == 'main': - print("Compute", time.perf_counter() - tic) - tic = time.perf_counter() - - k_scale = 1 - rotate_angle = 0 - dx = 0 - dy = 0 - if args.extend_movement: - k_scale = position_vector[2] * math.sqrt(args.extend_movement) + 1 - rotate_angle = -position_vector[0] * 10 * args.extend_movement - dx = position_vector[0] * 400 * k_scale * args.extend_movement - dy = -position_vector[1] * 600 * k_scale * args.extend_movement - if args.bongo: - rotate_angle -= 5 - - rm = cv2.getRotationMatrix2D((args.model_output_size / 2, args.model_output_size / 2), rotate_angle, k_scale) - rm[0, 2] += dx + args.model_output_size / 2 - args.model_output_size / 2 - rm[1, 2] += dy + args.model_output_size / 2 - args.model_output_size / 2 - - postprocessed_image = cv2.warpAffine( - postprocessed_image, - rm, - (args.model_output_size, args.model_output_size)) - - if args.perf == 'main': - print("postprocess", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - - output_fps_number = output_fps() - - if args.anime4k: - alpha_channel = postprocessed_image[:, :, 3] - alpha_channel = cv2.resize(alpha_channel, None, fx=2, fy=2) - - # a.load_image_from_numpy(cv2.cvtColor(postprocessed_image, cv2.COLOR_RGBA2RGB), input_type=ac.AC_INPUT_RGB) - # img = cv2.imread("character/test41.png") - img1 = cv2.cvtColor(postprocessed_image, cv2.COLOR_BGRA2BGR) - # a.load_image_from_numpy(img, input_type=ac.AC_INPUT_BGR) - a.load_image_from_numpy(img1, input_type=ac.AC_INPUT_BGR) - a.process() - postprocessed_image = a.save_image_to_numpy() - postprocessed_image = cv2.merge((postprocessed_image, alpha_channel)) - if args.perf == 'main': - print("anime4k", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - if args.alpha_split: - alpha_image = cv2.merge( - [postprocessed_image[:, :, 3], postprocessed_image[:, :, 3], postprocessed_image[:, :, 3]]) - alpha_image = cv2.cvtColor(alpha_image, cv2.COLOR_BGR2BGRA) - postprocessed_image = cv2.hconcat([postprocessed_image, alpha_image]) - - if args.debug: - if args.use_sr: - time.sleep( - 0.000001) # Opencv is using different strategy when putting 512x512 and 1024x1024 image to imshow, - # in a way that 1024x1024 is much faster than 512x512 14msvs0.2ms(maybe using zero-copy?) so when using SR, add a delay to prevent overan - # Due to windows system limitation this actually sleep 16ms - output_frame = postprocessed_image - # resized_frame = cv2.resize(output_frame, (np.min(debug_image.shape[:2]), np.min(debug_image.shape[:2]))) - # output_frame = np.concatenate([debug_image, resized_frame], axis=1) - cv2.putText(output_frame, str('OUT_FPS:%.1f' % output_fps_number), (0, 16), cv2.FONT_HERSHEY_PLAIN, 1, - (0, 255, 0), 1) - cv2.putText(output_frame, str( - 'GPU_FPS:%.1f' % ( - model_process.model_fps_number.value if not args.use_interpolation else model_process.model_fps_number.value * args.interpolation_scale)), - (0, 32), - cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) - if args.ifm is not None: - cv2.putText(output_frame, str('IFM_FPS:%.1f' % client_process.ifm_fps_number.value), (0, 48), - cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) - if args.max_cache_len > 0: - cv2.putText(output_frame, str('MEMCACHED:%.1f%%' % (model_process.cache_hit_ratio.value * 100)), - (0, 64), - cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) - if args.max_gpu_cache_len > 0.0: - cv2.putText(output_frame, str('GPUCACHED:%.1f%%' % (model_process.gpu_cache_hit_ratio.value * 100)), - (0, 80), - cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) - cv2.imshow("frame", output_frame) - # cv2.imshow("camera", debug_image) - cv2.waitKey(1) - if args.output_webcam: - # result_image = np.zeros([720, 1280, 3], dtype=np.uint8) - # result_image[720 - 512:, 1280 // 2 - 256:1280 // 2 + 256] = cv2.resize( - # cv2.cvtColor(postprocessing_image(output_image.cpu()), cv2.COLOR_RGBA2RGB), (512, 512)) - result_image = postprocessed_image - if args.output_webcam == 'spout': - sender.send_image(cv2.cvtColor(postprocessed_image, cv2.COLOR_BGRA2RGBA), False) - time.sleep(0.000001) - else: - if args.output_webcam == 'obs': - result_image = cv2.cvtColor(result_image, cv2.COLOR_BGRA2RGB) - cam.send(result_image) - cam.sleep_until_next_frame() - if args.perf == 'main': - print("output", (time.perf_counter() - tic) * 1000) - tic = time.perf_counter() - - -if __name__ == '__main__': - main() 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 new file mode 100644 index 0000000..922dd7b --- /dev/null +++ b/pack_release.bat @@ -0,0 +1,54 @@ +@echo on +REM ======================================================================== +REM IMPORTANT: Before running this script, please make sure to: +REM 1. Checkout to the release_package branch: git checkout release_package +REM 2. Rebase with the latest changes: git rebase main +REM ======================================================================== + +cd /D "%~dp0" + +if not exist "envs" mkdir envs +if not exist "envs\miniconda3" mkdir envs\miniconda3 + +IF not EXIST %~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin ( + @RD /S /Q %~dp0envs\TensorRT-RTX-1.3.0.35_cu129 + cd /D "%~dp0\envs" + curl -L -o trt_rtx.zip https://developer.nvidia.com/downloads/trt/rtx_sdk/secure/1.3/TensorRT-RTX-1.3.0.35-win10-amd64-cuda-12.9-Release-external.zip + tar -xf trt_rtx.zip "TensorRT-RTX-1.3.0.35" && ren "TensorRT-RTX-1.3.0.35" "TensorRT-RTX-1.3.0.35_cu129" && del trt_rtx.zip + cd /D "%~dp0" +) + +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'}" + + 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 +) + +SET PATH=%~dp0envs\TensorRT-RTX-1.3.0.35_cu129\bin;%~dp0envs\miniconda3\Scripts;%PATH% + +call activate +call conda env list +call conda update -y --all +call conda create -n ezvtb_rt_venv_release python=3.10 -y +call conda activate ezvtb_rt_venv_release +call conda env list + +call conda install -y conda-pack +call conda install -y conda-forge::pycuda +call conda install -y -c nvidia/label/cuda-12.9.1 cuda-nvcc-dev_win-64 cudnn cuda-runtime + +call conda-pack -n ezvtb_rt_venv_release -o %~dp0envs\python_embedded --format no-archive + + +SET PATH=%~dp0envs\python_embedded;%~dp0envs\python_embedded\Scripts;%~dp0envs\python_embedded\Library\bin;%PATH% +call python -m pip install %~dp0envs\TensorRT-RTX-1.3.0.35_cu129\python\tensorrt_rtx-1.3.0.35-cp310-none-win_amd64.whl +call python -m pip install -r requirements.txt --no-warn-script-location + +@RD /S /Q %~dp0envs\miniconda3 +pause \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 9bd7383..1503fd9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,17 +1,18 @@ -numpy==1.26.4 +numpy scipy -matplotlib -protobuf==3.20.1 -pynput==1.7.6 -mediapipe -opencv_python==4.5.5.64 -Pillow==9.1.0 -pyanime4k==2.5.2 +protobuf +pynput +mediapipe==0.10.11 +opencv_python +Pillow +pyanime4k==3.1.0 pyvirtualcam==0.9.1 pyopengl wxpython onnx onnxruntime-directml -turbojpeg tqdm -torch \ No newline at end of file +torch +OneEuroFilter +sounddevice +brotli \ No newline at end of file diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/args.py b/src/args.py new file mode 100644 index 0000000..4b2d98f --- /dev/null +++ b/src/args.py @@ -0,0 +1,97 @@ +import argparse +import re + + +def convert_to_byte(size): + result = re.search('(\d+\.?\d*)(b|kb|mb|gb|tb)', size.lower()) + if (result and result.groups()): + unit = result.groups()[1] + amount = float(result.groups()[0]) + index = ['b', 'kb', 'mb', 'gb', 'tb'].index(unit) + return amount * pow(1024, index) + raise ValueError("Invalid size provided, value is " + size) + + +parser = argparse.ArgumentParser() +parser.add_argument('--character', type=str, default='lambda_00') + +parser.add_argument('--debug_input', action='store_true') +parser.add_argument('--cam_input', action='store_true') +parser.add_argument('--mouse_input', type=str) +parser.add_argument('--ifm_input', type=str) +parser.add_argument('--osf_input', type=str) + +parser.add_argument('--mouse_audio_input', action='store_true') +parser.add_argument('--audio_sensitivity', type=float, default=0.02) +parser.add_argument('--audio_threshold', type=float, default=10.0) +parser.add_argument('--blink_interval', type=float, default=5.0) +parser.add_argument('--breath_cycle', type=float, default=float('inf')) + +parser.add_argument('--output_virtual_cam', action='store_true') +parser.add_argument('--output_spout2', action='store_true') +parser.add_argument('--output_debug', action='store_true') + +parser.add_argument('--alpha_split', action='store_true') +parser.add_argument('--bongo', action='store_true') +parser.add_argument('--extend_movement', action='store_true') + +parser.add_argument('--filter_min_cutoff', type=float, default=10.0) +parser.add_argument('--filter_beta', type=float, default=0.3) + +parser.add_argument('--simplify', type=int, default=1) + +parser.add_argument('--use_tensorrt', action='store_true') +parser.add_argument('--model_version', type=str, default='v3') +parser.add_argument('--model_name', type=str, default='') +parser.add_argument('--model_seperable', action='store_true') +parser.add_argument('--model_half', action='store_true') +parser.add_argument('--gpu_cache', type=str, default='512mb') +parser.add_argument('--eyebrow', action='store_true') + +parser.add_argument('--use_interpolation', action='store_true') +parser.add_argument('--interpolation_scale', type=int, default=1) +parser.add_argument('--interpolation_half', action='store_true') +parser.add_argument('--mark_interpolated', action='store_true', help='在插值帧左上角打红点便于确认补帧输出') + +parser.add_argument('--use_sr', action='store_true') +parser.add_argument('--sr_x4', action='store_true') +parser.add_argument('--sr_half', action='store_true') +parser.add_argument('--sr_a4k', action='store_true') + +parser.add_argument('--cache', type=str, default='256mb') + +parser.add_argument('--frame_rate_limit', type=int, default=30) + +parser.add_argument('--alpha_clean', action='store_true') + +args = parser.parse_args() + +if args.cache is not None: + args.max_ram_cache_len = convert_to_byte(args.cache) / pow(1024, 3) # In gigabytes +else: + args.max_ram_cache_len = 0 +if args.gpu_cache is not None: + args.max_gpu_cache_len = convert_to_byte(args.gpu_cache) / pow(1024, 3) # In gigabytes +else: + args.max_gpu_cache_len = 0 + +if args.simplify == 0: + args.max_ram_cache_len = 0 # Disable cacher if simplify is Off + +if not args.output_virtual_cam and not args.output_spout2: + args.output_debug = True # Default to debug output + +if args.output_spout2: + args.alpha_split = False # Disable alpha split for spout2 output + +if not args.cam_input and args.mouse_input is None \ + and args.ifm_input is None and args.osf_input is None: + args.debug_input = True # Default to debug input + +if args.sr_a4k: + args.use_sr = True + +if args.output_spout2: + args.alpha_split = False # Disable alpha split for spout2 output + +args.model_output_size = 1024 if args.use_sr else 512 diff --git a/src/debug_input_client.py b/src/debug_input_client.py new file mode 100644 index 0000000..48ad7b6 --- /dev/null +++ b/src/debug_input_client.py @@ -0,0 +1,65 @@ +from multiprocessing import Process, shared_memory, Value +import numpy as np +import time +import math +from .utils.timer_wait import wait_until +from .utils.shared_mem_guard import SharedMemoryGuard +from .args import args + +class DebugInputClientProcess(Process): + def __init__(self, pose_position_shm: shared_memory.SharedMemory): + super().__init__() + self.pose_position_shm = pose_position_shm + self.fps = Value('f', 60.0) + + def run(self): + last_time : float = time.perf_counter() + interval : float = 1.0 / 60 # 60 FPS + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) + np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) + # 呼吸循环参数 + breath_start_time = time.perf_counter() + while True: + eyebrow_vector = [0.0] * 12 + mouth_eye_vector = [0.0] * 27 + pose_vector = [0.0] * 6 + position_vector = [0, 0, 0, 1] + + # 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0） + breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle + # 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0 + # sin 在 0 到 π 之间从 0 到 1 到 0 + breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi) + + mouth_eye_vector[2] = math.sin(time.perf_counter() * 3) + mouth_eye_vector[3] = math.sin(time.perf_counter() * 3) + + mouth_eye_vector[14] = 0 + + mouth_eye_vector[25] = math.sin(time.perf_counter() * 2.2) * 0.2 + mouth_eye_vector[26] = math.sin(time.perf_counter() * 3.5) * 0.8 + + pose_vector[0] = math.sin(time.perf_counter() * 1.1) + pose_vector[1] = math.sin(time.perf_counter() * 1.2) + pose_vector[2] = math.sin(time.perf_counter() * 1.5) + pose_vector[3] = pose_vector[1] + pose_vector[4] = pose_vector[2] + pose_vector[5] = breath_value + eyebrow_vector[6] = math.sin(time.perf_counter() * 1.1) + eyebrow_vector[7] = math.sin(time.perf_counter() * 1.1) + + model_input_arr = eyebrow_vector + model_input_arr.extend(mouth_eye_vector) + model_input_arr.extend(pose_vector) + + position_vector[0] = math.sin(time.perf_counter() * 0.5) * 0.1 + position_vector[1] = math.sin(time.perf_counter() * 0.6) * 0.1 + position_vector[2] = math.sin(time.perf_counter() * 0.7) * 0.1 + position_vector[3] = 1 + + with pose_position_shm_guard.lock(): + np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) + np_position_shm[:] = np.array(position_vector, dtype=np.float32) + wait_until(last_time + interval) + last_time += interval \ No newline at end of file diff --git a/src/ezvtb_rt_interface.py b/src/ezvtb_rt_interface.py new file mode 100644 index 0000000..09204a2 --- /dev/null +++ b/src/ezvtb_rt_interface.py @@ -0,0 +1,72 @@ +import os +import sys +from .args import args +dir_path = os.path.dirname(os.path.realpath(__file__)) +ezvtb_path = os.path.join(dir_path, "..", "ezvtuber-rt") +ezvtb_main_path = os.path.join(dir_path, "..", 'ezvtuber-rt-main') + +project_path = '' +if os.path.exists(ezvtb_path): + project_path = ezvtb_path +else: + project_path = ezvtb_main_path + +if project_path not in sys.path: + sys.path.append(project_path) + +# Initialize model data path to point to data/models directory +import ezvtb_rt +models_path = os.path.join(dir_path, '..', 'data', 'models') +ezvtb_rt.init_model_path(models_path) + +def get_core( + use_tensorrt:bool = True, + #THA3 model setting + model_version:str = 'v3', #'v3' or 'v4' or 'v4_student' + model_name:str = '', + model_seperable:bool = True, + model_half:bool = True, #If using directml+half, there is small numerical error on Nvidia, and huge numerical error on AMD + model_cache_size:float = 1.0, #unit of GigaBytes, only works for tensorrt + model_use_eyebrow:bool = True, + #RIFE interpolation setting + use_interpolation:bool = True, + interpolation_scale:int = 2, + interpolation_half:bool = True, #If using directml+half, there is small numerical error on Nvidia, and huge numerical error on AMD + #Cacher setting + cacher_ram_size:float = 2.0,#unit of GigaBytes + #SR setting + use_sr:bool = False, + sr_x4:bool = True, + sr_half:bool = True, + sr_a4k:bool = False, + ): + if use_tensorrt: + try: + from ezvtb_rt.core_trt import CoreTRT as Core + except: + print("TensorRT is not available, fallback to ONNX Runtime.") + args.use_tensorrt = False + from ezvtb_rt.core_ort import CoreORT as Core + else: + from ezvtb_rt.core_ort import CoreORT as Core + core = Core( + tha_model_version=model_version, + tha_model_seperable=model_seperable, + tha_model_fp16=model_half, + tha_model_name=model_name, + + 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_fp16=sr_half, + sr_a4k=sr_a4k, + + vram_cache_size=model_cache_size, + cache_max_giga=cacher_ram_size, + use_eyebrow=model_use_eyebrow, + ) + return core + \ No newline at end of file diff --git a/src/face_mesh_client.py b/src/face_mesh_client.py new file mode 100644 index 0000000..bf752c4 --- /dev/null +++ b/src/face_mesh_client.py @@ -0,0 +1,100 @@ +import mediapipe as mp +from multiprocessing import Process, shared_memory, Value +import cv2 +import numpy as np +from .args import args +from .utils.shared_mem_guard import SharedMemoryGuard +from .utils.pose import get_pose +from .utils.fps import FPS +from .utils.filter import OneEuroFilterNumpy +from OneEuroFilter import OneEuroFilter +import time + + +class FaceMeshClientProcess(Process): + def __init__(self, pose_position_shm: shared_memory.SharedMemory): + super().__init__() + self.pose_position_shm = pose_position_shm + self.fps = Value('f', 0.0) + + def run(self): + facemesh = mp.solutions.face_mesh.FaceMesh(refine_landmarks=True) + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) + np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) + + cap = cv2.VideoCapture(0, cv2.CAP_DSHOW) + if not cap.isOpened(): + raise Exception("Can't open webcam") + + print("Warming up webcam...") + + frame_count = 0 + input_fps = FPS(60) + while frame_count < 5: + ret, frame = cap.read() + if not ret: + raise Exception("Can't receive frame (stream end?).") + input_fps() + frame_count += 1 + + # 呼吸循环参数 + breath_start_time = time.perf_counter() + pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + position_offset = None + print("Webcam Input Running at {:.2f} FPS".format(input_fps.view())) + position_vector = np.array([0, 0, 0, 1], dtype=np.float32) + while True: + ret, frame = cap.read() + if not ret: + raise Exception("Can't receive frame (stream end?).") + self.fps.value = input_fps() + rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + results = facemesh.process(rgb_frame) + if results.multi_face_landmarks is None: + continue + + # 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0） + breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle + # 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0 + # sin 在 0 到 π 之间从 0 到 1 到 0 + breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi) + + facial_landmarks = results.multi_face_landmarks[0].landmark + pose = get_pose(facial_landmarks) + eye_l_h_temp = pose[0] + eye_r_h_temp = pose[1] + mouth_ratio = pose[2] + eye_y_ratio = pose[3] + eye_x_ratio = pose[4] + x_angle = pose[5] + y_angle = pose[6] + z_angle = pose[7] + + eyebrow_vector = [0.0] * 12 + mouth_eye_vector = [0.0] * 27 + pose_vector = [0.0] * 6 + + mouth_eye_vector[2] = max(eye_l_h_temp, eye_r_h_temp) + mouth_eye_vector[3] = max(eye_l_h_temp, eye_r_h_temp) + + mouth_eye_vector[14] = mouth_ratio * 2.0 + + mouth_eye_vector[25] = 0.0 # keep iris stable for user demo + mouth_eye_vector[26] = 0.0 + if position_offset is None: + position_offset = [(x_angle - 1.5) * 1.6, y_angle * 2.0 , (z_angle + 1.5) * 2] + pose_vector[0] = (x_angle - 1.5) * 1.6 - position_offset[0] + pose_vector[1] = y_angle * 2.0 - position_offset[1] # temp weight + pose_vector[2] = (z_angle + 1.5) * 2 - position_offset[2] # temp weight + pose_vector[3] = pose_vector[1] + pose_vector[4] = pose_vector[2] + pose_vector[5] = breath_value + + model_input_arr = eyebrow_vector + model_input_arr.extend(mouth_eye_vector) + model_input_arr.extend(pose_vector) + + with pose_position_shm_guard.lock(): + np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) + np_position_shm[:] = position_vector \ No newline at end of file diff --git a/src/i_facial_mocap_client.py b/src/i_facial_mocap_client.py new file mode 100644 index 0000000..f812a31 --- /dev/null +++ b/src/i_facial_mocap_client.py @@ -0,0 +1,133 @@ +from multiprocessing import Process, shared_memory, Value +from .args import args +from .utils.timer_wait import wait_until +from .utils.shared_mem_guard import SharedMemoryGuard +import socket +import errno +import numpy as np +import time +import tha2.poser.modes.mode_20_wx +from tha2.mocap.ifacialmocap_constants import * +from .utils.fps import FPS +from .utils.filter import OneEuroFilterNumpy +ifm_converter = tha2.poser.modes.mode_20_wx.IFacialMocapPoseConverter20() + +class IFMClientProcess(Process): + def __init__(self, pose_position_shm: shared_memory.SharedMemory): + super().__init__() + self.pose_position_shm = pose_position_shm + self.address = args.ifm_input.split(':')[0] + self.port = int(args.ifm_input.split(':')[1]) + self.fps = Value('f', 0.0) + def run(self): + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) + np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) + + udpClntSock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + + data = "iFacialMocap_sahuasouryya9218sauhuiayeta91555dy3719" + + data = data.encode('utf-8') + + udpClntSock.sendto(data, (self.address, self.port)) + + self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.socket.bind(("", self.port)) + self.socket.settimeout(5.0) # give ifm app some time to respond + + print("Warming up iFacialMocap connection...") + + frame_count = 0 + input_fps = FPS(60) + while frame_count < 60: + socket_bytes = self.socket.recv(8192) + if not socket_bytes: + raise Exception("Can't receive iFacialMocap data (stream end?).") + input_fps() + frame_count += 1 + + self.socket.settimeout(0.5) # 500 ms + + # pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + # position_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + # 呼吸循环参数 + breath_start_time = time.perf_counter() + print("iFacialMocap Input Running at {:.2f} FPS".format(input_fps.view())) + while True: + socket_bytes = self.socket.recv(8192) + + if not socket_bytes: + raise Exception("Can't receive iFacialMocap data (stream end?).") + + self.fps.value = input_fps() + + socket_string = socket_bytes.decode("utf-8") + + try: + data = self.convert_from_blender_data(socket_string) + except Exception: + print("iFacialMocap data parse error:", socket_string) + continue + + # 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0） + breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle + # 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0 + # sin 在 0 到 π 之间从 0 到 1 到 0 + breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi) + + ifacialmocap_pose_converted = ifm_converter.convert(data) + position_vector = data[HEAD_BONE_QUAT] + + eyebrow_vector = [0.0] * 12 + mouth_eye_vector = [0.0] * 27 + pose_vector = [0.0] * 6 + for i in range(0, 12): + eyebrow_vector[i] = ifacialmocap_pose_converted[i] + for i in range(12, 39): + mouth_eye_vector[i - 12] = ifacialmocap_pose_converted[i] + for i in range(39, 42): + pose_vector[i - 39] = ifacialmocap_pose_converted[i] + pose_vector[3] = pose_vector[1] + pose_vector[4] = pose_vector[2] + pose_vector[5] = breath_value + + mouth_eye_vector[14] = mouth_eye_vector[14] * 1.5 + + model_input_arr = eyebrow_vector + model_input_arr.extend(mouth_eye_vector) + model_input_arr.extend(pose_vector) + + with pose_position_shm_guard.lock(): + # np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) + # np_position_shm[:] = position_filter(np.array(position_vector, dtype=np.float32)) + np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) + np_position_shm[:] = np.array(position_vector, dtype=np.float32) + + @staticmethod + def convert_from_blender_data(blender_data): + data = {} + + for item in blender_data.split('|'): + if item.find('#') != -1: + k, arr = item.split('#') + arr = [float(n) for n in arr.split(',')] + data[k.replace("_L", "Left").replace("_R", "Right")] = arr + elif item.find('-') != -1: + k, v = item.split("-") + data[k.replace("_L", "Left").replace("_R", "Right")] = float(v) / 100 + + to_rad = 57.3 + data[HEAD_BONE_X] = data["=head"][0] / to_rad + data[HEAD_BONE_Y] = data["=head"][1] / to_rad + data[HEAD_BONE_Z] = data["=head"][2] / to_rad + data[HEAD_BONE_QUAT] = [data["=head"][3], data["=head"][4], data["=head"][5], 1] + # print(data[HEAD_BONE_QUAT][2],min(data[EYE_BLINK_LEFT],data[EYE_BLINK_RIGHT])) + data[RIGHT_EYE_BONE_X] = data["rightEye"][0] / to_rad + data[RIGHT_EYE_BONE_Y] = data["rightEye"][1] / to_rad + data[RIGHT_EYE_BONE_Z] = data["rightEye"][2] / to_rad + data[LEFT_EYE_BONE_X] = data["leftEye"][0] / to_rad + data[LEFT_EYE_BONE_Y] = data["leftEye"][1] / to_rad + data[LEFT_EYE_BONE_Z] = data["leftEye"][2] / to_rad + + return data \ No newline at end of file diff --git a/src/main.py b/src/main.py new file mode 100644 index 0000000..2545b5e --- /dev/null +++ b/src/main.py @@ -0,0 +1,159 @@ +import numpy as np +import time +from .model_infer_client import ModelClientProcess +from .args import args +from .utils.preprocess import resize_to_512_center, apply_color_curves +import cv2 +from .utils.shared_mem_guard import SharedMemoryGuard +from multiprocessing import shared_memory +from .utils.timer_wait import wait_until +from PIL import Image +from .utils.fps import FPS +import pyvirtualcam +from OpenGL.GL import GL_RGBA + + +def main(): + # Load character image + 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) + + print("Character Image Loaded:", args.character) + + pose_position_shm = shared_memory.SharedMemory(create=True, + size=(45 + 4) * 4) # 45 floats for pose, 4 floats for position + input_process = None + + if args.cam_input: + from .face_mesh_client import FaceMeshClientProcess + input_process = FaceMeshClientProcess(pose_position_shm) + elif args.ifm_input is not None: + from .i_facial_mocap_client import IFMClientProcess + input_process = IFMClientProcess(pose_position_shm) + elif args.osf_input is not None: + from .open_see_face_client import OSFClientProcess + input_process = OSFClientProcess(pose_position_shm) + elif args.mouse_input is not None: + from .mouse_client import MouseClientProcess + input_process = MouseClientProcess(pose_position_shm) + else: + from .debug_input_client import DebugInputClientProcess + input_process = DebugInputClientProcess(pose_position_shm) + + input_fps = input_process.fps + input_process.daemon = True + input_process.start() + + infer_process = ModelClientProcess(input_image, pose_position_shm, input_fps) + infer_process.daemon = True + infer_process.start() + + cam_width_scale = 2 if args.alpha_split else 1 + ret_channels = 3 if args.output_virtual_cam or args.output_debug else 4 + ret_batch_shm_channels = [ + SharedMemoryGuard(infer_process.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}") + for i in range(args.interpolation_scale) + ] + np_ret_shms = [ + np.ndarray((args.model_output_size, cam_width_scale * args.model_output_size, ret_channels), dtype=np.uint8, + buffer=infer_process.ret_shared_mem.buf[ + i * cam_width_scale * args.model_output_size * args.model_output_size * ret_channels: + (i + 1) * cam_width_scale * args.model_output_size * args.model_output_size * ret_channels]) + for i in range(args.interpolation_scale) + ] + + last_time: float = time.perf_counter() + interval: float = 1.0 / args.frame_rate_limit if args.frame_rate_limit > 0 else 0.0 + + if args.output_virtual_cam: + virtual_cam = pyvirtualcam.Camera(width=cam_width_scale * args.model_output_size, + height=args.model_output_size, + fps=args.frame_rate_limit, + backend='obs', + fmt=pyvirtualcam.PixelFormat.RGB) + print(f'Using virtual camera: {virtual_cam.device}') + elif args.output_spout2: + from PySpout import SpoutSender + spout_sender = SpoutSender("EasyVtuber", cam_width_scale * args.model_output_size, + args.model_output_size, GL_RGBA) + else: + print("Using OpenCV windows for output display.") + + pipeline_fps = FPS() + last_frame_time = None # 上一帧输出时间，用于打印帧时间差 + last_batch_start_time = None # 上一批就绪时间，用于周期估计 + n_frames = args.interpolation_scale + min_period = n_frames * interval if interval > 0 else n_frames / 60.0 # 60fps 下本批最少占用时间 + default_period = 1.0 / 15.0 # 约 15fps 推理时的周期，首包无历史时使用 + + print("Interval set to {:.3f} seconds".format(interval)) + while True: + infer_process.finish_event.wait() + infer_process.finish_event.clear() + for i in range(n_frames): + ret_batch_shm_channels[i].acquire() + + # 动态周期：本批就绪与上一批就绪的时间间隔，用于本批内均匀排期 + batch_start_time = time.perf_counter() + if last_batch_start_time is not None: + observed_period = batch_start_time - last_batch_start_time + period = max(min_period, min(observed_period, 1.0)) + else: + period = max(min_period, default_period) + last_batch_start_time = batch_start_time + + for i in range(n_frames): + # 均匀排期 + frame_rate_limit：取两者中较晚的时间发送 + target_send_time = batch_start_time + i * (period / n_frames) + if interval > 0: + target_send_time = max(target_send_time, last_time) + wait_until(target_send_time) + + if args.output_virtual_cam: + virtual_cam.send(np_ret_shms[i]) + elif args.output_spout2: + spout_sender.send_image(np_ret_shms[i], False) + else: + cv2.imshow("EasyVtuber Debug Frame", np_ret_shms[i]) + cv2.waitKey(1) + now_send = time.perf_counter() + last_frame_time = now_send + # 限速：下一帧最早在 last_time + interval，若已落后于当前时间则对齐到 now + if interval > 0: + last_time += interval + if last_time < now_send: + last_time = now_send + ret_batch_shm_channels[i].release() + output_pipeline_fps_val = pipeline_fps() * args.interpolation_scale + infer_process.output_pipeline_fps.value = output_pipeline_fps_val + print( + "Infer Process FPS: {:.2f}, Input FPS: {:.2f}, Model Avg Interval: {:.2f} ms, Cache Hit Ratio: {:.2f}%, GPU Cache Hit Ratio: {:.2f}%, Output Pipeline FPS {:.5f}".format( + infer_process.pipeline_fps_number.value, + input_fps.value, + infer_process.average_model_interval.value * 1000, + infer_process.cache_hit_ratio.value * 100, + infer_process.gpu_cache_hit_ratio.value * 100, + output_pipeline_fps_val + ), end='\r', flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/model_infer_client.py b/src/model_infer_client.py new file mode 100644 index 0000000..e41566c --- /dev/null +++ b/src/model_infer_client.py @@ -0,0 +1,194 @@ +import math +import os +from multiprocessing import Process, Value, shared_memory, Event +import cv2 +import numpy as np +import time +from .ezvtb_rt_interface import get_core +from .args import args +from .utils.shared_mem_guard import SharedMemoryGuard +from .utils.pose_simplify import pose_simplify +from .utils.fps import FPS, Interval +from typing import List + +class ModelClientProcess(Process): + def __init__(self, input_image, pose_position_shm: shared_memory.SharedMemory , input_fps): + super().__init__() + self.input_image = input_image + self.pose_position_shm = pose_position_shm # 45 floats for pose, 4 floats for position + + self.alpha_width_scale = 2 if args.alpha_split else 1 + self.ret_channels = 3 if args.output_virtual_cam or args.output_debug else 4 + self.ret_shape = (args.interpolation_scale, args.model_output_size, self.alpha_width_scale * args.model_output_size, self.ret_channels) + self.ret_nbytes = self.alpha_width_scale * args.interpolation_scale * args.model_output_size * args.model_output_size * self.ret_channels # RGBA + self.ret_shared_mem = shared_memory.SharedMemory(create=True, size=self.ret_nbytes) + + self.last_model_interval = Value('f', 0.0) + self.average_model_interval = Value('f', 0.0) + self.cache_hit_ratio = Value('f', 0.0) + self.gpu_cache_hit_ratio = Value('f', 0.0) + self.pipeline_fps_number = Value('f', 0.0) + self.output_pipeline_fps = Value('f', 0.0) # 由 main 进程更新，与 main print 一致 + self.input_fps = input_fps + + self.finish_event = Event() + + def run(self): + # 插值帧红点标记：由 --mark_interpolated 控制，供 ezvtb_rt 读取 + if getattr(args, 'mark_interpolated', False): + os.environ['EZVTB_MARK_INTERPOLATED'] = '1' + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) + np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) + + ret_batch_shm_guard = [ + SharedMemoryGuard(self.ret_shared_mem, ctrl_name=f"ret_shm_ctrl_batch_{i}") + for i in range(args.interpolation_scale) + ] + np_ret_shms = [ + np.ndarray((args.model_output_size, self.alpha_width_scale * args.model_output_size, self.ret_channels), dtype=np.uint8, + buffer=self.ret_shared_mem.buf[i * self.alpha_width_scale * args.model_output_size * args.model_output_size * self.ret_channels: + (i + 1) * self.alpha_width_scale * args.model_output_size * args.model_output_size * self.ret_channels]) + for i in range(args.interpolation_scale) + ] + + model_infer_average_interval: Interval = Interval() + pipeline_fps = FPS() + + # Use unified ezvtb_rt interface for both THA3 and THA4 + model = get_core(use_tensorrt=args.use_tensorrt, + model_version=args.model_version, + model_name=args.model_name, + + 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, + + cacher_ram_size=args.max_ram_cache_len, + + use_sr=args.use_sr, + sr_half=args.sr_half, + sr_x4=args.sr_x4, + sr_a4k=args.sr_a4k, + ) + model.setImage(self.input_image) + model_infer_average_interval.start() + model.inference([np.zeros((1, 45), dtype=np.float32)]) # Warm up + model_infer_average_interval.stop() + self.last_model_interval.value = model_infer_average_interval.last() + + last_pose = np.zeros((45,), dtype=np.float32) + + print("Model Inference Ready") + while True: + with pose_position_shm_guard.lock(): + np_pose = np_pose_shm.copy() + np_position = np_position_shm.copy() + + input_poses = [] + increment = (np_pose - last_pose) / args.interpolation_scale + for i in range(args.interpolation_scale): + input_poses.append(pose_simplify(last_pose + increment * (i + 1))) + last_pose = np_pose + + model_infer_average_interval.start() + output_images = model.inference(input_poses) + + if args.max_ram_cache_len > 0: + hits = model.cacher.hits + miss = model.cacher.miss + if args.use_sr: + hits += model.sr_cacher.hits + miss += model.sr_cacher.miss + total = hits + miss + self.cache_hit_ratio.value = (hits / total) if total > 0 else 0.0 + + if args.use_tensorrt and args.max_gpu_cache_len > 0: + hits = model.tha.cacher.hits + miss = model.tha.cacher.miss + total = hits + miss + self.gpu_cache_hit_ratio.value = (hits / total) if total > 0 else 0.0 + + output_images = self.post_process_ret(np_position, output_images) + + self.average_model_interval.value = model_infer_average_interval.stop() + self.last_model_interval.value = model_infer_average_interval.last() + + self.pipeline_fps_number.value = pipeline_fps() + for i in range(args.interpolation_scale): + with ret_batch_shm_guard[i].lock(): # get pressure from main process if ret not consumed + np_ret_shms[i][:, :, :] = output_images[i] + + self.finish_event.set() # Back pressure main process loop if infer slow + + def post_process_ret(self, np_position: np.ndarray, output_images: np.ndarray) -> List[np.ndarray]: + k_scale = 1 + rotate_angle = 0 + dx = 0 + dy = 0 + if args.extend_movement: + k_scale = np_position[2] * math.sqrt(args.extend_movement) + 1 + rotate_angle = -np_position[0] * 10 * args.extend_movement + dx = np_position[0] * 400 * k_scale * args.extend_movement + dy = -np_position[1] * 600 * k_scale * args.extend_movement + if args.bongo: + rotate_angle -= 5 + + rm = cv2.getRotationMatrix2D((output_images[0].shape[1] / 2, output_images[0].shape[0] / 2), rotate_angle, k_scale) + rm[0, 2] += dx + output_images[0].shape[1] / 2 - output_images[0].shape[1] / 2 + rm[1, 2] += dy + output_images[0].shape[0] / 2 - output_images[0].shape[0] / 2 + + ret = [] + for i in range(output_images.shape[0]): + bgra_image = output_images[i] + bgra_image = cv2.warpAffine( + bgra_image, + rm, + (bgra_image.shape[1], bgra_image.shape[0])) + + if args.output_debug: + # 与 main.py 输出格式一致 + y = 16 + cv2.putText(bgra_image, 'INFER/S: {:.4f}'.format(self.pipeline_fps_number.value), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + y += 16 + cv2.putText(bgra_image, 'INPUT/S: {:.4f}'.format(self.input_fps.value), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + y += 16 + cv2.putText(bgra_image, 'OUTPUT/S: {:.4f}'.format(self.output_pipeline_fps.value), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + y += 16 + cv2.putText(bgra_image, 'CALC: {:.2f} ms'.format(self.average_model_interval.value * 1000), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + y += 16 + if args.max_ram_cache_len > 0: + cv2.putText(bgra_image, 'MEM CACHE: {:.2f}%'.format(self.cache_hit_ratio.value * 100), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + y += 16 + if args.max_gpu_cache_len > 0: + cv2.putText(bgra_image, 'GPU CACHE: {:.2f}%'.format(self.gpu_cache_hit_ratio.value * 100), (0, y), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + + if args.alpha_split: + rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA) + alpha_channel = rgba_image[:, :, 3] + rgb_channels = rgba_image[:,:,:3] + alpha_image = cv2.cvtColor(alpha_channel, cv2.COLOR_GRAY2RGB) + rgb_channels = cv2.hconcat([rgb_channels, alpha_image]) + + if args.output_debug: + if args.alpha_split: + bgr_channels = cv2.cvtColor(rgb_channels, cv2.COLOR_RGB2BGR) + else: + bgr_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2BGR) + ret.append(bgr_channels) + elif args.output_virtual_cam: + if not args.alpha_split: + rgb_channels = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGB) + ret.append(rgb_channels) + else: + rgba_image = cv2.cvtColor(bgra_image, cv2.COLOR_BGRA2RGBA) + ret.append(rgba_image) + return ret + +if __name__ == "__main__": + pass \ No newline at end of file diff --git a/src/mouse_client.py b/src/mouse_client.py new file mode 100644 index 0000000..62c92f8 --- /dev/null +++ b/src/mouse_client.py @@ -0,0 +1,183 @@ +from multiprocessing import Process, shared_memory, Value +from pynput.mouse import Controller +import numpy as np +import time +import sounddevice as sd +from .args import args +from .utils.shared_mem_guard import SharedMemoryGuard +from .utils.timer_wait import wait_until +from OneEuroFilter import OneEuroFilter +from .utils.fps import FPS +from .utils.timer_wait import wait_until + +class MouseClientProcess(Process): + def __init__(self, pose_position_shm: shared_memory.SharedMemory): + super().__init__() + self.pose_position_shm = pose_position_shm + self.fps = Value('f', 60.0) + self.audio_volume = Value('f', 0.0) + self.audio_running = False + self.audio_callback_fps = FPS(60) + + def audio_callback(self, indata, frames, time_info, status): + """音频回调函数，计算音量""" + if status: + print(f"Audio status: {status}") + + self.audio_callback_fps() + # 计算 RMS (均方根) 音量 + volume_norm = np.linalg.norm(indata) * 10 + + # 应用阈值和灵敏度 + if volume_norm < args.audio_threshold: + volume_norm = 0 + else: + volume_norm = (volume_norm - args.audio_threshold) * args.audio_sensitivity + # 限制在 0-1 范围 + volume_norm = np.clip(volume_norm, 0, 1) + self.audio_volume.value = volume_norm + + def run(self): + mouse = Controller() + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) + np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) + + # posLimit = [0, 0, 1920, 1080] + posLimit = [int(x) for x in args.mouse_input.split(',')] + + if args.mouse_audio_input: + # 启动音频流 (使用 WASAPI) + try: + # 获取默认输入设备（麦克风或系统音频） + # 如果要捕获系统音频输出，需要使用 loopback 设备 + audio_stream = sd.InputStream( + callback=self.audio_callback, + channels=1, + samplerate=16000, + blocksize=1024, + ) + audio_stream.start() + self.audio_running = True + wait_until(time.perf_counter() + 1.0) # 等待一秒以稳定音频输入 + print("Audio capture started (WASAPI)") + except Exception as e: + print(f"Failed to start audio capture: {e}") + print("Continuing without audio input...") + self.audio_running = False + audio_filter = OneEuroFilter(freq=self.audio_callback_fps.view(), mincutoff=10.0, beta=0.0) + + prev = { + 'eye_l_h_temp': 0, + 'eye_r_h_temp': 0, + 'mouth_ratio': 0, + 'eye_y_ratio': 0, + 'eye_x_ratio': 0, + 'x_angle': 0, + 'y_angle': 0, + 'z_angle': 0, + } + last_time : float = time.perf_counter() + interval : float = 1.0 / 60 # 60 FPS + + # 眨眼相关参数 + blink_interval = args.blink_interval # 眨眼间隔时间（秒） + blink_duration = 0.25 # 眨眼持续时间（秒） + last_blink_time = time.perf_counter() + + # 呼吸循环参数 + breath_start_time = time.perf_counter() + + print("Mouse Input Running at 60 FPS") + position_vector = np.array([0, 0, 0, 1], dtype=np.float32) + while True: + pos = mouse.position + # print(pos) + eye_limit = [0.8, 0.5] + head_eye_reduce = 0.6 + head_slowness = 0.2 + + # 计算眨眼效果 + current_time = time.perf_counter() + time_since_last_blink = current_time - last_blink_time + + # 每隔 blink_interval 秒眨一次眼 + if time_since_last_blink >= blink_interval: + last_blink_time = current_time + time_since_last_blink = 0 + + # 计算眨眼动画（使用尖锐的 sin 函数） + if time_since_last_blink < blink_duration: + # 在眨眼持续时间内，使用 sin 函数从 0 升到接近 1 再返回 0 + # 使用 sin^2 使曲线更尖锐 + blink_progress = time_since_last_blink / blink_duration * np.pi + blink_value = np.sin(blink_progress) ** 2 + eye_close = blink_value * 1.0 # 升到1，返回到 0 + else: + eye_close = 0.0 + + # 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0） + breath_elapsed = (current_time - breath_start_time) % args.breath_cycle + # 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0 + # sin 在 0 到 π 之间从 0 到 1 到 0 + breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi) + + # 获取音频音量并映射到 mouth_ratio + current_mouth_ratio = audio_filter(self.audio_volume.value, time.perf_counter()) if self.audio_running else 0 + mouse_data = { + 'eye_l_h_temp': eye_close, + 'eye_r_h_temp': eye_close, + 'mouth_ratio': current_mouth_ratio, + 'eye_y_ratio': np.interp(pos[1], [posLimit[1], posLimit[3]], [1, -1]) * eye_limit[1], + 'eye_x_ratio': np.interp(pos[0], [posLimit[0], posLimit[2]], [1, -1]) * eye_limit[0], + 'x_angle': np.interp(pos[1], [posLimit[1], posLimit[3]], [1, -1]), + 'y_angle': np.interp(pos[0], [posLimit[0], posLimit[2]], [1, -1]), + 'z_angle': 0, + } + mouse_data['x_angle'] = np.interp(head_slowness, [0, 1], [prev['x_angle'], mouse_data['x_angle']]) + mouse_data['y_angle'] = np.interp(head_slowness, [0, 1], [prev['y_angle'], mouse_data['y_angle']]) + mouse_data['eye_y_ratio'] -= mouse_data['x_angle'] * eye_limit[1] * head_eye_reduce + mouse_data['eye_x_ratio'] -= mouse_data['y_angle'] * eye_limit[0] * head_eye_reduce + if args.bongo: + mouse_data['y_angle'] += 0.05 + mouse_data['x_angle'] += 0.05 + prev = mouse_data + + eye_l_h_temp = mouse_data['eye_l_h_temp'] + eye_r_h_temp = mouse_data['eye_r_h_temp'] + mouth_ratio = mouse_data['mouth_ratio'] + eye_y_ratio = mouse_data['eye_y_ratio'] + eye_x_ratio = mouse_data['eye_x_ratio'] + x_angle = mouse_data['x_angle'] + y_angle = mouse_data['y_angle'] + z_angle = mouse_data['z_angle'] + + eyebrow_vector = [0.0] * 12 + mouth_eye_vector = [0.0] * 27 + pose_vector = [0.0] * 6 + + mouth_eye_vector[2] = eye_l_h_temp + mouth_eye_vector[3] = eye_r_h_temp + + mouth_eye_vector[14] = mouth_ratio * 1.5 + + mouth_eye_vector[25] = eye_y_ratio + mouth_eye_vector[26] = eye_x_ratio + + pose_vector[0] = x_angle + pose_vector[1] = y_angle + pose_vector[2] = z_angle + pose_vector[3] = pose_vector[1] + pose_vector[4] = pose_vector[2] + pose_vector[5] = breath_value + + model_input_arr = eyebrow_vector + model_input_arr.extend(mouth_eye_vector) + model_input_arr.extend(pose_vector) + + with pose_position_shm_guard.lock(): + np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) + np_position_shm[:] = position_vector + + wait_until(last_time + interval) + last_time += interval \ No newline at end of file diff --git a/src/open_see_face_client.py b/src/open_see_face_client.py new file mode 100644 index 0000000..ff4f42a --- /dev/null +++ b/src/open_see_face_client.py @@ -0,0 +1,181 @@ +from multiprocessing import Process, Value, shared_memory +from .args import args +import socket +import struct +import numpy as np +from .utils.shared_mem_guard import SharedMemoryGuard +from .utils.fps import FPS +from .utils.filter import OneEuroFilterNumpy +from OneEuroFilter import OneEuroFilter +import time + +class OSFClientProcess(Process): + def __init__(self, pose_position_shm: shared_memory.SharedMemory): + super().__init__() + self.pose_position_shm = pose_position_shm + self.address = args.osf_input.split(':')[0] + self.port = int(args.osf_input.split(':')[1]) + self.fps = Value('f', 0.0) + + def run(self): + pose_position_shm_guard = SharedMemoryGuard(self.pose_position_shm, ctrl_name="pose_position_shm_ctrl") + np_pose_shm = np.ndarray((45,), dtype=np.float32, buffer=self.pose_position_shm.buf[:45 * 4]) + np_position_shm = np.ndarray((4,), dtype=np.float32, buffer=self.pose_position_shm.buf[45 * 4:45 * 4 + 4 * 4]) + + self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.socket.bind(("", self.port)) + self.socket.settimeout(None) # wait indefinitely + + print("Warming up OpenSeeFace connection...") + + frame_count = 0 + input_fps = FPS(60) + while frame_count < 5: + socket_bytes = self.socket.recv(8192) + if not socket_bytes: + raise Exception("Can't receive OpenSeeFace data (stream end?).") + input_fps() + frame_count += 1 + + position_vector_0 = None + position_vector = [0, 0, 0, 1] + # 呼吸循环参数 + breath_start_time = time.perf_counter() + pose_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + position_filter = OneEuroFilterNumpy(freq=input_fps.view(), mincutoff=args.filter_min_cutoff, beta=args.filter_beta) + iris_x_filter = OneEuroFilter(freq=input_fps.view(), mincutoff=0.001, beta=1.0) # extra filter for iris movement + iris_y_filter = OneEuroFilter(freq=input_fps.view(), mincutoff=0.001, beta=1.0) # extra filter for iris movement + rotation_offset = None + print("OpenSeeFace Input Running at {:.2f} FPS".format(input_fps.view())) + while True: + socket_bytes = self.socket.recv(8192) + + if not socket_bytes: + raise Exception("Can't receive OpenSeeFace data (stream end?).") + + self.fps.value = input_fps() + + try: + osf_raw = (struct.unpack('=di2f2fB1f4f3f3f68f136f210f14f', socket_bytes)) + # print(osf_raw[432:]) + data = {} + OpenSeeDataIndex = [ + 'time', + 'id', + 'cameraResolutionW', + 'cameraResolutionH', + 'rightEyeOpen', + 'leftEyeOpen', + 'got3DPoints', + 'fit3DError', + 'rawQuaternionX', + 'rawQuaternionY', + 'rawQuaternionZ', + 'rawQuaternionW', + 'rawEulerX', + 'rawEulerY', + 'rawEulerZ', + 'translationY', + 'translationX', + 'translationZ', + ] + for i in range(len(OpenSeeDataIndex)): + data[OpenSeeDataIndex[i]] = osf_raw[i] + data['translationY'] *= -1 + data['translationZ'] *= -1 + data['rotationY'] = data['rawEulerY'] - 10 + data['rotationX'] = (-data['rawEulerX'] + 360) % 360 - 180 + data['rotationZ'] = (data['rawEulerZ'] - 90) + OpenSeeFeatureIndex = [ + 'EyeLeft', + 'EyeRight', + 'EyebrowSteepnessLeft', + 'EyebrowUpDownLeft', + 'EyebrowQuirkLeft', + 'EyebrowSteepnessRight', + 'EyebrowUpDownRight', + 'EyebrowQuirkRight', + 'MouthCornerUpDownLeft', + 'MouthCornerInOutLeft', + 'MouthCornerUpDownRight', + 'MouthCornerInOutRight', + 'MouthOpen', + 'MouthWide' + ] + except Exception: + print("OpenSeeFace data parse error:", socket_bytes) + continue + + for i in range(68): + data['confidence' + str(i)] = osf_raw[i + 18] + for i in range(68): + data['pointsX' + str(i)] = osf_raw[i * 2 + 18 + 68] + data['pointsY' + str(i)] = osf_raw[i * 2 + 18 + 68 + 1] + for i in range(70): + data['points3DX' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2] + data['points3DY' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2 + 1] + data['points3DZ' + str(i)] = osf_raw[i * 3 + 18 + 68 + 68 * 2 + 2] + + for i in range(len(OpenSeeFeatureIndex)): + data[OpenSeeFeatureIndex[i]] = osf_raw[i + 432] + # print(data['rotationX'],data['rotationY'],data['rotationZ']) + + # 计算呼吸效果（使用 sin 函数，在 breath_cycle 时间内从 0 到 1 再到 0） + breath_elapsed = (time.perf_counter() - breath_start_time) % args.breath_cycle + # 使用 sin 函数，让值在一个周期内从 0 -> 1 -> 0 + # sin 在 0 到 π 之间从 0 到 1 到 0 + breath_value = np.sin(breath_elapsed / args.breath_cycle * np.pi) + + a = np.array([ + data['points3DX66'] - data['points3DX68'] + data['points3DX67'] - data['points3DX69'], + data['points3DY66'] - data['points3DY68'] + data['points3DY67'] - data['points3DY69'], + data['points3DZ66'] - data['points3DZ68'] + data['points3DZ67'] - data['points3DZ69'] + ]) + a = (a / np.linalg.norm(a)) + data['eyeRotationX'] = a[0] + data['eyeRotationY'] = a[1] + + eyebrow_vector = [0.0] * 12 + mouth_eye_vector = [0.0] * 27 + pose_vector = [0.0] * 6 + + mouth_eye_vector[2] = 1 - data['leftEyeOpen'] + mouth_eye_vector[3] = 1 - data['rightEyeOpen'] + + mouth_eye_vector[14] = max(data['MouthOpen'], 0) * 2 # Open larger mouth + # print(mouth_eye_vector[14]) + + mouth_eye_vector[25] = iris_x_filter(-data['eyeRotationY'] * 3 - (data['rotationX']) / 57.3 * 1.5, timestamp=time.perf_counter()) + mouth_eye_vector[26] = iris_y_filter(data['eyeRotationX'] * 3 + (data['rotationY']) / 57.3, timestamp=time.perf_counter()) + # print(mouth_eye_vector[25:27]) + eyebrow_vector[6] = data['EyebrowUpDownLeft'] + eyebrow_vector[7] = data['EyebrowUpDownRight'] + # print(data['EyebrowUpDownLeft'],data['EyebrowUpDownRight']) + if rotation_offset is None: + rotation_offset = [data['rotationX'], data['rotationY'], data['rotationZ']] + pose_vector[0] = (data['rotationX'] - rotation_offset[0]) / 57.3 * 3 + pose_vector[1] = -(data['rotationY'] - rotation_offset[1]) / 57.3 * 3 + pose_vector[2] = (data['rotationZ'] - rotation_offset[2]) / 57.3 * 2 + pose_vector[3] = pose_vector[1] + pose_vector[4] = pose_vector[2] + pose_vector[5] = breath_value + + if position_vector_0 == None: #Provide an initial reference point + position_vector_0 = [0, 0, 0, 1] + position_vector_0[0] = data['translationX'] + position_vector_0[1] = data['translationY'] + position_vector_0[2] = data['translationZ'] + #Compute relative translation + position_vector[0] = -(data['translationX'] - position_vector_0[0]) * 0.1 + position_vector[1] = -(data['translationY'] - position_vector_0[1]) * 0.1 + position_vector[2] = -(data['translationZ'] - position_vector_0[2]) * 0.1 + + model_input_arr = eyebrow_vector + model_input_arr.extend(mouth_eye_vector) + model_input_arr.extend(pose_vector) + + with pose_position_shm_guard.lock(): + np_pose_shm[:] = pose_filter(np.array(model_input_arr, dtype=np.float32)) + np_position_shm[:] = position_filter(np.array(position_vector, dtype=np.float32)) + # np_pose_shm[:] = np.array(model_input_arr, dtype=np.float32) + # np_position_shm[:] = np.array(position_vector, dtype=np.float32) diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/facial_points.py b/src/utils/facial_points.py similarity index 100% rename from facial_points.py rename to src/utils/facial_points.py diff --git a/src/utils/filter.py b/src/utils/filter.py new file mode 100644 index 0000000..2852c99 --- /dev/null +++ b/src/utils/filter.py @@ -0,0 +1,155 @@ +""" +Numpy-compatible wrapper for OneEuroFilter. + +Provides an adaptive One Euro Filter that automatically handles numpy arrays +by maintaining separate filter instances for each element. +""" + +import numpy as np +from OneEuroFilter import OneEuroFilter +from typing import Optional +import time + +class OneEuroFilterNumpy: + """ + Auto-initializing One Euro Filter for numpy arrays. + + Automatically determines shape from first input and maintains independent + OneEuroFilter instances for each element in the array. + + Example: + >>> filter = AdaptiveOneEuroFilterNumpy(freq=30, mincutoff=1.0, beta=0.01) + >>> noisy_data = np.array([1.0, 2.0, 3.0]) + >>> filtered = filter(noisy_data) + >>> filtered = filter(noisy_data + np.random.randn(3) * 0.1) + """ + + def __init__( + self, + freq: float, + mincutoff: float = 1.0, + beta: float = 0.0, + dcutoff: float = 1.0 + ) -> None: + """ + Initialize an adaptive One Euro Filter. + + Args: + freq: An estimate of the frequency in Hz of the signal (> 0) + mincutoff: Min cutoff frequency in Hz (> 0). Lower values remove more jitter + beta: Parameter to reduce latency (>= 0) + dcutoff: Cutoff frequency for derivatives (> 0) + """ + self.freq = freq + self.mincutoff = mincutoff + self.beta = beta + self.dcutoff = dcutoff + self._filters = None + self._shape = None + + def __call__( + self, + x: np.ndarray, + timestamp: Optional[float] = None + ) -> np.ndarray: + """ + Filter a noisy numpy array. + + On first call, initializes filters based on input shape. + + Args: + x: Noisy array to filter + timestamp: Timestamp in seconds (optional) + + Returns: + Filtered array with same shape as input + """ + # Initialize filters on first call + if self._filters is None: + self._shape = x.shape + size = np.prod(self._shape).astype(int) + self._filters = [ + OneEuroFilter( + freq=self.freq, + mincutoff=self.mincutoff, + beta=self.beta, + dcutoff=self.dcutoff + ) + for _ in range(size) + ] + + # Validate shape + if x.shape != self._shape: + raise ValueError( + f"Input shape {x.shape} doesn't match initialized shape {self._shape}" + ) + + if timestamp is None: + timestamp = time.perf_counter() + + # Flatten, filter, reshape + x_flat = x.flatten() + filtered_flat = np.array([ + self._filters[i](float(val), timestamp) + for i, val in enumerate(x_flat) + ]) + + return filtered_flat.reshape(self._shape) + + def filter( + self, + x: np.ndarray, + timestamp: Optional[float] = None + ) -> np.ndarray: + """Filter a noisy numpy array (alias for __call__).""" + return self.__call__(x, timestamp) + + def reset(self) -> None: + """Reset all internal filter states.""" + if self._filters is not None: + for f in self._filters: + f.reset() + + def setFrequency(self, freq: float) -> None: + """Set the frequency for all filters.""" + self.freq = freq + if self._filters is not None: + for f in self._filters: + f.setFrequency(freq) + + def setMinCutoff(self, mincutoff: float) -> None: + """Set the min cutoff frequency for all filters.""" + self.mincutoff = mincutoff + if self._filters is not None: + for f in self._filters: + f.setMinCutoff(mincutoff) + + def setBeta(self, beta: float) -> None: + """Set the beta parameter for all filters.""" + self.beta = beta + if self._filters is not None: + for f in self._filters: + f.setBeta(beta) + + def setDerivateCutoff(self, dcutoff: float) -> None: + """Set the derivative cutoff frequency for all filters.""" + self.dcutoff = dcutoff + if self._filters is not None: + for f in self._filters: + f.setDerivateCutoff(dcutoff) + + def setParameters( + self, + freq: float, + mincutoff: float = 1.0, + beta: float = 0.0, + dcutoff: float = 1.0 + ) -> None: + """Set all parameters for all filters.""" + self.freq = freq + self.mincutoff = mincutoff + self.beta = beta + self.dcutoff = dcutoff + if self._filters is not None: + for f in self._filters: + f.setParameters(freq, mincutoff, beta, dcutoff) \ No newline at end of file diff --git a/src/utils/fps.py b/src/utils/fps.py new file mode 100644 index 0000000..a7db58f --- /dev/null +++ b/src/utils/fps.py @@ -0,0 +1,55 @@ +import time +import collections + +class FPS: + def __init__(self, avarageof=50): + self.frametimestamps = collections.deque(maxlen=avarageof) + + def __call__(self): + self.frametimestamps.append(time.perf_counter()) + return self.view() + + def view(self): + if len(self.frametimestamps) > 1: + return len(self.frametimestamps) / (self.frametimestamps[-1] - self.frametimestamps[0] + 1e-10) + else: + return 0.0 + + def last(self): + if len(self.frametimestamps): + return self.frametimestamps[-1] + else: + return time.perf_counter() + +class Interval: + def __init__(self, avarageof=50): + self.frame_intervals = collections.deque(maxlen=avarageof) + self.start_time = None + self.interval = 0.0 + self.moving_sum = 0.0 + + def start(self): + self.start_time = time.perf_counter() + + def stop(self): + if self.start_time is not None: + last_interval = time.perf_counter() - self.start_time + + if len(self.frame_intervals) >= self.frame_intervals.maxlen: + self.moving_sum -= self.frame_intervals[0] + self.moving_sum += last_interval + + self.frame_intervals.append(last_interval) + self.start_time = None + + if len(self.frame_intervals) > 0: + self.interval = self.moving_sum / len(self.frame_intervals) + return self.interval + else: + raise RuntimeError("FPSInterval.stop() called before start()") + + def view(self): + return self.interval + + def last(self): + return self.frame_intervals[-1] if len(self.frame_intervals) > 0 else 0.0 \ No newline at end of file diff --git a/pose.py b/src/utils/pose.py similarity index 96% rename from pose.py rename to src/utils/pose.py index 04074da..3fa76d6 100644 --- a/pose.py +++ b/src/utils/pose.py @@ -12,10 +12,10 @@ from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_EYEB from mediapipe.python.solutions.face_mesh_connections import FACEMESH_RIGHT_IRIS from mediapipe.python.solutions.face_mesh_connections import FACEMESH_TESSELATION -from facial_points import MOUTH_TOP, MOUTH_BOTTOM, MOUTH_LEFT1, MOUTH_LEFT2, MOUTH_RIGHT -from facial_points import IRIS_L_TOP, IRIS_L_BOTTOM, IRIS_L_LEFT, IRIS_L_RIGHT -from facial_points import IRIS_R_TOP, IRIS_R_BOTTOM, IRIS_R_LEFT, IRIS_R_RIGHT -from utils import get_distance +from .facial_points import MOUTH_TOP, MOUTH_BOTTOM, MOUTH_LEFT1, MOUTH_LEFT2, MOUTH_RIGHT +from .facial_points import IRIS_L_TOP, IRIS_L_BOTTOM, IRIS_L_LEFT, IRIS_L_RIGHT +from .facial_points import IRIS_R_TOP, IRIS_R_BOTTOM, IRIS_R_LEFT, IRIS_R_RIGHT +from .utils import get_distance # pose vector index ------------------------------------------ diff --git a/src/utils/pose_simplify.py b/src/utils/pose_simplify.py new file mode 100644 index 0000000..4da1798 --- /dev/null +++ b/src/utils/pose_simplify.py @@ -0,0 +1,130 @@ +from ..args import args +import numpy as np +import math + +# Delay import to avoid loading torch during module initialization +ifm_converter = None + +def _get_ifm_converter(): + global ifm_converter + if ifm_converter is None: + import tha2.poser.modes.mode_20_wx + ifm_converter = tha2.poser.modes.mode_20_wx.IFacialMocapPoseConverter20() + return ifm_converter + +def pose_simplify(model_input): + ifm_converter = _get_ifm_converter() + simplify_arr = [1000] * ifm_converter.pose_size + if args.simplify >= 1: + simplify_arr = [200] * ifm_converter.pose_size + simplify_arr[ifm_converter.eye_wink_left_index] = 50 + simplify_arr[ifm_converter.eye_wink_right_index] = 50 + simplify_arr[ifm_converter.eye_happy_wink_left_index] = 50 + simplify_arr[ifm_converter.eye_happy_wink_right_index] = 50 + simplify_arr[ifm_converter.eye_surprised_left_index] = 30 + simplify_arr[ifm_converter.eye_surprised_right_index] = 30 + simplify_arr[ifm_converter.iris_rotation_x_index] = 25 + simplify_arr[ifm_converter.iris_rotation_y_index] = 25 + simplify_arr[ifm_converter.eye_raised_lower_eyelid_left_index] = 10 + simplify_arr[ifm_converter.eye_raised_lower_eyelid_right_index] = 10 + simplify_arr[ifm_converter.mouth_lowered_corner_left_index] = 5 + simplify_arr[ifm_converter.mouth_lowered_corner_right_index] = 5 + simplify_arr[ifm_converter.mouth_raised_corner_left_index] = 5 + simplify_arr[ifm_converter.mouth_raised_corner_right_index] = 5 + if args.simplify >= 2: + simplify_arr[ifm_converter.head_x_index] = 100 + simplify_arr[ifm_converter.head_y_index] = 100 + simplify_arr[ifm_converter.eye_surprised_left_index] = 10 + simplify_arr[ifm_converter.eye_surprised_right_index] = 10 + model_input[ifm_converter.eye_wink_left_index] += model_input[ + ifm_converter.eye_happy_wink_left_index] + model_input[ifm_converter.eye_happy_wink_left_index] = model_input[ + ifm_converter.eye_wink_left_index] / 2 + model_input[ifm_converter.eye_wink_left_index] = model_input[ + ifm_converter.eye_wink_left_index] / 2 + model_input[ifm_converter.eye_wink_right_index] += model_input[ + ifm_converter.eye_happy_wink_right_index] + model_input[ifm_converter.eye_happy_wink_right_index] = model_input[ + ifm_converter.eye_wink_right_index] / 2 + model_input[ifm_converter.eye_wink_right_index] = model_input[ + ifm_converter.eye_wink_right_index] / 2 + + uosum = model_input[ifm_converter.mouth_uuu_index] + \ + model_input[ifm_converter.mouth_ooo_index] + model_input[ifm_converter.mouth_ooo_index] = uosum + model_input[ifm_converter.mouth_uuu_index] = 0 + is_open = (model_input[ifm_converter.mouth_aaa_index] + model_input[ + ifm_converter.mouth_iii_index] + uosum) > 0 + model_input[ifm_converter.mouth_lowered_corner_left_index] = 0 + model_input[ifm_converter.mouth_lowered_corner_right_index] = 0 + model_input[ifm_converter.mouth_raised_corner_left_index] = 0.5 if is_open else 0 + model_input[ifm_converter.mouth_raised_corner_right_index] = 0.5 if is_open else 0 + simplify_arr[ifm_converter.mouth_lowered_corner_left_index] = 0 + simplify_arr[ifm_converter.mouth_lowered_corner_right_index] = 0 + simplify_arr[ifm_converter.mouth_raised_corner_left_index] = 0 + simplify_arr[ifm_converter.mouth_raised_corner_right_index] = 0 + if args.simplify >= 3: + simplify_arr[ifm_converter.iris_rotation_x_index] = 20 + simplify_arr[ifm_converter.iris_rotation_y_index] = 20 + simplify_arr[ifm_converter.eye_wink_left_index] = 32 + simplify_arr[ifm_converter.eye_wink_right_index] = 32 + simplify_arr[ifm_converter.eye_happy_wink_left_index] = 32 + simplify_arr[ifm_converter.eye_happy_wink_right_index] = 32 + if args.simplify >= 4: + simplify_arr[ifm_converter.head_x_index] = 50 + simplify_arr[ifm_converter.head_y_index] = 50 + simplify_arr[ifm_converter.neck_z_index] = 100 + model_input[ifm_converter.eye_raised_lower_eyelid_left_index] = 0 + model_input[ifm_converter.eye_raised_lower_eyelid_right_index] = 0 + simplify_arr[ifm_converter.iris_rotation_x_index] = 10 + simplify_arr[ifm_converter.iris_rotation_y_index] = 10 + simplify_arr[ifm_converter.eye_wink_left_index] = 24 + simplify_arr[ifm_converter.eye_wink_right_index] = 24 + simplify_arr[ifm_converter.eye_happy_wink_left_index] = 24 + simplify_arr[ifm_converter.eye_happy_wink_right_index] = 24 + simplify_arr[ifm_converter.eye_surprised_left_index] = 8 + simplify_arr[ifm_converter.eye_surprised_right_index] = 8 + model_input[ifm_converter.eye_wink_left_index] += model_input[ + ifm_converter.eye_wink_right_index] + model_input[ifm_converter.eye_wink_right_index] = model_input[ + ifm_converter.eye_wink_left_index] / 2 + model_input[ifm_converter.eye_wink_left_index] = model_input[ + ifm_converter.eye_wink_left_index] / 2 + + model_input[ifm_converter.eye_surprised_left_index] += model_input[ + ifm_converter.eye_surprised_right_index] + model_input[ifm_converter.eye_surprised_right_index] = model_input[ + ifm_converter.eye_surprised_left_index] / 2 + model_input[ifm_converter.eye_surprised_left_index] = model_input[ + ifm_converter.eye_surprised_left_index] / 2 + + model_input[ifm_converter.eye_happy_wink_left_index] += model_input[ + ifm_converter.eye_happy_wink_right_index] + model_input[ifm_converter.eye_happy_wink_right_index] = model_input[ + ifm_converter.eye_happy_wink_left_index] / 2 + model_input[ifm_converter.eye_happy_wink_left_index] = model_input[ + ifm_converter.eye_happy_wink_left_index] / 2 + model_input[ifm_converter.mouth_aaa_index] = min( + model_input[ifm_converter.mouth_aaa_index] + + model_input[ifm_converter.mouth_ooo_index] / 2 + + model_input[ifm_converter.mouth_iii_index] / 2 + + model_input[ifm_converter.mouth_uuu_index] / 2, 1 + ) + model_input[ifm_converter.mouth_ooo_index] = 0 + model_input[ifm_converter.mouth_iii_index] = 0 + model_input[ifm_converter.mouth_uuu_index] = 0 + for i in range(4, args.simplify): + simplify_arr = [max(math.ceil(x * 0.8), 5) for x in simplify_arr] + for i in range(0, len(simplify_arr)): + if simplify_arr[i] > 0: + model_input[i] = round(model_input[i] * simplify_arr[i]) / simplify_arr[i] + + input_pose = np.zeros((1, 45), dtype=np.float32) + if args.eyebrow: + for i in range(12): + input_pose[0, i] = model_input[i] + for i in range(27): + input_pose[0, i + 12] = model_input[i + 12] + for i in range(6): + input_pose[0, i + 12 + 27] = model_input[i + 27 + 12] + return input_pose \ No newline at end of file diff --git a/preprocess.py b/src/utils/preprocess.py similarity index 100% rename from preprocess.py rename to src/utils/preprocess.py diff --git a/src/utils/shared_mem_guard.py b/src/utils/shared_mem_guard.py new file mode 100644 index 0000000..c349d3e --- /dev/null +++ b/src/utils/shared_mem_guard.py @@ -0,0 +1,149 @@ +# windows_only_shm_lock.py +from __future__ import annotations +import ctypes +import ctypes.wintypes as wt +import hashlib +import time +from contextlib import contextmanager +from multiprocessing import shared_memory +from typing import Optional, Union + +# ---- WinAPI for Named Mutex ---- +kernel32 = ctypes.WinDLL('kernel32', use_last_error=True) + +# HANDLE CreateMutexW(LPSECURITY_ATTRIBUTES, BOOL, LPCWSTR); +_CreateMutexW = kernel32.CreateMutexW +_CreateMutexW.argtypes = [wt.LPVOID, wt.BOOL, wt.LPCWSTR] +_CreateMutexW.restype = wt.HANDLE + +# HANDLE OpenMutexW(DWORD, BOOL, LPCWSTR); +_OpenMutexW = kernel32.OpenMutexW +_OpenMutexW.argtypes = [wt.DWORD, wt.BOOL, wt.LPCWSTR] +_OpenMutexW.restype = wt.HANDLE + +# DWORD WaitForSingleObject(HANDLE, DWORD); +_WaitForSingleObject = kernel32.WaitForSingleObject +_WaitForSingleObject.argtypes = [wt.HANDLE, wt.DWORD] +_WaitForSingleObject.restype = wt.DWORD + +# BOOL ReleaseMutex(HANDLE); +_ReleaseMutex = kernel32.ReleaseMutex +_ReleaseMutex.argtypes = [wt.HANDLE] +_ReleaseMutex.restype = wt.BOOL + +# BOOL CloseHandle(HANDLE); +_CloseHandle = kernel32.CloseHandle +_CloseHandle.argtypes = [wt.HANDLE] +_CloseHandle.restype = wt.BOOL + +# Constants +INFINITE = 0xFFFFFFFF +WAIT_OBJECT_0 = 0x00000000 +WAIT_TIMEOUT = 0x00000102 +WAIT_ABANDONED = 0x00000080 +WAIT_FAILED = 0xFFFFFFFF +MUTEX_ALL_ACCESS = 0x1F0001 +ERROR_ALREADY_EXISTS = 183 + + +class SharedMemoryGuard: + """ + Two-process exclusive access wrapper using Windows Named Mutex: + - payload SharedMemory is provided externally + - Uses Windows Named Mutex for cross-process synchronization + - More direct and simpler than manual InterlockedCompareExchange implementation + """ + + def __init__( + self, + payload: shared_memory.SharedMemory, + ctrl_name: str + ): + # 1) attach payload (external) + self.payload_shm = payload + + # 2) create/open named mutex for synchronization + self.mutex_name = f"Global\\{ctrl_name}_mutex" + self._mutex_handle = self._create_or_open_mutex(self.mutex_name) + if self._mutex_handle is None or self._mutex_handle == 0: + error = ctypes.get_last_error() + raise OSError(f"Failed to create/open mutex {self.mutex_name}, error: {error}") + + @staticmethod + def _create_or_open_mutex(name: str) -> wt.HANDLE: + """Create or open a named mutex.""" + # Try to create the mutex + handle = _CreateMutexW(None, False, name) + if handle is None or handle == 0: + # If creation failed, try to open existing mutex + handle = _OpenMutexW(MUTEX_ALL_ACCESS, False, name) + return handle + + # ---- payload view: zero-copy for numpy ---- + def payload_view(self) -> memoryview: + return self.payload_shm.buf # user can slice if needed + + # ---- lock API ---- + def acquire(self, timeout_ms: Optional[int] = None) -> bool: + """ + Acquire exclusive access using Windows Mutex. + - timeout_ms=None => infinite wait + - spin parameter is kept for API compatibility but ignored (Mutex handles waiting efficiently) + + Returns True if lock acquired, False if timeout. + """ + if timeout_ms is None: + wait_time = INFINITE + else: + wait_time = max(0, timeout_ms) + + result = _WaitForSingleObject(self._mutex_handle, wait_time) + + if result == WAIT_OBJECT_0: + # Successfully acquired the mutex + return True + elif result == WAIT_ABANDONED: + # Previous owner died without releasing, but we now own it + # This is still considered a successful acquisition + return True + elif result == WAIT_TIMEOUT: + # Timeout occurred + return False + else: # WAIT_FAILED or other error + error = ctypes.get_last_error() + raise OSError(f"WaitForSingleObject failed with error: {error}") + + def release(self) -> None: + """Release the mutex.""" + if not _ReleaseMutex(self._mutex_handle): + error = ctypes.get_last_error() + raise OSError(f"ReleaseMutex failed with error: {error}") + + @contextmanager + def lock(self, timeout_ms: Optional[int] = None): + """ + Context manager for acquiring and releasing the lock. + - timeout_ms: timeout in milliseconds (None for infinite) + - spin: kept for API compatibility but ignored + """ + if not self.acquire(timeout_ms=timeout_ms): + raise TimeoutError("acquire() timeout") + try: + yield self + finally: + self.release() + + # ---- lifecycle ---- + def close(self) -> None: + """Close handles and shared memory.""" + if self._mutex_handle and self._mutex_handle != 0: + _CloseHandle(self._mutex_handle) + self._mutex_handle = None + self.payload_shm.close() + + def unlink_control_if_owner(self) -> None: + """ + Named mutex is automatically cleaned up by Windows when all handles are closed. + This method is kept for API compatibility but does nothing. + """ + pass diff --git a/src/utils/timer_wait.py b/src/utils/timer_wait.py new file mode 100644 index 0000000..e966548 --- /dev/null +++ b/src/utils/timer_wait.py @@ -0,0 +1,57 @@ +import ctypes +import time + +ntdll = ctypes.WinDLL("ntdll") + +desired = ctypes.c_ulong(5000) # 0.5 ms = 5000 * 100ns +current = ctypes.c_ulong() + +ntdll.NtSetTimerResolution(desired, True, ctypes.byref(current)) # Set timer resolution to 0.5ms + +# ------------------------------- +# 高精度等待函数 +# ------------------------------- +def wait_until( + target_time: float, + spin_threshold: float = 0.0005, + sleep_min: float = 0.001, +): + """ + 等待直到 perf_counter() >= target_time + + :param target_time: 目标时间（time.perf_counter() 的绝对时间，秒） + :param spin_threshold: 最后 busy-wait 窗口（秒，默认 0.5ms） + :param sleep_min: 允许 sleep 的最小时间（避免 0 sleep） + """ + + # 防止负时间 / NaN + if not isinstance(target_time, (int, float)): + return + + while True: + now = time.perf_counter() + remaining = target_time - now + + # 已经过点 + if remaining <= 0: + return + + # 进入 busy-wait 区间 + if remaining <= spin_threshold: + break + + # 可 sleep 区间 + sleep_time = remaining - spin_threshold + + # 防止 sleep(0) 或极小值 + if sleep_time >= sleep_min: + time.sleep(sleep_time) + else: + # 太小了，直接进入自旋 + break + + # -------- busy-wait 精对齐 -------- + # 使用局部变量减少属性查找开销 + perf = time.perf_counter + while perf() < target_time: + pass diff --git a/utils.py b/src/utils/utils.py similarity index 99% rename from utils.py rename to src/utils/utils.py index 30bf9fe..23773d3 100644 --- a/utils.py +++ b/src/utils/utils.py @@ -1,4 +1,3 @@ -import torch import numpy as np diff --git a/test.py b/test.py deleted file mode 100644 index 2ae533f..0000000 --- a/test.py +++ /dev/null @@ -1,258 +0,0 @@ -# from pyanime4k import ac -# import pyanime4k -# -# parameters = ac.Parameters() -# # enable HDN for ACNet -# parameters.HDN = True -# -# a = ac.AC( -# type=ac.ProcessorType.CPU_ACNet, -# ) -# a.set_arguments(parameters) -import struct - -import cv2 -# -# img = cv2.imread("character/test41.png") -# -# a.load_image_from_numpy(img,input_type=ac.AC_INPUT_BGR) -# -# a.get_processor_info() -# -# # start processing -# a.process() -# -# a.show_image() - -# import numpy as np -# import zlib -# import bz2 -# import lzma -# import time -# -# postprocessed_image=np.load('out.npy') -# tic=time.perf_counter() -# res=None -# for i in range(100): -# res=zlib.compress(postprocessed_image) -# print("zlib", len(res),(time.perf_counter() - tic) * 1000) -# tic = time.perf_counter() -# for i in range(100): -# zlib.decompress(res) -# print("zlib", len(res),(time.perf_counter() - tic) * 1000) -# tic = time.perf_counter() -# for i in range(100): -# res=bz2.compress(postprocessed_image) -# print("bz2",len(res), (time.perf_counter() - tic) * 1000) -# tic = time.perf_counter() -# for i in range(100): -# bz2.decompress(res) -# print("bz2", len(res),(time.perf_counter() - tic) * 1000) -# tic = time.perf_counter() -# for i in range(100): -# res=lzma.compress(postprocessed_image) -# print("lzma",len(res), (time.perf_counter() - tic) * 1000) -# tic = time.perf_counter() -# for i in range(100): -# lzma.decompress(res) -# print("lzma", len(res),(time.perf_counter() - tic) * 1000) -# tic = time.perf_counter() -# -# -# output_frame = cv2.cvtColor(postprocessed_image, cv2.COLOR_RGBA2BGRA) -# cv2.imshow("frame", output_frame) -# cv2.waitKey(100000) -# preview upscaled image -# img=a.save_image_to_numpy() -# -# print(img) -# import cv2 -# import numpy as np -# -# rgb=[ -# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]], -# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]], -# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]], -# [[1,2,3],[2,2,3],[3,2,3],[4,2,3]] -# ] -# a=[ -# [1,2,3,4], -# [1,2,3,4], -# [1,2,3,4], -# [1,2,3,4], -# ] -# -# print(cv2.merge((np.array(rgb),np.array(a)))) - - -from multiprocessing import Value, Process, Queue -import queue -import socket -import errno -import time -import tha2.poser.modes.mode_20_wx -from tha2.mocap.ifacialmocap_constants import * -import numpy as np - -ifm='127.0.0.1:11573' - - - -class ClientProcess(Process): - def __init__(self): - super().__init__() - self.queue = Queue() - self.should_terminate = Value('b', False) - self.address = ifm.split(':')[0] - self.port = int(ifm.split(':')[1]) - self.perf_time = 0 - # self.ifm_converter = tha2.poser.modes.mode_20_wx.create_ifacialmocap_pose_converter() - self.a_min=None - self.a_max=None - - def run(self): - self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) - self.socket.setblocking(False) - self.socket.bind(("", self.port)) - self.socket.settimeout(0.1) - while True: - if self.should_terminate.value: - break - try: - socket_bytes = self.socket.recv(8192) - except socket.error as e: - err = e.args[0] - if err == errno.EAGAIN or err == errno.EWOULDBLOCK or err == 'timed out': - continue - else: - raise e - - # socket_string = socket_bytes.decode("utf-8") - osf_raw=(struct.unpack('=di2f2fB1f4f3f3f68f136f210f14f',socket_bytes)) - # print(osf_raw[432:]) - data={} - OpenSeeDataIndex=[ - 'time', - 'id', - 'cameraResolutionW', - 'cameraResolutionH', - 'rightEyeOpen', - 'leftEyeOpen', - 'got3DPoints', - 'fit3DError', - 'rawQuaternionX', - 'rawQuaternionY', - 'rawQuaternionZ', - 'rawQuaternionW', - 'rawEulerX', - 'rawEulerY', - 'rawEulerZ', - 'translationY', - 'translationX', - 'translationZ', - ] - for i in range(len(OpenSeeDataIndex)): - data[OpenSeeDataIndex[i]]=osf_raw[i] - data['translationY']*=-1 - data['translationZ']*=-1 - data['rotationY']=data['rawEulerY'] - data['rotationX']=(-data['rawEulerX']+180)%360 - data['rotationZ']=(data['rawEulerZ']-90) - OpenSeeFeatureIndex=[ - 'EyeLeft', - 'EyeRight', - 'EyebrowSteepnessLeft', - 'EyebrowUpDownLeft', - 'EyebrowQuirkLeft', - 'EyebrowSteepnessRight', - 'EyebrowUpDownRight', - 'EyebrowQuirkRight', - 'MouthCornerUpDownLeft', - 'MouthCornerInOutLeft', - 'MouthCornerUpDownRight', - 'MouthCornerInOutRight', - 'MouthOpen', - 'MouthWide' - ] - - for i in range(68): - data['confidence'+str(i)]=osf_raw[i+18] - for i in range(68): - data['pointsX'+str(i)]=osf_raw[i*2+18+68] - data['pointsY'+str(i)]=osf_raw[i*2+18+68+1] - for i in range(70): - data['points3DX'+str(i)]=osf_raw[i*3+18+68+68*2] - data['points3DY'+str(i)]=osf_raw[i*3+18+68+68*2+1] - data['points3DZ'+str(i)]=osf_raw[i*3+18+68+68*2+2] - - for i in range(len(OpenSeeFeatureIndex)): - data[OpenSeeFeatureIndex[i]]=osf_raw[i+432] - # print(data['rotationX'],data['rotationY'],data['rotationZ']) - - a=np.array([ - data['points3DX66']-data['points3DX68']+data['points3DX67']-data['points3DX69'], - data['points3DY66']-data['points3DY68']+data['points3DY67']-data['points3DY69'], - data['points3DZ66']-data['points3DZ68']+data['points3DZ67']-data['points3DZ69'] - ]) - a=(a/np.linalg.norm(a)) - data['eyeRotationX']=a[0] - data['eyeRotationY']=a[1] - # blender_data = json.loads(socket_string) - # data = self.convert_from_blender_data(socket_string) - # data=self.ifm_converter.convert(data) - # print(data) - # if(self.a_max is None): - # self.a_max=[x for x in data] - # else: - # self.a_max=[max(data[i],self.a_max[i]) for i in range(len(data))] - # if(self.a_min is None): - # self.a_min=[x for x in data] - # else: - # self.a_min=[min(data[i],self.a_min[i]) for i in range(len(data))] - # print([[self.a_min[i],self.a_max[i]] for i in range(len(data))]) - cur_time = time.perf_counter() - fps = 1 / (cur_time - self.perf_time) - self.perf_time = cur_time - # print(fps) - try: - self.queue.put_nowait(data) - except queue.Full: - pass - self.queue.close() - self.socket.close() - @staticmethod - def convert_from_blender_data(blender_data): - data = {} - - for item in blender_data.split('|'): - if item.find('#') != -1: - k, arr = item.split('#') - arr = [float(n) for n in arr.split(',')] - data[k.replace("_L", "Left").replace("_R", "Right")] = arr - elif item.find('-') != -1: - k, v = item.split("-") - data[k.replace("_L", "Left").replace("_R", "Right")] = float(v) / 100 - - to_rad = 57.3 - data[HEAD_BONE_X] = data["=head"][0] / to_rad - data[HEAD_BONE_Y] = data["=head"][1] / to_rad - data[HEAD_BONE_Z] = data["=head"][2] / to_rad - data[HEAD_BONE_QUAT] = [data["=head"][3], data["=head"][4], data["=head"][5], 1] - # print(data[HEAD_BONE_QUAT][2],min(data[EYE_BLINK_LEFT],data[EYE_BLINK_RIGHT])) - data[RIGHT_EYE_BONE_X] = data["rightEye"][0] / to_rad - data[RIGHT_EYE_BONE_Y] = data["rightEye"][1] / to_rad - data[RIGHT_EYE_BONE_Z] = data["rightEye"][2] / to_rad - data[LEFT_EYE_BONE_X] = data["leftEye"][0] / to_rad - data[LEFT_EYE_BONE_Y] = data["leftEye"][1] / to_rad - data[LEFT_EYE_BONE_Z] = data["leftEye"][2] / to_rad - - return data - - -if __name__ == '__main__': - client_process = ClientProcess() - client_process.daemon = True - client_process.start() - - while True: - pass \ No newline at end of file diff --git a/tha4/onnx_export_tha4.py b/tha4/onnx_export_tha4.py new file mode 100644 index 0000000..b346189 --- /dev/null +++ b/tha4/onnx_export_tha4.py @@ -0,0 +1,232 @@ +import torch +from os.path import join +from pathlib import Path +import sys +import os +sys.path.append(os.getcwd()) +sys.path.append(join(os.getcwd(), 'tha4', 'src')) +import numpy as np +import PIL.Image +import cv2 +import onnx +from onnxsim import simplify +import onnxruntime as ort +from torch import Tensor +from torch.nn import Module +from typing import List, Optional +from torch.nn.functional import interpolate +import matplotlib.pyplot as plt + +# ================================================================================= +# START: Monkey Patch for torch.nn.functional.affine_grid (CORRECTED, DEVICE-AWARE) +# ================================================================================= +import torch.nn.functional + +def linspace_from_neg_one(num_steps, dtype=torch.float32, device='cpu'): + r = torch.linspace(-1, 1, num_steps, dtype=torch.float32, device=device) + if num_steps > 1: r = r * (num_steps - 1) / num_steps + return r +def create_grid(N, C, H, W, device='cpu'): + grid = torch.empty((N, H, W, 3), dtype=torch.float32, device=device) + grid.select(-1, 0).copy_(linspace_from_neg_one(W, device=device)) + grid.select(-1, 1).copy_(linspace_from_neg_one(H, device=device).unsqueeze_(-1)) + grid.select(-1, 2).fill_(1); return grid +def custom_affine_grid(theta, size, align_corners=False): + N, C, H, W = size; device = theta.device + grid = create_grid(N, C, H, W, device=device) + grid = grid.view(N, H * W, 3).bmm(theta.transpose(1, 2)) + grid = grid.view(N, H, W, 2); return grid +def apply_affine_grid_patch(): + torch.nn.functional.affine_grid = custom_affine_grid + print("✅ Applied monkey patch to torch.nn.functional.affine_grid (Device-aware version).") +# ================================================================================= +# END: Monkey Patch +# ================================================================================= + +apply_affine_grid_patch() + +if len(sys.argv) != 2: + print("Usage: python onnx_export_tha4.py "); print("Available models: float, half") + raise ValueError('Do not get model name') + +MODEL_NAME = sys.argv[1] +HALF = (MODEL_NAME == 'half') +DEVICE_NAME = 'cuda:0' +IMAGE_INPUT = os.path.join('.', 'data', 'images', 'character.png') + +TMP_DIR = join('.', 'onnx_model_tha4', 'tmp'); Path(TMP_DIR).mkdir(parents=True, exist_ok=True) +MODEL_DIR = join('.', 'onnx_model_tha4', 'fp32' if not HALF else 'fp16'); Path(MODEL_DIR).mkdir(parents=True, exist_ok=True) +TMP_FILE_WRITE = join(TMP_DIR, 'tmp.onnx') + +gpu_device = torch.device(DEVICE_NAME); cpu_device = torch.device('cpu') +gpu_dtype = torch.float16 if HALF else torch.float32 +cpu_dtype = torch.float32 + +print(f"Using GPU device for validation: {gpu_device}, dtype: {gpu_dtype}") +print(f"Using CPU device for ONNX export, dtype: {cpu_dtype}") + +def torch_srgb_to_linear(x: torch.Tensor): + x = torch.clip(x, 0.0, 1.0); return torch.where(torch.le(x, 0.04045), x / 12.92, ((x + 0.055) / 1.055) ** 2.4) +def torch_linear_to_srgb(x): + x = torch.clip(x, 0.0, 1.0); return torch.where(torch.le(x, 0.003130804953560372), x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) +def numpy_linear_to_srgb(x): + x = np.clip(x, 0.0, 1.0); return np.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) +def load_tha4_poser(device: torch.device, use_half: bool): + from tha4.poser.modes.mode_07 import create_poser + m_files = { 'eyebrow_decomposer': 'data/models/tha4/eyebrow_decomposer.pt', 'eyebrow_morphing_combiner': 'data/models/tha4/eyebrow_morphing_combiner.pt', 'face_morpher': 'data/models/tha4/face_morpher.pt', 'body_morpher': 'data/models/tha4/body_morpher.pt', 'upscaler': 'data/models/tha4/upscaler.pt' } + print(f"Loading THA4 poser onto {device}..."); poser = create_poser(device=device, module_file_names=m_files) + if use_half and device.type == 'cuda': print("Converting poser modules to half precision (FP16)..."); poser.half() + return poser + +gpu_poser = load_tha4_poser(gpu_device, HALF) +cpu_poser = load_tha4_poser(cpu_device, False) +pose_size = gpu_poser.get_num_parameters(); print(f"Pose parameters: {pose_size}") + +def resize_PIL_image(img, size=(512, 512)): return img.resize(size, PIL.Image.Resampling.LANCZOS) +def extract_PIL_image_from_filelike(path): return PIL.Image.open(path) +def extract_pytorch_image_from_PIL_image(img): + arr = np.array(img) + if arr.shape[2] != 4: raise ValueError("Input image must be RGBA") + t = torch.from_numpy(arr).float() / 255.0 + t[:, :, :3] = torch_srgb_to_linear(t[:, :, :3]) + t = t.permute(2, 0, 1) * 2.0 - 1.0; return t + +pil_image = resize_PIL_image(extract_PIL_image_from_filelike(IMAGE_INPUT)) +pt_img_gpu = extract_pytorch_image_from_PIL_image(pil_image).to(gpu_dtype).unsqueeze(0).to(gpu_device) +zero_pose_gpu = torch.zeros(1, pose_size, dtype=gpu_dtype, device=gpu_device) +with torch.no_grad(): poser_torch_res = gpu_poser.pose(pt_img_gpu, zero_pose_gpu) +print(f"Reference GPU output shape: {poser_torch_res.shape}") + +print("\n" + "="*80); print("Starting ONNX export process on CPU..."); print("="*80) +cpu_modules = cpu_poser.get_modules() +cv_img_numpy = cv2.imread(IMAGE_INPUT, cv2.IMREAD_UNCHANGED) +cv_img_tensor_cpu = torch.from_numpy(cv_img_numpy) + +class EyebrowDecomposerWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor) -> List[Tensor]: + img_f = img.to(cpu_dtype) / 255.0; img_f[:, :, :3] = torch_srgb_to_linear(img_f[:, :, :3]) + # --- THE FIX: Swap BGR to RGB --- + img_f = img_f[:, :, [2, 1, 0, 3]] + img_c = img_f.permute(2, 0, 1).unsqueeze(0) * 2.0 - 1.0 + crop = img_c[:, :, 64:192, 192:320]; res = self.m(crop) + return [res[3], res[0], img_c] +eyebrow_decomposer_wrapper = EyebrowDecomposerWrapper(cpu_modules['eyebrow_decomposer']).eval() +with torch.no_grad(): decomposer_res_cpu = eyebrow_decomposer_wrapper(cv_img_tensor_cpu) +torch.onnx.export(eyebrow_decomposer_wrapper, cv_img_tensor_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['input_image'], output_names=["background_layer", "eyebrow_layer", "image_prepared"]) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "decomposer.onnx")); print("✓ Decomposer exported") + +eyebrow_pose_cpu = torch.zeros(1, 12, dtype=cpu_dtype) +class EyebrowCombinerWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor, bg: Tensor, layer: Tensor, pose: Tensor) -> Tensor: + crop = img[:, :, 32:224, 160:352].clone() + morphed = self.m(bg, layer, pose)[2] + crop[:, :, 32:160, 32:160] = morphed; return crop +eyebrow_combiner_wrapper = EyebrowCombinerWrapper(cpu_modules['eyebrow_morphing_combiner']).eval() +combiner_inputs_cpu = (decomposer_res_cpu[2], decomposer_res_cpu[0], decomposer_res_cpu[1], eyebrow_pose_cpu) +with torch.no_grad(): combiner_res_cpu = eyebrow_combiner_wrapper(*combiner_inputs_cpu) +torch.onnx.export(eyebrow_combiner_wrapper, combiner_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['image_prepared', 'eyebrow_background_layer', "eyebrow_layer", 'eyebrow_pose'], output_names=['eyebrow_image']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "combiner.onnx")); print("✓ Combiner exported") + +face_pose_cpu = torch.zeros(1, 27, dtype=cpu_dtype) +class FaceMorpherWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor, crop: Tensor, pose: Tensor) -> List[Tensor]: + morphed = self.m(crop, pose)[0]; full = img.clone() + full[:, :, 32:224, 160:352] = morphed + half = interpolate(full, size=(256, 256), mode='bilinear', align_corners=False); return [full, half] +face_morpher_wrapper = FaceMorpherWrapper(cpu_modules['face_morpher']).eval() +morpher_inputs_cpu = (decomposer_res_cpu[2], combiner_res_cpu, face_pose_cpu) +with torch.no_grad(): morpher_res_cpu = face_morpher_wrapper(*morpher_inputs_cpu) +torch.onnx.export(face_morpher_wrapper, morpher_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['image_prepared', 'im_morpher_crop', 'face_pose'], output_names=['face_morphed_full', 'face_morphed_half']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "morpher.onnx")); print("✓ Face Morpher exported") + +rotation_pose_cpu = torch.zeros(1, 6, dtype=cpu_dtype) +class BodyMorpherWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, img: Tensor, pose: Tensor) -> List[Tensor]: + res = self.m(img, pose); return [res[0], res[3]] +body_morpher_wrapper = BodyMorpherWrapper(cpu_modules['body_morpher']).eval() +body_morpher_inputs_cpu = (morpher_res_cpu[1], rotation_pose_cpu) +with torch.no_grad(): body_morpher_res_cpu = body_morpher_wrapper(*body_morpher_inputs_cpu) +torch.onnx.export(body_morpher_wrapper, body_morpher_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['face_morphed_half', 'rotation_pose'], output_names=['half_res_posed_image', 'half_res_grid_change']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "body_morpher.onnx")); print("✓ Body Morpher exported") + +class UpscalerWrapper(Module): + def __init__(self, obj): super().__init__(); self.m = obj + def forward(self, rest_img: Tensor, posed_half: Tensor, grid_half: Tensor, pose: Tensor) -> List[Tensor]: + posed_full = interpolate(posed_half, size=(512, 512), mode='bilinear', align_corners=False) + grid_full = interpolate(grid_half, size=(512, 512), mode='bilinear', align_corners=False) + res = self.m(rest_img, posed_full, grid_full, pose)[0] + output = res / 2.0 + 0.5; output[:, :3, :, :] = torch_linear_to_srgb(output[:, :3, :, :]) + cv_res = (output.squeeze(0).permute(1, 2, 0)[:, :, [2, 1, 0, 3]] * 255).clamp(0, 255).to(torch.uint8) + return [output, cv_res] +upscaler_wrapper = UpscalerWrapper(cpu_modules['upscaler']).eval() +upscaler_inputs_cpu = (morpher_res_cpu[0], body_morpher_res_cpu[0], body_morpher_res_cpu[1], rotation_pose_cpu) +with torch.no_grad(): upscaler_res_cpu = upscaler_wrapper(*upscaler_inputs_cpu) +torch.onnx.export(upscaler_wrapper, upscaler_inputs_cpu, TMP_FILE_WRITE, opset_version=16, input_names=['rest_image', 'half_res_posed_image', 'half_res_grid_change', 'rotation_pose'], output_names=['result', 'cv_result']) +onnx_model = onnx.load(TMP_FILE_WRITE); onnx.checker.check_model(onnx_model); onnx_model_sim, check = simplify(onnx_model) +onnx.save(onnx_model_sim if check else onnx_model, join(MODEL_DIR, "upscaler.onnx")); print("✓ Upscaler exported") + +print("\n" + "="*80); print("Validation: Running ONNX inference..."); print("="*80) +class THA4ONNXRunner: + def __init__(self, model_dir, use_gpu=True): + providers = [("CUDAExecutionProvider", {"device_id": 0})] if use_gpu else ['CPUExecutionProvider'] + self.decomposer_sess = self._create_session(join(model_dir, 'decomposer.onnx'), providers) + self.combiner_sess = self._create_session(join(model_dir, "combiner.onnx"), providers) + self.morpher_sess = self._create_session(join(model_dir, "morpher.onnx"), providers) + self.body_morpher_sess = self._create_session(join(model_dir, "body_morpher.onnx"), providers) + self.upscaler_sess = self._create_session(join(model_dir, "upscaler.onnx"), providers) + print("✓ All ONNX models loaded") + def _create_session(self, path, providers): + try: return ort.InferenceSession(path, providers=providers) + except Exception: print(f"Failed with {providers}, falling back to CPU for {os.path.basename(path)}"); return ort.InferenceSession(path, providers=['CPUExecutionProvider']) + def run(self, img_path, p_eye, p_face, p_rot): + img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) + d_res = self.decomposer_sess.run(None, {'input_image': img}) + c_res = self.combiner_sess.run(None, {'image_prepared': d_res[2], 'eyebrow_background_layer': d_res[0], "eyebrow_layer": d_res[1], 'eyebrow_pose': p_eye}) + m_res = self.morpher_sess.run(None, {'image_prepared': d_res[2], 'im_morpher_crop': c_res[0], 'face_pose': p_face}) + b_res = self.body_morpher_sess.run(None, {'face_morphed_half': m_res[1], 'rotation_pose': p_rot}) + u_res = self.upscaler_sess.run(None, {'rest_image': m_res[0], 'half_res_posed_image': b_res[0], 'half_res_grid_change': b_res[1], 'rotation_pose': p_rot}) + return u_res[0] + +def save_comparison_images(onnx_img_np, torch_img_np, prefix=""): + onnx_hwc = np.transpose(onnx_img_np.squeeze(0), (1, 2, 0))[:, :, :3] + torch_hwc = np.transpose(torch_img_np.squeeze(0), (1, 2, 0))[:, :, :3] + onnx_hwc = np.clip(onnx_hwc, 0, 1) + torch_hwc = np.clip(torch_hwc, 0, 1) + onnx_path = f"{prefix}onnx_output.png"; torch_path = f"{prefix}pytorch_output.png" + plt.imsave(onnx_path, onnx_hwc); plt.imsave(torch_path, torch_hwc) + print(f"✅ Saved comparison images: {onnx_path} and {torch_path}") + +runner = THA4ONNXRunner(MODEL_DIR) +p_eye_z, p_face_z, p_rot_z = np.zeros((1, 12), dtype=np.float32), np.zeros((1, 27), dtype=np.float32), np.zeros((1, 6), dtype=np.float32) + +print("\nRunning ONNX inference with zero pose...") +onnx_output = runner.run(IMAGE_INPUT, p_eye_z, p_face_z, p_rot_z) +ref_np = poser_torch_res.cpu().numpy() / 2.0 + 0.5; ref_np[0, :3, :, :] = numpy_linear_to_srgb(ref_np[0, :3, :, :]) +save_comparison_images(onnx_output, ref_np, "zero_pose_") +mse = ((onnx_output - ref_np) ** 2).mean(); print(f"\n{'='*80}\nMSE (zero pose): {mse:.10f}\n{'='*80}") +ACCEPTABLE_MSE_THRESHOLD = 1e-2 +if mse < ACCEPTABLE_MSE_THRESHOLD: print(f"✅ Validation PASSED! MSE is within the acceptable range of {ACCEPTABLE_MSE_THRESHOLD}.") +else: print(f"⚠️ Warning: MSE is higher than expected! But this can be normal due to precision differences.") + +print("\n" + "="*80); print("Testing with random poses..."); print("="*80) +p_eye_r = np.random.randn(1, 12).astype(np.float32) * 0.5 +p_face_r = np.random.randn(1, 27).astype(np.float32) * 0.5 +p_rot_r = np.random.randn(1, 6).astype(np.float32) * 0.3 +onnx_output_r = runner.run(IMAGE_INPUT, p_eye_r, p_face_r, p_rot_r) +pose_r_gpu = torch.from_numpy(np.concatenate([p_eye_r, p_face_r, p_rot_r], axis=1)).to(gpu_dtype).to(gpu_device) +with torch.no_grad(): ref_output_r = gpu_poser.pose(pt_img_gpu, pose_r_gpu) +ref_np_r = ref_output_r.cpu().numpy() / 2.0 + 0.5; ref_np_r[0, :3, :, :] = numpy_linear_to_srgb(ref_np_r[0, :3, :, :]) +save_comparison_images(onnx_output_r, ref_np_r, "random_pose_") +mse_r = ((onnx_output_r - ref_np_r) ** 2).mean(); print(f"MSE (random pose): {mse_r:.10f}") +if mse_r < ACCEPTABLE_MSE_THRESHOLD: print(f"✅ Random pose test PASSED! (Threshold: {ACCEPTABLE_MSE_THRESHOLD})") +else: print(f"⚠️ Warning: High MSE with random pose! But this can be normal.") + +print(f"\n{'='*80}\nExport and validation completed!\nModels saved to: {MODEL_DIR}\n{'='*80}") \ No newline at end of file diff --git a/tha4/onnx_export_tha4_student.py b/tha4/onnx_export_tha4_student.py new file mode 100644 index 0000000..b6ca3a7 --- /dev/null +++ b/tha4/onnx_export_tha4_student.py @@ -0,0 +1,470 @@ +import torch +from os.path import join +from pathlib import Path +import sys +import os +sys.path.append(os.getcwd()) +sys.path.append(join(os.getcwd(), 'tha4', 'src')) +import numpy as np +import PIL.Image +import cv2 +import onnx +from onnxsim import simplify +import onnxruntime as ort +from torch import Tensor +from torch.nn import Module +from typing import List +import matplotlib.pyplot as plt +from onnxconverter_common import float16 + +# ================================================================================= +# START: Monkey Patch for torch.nn.functional.affine_grid (DEVICE-AWARE) +# ================================================================================= +import torch.nn.functional + +def linspace_from_neg_one(num_steps, dtype=torch.float32, device='cpu'): + r = torch.linspace(-1, 1, num_steps, dtype=torch.float32, device=device) + if num_steps > 1: + r = r * (num_steps - 1) / num_steps + return r + +def create_grid(N, C, H, W, device='cpu'): + grid = torch.empty((N, H, W, 3), dtype=torch.float32, device=device) + grid.select(-1, 0).copy_(linspace_from_neg_one(W, device=device)) + grid.select(-1, 1).copy_(linspace_from_neg_one(H, device=device).unsqueeze_(-1)) + grid.select(-1, 2).fill_(1) + return grid + +def custom_affine_grid(theta, size, align_corners=False): + N, C, H, W = size + device = theta.device + grid = create_grid(N, C, H, W, device=device) + grid = grid.view(N, H * W, 3).bmm(theta.transpose(1, 2)) + grid = grid.view(N, H, W, 2) + return grid + +def apply_affine_grid_patch(): + torch.nn.functional.affine_grid = custom_affine_grid + print("✅ Applied monkey patch to torch.nn.functional.affine_grid (Device-aware version).") + +# ================================================================================= +# END: Monkey Patch +# ================================================================================= + +apply_affine_grid_patch() + +# Parse command line arguments +if len(sys.argv) < 2: + print("Usage: python onnx_export_tha4_student.py [--fp32 or --fp16] [--cpu]") + print("Example: python onnx_export_tha4_student.py kanori_tha4 --fp32") + print(" python onnx_export_tha4_student.py kanori_tha4 --fp32 --cpu") + raise ValueError('Missing model name') + +MODEL_NAME = sys.argv[1] +HALF = False +FORCE_CPU = '--cpu' in sys.argv + +# Check CUDA availability +if FORCE_CPU: + DEVICE_NAME = 'cpu' + print("⚠️ Forced CPU mode enabled") +elif not torch.cuda.is_available(): + DEVICE_NAME = 'cpu' + print("⚠️ CUDA not available, automatically using CPU mode") +else: + DEVICE_NAME = 'cuda:0' + +# 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') + +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 = MODEL_BASE_PATH # Save directly in model folder +TMP_FILE_WRITE = join(TMP_DIR, 'tmp.onnx') + +validation_device = torch.device(DEVICE_NAME) +export_device = torch.device('cpu') +validation_dtype = torch.float16 if (HALF and DEVICE_NAME != 'cpu') else torch.float32 +export_dtype = torch.float32 + +if DEVICE_NAME == 'cpu': + print(f"Using CPU device for both validation and ONNX export, dtype: {export_dtype}") +else: + print(f"Using GPU device for validation: {validation_device}, dtype: {validation_dtype}") + print(f"Using CPU device for ONNX export, dtype: {export_dtype}") + +# Color space conversion utilities +def torch_srgb_to_linear(x: torch.Tensor): + x = torch.clip(x, 0.0, 1.0) + return torch.where(torch.le(x, 0.04045), x / 12.92, ((x + 0.055) / 1.055) ** 2.4) + +def torch_linear_to_srgb(x): + x = torch.clip(x, 0.0, 1.0) + return torch.where(torch.le(x, 0.003130804953560372), x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) + +def numpy_linear_to_srgb(x): + x = np.clip(x, 0.0, 1.0) + return np.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) + +# Load THA4 Student Model (mode_14: 2-step architecture) +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 + m_files = { + 'face_morpher': join(model_base_path, 'face_morpher.pt'), + 'body_morpher': join(model_base_path, 'body_morpher.pt') + } + + print(f"Loading THA4 Student poser from {model_base_path}...") + print(f" - Face Morpher: {m_files['face_morpher']}") + print(f" - Body Morpher: {m_files['body_morpher']}") + + poser = create_poser(device=device, module_file_names=m_files) + + if use_half and device.type == 'cuda': + print("Converting poser modules to half precision (FP16)...") + poser.half() + + return poser + +# Load posers for validation and export +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_BASE_PATH) + +pose_size = validation_poser.get_num_parameters() +print(f"Pose parameters: {pose_size}") + +# Image loading and preprocessing utilities +def resize_PIL_image(img, size=(512, 512)): + return img.resize(size, PIL.Image.Resampling.LANCZOS) + +def extract_PIL_image_from_filelike(path): + return PIL.Image.open(path) + +def extract_pytorch_image_from_PIL_image(img): + arr = np.array(img) + if arr.shape[2] != 4: + raise ValueError("Input image must be RGBA") + t = torch.from_numpy(arr).float() / 255.0 + t[:, :, :3] = torch_srgb_to_linear(t[:, :, :3]) + t = t.permute(2, 0, 1) * 2.0 - 1.0 + return t + +# Load test image and run reference inference +pil_image = resize_PIL_image(extract_PIL_image_from_filelike(IMAGE_INPUT)) +pt_img_validation = extract_pytorch_image_from_PIL_image(pil_image).to(validation_dtype).unsqueeze(0).to(validation_device) +zero_pose_validation = torch.zeros(1, pose_size, dtype=validation_dtype, device=validation_device) + +with torch.no_grad(): + poser_torch_res = validation_poser.pose(pt_img_validation, zero_pose_validation) +print(f"Reference output shape: {poser_torch_res.shape}") + +print("\n" + "="*80) +print("Starting ONNX export process...") +print("="*80) + +# Get export modules +export_modules = export_poser.get_modules() +cv_img_numpy = cv2.imread(IMAGE_INPUT, cv2.IMREAD_UNCHANGED) +cv_img_tensor_export = torch.from_numpy(cv_img_numpy) + +# ================================================================================= +# Export Face Morpher (Student Model Network 1) +# ================================================================================= +print("\n[1/2] Exporting Face Morpher...") + +class FaceMorpherWrapper(Module): + """ + Face Morpher for Student Model (mode_14) + - Input: 45-dimensional pose (first 39 used for face) + - Output: 128x128 face region + Note: SIREN-based face morpher doesn't need input image, only pose + """ + def __init__(self, face_morpher_module): + super().__init__() + self.face_morpher = face_morpher_module + + def forward(self, pose: Tensor) -> Tensor: + """ + Args: + pose: Full pose parameters [1, 39] + Returns: + face_output: Face morphed image [1, 4, 128, 128] in normalized format + """ + # Run face morpher (SIREN generates image from pose only) + face_output = self.face_morpher.forward(pose) + + return face_output + +face_morpher_wrapper = FaceMorpherWrapper(export_modules['face_morpher']).eval() + +# Test inputs +face_pose_export = torch.zeros(1, 39, dtype=export_dtype) + +# Export +with torch.no_grad(): + face_morpher_res_export = face_morpher_wrapper(face_pose_export) + +torch.onnx.export( + face_morpher_wrapper, + face_pose_export, + TMP_FILE_WRITE, + opset_version=16, + input_names=['pose'], + output_names=['face_morphed'], + external_data =False, + optimize =True, + verify =True, + dynamo =False, +) + +onnx_model = onnx.load(TMP_FILE_WRITE) +onnx.checker.check_model(onnx_model) +print(f" Original model inputs: {[inp.name for inp in onnx_model.graph.input]}") + +# onnx_model = float16.convert_float_to_float16(onnx_model, keep_io_types=True) +# Simplify with input/output names preservation +onnx_model_sim, check = simplify(onnx_model, skip_fuse_bn=True) +if check: + print(f" Simplified model inputs: {[inp.name for inp in onnx_model_sim.graph.input]}") + onnx.save(onnx_model_sim, join(MODEL_DIR, "face_morpher.onnx")) +else: + print(" Simplification failed, using original model") + onnx.save(onnx_model, join(MODEL_DIR, "face_morpher.onnx")) + +print("✓ Face Morpher exported") + +# ================================================================================= +# Export Body Morpher (Student Model Network 2) +# ================================================================================= +print("\n[2/2] Exporting Body Morpher...") + +class BodyMorpherWrapper(Module): + """ + Body Morpher for Student Model (mode_14) + - Input: Full 512x512 image with face morpher output embedded + - Output: Final morphed image + """ + def __init__(self, body_morpher_module): + super().__init__() + self.body_morpher = body_morpher_module + + def forward(self, img: Tensor, face_morphed: Tensor, pose: Tensor) -> List[Tensor]: + """ + Args: + img: Input image [H, W, 4] in uint8 BGR+A format + face_morphed: Face morphed output [1, 4, 128, 128] from face morpher + pose: Full pose parameters [1, 45] + Returns: + result: Final output in [0, 255] uint8 format (OpenCV compatible) + """ + # Convert image to normalized format + img_f = img.to(export_dtype) / 255.0 + img_f[0, :, :, :3] = torch_srgb_to_linear(img_f[0, :, :, :3]) + img_f = img_f[:, :, :, [2, 1, 0, 3]] # BGR to RGB + img_c = img_f.permute(0, 3, 1, 2) * 2.0 - 1.0 + + # Embed face morpher output into the full image + center_x = 256 + center_y = 128 + 16 # 144 + img_with_face = img_c.clone() + img_with_face[:, :, center_y-64:center_y+64, center_x-64:center_x+64] = face_morphed + + # Run body morpher + body_output = self.body_morpher.forward(img_with_face, pose) + + # Take the first output (blended image) + result = body_output[0] + + # Convert back to [0, 1] range + output = result / 2.0 + 0.5 + output[:, :3, :, :] = torch_linear_to_srgb(output[:, :3, :, :]) + + # Convert to OpenCV format (BGR, uint8) + cv_result = (output.squeeze(0).permute(1, 2, 0)[:, :, [2, 1, 0, 3]] * 255).clamp(0, 255).to(torch.uint8) + + return [output, cv_result] + +body_morpher_wrapper = BodyMorpherWrapper(export_modules['body_morpher']).eval() + +full_pose_export = torch.zeros(1, 45, dtype=export_dtype) + +cv_img_tensor_export = cv_img_tensor_export.unsqueeze(0) +# Export +with torch.no_grad(): + body_morpher_res_export = body_morpher_wrapper(cv_img_tensor_export, face_morpher_res_export, full_pose_export) + +torch.onnx.export( + body_morpher_wrapper, + (cv_img_tensor_export, face_morpher_res_export, full_pose_export), + TMP_FILE_WRITE, + opset_version=16, + input_names=['input_image', 'face_morphed', 'pose'], + output_names=['result', 'cv_result'], + external_data =False, + optimize =True, + verify =True, + dynamo =False, +) + +onnx_model = onnx.load(TMP_FILE_WRITE) +onnx.checker.check_model(onnx_model) +print(f" Original model inputs: {[inp.name for inp in onnx_model.graph.input]}") + +# onnx_model = float16.convert_float_to_float16(onnx_model, keep_io_types=True) +# for node in onnx_model.graph.node: +# if node.name == "/Cast": +# for attr in node.attribute: +# if attr.name == "to": +# attr.i = onnx.TensorProto.FLOAT16 # Set to 10 (FP16) +# break +# if node.name == "/Squeeze": +# node.input[0] = "/ScatterND_14_cast_to_result" + +# Simplify with input/output names preservation +onnx_model_sim, check = simplify(onnx_model, skip_fuse_bn=True) +if check: + print(f" Simplified model inputs: {[inp.name for inp in onnx_model_sim.graph.input]}") + onnx.save(onnx_model_sim, join(MODEL_DIR, "body_morpher.onnx")) +else: + print(" Simplification failed, using original model") + onnx.save(onnx_model, join(MODEL_DIR, "body_morpher.onnx")) + +print("✓ Body Morpher exported") + +# ================================================================================= +# Validation: ONNX Inference +# ================================================================================= +print("\n" + "="*80) +print("Validation: Running ONNX inference...") +print("="*80) + +class THA4StudentONNXRunner: + def __init__(self, model_dir, use_gpu=True): + providers = [("CUDAExecutionProvider", {"device_id": 0})] if use_gpu else ['CPUExecutionProvider'] + + self.face_morpher_sess = self._create_session(join(model_dir, 'face_morpher.onnx'), providers) + self.body_morpher_sess = self._create_session(join(model_dir, 'body_morpher.onnx'), providers) + + # Print actual input names for debugging + face_inputs = [inp.name for inp in self.face_morpher_sess.get_inputs()] + body_inputs = [inp.name for inp in self.body_morpher_sess.get_inputs()] + print(f"✓ All ONNX models loaded") + print(f" Face Morpher inputs: {face_inputs}") + print(f" Body Morpher inputs: {body_inputs}") + + def _create_session(self, path, providers): + try: + return ort.InferenceSession(path, providers=providers) + except Exception as e: + print(f"Failed with {providers}, falling back to CPU for {os.path.basename(path)}") + return ort.InferenceSession(path, providers=['CPUExecutionProvider']) + + def run(self, img_path, pose): + """ + Run inference with student model + + Args: + img_path: Path to input image + pose: Pose parameters [1, 45] + """ + img = cv2.imread(img_path, cv2.IMREAD_UNCHANGED) + + # Step 1: Face Morpher (SIREN only needs pose, generates face from scratch) + face_res = self.face_morpher_sess.run(None, { + 'pose': pose[:, :39] + }) + + # Step 2: Body Morpher (needs original image, face result, and pose) + body_res = self.body_morpher_sess.run(None, { + 'input_image': img, + 'face_morphed': face_res[0], + 'pose': pose + }) + + return body_res[0] + +def save_comparison_images(onnx_img_np, torch_img_np, prefix=""): + onnx_hwc = np.transpose(onnx_img_np.squeeze(0), (1, 2, 0))[:, :, :3] + torch_hwc = np.transpose(torch_img_np.squeeze(0), (1, 2, 0))[:, :, :3] + + onnx_hwc = np.clip(onnx_hwc, 0, 1) + torch_hwc = np.clip(torch_hwc, 0, 1) + + onnx_path = join(MODEL_DIR, f"{prefix}onnx_output.png") + torch_path = join(MODEL_DIR, f"{prefix}pytorch_output.png") + + plt.imsave(onnx_path, onnx_hwc) + plt.imsave(torch_path, torch_hwc) + + print(f"✅ Saved comparison images: {onnx_path} and {torch_path}") + +# Initialize ONNX runner +runner = THA4StudentONNXRunner(MODEL_DIR) +pose_zero = np.zeros((1, 45), dtype=np.float32) + +# Test with zero pose +print("\nRunning ONNX inference with zero pose...") +onnx_output = runner.run(IMAGE_INPUT, pose_zero) + +# Convert reference output to numpy +ref_np = poser_torch_res.cpu().numpy() / 2.0 + 0.5 +ref_np[0, :3, :, :] = numpy_linear_to_srgb(ref_np[0, :3, :, :]) + +# Save comparison images +save_comparison_images(onnx_output, ref_np, "zero_pose_") + +# Calculate MSE +mse = ((onnx_output - ref_np) ** 2).mean() +print(f"\n{'='*80}\nMSE (zero pose): {mse:.10f}\n{'='*80}") + +ACCEPTABLE_MSE_THRESHOLD = 1e-2 +if mse < ACCEPTABLE_MSE_THRESHOLD: + print(f"✅ Validation PASSED! MSE is within the acceptable range of {ACCEPTABLE_MSE_THRESHOLD}.") +else: + print(f"⚠️ Warning: MSE is higher than expected! But this can be normal due to precision differences.") + +# Test with random pose +print("\n" + "="*80) +print("Testing with random poses...") +print("="*80) + +pose_random = np.random.randn(1, 45).astype(np.float32) * 0.3 +onnx_output_r = runner.run(IMAGE_INPUT, pose_random) + +pose_r_validation = torch.from_numpy(pose_random).to(validation_dtype).to(validation_device) +with torch.no_grad(): + ref_output_r = validation_poser.pose(pt_img_validation, pose_r_validation) + +ref_np_r = ref_output_r.cpu().numpy() / 2.0 + 0.5 +ref_np_r[0, :3, :, :] = numpy_linear_to_srgb(ref_np_r[0, :3, :, :]) + +save_comparison_images(onnx_output_r, ref_np_r, "random_pose_") + +mse_r = ((onnx_output_r - ref_np_r) ** 2).mean() +print(f"MSE (random pose): {mse_r:.10f}") + +if mse_r < ACCEPTABLE_MSE_THRESHOLD: + print(f"✅ Random pose test PASSED! (Threshold: {ACCEPTABLE_MSE_THRESHOLD})") +else: + print(f"⚠️ Warning: High MSE with random pose! But this can be normal.") + +print(f"\n{'='*80}") +print("Export and validation completed!") +print(f"Models saved to: {MODEL_DIR}") +print("Generated files:") +print(f" - face_morpher.onnx") +print(f" - body_morpher.onnx") +print(f"{'='*80}") diff --git a/tha4/src/tha4/__init__.py b/tha4/src/tha4/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/app/__init__.py b/tha4/src/tha4/app/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/app/character_model_ifacialmocap_puppeteer.py b/tha4/src/tha4/app/character_model_ifacialmocap_puppeteer.py new file mode 100644 index 0000000..872d8d9 --- /dev/null +++ b/tha4/src/tha4/app/character_model_ifacialmocap_puppeteer.py @@ -0,0 +1,417 @@ +import os +import socket +import sys +import threading +import time +from typing import Optional + +import PIL.Image + +from tha4.shion.base.image_util import torch_linear_to_srgb +from tha4.image_util import convert_linear_to_srgb +from tha4.mocap.ifacialmocap_pose_converter_25 import create_ifacialmocap_pose_converter +from tha4.app.full_manual_poser import resize_PIL_image +from tha4.charmodel.character_model import CharacterModel + +sys.path.append(os.getcwd()) + +from tha4.mocap.ifacialmocap_pose import create_default_ifacialmocap_pose +from tha4.mocap.ifacialmocap_v2 import IFACIALMOCAP_PORT, IFACIALMOCAP_START_STRING, parse_ifacialmocap_v2_pose + +import torch +import wx + +from tha4.mocap.ifacialmocap_constants import * +from tha4.mocap.ifacialmocap_pose_converter import IFacialMocapPoseConverter + + +class FpsStatistics: + def __init__(self): + self.count = 100 + self.fps = [] + + def add_fps(self, fps): + self.fps.append(fps) + while len(self.fps) > self.count: + del self.fps[0] + + def get_average_fps(self): + if len(self.fps) == 0: + return 0.0 + else: + return sum(self.fps) / len(self.fps) + + +class MainFrame(wx.Frame): + IMAGE_SIZE = 512 + + def __init__(self, pose_converter: IFacialMocapPoseConverter, device: torch.device): + super().__init__(None, wx.ID_ANY, "iFacialMocap Puppeteer (Fuji)") + self.poser = None + self.pose_converter = pose_converter + self.device = device + + self.ifacialmocap_pose = create_default_ifacialmocap_pose() + self.source_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE) + self.result_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE) + self.wx_source_image = None + self.torch_source_image = None + self.last_pose = None + self.fps_statistics = FpsStatistics() + self.last_update_time = None + + self.create_receiving_socket() + self.create_ui() + self.create_timers() + self.Bind(wx.EVT_CLOSE, self.on_close) + + self.update_source_image_bitmap() + self.update_result_image_bitmap() + + def create_receiving_socket(self): + self.receiving_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + self.receiving_socket.bind(("", IFACIALMOCAP_PORT)) + self.receiving_socket.setblocking(False) + + def create_timers(self): + self.capture_timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_capture_panel, id=self.capture_timer.GetId()) + self.animation_timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_result_image_bitmap, id=self.animation_timer.GetId()) + + def on_close(self, event: wx.Event): + # Stop the timers + self.animation_timer.Stop() + self.capture_timer.Stop() + + # Close receiving socket + self.receiving_socket.close() + + # Destroy the windows + self.Destroy() + event.Skip() + + def on_start_capture(self, event: wx.Event): + capture_device_ip_address = self.capture_device_ip_text_ctrl.GetValue() + out_socket = None + try: + address = (capture_device_ip_address, IFACIALMOCAP_PORT) + out_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + out_socket.sendto(IFACIALMOCAP_START_STRING, address) + except Exception as e: + message_dialog = wx.MessageDialog(self, str(e), "Error!", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + finally: + if out_socket is not None: + out_socket.close() + + def read_ifacialmocap_pose(self): + if not self.animation_timer.IsRunning(): + return self.ifacialmocap_pose + socket_bytes = None + while True: + try: + socket_bytes = self.receiving_socket.recv(8192) + except socket.error as e: + break + if socket_bytes is not None: + socket_string = socket_bytes.decode("utf-8") + self.ifacialmocap_pose = parse_ifacialmocap_v2_pose(socket_string) + return self.ifacialmocap_pose + + def on_erase_background(self, event: wx.Event): + pass + + def create_animation_panel(self, parent): + self.animation_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.animation_panel_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.animation_panel.SetSizer(self.animation_panel_sizer) + self.animation_panel.SetAutoLayout(1) + + image_size = MainFrame.IMAGE_SIZE + + if True: + self.input_panel = wx.Panel(self.animation_panel, size=(image_size, image_size + 128), + style=wx.SIMPLE_BORDER) + self.input_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.input_panel.SetSizer(self.input_panel_sizer) + self.input_panel.SetAutoLayout(1) + self.animation_panel_sizer.Add(self.input_panel, 0, wx.FIXED_MINSIZE) + + self.source_image_panel = wx.Panel(self.input_panel, size=(image_size, image_size), style=wx.SIMPLE_BORDER) + self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel) + self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.input_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE) + + self.load_model_button = wx.Button(self.input_panel, wx.ID_ANY, "Load Model") + self.input_panel_sizer.Add(self.load_model_button, 1, wx.EXPAND) + self.load_model_button.Bind(wx.EVT_BUTTON, self.load_model) + + self.input_panel_sizer.Fit(self.input_panel) + + if True: + self.pose_converter.init_pose_converter_panel(self.animation_panel) + + if True: + self.animation_left_panel = wx.Panel(self.animation_panel, style=wx.SIMPLE_BORDER) + self.animation_left_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.animation_left_panel.SetSizer(self.animation_left_panel_sizer) + self.animation_left_panel.SetAutoLayout(1) + self.animation_panel_sizer.Add(self.animation_left_panel, 0, wx.EXPAND) + + self.result_image_panel = wx.Panel(self.animation_left_panel, size=(image_size, image_size), + style=wx.SIMPLE_BORDER) + self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel) + self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.animation_left_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE) + + separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5)) + self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND) + + background_text = wx.StaticText(self.animation_left_panel, label="--- Background ---", + style=wx.ALIGN_CENTER) + self.animation_left_panel_sizer.Add(background_text, 0, wx.EXPAND) + + self.output_background_choice = wx.Choice( + self.animation_left_panel, + choices=[ + "TRANSPARENT", + "GREEN", + "BLUE", + "BLACK", + "WHITE" + ]) + self.output_background_choice.SetSelection(0) + self.animation_left_panel_sizer.Add(self.output_background_choice, 0, wx.EXPAND) + + separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5)) + self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND) + + self.fps_text = wx.StaticText(self.animation_left_panel, label="") + self.animation_left_panel_sizer.Add(self.fps_text, wx.SizerFlags().Border()) + + self.animation_left_panel_sizer.Fit(self.animation_left_panel) + + self.animation_panel_sizer.Fit(self.animation_panel) + + def create_ui(self): + self.main_sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.main_sizer) + self.SetAutoLayout(1) + + self.capture_pose_lock = threading.Lock() + + self.create_connection_panel(self) + self.main_sizer.Add(self.connection_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.create_animation_panel(self) + self.main_sizer.Add(self.animation_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.create_capture_panel(self) + self.main_sizer.Add(self.capture_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.main_sizer.Fit(self) + + def create_connection_panel(self, parent): + self.connection_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.connection_panel_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.connection_panel.SetSizer(self.connection_panel_sizer) + self.connection_panel.SetAutoLayout(1) + + capture_device_ip_text = wx.StaticText(self.connection_panel, label="Capture Device IP:", style=wx.ALIGN_RIGHT) + self.connection_panel_sizer.Add(capture_device_ip_text, wx.SizerFlags(0).FixedMinSize().Border(wx.ALL, 3)) + + self.capture_device_ip_text_ctrl = wx.TextCtrl(self.connection_panel, value="192.168.0.1") + self.connection_panel_sizer.Add(self.capture_device_ip_text_ctrl, wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + self.start_capture_button = wx.Button(self.connection_panel, label="START CAPTURE!") + self.connection_panel_sizer.Add(self.start_capture_button, wx.SizerFlags(0).FixedMinSize().Border(wx.ALL, 3)) + self.start_capture_button.Bind(wx.EVT_BUTTON, self.on_start_capture) + + def create_capture_panel(self, parent): + self.capture_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.capture_panel_sizer = wx.FlexGridSizer(cols=5) + for i in range(5): + self.capture_panel_sizer.AddGrowableCol(i) + self.capture_panel.SetSizer(self.capture_panel_sizer) + self.capture_panel.SetAutoLayout(1) + + self.rotation_labels = {} + self.rotation_value_labels = {} + rotation_column_0 = self.create_rotation_column(self.capture_panel, RIGHT_EYE_BONE_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column_0, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + rotation_column_1 = self.create_rotation_column(self.capture_panel, LEFT_EYE_BONE_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column_1, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + rotation_column_2 = self.create_rotation_column(self.capture_panel, HEAD_BONE_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column_2, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + + def create_rotation_column(self, parent, rotation_names): + column_panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + column_panel_sizer = wx.FlexGridSizer(cols=2) + column_panel_sizer.AddGrowableCol(1) + column_panel.SetSizer(column_panel_sizer) + column_panel.SetAutoLayout(1) + + for rotation_name in rotation_names: + self.rotation_labels[rotation_name] = wx.StaticText( + column_panel, label=rotation_name, style=wx.ALIGN_RIGHT) + column_panel_sizer.Add(self.rotation_labels[rotation_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + self.rotation_value_labels[rotation_name] = wx.TextCtrl( + column_panel, style=wx.TE_RIGHT) + self.rotation_value_labels[rotation_name].SetValue("0.00") + self.rotation_value_labels[rotation_name].Disable() + column_panel_sizer.Add(self.rotation_value_labels[rotation_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + column_panel.GetSizer().Fit(column_panel) + return column_panel + + def paint_capture_panel(self, event: wx.Event): + self.update_capture_panel(event) + + def update_capture_panel(self, event: wx.Event): + data = self.ifacialmocap_pose + for rotation_name in ROTATION_NAMES: + value = data[rotation_name] + self.rotation_value_labels[rotation_name].SetValue("%0.2f" % value) + + @staticmethod + def convert_to_100(x): + return int(max(0.0, min(1.0, x)) * 100) + + def paint_source_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap) + + def update_source_image_bitmap(self): + dc = wx.MemoryDC() + dc.SelectObject(self.source_image_bitmap) + if self.wx_source_image is None: + self.draw_nothing_yet_string(dc) + else: + dc.Clear() + dc.DrawBitmap(self.wx_source_image, 0, 0, True) + del dc + + def draw_nothing_yet_string(self, dc): + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent("Nothing yet!") + dc.DrawText("Nothing yet!", (MainFrame.IMAGE_SIZE - w) // 2, (MainFrame.IMAGE_SIZE - h) // 2) + + def paint_result_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap) + + def update_result_image_bitmap(self, event: Optional[wx.Event] = None): + ifacialmocap_pose = self.read_ifacialmocap_pose() + current_pose = self.pose_converter.convert(ifacialmocap_pose) + if self.last_pose is not None and self.last_pose == current_pose: + return + self.last_pose = current_pose + + if self.torch_source_image is None or self.poser is None: + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + self.draw_nothing_yet_string(dc) + del dc + return + + pose = torch.tensor(current_pose, device=self.device, dtype=self.poser.get_dtype()) + + with torch.no_grad(): + output_image = self.poser.pose(self.torch_source_image, pose)[0].float() + output_image = torch.clip((output_image + 1.0) / 2.0, 0.0, 1.0) + output_image = convert_linear_to_srgb(output_image) + + background_choice = self.output_background_choice.GetSelection() + if background_choice == 0: + pass + else: + background = torch.zeros(4, output_image.shape[1], output_image.shape[2], device=self.device) + background[3, :, :] = 1.0 + if background_choice == 1: + background[1, :, :] = 1.0 + output_image = self.blend_with_background(output_image, background) + elif background_choice == 2: + background[2, :, :] = 1.0 + output_image = self.blend_with_background(output_image, background) + elif background_choice == 3: + output_image = self.blend_with_background(output_image, background) + else: + background[0:3, :, :] = 1.0 + output_image = self.blend_with_background(output_image, background) + + c, h, w = output_image.shape + output_image = 255.0 * torch.transpose(output_image.reshape(c, h * w), 0, 1).reshape(h, w, c) + output_image = output_image.byte() + + numpy_image = output_image.detach().cpu().numpy() + wx_image = wx.ImageFromBuffer(numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + dc.Clear() + dc.DrawBitmap(wx_bitmap, + (MainFrame.IMAGE_SIZE - numpy_image.shape[0]) // 2, + (MainFrame.IMAGE_SIZE - numpy_image.shape[1]) // 2, True) + del dc + + time_now = time.time_ns() + if self.last_update_time is not None: + elapsed_time = time_now - self.last_update_time + fps = 1.0 / (elapsed_time / 10 ** 9) + if self.torch_source_image is not None: + self.fps_statistics.add_fps(fps) + self.fps_text.SetLabelText("FPS = %0.2f" % self.fps_statistics.get_average_fps()) + self.last_update_time = time_now + + self.Refresh() + + def blend_with_background(self, numpy_image, background): + alpha = numpy_image[3:4, :, :] + color = numpy_image[0:3, :, :] + new_color = color * alpha + (1.0 - alpha) * background[0:3, :, :] + return torch.cat([new_color, background[3:4, :, :]], dim=0) + + def load_model(self, event: wx.Event): + dir_name = "data/character_models" + file_dialog = wx.FileDialog(self, "Choose a model", dir_name, "", "*.yaml", wx.FD_OPEN) + if file_dialog.ShowModal() == wx.ID_OK: + character_model_json_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + self.character_model = CharacterModel.load(character_model_json_file_name) + self.torch_source_image = self.character_model.get_character_image(self.device) + pil_image = resize_PIL_image( + PIL.Image.open(self.character_model.character_image_file_name), + (MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)) + w, h = pil_image.size + self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes()) + self.update_source_image_bitmap() + self.poser = self.character_model.get_poser(self.device) + except Exception: + message_dialog = wx.MessageDialog( + self, "Could not load character model " + character_model_json_file_name, "Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + self.Refresh() + + +if __name__ == "__main__": + device = torch.device('cuda:0') + + pose_converter = create_ifacialmocap_pose_converter() + + app = wx.App() + main_frame = MainFrame(pose_converter, device) + main_frame.Show(True) + main_frame.capture_timer.Start(10) + main_frame.animation_timer.Start(10) + app.MainLoop() diff --git a/tha4/src/tha4/app/character_model_manual_poser.py b/tha4/src/tha4/app/character_model_manual_poser.py new file mode 100644 index 0000000..53b6757 --- /dev/null +++ b/tha4/src/tha4/app/character_model_manual_poser.py @@ -0,0 +1,438 @@ +import logging +import os +import sys +import time +from typing import List + +from tha4.charmodel.character_model import CharacterModel +from tha4.image_util import resize_PIL_image, convert_output_image_from_torch_to_numpy +from tha4.poser.modes.mode_14 import get_pose_parameters + +sys.path.append(os.getcwd()) + +import PIL.Image +import torch +import wx + +from tha4.poser.poser import PoseParameterCategory, PoseParameterGroup + + +class MorphCategoryControlPanel(wx.Panel): + def __init__(self, + parent, + title: str, + pose_param_category: PoseParameterCategory, + param_groups: List[PoseParameterGroup]): + super().__init__(parent, style=wx.SIMPLE_BORDER) + self.pose_param_category = pose_param_category + self.sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.sizer) + self.SetAutoLayout(1) + + title_text = wx.StaticText(self, label=title, style=wx.ALIGN_CENTER) + self.sizer.Add(title_text, 0, wx.EXPAND) + + self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category] + self.choice = wx.Choice(self, choices=[group.get_group_name() for group in self.param_groups]) + if len(self.param_groups) > 0: + self.choice.SetSelection(0) + self.choice.Bind(wx.EVT_CHOICE, self.on_choice_updated) + self.sizer.Add(self.choice, 0, wx.EXPAND) + + self.left_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL) + self.sizer.Add(self.left_slider, 0, wx.EXPAND) + + self.right_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL) + self.sizer.Add(self.right_slider, 0, wx.EXPAND) + + self.checkbox = wx.CheckBox(self, label="Show") + self.checkbox.SetValue(True) + self.sizer.Add(self.checkbox, 0, wx.SHAPED | wx.ALIGN_CENTER) + + self.update_ui() + + self.sizer.Fit(self) + + def update_ui(self): + param_group = self.param_groups[self.choice.GetSelection()] + if param_group.is_discrete(): + self.left_slider.Enable(False) + self.right_slider.Enable(False) + self.checkbox.Enable(True) + elif param_group.get_arity() == 1: + self.left_slider.Enable(True) + self.right_slider.Enable(False) + self.checkbox.Enable(False) + else: + self.left_slider.Enable(True) + self.right_slider.Enable(True) + self.checkbox.Enable(False) + + def on_choice_updated(self, event: wx.Event): + param_group = self.param_groups[self.choice.GetSelection()] + if param_group.is_discrete(): + self.checkbox.SetValue(True) + self.update_ui() + + def set_param_value(self, pose: List[float]): + if len(self.param_groups) == 0: + return + selected_morph_index = self.choice.GetSelection() + param_group = self.param_groups[selected_morph_index] + param_index = param_group.get_parameter_index() + if param_group.is_discrete(): + if self.checkbox.GetValue(): + for i in range(param_group.get_arity()): + pose[param_index + i] = 1.0 + else: + param_range = param_group.get_range() + alpha = (self.left_slider.GetValue() + 1000) / 2000.0 + pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha + if param_group.get_arity() == 2: + alpha = (self.right_slider.GetValue() + 1000) / 2000.0 + pose[param_index + 1] = param_range[0] + (param_range[1] - param_range[0]) * alpha + + +class SimpleParamGroupsControlPanel(wx.Panel): + def __init__(self, parent, + pose_param_category: PoseParameterCategory, + param_groups: List[PoseParameterGroup]): + super().__init__(parent, style=wx.SIMPLE_BORDER) + self.sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.sizer) + self.SetAutoLayout(1) + + self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category] + for param_group in self.param_groups: + assert not param_group.is_discrete() + assert param_group.get_arity() == 1 + + self.sliders = [] + for param_group in self.param_groups: + static_text = wx.StaticText( + self, + label=" ------------ %s ------------ " % param_group.get_group_name(), style=wx.ALIGN_CENTER) + self.sizer.Add(static_text, 0, wx.EXPAND) + range = param_group.get_range() + min_value = int(range[0] * 1000) + max_value = int(range[1] * 1000) + slider = wx.Slider(self, minValue=min_value, maxValue=max_value, value=0, style=wx.HORIZONTAL) + self.sizer.Add(slider, 0, wx.EXPAND) + self.sliders.append(slider) + + self.sizer.Fit(self) + + def set_param_value(self, pose: List[float]): + if len(self.param_groups) == 0: + return + for param_group_index in range(len(self.param_groups)): + param_group = self.param_groups[param_group_index] + slider = self.sliders[param_group_index] + param_range = param_group.get_range() + param_index = param_group.get_parameter_index() + alpha = (slider.GetValue() - slider.GetMin()) * 1.0 / (slider.GetMax() - slider.GetMin()) + pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha + + +class MainFrame(wx.Frame): + IMAGE_SIZE = 512 + OUTPUT_LENGTH = 6 + NUM_PARAMETERS = 45 + + def __init__(self, device: torch.device): + super().__init__(None, wx.ID_ANY, "Poser") + self.poser = None + self.device = device + + self.wx_source_image = None + self.torch_source_image = None + + self.main_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.SetSizer(self.main_sizer) + self.SetAutoLayout(1) + self.init_left_panel() + self.init_control_panel() + self.init_right_panel() + self.main_sizer.Fit(self) + + self.timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_images, self.timer) + + save_image_id = wx.NewIdRef() + self.Bind(wx.EVT_MENU, self.on_save_image, id=save_image_id) + accelerator_table = wx.AcceleratorTable([ + (wx.ACCEL_CTRL, ord('S'), save_image_id) + ]) + self.SetAcceleratorTable(accelerator_table) + + self.last_pose = None + self.last_output_index = self.output_index_choice.GetSelection() + self.last_output_numpy_image = None + + self.wx_source_image = None + self.torch_source_image = None + self.source_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE) + self.result_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE) + self.source_image_dirty = True + + def init_left_panel(self): + self.control_panel = wx.Panel(self, style=wx.SIMPLE_BORDER, size=(MainFrame.IMAGE_SIZE, -1)) + self.left_panel = wx.Panel(self, style=wx.SIMPLE_BORDER) + left_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.left_panel.SetSizer(left_panel_sizer) + self.left_panel.SetAutoLayout(1) + + self.source_image_panel = wx.Panel(self.left_panel, size=(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE), + style=wx.SIMPLE_BORDER) + self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel) + self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + left_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE) + + self.load_model_button = wx.Button(self.left_panel, wx.ID_ANY, "\nLoad Model\n\n") + left_panel_sizer.Add(self.load_model_button, 1, wx.EXPAND) + self.load_model_button.Bind(wx.EVT_BUTTON, self.load_model) + + left_panel_sizer.Fit(self.left_panel) + self.main_sizer.Add(self.left_panel, 0, wx.FIXED_MINSIZE) + + def on_erase_background(self, event: wx.Event): + pass + + def init_control_panel(self): + self.control_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.control_panel.SetSizer(self.control_panel_sizer) + self.control_panel.SetMinSize(wx.Size(256, 1)) + + morph_categories = [ + PoseParameterCategory.EYEBROW, + PoseParameterCategory.EYE, + PoseParameterCategory.MOUTH, + PoseParameterCategory.IRIS_MORPH + ] + morph_category_titles = { + PoseParameterCategory.EYEBROW: " ------------ Eyebrow ------------ ", + PoseParameterCategory.EYE: " ------------ Eye ------------ ", + PoseParameterCategory.MOUTH: " ------------ Mouth ------------ ", + PoseParameterCategory.IRIS_MORPH: " ------------ Iris morphs ------------ ", + } + self.morph_control_panels = {} + param_groups = get_pose_parameters().get_pose_parameter_groups() + for category in morph_categories: + filtered_param_groups = [group for group in param_groups if group.get_category() == category] + if len(filtered_param_groups) == 0: + continue + control_panel = MorphCategoryControlPanel( + self.control_panel, + morph_category_titles[category], + category, + param_groups) + self.morph_control_panels[category] = control_panel + self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND) + + self.non_morph_control_panels = {} + non_morph_categories = [ + PoseParameterCategory.IRIS_ROTATION, + PoseParameterCategory.FACE_ROTATION, + PoseParameterCategory.BODY_ROTATION, + PoseParameterCategory.BREATHING + ] + for category in non_morph_categories: + filtered_param_groups = [group for group in param_groups if group.get_category() == category] + if len(filtered_param_groups) == 0: + continue + control_panel = SimpleParamGroupsControlPanel( + self.control_panel, + category, + param_groups) + self.non_morph_control_panels[category] = control_panel + self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND) + + self.control_panel_sizer.Fit(self.control_panel) + self.main_sizer.Add(self.control_panel, 1, wx.FIXED_MINSIZE) + + def init_right_panel(self): + self.right_panel = wx.Panel(self, style=wx.SIMPLE_BORDER) + right_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.right_panel.SetSizer(right_panel_sizer) + self.right_panel.SetAutoLayout(1) + + self.result_image_panel = wx.Panel(self.right_panel, + size=(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE), + style=wx.SIMPLE_BORDER) + self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel) + self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.output_index_choice = wx.Choice( + self.right_panel, + choices=[str(i) for i in range(MainFrame.OUTPUT_LENGTH)]) + self.output_index_choice.SetSelection(0) + right_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE) + right_panel_sizer.Add(self.output_index_choice, 0, wx.EXPAND) + + self.save_image_button = wx.Button(self.right_panel, wx.ID_ANY, "\nSave Image\n\n") + right_panel_sizer.Add(self.save_image_button, 1, wx.EXPAND) + self.save_image_button.Bind(wx.EVT_BUTTON, self.on_save_image) + + right_panel_sizer.Fit(self.right_panel) + self.main_sizer.Add(self.right_panel, 0, wx.FIXED_MINSIZE) + + def create_param_category_choice(self, param_category: PoseParameterCategory): + params = [] + for param_group in self.poser.get_pose_parameter_groups(): + if param_group.get_category() == param_category: + params.append(param_group.get_group_name()) + choice = wx.Choice(self.control_panel, choices=params) + if len(params) > 0: + choice.SetSelection(0) + return choice + + def load_model(self, event: wx.Event): + dir_name = "data/character_models" + file_dialog = wx.FileDialog(self, "Choose a model", dir_name, "", "*.yaml", wx.FD_OPEN) + if file_dialog.ShowModal() == wx.ID_OK: + character_model_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + self.character_model = CharacterModel.load(character_model_file_name) + self.torch_source_image = self.character_model.get_character_image(self.device) + pil_image = resize_PIL_image( + PIL.Image.open(self.character_model.character_image_file_name), + (MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)) + w, h = pil_image.size + self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes()) + self.poser = self.character_model.get_poser(self.device) + self.source_image_dirty = True + self.Refresh() + self.Update() + except RuntimeError as e: + message_dialog = wx.MessageDialog( + self, "Could not load character model " + character_model_file_name, "Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + + def paint_source_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap) + + def paint_result_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap) + + def draw_nothing_yet_string_to_bitmap(self, bitmap): + dc = wx.MemoryDC() + dc.SelectObject(bitmap) + + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent("Nothing yet!") + dc.DrawText("Nothing yet!", (MainFrame.IMAGE_SIZE - w) // 2, (MainFrame.IMAGE_SIZE - - h) // 2) + + del dc + + def get_current_pose(self): + current_pose = [0.0 for i in range(MainFrame.NUM_PARAMETERS)] + for morph_control_panel in self.morph_control_panels.values(): + morph_control_panel.set_param_value(current_pose) + for rotation_control_panel in self.non_morph_control_panels.values(): + rotation_control_panel.set_param_value(current_pose) + return current_pose + + def update_images(self, event: wx.Event): + current_pose = self.get_current_pose() + if not self.source_image_dirty \ + and self.last_pose is not None \ + and self.last_pose == current_pose \ + and self.last_output_index == self.output_index_choice.GetSelection(): + return + self.last_pose = current_pose + self.last_output_index = self.output_index_choice.GetSelection() + + if self.torch_source_image is None or self.poser is None: + self.draw_nothing_yet_string_to_bitmap(self.source_image_bitmap) + self.draw_nothing_yet_string_to_bitmap(self.result_image_bitmap) + self.source_image_dirty = False + self.Refresh() + self.Update() + return + + if self.source_image_dirty: + dc = wx.MemoryDC() + dc.SelectObject(self.source_image_bitmap) + dc.Clear() + dc.DrawBitmap(self.wx_source_image, 0, 0) + self.source_image_dirty = False + + pose = torch.tensor(current_pose, device=self.device) + output_index = self.output_index_choice.GetSelection() + with torch.no_grad(): + start_cuda_event = torch.cuda.Event(enable_timing=True) + end_cuda_event = torch.cuda.Event(enable_timing=True) + start_cuda_event.record() + start_time = time.time() + + output_image = self.poser.pose(self.torch_source_image, pose, output_index)[0].detach().cpu() + + end_time = time.time() + end_cuda_event.record() + torch.cuda.synchronize() + print("cuda time (ms):", start_cuda_event.elapsed_time(end_cuda_event)) + print("elapsed time (ms):", (end_time - start_time) * 1000.0) + + numpy_image = convert_output_image_from_torch_to_numpy(output_image) + self.last_output_numpy_image = numpy_image + wx_image = wx.ImageFromBuffer( + numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + dc.Clear() + dc.DrawBitmap(wx_bitmap, + (MainFrame.IMAGE_SIZE - numpy_image.shape[0]) // 2, + (MainFrame.IMAGE_SIZE - numpy_image.shape[1]) // 2, + True) + del dc + + self.Refresh() + self.Update() + + def on_save_image(self, event: wx.Event): + if self.last_output_numpy_image is None: + logging.info("There is no output image to save!!!") + return + + dir_name = "data/images" + file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_SAVE) + if file_dialog.ShowModal() == wx.ID_OK: + image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + if os.path.exists(image_file_name): + message_dialog = wx.MessageDialog(self, f"Override {image_file_name}", "Manual Poser", + wx.YES_NO | wx.ICON_QUESTION) + result = message_dialog.ShowModal() + if result == wx.ID_YES: + self.save_last_numpy_image(image_file_name) + message_dialog.Destroy() + else: + self.save_last_numpy_image(image_file_name) + except: + message_dialog = wx.MessageDialog(self, f"Could not save {image_file_name}", "Manual Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + + def save_last_numpy_image(self, image_file_name): + numpy_image = self.last_output_numpy_image + pil_image = PIL.Image.fromarray(numpy_image, mode='RGBA') + os.makedirs(os.path.dirname(image_file_name), exist_ok=True) + pil_image.save(image_file_name) + + +if __name__ == "__main__": + device = torch.device('cuda:0') + app = wx.App() + main_frame = MainFrame(device) + main_frame.Show(True) + main_frame.timer.Start(16) + app.MainLoop() diff --git a/tha4/src/tha4/app/character_model_mediapipe_puppeteer.py b/tha4/src/tha4/app/character_model_mediapipe_puppeteer.py new file mode 100644 index 0000000..e030d3a --- /dev/null +++ b/tha4/src/tha4/app/character_model_mediapipe_puppeteer.py @@ -0,0 +1,431 @@ +import os +import sys +import threading +import time +from typing import Optional +import PIL.Image + +import cv2 +import mediapipe +from scipy.spatial.transform import Rotation + +from tha4.shion.base.image_util import resize_PIL_image +from tha4.charmodel.character_model import CharacterModel +from tha4.image_util import convert_linear_to_srgb +from tha4.mocap.mediapipe_constants import HEAD_ROTATIONS, HEAD_X, HEAD_Y, HEAD_Z +from tha4.mocap.mediapipe_face_pose import MediaPipeFacePose +from tha4.mocap.mediapipe_face_pose_converter_00 import MediaPoseFacePoseConverter00 + +sys.path.append(os.getcwd()) + +import torch +import wx + + +class FpsStatistics: + def __init__(self): + self.count = 100 + self.fps = [] + + def add_fps(self, fps): + self.fps.append(fps) + while len(self.fps) > self.count: + del self.fps[0] + + def get_average_fps(self): + if len(self.fps) == 0: + return 0.0 + else: + return sum(self.fps) / len(self.fps) + + +class MainFrame(wx.Frame): + IMAGE_SIZE = 512 + + def __init__(self, + pose_converter: MediaPoseFacePoseConverter00, + video_capture, + face_landmarker, + device: torch.device): + super().__init__(None, wx.ID_ANY, "THA4 Character Model MediaPipe Puppeteer") + self.face_landmarker = face_landmarker + self.video_capture = video_capture + self.pose_converter = pose_converter + self.device = device + + self.source_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE) + self.result_image_bitmap = wx.Bitmap(MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE) + self.webcam_capture_bitmap = wx.Bitmap(256, 192) + self.wx_source_image = None + self.torch_source_image = None + self.last_pose = None + self.mediapipe_face_pose = None + self.fps_statistics = FpsStatistics() + self.last_update_time = None + self.character_model = None + self.poser = None + + self.create_ui() + self.create_timers() + self.Bind(wx.EVT_CLOSE, self.on_close) + + self.update_source_image_bitmap() + self.update_result_image_bitmap() + + def create_timers(self): + self.capture_timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_capture_panel, id=self.capture_timer.GetId()) + self.animation_timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_result_image_bitmap, id=self.animation_timer.GetId()) + + def on_close(self, event: wx.Event): + # Stop the timers + self.animation_timer.Stop() + self.capture_timer.Stop() + + # Destroy the windows + self.Destroy() + event.Skip() + + def on_erase_background(self, event: wx.Event): + pass + + def create_animation_panel(self, parent): + self.animation_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.animation_panel_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.animation_panel.SetSizer(self.animation_panel_sizer) + self.animation_panel.SetAutoLayout(1) + + image_size = MainFrame.IMAGE_SIZE + + if True: + self.input_panel = wx.Panel(self.animation_panel, size=(image_size, image_size + 128), + style=wx.SIMPLE_BORDER) + self.input_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.input_panel.SetSizer(self.input_panel_sizer) + self.input_panel.SetAutoLayout(1) + self.animation_panel_sizer.Add(self.input_panel, 0, wx.FIXED_MINSIZE) + + self.source_image_panel = wx.Panel(self.input_panel, size=(image_size, image_size), style=wx.SIMPLE_BORDER) + self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel) + self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.input_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE) + + self.load_model_button = wx.Button(self.input_panel, wx.ID_ANY, "Load Model") + self.input_panel_sizer.Add(self.load_model_button, 1, wx.EXPAND) + self.load_model_button.Bind(wx.EVT_BUTTON, self.load_model) + + self.input_panel_sizer.Fit(self.input_panel) + + if True: + def current_pose_supplier() -> Optional[MediaPipeFacePose]: + return self.mediapipe_face_pose + + self.pose_converter.init_pose_converter_panel(self.animation_panel, current_pose_supplier) + + if True: + self.animation_left_panel = wx.Panel(self.animation_panel, style=wx.SIMPLE_BORDER) + self.animation_left_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.animation_left_panel.SetSizer(self.animation_left_panel_sizer) + self.animation_left_panel.SetAutoLayout(1) + self.animation_panel_sizer.Add(self.animation_left_panel, 0, wx.EXPAND) + + self.result_image_panel = wx.Panel(self.animation_left_panel, size=(image_size, image_size), + style=wx.SIMPLE_BORDER) + self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel) + self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.animation_left_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE) + + separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5)) + self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND) + + background_text = wx.StaticText(self.animation_left_panel, label="--- Background ---", + style=wx.ALIGN_CENTER) + self.animation_left_panel_sizer.Add(background_text, 0, wx.EXPAND) + + self.output_background_choice = wx.Choice( + self.animation_left_panel, + choices=[ + "TRANSPARENT", + "GREEN", + "BLUE", + "BLACK", + "WHITE" + ]) + self.output_background_choice.SetSelection(0) + self.animation_left_panel_sizer.Add(self.output_background_choice, 0, wx.EXPAND) + + separator = wx.StaticLine(self.animation_left_panel, -1, size=(256, 5)) + self.animation_left_panel_sizer.Add(separator, 0, wx.EXPAND) + + self.fps_text = wx.StaticText(self.animation_left_panel, label="") + self.animation_left_panel_sizer.Add(self.fps_text, wx.SizerFlags().Border()) + + self.animation_left_panel_sizer.Fit(self.animation_left_panel) + + self.animation_panel_sizer.Fit(self.animation_panel) + + def create_ui(self): + self.main_sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.main_sizer) + self.SetAutoLayout(1) + + self.capture_pose_lock = threading.Lock() + + self.create_animation_panel(self) + self.main_sizer.Add(self.animation_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.create_capture_panel(self) + self.main_sizer.Add(self.capture_panel, wx.SizerFlags(0).Expand().Border(wx.ALL, 5)) + + self.main_sizer.Fit(self) + + def create_capture_panel(self, parent): + self.capture_panel = wx.Panel(parent, style=wx.RAISED_BORDER) + self.capture_panel_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.capture_panel.SetSizer(self.capture_panel_sizer) + self.capture_panel.SetAutoLayout(1) + + self.webcam_capture_panel = wx.Panel(self.capture_panel, size=(256, 192), style=wx.SIMPLE_BORDER) + self.webcam_capture_panel.Bind(wx.EVT_PAINT, self.paint_webcam_capture_panel) + self.webcam_capture_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.capture_panel_sizer.Add(self.webcam_capture_panel, wx.SizerFlags(0).FixedMinSize().Border(wx.ALL, 5)) + + self.rotation_labels = {} + self.rotation_value_labels = {} + rotation_column = self.create_rotation_column(self.capture_panel, HEAD_ROTATIONS) + self.capture_panel_sizer.Add(rotation_column, wx.SizerFlags(0).Expand().Border(wx.ALL, 3)) + + def paint_webcam_capture_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.webcam_capture_panel, self.webcam_capture_bitmap) + + def create_rotation_column(self, parent, rotation_names): + column_panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + column_panel_sizer = wx.FlexGridSizer(cols=2) + column_panel_sizer.AddGrowableCol(1) + column_panel.SetSizer(column_panel_sizer) + column_panel.SetAutoLayout(1) + + for rotation_name in rotation_names: + self.rotation_labels[rotation_name] = wx.StaticText( + column_panel, label=rotation_name, style=wx.ALIGN_RIGHT) + column_panel_sizer.Add(self.rotation_labels[rotation_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + self.rotation_value_labels[rotation_name] = wx.TextCtrl( + column_panel, style=wx.TE_RIGHT) + self.rotation_value_labels[rotation_name].SetValue("0.00") + self.rotation_value_labels[rotation_name].Disable() + column_panel_sizer.Add(self.rotation_value_labels[rotation_name], + wx.SizerFlags(1).Expand().Border(wx.ALL, 3)) + + column_panel.GetSizer().Fit(column_panel) + return column_panel + + def update_capture_panel(self, event: wx.Event): + there_is_frame, frame = self.video_capture.read() + if not there_is_frame: + dc = wx.MemoryDC() + dc.SelectObject(self.webcam_capture_bitmap) + self.draw_nothing_yet_string(dc) + del dc + return + + rgb_frame = cv2.flip(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB), 1) + # Use INTER_AREA for downscaling (better quality than INTER_LINEAR) + resized_frame = cv2.resize( + rgb_frame, (256, 192), + interpolation=cv2.INTER_AREA + ) + wx_image = wx.ImageFromBuffer(256, 192, resized_frame.tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.MemoryDC() + dc.SelectObject(self.webcam_capture_bitmap) + dc.Clear() + dc.DrawBitmap(wx_bitmap, 0, 0, True) + del dc + + self.webcam_capture_panel.Refresh() + + time_ms = int(time.time() * 1000) + mediapipe_image = mediapipe.Image(image_format=mediapipe.ImageFormat.SRGB, data=rgb_frame) + detection_result = self.face_landmarker.detect_for_video(mediapipe_image, time_ms) + self.update_mediapipe_face_pose(detection_result) + + def update_mediapipe_face_pose(self, detection_result): + if len(detection_result.facial_transformation_matrixes) == 0: + return + + xform_matrix = detection_result.facial_transformation_matrixes[0] + blendshape_params = {} + for item in detection_result.face_blendshapes[0]: + blendshape_params[item.category_name] = item.score + M = xform_matrix[0:3, 0:3] + rot = Rotation.from_matrix(M) + euler_angles = rot.as_euler('xyz', degrees=True) + + self.rotation_value_labels[HEAD_X].SetValue("%0.2f" % euler_angles[0]) + self.rotation_value_labels[HEAD_X].Refresh() + self.rotation_value_labels[HEAD_Y].SetValue("%0.2f" % euler_angles[1]) + self.rotation_value_labels[HEAD_Y].Refresh() + self.rotation_value_labels[HEAD_Z].SetValue("%0.2f" % euler_angles[2]) + self.rotation_value_labels[HEAD_Z].Refresh() + + self.mediapipe_face_pose = MediaPipeFacePose(blendshape_params, xform_matrix) + + @staticmethod + def convert_to_100(x): + return int(max(0.0, min(1.0, x)) * 100) + + def paint_source_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap) + + def update_source_image_bitmap(self): + dc = wx.MemoryDC() + dc.SelectObject(self.source_image_bitmap) + if self.wx_source_image is None: + self.draw_nothing_yet_string(dc) + else: + dc.Clear() + dc.DrawBitmap(self.wx_source_image, 0, 0, True) + del dc + + def draw_nothing_yet_string(self, dc): + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent("Nothing yet!") + dc.DrawText("Nothing yet!", (MainFrame.IMAGE_SIZE - w) // 2, (MainFrame.IMAGE_SIZE - h) // 2) + + def paint_result_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap) + + def update_result_image_bitmap(self, event: Optional[wx.Event] = None): + if self.mediapipe_face_pose is None or self.poser is None: + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + self.draw_nothing_yet_string(dc) + del dc + return + + current_pose = self.pose_converter.convert(self.mediapipe_face_pose) + if self.last_pose is not None and self.last_pose == current_pose: + return + self.last_pose = current_pose + + if self.torch_source_image is None: + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + self.draw_nothing_yet_string(dc) + del dc + return + + pose = torch.tensor(current_pose, device=self.device, dtype=self.poser.get_dtype()) + + with torch.no_grad(): + output_image = self.poser.pose(self.torch_source_image, pose)[0].float() + output_image = torch.clip((output_image + 1.0) / 2.0, 0.0, 1.0) + output_image = convert_linear_to_srgb(output_image) + + background_choice = self.output_background_choice.GetSelection() + if background_choice == 0: + pass + else: + background = torch.zeros(4, output_image.shape[1], output_image.shape[2], device=self.device) + background[3, :, :] = 1.0 + if background_choice == 1: + background[1, :, :] = 1.0 + output_image = self.blend_with_background(output_image, background) + elif background_choice == 2: + background[2, :, :] = 1.0 + output_image = self.blend_with_background(output_image, background) + elif background_choice == 3: + output_image = self.blend_with_background(output_image, background) + else: + background[0:3, :, :] = 1.0 + output_image = self.blend_with_background(output_image, background) + + c, h, w = output_image.shape + output_image = 255.0 * torch.transpose(output_image.reshape(c, h * w), 0, 1).reshape(h, w, c) + output_image = output_image.byte() + + numpy_image = output_image.detach().cpu().numpy() + wx_image = wx.ImageFromBuffer(numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + dc.Clear() + dc.DrawBitmap(wx_bitmap, + (MainFrame.IMAGE_SIZE - numpy_image.shape[0]) // 2, + (MainFrame.IMAGE_SIZE - numpy_image.shape[1]) // 2, True) + del dc + + time_now = time.time_ns() + if self.last_update_time is not None: + elapsed_time = time_now - self.last_update_time + fps = 1.0 / (elapsed_time / 10 ** 9) + if self.torch_source_image is not None: + self.fps_statistics.add_fps(fps) + self.fps_text.SetLabelText("FPS = %0.2f" % self.fps_statistics.get_average_fps()) + self.last_update_time = time_now + + self.Refresh() + + def blend_with_background(self, numpy_image, background): + alpha = numpy_image[3:4, :, :] + color = numpy_image[0:3, :, :] + new_color = color * alpha + (1.0 - alpha) * background[0:3, :, :] + return torch.cat([new_color, background[3:4, :, :]], dim=0) + + def load_model(self, event: wx.Event): + dir_name = "data/character_models" + file_dialog = wx.FileDialog(self, "Choose a model", dir_name, "", "*.yaml", wx.FD_OPEN) + if file_dialog.ShowModal() == wx.ID_OK: + character_model_json_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + self.character_model = CharacterModel.load(character_model_json_file_name) + self.torch_source_image = self.character_model.get_character_image(self.device) + pil_image = resize_PIL_image( + PIL.Image.open(self.character_model.character_image_file_name), + (MainFrame.IMAGE_SIZE, MainFrame.IMAGE_SIZE)) + w, h = pil_image.size + self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes()) + self.update_source_image_bitmap() + self.poser = self.character_model.get_poser(self.device) + except Exception: + message_dialog = wx.MessageDialog( + self, "Could not load character model " + character_model_json_file_name, "Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + self.Refresh() + + +if __name__ == "__main__": + device = torch.device("cuda:0") + + pose_converter = MediaPoseFacePoseConverter00() + + face_landmarker_base_options = mediapipe.tasks.BaseOptions( + model_asset_path='data/thirdparty/mediapipe/face_landmarker_v2_with_blendshapes.task') + options = mediapipe.tasks.vision.FaceLandmarkerOptions( + base_options=face_landmarker_base_options, + running_mode=mediapipe.tasks.vision.RunningMode.VIDEO, + output_face_blendshapes=True, + output_facial_transformation_matrixes=True, + num_faces=1) + face_landmarker = mediapipe.tasks.vision.FaceLandmarker.create_from_options(options) + + video_capture = cv2.VideoCapture(0) + + app = wx.App() + main_frame = MainFrame(pose_converter, video_capture, face_landmarker, device) + main_frame.Show(True) + main_frame.capture_timer.Start(30) + main_frame.animation_timer.Start(30) + app.MainLoop() diff --git a/tha4/src/tha4/app/distill.py b/tha4/src/tha4/app/distill.py new file mode 100644 index 0000000..8a81842 --- /dev/null +++ b/tha4/src/tha4/app/distill.py @@ -0,0 +1,25 @@ +import argparse +import logging + +from tha4.distiller.distiller_config import DistillerConfig +from tha4.pytasuku.workspace import Workspace + + +def run_config(config_file_name: str): + config = DistillerConfig.load(config_file_name) + + logging.basicConfig(level=logging.INFO, force=True) + workspace = Workspace() + config.define_tasks(workspace) + + workspace.start_session() + workspace.run(f"{config.prefix}/all") + workspace.end_session() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Training script.') + parser.add_argument("--config_file", type=str, required=True, + help="The name of the config file for the distillation process.") + args = parser.parse_args() + run_config(args.config_file) diff --git a/tha4/src/tha4/app/distiller_ui.py b/tha4/src/tha4/app/distiller_ui.py new file mode 100644 index 0000000..a9d991d --- /dev/null +++ b/tha4/src/tha4/app/distiller_ui.py @@ -0,0 +1,13 @@ +from tha4.app.distill import run_config +from tha4.distiller.ui.distiller_ui_main_frame import DistillerUiMainFrame + +import wx + +if __name__ == "__main__": + app = wx.App() + main_frame = DistillerUiMainFrame() + main_frame.Show(True) + app.MainLoop() + + if main_frame.config_file_to_run is not None: + run_config(main_frame.config_file_to_run) diff --git a/tha4/src/tha4/app/full_manual_poser.py b/tha4/src/tha4/app/full_manual_poser.py new file mode 100644 index 0000000..79d225c --- /dev/null +++ b/tha4/src/tha4/app/full_manual_poser.py @@ -0,0 +1,468 @@ +import logging +import os +import sys +import time +from typing import List + +from tha4.shion.base.image_util import extract_pytorch_image_from_PIL_image, pytorch_rgba_to_numpy_image, \ + pytorch_rgb_to_numpy_image +from tha4.image_util import grid_change_to_numpy_image, resize_PIL_image + +sys.path.append(os.getcwd()) + +import PIL.Image +import numpy +import torch +import wx + +from tha4.poser.poser import Poser, PoseParameterCategory, PoseParameterGroup + + +class MorphCategoryControlPanel(wx.Panel): + def __init__(self, + parent, + title: str, + pose_param_category: PoseParameterCategory, + param_groups: List[PoseParameterGroup]): + super().__init__(parent, style=wx.SIMPLE_BORDER) + self.pose_param_category = pose_param_category + self.sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.sizer) + self.SetAutoLayout(1) + + title_text = wx.StaticText(self, label=title, style=wx.ALIGN_CENTER) + self.sizer.Add(title_text, 0, wx.EXPAND) + + self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category] + self.choice = wx.Choice(self, choices=[group.get_group_name() for group in self.param_groups]) + if len(self.param_groups) > 0: + self.choice.SetSelection(0) + self.choice.Bind(wx.EVT_CHOICE, self.on_choice_updated) + self.sizer.Add(self.choice, 0, wx.EXPAND) + + self.left_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL) + self.sizer.Add(self.left_slider, 0, wx.EXPAND) + + self.right_slider = wx.Slider(self, minValue=-1000, maxValue=1000, value=-1000, style=wx.HORIZONTAL) + self.sizer.Add(self.right_slider, 0, wx.EXPAND) + + self.checkbox = wx.CheckBox(self, label="Show") + self.checkbox.SetValue(True) + self.sizer.Add(self.checkbox, 0, wx.SHAPED | wx.ALIGN_CENTER) + + self.update_ui() + + self.sizer.Fit(self) + + def update_ui(self): + param_group = self.param_groups[self.choice.GetSelection()] + if param_group.is_discrete(): + self.left_slider.Enable(False) + self.right_slider.Enable(False) + self.checkbox.Enable(True) + elif param_group.get_arity() == 1: + self.left_slider.Enable(True) + self.right_slider.Enable(False) + self.checkbox.Enable(False) + else: + self.left_slider.Enable(True) + self.right_slider.Enable(True) + self.checkbox.Enable(False) + + def on_choice_updated(self, event: wx.Event): + param_group = self.param_groups[self.choice.GetSelection()] + if param_group.is_discrete(): + self.checkbox.SetValue(True) + self.update_ui() + + def set_param_value(self, pose: List[float]): + if len(self.param_groups) == 0: + return + selected_morph_index = self.choice.GetSelection() + param_group = self.param_groups[selected_morph_index] + param_index = param_group.get_parameter_index() + if param_group.is_discrete(): + if self.checkbox.GetValue(): + for i in range(param_group.get_arity()): + pose[param_index + i] = 1.0 + else: + param_range = param_group.get_range() + alpha = (self.left_slider.GetValue() + 1000) / 2000.0 + pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha + if param_group.get_arity() == 2: + alpha = (self.right_slider.GetValue() + 1000) / 2000.0 + pose[param_index + 1] = param_range[0] + (param_range[1] - param_range[0]) * alpha + + +class SimpleParamGroupsControlPanel(wx.Panel): + def __init__(self, parent, + pose_param_category: PoseParameterCategory, + param_groups: List[PoseParameterGroup]): + super().__init__(parent, style=wx.SIMPLE_BORDER) + self.sizer = wx.BoxSizer(wx.VERTICAL) + self.SetSizer(self.sizer) + self.SetAutoLayout(1) + + self.param_groups = [group for group in param_groups if group.get_category() == pose_param_category] + for param_group in self.param_groups: + assert not param_group.is_discrete() + assert param_group.get_arity() == 1 + + self.sliders = [] + for param_group in self.param_groups: + static_text = wx.StaticText( + self, + label=" ------------ %s ------------ " % param_group.get_group_name(), style=wx.ALIGN_CENTER) + self.sizer.Add(static_text, 0, wx.EXPAND) + range = param_group.get_range() + min_value = int(range[0] * 1000) + max_value = int(range[1] * 1000) + slider = wx.Slider(self, minValue=min_value, maxValue=max_value, value=0, style=wx.HORIZONTAL) + self.sizer.Add(slider, 0, wx.EXPAND) + self.sliders.append(slider) + + self.sizer.Fit(self) + + def set_param_value(self, pose: List[float]): + if len(self.param_groups) == 0: + return + for param_group_index in range(len(self.param_groups)): + param_group = self.param_groups[param_group_index] + slider = self.sliders[param_group_index] + param_range = param_group.get_range() + param_index = param_group.get_parameter_index() + alpha = (slider.GetValue() - slider.GetMin()) * 1.0 / (slider.GetMax() - slider.GetMin()) + pose[param_index] = param_range[0] + (param_range[1] - param_range[0]) * alpha + + +def convert_output_image_from_torch_to_numpy(output_image): + if output_image.shape[2] == 2: + h, w, c = output_image.shape + numpy_image = torch.transpose(output_image.reshape(h * w, c), 0, 1).reshape(c, h, w) + elif output_image.shape[0] == 4: + numpy_image = pytorch_rgba_to_numpy_image(output_image) + elif output_image.shape[0] == 3: + numpy_image = pytorch_rgb_to_numpy_image(output_image) + elif output_image.shape[0] == 1: + c, h, w = output_image.shape + alpha_image = torch.cat([output_image.repeat(3, 1, 1) * 2.0 - 1.0, torch.ones(1, h, w)], dim=0) + numpy_image = pytorch_rgba_to_numpy_image(alpha_image) + elif output_image.shape[0] == 2: + numpy_image = grid_change_to_numpy_image(output_image, num_channels=4) + else: + raise RuntimeError("Unsupported # image channels: %d" % output_image.shape[0]) + numpy_image = numpy.uint8(numpy.rint(numpy_image * 255.0)) + return numpy_image + + +class MainFrame(wx.Frame): + def __init__(self, poser: Poser, device: torch.device): + super().__init__(None, wx.ID_ANY, "Poser") + self.poser = poser + self.dtype = self.poser.get_dtype() + self.device = device + self.image_size = self.poser.get_image_size() + + self.wx_source_image = None + self.torch_source_image = None + + self.main_sizer = wx.BoxSizer(wx.HORIZONTAL) + self.SetSizer(self.main_sizer) + self.SetAutoLayout(1) + self.init_left_panel() + self.init_control_panel() + self.init_right_panel() + self.main_sizer.Fit(self) + + self.timer = wx.Timer(self, wx.ID_ANY) + self.Bind(wx.EVT_TIMER, self.update_images, self.timer) + + save_image_id = wx.NewIdRef() + self.Bind(wx.EVT_MENU, self.on_save_image, id=save_image_id) + accelerator_table = wx.AcceleratorTable([ + (wx.ACCEL_CTRL, ord('S'), save_image_id) + ]) + self.SetAcceleratorTable(accelerator_table) + + self.last_pose = None + self.last_output_index = self.output_index_choice.GetSelection() + self.last_output_numpy_image = None + + self.wx_source_image = None + self.torch_source_image = None + self.source_image_bitmap = wx.Bitmap(self.image_size, self.image_size) + self.result_image_bitmap = wx.Bitmap(self.image_size, self.image_size) + self.source_image_dirty = True + + def init_left_panel(self): + self.control_panel = wx.Panel(self, style=wx.SIMPLE_BORDER, size=(self.image_size, -1)) + self.left_panel = wx.Panel(self, style=wx.SIMPLE_BORDER) + left_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.left_panel.SetSizer(left_panel_sizer) + self.left_panel.SetAutoLayout(1) + + self.source_image_panel = wx.Panel(self.left_panel, size=(self.image_size, self.image_size), + style=wx.SIMPLE_BORDER) + self.source_image_panel.Bind(wx.EVT_PAINT, self.paint_source_image_panel) + self.source_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + left_panel_sizer.Add(self.source_image_panel, 0, wx.FIXED_MINSIZE) + + self.load_image_button = wx.Button(self.left_panel, wx.ID_ANY, "\nLoad Image\n\n") + left_panel_sizer.Add(self.load_image_button, 1, wx.EXPAND) + self.load_image_button.Bind(wx.EVT_BUTTON, self.load_image) + + left_panel_sizer.Fit(self.left_panel) + self.main_sizer.Add(self.left_panel, 0, wx.FIXED_MINSIZE) + + def on_erase_background(self, event: wx.Event): + pass + + def init_control_panel(self): + self.control_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.control_panel.SetSizer(self.control_panel_sizer) + self.control_panel.SetMinSize(wx.Size(256, 1)) + + morph_categories = [ + PoseParameterCategory.EYEBROW, + PoseParameterCategory.EYE, + PoseParameterCategory.MOUTH, + PoseParameterCategory.IRIS_MORPH + ] + morph_category_titles = { + PoseParameterCategory.EYEBROW: " ------------ Eyebrow ------------ ", + PoseParameterCategory.EYE: " ------------ Eye ------------ ", + PoseParameterCategory.MOUTH: " ------------ Mouth ------------ ", + PoseParameterCategory.IRIS_MORPH: " ------------ Iris morphs ------------ ", + } + self.morph_control_panels = {} + for category in morph_categories: + param_groups = self.poser.get_pose_parameter_groups() + filtered_param_groups = [group for group in param_groups if group.get_category() == category] + if len(filtered_param_groups) == 0: + continue + control_panel = MorphCategoryControlPanel( + self.control_panel, + morph_category_titles[category], + category, + self.poser.get_pose_parameter_groups()) + self.morph_control_panels[category] = control_panel + self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND) + + self.non_morph_control_panels = {} + non_morph_categories = [ + PoseParameterCategory.IRIS_ROTATION, + PoseParameterCategory.FACE_ROTATION, + PoseParameterCategory.BODY_ROTATION, + PoseParameterCategory.BREATHING + ] + for category in non_morph_categories: + param_groups = self.poser.get_pose_parameter_groups() + filtered_param_groups = [group for group in param_groups if group.get_category() == category] + if len(filtered_param_groups) == 0: + continue + control_panel = SimpleParamGroupsControlPanel( + self.control_panel, + category, + self.poser.get_pose_parameter_groups()) + self.non_morph_control_panels[category] = control_panel + self.control_panel_sizer.Add(control_panel, 0, wx.EXPAND) + + self.control_panel_sizer.Fit(self.control_panel) + self.main_sizer.Add(self.control_panel, 1, wx.FIXED_MINSIZE) + + def init_right_panel(self): + self.right_panel = wx.Panel(self, style=wx.SIMPLE_BORDER) + right_panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.right_panel.SetSizer(right_panel_sizer) + self.right_panel.SetAutoLayout(1) + + self.result_image_panel = wx.Panel(self.right_panel, + size=(self.image_size, self.image_size), + style=wx.SIMPLE_BORDER) + self.result_image_panel.Bind(wx.EVT_PAINT, self.paint_result_image_panel) + self.result_image_panel.Bind(wx.EVT_ERASE_BACKGROUND, self.on_erase_background) + self.output_index_choice = wx.Choice( + self.right_panel, + choices=[str(i) for i in range(self.poser.get_output_length())]) + self.output_index_choice.SetSelection(0) + right_panel_sizer.Add(self.result_image_panel, 0, wx.FIXED_MINSIZE) + right_panel_sizer.Add(self.output_index_choice, 0, wx.EXPAND) + + self.save_image_button = wx.Button(self.right_panel, wx.ID_ANY, "\nSave Image\n\n") + right_panel_sizer.Add(self.save_image_button, 1, wx.EXPAND) + self.save_image_button.Bind(wx.EVT_BUTTON, self.on_save_image) + + right_panel_sizer.Fit(self.right_panel) + self.main_sizer.Add(self.right_panel, 0, wx.FIXED_MINSIZE) + + def create_param_category_choice(self, param_category: PoseParameterCategory): + params = [] + for param_group in self.poser.get_pose_parameter_groups(): + if param_group.get_category() == param_category: + params.append(param_group.get_group_name()) + choice = wx.Choice(self.control_panel, choices=params) + if len(params) > 0: + choice.SetSelection(0) + return choice + + def load_image(self, event: wx.Event): + dir_name = "data/images" + file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_OPEN) + if file_dialog.ShowModal() == wx.ID_OK: + image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + pil_image = resize_PIL_image(PIL.Image.open(image_file_name), + (self.poser.get_image_size(), self.poser.get_image_size())) + w, h = pil_image.size + if pil_image.mode != 'RGBA': + self.source_image_string = "Image must have alpha channel!" + self.wx_source_image = None + self.torch_source_image = None + else: + self.wx_source_image = wx.Bitmap.FromBufferRGBA(w, h, pil_image.convert("RGBA").tobytes()) + self.torch_source_image = extract_pytorch_image_from_PIL_image(pil_image) \ + .to(self.device).to(self.dtype) + self.source_image_dirty = True + self.Refresh() + self.Update() + except: + message_dialog = wx.MessageDialog(self, "Could not load image " + image_file_name, "Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + + def paint_source_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.source_image_panel, self.source_image_bitmap) + + def paint_result_image_panel(self, event: wx.Event): + wx.BufferedPaintDC(self.result_image_panel, self.result_image_bitmap) + + def draw_nothing_yet_string_to_bitmap(self, bitmap): + dc = wx.MemoryDC() + dc.SelectObject(bitmap) + + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent("Nothing yet!") + dc.DrawText("Nothing yet!", (self.image_size - w) // 2, (self.image_size - - h) // 2) + + del dc + + def get_current_pose(self): + current_pose = [0.0 for i in range(self.poser.get_num_parameters())] + for morph_control_panel in self.morph_control_panels.values(): + morph_control_panel.set_param_value(current_pose) + for rotation_control_panel in self.non_morph_control_panels.values(): + rotation_control_panel.set_param_value(current_pose) + return current_pose + + def update_images(self, event: wx.Event): + current_pose = self.get_current_pose() + if not self.source_image_dirty \ + and self.last_pose is not None \ + and self.last_pose == current_pose \ + and self.last_output_index == self.output_index_choice.GetSelection(): + return + self.last_pose = current_pose + self.last_output_index = self.output_index_choice.GetSelection() + + if self.torch_source_image is None: + self.draw_nothing_yet_string_to_bitmap(self.source_image_bitmap) + self.draw_nothing_yet_string_to_bitmap(self.result_image_bitmap) + self.source_image_dirty = False + self.Refresh() + self.Update() + return + + if self.source_image_dirty: + dc = wx.MemoryDC() + dc.SelectObject(self.source_image_bitmap) + dc.Clear() + dc.DrawBitmap(self.wx_source_image, 0, 0) + self.source_image_dirty = False + + pose = torch.tensor(current_pose, device=self.device, dtype=self.dtype) + output_index = self.output_index_choice.GetSelection() + with torch.no_grad(): + start_cuda_event = torch.cuda.Event(enable_timing=True) + end_cuda_event = torch.cuda.Event(enable_timing=True) + start_cuda_event.record() + start_time = time.time() + + output_image = self.poser.pose(self.torch_source_image, pose, output_index)[0].detach().cpu() + + end_time = time.time() + end_cuda_event.record() + torch.cuda.synchronize() + print("cuda time (ms):", start_cuda_event.elapsed_time(end_cuda_event)) + print("elapsed time (ms):", (end_time - start_time) * 1000.0) + + numpy_image = convert_output_image_from_torch_to_numpy(output_image) + self.last_output_numpy_image = numpy_image + wx_image = wx.ImageFromBuffer( + numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + wx_bitmap = wx_image.ConvertToBitmap() + + dc = wx.MemoryDC() + dc.SelectObject(self.result_image_bitmap) + dc.Clear() + dc.DrawBitmap(wx_bitmap, + (self.image_size - numpy_image.shape[0]) // 2, + (self.image_size - numpy_image.shape[1]) // 2, + True) + del dc + + self.Refresh() + self.Update() + + def on_save_image(self, event: wx.Event): + if self.last_output_numpy_image is None: + logging.info("There is no output image to save!!!") + return + + dir_name = "data/images" + file_dialog = wx.FileDialog(self, "Choose an image", dir_name, "", "*.png", wx.FD_SAVE) + if file_dialog.ShowModal() == wx.ID_OK: + image_file_name = os.path.join(file_dialog.GetDirectory(), file_dialog.GetFilename()) + try: + if os.path.exists(image_file_name): + message_dialog = wx.MessageDialog(self, f"Override {image_file_name}", "Manual Poser", + wx.YES_NO | wx.ICON_QUESTION) + result = message_dialog.ShowModal() + if result == wx.ID_YES: + self.save_last_numpy_image(image_file_name) + message_dialog.Destroy() + else: + self.save_last_numpy_image(image_file_name) + except: + message_dialog = wx.MessageDialog(self, f"Could not save {image_file_name}", "Manual Poser", wx.OK) + message_dialog.ShowModal() + message_dialog.Destroy() + file_dialog.Destroy() + + def save_last_numpy_image(self, image_file_name): + numpy_image = self.last_output_numpy_image + pil_image = PIL.Image.fromarray(numpy_image, mode='RGBA') + os.makedirs(os.path.dirname(image_file_name), exist_ok=True) + pil_image.save(image_file_name) + + +if __name__ == "__main__": + device = torch.device('cuda:0') + try: + import tha4.poser.modes.mode_07 + + poser = tha4.poser.modes.mode_07.create_poser(device) + except RuntimeError as e: + print(e) + sys.exit() + + app = wx.App() + main_frame = MainFrame(poser, device) + main_frame.Show(True) + main_frame.timer.Start(16) + app.MainLoop() diff --git a/tha4/src/tha4/charmodel/__init__.py b/tha4/src/tha4/charmodel/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/charmodel/character_model.py b/tha4/src/tha4/charmodel/character_model.py new file mode 100644 index 0000000..15b3012 --- /dev/null +++ b/tha4/src/tha4/charmodel/character_model.py @@ -0,0 +1,69 @@ +import json +import os.path + +import PIL.Image +import torch +from omegaconf import OmegaConf + +from tha4.shion.base.image_util import extract_pytorch_image_from_PIL_image +from tha4.poser.modes.mode_14 import create_poser, KEY_FACE_MORPHER, KEY_BODY_MORPHER + + +class CharacterModel: + def __init__(self, + character_image_file_name: str, + face_morpher_file_name: str, + body_morpher_file_name: str): + self.body_morpher_file_name = body_morpher_file_name + self.face_morpher_file_name = face_morpher_file_name + self.character_image_file_name = character_image_file_name + self.poser = None + self.character_image = None + + def get_poser(self, device: torch.device): + if self.poser is not None: + self.poser.to(device) + else: + self.poser = create_poser( + device, + module_file_names={ + KEY_FACE_MORPHER: self.face_morpher_file_name, + KEY_BODY_MORPHER: self.body_morpher_file_name + }) + return self.poser + + def get_character_image(self, device: torch.device): + if self.character_image is None: + pil_image = PIL.Image.open(self.character_image_file_name) + if pil_image.mode != 'RGBA': + raise RuntimeError("Character image is not an RGBA image!") + self.character_image = extract_pytorch_image_from_PIL_image(pil_image) + self.character_image = self.character_image.to(device) + return self.character_image + + def save(self, file_name: str): + dir = os.path.dirname(file_name) + rel_char_image_file_name = os.path.relpath(self.character_image_file_name, dir) + rel_face_morpher_file_name = os.path.relpath(self.face_morpher_file_name, dir) + rel_body_morpher_file_name = os.path.relpath(self.body_morpher_file_name, dir) + data = { + "character_image_file_name": rel_char_image_file_name, + "face_morpher_file_name": rel_face_morpher_file_name, + "body_morpher_file_name": rel_body_morpher_file_name, + } + conf = OmegaConf.create(data) + os.makedirs(dir, exist_ok=True) + with open(file_name, "wt") as fout: + fout.write(OmegaConf.to_yaml(conf)) + + @staticmethod + def load(file_name: str): + conf = OmegaConf.to_container(OmegaConf.load(file_name)) + dir = os.path.dirname(file_name) + character_image_file_name = os.path.join(dir, conf["character_image_file_name"]) + face_morpher_file_name = os.path.join(dir, conf["face_morpher_file_name"]) + body_morpher_file_name = os.path.join(dir, conf["body_morpher_file_name"]) + return CharacterModel( + character_image_file_name, + face_morpher_file_name, + body_morpher_file_name) diff --git a/tha4/src/tha4/dataset/__init__.py b/tha4/src/tha4/dataset/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/dataset/image_poses_and_aother_images_dataset.py b/tha4/src/tha4/dataset/image_poses_and_aother_images_dataset.py new file mode 100644 index 0000000..9832ec3 --- /dev/null +++ b/tha4/src/tha4/dataset/image_poses_and_aother_images_dataset.py @@ -0,0 +1,35 @@ +from typing import List, Callable + +from torch import Tensor +from torch.utils.data import Dataset + + +class ImagePosesAndOtherImagesDataset(Dataset): + def __init__(self, + main_image_func: Callable[[], Tensor], + pose_dataset: Dataset, + other_image_funcs: List[Callable[[], Tensor]]): + self.main_image_func = main_image_func + self.other_image_funcs = other_image_funcs + self.pose_dataset = pose_dataset + self.main_image = None + self.other_images = [None for i in range(len(self.other_image_funcs))] + + def get_main_image(self): + if self.main_image is None: + self.main_image = self.main_image_func() + return self.main_image + + def get_other_image(self, image_index: int): + if self.other_images[image_index] is None: + self.other_images[image_index] = self.other_image_funcs[image_index]() + return self.other_images[image_index] + + def __len__(self): + return len(self.pose_dataset) + + def __getitem__(self, index): + main_image = self.get_main_image() + pose = self.pose_dataset[index][0] + other_images = [self.get_other_image(i) for i in range(len(self.other_image_funcs))] + return [main_image, pose] + other_images diff --git a/tha4/src/tha4/distiller/__init__.py b/tha4/src/tha4/distiller/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/distiller/config_based_training_tasks.py b/tha4/src/tha4/distiller/config_based_training_tasks.py new file mode 100644 index 0000000..2e13b1e --- /dev/null +++ b/tha4/src/tha4/distiller/config_based_training_tasks.py @@ -0,0 +1,98 @@ +import logging +import os +import sys +from typing import Callable, List, Optional + +from tha4.pytasuku.workspace import Workspace +from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer +from tha4.shion.core.training.distrib.distributed_training_states import DistributedTrainingState + + +def get_torchrun_executable(): + return os.path.dirname(sys.executable) + os.path.sep + "torchrun" + + +class RdzvConfig: + def __init__(self, id: int, port: int): + self.port = port + self.id = id + + +def run_standalone_config_based_training_script( + training_script_file_name: str, + config_file_name: str, + num_proc_per_node: int, + target_checkpoint_examples: Optional[int] = None, + rdzv_config: Optional[RdzvConfig] = None): + command = f"{get_torchrun_executable()} " \ + f"--nnodes=1 " \ + f"--nproc_per_node={num_proc_per_node} " + if rdzv_config is not None: + command += f"--rdzv_endpoint=localhost:{rdzv_config.port} " + command += "--rdzv_backend=c10d " + command += f"--rdzv_id={rdzv_config.id} " + else: + command += "--standalone " + command += f"{training_script_file_name} " + if target_checkpoint_examples is not None: + command += f"--target_checkpoint_examples {target_checkpoint_examples} " + command += f"--config_file={config_file_name} " + logging.info(f"Executing -- {command}") + os.system(command) + + +def define_standalone_config_based_training_tasks( + workspace: Workspace, + distributed_trainer_func: Callable[[], DistributedTrainer], + training_script_file_name: str, + config_file_name: str, + num_proc_per_node: int, + dependencies: Optional[List[str]] = None, + rdzv_config: Optional[RdzvConfig] = None): + trainer = distributed_trainer_func() + checkpoint_examples = trainer.training_protocol.get_checkpoint_examples() + assert len(checkpoint_examples) >= 1 + assert checkpoint_examples[0] > 0 + checkpoint_examples = [0] + checkpoint_examples + + if dependencies is None: + dependencies = [] + module_file_dependencies = dependencies[:] + for module_name in trainer.pretrained_module_file_names: + module_file_dependencies.append(trainer.pretrained_module_file_names[module_name]) + + def create_train_func(target_checkpoint_examples: int): + return lambda: run_standalone_config_based_training_script( + training_script_file_name, + config_file_name, + num_proc_per_node, + target_checkpoint_examples, + rdzv_config=rdzv_config) + + train_tasks = [] + for checkpoint_index in range(0, len(checkpoint_examples)): + for module_name in trainer.module_names: + module_file_name = DistributedTrainingState.get_module_file_name( + trainer.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + module_file_name, + module_file_dependencies, + create_train_func(trainer.checkpoint_examples[checkpoint_index])) + for module_name in trainer.accumulators: + accumulated_module_file_name = DistributedTrainingState.get_accumulated_module_file_name( + trainer.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + accumulated_module_file_name, + module_file_dependencies, + create_train_func(checkpoint_examples[checkpoint_index])) + workspace.create_command_task( + trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standalone", + module_file_dependencies, + create_train_func(checkpoint_examples[checkpoint_index])) + train_tasks.append(trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standlone") + workspace.create_file_task( + trainer.prefix + "/train_standalone", + module_file_dependencies, + create_train_func(checkpoint_examples[-1])) diff --git a/tha4/src/tha4/distiller/distill_body_morpher.py b/tha4/src/tha4/distiller/distill_body_morpher.py new file mode 100644 index 0000000..ae6752f --- /dev/null +++ b/tha4/src/tha4/distiller/distill_body_morpher.py @@ -0,0 +1,16 @@ +import logging + +from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer +from tha4.distiller.distiller_config import DistillerConfig + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + + parser = DistributedTrainer.get_default_arg_parser() + parser.add_argument('--config_file', type=str) + args = parser.parse_args() + + config_file_name = args.config_file + config = DistillerConfig.load(config_file_name) + + DistributedTrainer.run_with_args(config.get_body_morpher_trainer, args) diff --git a/tha4/src/tha4/distiller/distill_face_morpher.py b/tha4/src/tha4/distiller/distill_face_morpher.py new file mode 100644 index 0000000..2979fe6 --- /dev/null +++ b/tha4/src/tha4/distiller/distill_face_morpher.py @@ -0,0 +1,16 @@ +import logging + +from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer +from tha4.distiller.distiller_config import DistillerConfig + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + + parser = DistributedTrainer.get_default_arg_parser() + parser.add_argument('--config_file', type=str) + args = parser.parse_args() + + config_file_name = args.config_file + config = DistillerConfig.load(config_file_name) + + DistributedTrainer.run_with_args(config.get_face_morpher_trainer, args) diff --git a/tha4/src/tha4/distiller/distiller_config.py b/tha4/src/tha4/distiller/distiller_config.py new file mode 100644 index 0000000..08e6edc --- /dev/null +++ b/tha4/src/tha4/distiller/distiller_config.py @@ -0,0 +1,310 @@ +import os.path +import shutil +import PIL.Image +from dataclasses import dataclass +from typing import Optional + +from omegaconf import OmegaConf +from tha4.charmodel.character_model import CharacterModel +from tha4.pytasuku.workspace import Workspace, file_task +from tha4.distiller.config_based_training_tasks import define_standalone_config_based_training_tasks +from tha4.nn.siren.face_morpher.siren_face_morpher_00_trainer import SirenFaceMorpher00TrainerArgs +from tha4.nn.siren.morpher.siren_morpher_03_trainer import SirenMorpher03TrainerArgs, TrainingPhases, TrainingPhase, \ + LossWeights, LossTerm +from tha4.shion.base.image_util import pil_image_has_transparency + +POSE_DATASET_FILE_NAME = 'data/pose_dataset.pt' + + +def copy_file(source_file_name: str, dest_file_name): + os.makedirs(os.path.dirname(dest_file_name), exist_ok=True) + shutil.copyfile(source_file_name, dest_file_name) + + +@dataclass +class DistillerConfig: + prefix: str + character_image_file_name: str + face_mask_image_file_name: str + + face_morpher_random_seed_0: int = 12771885812175595441 + face_morpher_random_seed_1: int = 14367217090963479175 + face_morpher_num_training_examples_per_sample_output: Optional[int] = 10_000 + face_morpher_batch_size: int = 8 + + body_morpher_random_seed_0: int = 2892221210020292507 + body_morpher_random_seed_1: int = 9998918537095922080 + body_morpher_num_training_examples_per_sample_output: Optional[int] = 10_000 + body_morpher_batch_size: int = 8 + + num_cpu_workers: int = 1 + num_gpus: int = 1 + + def check(self): + DistillerConfig.check_prefix(self.prefix) + DistillerConfig.check_character_image_file_name(self.character_image_file_name) + DistillerConfig.check_face_mask_image_file_name(self.face_mask_image_file_name) + + DistillerConfig.check_num_cpu_workers(self.num_cpu_workers) + DistillerConfig.check_num_gpus(self.num_gpus) + + DistillerConfig.check_random_seed(self.face_morpher_random_seed_0, "face_morpher_random_seed_0") + DistillerConfig.check_random_seed(self.face_morpher_random_seed_1, "face_morpher_random_seed_1") + DistillerConfig.check_batch_size(self.face_morpher_batch_size, "face_morpher_batch_size") + DistillerConfig.check_num_training_examples_per_sample_output( + self.face_morpher_num_training_examples_per_sample_output, + "face_morpher_num_training_examples_per_sample_output") + + DistillerConfig.check_random_seed(self.body_morpher_random_seed_0, "body_morpher_random_seed_0") + DistillerConfig.check_random_seed(self.body_morpher_random_seed_1, "body_morpher_random_seed_1") + DistillerConfig.check_batch_size(self.body_morpher_batch_size, "body_morpher_batch_size") + DistillerConfig.check_num_training_examples_per_sample_output( + self.body_morpher_num_training_examples_per_sample_output, + "body_morpher_num_training_examples_per_sample_output") + + @staticmethod + def check_prefix(prefix): + assert os.path.isdir(prefix), "The 'prefix' must be a directory." + assert os.path.exists(prefix), f"The {prefix} directory does not exist." + + @staticmethod + def check_character_image_file_name(file_name): + _, ext = os.path.splitext(file_name) + assert os.path.isfile(file_name), \ + f"The specified character image file name, {file_name}, does not point to a file." + assert ext.lower() == ".png", "The character image file name must have extension '.png'." + + image = PIL.Image.open(file_name) + assert pil_image_has_transparency(image), "The character image must have an alpha channel." + assert image.width == 512 and image.height == 512, "The character image must be 512x512." + image.close() + + @staticmethod + def check_face_mask_image_file_name(file_name): + _, ext = os.path.splitext(file_name) + assert os.path.isfile(file_name), \ + f"The specified face mask image file name, {file_name}, does not point to a file." + assert ext.lower() == ".png", "The face mask image file name must have extension '.png'." + + image = PIL.Image.open(file_name) + assert image.width == 512 and image.height == 512, "The face mask image must be 512x512." + assert image.mode == "RGB", "The face mask image must be an RGB image." + for x in range(512): + for y in range(512): + r, g, b = image.getpixel((x, y)) + assert (r == 0) or (r == 255), "The R channel of the face mask image must be 0 or 255" + assert (g == 0) or (g == 255), "The G channel of the face mask image must be 0 or 255" + assert (b == 0) or (b == 255), "The B channel of the face mask image must be 0 or 255" + image.close() + + @staticmethod + def check_batch_size(value, field_name: str): + assert isinstance(value, int), f"The {field_name} must be an integer." + assert value >= 1, f"The {field_name} must be at least 1." + assert value <= 8, f"The {field_name} must be at most 8." + + @staticmethod + def check_num_cpu_workers(value): + assert value >= 1, "The value of 'num_cpu_workers must be at least 1." + + @staticmethod + def check_num_gpus(value): + assert value >= 1, "The value of 'num_gpus' must be at least 1." + + @staticmethod + def check_random_seed(value, field_name: str): + assert isinstance(value, int), f"The {field_name} must be an integer." + assert value >= 0 and value <= 0x_ffff_ffff_ffff_ffff, "A random seed must be between 0 and 2**64-1." + + @staticmethod + def check_num_training_examples_per_sample_output(value, field_name): + assert value in [10_000, 100_000, 1_000_000, + None], f"The {field_name} must be 10_000, 100_00, 1_000_000_000, or None." + + def save(self, file_name: str): + conf = OmegaConf.structured(self) + os.makedirs(self.prefix, exist_ok=True) + with open(file_name, "wt") as fout: + fout.write(OmegaConf.to_yaml(conf)) + + def config_yaml_file_name(self): + return f"{self.prefix}/config.yaml" + + def create_config_yaml_file(self): + if os.path.exists(self.config_yaml_file_name()): + return + self.save(self.config_yaml_file_name()) + + @staticmethod + def load(file_name: str) -> 'DistillerConfig': + conf = OmegaConf.to_container(OmegaConf.load(file_name)) + args = DistillerConfig(**conf) + args.check() + return args + + def face_morpher_prefix(self): + return f"{self.prefix}/face_morpher" + + def get_face_morpher_trainer(self, world_size: Optional[int] = None, backend: str = 'gloo'): + if world_size is None: + world_size = self.num_gpus + args = SirenFaceMorpher00TrainerArgs( + character_file_name=self.character_image_file_name, + face_mask_file_name=self.face_mask_image_file_name, + pose_dataset_file_name=POSE_DATASET_FILE_NAME, + total_worker=self.num_cpu_workers, + num_training_examples_per_sample_output=self.face_morpher_num_training_examples_per_sample_output, + total_batch_size=self.face_morpher_batch_size, + training_random_seed=self.face_morpher_random_seed_0, + sample_output_random_seed=self.face_morpher_random_seed_1) + return args.create_trainer(self.face_morpher_prefix(), world_size, backend) + + def body_morpher_prefix(self): + return f"{self.prefix}/body_morpher" + + def get_body_morpher_trainer(self, world_size: Optional[int] = None, backend: str = 'gloo'): + if world_size is None: + world_size = self.num_gpus + args = SirenMorpher03TrainerArgs( + character_file_name=self.character_image_file_name, + pose_dataset_file_name=POSE_DATASET_FILE_NAME, + total_worker=self.num_cpu_workers, + num_training_examples_per_sample_output=self.body_morpher_num_training_examples_per_sample_output, + training_random_seed=self.body_morpher_random_seed_0, + sample_output_random_seed=self.body_morpher_random_seed_1, + total_batch_size=self.body_morpher_batch_size, + sample_output_batch_size=1, + training_phases=TrainingPhases([ + TrainingPhase( + num_examples_upper_bound=200_000, + learning_rate=1e-4, + loss_weights=LossWeights(weights={ + LossTerm.full_blended: 0.25, + LossTerm.full_warped: 0.25, + LossTerm.full_grid_change: 0.5, + LossTerm.full_color_change: 2.0, + })), + TrainingPhase( + num_examples_upper_bound=400_000, + learning_rate=3e-5, + loss_weights=LossWeights(weights={ + LossTerm.full_blended: 0.25, + LossTerm.full_warped: 0.25, + LossTerm.full_grid_change: 0.5, + LossTerm.full_color_change: 2.0, + })), + TrainingPhase( + num_examples_upper_bound=600_000, + learning_rate=3e-5, + loss_weights=LossWeights(weights={ + LossTerm.full_blended: 1.0, + LossTerm.full_warped: 2.5, + LossTerm.full_grid_change: 5.0, + LossTerm.full_color_change: 1.0, + })), + TrainingPhase( + num_examples_upper_bound=800_000, + learning_rate=1e-5, + loss_weights=LossWeights(weights={ + LossTerm.full_blended: 1.0, + LossTerm.full_warped: 2.5, + LossTerm.full_grid_change: 5.0, + LossTerm.full_color_change: 1.0, + })), + TrainingPhase( + num_examples_upper_bound=1_300_000, + learning_rate=1e-5, + loss_weights=LossWeights(weights={ + LossTerm.full_blended: 10.0, + LossTerm.full_warped: 1.0, + LossTerm.full_grid_change: 1.0, + LossTerm.full_color_change: 1.0, + })), + TrainingPhase( + num_examples_upper_bound=1_500_000, + learning_rate=3e-6, + loss_weights=LossWeights(weights={ + LossTerm.full_blended: 10.0, + LossTerm.full_warped: 1.0, + LossTerm.full_grid_change: 1.0, + LossTerm.full_color_change: 1.0, + })), + ])) + return args.create_trainer(self.body_morpher_prefix(), world_size, backend) + + def character_model_prefix(self): + return f"{self.prefix}/character_model" + + def character_model_face_morpher_file_name(self): + return f"{self.character_model_prefix()}/face_morpher.pt" + + def character_model_body_morpher_file_name(self): + return f"{self.character_model_prefix()}/body_morpher.pt" + + def character_model_character_png_file_name(self): + return f"{self.character_model_prefix()}/character.png" + + def character_model_yaml_file_name(self): + return f"{self.character_model_prefix()}/character_model.yaml" + + def define_tasks(self, workspace: Workspace): + workspace.create_file_task(self.config_yaml_file_name(), [], self.create_config_yaml_file) + + define_standalone_config_based_training_tasks( + workspace, + self.get_face_morpher_trainer, + "src/tha4/distiller/distill_face_morpher.py", + self.config_yaml_file_name(), + num_proc_per_node=self.num_gpus, + dependencies=[ + self.config_yaml_file_name(), + ]) + + define_standalone_config_based_training_tasks( + workspace, + self.get_body_morpher_trainer, + "src/tha4/distiller/distill_body_morpher.py", + self.config_yaml_file_name(), + num_proc_per_node=self.num_gpus, + dependencies=[ + self.config_yaml_file_name(), + ]) + + @file_task(workspace, self.character_model_character_png_file_name(), [self.character_image_file_name]) + def copy_character_image_file_name(): + copy_file(self.character_image_file_name, self.character_model_character_png_file_name()) + + @file_task(workspace, self.character_model_face_morpher_file_name(), [ + f"{self.face_morpher_prefix()}/checkpoint/0010/module_module.pt", + ]) + def copy_face_morpher(): + copy_file( + f"{self.face_morpher_prefix()}/checkpoint/0010/module_module.pt", + self.character_model_face_morpher_file_name()) + + @file_task(workspace, self.character_model_body_morpher_file_name(), [ + f"{self.body_morpher_prefix()}/checkpoint/0015/module_module.pt", + ]) + def copy_face_morpher(): + copy_file( + f"{self.body_morpher_prefix()}/checkpoint/0015/module_module.pt", + self.character_model_body_morpher_file_name()) + + @file_task(workspace, self.character_model_yaml_file_name(), []) + def create_character_model_yaml_file(): + character_model = CharacterModel( + self.character_model_character_png_file_name(), + self.character_model_face_morpher_file_name(), + self.character_model_body_morpher_file_name()) + character_model.save(self.character_model_yaml_file_name()) + + workspace.create_command_task( + f"{self.prefix}/all", + [ + f"{self.face_morpher_prefix()}/train_standalone", + f"{self.body_morpher_prefix()}/train_standalone", + self.character_model_character_png_file_name(), + self.character_model_face_morpher_file_name(), + self.character_model_body_morpher_file_name(), + self.character_model_yaml_file_name(), + ]) diff --git a/tha4/src/tha4/distiller/ui/__init__.py b/tha4/src/tha4/distiller/ui/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/distiller/ui/distiller_config_state.py b/tha4/src/tha4/distiller/ui/distiller_config_state.py new file mode 100644 index 0000000..a3cc822 --- /dev/null +++ b/tha4/src/tha4/distiller/ui/distiller_config_state.py @@ -0,0 +1,142 @@ +import os.path +from contextlib import contextmanager +from pathlib import PurePath, Path +from typing import Callable, Any, Optional + +from tha4.distiller.distiller_config import DistillerConfig + + +class DistillerConfigState: + def __init__(self): + self.config = DistillerConfig(prefix="", character_image_file_name="", face_mask_image_file_name="") + self.last_saved_timestamp = None + self.dirty = False + + def load(self, file_name): + self.config = DistillerConfig.load(file_name) + if os.path.exists(self.config.config_yaml_file_name()): + self.last_saved_timestamp = os.path.getmtime(self.config.config_yaml_file_name()) + else: + self.last_saved_timestamp = None + self.dirty = False + + def need_to_check_overwrite(self): + if self.last_saved_timestamp is None: + return True + if not os.path.exists(self.config.config_yaml_file_name()): + return False + if self.last_saved_timestamp < os.path.getmtime(self.config.config_yaml_file_name()): + return True + return False + + def save(self): + self.config.save(self.config.config_yaml_file_name()) + self.dirty = False + self.last_saved_timestamp = os.path.getmtime(self.config.config_yaml_file_name()) + + @contextmanager + def updating_value(self, value_func: Callable[[], Any]): + old_value = value_func() + yield + new_value = value_func() + if new_value != old_value: + self.dirty = True + + def set_prefix(self, new_value): + with self.updating_value(lambda: self.config.prefix): + new_relative_path = self.get_relative_path_to_cwd( + new_value, + "The prefix directory must be a subdirectory of the talking-head-anime-4-demo's source code directory.") + DistillerConfig.check_prefix(new_relative_path) + self.config.prefix = new_relative_path + + def set_character_image_file_name(self, new_value): + with self.updating_value(lambda: self.config.character_image_file_name): + new_relative_path = self.get_relative_path_to_cwd( + new_value, + "The character image file must be under talking-head-anime-4-demo's source code directory.") + DistillerConfig.check_character_image_file_name(new_relative_path) + self.config.character_image_file_name = new_relative_path + + def set_face_mask_image_file_name(self, new_value): + with self.updating_value(lambda: self.config.face_mask_image_file_name): + new_relative_path = self.get_relative_path_to_cwd( + new_value, + "The face mask image file must be under talking-head-anime-4-demo's source code directory.") + DistillerConfig.check_face_mask_image_file_name(new_relative_path) + self.config.face_mask_image_file_name = new_relative_path + + def set_num_cpu_workers(self, new_value: int): + with self.updating_value(lambda: self.config.num_cpu_workers): + DistillerConfig.check_num_cpu_workers(new_value) + self.config.num_cpu_workers = new_value + + def set_num_gpus(self, new_value: int): + with self.updating_value(lambda: self.config.num_gpus): + DistillerConfig.check_num_cpu_workers(new_value) + self.config.num_gpus = new_value + + def set_face_morpher_random_seed_0(self, new_value: int): + with self.updating_value(lambda: self.config.face_morpher_random_seed_0): + DistillerConfig.check_random_seed(new_value, "face_morpher_random_seed_0") + self.config.face_morpher_random_seed_0 = new_value + + def set_face_morpher_random_seed_1(self, new_value: int): + with self.updating_value(lambda: self.config.face_morpher_random_seed_1): + DistillerConfig.check_random_seed(new_value, "face_morpher_random_seed_1") + self.config.face_morpher_random_seed_1 = new_value + + def set_face_morpher_num_training_examples_per_sample_output(self, new_value: Optional[int]): + with self.updating_value(lambda: self.config.face_morpher_num_training_examples_per_sample_output): + DistillerConfig.check_num_training_examples_per_sample_output( + new_value, "face_morpher_num_training_examples_per_sample_output") + self.config.face_morpher_num_training_examples_per_sample_output = new_value + + def set_face_morpher_batch_size(self, new_value: int): + with self.updating_value(lambda: self.config.face_morpher_batch_size): + DistillerConfig.check_batch_size(new_value, "face_morpher_batch_size") + self.config.face_morpher_batch_size = new_value + + def set_body_morpher_random_seed_0(self, new_value: int): + with self.updating_value(lambda: self.config.body_morpher_random_seed_0): + DistillerConfig.check_random_seed(new_value, "body_morpher_random_seed_0") + self.config.body_morpher_random_seed_0 = new_value + + def set_body_morpher_random_seed_1(self, new_value: int): + with self.updating_value(lambda: self.config.body_morpher_random_seed_1): + DistillerConfig.check_random_seed(new_value, "body_morpher_random_seed_1") + self.config.body_morpher_random_seed_1 = new_value + + def set_body_morpher_num_training_examples_per_sample_output(self, new_value: Optional[int]): + with self.updating_value(lambda: self.config.body_morpher_num_training_examples_per_sample_output): + DistillerConfig.check_num_training_examples_per_sample_output( + new_value, "body_morpher_num_training_examples_per_sample_output") + self.config.body_morpher_num_training_examples_per_sample_output = new_value + + def set_body_morpher_batch_size(self, new_value: int): + with self.updating_value(lambda: self.config.body_morpher_batch_size): + DistillerConfig.check_batch_size(new_value, "body_morpher_batch_size") + self.config.body_morpher_batch_size = new_value + + def get_relative_path_to_cwd(self, file_name: str, message: str): + cwd = os.getcwd() + assert os.path.commonprefix([cwd, file_name]) == cwd, message + cwd_path = Path(cwd).as_posix() + new_path = Path(file_name).as_posix() + new_relative_path = os.path.relpath(str(new_path), cwd_path) + new_relative_path = str(Path(new_relative_path).as_posix()) + return new_relative_path + + def can_show_character_image(self): + return os.path.isfile(self.config.character_image_file_name) + + def can_show_face_mask_image(self): + return os.path.isfile(self.config.face_mask_image_file_name) + + def can_show_mask_on_face_image(self): + return self.can_show_character_image() and self.can_show_face_mask_image() + + def can_save(self): + return os.path.isdir(self.config.prefix) \ + and os.path.isfile(self.config.character_image_file_name) \ + and os.path.isfile(self.config.face_mask_image_file_name) diff --git a/tha4/src/tha4/distiller/ui/distiller_ui_main_frame.py b/tha4/src/tha4/distiller/ui/distiller_ui_main_frame.py new file mode 100644 index 0000000..20bb7d3 --- /dev/null +++ b/tha4/src/tha4/distiller/ui/distiller_ui_main_frame.py @@ -0,0 +1,761 @@ +import multiprocessing +import random +from contextlib import contextmanager +from typing import Callable +import PIL.Image + +import torch +import wx +import wx.html +import wx.lib.intctrl +from tha4.distiller.ui.distiller_config_state import DistillerConfigState +from tha4.image_util import convert_output_image_from_torch_to_numpy +from tha4.shion.base.image_util import extract_pytorch_image_from_PIL_image + + +def wx_bind_event(widget, evt): + def f(handler): + widget.Bind(evt, handler) + return handler + + return f + + +class DistillerUiMainFrame(wx.Frame): + PARAM_NAME_STATIC_TEXT_MIN_WIDTH = 400 + NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES = [ + "10_000", "100_000", "1_000_000", "Do not generate sample outputs"] + + def __init__(self): + super().__init__(None, wx.ID_ANY, "Distiller UI") + + self.init_ui() + self.init_menus() + self.init_bitmaps() + self.Bind(wx.EVT_CLOSE, self.on_close) + + self.state = DistillerConfigState() + self.update_ui() + + self.config_file_to_run = None + + def init_ui(self): + main_sizer = wx.BoxSizer(wx.HORIZONTAL) + + self.SetSizer(main_sizer) + self.SetAutoLayout(1) + + left_panel = self.init_left_panel(self) + main_sizer.Add(left_panel, 0, wx.FIXED_MINSIZE) + + middle_panel = self.init_middle_panel(self) + main_sizer.Add(middle_panel, 0, wx.EXPAND) + + right_panel = self.init_right_panel(self) + main_sizer.Add(right_panel, 1, wx.EXPAND) + + main_sizer.Fit(self) + + def init_menus(self): + self.file_menu = wx.Menu() + + self.new_menu_id = wx.Window.NewControlId() + self.file_menu.Append( + self.new_menu_id, item="&New\tCTRL+N", helpString="Create a new distiller configuration.") + self.Bind(wx.EVT_MENU, self.on_new, id=self.new_menu_id) + + self.open_menu_id = wx.Window.NewControlId() + self.file_menu.Append( + self.open_menu_id, item="&Open\tCTRL+O", helpString="Open a distiller confuguration.") + self.Bind(wx.EVT_MENU, self.on_open, id=self.open_menu_id) + + self.save_menu_id = wx.Window.NewControlId() + self.save_menu_item = wx.MenuItem( + self.file_menu, id=self.save_menu_id, text="&Save\tCTRL+S", + helpString="Save the current distiller configuration. Error message will be shown it it is not well formed.") + self.Bind(wx.EVT_MENU, self.on_save, id=self.save_menu_id) + self.file_menu.Append(self.save_menu_item) + + self.file_menu.AppendSeparator() + + self.exit_menu_id = wx.ID_EXIT + self.file_menu.Append( + self.exit_menu_id, item="E&xit\tCTRL+Q", helpString="Exit the application.") + self.Bind(wx.EVT_MENU, self.on_close, id=self.exit_menu_id) + + self.menu_bar = wx.MenuBar() + self.menu_bar.Append(self.file_menu, "&File") + + self.SetMenuBar(self.menu_bar) + + def init_bitmaps(self): + self.face_image_bitmap = wx.Bitmap(128, 128) + self.face_image_pytorch = None + self.face_mask_image_bitmap = wx.Bitmap(128, 128) + self.face_mask_image_pytorch = None + self.mask_on_face_image_bitmap = wx.Bitmap(128, 128) + self.draw_nothing_yet_string_to_bitmap(self.face_image_bitmap, 128, 128) + self.draw_nothing_yet_string_to_bitmap(self.face_mask_image_bitmap, 128, 128) + self.draw_nothing_yet_string_to_bitmap(self.mask_on_face_image_bitmap, 128, 128) + + @contextmanager + def create_panel(self, parent, sizer, *args, **kwargs): + panel = wx.Panel(parent, *args, **kwargs) + panel.SetSizer(sizer) + panel.SetAutoLayout(1) + + try: + yield panel, sizer + finally: + sizer.Fit(panel) + + def init_left_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, sizer): + self.face_image_panel = wx.Panel(panel, size=(128, 128), style=wx.SIMPLE_BORDER) + self.face_image_panel.Bind(wx.EVT_PAINT, self.on_face_image_panel_paint) + sizer.Add(self.face_image_panel, 0, wx.EXPAND) + + static_text = wx.StaticText(panel, label="Face", style=wx.ALIGN_CENTER) + sizer.Add(static_text, 0, wx.EXPAND) + + self.face_mask_image_panel = wx.Panel(panel, size=(128, 128), style=wx.SIMPLE_BORDER) + self.face_mask_image_panel.Bind(wx.EVT_PAINT, self.on_face_mask_image_panel_paint) + sizer.Add(self.face_mask_image_panel, 0, wx.EXPAND) + + static_text = wx.StaticText(panel, label="Face mask", style=wx.ALIGN_CENTER) + sizer.Add(static_text, 0, wx.EXPAND) + + self.mask_on_face_image_panel = wx.Panel(panel, size=(128, 128), style=wx.SIMPLE_BORDER) + self.mask_on_face_image_panel.Bind(wx.EVT_PAINT, self.on_mask_on_face_image_panel_paint) + sizer.Add(self.mask_on_face_image_panel, 0, wx.EXPAND) + + static_text = wx.StaticText(panel, label="Mask upon face", style=wx.ALIGN_CENTER) + sizer.Add(static_text, 0, wx.EXPAND) + + return panel + + def on_erase_background(self, event): + pass + + def on_face_image_panel_paint(self, event): + wx.BufferedPaintDC(self.face_image_panel, self.face_image_bitmap) + + def on_face_mask_image_panel_paint(self, event): + wx.BufferedPaintDC(self.face_mask_image_panel, self.face_mask_image_bitmap) + + def on_mask_on_face_image_panel_paint(self, event): + wx.BufferedPaintDC(self.mask_on_face_image_panel, self.mask_on_face_image_bitmap) + + def init_middle_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, sizer): + sizer.Add(self.init_prefix_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_character_image_file_name_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_face_mask_image_file_name_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_num_cpu_workers_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_num_gpus_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_face_morpher_random_seed_0_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_face_morpher_random_seed_1_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_face_morpher_batch_size_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_body_morpher_random_seed_0_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_body_morpher_random_seed_1_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_body_morpher_batch_size_panel(panel), 0, wx.EXPAND) + sizer.Add(self.init_num_training_examples_per_sample_output_panel(panel), 0, wx.EXPAND) + + self.run_button = wx.Button(panel, label="RUN") + self.run_button.SetMinSize((-1, 64)) + self.run_button.Bind(wx.EVT_BUTTON, self.on_run) + sizer.Add(self.run_button, 1, wx.EXPAND) + + return panel + + def init_prefix_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "prefix (i.e. project directory)", + self.create_help_button_func("distiller-ui-doc/params/prefix.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) \ + as (prefix_panel, prefix_sizer): + self.prefix_text_ctrl = wx.TextCtrl(prefix_panel, value="") + self.prefix_text_ctrl.SetEditable(False) + prefix_sizer.Add(self.prefix_text_ctrl, 1, wx.EXPAND) + + self.prefix_change_button = wx.Button(prefix_panel, label="Change...") + self.prefix_change_button.Bind(wx.EVT_BUTTON, self.on_prefix_change_button) + prefix_sizer.Add(self.prefix_change_button, 0, wx.EXPAND) + panel_sizer.Add(prefix_panel, 1, wx.EXPAND) + + return panel + + def on_prefix_change_button(self, event): + dir_dialog = wx.DirDialog(self, "Choose a directory.", style=wx.DD_DEFAULT_STYLE | wx.DD_NEW_DIR_BUTTON) + if dir_dialog.ShowModal() != wx.ID_OK: + return + prefix_value = dir_dialog.GetPath() + try: + self.state.set_prefix(prefix_value) + self.update_ui() + except Exception as e: + message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + + def init_character_image_file_name_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "character_image_file_name", + self.create_help_button_func("distiller-ui-doc/params/character_image_file_name.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer): + self.character_image_file_name_text_ctrl = wx.TextCtrl(sub_panel, value="") + self.character_image_file_name_text_ctrl.SetEditable(False) + sub_sizer.Add(self.character_image_file_name_text_ctrl, 1, wx.EXPAND) + + self.character_image_change_button = wx.Button(sub_panel, label="Change...") + self.character_image_change_button.Bind(wx.EVT_BUTTON, self.on_character_image_change_button) + sub_sizer.Add(self.character_image_change_button, 0, wx.EXPAND) + panel_sizer.Add(sub_panel, 1, wx.EXPAND) + + return panel + + def on_character_image_change_button(self, event): + file_dialog = wx.FileDialog(self, "Choose a PNG file", wildcard="*.png", style=wx.FD_OPEN) + if file_dialog.ShowModal() != wx.ID_OK: + return + file_name = file_dialog.GetPath() + try: + self.state.set_character_image_file_name(file_name) + self.update_face_image_bitmap(file_name) + self.update_ui() + except Exception as e: + message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + + def update_face_image_bitmap(self, new_file_name: str): + pil_image = PIL.Image.open(new_file_name) + subimage = pil_image.crop((256 - 64, 80, 256 + 64, 208)) + self.face_image_bitmap = wx.Bitmap.FromBufferRGBA(128, 128, subimage.convert("RGBA").tobytes()) + self.face_image_pytorch = extract_pytorch_image_from_PIL_image(subimage).to(torch.float) + self.update_mask_on_face_image_bitmap() + + def init_face_mask_image_file_name_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "face_mask_image_file_name", + self.create_help_button_func("distiller-ui-doc/params/face_mask_image_file_name.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer): + self.face_mask_image_file_name_text_ctrl = wx.TextCtrl(sub_panel, value="") + self.face_mask_image_file_name_text_ctrl.SetEditable(False) + sub_sizer.Add(self.face_mask_image_file_name_text_ctrl, 1, wx.EXPAND) + + self.face_mask_image_file_name_change_button = wx.Button(sub_panel, label="Change...") + self.face_mask_image_file_name_change_button.Bind(wx.EVT_BUTTON, self.on_face_mask_image_change_button) + sub_sizer.Add(self.face_mask_image_file_name_change_button, 0, wx.EXPAND) + panel_sizer.Add(sub_panel, 1, wx.EXPAND) + + return panel + + def on_face_mask_image_change_button(self, event): + file_dialog = wx.FileDialog(self, "Choose a PNG file", wildcard="*.png", style=wx.FD_OPEN) + if file_dialog.ShowModal() != wx.ID_OK: + return + file_name = file_dialog.GetPath() + try: + self.state.set_face_mask_image_file_name(file_name) + self.update_face_mask_image_bitmap(file_name) + self.update_ui() + except Exception as e: + message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + + def update_face_mask_image_bitmap(self, new_file_name): + pil_image = PIL.Image.open(new_file_name) + subimage = pil_image.crop((256 - 64, 80, 256 + 64, 208)) + self.face_mask_image_bitmap = wx.Bitmap.FromBufferRGBA(128, 128, subimage.convert("RGBA").tobytes()) + self.face_mask_image_pytorch = extract_pytorch_image_from_PIL_image(subimage).to(torch.float) + self.face_mask_image_pytorch = self.face_mask_image_pytorch[0:1, :, :] + self.update_mask_on_face_image_bitmap() + + def update_mask_on_face_image_bitmap(self): + if self.face_image_pytorch is None: + return + if self.face_mask_image_pytorch is None: + return + + mask_on_face_image = (0.5 * self.face_image_pytorch) + (0.5 * self.face_mask_image_pytorch) + numpy_image = convert_output_image_from_torch_to_numpy(mask_on_face_image) + wx_image = wx.ImageFromBuffer( + numpy_image.shape[0], + numpy_image.shape[1], + numpy_image[:, :, 0:3].tobytes(), + numpy_image[:, :, 3].tobytes()) + self.mask_on_face_image_bitmap = wx_image.ConvertToBitmap() + + def init_num_cpu_workers_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "num_cpu_workers", + self.create_help_button_func("distiller-ui-doc/params/num_cpu_workers.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + num_cpus = multiprocessing.cpu_count() + self.num_cpu_workers_spin_ctrl = wx.SpinCtrl(panel, initial=1, min=1, max=num_cpus) + + @wx_bind_event(self.num_cpu_workers_spin_ctrl, wx.EVT_SPINCTRL) + def on_num_cpu_workers_spin_ctrl(event): + self.state.set_num_cpu_workers(self.num_cpu_workers_spin_ctrl.GetValue()) + self.Refresh() + + panel_sizer.Add(self.num_cpu_workers_spin_ctrl, 1, wx.EXPAND) + + return panel + + def init_num_gpus_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "num_gpus", + self.create_help_button_func("distiller-ui-doc/params/num_gpus.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + num_gpus = torch.cuda.device_count() + self.num_gpus_spin_ctrl = wx.SpinCtrl(panel, initial=1, min=1, max=max(1, num_gpus)) + + @wx_bind_event(self.num_gpus_spin_ctrl, wx.EVT_SPINCTRL) + def on_num_cpu_workers_spin_ctrl(event): + self.state.set_num_gpus(self.num_gpus_spin_ctrl.GetValue()) + self.Refresh() + + panel_sizer.Add(self.num_gpus_spin_ctrl, 1, wx.EXPAND) + + return panel + + def init_face_morpher_random_seed_0_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "face_morpher_random_seed_0", + self.create_help_button_func("distiller-ui-doc/params/face_morpher_random_seed_0.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer): + initial_value = random.randint(0, 2 ** 64 - 1) + self.face_morpher_random_seed_0_int_ctrl = wx.lib.intctrl.IntCtrl( + sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff) + + @wx_bind_event(self.face_morpher_random_seed_0_int_ctrl, wx.EVT_TEXT) + def on_face_morpher_random_seed_0_int_ctrl_text(event): + self.state.set_face_morpher_random_seed_0(self.face_morpher_random_seed_0_int_ctrl.GetValue()) + + sub_sizer.Add(self.face_morpher_random_seed_0_int_ctrl, 1, wx.EXPAND) + + self.face_morpher_random_seed_0_randomize_button = wx.Button(sub_panel, label="Randomize") + + @wx_bind_event(self.face_morpher_random_seed_0_randomize_button, wx.EVT_BUTTON) + def on_face_morpher_random_seed_0_randomize_button(event): + new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff) + self.face_morpher_random_seed_0_int_ctrl.SetValue(new_value) + self.state.set_face_morpher_random_seed_0(new_value) + + sub_sizer.Add(self.face_morpher_random_seed_0_randomize_button, 0, wx.EXPAND) + panel_sizer.Add(sub_panel, 1, wx.EXPAND) + + return panel + + def init_face_morpher_random_seed_1_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "face_morpher_random_seed_1", + self.create_help_button_func("distiller-ui-doc/params/face_morpher_random_seed_1.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer): + initial_value = random.randint(0, 2 ** 64 - 1) + self.face_morpher_random_seed_1_int_ctrl = wx.lib.intctrl.IntCtrl( + sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff) + + @wx_bind_event(self.face_morpher_random_seed_1_int_ctrl, wx.EVT_TEXT) + def on_face_morpher_random_seed_1_int_ctrl_text(event): + self.state.set_face_morpher_random_seed_1(self.face_morpher_random_seed_1_int_ctrl.GetValue()) + + sub_sizer.Add(self.face_morpher_random_seed_1_int_ctrl, 1, wx.EXPAND) + + self.face_morpher_random_seed_1_randomize_button = wx.Button(sub_panel, label="Randomize") + + @wx_bind_event(self.face_morpher_random_seed_1_randomize_button, wx.EVT_BUTTON) + def on_face_morpher_random_seed_1_randomize_button(event): + new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff) + self.face_morpher_random_seed_1_int_ctrl.SetValue(new_value) + self.state.set_face_morpher_random_seed_1(new_value) + + sub_sizer.Add(self.face_morpher_random_seed_1_randomize_button, 0, wx.EXPAND) + panel_sizer.Add(sub_panel, 1, wx.EXPAND) + + return panel + + def init_face_morpher_batch_size_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "face_morpher_batch_size", + self.create_help_button_func("distiller-ui-doc/params/face_morpher_batch_size.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + self.face_morpher_batch_size_spin_ctrl = wx.SpinCtrl(panel, initial=8, min=1, max=8) + + @wx_bind_event(self.face_morpher_batch_size_spin_ctrl, wx.EVT_SPINCTRL) + def on_face_morpher_batch_size_spin_ctrl(event): + self.state.set_face_morpher_batch_size(self.face_morpher_batch_size_spin_ctrl.GetValue()) + + panel_sizer.Add(self.face_morpher_batch_size_spin_ctrl, 1, wx.EXPAND) + + return panel + + def init_body_morpher_random_seed_0_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "body_morpher_random_seed_0", + self.create_help_button_func("distiller-ui-doc/params/body_morpher_random_seed_0.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer): + initial_value = random.randint(0, 2 ** 64 - 1) + self.body_morpher_random_seed_0_int_ctrl = wx.lib.intctrl.IntCtrl( + sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff) + + @wx_bind_event(self.body_morpher_random_seed_0_int_ctrl, wx.EVT_TEXT) + def on_body_morpher_random_seed_0_int_ctrl_text(event): + self.state.set_body_morpher_random_seed_0(self.body_morpher_random_seed_0_int_ctrl.GetValue()) + + sub_sizer.Add(self.body_morpher_random_seed_0_int_ctrl, 1, wx.EXPAND) + + self.body_morpher_random_seed_0_randomize_button = wx.Button(sub_panel, label="Randomize") + + @wx_bind_event(self.body_morpher_random_seed_0_randomize_button, wx.EVT_BUTTON) + def on_body_morpher_random_seed_0_randomize_button(event): + new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff) + self.body_morpher_random_seed_0_int_ctrl.SetValue(new_value) + self.state.set_body_morpher_random_seed_0(new_value) + + sub_sizer.Add(self.body_morpher_random_seed_0_randomize_button, 0, wx.EXPAND) + panel_sizer.Add(sub_panel, 1, wx.EXPAND) + + return panel + + def init_body_morpher_random_seed_1_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "body_morpher_random_seed_1", + self.create_help_button_func("distiller-ui-doc/params/body_morpher_random_seed_1.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + with self.create_panel(panel, wx.BoxSizer(wx.HORIZONTAL), style=wx.BORDER_NONE) as (sub_panel, sub_sizer): + initial_value = random.randint(0, 2 ** 64 - 1) + self.body_morpher_random_seed_1_int_ctrl = wx.lib.intctrl.IntCtrl( + sub_panel, value=initial_value, min=0, max=0x_ffff_ffff_ffff_ffff) + + @wx_bind_event(self.body_morpher_random_seed_1_int_ctrl, wx.EVT_TEXT) + def on_body_morpher_random_seed_1_int_ctrl_text(event): + self.state.set_body_morpher_random_seed_1(self.body_morpher_random_seed_1_int_ctrl.GetValue()) + + sub_sizer.Add(self.body_morpher_random_seed_1_int_ctrl, 1, wx.EXPAND) + + self.body_morpher_random_seed_1_randomize_button = wx.Button(sub_panel, label="Randomize") + + @wx_bind_event(self.body_morpher_random_seed_1_randomize_button, wx.EVT_BUTTON) + def on_body_morpher_random_seed_1_randomize_button(event): + new_value = random.randint(0, 0x_ffff_ffff_ffff_ffff) + self.body_morpher_random_seed_1_int_ctrl.SetValue(new_value) + self.state.set_body_morpher_random_seed_1(new_value) + + sub_sizer.Add(self.body_morpher_random_seed_1_randomize_button, 0, wx.EXPAND) + panel_sizer.Add(sub_panel, 1, wx.EXPAND) + + return panel + + def init_body_morpher_batch_size_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "body_morpher_batch_size", + self.create_help_button_func("distiller-ui-doc/params/body_morpher_batch_size.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + self.body_morpher_batch_size_spin_ctrl = wx.SpinCtrl(panel, initial=8, min=1, max=8) + + @wx_bind_event(self.body_morpher_batch_size_spin_ctrl, wx.EVT_SPINCTRL) + def on_body_morpher_batch_size_spin_ctrl(event): + self.state.set_body_morpher_batch_size(self.body_morpher_batch_size_spin_ctrl.GetValue()) + + panel_sizer.Add(self.body_morpher_batch_size_spin_ctrl, 1, wx.EXPAND) + + return panel + + def init_num_training_examples_per_sample_output_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, panel_sizer): + prefix_param_name_panel = self.create_param_name_panel_with_help_button( + panel, + "num_training_examples_per_sample_output", + self.create_help_button_func("distiller-ui-doc/params/num_training_examples_per_sample_output.html")) + panel_sizer.Add(prefix_param_name_panel, 1, wx.EXPAND) + + self.num_training_examples_per_sample_output_combobox = \ + wx.ComboBox(panel, + value="10_000", + choices=DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES) + + @wx_bind_event(self.num_training_examples_per_sample_output_combobox, wx.EVT_COMBOBOX) + def on_num_training_examples_per_sample_output_combobox(event): + index = self.num_training_examples_per_sample_output_combobox.GetSelection() + if index == 3: + self.state.set_face_morpher_num_training_examples_per_sample_output(None) + self.state.set_body_morpher_num_training_examples_per_sample_output(None) + else: + selected = DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES[index] + new_value = int(selected) + self.state.set_face_morpher_num_training_examples_per_sample_output(new_value) + self.state.set_body_morpher_num_training_examples_per_sample_output(new_value) + + panel_sizer.Add(self.num_training_examples_per_sample_output_combobox, 1, wx.EXPAND) + + return panel + + def on_close(self, event): + if self.state.dirty: + confirmation_dialog = wx.MessageDialog( + parent=self, + message=f"You have not saved your work. Do you want to exit anyway?", + caption="Confirmation", + style=wx.YES_NO | wx.ICON_QUESTION) + result = confirmation_dialog.ShowModal() + if result == wx.ID_NO: + return + + self.Destroy() + + def create_help_button_func(self, html_file_name: str): + def init_help_button_func(parent): + button = wx.Button(parent, label="Help") + + @wx_bind_event(button, wx.EVT_BUTTON) + def on_prefix_button(event): + self.html_window.LoadPage(html_file_name) + self.Refresh() + + return button + + return init_help_button_func + + def create_param_name_panel_with_help_button( + self, parent, param_name: str, help_button_func: Callable[[wx.Window], wx.Button]): + with self.create_panel(parent, wx.BoxSizer(wx.HORIZONTAL), style=wx.NO_BORDER) \ + as (panel, sizer): + title_text_panel = self.create_vertically_centered_text_panel( + panel, param_name, DistillerUiMainFrame.PARAM_NAME_STATIC_TEXT_MIN_WIDTH) + sizer.Add(title_text_panel, 1, wx.EXPAND) + + help_button = help_button_func(panel) + sizer.Add(help_button, 0, wx.EXPAND) + return panel + + def create_vertically_centered_text_panel(self, parent, text: str, min_width: int): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.NO_BORDER) as (panel, sizer): + sizer.AddStretchSpacer(1) + text = wx.StaticText( + panel, + label=text, + style=wx.ALIGN_CENTER) + text.SetMinSize((min_width, -1)) + sizer.Add(text, 0, wx.EXPAND) + sizer.AddStretchSpacer(1) + return panel + + def init_right_panel(self, parent): + with self.create_panel(parent, wx.BoxSizer(wx.VERTICAL), style=wx.BORDER_SIMPLE) as (panel, sizer): + self.html_window = wx.html.HtmlWindow(panel) + self.html_window.SetMinSize((600, 600)) + self.html_window.SetFonts("Times New Roman", "Courier New", sizes=[10, 12, 14, 16, 18, 20, 24]) + self.html_window.LoadPage("distiller-ui-doc/index.html") + sizer.Add(self.html_window, 1, wx.EXPAND) + + go_to_main_documentation_button = wx.Button(panel, label="Go to Main Documentation") + sizer.Add(go_to_main_documentation_button, 0, wx.EXPAND) + + @wx_bind_event(go_to_main_documentation_button, wx.EVT_BUTTON) + def on_go_to_main_documentation_button(event): + self.html_window.LoadPage("distiller-ui-doc/index.html") + self.Refresh() + + return panel + + def populate_distiller_config(self): + self.state.config.prefix = self.prefix_text_ctrl.GetValue() + self.state.config.character_image_file_name = self.character_image_file_name_text_ctrl.GetValue() + self.state.config.face_mask_image_file_name = self.face_mask_image_file_name_text_ctrl.GetValue() + + self.state.config.num_cpu_workers = self.num_cpu_workers_spin_ctrl.GetValue() + self.state.config.num_gpus = self.num_gpus_spin_ctrl.GetValue() + + self.state.config.face_morpher_random_seed_0 = self.face_morpher_random_seed_0_int_ctrl.GetValue() + self.state.config.face_morpher_random_seed_1 = self.face_morpher_random_seed_1_int_ctrl.GetValue() + self.state.config.face_morpher_batch_size = self.face_morpher_batch_size_spin_ctrl.GetValue() + + self.state.config.body_morpher_random_seed_0 = self.body_morpher_random_seed_0_int_ctrl.GetValue() + self.state.config.body_morpher_random_seed_1 = self.body_morpher_random_seed_1_int_ctrl.GetValue() + self.state.config.body_morpher_batch_size = self.body_morpher_batch_size_spin_ctrl.GetValue() + + if self.num_training_examples_per_sample_output_combobox.GetValue() == \ + DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES[-1]: + self.state.config.face_morpher_num_training_examples_per_sample_output = None + self.state.config.body_morpher_num_training_examples_per_sample_output = None + else: + value = int(self.num_training_examples_per_sample_output_combobox.GetValue()) + self.state.config.face_morpher_num_training_examples_per_sample_output = value + self.state.config.body_morpher_num_training_examples_per_sample_output = value + + def update_ui(self): + self.prefix_text_ctrl.SetValue(self.state.config.prefix) + self.character_image_file_name_text_ctrl.SetValue(self.state.config.character_image_file_name) + self.face_mask_image_file_name_text_ctrl.SetValue(self.state.config.face_mask_image_file_name) + + if not self.state.can_show_character_image(): + self.draw_nothing_yet_string_to_bitmap(self.face_image_bitmap, 128, 128) + if not self.state.can_show_face_mask_image(): + self.draw_nothing_yet_string_to_bitmap(self.face_mask_image_bitmap, 128, 128) + if not self.state.can_show_mask_on_face_image(): + self.draw_nothing_yet_string_to_bitmap(self.mask_on_face_image_bitmap, 128, 128) + + self.num_cpu_workers_spin_ctrl.SetValue(self.state.config.num_cpu_workers) + self.num_gpus_spin_ctrl.SetValue(self.state.config.num_gpus) + + self.face_morpher_random_seed_0_int_ctrl.SetValue(self.state.config.face_morpher_random_seed_0) + self.face_morpher_random_seed_1_int_ctrl.SetValue(self.state.config.face_morpher_random_seed_1) + self.face_morpher_batch_size_spin_ctrl.SetValue(self.state.config.face_morpher_batch_size) + + self.body_morpher_random_seed_0_int_ctrl.SetValue(self.state.config.body_morpher_random_seed_0) + self.body_morpher_random_seed_1_int_ctrl.SetValue(self.state.config.body_morpher_random_seed_1) + self.body_morpher_batch_size_spin_ctrl.SetValue(self.state.config.body_morpher_batch_size) + + if self.state.config.body_morpher_num_training_examples_per_sample_output is None: + self.num_training_examples_per_sample_output_combobox.SetSelection(3) + else: + choices = [int(x) for x in DistillerUiMainFrame.NUM_TRAINING_EXAMPLES_PER_SAMPLE_OUTPUT_CHOICES[:-1]] + self.num_training_examples_per_sample_output_combobox.SetSelection( + choices.index(self.state.config.body_morpher_num_training_examples_per_sample_output)) + + self.save_menu_item.Enable(self.state.can_save()) + + self.Refresh() + + def draw_nothing_yet_string_to_bitmap(self, bitmap, width: int, height: int): + dc = wx.MemoryDC() + dc.SelectObject(bitmap) + + dc.Clear() + font = wx.Font(wx.FontInfo(14).Family(wx.FONTFAMILY_SWISS)) + dc.SetFont(font) + w, h = dc.GetTextExtent("Nothing yet!") + dc.DrawText("Nothing yet!", (width - w) // 2, (height - h) // 2) + + del dc + + def try_saving(self): + if not self.state.can_save(): + message_dialog = wx.MessageDialog( + self, + "Cannot save yet! Please make sure you set the prefix, character_image_file_name, " + "and face_mask_image_file_name first.", + "Error", + wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + return False + else: + if self.state.need_to_check_overwrite(): + confirmation_dialog = wx.MessageDialog( + parent=self, + message=f"Overwriting {self.state.config.config_yaml_file_name()}?", + caption="Confirmation", + style=wx.YES_NO | wx.CANCEL | wx.ICON_QUESTION) + result = confirmation_dialog.ShowModal() + if result == wx.ID_YES: + self.state.save() + return True + elif result == wx.ID_NO: + return False + else: + return False + else: + self.state.save() + return True + + def on_save(self, event): + return self.try_saving() + + def on_new(self, event): + if self.state.dirty: + confirmation_dialog = wx.MessageDialog( + parent=self, + message=f"You have not saved the current config. Do you want to proceed?", + caption="Confirmation", + style=wx.YES_NO | wx.ICON_QUESTION) + result = confirmation_dialog.ShowModal() + if result == wx.ID_NO: + return + self.state = DistillerConfigState() + self.update_ui() + + def on_open(self, event): + if self.state.dirty: + confirmation_dialog = wx.MessageDialog( + parent=self, + message=f"You have not saved the current config. Do you want to proceed?", + caption="Confirmation", + style=wx.YES_NO | wx.ICON_QUESTION) + result = confirmation_dialog.ShowModal() + if result == wx.ID_NO: + return + + file_dialog = wx.FileDialog(self, "Choose a YAML file", wildcard="*.yaml", style=wx.FD_OPEN) + if file_dialog.ShowModal() != wx.ID_OK: + return + file_name = file_dialog.GetPath() + try: + self.state.load(file_name) + self.face_image_pytorch = None + self.face_mask_image_pytorch = None + self.update_face_image_bitmap(self.state.config.character_image_file_name) + self.update_face_mask_image_bitmap(self.state.config.face_mask_image_file_name) + self.update_ui() + except Exception as e: + message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + + def on_run(self, event): + try: + self.state.config.check() + except Exception as e: + message_dialog = wx.MessageDialog(self, str(e), "Error", wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + return + + if self.state.dirty: + message_dialog = wx.MessageDialog( + self, + "Please save the configuration first.", + "Error", + wx.OK | wx.ICON_ERROR) + message_dialog.ShowModal() + return + + self.config_file_to_run = self.state.config.config_yaml_file_name() + self.Destroy() diff --git a/tha4/src/tha4/image_util.py b/tha4/src/tha4/image_util.py new file mode 100644 index 0000000..cbe8482 --- /dev/null +++ b/tha4/src/tha4/image_util.py @@ -0,0 +1,58 @@ +import math + +import PIL.Image +import numpy +import torch +from matplotlib import cm +from tha4.shion.base.image_util import numpy_linear_to_srgb, pytorch_rgba_to_numpy_image, pytorch_rgb_to_numpy_image, \ + torch_linear_to_srgb + + +def grid_change_to_numpy_image(torch_image, num_channels=3): + height = torch_image.shape[1] + width = torch_image.shape[2] + size_image = (torch_image[0, :, :] ** 2 + torch_image[1, :, :] ** 2).sqrt().view(height, width, 1).numpy() + hsv = cm.get_cmap('hsv') + angle_image = hsv(((torch.atan2( + torch_image[0, :, :].view(height * width), + torch_image[1, :, :].view(height * width)).view(height, width) + math.pi) / (2 * math.pi)).numpy()) * 3 + numpy_image = size_image * angle_image[:, :, 0:3] + rgb_image = numpy_linear_to_srgb(numpy_image) + if num_channels == 3: + return rgb_image + elif num_channels == 4: + return numpy.concatenate([rgb_image, numpy.ones_like(size_image)], axis=2) + else: + raise RuntimeError("Unsupported num_channels: " + str(num_channels)) + + +def resize_PIL_image(pil_image, size=(256, 256)): + w, h = pil_image.size + d = min(w, h) + r = ((w - d) // 2, (h - d) // 2, (w + d) // 2, (h + d) // 2) + return pil_image.resize(size, resample=PIL.Image.LANCZOS, box=r) + + +def convert_output_image_from_torch_to_numpy(output_image): + if output_image.shape[2] == 2: + h, w, c = output_image.shape + numpy_image = torch.transpose(output_image.reshape(h * w, c), 0, 1).reshape(c, h, w) + elif output_image.shape[0] == 4: + numpy_image = pytorch_rgba_to_numpy_image(output_image) + elif output_image.shape[0] == 3: + numpy_image = pytorch_rgb_to_numpy_image(output_image) + elif output_image.shape[0] == 1: + c, h, w = output_image.shape + alpha_image = torch.cat([output_image.repeat(3, 1, 1) * 2.0 - 1.0, torch.ones(1, h, w)], dim=0) + numpy_image = pytorch_rgba_to_numpy_image(alpha_image) + elif output_image.shape[0] == 2: + numpy_image = grid_change_to_numpy_image(output_image, num_channels=4) + else: + raise RuntimeError("Unsupported # image channels: %d" % output_image.shape[0]) + numpy_image = numpy.uint8(numpy.rint(numpy_image * 255.0)) + return numpy_image + + +def convert_linear_to_srgb(image: torch.Tensor) -> torch.Tensor: + rgb_image = torch_linear_to_srgb(image[0:3, :, :]) + return torch.cat([rgb_image, image[3:4, :, :]], dim=0) diff --git a/tha4/src/tha4/mocap/__init__.py b/tha4/src/tha4/mocap/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/mocap/ifacialmocap_constants.py b/tha4/src/tha4/mocap/ifacialmocap_constants.py new file mode 100644 index 0000000..27031ac --- /dev/null +++ b/tha4/src/tha4/mocap/ifacialmocap_constants.py @@ -0,0 +1,239 @@ +EYE_LOOK_IN_LEFT = "eyeLookInLeft" +EYE_LOOK_OUT_LEFT = "eyeLookOutLeft" +EYE_LOOK_DOWN_LEFT = "eyeLookDownLeft" +EYE_LOOK_UP_LEFT = "eyeLookUpLeft" +EYE_BLINK_LEFT = "eyeBlinkLeft" +EYE_SQUINT_LEFT = "eyeSquintLeft" +EYE_WIDE_LEFT = "eyeWideLeft" +EYE_LOOK_IN_RIGHT = "eyeLookInRight" +EYE_LOOK_OUT_RIGHT = "eyeLookOutRight" +EYE_LOOK_DOWN_RIGHT = "eyeLookDownRight" +EYE_LOOK_UP_RIGHT = "eyeLookUpRight" +EYE_BLINK_RIGHT = "eyeBlinkRight" +EYE_SQUINT_RIGHT = "eyeSquintRight" +EYE_WIDE_RIGHT = "eyeWideRight" +BROW_DOWN_LEFT = "browDownLeft" +BROW_OUTER_UP_LEFT = "browOuterUpLeft" +BROW_DOWN_RIGHT = "browDownRight" +BROW_OUTER_UP_RIGHT = "browOuterUpRight" +BROW_INNER_UP = "browInnerUp" +NOSE_SNEER_LEFT = "noseSneerLeft" +NOSE_SNEER_RIGHT = "noseSneerRight" +CHEEK_SQUINT_LEFT = "cheekSquintLeft" +CHEEK_SQUINT_RIGHT = "cheekSquintRight" +CHEEK_PUFF = "cheekPuff" +MOUTH_LEFT = "mouthLeft" +MOUTH_DIMPLE_LEFT = "mouthDimpleLeft" +MOUTH_FROWN_LEFT = "mouthFrownLeft" +MOUTH_LOWER_DOWN_LEFT = "mouthLowerDownLeft" +MOUTH_PRESS_LEFT = "mouthPressLeft" +MOUTH_SMILE_LEFT = "mouthSmileLeft" +MOUTH_STRETCH_LEFT = "mouthStretchLeft" +MOUTH_UPPER_UP_LEFT = "mouthUpperUpLeft" +MOUTH_RIGHT = "mouthRight" +MOUTH_DIMPLE_RIGHT = "mouthDimpleRight" +MOUTH_FROWN_RIGHT = "mouthFrownRight" +MOUTH_LOWER_DOWN_RIGHT = "mouthLowerDownRight" +MOUTH_PRESS_RIGHT = "mouthPressRight" +MOUTH_SMILE_RIGHT = "mouthSmileRight" +MOUTH_STRETCH_RIGHT = "mouthStretchRight" +MOUTH_UPPER_UP_RIGHT = "mouthUpperUpRight" +MOUTH_CLOSE = "mouthClose" +MOUTH_FUNNEL = "mouthFunnel" +MOUTH_PUCKER = "mouthPucker" +MOUTH_ROLL_LOWER = "mouthRollLower" +MOUTH_ROLL_UPPER = "mouthRollUpper" +MOUTH_SHRUG_LOWER = "mouthShrugLower" +MOUTH_SHRUG_UPPER = "mouthShrugUpper" +JAW_LEFT = "jawLeft" +JAW_RIGHT = "jawRight" +JAW_FORWARD = "jawForward" +JAW_OPEN = "jawOpen" +TONGUE_OUT = "tongueOut" + +BLENDSHAPE_NAMES = [ + EYE_LOOK_IN_LEFT, # 0 + EYE_LOOK_OUT_LEFT, # 1 + EYE_LOOK_DOWN_LEFT, # 2 + EYE_LOOK_UP_LEFT, # 3 + EYE_BLINK_LEFT, # 4 + EYE_SQUINT_LEFT, # 5 + EYE_WIDE_LEFT, # 6 + EYE_LOOK_IN_RIGHT, # 7 + EYE_LOOK_OUT_RIGHT, # 8 + EYE_LOOK_DOWN_RIGHT, # 9 + EYE_LOOK_UP_RIGHT, # 10 + EYE_BLINK_RIGHT, # 11 + EYE_SQUINT_RIGHT, # 12 + EYE_WIDE_RIGHT, # 13 + BROW_DOWN_LEFT, # 14 + BROW_OUTER_UP_LEFT, # 15 + BROW_DOWN_RIGHT, # 16 + BROW_OUTER_UP_RIGHT, # 17 + BROW_INNER_UP, # 18 + NOSE_SNEER_LEFT, # 19 + NOSE_SNEER_RIGHT, # 20 + CHEEK_SQUINT_LEFT, # 21 + CHEEK_SQUINT_RIGHT, # 22 + CHEEK_PUFF, # 23 + MOUTH_LEFT, # 24 + MOUTH_DIMPLE_LEFT, # 25 + MOUTH_FROWN_LEFT, # 26 + MOUTH_LOWER_DOWN_LEFT, # 27 + MOUTH_PRESS_LEFT, # 28 + MOUTH_SMILE_LEFT, # 29 + MOUTH_STRETCH_LEFT, # 30 + MOUTH_UPPER_UP_LEFT, # 31 + MOUTH_RIGHT, # 32 + MOUTH_DIMPLE_RIGHT, # 33 + MOUTH_FROWN_RIGHT, # 34 + MOUTH_LOWER_DOWN_RIGHT, # 35 + MOUTH_PRESS_RIGHT, # 36 + MOUTH_SMILE_RIGHT, # 37 + MOUTH_STRETCH_RIGHT, # 38 + MOUTH_UPPER_UP_RIGHT, # 39 + MOUTH_CLOSE, # 40 + MOUTH_FUNNEL, # 41 + MOUTH_PUCKER, # 42 + MOUTH_ROLL_LOWER, # 43 + MOUTH_ROLL_UPPER, # 44 + MOUTH_SHRUG_LOWER, # 45 + MOUTH_SHRUG_UPPER, # 46 + JAW_LEFT, # 47 + JAW_RIGHT, # 48 + JAW_FORWARD, # 49 + JAW_OPEN, # 50 + TONGUE_OUT, # 51 +] + +EYE_LEFT_BLENDSHAPES = [ + EYE_LOOK_IN_LEFT, # 0 + EYE_LOOK_OUT_LEFT, # 1 + EYE_LOOK_DOWN_LEFT, # 2 + EYE_LOOK_UP_LEFT, # 3 + EYE_BLINK_LEFT, # 4 + EYE_SQUINT_LEFT, # 5 + EYE_WIDE_LEFT, # 6 +] + +EYE_RIGHT_BLENDSHAPES = [ + EYE_LOOK_IN_RIGHT, # 7 + EYE_LOOK_OUT_RIGHT, # 8 + EYE_LOOK_DOWN_RIGHT, # 9 + EYE_LOOK_UP_RIGHT, # 10 + EYE_BLINK_RIGHT, # 11 + EYE_SQUINT_RIGHT, # 12 + EYE_WIDE_RIGHT, # 13 +] + +BROW_LEFT_BLENDSHAPES = [ + BROW_DOWN_LEFT, # 14 + BROW_OUTER_UP_LEFT, # 15 + +] + +BROW_RIGHT_BLENDSHAPES = [ + BROW_DOWN_RIGHT, # 16 + BROW_OUTER_UP_RIGHT, # 17 + +] + +BROW_BOTH_BLENDSHAPES = [ + BROW_INNER_UP, # 18 +] + +NOSE_BLENDSHAPES = [ + NOSE_SNEER_LEFT, # 19 + NOSE_SNEER_RIGHT, # 20 +] + +CHECK_BLENDSHAPES = [ + CHEEK_SQUINT_LEFT, # 21 + CHEEK_SQUINT_RIGHT, # 22 + CHEEK_PUFF, # 23 +] + +MOUTH_LEFT_BLENDSHAPES = [ + MOUTH_LEFT, # 24 + MOUTH_DIMPLE_LEFT, # 25 + MOUTH_FROWN_LEFT, # 26 + MOUTH_LOWER_DOWN_LEFT, # 27 + MOUTH_PRESS_LEFT, # 28 + MOUTH_SMILE_LEFT, # 29 + MOUTH_STRETCH_LEFT, # 30 + MOUTH_UPPER_UP_LEFT, # 31 +] + +MOUTH_RIGHT_BLENDSHAPES = [ + MOUTH_RIGHT, # 32 + MOUTH_DIMPLE_RIGHT, # 33 + MOUTH_FROWN_RIGHT, # 34 + MOUTH_LOWER_DOWN_RIGHT, # 35 + MOUTH_PRESS_RIGHT, # 36 + MOUTH_SMILE_RIGHT, # 37 + MOUTH_STRETCH_RIGHT, # 38 + MOUTH_UPPER_UP_RIGHT, # 39 +] + +MOUTH_BOTH_BLENDSHAPES = [ + MOUTH_CLOSE, # 40 + MOUTH_FUNNEL, # 41 + MOUTH_PUCKER, # 42 + MOUTH_ROLL_LOWER, # 43 + MOUTH_ROLL_UPPER, # 44 + MOUTH_SHRUG_LOWER, # 45 + MOUTH_SHRUG_UPPER, # 46 +] + +JAW_BLENDSHAPES = [ + JAW_LEFT, # 47 + JAW_RIGHT, # 48 + JAW_FORWARD, # 49 + JAW_OPEN, # 50 +] + +TONGUE_BLENDSHAPES = [ + TONGUE_OUT, # 51 +] + +COLUMN_0_BLENDSHAPES = EYE_RIGHT_BLENDSHAPES + BROW_RIGHT_BLENDSHAPES + [NOSE_SNEER_RIGHT, CHEEK_SQUINT_RIGHT] +COLUMN_1_BLENDSHAPES = EYE_LEFT_BLENDSHAPES + BROW_LEFT_BLENDSHAPES + [NOSE_SNEER_LEFT, CHEEK_SQUINT_LEFT] +COLUMN_2_BLENDSHAPES = MOUTH_RIGHT_BLENDSHAPES + [JAW_RIGHT] +COLUMN_3_BLENDSHAPES = MOUTH_LEFT_BLENDSHAPES + [JAW_LEFT] +COLUMN_4_BLENDSHAPES = [BROW_INNER_UP, CHEEK_PUFF] + MOUTH_BOTH_BLENDSHAPES + [JAW_FORWARD, JAW_OPEN, TONGUE_OUT] + +BLENDSHAPE_COLUMNS = [ + COLUMN_0_BLENDSHAPES, + COLUMN_1_BLENDSHAPES, + COLUMN_2_BLENDSHAPES, + COLUMN_3_BLENDSHAPES, + COLUMN_4_BLENDSHAPES, +] + +RIGHT_EYE_BONE_X = "rightEyeBoneX" +RIGHT_EYE_BONE_Y = "rightEyeBoneY" +RIGHT_EYE_BONE_Z = "rightEyeBoneZ" +RIGHT_EYE_BONE_ROTATIONS = [RIGHT_EYE_BONE_X, RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z] + +LEFT_EYE_BONE_X = "leftEyeBoneX" +LEFT_EYE_BONE_Y = "leftEyeBoneY" +LEFT_EYE_BONE_Z = "leftEyeBoneZ" +LEFT_EYE_BONE_ROTATIONS = [LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z] + +HEAD_BONE_X = "headBoneX" +HEAD_BONE_Y = "headBoneY" +HEAD_BONE_Z = "headBoneZ" +HEAD_BONE_ROTATIONS = [HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z] + +ROTATION_NAMES = RIGHT_EYE_BONE_ROTATIONS + LEFT_EYE_BONE_ROTATIONS + HEAD_BONE_ROTATIONS + +RIGHT_EYE_BONE_QUAT = "rightEyeBoneQuat" +LEFT_EYE_BONE_QUAT = "leftEyeBoneQuat" +HEAD_BONE_QUAT = "headBoneQuat" +QUATERNION_NAMES = [ + RIGHT_EYE_BONE_QUAT, + LEFT_EYE_BONE_QUAT, + HEAD_BONE_QUAT +] + +IFACIALMOCAP_DATETIME_FORMAT = "%Y/%m/%d-%H:%M:%S.%f" diff --git a/tha4/src/tha4/mocap/ifacialmocap_pose.py b/tha4/src/tha4/mocap/ifacialmocap_pose.py new file mode 100644 index 0000000..595b3c8 --- /dev/null +++ b/tha4/src/tha4/mocap/ifacialmocap_pose.py @@ -0,0 +1,27 @@ +from tha4.mocap.ifacialmocap_constants import BLENDSHAPE_NAMES, HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z, \ + HEAD_BONE_QUAT, LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z, LEFT_EYE_BONE_QUAT, RIGHT_EYE_BONE_X, \ + RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z, RIGHT_EYE_BONE_QUAT + + +def create_default_ifacialmocap_pose(): + data = {} + + for blendshape_name in BLENDSHAPE_NAMES: + data[blendshape_name] = 0.0 + + data[HEAD_BONE_X] = 0.0 + data[HEAD_BONE_Y] = 0.0 + data[HEAD_BONE_Z] = 0.0 + data[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + + data[LEFT_EYE_BONE_X] = 0.0 + data[LEFT_EYE_BONE_Y] = 0.0 + data[LEFT_EYE_BONE_Z] = 0.0 + data[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + + data[RIGHT_EYE_BONE_X] = 0.0 + data[RIGHT_EYE_BONE_Y] = 0.0 + data[RIGHT_EYE_BONE_Z] = 0.0 + data[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + + return data \ No newline at end of file diff --git a/tha4/src/tha4/mocap/ifacialmocap_pose_converter.py b/tha4/src/tha4/mocap/ifacialmocap_pose_converter.py new file mode 100644 index 0000000..390460b --- /dev/null +++ b/tha4/src/tha4/mocap/ifacialmocap_pose_converter.py @@ -0,0 +1,12 @@ +from abc import ABC, abstractmethod +from typing import Dict, List + + +class IFacialMocapPoseConverter(ABC): + @abstractmethod + def convert(self, ifacialmocap_pose: Dict[str, float]) -> List[float]: + pass + + @abstractmethod + def init_pose_converter_panel(self, parent): + pass \ No newline at end of file diff --git a/tha4/src/tha4/mocap/ifacialmocap_pose_converter_25.py b/tha4/src/tha4/mocap/ifacialmocap_pose_converter_25.py new file mode 100644 index 0000000..25a7bab --- /dev/null +++ b/tha4/src/tha4/mocap/ifacialmocap_pose_converter_25.py @@ -0,0 +1,614 @@ +import math +import time +from enum import Enum +from typing import Optional, Dict, List, Callable + +import numpy +import scipy.optimize +import wx + +from tha4.mocap.ifacialmocap_constants import MOUTH_SMILE_LEFT, MOUTH_SHRUG_UPPER, MOUTH_SMILE_RIGHT, \ + BROW_INNER_UP, BROW_OUTER_UP_RIGHT, BROW_OUTER_UP_LEFT, BROW_DOWN_LEFT, BROW_DOWN_RIGHT, EYE_WIDE_LEFT, \ + EYE_WIDE_RIGHT, EYE_BLINK_LEFT, EYE_BLINK_RIGHT, CHEEK_SQUINT_LEFT, CHEEK_SQUINT_RIGHT, EYE_LOOK_IN_LEFT, \ + EYE_LOOK_OUT_LEFT, EYE_LOOK_IN_RIGHT, EYE_LOOK_OUT_RIGHT, EYE_LOOK_UP_LEFT, EYE_LOOK_UP_RIGHT, EYE_LOOK_DOWN_RIGHT, \ + EYE_LOOK_DOWN_LEFT, HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z, JAW_OPEN, MOUTH_FROWN_LEFT, MOUTH_FROWN_RIGHT, \ + MOUTH_LOWER_DOWN_LEFT, MOUTH_LOWER_DOWN_RIGHT, MOUTH_FUNNEL, MOUTH_PUCKER +from tha4.mocap.ifacialmocap_pose_converter import IFacialMocapPoseConverter +from tha4.poser.modes.pose_parameters import get_pose_parameters + + +class EyebrowDownMode(Enum): + TROUBLED = 1 + ANGRY = 2 + LOWERED = 3 + SERIOUS = 4 + + +class WinkMode(Enum): + NORMAL = 1 + RELAXED = 2 + + +def rad_to_deg(rad): + return rad * 180.0 / math.pi + + +def deg_to_rad(deg): + return deg * math.pi / 180.0 + + +def clamp(x, min_value, max_value): + return max(min_value, min(max_value, x)) + + +class IFacialMocapPoseConverter25Args: + def __init__(self, + smile_threshold_min: float = 0.4, + smile_threshold_max: float = 0.6, + eyebrow_down_mode: EyebrowDownMode = EyebrowDownMode.ANGRY, + wink_mode: WinkMode = WinkMode.NORMAL, + eye_surprised_max: float = 0.5, + eye_blink_max: float = 0.8, + eyebrow_down_max: float = 0.4, + cheek_squint_min: float = 0.1, + cheek_squint_max: float = 0.7, + eye_rotation_factor: float = 1.0 / 0.75, + jaw_open_min: float = 0.1, + jaw_open_max: float = 0.4, + mouth_frown_max: float = 0.6, + mouth_funnel_min: float = 0.25, + mouth_funnel_max: float = 0.5, + iris_small_left=0.0, + iris_small_right=0.0): + self.iris_small_right = iris_small_left + self.iris_small_left = iris_small_right + + self.wink_mode = wink_mode + + self.mouth_funnel_max = mouth_funnel_max + self.mouth_funnel_min = mouth_funnel_min + self.mouth_frown_max = mouth_frown_max + + self.jaw_open_max = jaw_open_max + self.jaw_open_min = jaw_open_min + + self.eye_rotation_factor = eye_rotation_factor + + self.cheek_squint_max = cheek_squint_max + self.cheek_squint_min = cheek_squint_min + + self.eyebrow_down_max = eyebrow_down_max + + self.eye_blink_max = eye_blink_max + self.eye_surprised_max = eye_surprised_max + self.eyebrow_down_mode = eyebrow_down_mode + + self.smile_threshold_min = smile_threshold_min + self.smile_threshold_max = smile_threshold_max + + def set_smile_threshold_min(self, new_value: float): + self.smile_threshold_min = new_value + + def set_smile_threshold_max(self, new_value: float): + self.smile_threshold_max = new_value + + def set_eye_surprised_max(self, new_value: float): + self.eye_surprised_max = new_value + + def set_eye_blink_max(self, new_value: float): + self.eye_blink_max = new_value + + def set_eyebrow_down_max(self, new_value: float): + self.eyebrow_down_max = new_value + + def set_cheek_squint_min(self, new_value: float): + self.cheek_squint_min = new_value + + def set_cheek_squint_max(self, new_value: float): + self.cheek_squint_max = new_value + + def set_jaw_open_min(self, new_value: float): + self.jaw_open_min = new_value + + def set_jaw_open_max(self, new_value: float): + self.jaw_open_max = new_value + + def set_mouth_frown_max(self, new_value: float): + self.mouth_frown_max = new_value + + def set_mouth_funnel_min(self, new_value: float): + self.mouth_funnel_min = new_value + + def set_mouth_funnel_max(self, new_value: float): + self.mouth_funnel_min = new_value + + +class IFacialMocapPoseConverter25(IFacialMocapPoseConverter): + def __init__(self, args: Optional[IFacialMocapPoseConverter25Args] = None): + super().__init__() + if args is None: + args = IFacialMocapPoseConverter25Args() + self.args = args + pose_parameters = get_pose_parameters() + self.pose_size = 45 + + self.eyebrow_troubled_left_index = pose_parameters.get_parameter_index("eyebrow_troubled_left") + self.eyebrow_troubled_right_index = pose_parameters.get_parameter_index("eyebrow_troubled_right") + self.eyebrow_angry_left_index = pose_parameters.get_parameter_index("eyebrow_angry_left") + self.eyebrow_angry_right_index = pose_parameters.get_parameter_index("eyebrow_angry_right") + self.eyebrow_happy_left_index = pose_parameters.get_parameter_index("eyebrow_happy_left") + self.eyebrow_happy_right_index = pose_parameters.get_parameter_index("eyebrow_happy_right") + self.eyebrow_raised_left_index = pose_parameters.get_parameter_index("eyebrow_raised_left") + self.eyebrow_raised_right_index = pose_parameters.get_parameter_index("eyebrow_raised_right") + self.eyebrow_lowered_left_index = pose_parameters.get_parameter_index("eyebrow_lowered_left") + self.eyebrow_lowered_right_index = pose_parameters.get_parameter_index("eyebrow_lowered_right") + self.eyebrow_serious_left_index = pose_parameters.get_parameter_index("eyebrow_serious_left") + self.eyebrow_serious_right_index = pose_parameters.get_parameter_index("eyebrow_serious_right") + + self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left") + self.eye_surprised_right_index = pose_parameters.get_parameter_index("eye_surprised_right") + self.eye_wink_left_index = pose_parameters.get_parameter_index("eye_wink_left") + self.eye_wink_right_index = pose_parameters.get_parameter_index("eye_wink_right") + self.eye_happy_wink_left_index = pose_parameters.get_parameter_index("eye_happy_wink_left") + self.eye_happy_wink_right_index = pose_parameters.get_parameter_index("eye_happy_wink_right") + self.eye_relaxed_left_index = pose_parameters.get_parameter_index("eye_relaxed_left") + self.eye_relaxed_right_index = pose_parameters.get_parameter_index("eye_relaxed_right") + self.eye_raised_lower_eyelid_left_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_left") + self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right") + + self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left") + self.iris_small_right_index = pose_parameters.get_parameter_index("iris_small_right") + + self.iris_rotation_x_index = pose_parameters.get_parameter_index("iris_rotation_x") + self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y") + + self.head_x_index = pose_parameters.get_parameter_index("head_x") + self.head_y_index = pose_parameters.get_parameter_index("head_y") + self.neck_z_index = pose_parameters.get_parameter_index("neck_z") + + self.mouth_aaa_index = pose_parameters.get_parameter_index("mouth_aaa") + self.mouth_iii_index = pose_parameters.get_parameter_index("mouth_iii") + self.mouth_uuu_index = pose_parameters.get_parameter_index("mouth_uuu") + self.mouth_eee_index = pose_parameters.get_parameter_index("mouth_eee") + self.mouth_ooo_index = pose_parameters.get_parameter_index("mouth_ooo") + + self.mouth_lowered_corner_left_index = pose_parameters.get_parameter_index("mouth_lowered_corner_left") + self.mouth_lowered_corner_right_index = pose_parameters.get_parameter_index("mouth_lowered_corner_right") + self.mouth_raised_corner_left_index = pose_parameters.get_parameter_index("mouth_raised_corner_left") + self.mouth_raised_corner_right_index = pose_parameters.get_parameter_index("mouth_raised_corner_right") + + self.body_y_index = pose_parameters.get_parameter_index("body_y") + self.body_z_index = pose_parameters.get_parameter_index("body_z") + self.breathing_index = pose_parameters.get_parameter_index("breathing") + + self.breathing_start_time = time.time() + + self.panel = None + + def init_pose_converter_panel(self, parent): + self.panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + self.panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.panel.SetSizer(self.panel_sizer) + self.panel.SetAutoLayout(1) + parent.GetSizer().Add(self.panel, 0, wx.EXPAND) + + if True: + eyebrow_down_mode_text = wx.StaticText(self.panel, label=" --- Eyebrow Down Mode --- ", + style=wx.ALIGN_CENTER) + self.panel_sizer.Add(eyebrow_down_mode_text, 0, wx.EXPAND) + + self.eyebrow_down_mode_choice = wx.Choice( + self.panel, + choices=[ + "ANGRY", + "TROUBLED", + "SERIOUS", + "LOWERED", + ]) + self.eyebrow_down_mode_choice.SetSelection(0) + self.panel_sizer.Add(self.eyebrow_down_mode_choice, 0, wx.EXPAND) + self.eyebrow_down_mode_choice.Bind(wx.EVT_CHOICE, self.change_eyebrow_down_mode) + + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + if True: + wink_mode_text = wx.StaticText(self.panel, label=" --- Wink Mode --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(wink_mode_text, 0, wx.EXPAND) + + self.wink_mode_choice = wx.Choice( + self.panel, + choices=[ + "NORMAL", + "RELAXED", + ]) + self.wink_mode_choice.SetSelection(0) + self.panel_sizer.Add(self.wink_mode_choice, 0, wx.EXPAND) + self.wink_mode_choice.Bind(wx.EVT_CHOICE, self.change_wink_mode) + + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + if True: + iris_size_text = wx.StaticText(self.panel, label=" --- Iris Size --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(iris_size_text, 0, wx.EXPAND) + + self.iris_left_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.iris_left_slider, 0, wx.EXPAND) + self.iris_left_slider.Bind(wx.EVT_SLIDER, self.change_iris_size) + + self.iris_right_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.iris_right_slider, 0, wx.EXPAND) + self.iris_right_slider.Bind(wx.EVT_SLIDER, self.change_iris_size) + self.iris_right_slider.Enable(False) + + self.link_left_right_irises = wx.CheckBox( + self.panel, label="Use same value for both sides") + self.link_left_right_irises.SetValue(True) + self.panel_sizer.Add(self.link_left_right_irises, wx.SizerFlags().CenterHorizontal().Border()) + self.link_left_right_irises.Bind(wx.EVT_CHECKBOX, self.link_left_right_irises_clicked) + + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + if True: + breathing_frequency_text = wx.StaticText( + self.panel, label=" --- Breathing --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(breathing_frequency_text, 0, wx.EXPAND) + + self.restart_breathing_cycle_button = wx.Button(self.panel, label="Restart Breathing Cycle") + self.restart_breathing_cycle_button.Bind(wx.EVT_BUTTON, self.restart_breathing_cycle_clicked) + self.panel_sizer.Add(self.restart_breathing_cycle_button, 0, wx.EXPAND) + + self.breathing_frequency_slider = wx.Slider( + self.panel, minValue=0, maxValue=60, value=20, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.breathing_frequency_slider, 0, wx.EXPAND) + + self.breathing_gauge = wx.Gauge(self.panel, style=wx.GA_HORIZONTAL, range=1000) + self.panel_sizer.Add(self.breathing_gauge, 0, wx.EXPAND) + + if True: + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + convertion_parameters_text = wx.StaticText( + self.panel, label="--- Conversion Parameters ---", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(convertion_parameters_text, 0, wx.EXPAND) + + conversion_param_panel = wx.Panel(self.panel) + self.panel_sizer.Add(conversion_param_panel, 0, wx.EXPAND) + conversion_panel_sizer = wx.FlexGridSizer(cols=2) + conversion_panel_sizer.AddGrowableCol(1) + conversion_param_panel.SetSizer(conversion_panel_sizer) + conversion_param_panel.SetAutoLayout(1) + + self.smile_thresold_min_spin = self.create_spin_control( + conversion_param_panel, + "Smile Threshold Min:", self.args.smile_threshold_min, self.args.set_smile_threshold_min) + self.smile_thresold_max_spin = self.create_spin_control( + conversion_param_panel, + "Smile Threshold Max:", self.args.smile_threshold_max, self.args.set_smile_threshold_max) + self.eye_surprised_max_spin = self.create_spin_control( + conversion_param_panel, + "Eye Surprised Max:", self.args.eye_surprised_max, self.args.set_eye_surprised_max) + self.eye_blink_max_spin = self.create_spin_control( + conversion_param_panel, + "Eye Blink Max:", self.args.eye_blink_max, self.args.set_eye_blink_max) + self.eyebrow_down_max_spin = self.create_spin_control( + conversion_param_panel, + "Eyebrow Down Max:", self.args.eyebrow_down_max, self.args.set_eyebrow_down_max) + self.cheek_squint_min_spin = self.create_spin_control( + conversion_param_panel, + "Cheek Squint Min:", self.args.cheek_squint_min, self.args.set_cheek_squint_min) + self.cheek_squint_max_spin = self.create_spin_control( + conversion_param_panel, + "Cheek Squint Max:", self.args.cheek_squint_max, self.args.set_cheek_squint_max) + self.jaw_open_min_spin = self.create_spin_control( + conversion_param_panel, + "Jaw Open Min:", self.args.jaw_open_min, self.args.set_jaw_open_min) + self.jaw_open_max_spin = self.create_spin_control( + conversion_param_panel, + "Jaw Open Max:", self.args.jaw_open_max, self.args.set_jaw_open_max) + self.mouth_frown_max_spin = self.create_spin_control( + conversion_param_panel, + "Mouth Frown Max:", self.args.mouth_frown_max, self.args.set_mouth_frown_max) + self.mouth_funnel_min_spin = self.create_spin_control( + conversion_param_panel, + "Mouth Funnel Min:", self.args.mouth_funnel_min, self.args.set_mouth_funnel_min) + self.mouth_funnel_max_spin = self.create_spin_control( + conversion_param_panel, + "Mouth Funnel Max:", self.args.mouth_funnel_max, self.args.set_mouth_funnel_max) + + self.panel_sizer.Fit(self.panel) + + def create_spin_control(self, parent, label: str, initial_value: float, set_func: Callable[[float], None]): + sizer = parent.GetSizer() + + text = wx.StaticText(parent, label=label) + sizer.Add(text, wx.SizerFlags().Right().Border(wx.ALL, 2)) + + spin_ctrl = wx.SpinCtrlDouble( + parent, + wx.ID_ANY, + min=0.0, + max=1.0, + initial=initial_value, + inc=0.01) + sizer.Add(spin_ctrl, wx.SizerFlags().Border(wx.ALL, 2).Expand()) + + def handler(event: wx.Event): + new_value = spin_ctrl.GetValue() + set_func(new_value) + + spin_ctrl.Bind(wx.EVT_SPINCTRLDOUBLE, handler) + + return spin_ctrl + + def restart_breathing_cycle_clicked(self, event: wx.Event): + self.breathing_start_time = time.time() + + def change_eyebrow_down_mode(self, event: wx.Event): + selected_index = self.eyebrow_down_mode_choice.GetSelection() + if selected_index == 0: + self.args.eyebrow_down_mode = EyebrowDownMode.ANGRY + elif selected_index == 1: + self.args.eyebrow_down_mode = EyebrowDownMode.TROUBLED + elif selected_index == 2: + self.args.eyebrow_down_mode = EyebrowDownMode.SERIOUS + else: + self.args.eyebrow_down_mode = EyebrowDownMode.LOWERED + + def change_wink_mode(self, event: wx.Event): + selected_index = self.wink_mode_choice.GetSelection() + if selected_index == 0: + self.args.wink_mode = WinkMode.NORMAL + else: + self.args.wink_mode = WinkMode.RELAXED + + def change_iris_size(self, event: wx.Event): + if self.link_left_right_irises.GetValue(): + left_value = self.iris_left_slider.GetValue() + right_value = self.iris_right_slider.GetValue() + if left_value != right_value: + self.iris_right_slider.SetValue(left_value) + self.args.iris_small_left = left_value / 1000.0 + self.args.iris_small_right = left_value / 1000.0 + else: + self.args.iris_small_left = self.iris_left_slider.GetValue() / 1000.0 + self.args.iris_small_right = self.iris_right_slider.GetValue() / 1000.0 + + def link_left_right_irises_clicked(self, event: wx.Event): + if self.link_left_right_irises.GetValue(): + self.iris_right_slider.Enable(False) + else: + self.iris_right_slider.Enable(True) + self.change_iris_size(event) + + def decompose_head_body_param(self, param, threshold=2.0 / 3): + if abs(param) < threshold: + return (param, 0.0) + else: + if param < 0: + sign = -1.0 + else: + sign = 1.0 + return (threshold * sign, (abs(param) - threshold) * sign) + + def convert(self, ifacialmocap_pose: Dict[str, float]) -> List[float]: + pose = [0.0 for i in range(self.pose_size)] + + smile_value = \ + (ifacialmocap_pose[MOUTH_SMILE_LEFT] + ifacialmocap_pose[MOUTH_SMILE_RIGHT]) / 2.0 \ + + ifacialmocap_pose[MOUTH_SHRUG_UPPER] + if self.args.smile_threshold_min >= self.args.smile_threshold_max: + smile_degree = 0.0 + else: + if smile_value < self.args.smile_threshold_min: + smile_degree = 0.0 + elif smile_value > self.args.smile_threshold_max: + smile_degree = 1.0 + else: + smile_degree = (smile_value - self.args.smile_threshold_min) / ( + self.args.smile_threshold_max - self.args.smile_threshold_min) + + # Eyebrow + if True: + brow_inner_up = ifacialmocap_pose[BROW_INNER_UP] + brow_outer_up_right = ifacialmocap_pose[BROW_OUTER_UP_RIGHT] + brow_outer_up_left = ifacialmocap_pose[BROW_OUTER_UP_LEFT] + + brow_up_left = clamp(brow_inner_up + brow_outer_up_left, 0.0, 1.0) + brow_up_right = clamp(brow_inner_up + brow_outer_up_right, 0.0, 1.0) + pose[self.eyebrow_raised_left_index] = brow_up_left + pose[self.eyebrow_raised_right_index] = brow_up_right + + if self.args.eyebrow_down_max <= 0.0: + brow_down_left = 0.0 + brow_down_right = 0.0 + else: + brow_down_left = (1.0 - smile_degree) \ + * clamp(ifacialmocap_pose[BROW_DOWN_LEFT] / self.args.eyebrow_down_max, 0.0, 1.0) + brow_down_right = (1.0 - smile_degree) \ + * clamp(ifacialmocap_pose[BROW_DOWN_RIGHT] / self.args.eyebrow_down_max, 0.0, 1.0) + + if self.args.eyebrow_down_mode == EyebrowDownMode.TROUBLED: + pose[self.eyebrow_troubled_left_index] = brow_down_left + pose[self.eyebrow_troubled_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.ANGRY: + pose[self.eyebrow_angry_left_index] = brow_down_left + pose[self.eyebrow_angry_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.LOWERED: + pose[self.eyebrow_lowered_left_index] = brow_down_left + pose[self.eyebrow_lowered_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.SERIOUS: + pose[self.eyebrow_serious_left_index] = brow_down_left + pose[self.eyebrow_serious_right_index] = brow_down_right + + brow_happy_value = clamp(smile_value, 0.0, 1.0) * smile_degree + pose[self.eyebrow_happy_left_index] = brow_happy_value + pose[self.eyebrow_happy_right_index] = brow_happy_value + + # Eye + if True: + # Surprised + if self.args.eye_surprised_max <= 0.0: + pose[self.eye_surprised_left_index] = 0.0 + pose[self.eye_surprised_right_index] = 0.0 + else: + pose[self.eye_surprised_left_index] = clamp( + ifacialmocap_pose[EYE_WIDE_LEFT] / self.args.eye_surprised_max, 0.0, 1.0) + pose[self.eye_surprised_right_index] = clamp( + ifacialmocap_pose[EYE_WIDE_RIGHT] / self.args.eye_surprised_max, 0.0, 1.0) + + # Wink + if self.args.wink_mode == WinkMode.NORMAL: + wink_left_index = self.eye_wink_left_index + wink_right_index = self.eye_wink_right_index + else: + wink_left_index = self.eye_relaxed_left_index + wink_right_index = self.eye_relaxed_right_index + if self.args.eye_blink_max <= 0: + pose[wink_left_index] = 0.0 + pose[wink_right_index] = 0.0 + pose[self.eye_happy_wink_left_index] = 0.0 + pose[self.eye_happy_wink_right_index] = 0.0 + else: + pose[wink_left_index] = (1.0 - smile_degree) * clamp( + ifacialmocap_pose[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0) + pose[wink_right_index] = (1.0 - smile_degree) * clamp( + ifacialmocap_pose[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0) + pose[self.eye_happy_wink_left_index] = smile_degree * clamp( + ifacialmocap_pose[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0) + pose[self.eye_happy_wink_right_index] = smile_degree * clamp( + ifacialmocap_pose[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0) + + # Lower eyelid + cheek_squint_denom = self.args.cheek_squint_max - self.args.cheek_squint_min + if cheek_squint_denom <= 0.0: + pose[self.eye_raised_lower_eyelid_left_index] = 0.0 + pose[self.eye_raised_lower_eyelid_right_index] = 0.0 + else: + pose[self.eye_raised_lower_eyelid_left_index] = \ + clamp( + (ifacialmocap_pose[CHEEK_SQUINT_LEFT] - self.args.cheek_squint_min) / cheek_squint_denom, + 0.0, 1.0) + pose[self.eye_raised_lower_eyelid_right_index] = \ + clamp( + (ifacialmocap_pose[CHEEK_SQUINT_RIGHT] - self.args.cheek_squint_min) / cheek_squint_denom, + 0.0, 1.0) + + # Iris rotation + if True: + eye_rotation_y = (ifacialmocap_pose[EYE_LOOK_IN_LEFT] + - ifacialmocap_pose[EYE_LOOK_OUT_LEFT] + - ifacialmocap_pose[EYE_LOOK_IN_RIGHT] + + ifacialmocap_pose[EYE_LOOK_OUT_RIGHT]) / 2.0 * self.args.eye_rotation_factor + pose[self.iris_rotation_y_index] = clamp(eye_rotation_y, -1.0, 1.0) + + eye_rotation_x = (ifacialmocap_pose[EYE_LOOK_UP_LEFT] + + ifacialmocap_pose[EYE_LOOK_UP_RIGHT] + - ifacialmocap_pose[EYE_LOOK_DOWN_LEFT] + - ifacialmocap_pose[EYE_LOOK_DOWN_RIGHT]) / 2.0 * self.args.eye_rotation_factor + pose[self.iris_rotation_x_index] = clamp(eye_rotation_x, -1.0, 1.0) + + # Iris size + if True: + pose[self.iris_small_left_index] = self.args.iris_small_left + pose[self.iris_small_right_index] = self.args.iris_small_right + + # Head rotation + if True: + x_param = clamp(-ifacialmocap_pose[HEAD_BONE_X] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + pose[self.head_x_index] = x_param + + y_param = clamp(-ifacialmocap_pose[HEAD_BONE_Y] * 180.0 / math.pi, -10.0, 10.0) / 10.0 + pose[self.head_y_index] = y_param + pose[self.body_y_index] = y_param + + z_param = clamp(ifacialmocap_pose[HEAD_BONE_Z] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + pose[self.neck_z_index] = z_param + pose[self.body_z_index] = z_param + + # Mouth + if True: + jaw_open_denom = self.args.jaw_open_max - self.args.jaw_open_min + if jaw_open_denom <= 0: + mouth_open = 0.0 + else: + mouth_open = clamp((ifacialmocap_pose[JAW_OPEN] - self.args.jaw_open_min) / jaw_open_denom, 0.0, 1.0) + pose[self.mouth_aaa_index] = mouth_open + pose[self.mouth_raised_corner_left_index] = clamp(smile_value, 0.0, 1.0) + pose[self.mouth_raised_corner_right_index] = clamp(smile_value, 0.0, 1.0) + + is_mouth_open = mouth_open > 0.0 + if not is_mouth_open: + if self.args.mouth_frown_max > 0: + mouth_frown_value = 0.0 + else: + mouth_frown_value = clamp( + (ifacialmocap_pose[MOUTH_FROWN_LEFT] + ifacialmocap_pose[ + MOUTH_FROWN_RIGHT]) / self.args.mouth_frown_max, 0.0, 1.0) + pose[self.mouth_lowered_corner_left_index] = mouth_frown_value + pose[self.mouth_lowered_corner_right_index] = mouth_frown_value + else: + mouth_lower_down = clamp( + ifacialmocap_pose[MOUTH_LOWER_DOWN_LEFT] + ifacialmocap_pose[MOUTH_LOWER_DOWN_RIGHT], 0.0, 1.0) + mouth_funnel = ifacialmocap_pose[MOUTH_FUNNEL] + mouth_pucker = ifacialmocap_pose[MOUTH_PUCKER] + + mouth_point = [mouth_open, mouth_lower_down, mouth_funnel, mouth_pucker] + + aaa_point = [1.0, 1.0, 0.0, 0.0] + iii_point = [0.0, 1.0, 0.0, 0.0] + uuu_point = [0.5, 0.3, 0.25, 0.75] + ooo_point = [1.0, 0.5, 0.5, 0.4] + + decomp = numpy.array([0, 0, 0, 0]) + M = numpy.array([ + aaa_point, + iii_point, + uuu_point, + ooo_point + ]) + + def loss(decomp): + return numpy.linalg.norm(numpy.matmul(decomp, M) - mouth_point) \ + + 0.01 * numpy.linalg.norm(decomp, ord=1) + + opt_result = scipy.optimize.minimize( + loss, decomp, bounds=[(0.0, 1.0), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)]) + decomp = opt_result["x"] + restricted_decomp = [decomp.item(0), decomp.item(1), decomp.item(2), decomp.item(3)] + pose[self.mouth_aaa_index] = restricted_decomp[0] + pose[self.mouth_iii_index] = restricted_decomp[1] + mouth_funnel_denom = self.args.mouth_funnel_max - self.args.mouth_funnel_min + if mouth_funnel_denom <= 0: + ooo_alpha = 0.0 + uo_value = 0.0 + else: + ooo_alpha = clamp((mouth_funnel - self.args.mouth_funnel_min) / mouth_funnel_denom, 0.0, 1.0) + uo_value = clamp(restricted_decomp[2] + restricted_decomp[3], 0.0, 1.0) + pose[self.mouth_uuu_index] = uo_value * (1.0 - ooo_alpha) + pose[self.mouth_ooo_index] = uo_value * ooo_alpha + + if self.panel is not None: + frequency = self.breathing_frequency_slider.GetValue() + if frequency == 0: + value = 0.0 + pose[self.breathing_index] = value + self.breathing_start_time = time.time() + else: + period = 60.0 / frequency + now = time.time() + diff = now - self.breathing_start_time + frac = (diff % period) / period + value = (-math.cos(2 * math.pi * frac) + 1.0) / 2.0 + pose[self.breathing_index] = value + self.breathing_gauge.SetValue(int(1000 * value)) + + return pose + + +def create_ifacialmocap_pose_converter( + args: Optional[IFacialMocapPoseConverter25Args] = None) -> IFacialMocapPoseConverter: + return IFacialMocapPoseConverter25(args) diff --git a/tha4/src/tha4/mocap/ifacialmocap_v2.py b/tha4/src/tha4/mocap/ifacialmocap_v2.py new file mode 100644 index 0000000..f61a34b --- /dev/null +++ b/tha4/src/tha4/mocap/ifacialmocap_v2.py @@ -0,0 +1,89 @@ +import math + +from tha4.mocap.ifacialmocap_constants import BLENDSHAPE_NAMES, HEAD_BONE_X, HEAD_BONE_Y, HEAD_BONE_Z, \ + RIGHT_EYE_BONE_X, RIGHT_EYE_BONE_Y, RIGHT_EYE_BONE_Z, LEFT_EYE_BONE_X, LEFT_EYE_BONE_Y, LEFT_EYE_BONE_Z, \ + HEAD_BONE_QUAT, LEFT_EYE_BONE_QUAT, RIGHT_EYE_BONE_QUAT + +IFACIALMOCAP_PORT = 49983 +IFACIALMOCAP_START_STRING = "iFacialMocap_sahuasouryya9218sauhuiayeta91555dy3719|sendDataVersion=v2".encode('utf-8') + + +def parse_ifacialmocap_v2_pose(ifacialmocap_output): + output = {} + parts = ifacialmocap_output.split("|") + for part in parts: + part = part.strip() + if len(part) == 0: + continue + if "&" in part: + components = part.split("&") + assert len(components) == 2 + key = components[0] + value = float(components[1]) / 100.0 + if key.endswith("_L"): + key = key[:-2] + "Left" + elif key.endswith("_R"): + key = key[:-2] + "Right" + if key in BLENDSHAPE_NAMES: + output[key] = value + elif part.startswith("=head#"): + components = part[len("=head#"):].split(",") + assert len(components) == 6 + output[HEAD_BONE_X] = float(components[0]) * math.pi / 180 + output[HEAD_BONE_Y] = float(components[1]) * math.pi / 180 + output[HEAD_BONE_Z] = float(components[2]) * math.pi / 180 + elif part.startswith("rightEye#"): + components = part[len("rightEye#"):].split(",") + output[RIGHT_EYE_BONE_X] = float(components[0]) * math.pi / 180 + output[RIGHT_EYE_BONE_Y] = float(components[1]) * math.pi / 180 + output[RIGHT_EYE_BONE_Z] = float(components[2]) * math.pi / 180 + elif part.startswith("leftEye#"): + components = part[len("leftEye#"):].split(",") + output[LEFT_EYE_BONE_X] = float(components[0]) * math.pi / 180 + output[LEFT_EYE_BONE_Y] = float(components[1]) * math.pi / 180 + output[LEFT_EYE_BONE_Z] = float(components[2]) * math.pi / 180 + output[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + output[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + output[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + return output + + +def parse_ifacialmocap_v1_pose(ifacialmocap_output): + output = {} + parts = ifacialmocap_output.split("|") + for part in parts: + part = part.strip() + if len(part) == 0: + continue + if part.startswith("=head#"): + components = part[len("=head#"):].split(",") + assert len(components) == 6 + output[HEAD_BONE_X] = float(components[0]) * math.pi / 180 + output[HEAD_BONE_Y] = float(components[1]) * math.pi / 180 + output[HEAD_BONE_Z] = float(components[2]) * math.pi / 180 + elif part.startswith("rightEye#"): + components = part[len("rightEye#"):].split(",") + output[RIGHT_EYE_BONE_X] = float(components[0]) * math.pi / 180 + output[RIGHT_EYE_BONE_Y] = float(components[1]) * math.pi / 180 + output[RIGHT_EYE_BONE_Z] = float(components[2]) * math.pi / 180 + elif part.startswith("leftEye#"): + components = part[len("leftEye#"):].split(",") + output[LEFT_EYE_BONE_X] = float(components[0]) * math.pi / 180 + output[LEFT_EYE_BONE_Y] = float(components[1]) * math.pi / 180 + output[LEFT_EYE_BONE_Z] = float(components[2]) * math.pi / 180 + else: + components = part.split("-") + assert len(components) == 2 + key = components[0] + value = float(components[1]) / 100.0 + if key.endswith("_L"): + key = key[:-2] + "Left" + elif key.endswith("_R"): + key = key[:-2] + "Right" + if key in BLENDSHAPE_NAMES: + output[key] = value + output[HEAD_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + output[LEFT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + output[RIGHT_EYE_BONE_QUAT] = [0.0, 0.0, 0.0, 1.0] + return output + diff --git a/tha4/src/tha4/mocap/mediapipe_constants.py b/tha4/src/tha4/mocap/mediapipe_constants.py new file mode 100644 index 0000000..0fc61ee --- /dev/null +++ b/tha4/src/tha4/mocap/mediapipe_constants.py @@ -0,0 +1,216 @@ +EYE_LOOK_IN_LEFT = "eyeLookInLeft" +EYE_LOOK_OUT_LEFT = "eyeLookOutLeft" +EYE_LOOK_DOWN_LEFT = "eyeLookDownLeft" +EYE_LOOK_UP_LEFT = "eyeLookUpLeft" +EYE_BLINK_LEFT = "eyeBlinkLeft" +EYE_SQUINT_LEFT = "eyeSquintLeft" +EYE_WIDE_LEFT = "eyeWideLeft" +EYE_LOOK_IN_RIGHT = "eyeLookInRight" +EYE_LOOK_OUT_RIGHT = "eyeLookOutRight" +EYE_LOOK_DOWN_RIGHT = "eyeLookDownRight" +EYE_LOOK_UP_RIGHT = "eyeLookUpRight" +EYE_BLINK_RIGHT = "eyeBlinkRight" +EYE_SQUINT_RIGHT = "eyeSquintRight" +EYE_WIDE_RIGHT = "eyeWideRight" +BROW_DOWN_LEFT = "browDownLeft" +BROW_OUTER_UP_LEFT = "browOuterUpLeft" +BROW_DOWN_RIGHT = "browDownRight" +BROW_OUTER_UP_RIGHT = "browOuterUpRight" +BROW_INNER_UP = "browInnerUp" +NOSE_SNEER_LEFT = "noseSneerLeft" +NOSE_SNEER_RIGHT = "noseSneerRight" +CHEEK_SQUINT_LEFT = "cheekSquintLeft" +CHEEK_SQUINT_RIGHT = "cheekSquintRight" +CHEEK_PUFF = "cheekPuff" +MOUTH_LEFT = "mouthLeft" +MOUTH_DIMPLE_LEFT = "mouthDimpleLeft" +MOUTH_FROWN_LEFT = "mouthFrownLeft" +MOUTH_LOWER_DOWN_LEFT = "mouthLowerDownLeft" +MOUTH_PRESS_LEFT = "mouthPressLeft" +MOUTH_SMILE_LEFT = "mouthSmileLeft" +MOUTH_STRETCH_LEFT = "mouthStretchLeft" +MOUTH_UPPER_UP_LEFT = "mouthUpperUpLeft" +MOUTH_RIGHT = "mouthRight" +MOUTH_DIMPLE_RIGHT = "mouthDimpleRight" +MOUTH_FROWN_RIGHT = "mouthFrownRight" +MOUTH_LOWER_DOWN_RIGHT = "mouthLowerDownRight" +MOUTH_PRESS_RIGHT = "mouthPressRight" +MOUTH_SMILE_RIGHT = "mouthSmileRight" +MOUTH_STRETCH_RIGHT = "mouthStretchRight" +MOUTH_UPPER_UP_RIGHT = "mouthUpperUpRight" +MOUTH_CLOSE = "mouthClose" +MOUTH_FUNNEL = "mouthFunnel" +MOUTH_PUCKER = "mouthPucker" +MOUTH_ROLL_LOWER = "mouthRollLower" +MOUTH_ROLL_UPPER = "mouthRollUpper" +MOUTH_SHRUG_LOWER = "mouthShrugLower" +MOUTH_SHRUG_UPPER = "mouthShrugUpper" +JAW_LEFT = "jawLeft" +JAW_RIGHT = "jawRight" +JAW_FORWARD = "jawForward" +JAW_OPEN = "jawOpen" +NEUTRAL = "_neutral" + +BLENDSHAPE_NAMES = [ + EYE_LOOK_IN_LEFT, # 0 + EYE_LOOK_OUT_LEFT, # 1 + EYE_LOOK_DOWN_LEFT, # 2 + EYE_LOOK_UP_LEFT, # 3 + EYE_BLINK_LEFT, # 4 + EYE_SQUINT_LEFT, # 5 + EYE_WIDE_LEFT, # 6 + EYE_LOOK_IN_RIGHT, # 7 + EYE_LOOK_OUT_RIGHT, # 8 + EYE_LOOK_DOWN_RIGHT, # 9 + EYE_LOOK_UP_RIGHT, # 10 + EYE_BLINK_RIGHT, # 11 + EYE_SQUINT_RIGHT, # 12 + EYE_WIDE_RIGHT, # 13 + BROW_DOWN_LEFT, # 14 + BROW_OUTER_UP_LEFT, # 15 + BROW_DOWN_RIGHT, # 16 + BROW_OUTER_UP_RIGHT, # 17 + BROW_INNER_UP, # 18 + NOSE_SNEER_LEFT, # 19 + NOSE_SNEER_RIGHT, # 20 + CHEEK_SQUINT_LEFT, # 21 + CHEEK_SQUINT_RIGHT, # 22 + CHEEK_PUFF, # 23 + MOUTH_LEFT, # 24 + MOUTH_DIMPLE_LEFT, # 25 + MOUTH_FROWN_LEFT, # 26 + MOUTH_LOWER_DOWN_LEFT, # 27 + MOUTH_PRESS_LEFT, # 28 + MOUTH_SMILE_LEFT, # 29 + MOUTH_STRETCH_LEFT, # 30 + MOUTH_UPPER_UP_LEFT, # 31 + MOUTH_RIGHT, # 32 + MOUTH_DIMPLE_RIGHT, # 33 + MOUTH_FROWN_RIGHT, # 34 + MOUTH_LOWER_DOWN_RIGHT, # 35 + MOUTH_PRESS_RIGHT, # 36 + MOUTH_SMILE_RIGHT, # 37 + MOUTH_STRETCH_RIGHT, # 38 + MOUTH_UPPER_UP_RIGHT, # 39 + MOUTH_CLOSE, # 40 + MOUTH_FUNNEL, # 41 + MOUTH_PUCKER, # 42 + MOUTH_ROLL_LOWER, # 43 + MOUTH_ROLL_UPPER, # 44 + MOUTH_SHRUG_LOWER, # 45 + MOUTH_SHRUG_UPPER, # 46 + JAW_LEFT, # 47 + JAW_RIGHT, # 48 + JAW_FORWARD, # 49 + JAW_OPEN, # 50 + NEUTRAL, # 51 +] + +EYE_LEFT_BLENDSHAPES = [ + EYE_LOOK_IN_LEFT, # 0 + EYE_LOOK_OUT_LEFT, # 1 + EYE_LOOK_DOWN_LEFT, # 2 + EYE_LOOK_UP_LEFT, # 3 + EYE_BLINK_LEFT, # 4 + EYE_SQUINT_LEFT, # 5 + EYE_WIDE_LEFT, # 6 +] + +EYE_RIGHT_BLENDSHAPES = [ + EYE_LOOK_IN_RIGHT, # 7 + EYE_LOOK_OUT_RIGHT, # 8 + EYE_LOOK_DOWN_RIGHT, # 9 + EYE_LOOK_UP_RIGHT, # 10 + EYE_BLINK_RIGHT, # 11 + EYE_SQUINT_RIGHT, # 12 + EYE_WIDE_RIGHT, # 13 +] + +BROW_LEFT_BLENDSHAPES = [ + BROW_DOWN_LEFT, # 14 + BROW_OUTER_UP_LEFT, # 15 + +] + +BROW_RIGHT_BLENDSHAPES = [ + BROW_DOWN_RIGHT, # 16 + BROW_OUTER_UP_RIGHT, # 17 + +] + +BROW_BOTH_BLENDSHAPES = [ + BROW_INNER_UP, # 18 +] + +NOSE_BLENDSHAPES = [ + NOSE_SNEER_LEFT, # 19 + NOSE_SNEER_RIGHT, # 20 +] + +CHECK_BLENDSHAPES = [ + CHEEK_SQUINT_LEFT, # 21 + CHEEK_SQUINT_RIGHT, # 22 + CHEEK_PUFF, # 23 +] + +MOUTH_LEFT_BLENDSHAPES = [ + MOUTH_LEFT, # 24 + MOUTH_DIMPLE_LEFT, # 25 + MOUTH_FROWN_LEFT, # 26 + MOUTH_LOWER_DOWN_LEFT, # 27 + MOUTH_PRESS_LEFT, # 28 + MOUTH_SMILE_LEFT, # 29 + MOUTH_STRETCH_LEFT, # 30 + MOUTH_UPPER_UP_LEFT, # 31 +] + +MOUTH_RIGHT_BLENDSHAPES = [ + MOUTH_RIGHT, # 32 + MOUTH_DIMPLE_RIGHT, # 33 + MOUTH_FROWN_RIGHT, # 34 + MOUTH_LOWER_DOWN_RIGHT, # 35 + MOUTH_PRESS_RIGHT, # 36 + MOUTH_SMILE_RIGHT, # 37 + MOUTH_STRETCH_RIGHT, # 38 + MOUTH_UPPER_UP_RIGHT, # 39 +] + +MOUTH_BOTH_BLENDSHAPES = [ + MOUTH_CLOSE, # 40 + MOUTH_FUNNEL, # 41 + MOUTH_PUCKER, # 42 + MOUTH_ROLL_LOWER, # 43 + MOUTH_ROLL_UPPER, # 44 + MOUTH_SHRUG_LOWER, # 45 + MOUTH_SHRUG_UPPER, # 46 +] + +JAW_BLENDSHAPES = [ + JAW_LEFT, # 47 + JAW_RIGHT, # 48 + JAW_FORWARD, # 49 + JAW_OPEN, # 50 +] + +NEUTRAL_BLENDSHAPES = [ + NEUTRAL, # 51 +] + +COLUMN_0_BLENDSHAPES = EYE_RIGHT_BLENDSHAPES + BROW_RIGHT_BLENDSHAPES + [NOSE_SNEER_RIGHT, CHEEK_SQUINT_RIGHT] +COLUMN_1_BLENDSHAPES = EYE_LEFT_BLENDSHAPES + BROW_LEFT_BLENDSHAPES + [NOSE_SNEER_LEFT, CHEEK_SQUINT_LEFT] +COLUMN_2_BLENDSHAPES = MOUTH_RIGHT_BLENDSHAPES + [JAW_RIGHT] +COLUMN_3_BLENDSHAPES = MOUTH_LEFT_BLENDSHAPES + [JAW_LEFT] +COLUMN_4_BLENDSHAPES = [BROW_INNER_UP, CHEEK_PUFF] + MOUTH_BOTH_BLENDSHAPES + [JAW_FORWARD, JAW_OPEN, NEUTRAL] + +BLENDSHAPE_COLUMNS = [ + COLUMN_0_BLENDSHAPES, + COLUMN_1_BLENDSHAPES, + COLUMN_2_BLENDSHAPES, + COLUMN_3_BLENDSHAPES, + COLUMN_4_BLENDSHAPES, +] + +HEAD_X = "headX" +HEAD_Y = "headY" +HEAD_Z = "headZ" +HEAD_ROTATIONS = [HEAD_X, HEAD_Y, HEAD_Z] \ No newline at end of file diff --git a/tha4/src/tha4/mocap/mediapipe_face_pose.py b/tha4/src/tha4/mocap/mediapipe_face_pose.py new file mode 100644 index 0000000..305b4d6 --- /dev/null +++ b/tha4/src/tha4/mocap/mediapipe_face_pose.py @@ -0,0 +1,41 @@ +import json +import os +from typing import Optional, Dict + +import numpy + + +class MediaPipeFacePose: + KEY_BLENDSHAPE_PARAMS = "blendshape_params" + KEY_XFORM_MATRIX = "xform_matrix" + + def __init__(self, blendshape_params: Optional[Dict[str, float]], xform_matrix: Optional[numpy.ndarray]): + if blendshape_params is None: + blendshape_params = {} + if xform_matrix is None: + self.xform_matrix = numpy.zeros(4, 4) + for i in range(4): + self.xform_matrix[i, i] = 1.0 + + self.blendshape_params = blendshape_params + self.xform_matrix = xform_matrix + + def get_json(self): + return { + MediaPipeFacePose.KEY_BLENDSHAPE_PARAMS: self.blendshape_params.copy(), + MediaPipeFacePose.KEY_XFORM_MATRIX: self.xform_matrix.tolist() + } + + def save(self, file_name: str): + os.makedirs(os.path.dirname(file_name), exist_ok=True) + with open(file_name, "wt") as fout: + fout.write(json.dumps(self.get_json())) + + @staticmethod + def load(file_name: str): + with open(file_name, "rt") as fin: + s = fin.read() + json_data = json.loads(s) + return MediaPipeFacePose( + json_data[MediaPipeFacePose.KEY_BLENDSHAPE_PARAMS], + xform_matrix = numpy.array(json_data[MediaPipeFacePose.KEY_XFORM_MATRIX])) diff --git a/tha4/src/tha4/mocap/mediapipe_face_pose_converter.py b/tha4/src/tha4/mocap/mediapipe_face_pose_converter.py new file mode 100644 index 0000000..0cdc389 --- /dev/null +++ b/tha4/src/tha4/mocap/mediapipe_face_pose_converter.py @@ -0,0 +1,17 @@ +from abc import ABC, abstractmethod +from typing import List, Callable, Optional + +from tha4.mocap.mediapipe_face_pose import MediaPipeFacePose + + +class MediaPipeFacePoseConverter(ABC): + @abstractmethod + def convert(self, mediapipe_face_pose: MediaPipeFacePose) -> List[float]: + pass + + @abstractmethod + def init_pose_converter_panel( + self, + parent, + current_pose_supplier: Callable[[], Optional[MediaPipeFacePose]]): + pass \ No newline at end of file diff --git a/tha4/src/tha4/mocap/mediapipe_face_pose_converter_00.py b/tha4/src/tha4/mocap/mediapipe_face_pose_converter_00.py new file mode 100644 index 0000000..c2b6708 --- /dev/null +++ b/tha4/src/tha4/mocap/mediapipe_face_pose_converter_00.py @@ -0,0 +1,662 @@ +import math +import time +from enum import Enum +from typing import Optional, List, Callable + +import numpy +import scipy.optimize +import wx +from scipy.spatial.transform import Rotation + +from tha4.poser.modes.pose_parameters import get_pose_parameters +from tha4.mocap.mediapipe_constants import MOUTH_SMILE_LEFT, MOUTH_SHRUG_UPPER, MOUTH_SMILE_RIGHT, \ + BROW_INNER_UP, BROW_OUTER_UP_RIGHT, BROW_OUTER_UP_LEFT, BROW_DOWN_LEFT, BROW_DOWN_RIGHT, EYE_WIDE_LEFT, \ + EYE_WIDE_RIGHT, EYE_BLINK_LEFT, EYE_BLINK_RIGHT, CHEEK_SQUINT_LEFT, CHEEK_SQUINT_RIGHT, EYE_LOOK_IN_LEFT, \ + EYE_LOOK_OUT_LEFT, EYE_LOOK_IN_RIGHT, EYE_LOOK_OUT_RIGHT, EYE_LOOK_UP_LEFT, EYE_LOOK_UP_RIGHT, EYE_LOOK_DOWN_RIGHT, \ + EYE_LOOK_DOWN_LEFT, JAW_OPEN, MOUTH_FROWN_LEFT, MOUTH_FROWN_RIGHT, \ + MOUTH_LOWER_DOWN_LEFT, MOUTH_LOWER_DOWN_RIGHT, MOUTH_FUNNEL, MOUTH_PUCKER +from tha4.mocap.mediapipe_face_pose import MediaPipeFacePose +from tha4.mocap.mediapipe_face_pose_converter import MediaPipeFacePoseConverter + + +class EyebrowDownMode(Enum): + TROUBLED = 1 + ANGRY = 2 + LOWERED = 3 + SERIOUS = 4 + + +class WinkMode(Enum): + NORMAL = 1 + RELAXED = 2 + + +def rad_to_deg(rad): + return rad * 180.0 / math.pi + + +def deg_to_rad(deg): + return deg * math.pi / 180.0 + + +def clamp(x, min_value, max_value): + return max(min_value, min(max_value, x)) + + +class MediaPipeFacePoseConverter00Args: + def __init__(self, + smile_threshold_min: float = 0.4, + smile_threshold_max: float = 0.6, + eyebrow_down_mode: EyebrowDownMode = EyebrowDownMode.ANGRY, + wink_mode: WinkMode = WinkMode.NORMAL, + eye_surprised_max: float = 0.5, + eye_blink_max: float = 0.8, + eyebrow_down_max: float = 0.4, + cheek_squint_min: float = 0.1, + cheek_squint_max: float = 0.7, + eye_rotation_factor: float = 1.0 / 0.75, + jaw_open_min: float = 0.1, + jaw_open_max: float = 0.4, + mouth_frown_max: float = 0.6, + mouth_funnel_min: float = 0.25, + mouth_funnel_max: float = 0.5, + iris_small_left=0.0, + iris_small_right=0.0, + head_x_offset=0.0, + head_y_offset=0.0, + head_z_offset=0.0): + self.iris_small_right = iris_small_left + self.iris_small_left = iris_small_right + + self.wink_mode = wink_mode + + self.mouth_funnel_max = mouth_funnel_max + self.mouth_funnel_min = mouth_funnel_min + self.mouth_frown_max = mouth_frown_max + + self.jaw_open_max = jaw_open_max + self.jaw_open_min = jaw_open_min + + self.eye_rotation_factor = eye_rotation_factor + + self.cheek_squint_max = cheek_squint_max + self.cheek_squint_min = cheek_squint_min + + self.eyebrow_down_max = eyebrow_down_max + + self.eye_blink_max = eye_blink_max + self.eye_surprised_max = eye_surprised_max + + self.smile_threshold_min = smile_threshold_min + self.smile_threshold_max = smile_threshold_max + + self.head_z_offset = head_z_offset + self.head_y_offset = head_y_offset + self.head_x_offset = head_x_offset + + self.eyebrow_down_mode = eyebrow_down_mode + + def set_smile_threshold_min(self, new_value: float): + self.smile_threshold_min = new_value + + def set_smile_threshold_max(self, new_value: float): + self.smile_threshold_max = new_value + + def set_eye_surprised_max(self, new_value: float): + self.eye_surprised_max = new_value + + def set_eye_blink_max(self, new_value: float): + self.eye_blink_max = new_value + + def set_eyebrow_down_max(self, new_value: float): + self.eyebrow_down_max = new_value + + def set_cheek_squint_min(self, new_value: float): + self.cheek_squint_min = new_value + + def set_cheek_squint_max(self, new_value: float): + self.cheek_squint_max = new_value + + def set_jaw_open_min(self, new_value: float): + self.jaw_open_min = new_value + + def set_jaw_open_max(self, new_value: float): + self.jaw_open_max = new_value + + def set_mouth_frown_max(self, new_value: float): + self.mouth_frown_max = new_value + + def set_mouth_funnel_min(self, new_value: float): + self.mouth_funnel_min = new_value + + def set_mouth_funnel_max(self, new_value: float): + self.mouth_funnel_min = new_value + + +class MediaPoseFacePoseConverter00(MediaPipeFacePoseConverter): + def __init__(self, args: Optional[MediaPipeFacePoseConverter00Args] = None): + super().__init__() + if args is None: + args = MediaPipeFacePoseConverter00Args() + self.args = args + pose_parameters = get_pose_parameters() + self.pose_size = 45 + + self.eyebrow_troubled_left_index = pose_parameters.get_parameter_index("eyebrow_troubled_left") + self.eyebrow_troubled_right_index = pose_parameters.get_parameter_index("eyebrow_troubled_right") + self.eyebrow_angry_left_index = pose_parameters.get_parameter_index("eyebrow_angry_left") + self.eyebrow_angry_right_index = pose_parameters.get_parameter_index("eyebrow_angry_right") + self.eyebrow_happy_left_index = pose_parameters.get_parameter_index("eyebrow_happy_left") + self.eyebrow_happy_right_index = pose_parameters.get_parameter_index("eyebrow_happy_right") + self.eyebrow_raised_left_index = pose_parameters.get_parameter_index("eyebrow_raised_left") + self.eyebrow_raised_right_index = pose_parameters.get_parameter_index("eyebrow_raised_right") + self.eyebrow_lowered_left_index = pose_parameters.get_parameter_index("eyebrow_lowered_left") + self.eyebrow_lowered_right_index = pose_parameters.get_parameter_index("eyebrow_lowered_right") + self.eyebrow_serious_left_index = pose_parameters.get_parameter_index("eyebrow_serious_left") + self.eyebrow_serious_right_index = pose_parameters.get_parameter_index("eyebrow_serious_right") + + self.eye_surprised_left_index = pose_parameters.get_parameter_index("eye_surprised_left") + self.eye_surprised_right_index = pose_parameters.get_parameter_index("eye_surprised_right") + self.eye_wink_left_index = pose_parameters.get_parameter_index("eye_wink_left") + self.eye_wink_right_index = pose_parameters.get_parameter_index("eye_wink_right") + self.eye_happy_wink_left_index = pose_parameters.get_parameter_index("eye_happy_wink_left") + self.eye_happy_wink_right_index = pose_parameters.get_parameter_index("eye_happy_wink_right") + self.eye_relaxed_left_index = pose_parameters.get_parameter_index("eye_relaxed_left") + self.eye_relaxed_right_index = pose_parameters.get_parameter_index("eye_relaxed_right") + self.eye_raised_lower_eyelid_left_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_left") + self.eye_raised_lower_eyelid_right_index = pose_parameters.get_parameter_index("eye_raised_lower_eyelid_right") + + self.iris_small_left_index = pose_parameters.get_parameter_index("iris_small_left") + self.iris_small_right_index = pose_parameters.get_parameter_index("iris_small_right") + + self.iris_rotation_x_index = pose_parameters.get_parameter_index("iris_rotation_x") + self.iris_rotation_y_index = pose_parameters.get_parameter_index("iris_rotation_y") + + self.head_x_index = pose_parameters.get_parameter_index("head_x") + self.head_y_index = pose_parameters.get_parameter_index("head_y") + self.neck_z_index = pose_parameters.get_parameter_index("neck_z") + + self.mouth_aaa_index = pose_parameters.get_parameter_index("mouth_aaa") + self.mouth_iii_index = pose_parameters.get_parameter_index("mouth_iii") + self.mouth_uuu_index = pose_parameters.get_parameter_index("mouth_uuu") + self.mouth_eee_index = pose_parameters.get_parameter_index("mouth_eee") + self.mouth_ooo_index = pose_parameters.get_parameter_index("mouth_ooo") + + self.mouth_lowered_corner_left_index = pose_parameters.get_parameter_index("mouth_lowered_corner_left") + self.mouth_lowered_corner_right_index = pose_parameters.get_parameter_index("mouth_lowered_corner_right") + self.mouth_raised_corner_left_index = pose_parameters.get_parameter_index("mouth_raised_corner_left") + self.mouth_raised_corner_right_index = pose_parameters.get_parameter_index("mouth_raised_corner_right") + + self.body_y_index = pose_parameters.get_parameter_index("body_y") + self.body_z_index = pose_parameters.get_parameter_index("body_z") + self.breathing_index = pose_parameters.get_parameter_index("breathing") + + self.breathing_start_time = time.time() + + self.panel = None + self.current_pose_supplier = None + + def init_pose_converter_panel( + self, + parent, + current_pose_supplier: Callable[[], Optional[MediaPipeFacePose]]): + self.panel = wx.Panel(parent, style=wx.SIMPLE_BORDER) + self.panel_sizer = wx.BoxSizer(wx.VERTICAL) + self.panel.SetSizer(self.panel_sizer) + self.panel.SetAutoLayout(1) + parent.GetSizer().Add(self.panel, 0, wx.EXPAND) + + self.current_pose_supplier = current_pose_supplier + + if True: + eyebrow_down_mode_text = wx.StaticText(self.panel, label=" --- Eyebrow Down Mode --- ", + style=wx.ALIGN_CENTER) + self.panel_sizer.Add(eyebrow_down_mode_text, 0, wx.EXPAND) + + self.eyebrow_down_mode_choice = wx.Choice( + self.panel, + choices=[ + "ANGRY", + "TROUBLED", + "SERIOUS", + "LOWERED", + ]) + self.eyebrow_down_mode_choice.SetSelection(0) + self.panel_sizer.Add(self.eyebrow_down_mode_choice, 0, wx.EXPAND) + self.eyebrow_down_mode_choice.Bind(wx.EVT_CHOICE, self.change_eyebrow_down_mode) + + if True: + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + wink_mode_text = wx.StaticText(self.panel, label=" --- Wink Mode --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(wink_mode_text, 0, wx.EXPAND) + + self.wink_mode_choice = wx.Choice( + self.panel, + choices=[ + "NORMAL", + "RELAXED", + ]) + self.wink_mode_choice.SetSelection(0) + self.panel_sizer.Add(self.wink_mode_choice, 0, wx.EXPAND) + self.wink_mode_choice.Bind(wx.EVT_CHOICE, self.change_wink_mode) + + if True: + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + iris_size_text = wx.StaticText(self.panel, label=" --- Iris Size --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(iris_size_text, 0, wx.EXPAND) + + self.iris_left_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.iris_left_slider, 0, wx.EXPAND) + self.iris_left_slider.Bind(wx.EVT_SLIDER, self.change_iris_size) + + self.iris_right_slider = wx.Slider(self.panel, minValue=0, maxValue=1000, value=0, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.iris_right_slider, 0, wx.EXPAND) + self.iris_right_slider.Bind(wx.EVT_SLIDER, self.change_iris_size) + self.iris_right_slider.Enable(False) + + self.link_left_right_irises = wx.CheckBox( + self.panel, label="Use same value for both sides") + self.link_left_right_irises.SetValue(True) + self.panel_sizer.Add(self.link_left_right_irises, wx.SizerFlags().CenterHorizontal().Border()) + self.link_left_right_irises.Bind(wx.EVT_CHECKBOX, self.link_left_right_irises_clicked) + + if True: + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + breathing_frequency_text = wx.StaticText( + self.panel, label=" --- Breathing --- ", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(breathing_frequency_text, 0, wx.EXPAND) + + self.restart_breathing_cycle_button = wx.Button(self.panel, label="Restart Breathing Cycle") + self.restart_breathing_cycle_button.Bind(wx.EVT_BUTTON, self.restart_breathing_cycle_clicked) + self.panel_sizer.Add(self.restart_breathing_cycle_button, 0, wx.EXPAND) + + self.breathing_frequency_slider = wx.Slider( + self.panel, minValue=0, maxValue=60, value=20, style=wx.HORIZONTAL) + self.panel_sizer.Add(self.breathing_frequency_slider, 0, wx.EXPAND) + + self.breathing_gauge = wx.Gauge(self.panel, style=wx.GA_HORIZONTAL, range=1000) + self.panel_sizer.Add(self.breathing_gauge, 0, wx.EXPAND) + + if True: + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + face_orientation_text = wx.StaticText( + self.panel, label="--- Face Orientation ---", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(face_orientation_text, 0, wx.EXPAND) + + self.calibrate_face_orientation_button = wx.Button(self.panel, label="Calibrate (I'm looking forward)") + self.calibrate_face_orientation_button.Bind(wx.EVT_BUTTON, self.calibrate_face_orientation_clicked) + self.panel_sizer.Add(self.calibrate_face_orientation_button, 0, wx.EXPAND) + + if True: + separator = wx.StaticLine(self.panel, -1, size=(256, 5)) + self.panel_sizer.Add(separator, 0, wx.EXPAND) + + convertion_parameters_text = wx.StaticText( + self.panel, label="--- Conversion Parameters ---", style=wx.ALIGN_CENTER) + self.panel_sizer.Add(convertion_parameters_text, 0, wx.EXPAND) + + conversion_param_panel = wx.Panel(self.panel) + self.panel_sizer.Add(conversion_param_panel, 0, wx.EXPAND) + conversion_panel_sizer = wx.FlexGridSizer(cols=2) + conversion_panel_sizer.AddGrowableCol(1) + conversion_param_panel.SetSizer(conversion_panel_sizer) + conversion_param_panel.SetAutoLayout(1) + + self.smile_thresold_min_spin = self.create_spin_control( + conversion_param_panel, + "Smile Threshold Min:", self.args.smile_threshold_min, self.args.set_smile_threshold_min) + self.smile_thresold_max_spin = self.create_spin_control( + conversion_param_panel, + "Smile Threshold Max:", self.args.smile_threshold_max, self.args.set_smile_threshold_max) + self.eye_surprised_max_spin = self.create_spin_control( + conversion_param_panel, + "Eye Surprised Max:", self.args.eye_surprised_max, self.args.set_eye_surprised_max) + self.eye_blink_max_spin = self.create_spin_control( + conversion_param_panel, + "Eye Blink Max:", self.args.eye_blink_max, self.args.set_eye_blink_max) + self.eyebrow_down_max_spin = self.create_spin_control( + conversion_param_panel, + "Eyebrow Down Max:", self.args.eyebrow_down_max, self.args.set_eyebrow_down_max) + self.cheek_squint_min_spin = self.create_spin_control( + conversion_param_panel, + "Cheek Squint Min:", self.args.cheek_squint_min, self.args.set_cheek_squint_min) + self.cheek_squint_max_spin = self.create_spin_control( + conversion_param_panel, + "Cheek Squint Max:", self.args.cheek_squint_max, self.args.set_cheek_squint_max) + self.jaw_open_min_spin = self.create_spin_control( + conversion_param_panel, + "Jaw Open Min:", self.args.jaw_open_min, self.args.set_jaw_open_min) + self.jaw_open_max_spin = self.create_spin_control( + conversion_param_panel, + "Jaw Open Max:", self.args.jaw_open_max, self.args.set_jaw_open_max) + self.mouth_frown_max_spin = self.create_spin_control( + conversion_param_panel, + "Mouth Frown Max:", self.args.mouth_frown_max, self.args.set_mouth_frown_max) + self.mouth_funnel_min_spin = self.create_spin_control( + conversion_param_panel, + "Mouth Funnel Min:", self.args.mouth_funnel_min, self.args.set_mouth_funnel_min) + self.mouth_funnel_max_spin = self.create_spin_control( + conversion_param_panel, + "Mouth Funnel Max:", self.args.mouth_funnel_max, self.args.set_mouth_funnel_max) + + self.panel_sizer.Fit(self.panel) + + def create_spin_control(self, parent, label: str, initial_value: float, set_func: Callable[[float], None]): + sizer = parent.GetSizer() + + text = wx.StaticText(parent, label=label) + sizer.Add(text, wx.SizerFlags().Right().Border(wx.ALL, 2)) + + spin_ctrl = wx.SpinCtrlDouble( + parent, + wx.ID_ANY, + min=0.0, + max=1.0, + initial=initial_value, + inc=0.01) + sizer.Add(spin_ctrl, wx.SizerFlags().Border(wx.ALL, 2).Expand()) + + def handler(event: wx.Event): + new_value = spin_ctrl.GetValue() + set_func(new_value) + + spin_ctrl.Bind(wx.EVT_SPINCTRLDOUBLE, handler) + + return spin_ctrl + + def extract_euler_angles(self, mediapipe_face_pose: MediaPipeFacePose): + M = mediapipe_face_pose.xform_matrix[0:3, 0:3] + rot = Rotation.from_matrix(M) + return rot.as_euler('xyz', degrees=False) + + def calibrate_face_orientation_clicked(self, event: wx.Event): + if self.current_pose_supplier is None: + return + + mediapipe_face_pose = self.current_pose_supplier() + if mediapipe_face_pose is None: + return + + euler_angles = self.extract_euler_angles(mediapipe_face_pose) + self.args.head_x_offset = euler_angles[0] + self.args.head_y_offset = euler_angles[1] + self.args.head_z_offset = euler_angles[2] + + def restart_breathing_cycle_clicked(self, event: wx.Event): + self.breathing_start_time = time.time() + + def change_eyebrow_down_mode(self, event: wx.Event): + selected_index = self.eyebrow_down_mode_choice.GetSelection() + if selected_index == 0: + self.args.eyebrow_down_mode = EyebrowDownMode.ANGRY + elif selected_index == 1: + self.args.eyebrow_down_mode = EyebrowDownMode.TROUBLED + elif selected_index == 2: + self.args.eyebrow_down_mode = EyebrowDownMode.SERIOUS + else: + self.args.eyebrow_down_mode = EyebrowDownMode.LOWERED + + def change_wink_mode(self, event: wx.Event): + selected_index = self.wink_mode_choice.GetSelection() + if selected_index == 0: + self.args.wink_mode = WinkMode.NORMAL + else: + self.args.wink_mode = WinkMode.RELAXED + + def change_iris_size(self, event: wx.Event): + if self.link_left_right_irises.GetValue(): + left_value = self.iris_left_slider.GetValue() + right_value = self.iris_right_slider.GetValue() + if left_value != right_value: + self.iris_right_slider.SetValue(left_value) + self.args.iris_small_left = left_value / 1000.0 + self.args.iris_small_right = left_value / 1000.0 + else: + self.args.iris_small_left = self.iris_left_slider.GetValue() / 1000.0 + self.args.iris_small_right = self.iris_right_slider.GetValue() / 1000.0 + + def link_left_right_irises_clicked(self, event: wx.Event): + if self.link_left_right_irises.GetValue(): + self.iris_right_slider.Enable(False) + else: + self.iris_right_slider.Enable(True) + self.change_iris_size(event) + + def decompose_head_body_param(self, param, threshold=2.0 / 3): + if abs(param) < threshold: + return (param, 0.0) + else: + if param < 0: + sign = -1.0 + else: + sign = 1.0 + return (threshold * sign, (abs(param) - threshold) * sign) + + def convert(self, mediapipe_face_pose: MediaPipeFacePose) -> List[float]: + pose = [0.0 for i in range(self.pose_size)] + + blendshape_params = mediapipe_face_pose.blendshape_params + + smile_value = \ + (blendshape_params[MOUTH_SMILE_LEFT] + blendshape_params[MOUTH_SMILE_RIGHT]) / 2.0 \ + + blendshape_params[MOUTH_SHRUG_UPPER] + if self.args.smile_threshold_min >= self.args.smile_threshold_max: + smile_degree = 0.0 + else: + if smile_value < self.args.smile_threshold_min: + smile_degree = 0.0 + elif smile_value > self.args.smile_threshold_max: + smile_degree = 1.0 + else: + smile_degree = (smile_value - self.args.smile_threshold_min) / ( + self.args.smile_threshold_max - self.args.smile_threshold_min) + + # Eyebrow + if True: + brow_inner_up = blendshape_params[BROW_INNER_UP] + brow_outer_up_right = blendshape_params[BROW_OUTER_UP_RIGHT] + brow_outer_up_left = blendshape_params[BROW_OUTER_UP_LEFT] + + brow_up_left = clamp(brow_inner_up + brow_outer_up_left, 0.0, 1.0) + brow_up_right = clamp(brow_inner_up + brow_outer_up_right, 0.0, 1.0) + pose[self.eyebrow_raised_left_index] = brow_up_left + pose[self.eyebrow_raised_right_index] = brow_up_right + + if self.args.eyebrow_down_max <= 0.0: + brow_down_left = 0.0 + brow_down_right = 0.0 + else: + brow_down_left = (1.0 - smile_degree) \ + * clamp(blendshape_params[BROW_DOWN_LEFT] / self.args.eyebrow_down_max, 0.0, 1.0) + brow_down_right = (1.0 - smile_degree) \ + * clamp(blendshape_params[BROW_DOWN_RIGHT] / self.args.eyebrow_down_max, 0.0, 1.0) + + if self.args.eyebrow_down_mode == EyebrowDownMode.TROUBLED: + pose[self.eyebrow_troubled_left_index] = brow_down_left + pose[self.eyebrow_troubled_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.ANGRY: + pose[self.eyebrow_angry_left_index] = brow_down_left + pose[self.eyebrow_angry_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.LOWERED: + pose[self.eyebrow_lowered_left_index] = brow_down_left + pose[self.eyebrow_lowered_right_index] = brow_down_right + elif self.args.eyebrow_down_mode == EyebrowDownMode.SERIOUS: + pose[self.eyebrow_serious_left_index] = brow_down_left + pose[self.eyebrow_serious_right_index] = brow_down_right + + brow_happy_value = clamp(smile_value, 0.0, 1.0) * smile_degree + pose[self.eyebrow_happy_left_index] = brow_happy_value + pose[self.eyebrow_happy_right_index] = brow_happy_value + + # Eye + if True: + # Surprised + if self.args.eye_surprised_max <= 0.0: + pose[self.eye_surprised_left_index] = 0.0 + pose[self.eye_surprised_right_index] = 0.0 + else: + pose[self.eye_surprised_left_index] = clamp( + blendshape_params[EYE_WIDE_LEFT] / self.args.eye_surprised_max, 0.0, 1.0) + pose[self.eye_surprised_right_index] = clamp( + blendshape_params[EYE_WIDE_RIGHT] / self.args.eye_surprised_max, 0.0, 1.0) + + # Wink + if self.args.wink_mode == WinkMode.NORMAL: + wink_left_index = self.eye_wink_left_index + wink_right_index = self.eye_wink_right_index + else: + wink_left_index = self.eye_relaxed_left_index + wink_right_index = self.eye_relaxed_right_index + if self.args.eye_blink_max <= 0: + pose[wink_left_index] = 0.0 + pose[wink_right_index] = 0.0 + pose[self.eye_happy_wink_left_index] = 0.0 + pose[self.eye_happy_wink_right_index] = 0.0 + else: + pose[wink_left_index] = (1.0 - smile_degree) * clamp( + blendshape_params[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0) + pose[wink_right_index] = (1.0 - smile_degree) * clamp( + blendshape_params[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0) + pose[self.eye_happy_wink_left_index] = smile_degree * clamp( + blendshape_params[EYE_BLINK_LEFT] / self.args.eye_blink_max, 0.0, 1.0) + pose[self.eye_happy_wink_right_index] = smile_degree * clamp( + blendshape_params[EYE_BLINK_RIGHT] / self.args.eye_blink_max, 0.0, 1.0) + + # Lower eyelid + cheek_squint_denom = self.args.cheek_squint_max - self.args.cheek_squint_min + if cheek_squint_denom <= 0.0: + pose[self.eye_raised_lower_eyelid_left_index] = 0.0 + pose[self.eye_raised_lower_eyelid_right_index] = 0.0 + else: + pose[self.eye_raised_lower_eyelid_left_index] = \ + clamp( + (blendshape_params[CHEEK_SQUINT_LEFT] - self.args.cheek_squint_min) / cheek_squint_denom, + 0.0, 1.0) + pose[self.eye_raised_lower_eyelid_right_index] = \ + clamp( + (blendshape_params[CHEEK_SQUINT_RIGHT] - self.args.cheek_squint_min) / cheek_squint_denom, + 0.0, 1.0) + + # Iris rotation + if True: + eye_rotation_y = (blendshape_params[EYE_LOOK_IN_LEFT] + - blendshape_params[EYE_LOOK_OUT_LEFT] + - blendshape_params[EYE_LOOK_IN_RIGHT] + + blendshape_params[EYE_LOOK_OUT_RIGHT]) / 2.0 * self.args.eye_rotation_factor + pose[self.iris_rotation_y_index] = clamp(eye_rotation_y, -1.0, 1.0) + + eye_rotation_x = (blendshape_params[EYE_LOOK_UP_LEFT] + + blendshape_params[EYE_LOOK_UP_RIGHT] + - blendshape_params[EYE_LOOK_DOWN_LEFT] + - blendshape_params[EYE_LOOK_DOWN_RIGHT]) / 2.0 * self.args.eye_rotation_factor + pose[self.iris_rotation_x_index] = clamp(eye_rotation_x, -1.0, 1.0) + + # Iris size + if True: + pose[self.iris_small_left_index] = self.args.iris_small_left + pose[self.iris_small_right_index] = self.args.iris_small_right + + # Head rotation + if True: + euler_angles = self.extract_euler_angles(mediapipe_face_pose) + euler_angles[0] -= self.args.head_x_offset + euler_angles[1] -= self.args.head_y_offset + euler_angles[2] -= self.args.head_z_offset + + x_param = clamp(-euler_angles[0] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + pose[self.head_x_index] = x_param + + y_param = clamp(-euler_angles[1] * 180.0 / math.pi, -10.0, 10.0) / 10.0 + pose[self.head_y_index] = y_param + pose[self.body_y_index] = y_param + + z_param = clamp(euler_angles[2] * 180.0 / math.pi, -15.0, 15.0) / 15.0 + pose[self.neck_z_index] = z_param + pose[self.body_z_index] = z_param + + # Mouth + if True: + jaw_open_denom = self.args.jaw_open_max - self.args.jaw_open_min + if jaw_open_denom <= 0: + mouth_open = 0.0 + else: + mouth_open = clamp((blendshape_params[JAW_OPEN] - self.args.jaw_open_min) / jaw_open_denom, 0.0, 1.0) + pose[self.mouth_aaa_index] = mouth_open + pose[self.mouth_raised_corner_left_index] = clamp(smile_value, 0.0, 1.0) + pose[self.mouth_raised_corner_right_index] = clamp(smile_value, 0.0, 1.0) + + is_mouth_open = mouth_open > 0.0 + if not is_mouth_open: + if self.args.mouth_frown_max <= 0: + mouth_frown_value = 0.0 + else: + mouth_frown_value = clamp( + (blendshape_params[MOUTH_FROWN_LEFT] + blendshape_params[ + MOUTH_FROWN_RIGHT]) / self.args.mouth_frown_max, 0.0, 1.0) + pose[self.mouth_lowered_corner_left_index] = mouth_frown_value + pose[self.mouth_lowered_corner_right_index] = mouth_frown_value + else: + mouth_lower_down = clamp( + blendshape_params[MOUTH_LOWER_DOWN_LEFT] + blendshape_params[MOUTH_LOWER_DOWN_RIGHT], 0.0, 1.0) + mouth_funnel = blendshape_params[MOUTH_FUNNEL] + mouth_pucker = blendshape_params[MOUTH_PUCKER] + + mouth_point = [mouth_open, mouth_lower_down, mouth_funnel, mouth_pucker] + + aaa_point = [1.0, 1.0, 0.0, 0.0] + iii_point = [0.0, 1.0, 0.0, 0.0] + uuu_point = [0.5, 0.3, 0.25, 0.75] + ooo_point = [1.0, 0.5, 0.5, 0.4] + + decomp = numpy.array([0, 0, 0, 0]) + M = numpy.array([ + aaa_point, + iii_point, + uuu_point, + ooo_point + ]) + + def loss(decomp): + return numpy.linalg.norm(numpy.matmul(decomp, M) - mouth_point) \ + + 0.01 * numpy.linalg.norm(decomp, ord=1) + + opt_result = scipy.optimize.minimize( + loss, decomp, bounds=[(0.0, 1.0), (0.0, 1.0), (0.0, 1.0), (0.0, 1.0)]) + decomp = opt_result["x"] + restricted_decomp = [decomp.item(0), decomp.item(1), decomp.item(2), decomp.item(3)] + pose[self.mouth_aaa_index] = restricted_decomp[0] + pose[self.mouth_iii_index] = restricted_decomp[1] + mouth_funnel_denom = self.args.mouth_funnel_max - self.args.mouth_funnel_min + if mouth_funnel_denom <= 0: + ooo_alpha = 0.0 + uo_value = 0.0 + else: + ooo_alpha = clamp((mouth_funnel - self.args.mouth_funnel_min) / mouth_funnel_denom, 0.0, 1.0) + uo_value = clamp(restricted_decomp[2] + restricted_decomp[3], 0.0, 1.0) + pose[self.mouth_uuu_index] = uo_value * (1.0 - ooo_alpha) + pose[self.mouth_ooo_index] = uo_value * ooo_alpha + + if self.panel is not None: + frequency = self.breathing_frequency_slider.GetValue() + if frequency == 0: + value = 0.0 + pose[self.breathing_index] = value + self.breathing_start_time = time.time() + else: + period = 60.0 / frequency + now = time.time() + diff = now - self.breathing_start_time + frac = (diff % period) / period + value = (-math.cos(2 * math.pi * frac) + 1.0) / 2.0 + pose[self.breathing_index] = value + self.breathing_gauge.SetValue(int(1000 * value)) + + return pose diff --git a/tha4/src/tha4/nn/__init__.py b/tha4/src/tha4/nn/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/common/__init__.py b/tha4/src/tha4/nn/common/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/common/conv_block_factory.py b/tha4/src/tha4/nn/common/conv_block_factory.py new file mode 100644 index 0000000..d27d813 --- /dev/null +++ b/tha4/src/tha4/nn/common/conv_block_factory.py @@ -0,0 +1,55 @@ +from typing import Optional + +from tha4.nn.conv import create_conv7_block_from_block_args, create_conv3_block_from_block_args, \ + create_downsample_block_from_block_args, create_conv3 +from tha4.nn.resnet_block import ResnetBlock +from tha4.nn.resnet_block_seperable import ResnetBlockSeparable +from tha4.nn.separable_conv import create_separable_conv7_block, create_separable_conv3_block, \ + create_separable_downsample_block, create_separable_conv3 +from tha4.nn.util import BlockArgs + + +class ConvBlockFactory: + def __init__(self, + block_args: BlockArgs, + use_separable_convolution: bool = False): + self.use_separable_convolution = use_separable_convolution + self.block_args = block_args + + def create_conv3(self, + in_channels: int, + out_channels: int, + bias: bool, + initialization_method: Optional[str] = None): + if initialization_method is None: + initialization_method = self.block_args.initialization_method + if self.use_separable_convolution: + return create_separable_conv3( + in_channels, out_channels, bias, initialization_method, self.block_args.use_spectral_norm) + else: + return create_conv3( + in_channels, out_channels, bias, initialization_method, self.block_args.use_spectral_norm) + + def create_conv7_block(self, in_channels: int, out_channels: int): + if self.use_separable_convolution: + return create_separable_conv7_block(in_channels, out_channels, self.block_args) + else: + return create_conv7_block_from_block_args(in_channels, out_channels, self.block_args) + + def create_conv3_block(self, in_channels: int, out_channels: int): + if self.use_separable_convolution: + return create_separable_conv3_block(in_channels, out_channels, self.block_args) + else: + return create_conv3_block_from_block_args(in_channels, out_channels, self.block_args) + + def create_downsample_block(self, in_channels: int, out_channels: int, is_output_1x1: bool): + if self.use_separable_convolution: + return create_separable_downsample_block(in_channels, out_channels, is_output_1x1, self.block_args) + else: + return create_downsample_block_from_block_args(in_channels, out_channels, is_output_1x1) + + def create_resnet_block(self, num_channels: int, is_1x1: bool): + if self.use_separable_convolution: + return ResnetBlockSeparable.create(num_channels, is_1x1, block_args=self.block_args) + else: + return ResnetBlock.create(num_channels, is_1x1, block_args=self.block_args) \ No newline at end of file diff --git a/tha4/src/tha4/nn/common/poser_args.py b/tha4/src/tha4/nn/common/poser_args.py new file mode 100644 index 0000000..ba1c701 --- /dev/null +++ b/tha4/src/tha4/nn/common/poser_args.py @@ -0,0 +1,68 @@ +from typing import Optional + +from torch.nn import Sigmoid, Sequential, Tanh + +from tha4.nn.conv import create_conv3, create_conv3_from_block_args +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs + + +class PoserArgs00: + def __init__(self, + image_size: int, + input_image_channels: int, + output_image_channels: int, + start_channels: int, + num_pose_params: int, + block_args: Optional[BlockArgs] = None): + self.num_pose_params = num_pose_params + self.start_channels = start_channels + self.output_image_channels = output_image_channels + self.input_image_channels = input_image_channels + self.image_size = image_size + if block_args is None: + self.block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)) + else: + self.block_args = block_args + + def create_alpha_block(self): + from torch.nn import Sequential + return Sequential( + create_conv3( + in_channels=self.start_channels, + out_channels=1, + bias=True, + initialization_method=self.block_args.initialization_method, + use_spectral_norm=False), + Sigmoid()) + + def create_all_channel_alpha_block(self): + from torch.nn import Sequential + return Sequential( + create_conv3( + in_channels=self.start_channels, + out_channels=self.output_image_channels, + bias=True, + initialization_method=self.block_args.initialization_method, + use_spectral_norm=False), + Sigmoid()) + + def create_color_change_block(self): + return Sequential( + create_conv3_from_block_args( + in_channels=self.start_channels, + out_channels=self.output_image_channels, + bias=True, + block_args=self.block_args), + Tanh()) + + def create_grid_change_block(self): + return create_conv3( + in_channels=self.start_channels, + out_channels=2, + bias=False, + initialization_method='zero', + use_spectral_norm=False) \ No newline at end of file diff --git a/tha4/src/tha4/nn/common/poser_encoder_decoder_00.py b/tha4/src/tha4/nn/common/poser_encoder_decoder_00.py new file mode 100644 index 0000000..06f7e3d --- /dev/null +++ b/tha4/src/tha4/nn/common/poser_encoder_decoder_00.py @@ -0,0 +1,121 @@ +import math +from typing import Optional, List + +import torch +from torch import Tensor +from torch.nn import ModuleList, Module + +from tha4.nn.common.poser_args import PoserArgs00 +from tha4.nn.conv import create_conv3_block_from_block_args, create_downsample_block_from_block_args, \ + create_upsample_block_from_block_args +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.resnet_block import ResnetBlock +from tha4.nn.util import BlockArgs + + +class PoserEncoderDecoder00Args(PoserArgs00): + def __init__(self, + image_size: int, + input_image_channels: int, + output_image_channels: int, + num_pose_params: int , + start_channels: int, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels: int, + block_args: Optional[BlockArgs] = None): + super().__init__( + image_size, input_image_channels, output_image_channels, start_channels, num_pose_params, block_args) + self.max_channels = max_channels + self.num_bottleneck_blocks = num_bottleneck_blocks + self.bottleneck_image_size = bottleneck_image_size + assert bottleneck_image_size > 1 + + if block_args is None: + self.block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)) + else: + self.block_args = block_args + + +class PoserEncoderDecoder00(Module): + def __init__(self, args: PoserEncoderDecoder00Args): + super().__init__() + self.args = args + + self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1 + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append( + create_conv3_block_from_block_args( + args.input_image_channels, + args.start_channels, + args.block_args)) + current_image_size = args.image_size + current_num_channels = args.start_channels + while current_image_size > args.bottleneck_image_size: + next_image_size = current_image_size // 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.downsample_blocks.append(create_downsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + is_output_1x1=False, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + assert len(self.downsample_blocks) == self.num_levels + + self.bottleneck_blocks = ModuleList() + self.bottleneck_blocks.append(create_conv3_block_from_block_args( + in_channels=current_num_channels + args.num_pose_params, + out_channels=current_num_channels, + block_args=args.block_args)) + for i in range(1, args.num_bottleneck_blocks): + self.bottleneck_blocks.append( + ResnetBlock.create( + num_channels=current_num_channels, + is1x1=False, + block_args=args.block_args)) + + self.upsample_blocks = ModuleList() + while current_image_size < args.image_size: + next_image_size = current_image_size * 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.upsample_blocks.append(create_upsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + + def get_num_output_channels_from_level(self, level: int): + return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level)) + + def get_num_output_channels_from_image_size(self, image_size: int): + return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels) + + def forward(self, image: Tensor, pose: Optional[Tensor] = None) -> List[Tensor]: + if self.args.num_pose_params != 0: + assert pose is not None + else: + assert pose is None + outputs = [] + feature = image + outputs.append(feature) + for block in self.downsample_blocks: + feature = block(feature) + outputs.append(feature) + if pose is not None: + n, c = pose.shape + pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size) + feature = torch.cat([feature, pose], dim=1) + for block in self.bottleneck_blocks: + feature = block(feature) + outputs.append(feature) + for block in self.upsample_blocks: + feature = block(feature) + outputs.append(feature) + outputs.reverse() + return outputs diff --git a/tha4/src/tha4/nn/common/poser_encoder_decoder_00_separable.py b/tha4/src/tha4/nn/common/poser_encoder_decoder_00_separable.py new file mode 100644 index 0000000..52dceb9 --- /dev/null +++ b/tha4/src/tha4/nn/common/poser_encoder_decoder_00_separable.py @@ -0,0 +1,92 @@ +import math +from typing import Optional, List + +import torch +from torch import Tensor +from torch.nn import ModuleList, Module + +from tha4.nn.common.poser_encoder_decoder_00 import PoserEncoderDecoder00Args +from tha4.nn.resnet_block_seperable import ResnetBlockSeparable +from tha4.nn.separable_conv import create_separable_conv3_block, create_separable_downsample_block, \ + create_separable_upsample_block + + +class PoserEncoderDecoder00Separable(Module): + def __init__(self, args: PoserEncoderDecoder00Args): + super().__init__() + self.args = args + + self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1 + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append( + create_separable_conv3_block( + args.input_image_channels, + args.start_channels, + args.block_args)) + current_image_size = args.image_size + current_num_channels = args.start_channels + while current_image_size > args.bottleneck_image_size: + next_image_size = current_image_size // 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.downsample_blocks.append(create_separable_downsample_block( + in_channels=current_num_channels, + out_channels=next_num_channels, + is_output_1x1=False, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + assert len(self.downsample_blocks) == self.num_levels + + self.bottleneck_blocks = ModuleList() + self.bottleneck_blocks.append(create_separable_conv3_block( + in_channels=current_num_channels + args.num_pose_params, + out_channels=current_num_channels, + block_args=args.block_args)) + for i in range(1, args.num_bottleneck_blocks): + self.bottleneck_blocks.append( + ResnetBlockSeparable.create( + num_channels=current_num_channels, + is1x1=False, + block_args=args.block_args)) + + self.upsample_blocks = ModuleList() + while current_image_size < args.image_size: + next_image_size = current_image_size * 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.upsample_blocks.append(create_separable_upsample_block( + in_channels=current_num_channels, + out_channels=next_num_channels, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + + def get_num_output_channels_from_level(self, level: int): + return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level)) + + def get_num_output_channels_from_image_size(self, image_size: int): + return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels) + + def forward(self, image: Tensor, pose: Optional[Tensor] = None) -> List[Tensor]: + if self.args.num_pose_params != 0: + assert pose is not None + else: + assert pose is None + outputs = [] + feature = image + outputs.append(feature) + for block in self.downsample_blocks: + feature = block(feature) + outputs.append(feature) + if pose is not None: + n, c = pose.shape + pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size) + feature = torch.cat([feature, pose], dim=1) + for block in self.bottleneck_blocks: + feature = block(feature) + outputs.append(feature) + for block in self.upsample_blocks: + feature = block(feature) + outputs.append(feature) + outputs.reverse() + return outputs diff --git a/tha4/src/tha4/nn/common/resize_conv_encoder_decoder.py b/tha4/src/tha4/nn/common/resize_conv_encoder_decoder.py new file mode 100644 index 0000000..aaf0149 --- /dev/null +++ b/tha4/src/tha4/nn/common/resize_conv_encoder_decoder.py @@ -0,0 +1,100 @@ +import math +from typing import Optional, List + +import torch +from torch import Tensor +from torch.nn import Module, ModuleList, Sequential, Upsample + +from tha4.nn.common.conv_block_factory import ConvBlockFactory +from tha4.nn.nonlinearity_factory import LeakyReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs + + +class ResizeConvEncoderDecoderArgs: + def __init__(self, + image_size: int, + input_channels: int, + start_channels: int, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels: int, + block_args: Optional[BlockArgs] = None, + upsample_mode: str = 'bilinear', + use_separable_convolution=False): + self.use_separable_convolution = use_separable_convolution + self.upsample_mode = upsample_mode + self.block_args = block_args + self.max_channels = max_channels + self.num_bottleneck_blocks = num_bottleneck_blocks + self.bottleneck_image_size = bottleneck_image_size + self.start_channels = start_channels + self.image_size = image_size + self.input_channels = input_channels + + +class ResizeConvEncoderDecoder(Module): + def __init__(self, args: ResizeConvEncoderDecoderArgs): + super().__init__() + self.args = args + + self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1 + + conv_block_factory = ConvBlockFactory(args.block_args, args.use_separable_convolution) + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append(conv_block_factory.create_conv7_block(args.input_channels, args.start_channels)) + current_image_size = args.image_size + current_num_channels = args.start_channels + while current_image_size > args.bottleneck_image_size: + next_image_size = current_image_size // 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.downsample_blocks.append(conv_block_factory.create_downsample_block( + in_channels=current_num_channels, + out_channels=next_num_channels, + is_output_1x1=False)) + current_image_size = next_image_size + current_num_channels = next_num_channels + assert len(self.downsample_blocks) == self.num_levels + + self.bottleneck_blocks = ModuleList() + for i in range(args.num_bottleneck_blocks): + self.bottleneck_blocks.append(conv_block_factory.create_resnet_block(current_num_channels, is_1x1=False)) + + self.output_image_sizes = [current_image_size] + self.output_num_channels = [current_num_channels] + self.upsample_blocks = ModuleList() + if args.upsample_mode == 'nearest': + align_corners = None + else: + align_corners = False + while current_image_size < args.image_size: + next_image_size = current_image_size * 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.upsample_blocks.append( + Sequential( + Upsample(scale_factor=2, mode=args.upsample_mode, align_corners=align_corners), + conv_block_factory.create_conv3_block( + in_channels=current_num_channels, out_channels=next_num_channels))) + current_image_size = next_image_size + current_num_channels = next_num_channels + self.output_image_sizes.append(current_image_size) + self.output_num_channels.append(current_num_channels) + + def get_num_output_channels_from_level(self, level: int): + return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level)) + + def get_num_output_channels_from_image_size(self, image_size: int): + return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels) + + def forward(self, feature: Tensor) -> List[Tensor]: + outputs = [] + for block in self.downsample_blocks: + feature = block(feature) + for block in self.bottleneck_blocks: + feature = block(feature) + outputs.append(feature) + for block in self.upsample_blocks: + feature = block(feature) + outputs.append(feature) + return outputs \ No newline at end of file diff --git a/tha4/src/tha4/nn/common/resize_conv_unet.py b/tha4/src/tha4/nn/common/resize_conv_unet.py new file mode 100644 index 0000000..ca6ca08 --- /dev/null +++ b/tha4/src/tha4/nn/common/resize_conv_unet.py @@ -0,0 +1,107 @@ +from typing import Optional, List + +import torch +from torch import Tensor +from torch.nn import ModuleList, Module, Upsample + +from tha4.nn.common.conv_block_factory import ConvBlockFactory +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs + + +class ResizeConvUNetArgs: + def __init__(self, + image_size: int, + input_channels: int, + start_channels: int, + bottleneck_image_size: int, + num_bottleneck_blocks: int, + max_channels: int, + upsample_mode: str = 'bilinear', + block_args: Optional[BlockArgs] = None, + use_separable_convolution: bool = False): + if block_args is None: + block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=False)) + + self.use_separable_convolution = use_separable_convolution + self.block_args = block_args + self.upsample_mode = upsample_mode + self.max_channels = max_channels + self.num_bottleneck_blocks = num_bottleneck_blocks + self.bottleneck_image_size = bottleneck_image_size + self.input_channels = input_channels + self.start_channels = start_channels + self.image_size = image_size + + +class ResizeConvUNet(Module): + def __init__(self, args: ResizeConvUNetArgs): + super().__init__() + self.args = args + conv_block_factory = ConvBlockFactory(args.block_args, args.use_separable_convolution) + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append(conv_block_factory.create_conv3_block( + self.args.input_channels, + self.args.start_channels)) + current_channels = self.args.start_channels + current_size = self.args.image_size + + size_to_channel = { + current_size: current_channels + } + while current_size > self.args.bottleneck_image_size: + next_size = current_size // 2 + next_channels = min(self.args.max_channels, current_channels * 2) + self.downsample_blocks.append(conv_block_factory.create_downsample_block( + current_channels, + next_channels, + is_output_1x1=False)) + current_size = next_size + current_channels = next_channels + size_to_channel[current_size] = current_channels + + self.bottleneck_blocks = ModuleList() + for i in range(self.args.num_bottleneck_blocks): + self.bottleneck_blocks.append(conv_block_factory.create_resnet_block(current_channels, is_1x1=False)) + + self.output_image_sizes = [current_size] + self.output_num_channels = [current_channels] + self.upsample_blocks = ModuleList() + while current_size < self.args.image_size: + next_size = current_size * 2 + next_channels = size_to_channel[next_size] + self.upsample_blocks.append(conv_block_factory.create_conv3_block( + current_channels + next_channels, + next_channels)) + current_size = next_size + current_channels = next_channels + self.output_image_sizes.append(current_size) + self.output_num_channels.append(current_channels) + + if args.upsample_mode == 'nearest': + align_corners = None + else: + align_corners = False + self.double_resolution = Upsample(scale_factor=2, mode=args.upsample_mode, align_corners=align_corners) + + def forward(self, feature: Tensor) -> List[Tensor]: + downsampled_features = [] + for block in self.downsample_blocks: + feature = block(feature) + downsampled_features.append(feature) + + for block in self.bottleneck_blocks: + feature = block(feature) + + outputs = [feature] + for i in range(0, len(self.upsample_blocks)): + feature = self.double_resolution(feature) + feature = torch.cat([feature, downsampled_features[-i - 2]], dim=1) + feature = self.upsample_blocks[i](feature) + outputs.append(feature) + + return outputs \ No newline at end of file diff --git a/tha4/src/tha4/nn/common/unet.py b/tha4/src/tha4/nn/common/unet.py new file mode 100644 index 0000000..3049c21 --- /dev/null +++ b/tha4/src/tha4/nn/common/unet.py @@ -0,0 +1,658 @@ +import math +from enum import Enum +from typing import Optional, List + +import torch +from torch import zero_, Tensor +from torch.nn import Module, GroupNorm, Sequential, SiLU, Conv2d, AvgPool2d, Linear, Dropout, ModuleList +from torch.nn.functional import interpolate + +from tha4.shion.core.module_factory import ModuleFactory + + +class Identity(Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x + + +class IdentityFactory(ModuleFactory): + def create(self) -> Module: + return Identity() + + +def init_to_zero(module: Module): + with torch.no_grad(): + zero_(module.weight) + zero_(module.bias) + return module + + +class Upsample(Module): + def __init__(self, in_channels: int, out_channels: Optional[int] = None, use_conv: bool = False): + super().__init__() + if out_channels is None: + out_channels = in_channels + self.in_channels = in_channels + if use_conv or in_channels != out_channels: + self.postprocess = Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + else: + self.postprocess = Identity() + + def forward(self, x): + assert x.shape[1] == self.in_channels + return self.postprocess(interpolate(x, scale_factor=2, mode="nearest")) + + +class Downsample(Module): + def __init__(self, in_channels: int, out_channels: Optional[int] = None, use_conv: bool = False): + super().__init__() + if out_channels is None: + out_channels = in_channels + self.in_channels = in_channels + if use_conv or in_channels != out_channels: + self.op = Conv2d(in_channels, out_channels, kernel_size=3, stride=2, padding=1) + else: + self.op = AvgPool2d(kernel_size=2, stride=2) + + def forward(self, x): + assert x.shape[1] == self.in_channels + return self.op(x) + + +def GroupNorm32(channels): + return GroupNorm(min(32, channels), channels) + + +class SamplingMode(Enum): + SAME_RESOLUTION = 0 + UPSAMPLING = 1 + DOWNSAMPING = 2 + + +class ResBlockArgs: + def __init__(self, + dropout_prob: float, + use_cond0: bool = True, + use_cond1: bool = False, + init_conditioned_residual_to_zero: bool = False, + use_conv_on_skip_connection: bool = False): + assert not use_cond1 or use_cond0 + self.use_conv_on_skip_connection = use_conv_on_skip_connection + self.use_cond1 = use_cond1 + self.use_cond0 = use_cond0 + self.init_conditioned_residual_to_zero = init_conditioned_residual_to_zero + self.dropout_prob = dropout_prob + + +def apply_scaleshift(x: Tensor, scaleshift: Tensor, condition_bias: float = 1.0) -> Tensor: + assert len(scaleshift.shape) == 2 + assert len(x.shape) == 4 + assert x.shape[0] == scaleshift.shape[0] + assert 2 * x.shape[1] == scaleshift.shape[1] + scaleshift = scaleshift.reshape(scaleshift.shape[0], scaleshift.shape[1], 1, 1) + scale, shift = torch.chunk(scaleshift, 2, dim=1) + return x * (condition_bias + scale) + shift + + +class ResBlock(Module): + def __init__(self, + in_channels: int, + out_channels: int, + cond0_channels: Optional[int] = None, + cond1_channels: Optional[int] = None, + sampling_mode: SamplingMode = SamplingMode.SAME_RESOLUTION, + dropout_prob: float = 0.1, + condition_bias: float = 1.0): + super().__init__() + assert cond0_channels is not None or cond1_channels is None + + self.in_channels = in_channels + self.out_channels = out_channels + self.sampling_mode = sampling_mode + self.cond0_channels = cond0_channels + self.cond1_channels = cond1_channels + self.condition_bias = condition_bias + + if sampling_mode == SamplingMode.UPSAMPLING: + self.x_resample = Upsample(in_channels) + self.h_resample = Upsample(in_channels) + elif sampling_mode == SamplingMode.DOWNSAMPING: + self.x_resample = Downsample(in_channels) + self.h_resample = Downsample(in_channels) + else: + self.x_resample = Identity() + self.h_resample = Identity() + + self.nonlinear = SiLU() + + # Layers before conditioning + self.norm0 = GroupNorm32(in_channels) + self.conv0 = Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + + # Conditioning layers + if cond0_channels is not None: + self.cond0_layers = Sequential( + SiLU(), + Linear(cond0_channels, 2 * out_channels)) + self.norm1 = GroupNorm32(out_channels) + self.dropout = Dropout(dropout_prob) + self.conv1 = init_to_zero(Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)) + if cond1_channels is not None: + self.cond1_layers = Sequential( + SiLU(), + Linear(cond0_channels, 2 * out_channels)) + + # Skip layer + if in_channels == out_channels: + self.skip = Identity() + else: + self.skip = Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: Tensor, cond0: Optional[Tensor] = None, cond1: Optional[Tensor] = None) -> Tensor: + assert self.cond0_channels is None or cond0 is not None + assert self.cond1_channels is None or cond1 is not None + + h = self.conv0(self.h_resample(self.nonlinear(self.norm0(x)))) + if self.cond0_channels is not None: + h = self.norm1(h) + h = apply_scaleshift(h, self.cond0_layers(cond0), self.condition_bias) + if self.cond1_channels is not None: + h = apply_scaleshift(h, self.cond1_layers(cond1), self.condition_bias) + h = self.conv1(self.dropout(self.nonlinear(h))) + return self.skip(self.x_resample(x)) + h + + +class AttentionBlockArgs: + def __init__(self, + num_heads: Optional[int] = 1, + num_head_channels: Optional[int] = None, + use_new_attention_order: bool = False): + self.use_new_attention_order = use_new_attention_order + self.num_head_channels = num_head_channels + self.num_heads = num_heads + + +def qkv_attention_legacy(qkv: torch.Tensor, num_heads: int): + assert len(qkv.shape) == 3 + B, W, L = qkv.shape + H = num_heads + assert W % (3 * H) == 0 + C = W // (3 * H) + q, k, v = qkv.reshape(B * H, C * 3, L).split(C, dim=1) + scale = 1.0 / math.sqrt(math.sqrt(C)) + weight = torch.einsum('bct,bcs->bts', q * scale, k * scale) + weight = torch.softmax(weight, dim=-1) + output = torch.einsum("bts,bcs->bct", weight, v) + return output.reshape(B, H * C, L) + + +def qkv_attention(qkv: torch.Tensor, num_heads: int): + B, W, L = qkv.shape + H = num_heads + assert W % (3 * H) == 0 + C = W // (3 * H) + q, k, v = qkv.chunk(3, dim=1) + scale = 1.0 / math.sqrt(math.sqrt(C)) + weight = torch.einsum("bct,bcs->bts", (q * scale).view(B * H, C, L), (k * scale).view(B * H, C, L)) + weight = torch.softmax(weight, dim=-1) + output = torch.einsum("bts,bcs->bct", weight, v.reshape(B * H, C, L)) + return output.reshape(B, H * C, L) + + +class AttentionBlock(Module): + def __init__(self, + num_channels: int, + args: AttentionBlockArgs): + super().__init__() + self.use_new_attention_order = args.use_new_attention_order + + if args.num_head_channels is None: + assert args.num_heads is not None + assert num_channels % args.num_heads == 0 + self.num_heads = args.num_heads + self.num_head_channels = num_channels // self.num_heads + elif args.num_heads is None: + assert args.num_head_channels is not None + assert num_channels % args.num_head_channels == 0 + self.num_heads = num_channels // args.num_head_channels + self.num_head_channels = args.num_head_channels + + self.norm = GroupNorm32(num_channels) + self.qkv = Conv2d(num_channels, 3 * num_channels, kernel_size=1, stride=1, padding=0) + self.conv = Conv2d(num_channels, num_channels, kernel_size=1, stride=1, padding=0) + with torch.no_grad(): + zero_(self.conv.weight) + zero_(self.conv.bias) + + def forward(self, x: torch.Tensor): + assert len(x.shape) == 4 + B, C, H, W = x.shape + qkv = self.qkv(self.norm(x)).reshape(B, 3 * C, H * W) + if self.use_new_attention_order: + h = qkv_attention(qkv, self.num_heads) + else: + h = qkv_attention_legacy(qkv, self.num_heads) + h = self.conv(h.reshape(B, C, H, W)) + return x + h + + +class Arity3To1(Module): + def __init__(self, module: Module): + super().__init__() + self.module = module + + def forward(self, x: Tensor, y: Optional[Tensor] = None, z: Optional[Tensor] = None): + return self.module(x) + + +class DownsamplingBlock(Module): + def __init__(self, + in_channels: int, + out_channels: int, + cond0_channels: Optional[int], + cond1_channels: Optional[int], + num_res_blocks: int, + dropout_prob: float, + use_attention: bool, + perform_downsampling: bool, + resample_with_res_block: bool, + use_conv_to_resample: bool, + attention_block_args: AttentionBlockArgs, + condition_bias: float = 1.0): + super().__init__() + self.use_attention = use_attention + self.res_blocks = ModuleList() + self.attention_blocks = ModuleList() + self.perform_downsampling = perform_downsampling + self.output_channels = [] + for j in range(num_res_blocks): + self.res_blocks.append(ResBlock( + in_channels=in_channels if j == 0 else out_channels, + out_channels=out_channels, + cond0_channels=cond0_channels, + cond1_channels=cond1_channels, + dropout_prob=dropout_prob, + condition_bias=condition_bias)) + if use_attention: + self.attention_blocks.append(AttentionBlock(out_channels, attention_block_args)) + self.output_channels.append(out_channels) + if perform_downsampling: + if resample_with_res_block: + self.downsample = ResBlock( + in_channels=out_channels, + out_channels=out_channels, + cond0_channels=cond0_channels, + cond1_channels=cond1_channels, + dropout_prob=dropout_prob, + sampling_mode=SamplingMode.DOWNSAMPING, + condition_bias=condition_bias) + else: + self.downsample = Arity3To1(Downsample(out_channels, use_conv_to_resample)) + self.output_channels.append(out_channels) + + def forward(self, h: Tensor, cond0: Optional[Tensor] = None, cond1: Optional[Tensor] = None) -> List[Tensor]: + hs = [] + for i in range(len(self.res_blocks)): + h = self.res_blocks[i].forward(h, cond0, cond1) + if self.use_attention: + h = self.attention_blocks[i].forward(h) + hs.append(h) + if self.perform_downsampling: + hs.append(self.downsample(h, cond0, cond1)) + return hs + + +class UpsamplingBlock(Module): + def __init__(self, + in_channels: int, + out_channels: int, + cond0_channels: Optional[int], + cond1_channels: Optional[int], + num_resnet_blocks: int, + skip_channels: List[int], + dropout_prob: float, + use_attention: bool, + perform_upsampling: bool, + resample_with_res_block: bool, + use_conv_to_resample: bool, + attention_block_args: AttentionBlockArgs, + condition_bias: float = 1.0): + super().__init__() + self.use_attention = use_attention + self.resnet_blocks = ModuleList() + self.attention_blocks = ModuleList() + self.perform_upsampling = perform_upsampling + for i in range(num_resnet_blocks): + self.resnet_blocks.append(ResBlock( + in_channels=(in_channels if i == 0 else out_channels) + skip_channels[i], + out_channels=out_channels, + cond0_channels=cond0_channels, + cond1_channels=cond1_channels, + dropout_prob=dropout_prob, + condition_bias=condition_bias)) + if use_attention: + self.attention_blocks.append(AttentionBlock(out_channels, attention_block_args)) + if perform_upsampling: + if resample_with_res_block: + self.upsample = ResBlock( + in_channels=out_channels, + out_channels=out_channels, + cond0_channels=cond0_channels, + cond1_channels=cond1_channels, + sampling_mode=SamplingMode.UPSAMPLING, + dropout_prob=dropout_prob, + condition_bias=condition_bias) + else: + self.upsample = Arity3To1(Upsample(out_channels, use_conv_to_resample)) + + def forward(self, + h: Tensor, + skips: List[Tensor], + cond0: Optional[Tensor] = None, + cond1: Optional[Tensor] = None) -> Tensor: + for i in range(len(self.resnet_blocks)): + h = self.resnet_blocks[i].forward(torch.concat([h, skips[i]], dim=1), cond0, cond1) + if self.use_attention: + h = self.attention_blocks[i].forward(h) + if self.perform_upsampling: + h = self.upsample.forward(h, cond0, cond1) + return h + + +def compute_timestep_embedding(t: Tensor, out_channels: int): + assert len(t.shape) == 2 + b, c = t.shape + assert c == 1 + half_channels = out_channels // 2 + scale = -math.log(10000.0) / (half_channels - 1) + log_times = scale * torch.arange(0, half_channels, device=t.device) + times = torch.exp(log_times).reshape(1, half_channels) * t + t_emb = torch.cat([torch.cos(times), torch.sin(times)], dim=1) + if out_channels % 2 == 1: + t_emb = torch.nn.functional.pad(t_emb, (1, 1), mode='constant') + return t_emb + + +class TimeEmbedding(Module): + def __init__(self, out_channels: int): + super().__init__() + self.out_channels = out_channels + + def forward(self, t: Tensor): + return compute_timestep_embedding(t, self.out_channels) + + +class UnetArgs: + def __init__(self, + in_channels: int = 3, + out_channels: int = 3, + model_channels: int = 64, + level_channel_multipliers: Optional[List[int]] = None, + level_use_attention: Optional[List[bool]] = None, + num_res_blocks_per_level: int = 2, + num_middle_res_blocks: int = 2, + time_embedding_channels: Optional[int] = None, + cond_input_channels: int = 4, + cond_internal_channels: int = 512, + attention_block_args: Optional[AttentionBlockArgs] = None, + dropout_prob: float = 0.1, + resample_with_res_block: bool = True, + use_conv_to_resample=False, + condition_bias: float = 1.0): + assert len(level_channel_multipliers) == len(level_use_attention) + assert not use_conv_to_resample or not resample_with_res_block + + if time_embedding_channels is None: + time_embedding_channels = model_channels + if level_channel_multipliers is None: + level_channel_multipliers = [1, 2, 4, 8] + if level_use_attention is None: + level_use_attention = [False for _ in level_channel_multipliers] + if attention_block_args is None: + attention_block_args = AttentionBlockArgs( + num_heads=1, + num_head_channels=None, + use_new_attention_order=False) + + self.condition_bias = condition_bias + self.use_conv_to_resample = use_conv_to_resample + self.resample_with_res_block = resample_with_res_block + self.cond_internal_channels = cond_internal_channels + self.dropout_prob = dropout_prob + self.attention_block_args = attention_block_args + self.time_embedding_channels = time_embedding_channels + self.num_res_blocks_per_level = num_res_blocks_per_level + self.level_use_attention = level_use_attention + self.level_channel_multipliers = level_channel_multipliers + self.model_channels = model_channels + self.out_channels = out_channels + self.in_channels = in_channels + self.num_levels = len(level_channel_multipliers) + self.num_middle_res_blocks = num_middle_res_blocks + self.cond_input_channels = cond_input_channels + + +class Unet(Module): + def __init__(self, args: UnetArgs): + super().__init__() + self.args = args + + self.time_embed = Sequential( + TimeEmbedding(self.args.time_embedding_channels), + Linear(self.args.time_embedding_channels, self.args.cond_internal_channels), + SiLU(), + Linear(self.args.cond_internal_channels, self.args.cond_internal_channels)) + + self.cond_embed = Sequential( + Linear(self.args.cond_input_channels, self.args.cond_internal_channels), + SiLU(), + Linear(self.args.cond_internal_channels, self.args.cond_internal_channels)) + + self.first_conv = Conv2d(args.in_channels, args.model_channels, kernel_size=3, stride=1, padding=1) + current_channels = args.model_channels + channels = [current_channels] + + # Downsampling blocks + self.down_blocks = ModuleList() + for i in range(args.num_levels): + out_channels = args.model_channels * args.level_channel_multipliers[i] + perform_downsampling = i < args.num_levels - 1 + down_block = DownsamplingBlock( + in_channels=current_channels, + out_channels=out_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + num_res_blocks=args.num_res_blocks_per_level, + dropout_prob=args.dropout_prob, + use_attention=args.level_use_attention[i], + perform_downsampling=perform_downsampling, + attention_block_args=args.attention_block_args, + resample_with_res_block=args.resample_with_res_block, + use_conv_to_resample=args.use_conv_to_resample, + condition_bias=args.condition_bias) + self.down_blocks.append(down_block) + current_channels = out_channels + channels += down_block.output_channels + + # Middle blocks + self.middle_blocks = ModuleList() + for i in range(self.args.num_middle_res_blocks - 1): + self.middle_blocks.append(ResBlock( + in_channels=current_channels, + out_channels=current_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + dropout_prob=args.dropout_prob, + condition_bias=args.condition_bias)) + self.middle_blocks.append( + Arity3To1(AttentionBlock(num_channels=current_channels, args=args.attention_block_args))) + self.middle_blocks.append(ResBlock( + in_channels=current_channels, + out_channels=current_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + dropout_prob=args.dropout_prob, + condition_bias=args.condition_bias)) + + # Upsampling blocks + self.up_blocks = ModuleList() + for i in reversed(range(args.num_levels)): + skip_channels = [] + for j in range(args.num_res_blocks_per_level + 1): + skip_channels.append(channels.pop()) + perform_upsampling = i > 0 + out_channels = args.model_channels * args.level_channel_multipliers[i] + up_block = UpsamplingBlock( + in_channels=current_channels, + out_channels=out_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + num_resnet_blocks=args.num_res_blocks_per_level + 1, + skip_channels=skip_channels, + dropout_prob=args.dropout_prob, + use_attention=args.level_use_attention[i], + perform_upsampling=perform_upsampling, + attention_block_args=args.attention_block_args, + resample_with_res_block=args.resample_with_res_block, + use_conv_to_resample=args.use_conv_to_resample, + condition_bias=args.condition_bias) + self.up_blocks.append(up_block) + current_channels = out_channels + assert len(channels) == 0 + + self.last = Sequential( + GroupNorm32(current_channels), + SiLU(), + init_to_zero(Conv2d(current_channels, args.out_channels, kernel_size=3, stride=1, padding=1))) + + def forward(self, x: Tensor, t: Tensor, cond: Tensor): + t_emb = self.time_embed(t) + cond_emb = self.cond_embed(cond) + hs = [self.first_conv(x)] + for block in self.down_blocks: + hs += block.forward(hs[-1], t_emb, cond_emb) + h = hs[-1] + for block in self.middle_blocks: + h = block(h, t_emb, cond_emb) + for block in self.up_blocks: + skips = [] + for i in range(self.args.num_res_blocks_per_level + 1): + skips.append(hs.pop()) + h = block.forward(h, skips, t_emb, cond_emb) + assert len(hs) == 0 + return self.last(h) + + +class UnetWithFirstConvAddition(Module): + def __init__(self, args: UnetArgs): + super().__init__() + self.args = args + + self.time_embed = Sequential( + TimeEmbedding(self.args.time_embedding_channels), + Linear(self.args.time_embedding_channels, self.args.cond_internal_channels), + SiLU(), + Linear(self.args.cond_internal_channels, self.args.cond_internal_channels)) + + self.cond_embed = Sequential( + Linear(self.args.cond_input_channels, self.args.cond_internal_channels), + SiLU(), + Linear(self.args.cond_internal_channels, self.args.cond_internal_channels)) + + self.first_conv = Conv2d(args.in_channels, args.model_channels, kernel_size=3, stride=1, padding=1) + current_channels = args.model_channels + channels = [current_channels] + + # Downsampling blocks + self.down_blocks = ModuleList() + for i in range(args.num_levels): + out_channels = args.model_channels * args.level_channel_multipliers[i] + perform_downsampling = i < args.num_levels - 1 + down_block = DownsamplingBlock( + in_channels=current_channels, + out_channels=out_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + num_res_blocks=args.num_res_blocks_per_level, + dropout_prob=args.dropout_prob, + use_attention=args.level_use_attention[i], + perform_downsampling=perform_downsampling, + attention_block_args=args.attention_block_args, + resample_with_res_block=args.resample_with_res_block, + use_conv_to_resample=args.use_conv_to_resample, + condition_bias=args.condition_bias) + self.down_blocks.append(down_block) + current_channels = out_channels + channels += down_block.output_channels + + # Middle blocks + self.middle_blocks = ModuleList() + for i in range(self.args.num_middle_res_blocks - 1): + self.middle_blocks.append(ResBlock( + in_channels=current_channels, + out_channels=current_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + dropout_prob=args.dropout_prob, + condition_bias=args.condition_bias)) + self.middle_blocks.append( + Arity3To1(AttentionBlock(num_channels=current_channels, args=args.attention_block_args))) + self.middle_blocks.append(ResBlock( + in_channels=current_channels, + out_channels=current_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + dropout_prob=args.dropout_prob, + condition_bias=args.condition_bias)) + + # Upsampling blocks + self.up_blocks = ModuleList() + for i in reversed(range(args.num_levels)): + skip_channels = [] + for j in range(args.num_res_blocks_per_level + 1): + skip_channels.append(channels.pop()) + perform_upsampling = i > 0 + out_channels = args.model_channels * args.level_channel_multipliers[i] + up_block = UpsamplingBlock( + in_channels=current_channels, + out_channels=out_channels, + cond0_channels=args.cond_internal_channels, + cond1_channels=args.cond_internal_channels, + num_resnet_blocks=args.num_res_blocks_per_level + 1, + skip_channels=skip_channels, + dropout_prob=args.dropout_prob, + use_attention=args.level_use_attention[i], + perform_upsampling=perform_upsampling, + attention_block_args=args.attention_block_args, + resample_with_res_block=args.resample_with_res_block, + use_conv_to_resample=args.use_conv_to_resample, + condition_bias=args.condition_bias) + self.up_blocks.append(up_block) + current_channels = out_channels + assert len(channels) == 0 + + self.last = Sequential( + GroupNorm32(current_channels), + SiLU(), + init_to_zero(Conv2d(current_channels, args.out_channels, kernel_size=3, stride=1, padding=1))) + + def forward(self, x: Tensor, t: Tensor, cond: Tensor, first_conv_addition: Tensor): + t_emb = self.time_embed(t) + cond_emb = self.cond_embed(cond) + first_conv = self.first_conv(x) + hs = [first_conv + first_conv_addition] + for block in self.down_blocks: + hs += block.forward(hs[-1], t_emb, cond_emb) + h = hs[-1] + for block in self.middle_blocks: + h = block(h, t_emb, cond_emb) + for block in self.up_blocks: + skips = [] + for i in range(self.args.num_res_blocks_per_level + 1): + skips.append(hs.pop()) + h = block.forward(h, skips, t_emb, cond_emb) + assert len(hs) == 0 + return self.last(h) diff --git a/tha4/src/tha4/nn/conv.py b/tha4/src/tha4/nn/conv.py new file mode 100644 index 0000000..140de5d --- /dev/null +++ b/tha4/src/tha4/nn/conv.py @@ -0,0 +1,189 @@ +from typing import Optional, Union, Callable + +from torch.nn import Conv2d, Module, Sequential, ConvTranspose2d + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.nonlinearity_factory import resolve_nonlinearity_factory +from tha4.nn.normalization import NormalizationLayerFactory +from tha4.nn.util import wrap_conv_or_linear_module, BlockArgs + + +def create_conv7(in_channels: int, out_channels: int, + bias: bool = False, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + use_spectral_norm: bool = False) -> Module: + return wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=7, stride=1, padding=3, bias=bias), + initialization_method, + use_spectral_norm) + + +def create_conv7_from_block_args(in_channels: int, + out_channels: int, + bias: bool = False, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return create_conv7( + in_channels, out_channels, bias, + block_args.initialization_method, + block_args.use_spectral_norm) + + +def create_conv3(in_channels: int, + out_channels: int, + bias: bool = False, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + use_spectral_norm: bool = False) -> Module: + return wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias), + initialization_method, + use_spectral_norm) + + +def create_conv3_from_block_args(in_channels: int, out_channels: int, + bias: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return create_conv3(in_channels, out_channels, bias, + block_args.initialization_method, + block_args.use_spectral_norm) + + +def create_conv1(in_channels: int, out_channels: int, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + bias: bool = False, + use_spectral_norm: bool = False) -> Module: + return wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), + initialization_method, + use_spectral_norm) + + +def create_conv1_from_block_args(in_channels: int, + out_channels: int, + bias: bool = False, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return create_conv1( + in_channels=in_channels, + out_channels=out_channels, + initialization_method=block_args.initialization_method, + bias=bias, + use_spectral_norm=block_args.use_spectral_norm) + + +def create_conv7_block(in_channels: int, out_channels: int, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + return Sequential( + create_conv7(in_channels, out_channels, + bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_conv7_block_from_block_args( + in_channels: int, out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return create_conv7_block(in_channels, out_channels, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) + + +def create_conv3_block(in_channels: int, out_channels: int, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + return Sequential( + create_conv3(in_channels, out_channels, + bias=False, initialization_method=initialization_method, use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_conv3_block_from_block_args( + in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return create_conv3_block(in_channels, out_channels, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) + + +def create_downsample_block(in_channels: int, out_channels: int, + is_output_1x1: bool = False, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + if is_output_1x1: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + initialization_method, + use_spectral_norm), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + else: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + initialization_method, + use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_downsample_block_from_block_args(in_channels: int, out_channels: int, + is_output_1x1: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return create_downsample_block( + in_channels, out_channels, + is_output_1x1, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) + + +def create_upsample_block(in_channels: int, + out_channels: int, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + nonlinearity_factory: Optional[ModuleFactory] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False) -> Module: + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + return Sequential( + wrap_conv_or_linear_module( + ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + initialization_method, + use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(out_channels, affine=True), + resolve_nonlinearity_factory(nonlinearity_factory).create()) + + +def create_upsample_block_from_block_args(in_channels: int, + out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return create_upsample_block(in_channels, out_channels, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm) diff --git a/tha4/src/tha4/nn/eyebrow_decomposer/__init__.py b/tha4/src/tha4/nn/eyebrow_decomposer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/eyebrow_decomposer/eyebrow_decomposer_00.py b/tha4/src/tha4/nn/eyebrow_decomposer/eyebrow_decomposer_00.py new file mode 100644 index 0000000..6064f1e --- /dev/null +++ b/tha4/src/tha4/nn/eyebrow_decomposer/eyebrow_decomposer_00.py @@ -0,0 +1,81 @@ +from typing import List, Optional + +import torch +from torch import Tensor +from torch.nn import Module + +from tha4.nn.common.poser_encoder_decoder_00 import PoserEncoderDecoder00Args, PoserEncoderDecoder00 +from tha4.nn.image_processing_util import apply_color_change +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs + + +class EyebrowDecomposer00Args(PoserEncoderDecoder00Args): + def __init__(self, + image_size: int = 128, + image_channels: int = 4, + start_channels: int = 64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels: int = 512, + block_args: Optional[BlockArgs] = None): + super().__init__( + image_size, + image_channels, + image_channels, + 0, + start_channels, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels, + block_args) + + +class EyebrowDecomposer00(Module): + def __init__(self, args: EyebrowDecomposer00Args): + super().__init__() + self.args = args + self.body = PoserEncoderDecoder00(args) + self.background_layer_alpha = self.args.create_alpha_block() + self.background_layer_color_change = self.args.create_color_change_block() + self.eyebrow_layer_alpha = self.args.create_alpha_block() + self.eyebrow_layer_color_change = self.args.create_color_change_block() + + def forward(self, image: Tensor, *args) -> List[Tensor]: + feature = self.body(image)[0] + + background_layer_alpha = self.background_layer_alpha(feature) + background_layer_color_change = self.background_layer_color_change(feature) + background_layer_1 = apply_color_change(background_layer_alpha, background_layer_color_change, image) + + eyebrow_layer_alpha = self.eyebrow_layer_alpha(feature) + eyebrow_layer_color_change = self.eyebrow_layer_color_change(feature) + eyebrow_layer = apply_color_change(eyebrow_layer_alpha, image, eyebrow_layer_color_change) + + return [ + eyebrow_layer, # 0 + eyebrow_layer_alpha, # 1 + eyebrow_layer_color_change, # 2 + background_layer_1, # 3 + background_layer_alpha, # 4 + background_layer_color_change, # 5 + ] + + EYEBROW_LAYER_INDEX = 0 + EYEBROW_LAYER_ALPHA_INDEX = 1 + EYEBROW_LAYER_COLOR_CHANGE_INDEX = 2 + BACKGROUND_LAYER_INDEX = 3 + BACKGROUND_LAYER_ALPHA_INDEX = 4 + BACKGROUND_LAYER_COLOR_CHANGE_INDEX = 5 + OUTPUT_LENGTH = 6 + + +class EyebrowDecomposer00Factory(ModuleFactory): + def __init__(self, args: EyebrowDecomposer00Args): + super().__init__() + self.args = args + + def create(self) -> Module: + return EyebrowDecomposer00(self.args) diff --git a/tha4/src/tha4/nn/eyebrow_morphing_combiner/__init__.py b/tha4/src/tha4/nn/eyebrow_morphing_combiner/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/eyebrow_morphing_combiner/eyebrow_morphing_combiner_00.py b/tha4/src/tha4/nn/eyebrow_morphing_combiner/eyebrow_morphing_combiner_00.py new file mode 100644 index 0000000..eb579b5 --- /dev/null +++ b/tha4/src/tha4/nn/eyebrow_morphing_combiner/eyebrow_morphing_combiner_00.py @@ -0,0 +1,91 @@ +from typing import List, Optional + +import torch +from torch import Tensor +from torch.nn import Module + +from tha4.nn.common.poser_encoder_decoder_00 import PoserEncoderDecoder00Args, PoserEncoderDecoder00 +from tha4.nn.image_processing_util import apply_color_change, apply_grid_change, apply_rgb_change +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs + + +class EyebrowMorphingCombiner00Args(PoserEncoderDecoder00Args): + def __init__(self, + image_size: int = 128, + image_channels: int = 4, + num_pose_params: int = 12, + start_channels: int = 64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels: int = 512, + block_args: Optional[BlockArgs] = None): + super().__init__( + image_size, + 2 * image_channels, + image_channels, + num_pose_params, + start_channels, + bottleneck_image_size, + num_bottleneck_blocks, + max_channels, + block_args) + + +class EyebrowMorphingCombiner00(Module): + def __init__(self, args: EyebrowMorphingCombiner00Args): + super().__init__() + self.args = args + self.body = PoserEncoderDecoder00(args) + self.morphed_eyebrow_layer_grid_change = self.args.create_grid_change_block() + self.morphed_eyebrow_layer_alpha = self.args.create_alpha_block() + self.morphed_eyebrow_layer_color_change = self.args.create_color_change_block() + self.combine_alpha = self.args.create_alpha_block() + + def forward(self, background_layer: Tensor, eyebrow_layer: Tensor, pose: Tensor, *args) -> List[Tensor]: + combined_image = torch.cat([background_layer, eyebrow_layer], dim=1) + feature = self.body(combined_image, pose)[0] + + morphed_eyebrow_layer_grid_change = self.morphed_eyebrow_layer_grid_change(feature) + morphed_eyebrow_layer_alpha = self.morphed_eyebrow_layer_alpha(feature) + morphed_eyebrow_layer_color_change = self.morphed_eyebrow_layer_color_change(feature) + warped_eyebrow_layer = apply_grid_change(morphed_eyebrow_layer_grid_change, eyebrow_layer) + morphed_eyebrow_layer = apply_color_change( + morphed_eyebrow_layer_alpha, morphed_eyebrow_layer_color_change, warped_eyebrow_layer) + + combine_alpha = self.combine_alpha(feature) + eyebrow_image = apply_rgb_change(combine_alpha, morphed_eyebrow_layer, background_layer) + eyebrow_image_no_combine_alpha = apply_rgb_change( + (morphed_eyebrow_layer[:, 3:4, :, :] + 1.0) / 2.0, morphed_eyebrow_layer, background_layer) + + return [ + eyebrow_image, # 0 + combine_alpha, # 1 + eyebrow_image_no_combine_alpha, # 2 + morphed_eyebrow_layer, # 3 + morphed_eyebrow_layer_alpha, # 4 + morphed_eyebrow_layer_color_change, # 5 + warped_eyebrow_layer, # 6 + morphed_eyebrow_layer_grid_change, # 7 + ] + + EYEBROW_IMAGE_INDEX = 0 + COMBINE_ALPHA_INDEX = 1 + EYEBROW_IMAGE_NO_COMBINE_ALPHA_INDEX = 2 + MORPHED_EYEBROW_LAYER_INDEX = 3 + MORPHED_EYEBROW_LAYER_ALPHA_INDEX = 4 + MORPHED_EYEBROW_LAYER_COLOR_CHANGE_INDEX = 5 + WARPED_EYEBROW_LAYER_INDEX = 6 + MORPHED_EYEBROW_LAYER_GRID_CHANGE_INDEX = 7 + OUTPUT_LENGTH = 8 + + +class EyebrowMorphingCombiner00Factory(ModuleFactory): + def __init__(self, args: EyebrowMorphingCombiner00Args): + super().__init__() + self.args = args + + def create(self) -> Module: + return EyebrowMorphingCombiner00(self.args) diff --git a/tha4/src/tha4/nn/face_morpher/__init__.py b/tha4/src/tha4/nn/face_morpher/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/face_morpher/face_morpher_08.py b/tha4/src/tha4/nn/face_morpher/face_morpher_08.py new file mode 100644 index 0000000..d9bf5b1 --- /dev/null +++ b/tha4/src/tha4/nn/face_morpher/face_morpher_08.py @@ -0,0 +1,211 @@ +import math +from typing import List, Optional + +import torch +from torch import Tensor +from torch.nn import ModuleList, Sequential, Sigmoid, Tanh, Module +from torch.nn.functional import affine_grid, grid_sample + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.conv import create_conv3_block_from_block_args, \ + create_downsample_block_from_block_args, create_upsample_block_from_block_args, create_conv3_from_block_args, \ + create_conv3 +from tha4.nn.nonlinearity_factory import LeakyReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.resnet_block import ResnetBlock +from tha4.nn.util import BlockArgs + + +class FaceMorpher08Args: + def __init__(self, + image_size: int = 256, + image_channels: int = 4, + num_expression_params: int = 67, + start_channels: int = 16, + bottleneck_image_size=4, + num_bottleneck_blocks=3, + max_channels: int = 512, + block_args: Optional[BlockArgs] = None, + output_iris_mouth_grid_change: bool = False): + self.max_channels = max_channels + self.num_bottleneck_blocks = num_bottleneck_blocks + assert bottleneck_image_size > 1 + self.bottleneck_image_size = bottleneck_image_size + self.start_channels = start_channels + self.image_channels = image_channels + self.num_expression_params = num_expression_params + self.image_size = image_size + self.output_iris_mouth_grid_change = output_iris_mouth_grid_change + + if block_args is None: + self.block_args = BlockArgs( + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=LeakyReLUFactory(negative_slope=0.2, inplace=True)) + else: + self.block_args = block_args + + +class FaceMorpher08(Module): + def __init__(self, args: FaceMorpher08Args): + super().__init__() + self.args = args + self.num_levels = int(math.log2(args.image_size // args.bottleneck_image_size)) + 1 + + self.downsample_blocks = ModuleList() + self.downsample_blocks.append( + create_conv3_block_from_block_args( + args.image_channels, + args.start_channels, + args.block_args)) + current_image_size = args.image_size + current_num_channels = args.start_channels + while current_image_size > args.bottleneck_image_size: + next_image_size = current_image_size // 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.downsample_blocks.append(create_downsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + is_output_1x1=False, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + assert len(self.downsample_blocks) == self.num_levels + + self.bottleneck_blocks = ModuleList() + self.bottleneck_blocks.append(create_conv3_block_from_block_args( + in_channels=current_num_channels + args.num_expression_params, + out_channels=current_num_channels, + block_args=args.block_args)) + for i in range(1, args.num_bottleneck_blocks): + self.bottleneck_blocks.append( + ResnetBlock.create( + num_channels=current_num_channels, + is1x1=False, + block_args=args.block_args)) + + self.upsample_blocks = ModuleList() + while current_image_size < args.image_size: + next_image_size = current_image_size * 2 + next_num_channels = self.get_num_output_channels_from_image_size(next_image_size) + self.upsample_blocks.append(create_upsample_block_from_block_args( + in_channels=current_num_channels, + out_channels=next_num_channels, + block_args=args.block_args)) + current_image_size = next_image_size + current_num_channels = next_num_channels + + self.iris_mouth_grid_change = self.create_grid_change_block() + self.iris_mouth_color_change = self.create_color_change_block() + self.iris_mouth_alpha = self.create_alpha_block() + + self.eye_color_change = self.create_color_change_block() + self.eye_alpha = self.create_alpha_block() + + def create_alpha_block(self): + return Sequential( + create_conv3( + in_channels=self.args.start_channels, + out_channels=1, + bias=True, + initialization_method=self.args.block_args.initialization_method, + use_spectral_norm=False), + Sigmoid()) + + def create_color_change_block(self): + return Sequential( + create_conv3_from_block_args( + in_channels=self.args.start_channels, + out_channels=self.args.image_channels, + bias=True, + block_args=self.args.block_args), + Tanh()) + + def create_grid_change_block(self): + return create_conv3( + in_channels=self.args.start_channels, + out_channels=2, + bias=False, + initialization_method='zero', + use_spectral_norm=False) + + def get_num_output_channels_from_level(self, level: int): + return self.get_num_output_channels_from_image_size(self.args.image_size // (2 ** level)) + + def get_num_output_channels_from_image_size(self, image_size: int): + return min(self.args.start_channels * (self.args.image_size // image_size), self.args.max_channels) + + def merge_down(self, top_layer: Tensor, bottom_layer: Tensor): + top_layer_rgb = top_layer[:, 0:3, :, :] + top_layer_a = top_layer[:, 3:4, :, :] + return bottom_layer * (1 - top_layer_a) + torch.cat([top_layer_rgb * top_layer_a, top_layer_a], dim=1) + + def apply_grid_change(self, grid_change, image: Tensor) -> Tensor: + n, c, h, w = image.shape + device = grid_change.device + grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2) + identity = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + device=device, + dtype=grid_change.dtype).unsqueeze(0).repeat(n, 1, 1) + base_grid = affine_grid(identity, [n, c, h, w], align_corners=False) + grid = base_grid + grid_change + resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=False) + return resampled_image + + def apply_color_change(self, alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) + + def forward(self, image: Tensor, pose: Tensor, *args) -> List[Tensor]: + feature = image + for block in self.downsample_blocks: + feature = block(feature) + n, c = pose.shape + pose = pose.view(n, c, 1, 1).repeat(1, 1, self.args.bottleneck_image_size, self.args.bottleneck_image_size) + feature = torch.cat([feature, pose], dim=1) + for block in self.bottleneck_blocks: + feature = block(feature) + for block in self.upsample_blocks: + feature = block(feature) + + iris_mouth_grid_change = self.iris_mouth_grid_change(feature) + iris_mouth_image_0 = self.apply_grid_change(iris_mouth_grid_change, image) + iris_mouth_color_change = self.iris_mouth_color_change(feature) + iris_mouth_alpha = self.iris_mouth_alpha(feature) + iris_mouth_image_1 = self.apply_color_change(iris_mouth_alpha, iris_mouth_color_change, iris_mouth_image_0) + + eye_color_change = self.eye_color_change(feature) + eye_alpha = self.eye_alpha(feature) + output_image = self.apply_color_change(eye_alpha, eye_color_change, iris_mouth_image_1.detach()) + + outputs = [ + output_image, # 0 + eye_alpha, # 1 + eye_color_change, # 2 + iris_mouth_image_1, # 3 + iris_mouth_alpha, # 4 + iris_mouth_color_change, # 5 + iris_mouth_image_0, # 6 + ] + + if self.args.output_iris_mouth_grid_change: + outputs.append(iris_mouth_grid_change) + + return outputs + + OUTPUT_IMAGE_INDEX = 0 + EYE_ALPHA_INDEX = 1 + EYE_COLOR_CHANGE_INDEX = 2 + IRIS_MOUTH_IMAGE_1_INDEX = 3 + IRIS_MOUTH_ALPHA_INDEX = 4 + IRIS_MOUTH_COLOR_CHANGE_INDEX = 5 + IRIS_MOUTH_IMAGE_0_INDEX = 6 + IRIS_MOUTH_GRID_CHANGE_INDEX = 7 + + +class FaceMorpher08Factory(ModuleFactory): + def __init__(self, args: FaceMorpher08Args): + super().__init__() + self.args = args + + def create(self) -> Module: + return FaceMorpher08(self.args) diff --git a/tha4/src/tha4/nn/image_processing_util.py b/tha4/src/tha4/nn/image_processing_util.py new file mode 100644 index 0000000..e1d1a05 --- /dev/null +++ b/tha4/src/tha4/nn/image_processing_util.py @@ -0,0 +1,58 @@ +import torch +from torch import Tensor +from torch.nn.functional import affine_grid, grid_sample + + +def apply_rgb_change(alpha: Tensor, color_change: Tensor, image: Tensor): + image_rgb = image[:, 0:3, :, :] + color_change_rgb = color_change[:, 0:3, :, :] + output_rgb = color_change_rgb * alpha + image_rgb * (1 - alpha) + return torch.cat([output_rgb, image[:, 3:4, :, :]], dim=1) + + +def apply_grid_change(grid_change, image: Tensor) -> Tensor: + n, c, h, w = image.shape + device = grid_change.device + grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2) + identity = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=grid_change.dtype, + device=device).unsqueeze(0).repeat(n, 1, 1) + base_grid = affine_grid(identity, [n, c, h, w], align_corners=False) + grid = base_grid + grid_change + resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=False) + return resampled_image + + +class GridChangeApplier: + def __init__(self): + self.last_n = None + self.last_device = None + self.last_identity = None + + def apply(self, grid_change: Tensor, image: Tensor, align_corners: bool = False) -> Tensor: + n, c, h, w = image.shape + device = grid_change.device + grid_change = torch.transpose(grid_change.view(n, 2, h * w), 1, 2).view(n, h, w, 2) + + if n == self.last_n and device == self.last_device: + identity = self.last_identity + else: + identity = torch.tensor( + [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], + dtype=grid_change.dtype, + device=device, + requires_grad=False) \ + .unsqueeze(0).repeat(n, 1, 1) + self.last_identity = identity + self.last_n = n + self.last_device = device + base_grid = affine_grid(identity, [n, c, h, w], align_corners=align_corners) + + grid = base_grid + grid_change + resampled_image = grid_sample(image, grid, mode='bilinear', padding_mode='border', align_corners=align_corners) + return resampled_image + + +def apply_color_change(alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) diff --git a/tha4/src/tha4/nn/init_function.py b/tha4/src/tha4/nn/init_function.py new file mode 100644 index 0000000..b5061bd --- /dev/null +++ b/tha4/src/tha4/nn/init_function.py @@ -0,0 +1,76 @@ +from typing import Callable + +import torch +from torch import zero_ +from torch.nn import Module +from torch.nn.init import kaiming_normal_, xavier_normal_, normal_ + + +def create_init_function(method: str = 'none') -> Callable[[Module], Module]: + def init(module: Module): + if method == 'none': + return module + elif method == 'he': + kaiming_normal_(module.weight) + return module + elif method == 'xavier': + xavier_normal_(module.weight) + return module + elif method == 'dcgan': + normal_(module.weight, 0.0, 0.02) + return module + elif method == 'dcgan_001': + normal_(module.weight, 0.0, 0.01) + return module + elif method == "zero": + with torch.no_grad(): + zero_(module.weight) + return module + else: + raise ("Invalid initialization method %s" % method) + + return init + + +class HeInitialization: + def __init__(self, a: int = 0, mode: str = 'fan_in', nonlinearity: str = 'leaky_relu'): + self.nonlinearity = nonlinearity + self.mode = mode + self.a = a + + def __call__(self, module: Module) -> Module: + with torch.no_grad(): + kaiming_normal_(module.weight, a=self.a, mode=self.mode, nonlinearity=self.nonlinearity) + return module + + +class NormalInitialization: + def __init__(self, mean: float = 0.0, std: float = 1.0): + self.std = std + self.mean = mean + + def __call__(self, module: Module) -> Module: + with torch.no_grad(): + normal_(module.weight, self.mean, self.std) + return module + + +class XavierInitialization: + def __init__(self, gain: float = 1.0): + self.gain = gain + + def __call__(self, module: Module) -> Module: + with torch.no_grad(): + xavier_normal_(module.weight, self.gain) + return module + + +class ZeroInitialization: + def __call__(self, module: Module) -> Module: + with torch.no_grad: + zero_(module.weight) + return module + +class NoInitialization: + def __call__(self, module: Module) -> Module: + return module \ No newline at end of file diff --git a/tha4/src/tha4/nn/morpher/__init__.py b/tha4/src/tha4/nn/morpher/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/morpher/morpher_00.py b/tha4/src/tha4/nn/morpher/morpher_00.py new file mode 100644 index 0000000..b75ba2d --- /dev/null +++ b/tha4/src/tha4/nn/morpher/morpher_00.py @@ -0,0 +1,80 @@ +from typing import List + +import torch +from torch import Tensor +from torch.nn import Module + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.image_processing_util import GridChangeApplier +from tha4.nn.common.unet import UnetArgs, Unet, AttentionBlockArgs + + +def apply_color_change(alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) + + +class Morpher00Args: + def __init__(self, + image_size: int, + image_channels: int, + num_pose_parameters: int, + unet_args: UnetArgs): + assert unet_args.in_channels == image_channels + assert unet_args.out_channels == ( + image_channels + # direct + 2 + # warp + 1 # alpha + ) + assert unet_args.cond_input_channels == num_pose_parameters + self.image_channels = image_channels + self.image_size = image_size + self.num_pose_parameters = num_pose_parameters + self.unet_args = unet_args + + +class Morpher00(Module): + def __init__(self, args: Morpher00Args): + super().__init__() + self.args = args + self.body = Unet(args.unet_args) + self.grid_change_applier = GridChangeApplier() + + def forward(self, image: torch.Tensor, pose: torch.Tensor) -> List[Tensor]: + assert len(image.shape) == 4 + assert image.shape[1] == self.args.image_channels + assert image.shape[2] == self.args.image_size + assert image.shape[3] == self.args.image_size + assert len(pose.shape) == 2 + assert image.shape[0] == pose.shape[0] + assert pose.shape[1] == self.args.num_pose_parameters + + t = torch.zeros(image.shape[0], 1, device=image.device) + body_output = self.body(image, t, pose) + direct = body_output[:, 0:self.args.image_channels, :, :] + grid_change = body_output[:, self.args.image_channels:self.args.image_channels + 2, :, :] + alpha = torch.sigmoid(body_output[:, self.args.image_channels + 2:self.args.image_channels + 3, :, :]) + + warped = self.grid_change_applier.apply(grid_change, image) + merged = apply_color_change(alpha, direct, warped) + + return [ + merged, + alpha, + warped, + grid_change, + direct + ] + + INDEX_MERGED = 0 + INDEX_ALPHA = 1 + INDEX_WARPED = 2 + INDEX_GRID_CHANGE = 3 + INDEX_DIRECT = 4 + + +class Morpher00Factory(ModuleFactory): + def __init__(self, args: Morpher00Args): + self.args = args + + def create(self) -> Module: + return Morpher00(self.args) diff --git a/tha4/src/tha4/nn/nonlinearity_factory.py b/tha4/src/tha4/nn/nonlinearity_factory.py new file mode 100644 index 0000000..a935119 --- /dev/null +++ b/tha4/src/tha4/nn/nonlinearity_factory.py @@ -0,0 +1,72 @@ +from typing import Optional + +from torch.nn import Module, ReLU, LeakyReLU, ELU, ReLU6, Hardswish, SiLU, Tanh, Sigmoid + +from tha4.shion.core.module_factory import ModuleFactory + + +class ReLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return ReLU(self.inplace) + + +class LeakyReLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False, negative_slope: float = 1e-2): + self.negative_slope = negative_slope + self.inplace = inplace + + def create(self) -> Module: + return LeakyReLU(inplace=self.inplace, negative_slope=self.negative_slope) + + +class ELUFactory(ModuleFactory): + def __init__(self, inplace: bool = False, alpha: float = 1.0): + self.alpha = alpha + self.inplace = inplace + + def create(self) -> Module: + return ELU(inplace=self.inplace, alpha=self.alpha) + + +class ReLU6Factory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return ReLU6(inplace=self.inplace) + + +class SiLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return SiLU(inplace=self.inplace) + + +class HardswishFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return Hardswish(inplace=self.inplace) + + +class TanhFactory(ModuleFactory): + def create(self) -> Module: + return Tanh() + + +class SigmoidFactory(ModuleFactory): + def create(self) -> Module: + return Sigmoid() + + +def resolve_nonlinearity_factory(nonlinearity_fatory: Optional[ModuleFactory]) -> ModuleFactory: + if nonlinearity_fatory is None: + return ReLUFactory(inplace=False) + else: + return nonlinearity_fatory diff --git a/tha4/src/tha4/nn/normalization.py b/tha4/src/tha4/nn/normalization.py new file mode 100644 index 0000000..4d65584 --- /dev/null +++ b/tha4/src/tha4/nn/normalization.py @@ -0,0 +1,126 @@ +from abc import ABC, abstractmethod +from typing import Optional + +import torch +from torch import layer_norm +from torch.nn import Module, BatchNorm2d, InstanceNorm2d, Parameter +from torch.nn.init import normal_, constant_ + +from tha4.nn.pass_through import PassThrough + + +class PixelNormalization(Module): + def __init__(self, epsilon=1e-8): + super().__init__() + self.epsilon = epsilon + + def forward(self, x): + return x / torch.sqrt((x ** 2).mean(dim=1, keepdim=True) + self.epsilon) + + +class NormalizationLayerFactory(ABC): + def __init__(self): + super().__init__() + + @abstractmethod + def create(self, num_features: int, affine: bool = True) -> Module: + pass + + @staticmethod + def resolve_2d(factory: Optional['NormalizationLayerFactory']) -> 'NormalizationLayerFactory': + if factory is None: + return InstanceNorm2dFactory() + else: + return factory + + +class Bias2d(Module): + def __init__(self, num_features: int): + super().__init__() + self.num_features = num_features + self.bias = Parameter(torch.zeros(1, num_features, 1, 1)) + + def forward(self, x): + return x + self.bias + + +class NoNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + if affine: + return Bias2d(num_features) + else: + return PassThrough() + + +class BatchNorm2dFactory(NormalizationLayerFactory): + def __init__(self, + weight_mean: Optional[float] = None, + weight_std: Optional[float] = None, + bias: Optional[float] = None): + super().__init__() + self.bias = bias + self.weight_std = weight_std + self.weight_mean = weight_mean + + def get_weight_mean(self): + if self.weight_mean is None: + return 1.0 + else: + return self.weight_mean + + def get_weight_std(self): + if self.weight_std is None: + return 0.02 + else: + return self.weight_std + + def create(self, num_features: int, affine: bool = True) -> Module: + module = BatchNorm2d(num_features=num_features, affine=affine) + if affine: + if self.weight_mean is not None or self.weight_std is not None: + normal_(module.weight, self.get_weight_mean(), self.get_weight_std()) + if self.bias is not None: + constant_(module.bias, self.bias) + return module + + +class InstanceNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return InstanceNorm2d(num_features=num_features, affine=affine) + + +class PixelNormFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return PixelNormalization() + + +class LayerNorm2d(Module): + def __init__(self, channels: int, affine: bool = True): + super(LayerNorm2d, self).__init__() + self.channels = channels + self.affine = affine + + if self.affine: + self.weight = Parameter(torch.ones(1, channels, 1, 1)) + self.bias = Parameter(torch.zeros(1, channels, 1, 1)) + + def forward(self, x): + shape = x.size()[1:] + y = layer_norm(x, shape) * self.weight + self.bias + return y + +class LayerNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return LayerNorm2d(channels=num_features, affine=affine) diff --git a/tha4/src/tha4/nn/pass_through.py b/tha4/src/tha4/nn/pass_through.py new file mode 100644 index 0000000..c64d678 --- /dev/null +++ b/tha4/src/tha4/nn/pass_through.py @@ -0,0 +1,9 @@ +from torch.nn import Module + + +class PassThrough(Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x \ No newline at end of file diff --git a/tha4/src/tha4/nn/resnet_block.py b/tha4/src/tha4/nn/resnet_block.py new file mode 100644 index 0000000..43f80c0 --- /dev/null +++ b/tha4/src/tha4/nn/resnet_block.py @@ -0,0 +1,67 @@ +from typing import Optional + +import torch +from torch.nn import Module, Sequential, Parameter + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.conv import create_conv1, create_conv3 +from tha4.nn.nonlinearity_factory import resolve_nonlinearity_factory +from tha4.nn.normalization import NormalizationLayerFactory +from tha4.nn.util import BlockArgs + + +class ResnetBlock(Module): + @staticmethod + def create(num_channels: int, + is1x1: bool = False, + use_scale_parameters: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return ResnetBlock(num_channels, + is1x1, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm, + use_scale_parameters) + + def __init__(self, + num_channels: int, + is1x1: bool = False, + initialization_method: str = 'he', + nonlinearity_factory: ModuleFactory = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False, + use_scale_parameter: bool = False): + super().__init__() + self.use_scale_parameter = use_scale_parameter + if self.use_scale_parameter: + self.scale = Parameter(torch.zeros(1)) + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + if is1x1: + self.resnet_path = Sequential( + create_conv1(num_channels, num_channels, initialization_method, + bias=True, + use_spectral_norm=use_spectral_norm), + nonlinearity_factory.create(), + create_conv1(num_channels, num_channels, initialization_method, + bias=True, + use_spectral_norm=use_spectral_norm)) + else: + self.resnet_path = Sequential( + create_conv3(num_channels, num_channels, + bias=False, initialization_method=initialization_method, + use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(num_channels, affine=True), + nonlinearity_factory.create(), + create_conv3(num_channels, num_channels, + bias=False, initialization_method=initialization_method, + use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(num_channels, affine=True)) + + def forward(self, x): + if self.use_scale_parameter: + return x + self.scale * self.resnet_path(x) + else: + return x + self.resnet_path(x) diff --git a/tha4/src/tha4/nn/resnet_block_seperable.py b/tha4/src/tha4/nn/resnet_block_seperable.py new file mode 100644 index 0000000..4562a75 --- /dev/null +++ b/tha4/src/tha4/nn/resnet_block_seperable.py @@ -0,0 +1,71 @@ +from typing import Optional + +import torch +from torch.nn import Module, Sequential, Parameter + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.conv import create_conv1 +from tha4.nn.nonlinearity_factory import resolve_nonlinearity_factory +from tha4.nn.normalization import NormalizationLayerFactory +from tha4.nn.separable_conv import create_separable_conv3 +from tha4.nn.util import BlockArgs + + +class ResnetBlockSeparable(Module): + @staticmethod + def create(num_channels: int, + is1x1: bool = False, + use_scale_parameters: bool = False, + block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return ResnetBlockSeparable( + num_channels, + is1x1, + block_args.initialization_method, + block_args.nonlinearity_factory, + block_args.normalization_layer_factory, + block_args.use_spectral_norm, + use_scale_parameters) + + def __init__(self, + num_channels: int, + is1x1: bool = False, + initialization_method: str = 'he', + nonlinearity_factory: ModuleFactory = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + use_spectral_norm: bool = False, + use_scale_parameter: bool = False): + super().__init__() + self.use_scale_parameter = use_scale_parameter + if self.use_scale_parameter: + self.scale = Parameter(torch.zeros(1)) + nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + if is1x1: + self.resnet_path = Sequential( + create_conv1(num_channels, num_channels, initialization_method, + bias=True, + use_spectral_norm=use_spectral_norm), + nonlinearity_factory.create(), + create_conv1(num_channels, num_channels, initialization_method, + bias=True, + use_spectral_norm=use_spectral_norm)) + else: + self.resnet_path = Sequential( + create_separable_conv3( + num_channels, num_channels, + bias=False, initialization_method=initialization_method, + use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(num_channels, affine=True), + nonlinearity_factory.create(), + create_separable_conv3( + num_channels, num_channels, + bias=False, initialization_method=initialization_method, + use_spectral_norm=use_spectral_norm), + NormalizationLayerFactory.resolve_2d(normalization_layer_factory).create(num_channels, affine=True)) + + def forward(self, x): + if self.use_scale_parameter: + return x + self.scale * self.resnet_path(x) + else: + return x + self.resnet_path(x) diff --git a/tha4/src/tha4/nn/separable_conv.py b/tha4/src/tha4/nn/separable_conv.py new file mode 100644 index 0000000..c1ff1f4 --- /dev/null +++ b/tha4/src/tha4/nn/separable_conv.py @@ -0,0 +1,119 @@ +from typing import Optional + +from torch.nn import Sequential, Conv2d, ConvTranspose2d, Module + +from tha4.nn.normalization import NormalizationLayerFactory +from tha4.nn.util import BlockArgs, wrap_conv_or_linear_module + + +def create_separable_conv3(in_channels: int, out_channels: int, + bias: bool = False, + initialization_method='he', + use_spectral_norm: bool = False) -> Module: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=False, groups=in_channels), + initialization_method, + use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), + initialization_method, + use_spectral_norm)) + + +def create_separable_conv7(in_channels: int, out_channels: int, + bias: bool = False, + initialization_method='he', + use_spectral_norm: bool = False) -> Module: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, in_channels, kernel_size=7, stride=1, padding=3, bias=False, groups=in_channels), + initialization_method, + use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), + initialization_method, + use_spectral_norm)) + + +def create_separable_conv3_block( + in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=False, groups=in_channels), + block_args.initialization_method, + block_args.use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + block_args.initialization_method, + block_args.use_spectral_norm), + NormalizationLayerFactory.resolve_2d(block_args.normalization_layer_factory).create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) + + +def create_separable_conv7_block( + in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, in_channels, kernel_size=7, stride=1, padding=3, bias=False, groups=in_channels), + block_args.initialization_method, + block_args.use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + block_args.initialization_method, + block_args.use_spectral_norm), + NormalizationLayerFactory.resolve_2d(block_args.normalization_layer_factory).create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) + + +def create_separable_downsample_block( + in_channels: int, out_channels: int, is_output_1x1: bool, block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + if is_output_1x1: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, in_channels, kernel_size=4, stride=2, padding=1, bias=False, groups=in_channels), + block_args.initialization_method, + block_args.use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + block_args.initialization_method, + block_args.use_spectral_norm), + block_args.nonlinearity_factory.create()) + else: + return Sequential( + wrap_conv_or_linear_module( + Conv2d(in_channels, in_channels, kernel_size=4, stride=2, padding=1, bias=False, groups=in_channels), + block_args.initialization_method, + block_args.use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + block_args.initialization_method, + block_args.use_spectral_norm), + NormalizationLayerFactory.resolve_2d(block_args.normalization_layer_factory) + .create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) + + +def create_separable_upsample_block( + in_channels: int, out_channels: int, block_args: Optional[BlockArgs] = None): + if block_args is None: + block_args = BlockArgs() + return Sequential( + wrap_conv_or_linear_module( + ConvTranspose2d( + in_channels, in_channels, kernel_size=4, stride=2, padding=1, bias=False, groups=in_channels), + block_args.initialization_method, + block_args.use_spectral_norm), + wrap_conv_or_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=False), + block_args.initialization_method, + block_args.use_spectral_norm), + NormalizationLayerFactory.resolve_2d(block_args.normalization_layer_factory) + .create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) diff --git a/tha4/src/tha4/nn/siren/__init__.py b/tha4/src/tha4/nn/siren/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/siren/face_morpher/__init__.py b/tha4/src/tha4/nn/siren/face_morpher/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_00.py b/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_00.py new file mode 100644 index 0000000..ace842f --- /dev/null +++ b/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_00.py @@ -0,0 +1,59 @@ +from typing import Optional, List + +import torch +from torch import Tensor +from torch.nn import Module +from torch.nn.functional import affine_grid + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.siren.vanilla.siren import SirenArgs, Siren + + +class SirenFaceMorpher00Args: + def __init__(self, + image_size: int, + image_channels: int, + pose_size: int, + siren_args: SirenArgs): + assert siren_args.in_channels == pose_size + 2 + assert siren_args.out_channels == image_channels + assert not siren_args.use_tanh + + self.siren_args = siren_args + self.pose_size = pose_size + self.image_size = image_size + self.image_channels = image_channels + + +class SirenFaceMorpher00(Module): + def __init__(self, args: SirenFaceMorpher00Args): + super().__init__() + self.args = args + self.siren = Siren(self.args.siren_args) + + def forward(self, pose: Tensor, position: Optional[Tensor] = None) -> Tensor: + n, p = pose.shape[0], pose.shape[1] + device = pose.device + + if position is None: + h, w = self.args.image_size, self.args.image_size + identity = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], device=device).unsqueeze(0) + position = affine_grid(identity, [1, 1, h, w], align_corners=False) \ + .view(1, h * w, 2) + position = torch.transpose(position, dim0=1, dim1=2).view(1, 2, h, w) \ + .repeat(n, 1, 1, 1) + + h, w = position.shape[2], position.shape[3] + pose_image = pose.view(n, p, 1, 1).repeat(1, 1, h, w) + + siren_input = torch.cat([position, pose_image], dim=1) + + return self.siren.forward(siren_input) + + +class SirenFaceMorpher00Factory(ModuleFactory): + def __init__(self, args: SirenFaceMorpher00Args): + self.args = args + + def create(self) -> Module: + return SirenFaceMorpher00(self.args) diff --git a/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_00_trainer.py b/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_00_trainer.py new file mode 100644 index 0000000..f606c70 --- /dev/null +++ b/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_00_trainer.py @@ -0,0 +1,228 @@ +from typing import Dict, List, Optional, Callable + +import torch +from tha4.shion.base.dataset.lazy_tensor_dataset import LazyTensorDataset +from tha4.shion.base.image_util import extract_pytorch_image_from_filelike +from tha4.shion.base.loss.l1_loss import L1Loss, MaskedL1Loss +from tha4.shion.base.loss.sum_loss import SumLoss +from tha4.shion.base.optimizer_factories import AdamOptimizerFactory +from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer +from tha4.dataset.image_poses_and_aother_images_dataset import ImagePosesAndOtherImagesDataset +from tha4.nn.siren.face_morpher.siren_face_morpher_00 import SirenFaceMorpher00Factory, SirenFaceMorpher00Args +from tha4.nn.siren.face_morpher.siren_face_morpher_protocols_00 import SirenFaceMorpherComputationProtocol00, \ + SirenFaceMorpherSampleOutputProtocol00 +from tha4.nn.siren.morpher.siren_morpher_protocols_03 import SirenMorpherTrainingProtocol03 +from tha4.nn.siren.vanilla.siren import SirenArgs +from tha4.poser.poser import Poser +from torch import Tensor + +KEY_MODULE = "module" +KEY_POSER = "poser" + + +def get_poser(): + import tha4.poser.modes.mode_12 + poser = tha4.poser.modes.mode_12.create_poser(torch.device('cpu')) + return poser + + +class SirenFaceMorpher00TrainerArgs: + def __init__(self, + character_file_name: str, + face_mask_file_name: str, + pose_dataset_file_name: str, + num_training_total_examples: int = 1_000_000, + num_training_examples_per_checkpoint: int = 100_000, + num_training_examples_lr_boundaries: Optional[List[int]] = None, + num_training_examples_per_sample_output: Optional[int] = 5_000, + num_training_examples_per_snapshot: int = 10_000, + total_batch_size: int = 8, + training_random_seed: int = 2965603729, + sample_output_random_seed: int = 3522651501, + total_worker: int = 16, + poser_func: Optional[Callable[[], Poser]] = None, + base_learning_rate: float = 1e-4): + assert num_training_total_examples % num_training_examples_per_checkpoint == 0 + + if num_training_examples_lr_boundaries is None: + num_training_examples_lr_boundaries = [ + int(num_training_examples_per_checkpoint * 2), + int(num_training_examples_per_checkpoint * 5), + int(num_training_examples_per_checkpoint * 8), + ] + + for x in num_training_examples_lr_boundaries: + assert x % num_training_examples_per_snapshot == 0 + + if poser_func is None: + poser_func = get_poser + + self.face_mask_file_name = face_mask_file_name + self.base_learning_rate = base_learning_rate + self.poser_func = poser_func + self.total_worker = total_worker + self.num_training_examples_per_snapshot = num_training_examples_per_snapshot + self.num_training_examples_per_sample_output = num_training_examples_per_sample_output + self.sample_output_random_seed = sample_output_random_seed + self.training_random_seed = training_random_seed + self.total_batch_size = total_batch_size + self.num_training_total_examples = num_training_total_examples + self.num_training_examples_per_checkpoint = num_training_examples_per_checkpoint + self.num_training_examples_lr_boundaries = num_training_examples_lr_boundaries + self.pose_dataset_file_name = pose_dataset_file_name + self.character_file_name = character_file_name + + def get_character_image(self): + return extract_pytorch_image_from_filelike( + self.character_file_name, + scale=2.0, + offset=-1.0, + premultiply_alpha=True, + perform_srgb_to_linear=True) + + def get_face_mask_image(self): + loaded_image = extract_pytorch_image_from_filelike( + self.face_mask_file_name, + scale=1.0, + offset=0.0, + premultiply_alpha=True, + perform_srgb_to_linear=True) + output_image = torch.zeros(4, 128, 128) + center_x = 256 + center_y = 128 + 16 + for i in range(4): + output_image[i, :, :] = loaded_image[0, center_y - 64:center_y + 64, center_x - 64:center_x + 64] + return output_image + + def get_training_dataset(self): + return ImagePosesAndOtherImagesDataset( + main_image_func=self.get_character_image, + other_image_funcs=[self.get_face_mask_image], + pose_dataset=LazyTensorDataset(self.pose_dataset_file_name)) + + def get_module_factory(self): + return SirenFaceMorpher00Factory( + SirenFaceMorpher00Args( + image_size=128, + image_channels=4, + pose_size=39, + siren_args=SirenArgs( + in_channels=39 + 2, + out_channels=4, + intermediate_channels=128, + num_sine_layers=8))) + + def transform_pose_to_module_input(self, pose: Tensor): + return pose[:, 0:39] + + def transform_original_image_to_module_input(self, image: Tensor): + center_x = 256 + center_y = 128 + 16 + return image[:, :, center_y - 64:center_y + 64, center_x - 64:center_x + 64] + + def transform_poser_posed_image_to_groundtruth(self, image: Tensor): + center_x = 96 + center_y = 96 + 16 + return image[:, :, center_y - 64:center_y + 64, center_x - 64:center_x + 64] + + def get_training_computation_protocol(self): + return SirenFaceMorpherComputationProtocol00( + transform_pose_to_module_input_func=self.transform_pose_to_module_input, + transform_original_image_to_module_input_func=self.transform_original_image_to_module_input, + transform_poser_posed_image_to_groundtruth_func=self.transform_poser_posed_image_to_groundtruth) + + def get_learning_rate(self, examples_seen_so_far) -> Dict[str, float]: + if examples_seen_so_far < self.num_training_examples_lr_boundaries[0]: + return { + KEY_MODULE: self.base_learning_rate, + } + elif examples_seen_so_far < self.num_training_examples_lr_boundaries[1]: + return { + KEY_MODULE: self.base_learning_rate / 3.0, + } + elif examples_seen_so_far < self.num_training_examples_lr_boundaries[2]: + return { + KEY_MODULE: self.base_learning_rate / 10.0, + } + else: + return { + KEY_MODULE: self.base_learning_rate / 30.0, + } + + def get_optimizer_factories(self): + return { + KEY_MODULE: AdamOptimizerFactory(betas=(0.9, 0.999)), + } + + def get_poser(self): + return self.poser_func() + + def get_training_protocol(self, world_size: int): + total_examples = self.num_training_total_examples + per_checkpoint_examples = self.num_training_examples_per_checkpoint + num_checkpoints = total_examples // per_checkpoint_examples + batch_size = self.total_batch_size // world_size + return SirenMorpherTrainingProtocol03( + check_point_examples=[per_checkpoint_examples * (i + 1) for i in range(num_checkpoints)], + batch_size=batch_size, + learning_rate=self.get_learning_rate, + optimizer_factories=self.get_optimizer_factories(), + random_seed=self.training_random_seed, + poser_func=self.get_poser, + key_module=KEY_MODULE, + key_poser=KEY_POSER) + + def get_sample_output_protocol(self): + return SirenFaceMorpherSampleOutputProtocol00( + num_images=8, + image_size=128, + images_per_row=2, + examples_per_sample_output=self.num_training_examples_per_sample_output, + computation_protocol=self.get_training_computation_protocol(), + poser_func=self.get_poser, + random_seed=self.sample_output_random_seed) + + def get_loss(self): + protocol = self.get_training_computation_protocol() + return SumLoss([ + ( + 'full', + L1Loss( + expected_func=protocol.get_output_func(protocol.keys.groundtruth_posed_image), + actual_func=protocol.get_output_func(protocol.keys.predicted_posed_image), + weight=1.0) + ), + ( + 'eye_mouth', + MaskedL1Loss( + expected_func=protocol.get_output_func(protocol.keys.groundtruth_posed_image), + actual_func=protocol.get_output_func(protocol.keys.predicted_posed_image), + mask_func=protocol.get_output_func(protocol.keys.eye_mouth_mask), + weight=20.0) + ), + ]) + + def create_trainer(self, prefix: str, world_size: int, distrib_backend: str = 'gloo'): + if self.num_training_examples_per_sample_output is not None: + sample_output_protocol = self.get_sample_output_protocol() + else: + sample_output_protocol = None + + return DistributedTrainer( + prefix=prefix, + module_factories={ + KEY_MODULE: self.get_module_factory(), + }, + accumulators={}, + losses={ + KEY_MODULE: self.get_loss(), + }, + training_dataset=self.get_training_dataset(), + validation_dataset=self.get_training_dataset(), + training_protocol=self.get_training_protocol(world_size), + validation_protocol=None, + sample_output_protocol=sample_output_protocol, + pretrained_module_file_names={}, + example_per_snapshot=self.num_training_examples_per_snapshot, + num_data_loader_workers=max(1, self.total_worker // world_size), + distrib_backend=distrib_backend) diff --git a/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_protocols_00.py b/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_protocols_00.py new file mode 100644 index 0000000..83094f0 --- /dev/null +++ b/tha4/src/tha4/nn/siren/face_morpher/siren_face_morpher_protocols_00.py @@ -0,0 +1,220 @@ +import os +from dataclasses import dataclass +from typing import Dict, Any, Optional, Callable + +import PIL.Image +import numpy +import torch +from tha4.shion.base.dataset.util import get_indexed_batch +from tha4.shion.base.image_util import pytorch_rgba_to_numpy_image +from tha4.shion.core.cached_computation import CachedComputationProtocol, ComputationState, \ + ComposableCachedComputationProtocol, batch_indexing_func, add_step +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.poser.general_poser_02 import GeneralPoser02 +from torch import Tensor +from torch.nn import Module +from torch.utils.data import Dataset + +KEY_MODULE = "module" +KEY_POSER = "poser" + + +@dataclass +class SirenMorpherProtocol00Keys: + module: str = KEY_MODULE + module_output: str = "module_output" + + poser: str = KEY_POSER + poser_output: str = "poser_output" + + original_image: str = "original_image" + original_pose: str = "original_pose" + + module_input_image: str = "module_input_image" + module_input_pose: str = "module_input_pose" + + groundtruth_posed_image: str = 'groundtruth_posed_image' + predicted_posed_image: str = 'predicted_posed_image' + + eye_mouth_mask: str = 'eye_mouth_mask' + + +@dataclass +class SirenMorpherProtocol00Indices: + batch_original_image: int = 0 + batch_pose: int = 1 + batch_eye_mouth_mask: int = 2 + poser_posed_image: int = 0 + + +class SirenFaceMorpherComputationProtocol00(ComposableCachedComputationProtocol): + def __init__(self, + transform_pose_to_module_input_func: Callable[[Tensor], Tensor], + transform_original_image_to_module_input_func: Callable[[Tensor], Tensor], + transform_poser_posed_image_to_groundtruth_func: Callable[[Tensor], Tensor], + keys: Optional[SirenMorpherProtocol00Keys] = None, + indices: Optional[SirenMorpherProtocol00Indices] = None): + super().__init__() + + if keys is None: + keys = SirenMorpherProtocol00Keys() + if indices is None: + indices = SirenMorpherProtocol00Indices() + + self.keys = keys + self.indices = indices + self.transform_image_to_module_input_func = transform_original_image_to_module_input_func + self.transform_pose_to_module_input_func = transform_pose_to_module_input_func + self.transform_poser_posed_image_to_groundtruth_func = transform_poser_posed_image_to_groundtruth_func + + self.computation_steps[keys.original_image] = batch_indexing_func(indices.batch_original_image) + self.computation_steps[keys.original_pose] = batch_indexing_func(indices.batch_pose) + + @add_step(self.computation_steps, keys.module_input_pose) + def get_module_input_pose(protocol: CachedComputationProtocol, state: ComputationState): + original_pose = protocol.get_output(keys.original_pose, state) + return transform_pose_to_module_input_func(original_pose) + + @add_step(self.computation_steps, keys.module_input_image) + def get_module_input_image(protocol: CachedComputationProtocol, state: ComputationState): + original_image = protocol.get_output(keys.original_image, state) + return transform_original_image_to_module_input_func(original_image) + + @add_step(self.computation_steps, keys.poser_output) + def get_poser_output(protocol: CachedComputationProtocol, state: ComputationState): + with torch.no_grad(): + poser = state.modules[keys.poser] + pose = protocol.get_output(keys.original_pose, state) + image = protocol.get_output(keys.original_image, state) + return poser.get_posing_outputs(image, pose) + + @add_step(self.computation_steps, keys.groundtruth_posed_image) + def get_groundtruth_posed_image(protocol: CachedComputationProtocol, state: ComputationState): + poser_output = protocol.get_output(keys.poser_output, state) + poser_posed_image = poser_output[indices.poser_posed_image] + return transform_poser_posed_image_to_groundtruth_func(poser_posed_image) + + @add_step(self.computation_steps, keys.module_output) + def get_module_output(protocol: CachedComputationProtocol, state: ComputationState): + module_input_pose = protocol.get_output(keys.module_input_pose, state) + module = state.modules[keys.module] + return module.forward(module_input_pose) + + @add_step(self.computation_steps, keys.predicted_posed_image) + def get_predicted_image(protocol: CachedComputationProtocol, state: ComputationState): + return protocol.get_output(keys.module_output, state) + + self.computation_steps[keys.eye_mouth_mask] = batch_indexing_func(indices.batch_eye_mouth_mask) + + +class SirenFaceMorpherSampleOutputProtocol00(SampleOutputProtocol): + def __init__(self, + num_images: int, + image_size: int, + images_per_row: int, + examples_per_sample_output: int, + computation_protocol: SirenFaceMorpherComputationProtocol00, + poser_func: Callable[[], GeneralPoser02], + random_seed: int = 54859395058, + batch_size: Optional[int] = None): + if batch_size is None: + batch_size = num_images + + self.batch_size = batch_size + self.poser_func = poser_func + self.random_seed = random_seed + self.examples_per_sample_output = examples_per_sample_output + self.images_per_row = images_per_row + self.image_size = image_size + self.num_images = num_images + self.computation_protocol = computation_protocol + + def get_examples_per_sample_output(self) -> int: + return self.examples_per_sample_output + + def get_random_seed(self) -> int: + return self.random_seed + + def get_sample_output_data(self, validation_dataset: Dataset, device: torch.device) -> dict: + example_indices = torch.randint(0, len(validation_dataset), (self.num_images,)) + example_indices = [example_indices[i].item() for i in range(self.num_images)] + batch = get_indexed_batch(validation_dataset, example_indices, device) + poser = self.poser_func() + poser.to(device) + with torch.no_grad(): + ground_truth = poser.pose( + batch[self.computation_protocol.indices.batch_original_image], + batch[self.computation_protocol.indices.batch_pose]) + return { + 'batch': batch, + 'ground_truth': ground_truth + } + + def save_sample_output_data(self, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + sample_output_data: Any, + prefix: str, + examples_seen_so_far: int, + device: torch.device): + batch = sample_output_data['batch'] + ground_truth = sample_output_data['ground_truth'] + ground_truth = self.computation_protocol.transform_poser_posed_image_to_groundtruth_func(ground_truth) + + module = modules[self.computation_protocol.keys.module] + module.train(False) + + if self.batch_size == self.num_images: + with torch.no_grad(): + state = ComputationState( + modules=modules, + accumulated_modules=accumulated_modules, + batch=batch, + outputs={}) + poser_output_images = self.computation_protocol.get_output( + self.computation_protocol.keys.predicted_posed_image, state) + else: + poser_output_images_list = [] + start = 0 + while start < self.num_images: + end = start + self.batch_size + end = min(self.num_images, end) + minibatch = [batch[i][start:end] for i in range(len(batch))] + state = ComputationState( + modules=modules, + accumulated_modules=accumulated_modules, + batch=minibatch, + outputs={}) + with torch.no_grad(): + poser_output_images = self.computation_protocol.get_output( + self.computation_protocol.keys.predicted_posed_image, state) + poser_output_images_list.append(poser_output_images) + start = end + poser_output_images = torch.cat(poser_output_images_list, dim=0) + + num_rows = self.num_images // self.images_per_row + if self.num_images % self.images_per_row > 0: + num_rows += 1 + num_cols = 2 * self.images_per_row + + image_channels = 4 + output_image = numpy.zeros([self.image_size * num_rows, self.image_size * num_cols, image_channels]) + + for image_index in range(self.num_images): + row = image_index // self.images_per_row + start_row = row * self.image_size + + col = 2 * (image_index % self.images_per_row) + start_col = col * self.image_size + output_image[start_row:start_row + self.image_size, start_col:start_col + self.image_size, :] \ + = pytorch_rgba_to_numpy_image(ground_truth[image_index].detach().cpu()) + + start_col += self.image_size + output_image[start_row:start_row + self.image_size, start_col:start_col + self.image_size, :] \ + = pytorch_rgba_to_numpy_image(poser_output_images[image_index].detach().cpu()) + + file_name = "%s/sample_output_%010d.png" % (prefix, examples_seen_so_far) + os.makedirs(os.path.dirname(file_name), exist_ok=True) + pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(output_image * 255.0)), mode='RGBA') + pil_image.save(file_name) + print("Saved %s" % file_name) diff --git a/tha4/src/tha4/nn/siren/morpher/__init__.py b/tha4/src/tha4/nn/siren/morpher/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/siren/morpher/siren_morpher_03.py b/tha4/src/tha4/nn/siren/morpher/siren_morpher_03.py new file mode 100644 index 0000000..a4d18ff --- /dev/null +++ b/tha4/src/tha4/nn/siren/morpher/siren_morpher_03.py @@ -0,0 +1,153 @@ +from typing import List, Optional, Callable + +import torch +from torch import Tensor +from torch.nn import Module, ModuleList, Sequential, Conv2d +from torch.nn.functional import affine_grid, interpolate + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.nn00.initialization_funcs import HeInitialization +from tha4.nn.image_processing_util import GridChangeApplier +from tha4.nn.siren.vanilla.siren import SineLinearLayer + + +class SirenMorpherLevelArgs: + def __init__(self, + image_size: int, + intermediate_channels: int, + num_sine_layers: int): + assert num_sine_layers >= 2 + self.image_size = image_size + self.num_sine_layers = num_sine_layers + self.intermediate_channels = intermediate_channels + + +class SirenMorpher03Args: + def __init__(self, + image_size: int, + image_channels: int, + pose_size: int, + level_args: List[SirenMorpherLevelArgs], + init_func: Optional[Callable[[Module], Module]] = None): + assert len(level_args) >= 2 + if init_func is None: + init_func = HeInitialization() + self.image_size = image_size + self.init_func = init_func + self.level_args = level_args + self.pose_size = pose_size + self.image_channels = image_channels + + +class SirenMorpher03(Module): + def __init__(self, args: SirenMorpher03Args): + super().__init__() + self.args = args + + self.siren_layers = ModuleList() + + for i in range(len(args.level_args)): + level_args = args.level_args[i] + + layers = [] + + if i == 0: + layers.append(SineLinearLayer( + in_channels=args.pose_size + 2, + out_channels=level_args.intermediate_channels, + is_first=True)) + else: + layers.append(SineLinearLayer( + in_channels=level_args.intermediate_channels + args.pose_size + 2, + out_channels=level_args.intermediate_channels, + is_first=False)) + + for j in range(1, level_args.num_sine_layers - 1): + layers.append(SineLinearLayer( + in_channels=level_args.intermediate_channels, + out_channels=level_args.intermediate_channels, + is_first=False)) + + if i == len(args.level_args) - 1: + out_channels = level_args.intermediate_channels + else: + out_channels = args.level_args[i + 1].intermediate_channels + layers.append(SineLinearLayer( + in_channels=level_args.intermediate_channels, + out_channels=out_channels, + is_first=False)) + + self.siren_layers.append(Sequential(*layers)) + + self.last_linear = args.init_func(Conv2d( + args.level_args[-1].intermediate_channels, + args.image_channels + 2 + 1, + kernel_size=1, + stride=1, + padding=0, + bias=True)) + + self.grid_change_applier = GridChangeApplier() + + def get_position_grid(self, n: int, image_size: int, device: torch.device): + h, w = image_size, image_size + identity = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], device=device).unsqueeze(0) + position = affine_grid(identity, [1, 1, h, w], align_corners=False) \ + .view(1, h * w, 2) + position = torch.transpose(position, dim0=1, dim1=2).view(1, 2, h, w) \ + .repeat(n, 1, 1, 1) + return position + + def get_pose_image(self, pose: Tensor, image_size: int): + n, p = pose.shape[0], pose.shape[1] + h, w = image_size, image_size + pose_image = pose.view(n, p, 1, 1).repeat(1, 1, h, w) + return pose_image + + def forward(self, image: Tensor, pose: Tensor) -> List[Tensor]: + n = pose.shape[0] + device = pose.device + + x = None + for i in range(len(self.args.level_args)): + args = self.args.level_args[i] + position_and_pose = torch.cat([ + self.get_position_grid(n, args.image_size, device), + self.get_pose_image(pose, args.image_size) + ], dim=1) + if i == 0: + x = self.siren_layers[i].forward(position_and_pose) + else: + x = interpolate(x, size=(args.image_size, args.image_size), mode='bilinear') + x = torch.cat([x, position_and_pose], dim=1) + x = self.siren_layers[i].forward(x) + + siren_output = self.last_linear(x) + + grid_change = siren_output[:, 0:2, :, :] + alpha = siren_output[:, 2:3, :, :] + color_change = siren_output[:, 3:, :, :] + warped_image = self.grid_change_applier.apply(grid_change, image, align_corners=False) + blended_image = (1 - alpha) * warped_image + alpha * color_change + + return [ + blended_image, + alpha, + color_change, + warped_image, + grid_change + ] + + INDEX_BLENDED_IMAGE = 0 + INDEX_ALPHA = 1 + INDEX_COLOR_CHANGE = 2 + INDEX_WARPED_IMAGE = 3 + INDEX_GRID_CHANGE = 4 + + +class SirenMorpher03Factory(ModuleFactory): + def __init__(self, args: SirenMorpher03Args): + self.args = args + + def create(self): + return SirenMorpher03(self.args) diff --git a/tha4/src/tha4/nn/siren/morpher/siren_morpher_03_trainer.py b/tha4/src/tha4/nn/siren/morpher/siren_morpher_03_trainer.py new file mode 100644 index 0000000..5100105 --- /dev/null +++ b/tha4/src/tha4/nn/siren/morpher/siren_morpher_03_trainer.py @@ -0,0 +1,276 @@ +from enum import Enum +from typing import Dict, List, Optional, Callable + +import torch +from tha4.shion.base.dataset.lazy_tensor_dataset import LazyTensorDataset +from tha4.shion.base.image_util import extract_pytorch_image_from_filelike +from tha4.shion.base.loss.l1_loss import L1Loss +from tha4.shion.base.loss.sum_loss import SumLoss +from tha4.shion.base.loss.time_dependently_weighted_loss import TimeDependentlyWeightedLoss +from tha4.shion.base.optimizer_factories import AdamOptimizerFactory +from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer +from tha4.dataset.image_poses_and_aother_images_dataset import ImagePosesAndOtherImagesDataset +from tha4.nn.siren.morpher.siren_morpher_03 import SirenMorpherLevelArgs, SirenMorpher03Factory, SirenMorpher03Args +from tha4.nn.siren.morpher.siren_morpher_protocols_03 import SirenMorpherComputationProtocol03, \ + SirenMorpherProtocol03Indices, KEY_MODULE, KEY_POSER, KEY_EXAMPLES_SEEN_SO_FAR, SirenMorpherTrainingProtocol03, \ + SirenMorpherSampleOutputProtocol +from tha4.poser.poser import Poser + + +def get_poser(): + import tha4.poser.modes.mode_07 + poser = tha4.poser.modes.mode_07.create_poser(torch.device('cpu')) + return poser + + +class LossTerm(Enum): + full_blended = 1 + full_warped = 2 + full_grid_change = 3 + full_color_change = 4 + + def get_loss(self, protocol: SirenMorpherComputationProtocol03): + if self == LossTerm.full_blended: + return L1Loss( + expected_func=protocol.get_output_func(protocol.keys.groundtruth_posed_image), + actual_func=protocol.get_output_func(protocol.keys.predicted_posed_image)) + elif self == LossTerm.full_warped: + return L1Loss( + expected_func=protocol.get_output_func(protocol.keys.groundtruth_warped_image), + actual_func=protocol.get_output_func(protocol.keys.predicted_warped_image)) + elif self == LossTerm.full_grid_change: + return L1Loss( + expected_func=protocol.get_output_func(protocol.keys.groundtruth_grid_change), + actual_func=protocol.get_output_func(protocol.keys.predicted_grid_change)) + elif self == LossTerm.full_color_change: + return L1Loss( + expected_func=protocol.get_output_func(protocol.keys.groundtruth_posed_image), + actual_func=protocol.get_output_func(protocol.keys.predicted_color_change)) + else: + raise RuntimeError(f"Unsupported loss term {self}") + + +class LossWeights: + def __init__(self, weights: Optional[Dict[LossTerm, float]] = None): + self.weights = {} + for term in LossTerm: + self.weights[term] = 0.0 + if weights is not None: + for term in LossTerm: + if term in weights: + self.weights[term] = weights[term] + + +class TrainingPhase: + def __init__(self, + num_examples_upper_bound: int, + learning_rate: float, + loss_weights: LossWeights): + self.loss_weights = loss_weights + self.learning_rate = learning_rate + self.num_examples_upper_bound = num_examples_upper_bound + + +class LearningRateFunc: + def __init__(self, phases: List[TrainingPhase], keys: List[str]): + self.phases = phases + self.keys = keys + + def make_learning_rate_dict(self, keys: List[str], value: float): + output = {} + for key in keys: + output[key] = value + return output + + def __call__(self, examples_seen_so_far: int) -> Dict[str, float]: + for i in range(len(self.phases) - 1): + if examples_seen_so_far < self.phases[i].num_examples_upper_bound: + return self.make_learning_rate_dict(self.keys, self.phases[i].learning_rate) + return self.make_learning_rate_dict(self.keys, self.phases[-1].learning_rate) + + +class LossWeightFunc: + def __init__(self, phases: List[TrainingPhase], term: LossTerm): + self.term = term + self.phases = phases + + def __call__(self, examples_seen_so_far: int) -> float: + for i in range(len(self.phases) - 1): + if examples_seen_so_far < self.phases[i].num_examples_upper_bound: + return self.phases[i].loss_weights.weights[self.term] + return self.phases[-1].loss_weights.weights[self.term] + + +class TrainingPhases: + def __init__(self, phases: List[TrainingPhase]): + assert len(phases) > 0 + for i in range(1, len(phases)): + assert phases[i - 1].num_examples_upper_bound < phases[i].num_examples_upper_bound + + self.phases = phases + + def make_learning_rate_dict(self, keys: List[str], value: float): + output = {} + for key in keys: + output[key] = value + return output + + def get_learning_rate_func(self, keys: List[str]): + return LearningRateFunc(self.phases, keys) + + def get_loss_weight_func(self, term: LossTerm) -> Callable[[int], float]: + return LossWeightFunc(self.phases, term) + + +class SirenMorpher03TrainerArgs: + def __init__(self, + character_file_name: str, + pose_dataset_file_name: str, + training_phases: TrainingPhases, + num_training_examples_per_checkpoint: int = 100_000, + num_training_examples_per_sample_output: Optional[int] = 10_000, + num_training_examples_per_snapshot: int = 10_000, + total_batch_size: int = 8, + training_random_seed: int = 2965603729, + sample_output_random_seed: int = 3522651501, + total_worker: int = 8, + poser_func: Optional[Callable[[], Poser]] = None, + sample_output_batch_size: Optional[int] = None, + pretrained_module_file_name: Optional[str] = None): + for phase in training_phases.phases: + assert phase.num_examples_upper_bound % num_training_examples_per_checkpoint == 0 + + if poser_func is None: + poser_func = get_poser + + self.training_phases = training_phases + self.pretrained_module_file_name = pretrained_module_file_name + self.sample_output_batch_size = sample_output_batch_size + self.poser_func = poser_func + self.total_worker = total_worker + self.num_training_examples_per_snapshot = num_training_examples_per_snapshot + self.num_training_examples_per_sample_output = num_training_examples_per_sample_output + self.sample_output_random_seed = sample_output_random_seed + self.training_random_seed = training_random_seed + self.total_batch_size = total_batch_size + self.num_training_examples_per_checkpoint = num_training_examples_per_checkpoint + self.pose_dataset_file_name = pose_dataset_file_name + self.character_file_name = character_file_name + + def get_character_image(self): + return extract_pytorch_image_from_filelike( + self.character_file_name, + scale=2.0, + offset=-1.0, + premultiply_alpha=True, + perform_srgb_to_linear=True) + + def get_training_dataset(self): + return ImagePosesAndOtherImagesDataset( + main_image_func=self.get_character_image, + pose_dataset=LazyTensorDataset(self.pose_dataset_file_name), + other_image_funcs=[]) + + def get_module_factory(self): + return SirenMorpher03Factory( + SirenMorpher03Args( + image_size=512, + image_channels=4, + pose_size=45, + level_args=[ + SirenMorpherLevelArgs( + image_size=128, + intermediate_channels=360, + num_sine_layers=3), + SirenMorpherLevelArgs( + image_size=256, + intermediate_channels=180, + num_sine_layers=3), + SirenMorpherLevelArgs( + image_size=512, + intermediate_channels=90, + num_sine_layers=3), + ])) + + def get_training_computation_protocol(self): + return SirenMorpherComputationProtocol03( + indices=SirenMorpherProtocol03Indices( + batch_image=0, + batch_pose=1, + batch_face_mask=2)) + + def get_optimizer_factories(self): + return { + KEY_MODULE: AdamOptimizerFactory(betas=(0.9, 0.999)), + } + + def get_poser(self): + return self.poser_func() + + def get_training_protocol(self, world_size: int): + total_examples = self.training_phases.phases[-1].num_examples_upper_bound + per_checkpoint_examples = self.num_training_examples_per_checkpoint + num_checkpoints = total_examples // per_checkpoint_examples + batch_size = self.total_batch_size // world_size + return SirenMorpherTrainingProtocol03( + check_point_examples=[per_checkpoint_examples * (i + 1) for i in range(num_checkpoints)], + batch_size=batch_size, + learning_rate=self.training_phases.get_learning_rate_func([KEY_MODULE]), + optimizer_factories=self.get_optimizer_factories(), + random_seed=self.training_random_seed, + poser_func=self.get_poser, + key_module=KEY_MODULE, + key_poser=KEY_POSER) + + def get_sample_output_protocol(self): + return SirenMorpherSampleOutputProtocol( + num_images=4, + image_size=512, + examples_per_sample_output=self.num_training_examples_per_sample_output, + computation_protocol=self.get_training_computation_protocol(), + poser_func=self.get_poser, + random_seed=self.sample_output_random_seed, + batch_size=self.sample_output_batch_size, + batch_pose_index=1, + batch_image_index=0) + + def get_loss(self): + protocol = self.get_training_computation_protocol() + losses = [] + for term in LossTerm: + base_loss = term.get_loss(protocol) + loss = TimeDependentlyWeightedLoss( + base_loss, + examples_seen_so_far_func=lambda state: state.outputs[KEY_EXAMPLES_SEEN_SO_FAR], + weight_func=self.training_phases.get_loss_weight_func(term)) + losses.append((term.name, loss)) + return SumLoss(losses) + + def create_trainer(self, prefix: str, world_size: int, distrib_backend: str = 'gloo'): + if self.num_training_examples_per_sample_output is not None: + sample_output_protocol = self.get_sample_output_protocol() + else: + sample_output_protocol = None + + pretrained_module_file_names = {} + if self.pretrained_module_file_name is not None: + pretrained_module_file_names[KEY_MODULE] = self.pretrained_module_file_name + + return DistributedTrainer( + prefix=prefix, + module_factories={ + KEY_MODULE: self.get_module_factory(), + }, + accumulators={}, + losses={ + KEY_MODULE: self.get_loss(), + }, + training_dataset=self.get_training_dataset(), + validation_dataset=self.get_training_dataset(), + training_protocol=self.get_training_protocol(world_size), + validation_protocol=None, + sample_output_protocol=sample_output_protocol, + pretrained_module_file_names=pretrained_module_file_names, + example_per_snapshot=self.num_training_examples_per_snapshot, + num_data_loader_workers=max(1, self.total_worker // world_size), + distrib_backend=distrib_backend) diff --git a/tha4/src/tha4/nn/siren/morpher/siren_morpher_protocols_03.py b/tha4/src/tha4/nn/siren/morpher/siren_morpher_protocols_03.py new file mode 100644 index 0000000..c366e85 --- /dev/null +++ b/tha4/src/tha4/nn/siren/morpher/siren_morpher_protocols_03.py @@ -0,0 +1,352 @@ +from dataclasses import dataclass +from typing import Optional, List, Callable, Dict, Any + +import torch +from tha4.shion.base.dataset.util import get_indexed_batch + +from tha4.shion.core.cached_computation import output_array_indexing_func, add_step, ComputationState, \ + CachedComputationProtocol, ComposableCachedComputationProtocol, batch_indexing_func, proxy_func +from tha4.shion.core.loss import Loss +from tha4.shion.core.optimizer_factory import OptimizerFactory +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.shion.core.training.training_protocol import AbstractTrainingProtocol +from tha4.nn.image_processing_util import GridChangeApplier +from tha4.nn.siren.morpher.siren_morpher_03 import SirenMorpher03 +from tha4.poser.general_poser_02 import GeneralPoser02 +from tha4.sampleoutput.sample_image_creator import SampleImageSpec, ImageSource, ImageType, SampleImageSaver +from torch.nn import Module +from torch.optim import Optimizer +from torch.utils.data import Dataset + +KEY_MODULE = "module" +KEY_POSER = "poser" +KEY_EXAMPLES_SEEN_SO_FAR = "examples_seen_so_far" + + +@dataclass +class SirenMorpherProtocol03Keys: + module: str = KEY_MODULE + module_output: str = "module_output" + + poser: str = KEY_POSER + poser_output: str = "poser_output" + + image: str = "image" + pose: str = "pose" + face_mask: str = 'face_mask' + + groundtruth_posed_image: str = 'groundtruth_posed_image' + groundtruth_grid_change: str = 'groundtruth_grid_change' + groundtruth_posed_face_mask: str = 'groundtruth_posed_face_mask' + predicted_posed_image: str = 'predicted_posed_image' + + module_input_image: str = "module_input_image" + + predicted_grid_change: str = "predicted_grid_change" + predicted_color_change: str = "predicted_color_change" + predicted_warped_image: str = "predicted_warped_image" + predicted_alpha: str = "predicted_alpha" + + groundtruth_alpha: str = "groundtruth_alpha" + groundtruth_warped_image: str = "groundtruth_warped_image" + + zero: str = "zero" + + +@dataclass +class SirenMorpherProtocol03Indices: + batch_image: int = 0 + batch_face_mask: int = 1 + batch_pose: int = 2 + poser_posed_image: int = 0 + poser_grid_change: int = 3 + + poser_output_module_input_image_index: int = 5 + poser_alpha: int = 1 + poser_warped_image: int = 2 + + module_blended_image: int = SirenMorpher03.INDEX_BLENDED_IMAGE + module_grid_change: int = SirenMorpher03.INDEX_GRID_CHANGE + module_color_change: int = SirenMorpher03.INDEX_COLOR_CHANGE + module_warped_image: int = SirenMorpher03.INDEX_WARPED_IMAGE + module_alpha: int = SirenMorpher03.INDEX_ALPHA + + +class SirenMorpherComputationProtocol03(ComposableCachedComputationProtocol): + def __init__(self, + keys: Optional[SirenMorpherProtocol03Keys] = None, + indices: Optional[SirenMorpherProtocol03Indices] = None): + super().__init__() + + if keys is None: + keys = SirenMorpherProtocol03Keys() + if indices is None: + indices = SirenMorpherProtocol03Indices() + + self.keys = keys + self.indices = indices + + self.computation_steps[keys.image] = batch_indexing_func(indices.batch_image) + self.computation_steps[keys.pose] = batch_indexing_func(indices.batch_pose) + self.computation_steps[keys.face_mask] = batch_indexing_func(indices.batch_face_mask) + self.grid_change_applier = GridChangeApplier() + + @add_step(self.computation_steps, keys.module_output) + def get_module_output(protocol: CachedComputationProtocol, state: ComputationState): + pose = protocol.get_output(keys.pose, state) + module = state.modules[self.keys.module] + return module.forward(pose) + + self.computation_steps[keys.predicted_posed_image] = proxy_func(keys.module_output) + + @add_step(self.computation_steps, keys.poser_output) + def get_poser_output(protocol: CachedComputationProtocol, state: ComputationState): + with torch.no_grad(): + poser = state.modules[keys.poser] + pose = protocol.get_output(keys.pose, state) + image = protocol.get_output(keys.image, state) + return poser.get_posing_outputs(image, pose) + + @add_step(self.computation_steps, keys.groundtruth_posed_image) + def get_groundtruth_posed_image(protocol: CachedComputationProtocol, state: ComputationState): + return protocol.get_output(keys.poser_output, state)[indices.poser_posed_image] + + @add_step(self.computation_steps, keys.groundtruth_grid_change) + def get_groundtruth_posed_image(protocol: CachedComputationProtocol, state: ComputationState): + return protocol.get_output(keys.poser_output, state)[indices.poser_grid_change] + + @add_step(self.computation_steps, keys.groundtruth_posed_face_mask) + def get_groundtruth_posed_face_mask(protocol: CachedComputationProtocol, state: ComputationState): + face_mask = protocol.get_output(keys.face_mask, state) + groundtruth_grid_change = protocol.get_output(keys.groundtruth_grid_change, state) + with torch.no_grad(): + return self.grid_change_applier.apply(groundtruth_grid_change, face_mask) + + @add_step(self.computation_steps, keys.module_input_image) + def get_module_input_image(protocol: CachedComputationProtocol, state: ComputationState): + poser_output = protocol.get_output(keys.poser_output, state) + return poser_output[indices.poser_output_module_input_image_index] + + @add_step(self.computation_steps, keys.module_output) + def get_module_output(protocol: CachedComputationProtocol, state: ComputationState): + image = protocol.get_output(keys.module_input_image, state) + pose = protocol.get_output(keys.pose, state) + module = state.modules[self.keys.module] + return module.forward(image, pose) + + self.computation_steps[keys.predicted_posed_image] = output_array_indexing_func( + keys.module_output, indices.module_blended_image) + self.computation_steps[keys.predicted_grid_change] = output_array_indexing_func( + keys.module_output, indices.module_grid_change) + self.computation_steps[keys.predicted_color_change] = output_array_indexing_func( + keys.module_output, indices.module_color_change) + self.computation_steps[keys.predicted_warped_image] = output_array_indexing_func( + keys.module_output, indices.module_warped_image) + self.computation_steps[keys.predicted_alpha] = output_array_indexing_func( + keys.module_output, indices.module_alpha) + + self.computation_steps[keys.groundtruth_alpha] = output_array_indexing_func( + keys.poser_output, indices.poser_alpha) + self.computation_steps[keys.groundtruth_warped_image] = output_array_indexing_func( + keys.poser_output, indices.poser_warped_image) + + @add_step(self.computation_steps, keys.zero) + def get_zero(protocol: CachedComputationProtocol, state: ComputationState): + pose = protocol.get_output(keys.pose, state) + device = pose.device + return torch.zeros(1, device=device) + + +class SirenMorpherTrainingProtocol03(AbstractTrainingProtocol): + def __init__(self, + check_point_examples: List[int], + batch_size: int, + learning_rate: Callable[[int], Dict[str, float]], + optimizer_factories: Dict[str, OptimizerFactory], + random_seed: int, + poser_func: Callable[[], GeneralPoser02], + key_module: str, + key_poser: str = KEY_POSER, + key_examples_seen_so_far: str = KEY_EXAMPLES_SEEN_SO_FAR): + super().__init__(check_point_examples, batch_size, learning_rate, optimizer_factories, random_seed) + self.key_examples_seen_so_far = key_examples_seen_so_far + self.key_poser = key_poser + self.key_module = key_module + self.poser_func = poser_func + self.poser = None + + def run_training_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer], + losses: Dict[str, Loss], + create_log_func: Optional[Callable[[str, int], Callable[[str, float], None]]], + device: torch.device): + if self.poser is None: + self.poser = self.poser_func() + self.poser.to(device) + + module = modules[self.key_module] + module.train(True) + module_optimizer = optimizers[self.key_module] + module_optimizer.zero_grad(set_to_none=True) + + loss = losses[self.key_module] + if create_log_func is not None: + log_func = create_log_func(f"training_{self.key_module}", examples_seen_so_far) + else: + log_func = None + state = ComputationState( + modules={ + **modules, + self.key_poser: self.poser, + }, + accumulated_modules=accumulated_modules, + batch=batch, + outputs={ + self.key_examples_seen_so_far: examples_seen_so_far, + }) + loss_value = loss.compute(state, log_func) + loss_value.backward() + module_optimizer.step() + + +class SirenMorpherSampleOutputProtocol(SampleOutputProtocol): + def __init__(self, + num_images: int, + image_size: int, + examples_per_sample_output: int, + computation_protocol, + poser_func: Callable[[], GeneralPoser02], + random_seed: int = 54859395058, + batch_image_index: int = 0, + batch_pose_index: int = 2, + batch_size: Optional[int] = None, + sample_image_specs: Optional[List[SampleImageSpec]] = None, + cell_size: Optional[int] = None): + if batch_size is None: + batch_size = num_images + + if sample_image_specs is None: + sample_image_specs = [ + SampleImageSpec( + ImageSource.BATCH, + computation_protocol.indices.poser_posed_image, + ImageType.COLOR), + SampleImageSpec( + ImageSource.OUTPUT, + computation_protocol.indices.module_blended_image, + ImageType.COLOR), + SampleImageSpec( + ImageSource.OUTPUT, + computation_protocol.indices.module_alpha, + ImageType.ALPHA), + SampleImageSpec( + ImageSource.OUTPUT, + computation_protocol.indices.module_color_change, + ImageType.COLOR), + SampleImageSpec( + ImageSource.OUTPUT, + computation_protocol.indices.module_warped_image, + ImageType.COLOR), + SampleImageSpec( + ImageSource.BATCH, + computation_protocol.indices.poser_grid_change, + ImageType.GRID_CHANGE), + SampleImageSpec( + ImageSource.OUTPUT, + computation_protocol.indices.module_grid_change, + ImageType.GRID_CHANGE), + ] + + if cell_size is None: + cell_size = image_size + + self.batch_size = batch_size + self.batch_pose_index = batch_pose_index + self.batch_image_index = batch_image_index + self.poser_func = poser_func + self.random_seed = random_seed + self.examples_per_sample_output = examples_per_sample_output + self.image_size = image_size + self.num_images = num_images + self.computation_protocol = computation_protocol + self.cell_size = cell_size + + self.sample_image_saver = SampleImageSaver(image_size, cell_size, 4, sample_image_specs) + + def get_examples_per_sample_output(self) -> int: + return self.examples_per_sample_output + + def get_random_seed(self) -> int: + return self.random_seed + + def get_sample_output_data(self, validation_dataset: Dataset, device: torch.device) -> dict: + example_indices = torch.randint(0, len(validation_dataset), (self.num_images,)) + example_indices = [example_indices[i].item() for i in range(self.num_images)] + batch = get_indexed_batch(validation_dataset, example_indices, device) + poser = self.poser_func() + poser.to(device) + with torch.no_grad(): + ground_truth = poser.get_posing_outputs(batch[self.batch_image_index], batch[self.batch_pose_index]) + return { + 'batch': batch, + 'ground_truth': ground_truth + } + + def save_sample_output_data(self, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + sample_output_data: Any, + prefix: str, + examples_seen_so_far: int, + device: torch.device): + batch = sample_output_data['batch'] + ground_truth = sample_output_data['ground_truth'] + + module = modules[self.computation_protocol.keys.module] + module.train(False) + if self.batch_size == self.num_images: + with torch.no_grad(): + state = ComputationState( + modules=modules, + accumulated_modules=accumulated_modules, + batch=batch, + outputs={ + self.computation_protocol.keys.poser_output: ground_truth, + }) + module_outputs = self.computation_protocol.get_output( + self.computation_protocol.keys.module_output, state) + else: + module_outputs_list = [] + start = 0 + while start < self.num_images: + end = start + self.batch_size + end = min(self.num_images, end) + minibatch = [batch[i][start:end] for i in range(len(batch))] + ground_truth_batch = [ground_truth[i][start:end] for i in range(len(ground_truth))] + state = ComputationState( + modules=modules, + accumulated_modules=accumulated_modules, + batch=minibatch, + outputs={ + self.computation_protocol.keys.poser_output: ground_truth_batch + }) + with torch.no_grad(): + module_outputs = self.computation_protocol.get_output( + self.computation_protocol.keys.module_output, state) + module_outputs_list.append(module_outputs) + start = end + + module_outputs = [] + for i in range(len(module_outputs_list[0])): + tensor_list = [] + for j in range(len(module_outputs_list)): + tensor_list.append(module_outputs_list[j][i]) + module_output = torch.cat(tensor_list, dim=0) + module_outputs.append(module_output) + + self.sample_image_saver.save_sample_output_data(ground_truth, module_outputs, prefix, examples_seen_so_far) diff --git a/tha4/src/tha4/nn/siren/vanilla/__init__.py b/tha4/src/tha4/nn/siren/vanilla/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/siren/vanilla/siren.py b/tha4/src/tha4/nn/siren/vanilla/siren.py new file mode 100644 index 0000000..b951e11 --- /dev/null +++ b/tha4/src/tha4/nn/siren/vanilla/siren.py @@ -0,0 +1,100 @@ +import math +from typing import Callable, Optional, List + +import torch +from torch import Tensor +from torch.nn import Module, Conv2d, ModuleList + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.nn00.initialization_funcs import HeInitialization + + +class SineLinearLayer(Module): + def __init__(self, + in_channels: int, + out_channels: int, + is_first=False, + omega_0=30.0): + super().__init__() + self.out_channels = out_channels + self.in_channels = in_channels + self.omega_0 = omega_0 + self.is_first = is_first + self.linear = Conv2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=True) + with torch.no_grad(): + if self.is_first: + self.linear.weight.uniform_(-1 / in_channels, 1.0 / in_channels) + else: + self.linear.weight.uniform_( + -math.sqrt(6.0 / in_channels) / self.omega_0, + math.sqrt(6.0 / in_channels) / self.omega_0) + + def forward(self, x: Tensor): + return torch.sin(self.omega_0 * self.linear(x)) + + +class SirenArgs: + def __init__( + self, + in_channels: int, + out_channels: int, + intermediate_channels: int, + num_sine_layers: int, + use_tanh: bool = False, + init_func: Optional[Callable[[Module], Module]] = None): + if init_func is None: + init_func = HeInitialization() + self.init_func = init_func + self.use_tanh = use_tanh + assert num_sine_layers >= 1 + self.intermediate_channels = intermediate_channels + self.num_sine_layers = num_sine_layers + self.out_channels = out_channels + self.in_channels = in_channels + + +class Siren(Module): + def __init__(self, args: SirenArgs): + super().__init__() + self.args = args + self.sine_layers = ModuleList() + self.sine_layers.append( + SineLinearLayer( + in_channels=args.in_channels, out_channels=args.intermediate_channels, is_first=True)) + for i in range(args.num_sine_layers - 1): + self.sine_layers.append( + SineLinearLayer( + in_channels=args.intermediate_channels, + out_channels=args.intermediate_channels, + is_first=False)) + self.last_linear = args.init_func(Conv2d( + args.intermediate_channels, + args.out_channels, + kernel_size=1, + stride=1, + padding=0, + bias=True)) + + def forward(self, x: Tensor) -> Tensor: + for i in range(self.args.num_sine_layers): + x = self.sine_layers[i].forward(x) + x = self.last_linear(x) + if self.args.use_tanh: + return torch.tanh(x) + else: + return x + + +class SirenFactory(ModuleFactory): + def __init__(self, args: SirenArgs): + super().__init__() + self.args = args + + def create(self) -> Module: + return Siren(self.args) diff --git a/tha4/src/tha4/nn/spectral_norm.py b/tha4/src/tha4/nn/spectral_norm.py new file mode 100644 index 0000000..4f08dd4 --- /dev/null +++ b/tha4/src/tha4/nn/spectral_norm.py @@ -0,0 +1,9 @@ +from torch.nn import Module +from torch.nn.utils import spectral_norm + + +def apply_spectral_norm(module: Module, use_spectrial_norm: bool = False) -> Module: + if use_spectrial_norm: + return spectral_norm(module) + else: + return module diff --git a/tha4/src/tha4/nn/upscaler/__init__.py b/tha4/src/tha4/nn/upscaler/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/nn/upscaler/upscaler_02.py b/tha4/src/tha4/nn/upscaler/upscaler_02.py new file mode 100644 index 0000000..bf2de92 --- /dev/null +++ b/tha4/src/tha4/nn/upscaler/upscaler_02.py @@ -0,0 +1,110 @@ +from typing import List + +import torch +from torch import Tensor, zero_ +from torch.nn import Module, Conv2d + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.image_processing_util import GridChangeApplier +from tha4.nn.common.unet import UnetArgs, AttentionBlockArgs, UnetWithFirstConvAddition + + +class Upscaler02Args: + def __init__(self, + image_size: int, + image_channels: int, + num_pose_parameters: int, + unet_args: UnetArgs): + assert unet_args.in_channels == ( + image_channels + ) + assert unet_args.out_channels == ( + image_channels + # direct + 2 + # warp + 1 # alpha + ) + assert unet_args.cond_input_channels == num_pose_parameters + self.image_channels = image_channels + self.image_size = image_size + self.num_pose_parameters = num_pose_parameters + self.unet_args = unet_args + + +def apply_color_change(alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) + + +class Upscaler02(Module): + def __init__(self, args: Upscaler02Args): + super().__init__() + self.args = args + self.body = UnetWithFirstConvAddition(args.unet_args) + self.grid_change_applier = GridChangeApplier() + self.coarse_image_conv = Conv2d( + args.image_channels + args.image_channels + 2, + args.unet_args.model_channels, + kernel_size=3, + stride=1, + padding=1) + with torch.no_grad(): + zero_(self.coarse_image_conv.weight) + zero_(self.coarse_image_conv.bias) + + def check_image(self, image: torch.Tensor): + assert len(image.shape) == 4 + assert image.shape[1] == self.args.image_channels + assert image.shape[2] == self.args.image_size + assert image.shape[3] == self.args.image_size + + def forward(self, + rest_image: torch.Tensor, + coarse_posed_image: torch.Tensor, + coarse_grid_change: torch.Tensor, + pose: torch.Tensor) -> List[Tensor]: + self.check_image(rest_image) + self.check_image(coarse_posed_image) + + assert len(pose.shape) == 2 + assert rest_image.shape[0] == pose.shape[0] + assert coarse_posed_image.shape[0] == pose.shape[0] + assert coarse_grid_change.shape[0] == pose.shape[0] + assert coarse_grid_change.shape[1] == 2 + assert coarse_grid_change.shape[2] == self.args.image_size + assert coarse_grid_change.shape[3] == self.args.image_size + assert pose.shape[1] == self.args.num_pose_parameters + + warped_image = self.grid_change_applier.apply(coarse_grid_change, rest_image) + + t = torch.zeros(rest_image.shape[0], 1, device=rest_image.device) + feature = torch.cat([coarse_posed_image, warped_image, coarse_grid_change], dim=1) + first_conv_addition = self.coarse_image_conv(feature) + + body_output = self.body(rest_image, t, pose, first_conv_addition) + + direct = body_output[:, 0:self.args.image_channels, :, :] + grid_change = body_output[:, self.args.image_channels:self.args.image_channels + 2, :, :] + alpha = torch.sigmoid(body_output[:, self.args.image_channels + 2:self.args.image_channels + 3, :, :]) + warped = self.grid_change_applier.apply(grid_change, rest_image) + merged = apply_color_change(alpha, direct, warped) + + return [ + merged, + alpha, + warped, + grid_change, + direct + ] + + INDEX_MERGED = 0 + INDEX_ALPHA = 1 + INDEX_WARPED = 2 + INDEX_GRID_CHANGE = 3 + INDEX_DIRECT = 4 + + +class Upscaler02Factory(ModuleFactory): + def __init__(self, args: Upscaler02Args): + self.args = args + + def create(self) -> Module: + return Upscaler02(self.args) diff --git a/tha4/src/tha4/nn/util.py b/tha4/src/tha4/nn/util.py new file mode 100644 index 0000000..9f46989 --- /dev/null +++ b/tha4/src/tha4/nn/util.py @@ -0,0 +1,40 @@ +from typing import Optional, Callable, Union + +from torch.nn import Module + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.nn.init_function import create_init_function +from tha4.nn.nonlinearity_factory import resolve_nonlinearity_factory +from tha4.nn.normalization import NormalizationLayerFactory +from tha4.nn.spectral_norm import apply_spectral_norm + + +def wrap_conv_or_linear_module(module: Module, + initialization_method: Union[str, Callable[[Module], Module]], + use_spectral_norm: bool): + if isinstance(initialization_method, str): + init = create_init_function(initialization_method) + else: + init = initialization_method + return apply_spectral_norm(init(module), use_spectral_norm) + + +class BlockArgs: + def __init__(self, + initialization_method: Union[str, Callable[[Module], Module]] = 'he', + use_spectral_norm: bool = False, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + nonlinearity_factory: Optional[ModuleFactory] = None): + self.nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + self.normalization_layer_factory = normalization_layer_factory + self.use_spectral_norm = use_spectral_norm + self.initialization_method = initialization_method + + def wrap_module(self, module: Module) -> Module: + return wrap_conv_or_linear_module(module, self.get_init_func(), self.use_spectral_norm) + + def get_init_func(self) -> Callable[[Module], Module]: + if isinstance(self.initialization_method, str): + return create_init_function(self.initialization_method) + else: + return self.initialization_method diff --git a/tha4/src/tha4/poser/__init__.py b/tha4/src/tha4/poser/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/poser/general_poser_02.py b/tha4/src/tha4/poser/general_poser_02.py new file mode 100644 index 0000000..7535451 --- /dev/null +++ b/tha4/src/tha4/poser/general_poser_02.py @@ -0,0 +1,98 @@ +from typing import List, Optional, Tuple, Dict, Callable + +import torch +from tha4.shion.core.cached_computation import ComputationState +from tha4.poser.poser import PoseParameterGroup, Poser +from torch import Tensor +from torch.nn import Module + + +class GeneralPoser02(Poser): + def __init__(self, + module_loaders: Dict[str, Callable[[], Module]], + device: torch.device, + output_length: int, + pose_parameters: List[PoseParameterGroup], + output_list_func: Callable[[ComputationState], List[Tensor]], + subrect: Optional[Tuple[Tuple[int, int], Tuple[int, int]]] = None, + default_output_index: int = 0, + image_size: int = 256, + dtype: torch.dtype = torch.float): + self.dtype = dtype + self.image_size = image_size + self.default_output_index = default_output_index + self.output_list_func = output_list_func + self.subrect = subrect + self.pose_parameters = pose_parameters + self.device = device + self.module_loaders = module_loaders + + self.modules = None + + self.num_parameters = 0 + for pose_parameter in self.pose_parameters: + self.num_parameters += pose_parameter.get_arity() + + self.output_length = output_length + + def get_image_size(self) -> int: + return self.image_size + + def get_modules(self): + if self.modules is None: + self.modules = {} + for key in self.module_loaders: + module = self.module_loaders[key]() + self.modules[key] = module + module.to(self.device) + module.train(False) + return self.modules + + def get_pose_parameter_groups(self) -> List[PoseParameterGroup]: + return self.pose_parameters + + def get_num_parameters(self) -> int: + return self.num_parameters + + def pose(self, image: Tensor, pose: Tensor, output_index: Optional[int] = None) -> Tensor: + if output_index is None: + output_index = self.default_output_index + output_list = self.get_posing_outputs(image, pose) + return output_list[output_index] + + def get_posing_outputs(self, image: Tensor, pose: Tensor) -> List[Tensor]: + modules = self.get_modules() + + if len(image.shape) == 3: + image = image.unsqueeze(0) + if len(pose.shape) == 1: + pose = pose.unsqueeze(0) + if self.subrect is not None: + image = image[:, :, self.subrect[0][0]:self.subrect[0][1], self.subrect[1][0]:self.subrect[1][1]] + batch = [image, pose] + + state = ComputationState( + modules=modules, + accumulated_modules={}, + batch=batch, + outputs={}) + return self.output_list_func(state) + + def get_output_length(self) -> int: + return self.output_length + + def free(self): + self.modules = None + + def get_dtype(self) -> torch.dtype: + return self.dtype + + def to(self, device: torch.device) -> 'GeneralPoser02': + if device == self.device: + return self + modules = self.get_modules() + self.device = device + for key in modules: + module = modules[key] + module.to(self.device) + return self diff --git a/tha4/src/tha4/poser/modes/__init__.py b/tha4/src/tha4/poser/modes/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/poser/modes/mode_07.py b/tha4/src/tha4/poser/modes/mode_07.py new file mode 100644 index 0000000..e7e7fce --- /dev/null +++ b/tha4/src/tha4/poser/modes/mode_07.py @@ -0,0 +1,315 @@ +from enum import Enum +from typing import List, Dict, Optional + +import torch +from tha4.shion.core.cached_computation import CachedComputationProtocol, ComputationState +from tha4.shion.core.load_save import torch_load +from tha4.nn.eyebrow_decomposer.eyebrow_decomposer_00 import EyebrowDecomposer00, \ + EyebrowDecomposer00Factory, EyebrowDecomposer00Args +from tha4.nn.eyebrow_morphing_combiner.eyebrow_morphing_combiner_00 import \ + EyebrowMorphingCombiner00Factory, EyebrowMorphingCombiner00Args, EyebrowMorphingCombiner00 +from tha4.nn.face_morpher.face_morpher_08 import FaceMorpher08Args, FaceMorpher08Factory +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs +from tha4.nn.common.unet import UnetArgs, AttentionBlockArgs +from tha4.nn.morpher.morpher_00 import Morpher00Args, Morpher00 +from tha4.nn.upscaler.upscaler_02 import Upscaler02Args, Upscaler02 +from tha4.poser.general_poser_02 import GeneralPoser02 +from tha4.poser.modes.pose_parameters import get_pose_parameters +from torch import Tensor +from torch.nn.functional import interpolate + + +class Network(Enum): + eyebrow_decomposer = 1 + eyebrow_morphing_combiner = 2 + face_morpher = 3 + body_morpher = 4 + upscaler = 5 + + @property + def outputs_key(self): + return f"{self.name}_outputs" + + +class Branch(Enum): + face_morphed_half = 1 + face_morphed_full = 2 + all_outputs = 3 + + +NUM_EYEBROW_PARAMS = 12 +NUM_FACE_PARAMS = 27 +NUM_ROTATION_PARAMS = 6 + + +class FiveStepPoserComputationProtocol(CachedComputationProtocol): + def __init__(self, eyebrow_morphed_image_index: int): + super().__init__() + self.eyebrow_morphed_image_index = eyebrow_morphed_image_index + self.cached_batch_0 = None + self.cached_eyebrow_decomposer_output = None + + def compute_func(self): + def func(state: ComputationState) -> List[Tensor]: + if self.cached_batch_0 is None: + new_batch_0 = True + elif state.batch[0].shape[0] != self.cached_batch_0.shape[0]: + new_batch_0 = True + else: + new_batch_0 = torch.max((state.batch[0] - self.cached_batch_0).abs()).item() > 0 + if not new_batch_0: + state.outputs[Network.eyebrow_decomposer.outputs_key] = self.cached_eyebrow_decomposer_output + output = self.get_output(Branch.all_outputs.name, state) + if new_batch_0: + self.cached_batch_0 = state.batch[0] + self.cached_eyebrow_decomposer_output = state.outputs[Network.eyebrow_decomposer.outputs_key] + return output + + return func + + def compute_output(self, key: str, state: ComputationState) -> List[Tensor]: + if key == Network.eyebrow_decomposer.outputs_key: + input_image = state.batch[0][:, :, 64:192, 64 + 128:192 + 128] + return state.modules[Network.eyebrow_decomposer.name].forward(input_image) + elif key == Network.eyebrow_morphing_combiner.outputs_key: + eyebrow_decomposer_output = self.get_output(Network.eyebrow_decomposer.outputs_key, state) + background_layer = eyebrow_decomposer_output[EyebrowDecomposer00.BACKGROUND_LAYER_INDEX] + eyebrow_layer = eyebrow_decomposer_output[EyebrowDecomposer00.EYEBROW_LAYER_INDEX] + eyebrow_pose = state.batch[1][:, :NUM_EYEBROW_PARAMS] + return state.modules[Network.eyebrow_morphing_combiner.name].forward( + background_layer, + eyebrow_layer, + eyebrow_pose) + elif key == Network.face_morpher.outputs_key: + eyebrow_morphing_combiner_output = self.get_output( + Network.eyebrow_morphing_combiner.outputs_key, state) + eyebrow_morphed_image = eyebrow_morphing_combiner_output[self.eyebrow_morphed_image_index] + input_image = state.batch[0][:, :, 32:32 + 192, (32 + 128):(32 + 192 + 128)].clone() + input_image[:, :, 32:32 + 128, 32:32 + 128] = eyebrow_morphed_image + face_pose = state.batch[1][:, NUM_EYEBROW_PARAMS:NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS] + return state.modules[Network.face_morpher.name].forward(input_image, face_pose) + elif key == Branch.face_morphed_full.name: + face_morpher_output = self.get_output(Network.face_morpher.outputs_key, state) + face_morphed_image = face_morpher_output[0] + input_image = state.batch[0].clone() + input_image[:, :, 32:32 + 192, 32 + 128:32 + 192 + 128] = face_morphed_image + return [input_image] + elif key == Branch.face_morphed_half.name: + face_morphed_full = self.get_output(Branch.face_morphed_full.name, state)[0] + return [ + interpolate(face_morphed_full, size=(256, 256), mode='bilinear', align_corners=False) + ] + elif key == Network.body_morpher.outputs_key: + face_morphed_half = self.get_output(Branch.face_morphed_half.name, state)[0] + rotation_pose = state.batch[1][:, NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS:] + return state.modules[Network.body_morpher.name].forward(face_morphed_half, rotation_pose) + elif key == Network.upscaler.outputs_key: + rest_image = self.get_output(Branch.face_morphed_full.name, state)[0] + body_morpher_outputs = self.get_output( + Network.body_morpher.outputs_key, state) + half_res_posed_image = body_morpher_outputs[Morpher00.INDEX_MERGED] + half_res_grid_change = body_morpher_outputs[Morpher00.INDEX_GRID_CHANGE] + coarse_posed_image = interpolate(half_res_posed_image, size=(512, 512), mode='bilinear') + coarse_grid_change = interpolate(half_res_grid_change, size=(512, 512), mode='bilinear') + rotation_pose = state.batch[1][:, NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS:] + return state.modules[Network.upscaler.name].forward( + rest_image, coarse_posed_image, coarse_grid_change, rotation_pose) + elif key == Branch.all_outputs.name: + upscaler_output = self.get_output(Network.upscaler.outputs_key, state) + face_morphed_full = self.get_output(Branch.face_morphed_full.name, state) + body_morpher_output = self.get_output(Network.body_morpher.outputs_key, state) + face_morpher_output = self.get_output(Network.face_morpher.outputs_key, state) + eyebrow_morphing_combiner_output = self.get_output(Network.eyebrow_morphing_combiner.outputs_key, state) + eyebrow_decomposer_output = self.get_output(Network.eyebrow_decomposer.outputs_key, state) + output = upscaler_output \ + + face_morphed_full \ + + body_morpher_output \ + + face_morpher_output \ + + eyebrow_morphing_combiner_output \ + + eyebrow_decomposer_output + return output + else: + raise RuntimeError("Unsupported key: " + key) + + +def load_eyebrow_decomposer(file_name: str): + factory = EyebrowDecomposer00Factory( + EyebrowDecomposer00Args( + image_size=128, + image_channels=4, + start_channels=64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)))) + print("Loading the eyebrow decomposer ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + print("DONE!!!") + return module + + +def load_eyebrow_morphing_combiner(file_name: str): + factory = EyebrowMorphingCombiner00Factory( + EyebrowMorphingCombiner00Args( + image_size=128, + image_channels=4, + start_channels=64, + num_pose_params=12, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)))) + print("Loading the eyebrow morphing conbiner ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + print("DONE!!!") + return module + + +def load_face_morpher(file_name: str): + factory = FaceMorpher08Factory( + FaceMorpher08Args( + image_size=192, + image_channels=4, + num_expression_params=27, + start_channels=64, + bottleneck_image_size=24, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=False) + ), + output_iris_mouth_grid_change=True, + ) + ) + print("Loading the face morpher ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + print("DONE!!!") + return module + + +def apply_color_change(alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) + + +def load_morpher_00(file_name: str): + unet_args = UnetArgs( + in_channels=4, + out_channels=7, + model_channels=64, + level_channel_multipliers=[1, 2, 4, 4, 4], + level_use_attention=[False, False, False, False, True], + num_res_blocks_per_level=1, + num_middle_res_blocks=4, + time_embedding_channels=None, + cond_input_channels=6, + cond_internal_channels=256, + attention_block_args=AttentionBlockArgs( + num_heads=8, + use_new_attention_order=True), + dropout_prob=0.0) + morpher_00_args = Morpher00Args( + image_size=256, + image_channels=4, + num_pose_parameters=6, + unet_args=unet_args) + morpher_00 = Morpher00(morpher_00_args) + + print("Loading the body morpher ... ", end="") + morpher_00.load_state_dict(torch_load(file_name)) + print("DONE") + + morpher_00.train(False) + return morpher_00 + + +def load_upscaler_02(file_name: str): + unet_args = UnetArgs( + in_channels=4, + out_channels=7, + model_channels=32, + level_channel_multipliers=[1, 2, 4, 8, 8, 8], + level_use_attention=[False, False, False, False, False, True], + num_res_blocks_per_level=1, + num_middle_res_blocks=4, + time_embedding_channels=None, + cond_input_channels=6, + cond_internal_channels=256, + attention_block_args=AttentionBlockArgs( + num_heads=8, + use_new_attention_order=True), + dropout_prob=0.0) + upscaler_02_args = Upscaler02Args( + image_size=512, + image_channels=4, + num_pose_parameters=6, + unet_args=unet_args) + upscaler_02 = Upscaler02(upscaler_02_args) + + print("Loading the upscaler ... ", end="") + upscaler_02.load_state_dict(torch_load(file_name)) + print("DONE") + + upscaler_02.train(False) + return upscaler_02 + + +def create_poser( + device: torch.device, + module_file_names: Optional[Dict[str, str]] = None, + eyebrow_morphed_image_index: int = EyebrowMorphingCombiner00.EYEBROW_IMAGE_NO_COMBINE_ALPHA_INDEX, + default_output_index: int = 0) -> GeneralPoser02: + if module_file_names is None: + module_file_names = {} + if Network.eyebrow_decomposer.name not in module_file_names: + file_name = "data/tha4/eyebrow_decomposer.pt" + module_file_names[Network.eyebrow_decomposer.name] = file_name + if Network.eyebrow_morphing_combiner.name not in module_file_names: + file_name = "data/tha4/eyebrow_morphing_combiner.pt" + module_file_names[Network.eyebrow_morphing_combiner.name] = file_name + if Network.face_morpher.name not in module_file_names: + file_name = "data/tha4/face_morpher.pt" + module_file_names[Network.face_morpher.name] = file_name + if Network.body_morpher.name not in module_file_names: + file_name = "data/tha4/body_morpher.pt" + module_file_names[Network.body_morpher.name] = file_name + if Network.upscaler.name not in module_file_names: + file_name = "data/tha4/upscaler.pt" + module_file_names[Network.upscaler.name] = file_name + + loaders = { + Network.eyebrow_decomposer.name: + lambda: load_eyebrow_decomposer(module_file_names[Network.eyebrow_decomposer.name]), + Network.eyebrow_morphing_combiner.name: + lambda: load_eyebrow_morphing_combiner(module_file_names[Network.eyebrow_morphing_combiner.name]), + Network.face_morpher.name: + lambda: load_face_morpher(module_file_names[Network.face_morpher.name]), + Network.body_morpher.name: + lambda: load_morpher_00(module_file_names[Network.body_morpher.name]), + Network.upscaler.name: + lambda: load_upscaler_02(module_file_names[Network.upscaler.name]), + } + return GeneralPoser02( + image_size=512, + module_loaders=loaders, + pose_parameters=get_pose_parameters().get_pose_parameter_groups(), + output_list_func=FiveStepPoserComputationProtocol(eyebrow_morphed_image_index).compute_func(), + subrect=None, + device=device, + output_length=5 + 1 + 5 + 8 + 8 + 6, + default_output_index=default_output_index) diff --git a/tha4/src/tha4/poser/modes/mode_12.py b/tha4/src/tha4/poser/modes/mode_12.py new file mode 100644 index 0000000..908c1ec --- /dev/null +++ b/tha4/src/tha4/poser/modes/mode_12.py @@ -0,0 +1,202 @@ +from enum import Enum +from typing import List, Dict, Optional, Any + +import torch +from tha4.shion.core.cached_computation import CachedComputationProtocol, ComputationState +from tha4.shion.core.load_save import torch_load +from tha4.nn.eyebrow_decomposer.eyebrow_decomposer_00 import EyebrowDecomposer00, \ + EyebrowDecomposer00Factory, EyebrowDecomposer00Args +from tha4.nn.eyebrow_morphing_combiner.eyebrow_morphing_combiner_00 import \ + EyebrowMorphingCombiner00Factory, EyebrowMorphingCombiner00Args, EyebrowMorphingCombiner00 +from tha4.nn.face_morpher.face_morpher_08 import FaceMorpher08Args, FaceMorpher08Factory +from tha4.nn.nonlinearity_factory import ReLUFactory +from tha4.nn.normalization import InstanceNorm2dFactory +from tha4.nn.util import BlockArgs +from tha4.poser.general_poser_02 import GeneralPoser02 +from tha4.poser.modes.pose_parameters import get_pose_parameters +from torch import Tensor + + +class Network(Enum): + eyebrow_decomposer = 1 + eyebrow_morphing_combiner = 2 + face_morpher = 3 + + @property + def outputs_key(self): + return f"{self.name}_outputs" + + +class Branch(Enum): + face_morphed_half = 1 + face_morphed_full = 2 + all_outputs = 3 + + +NUM_EYEBROW_PARAMS = 12 +NUM_FACE_PARAMS = 27 +NUM_ROTATION_PARAMS = 6 + + +class FiveStepPoserComputationProtocol(CachedComputationProtocol): + def __init__(self, eyebrow_morphed_image_index: int): + super().__init__() + self.eyebrow_morphed_image_index = eyebrow_morphed_image_index + self.cached_batch_0 = None + self.cached_eyebrow_decomposer_output = None + + def compute_func(self): + def func(state: ComputationState) -> List[Tensor]: + if self.cached_batch_0 is None: + new_batch_0 = True + elif state.batch[0].shape[0] != self.cached_batch_0.shape[0]: + new_batch_0 = True + else: + new_batch_0 = torch.max((state.batch[0] - self.cached_batch_0).abs()).item() > 0 + if not new_batch_0: + state.outputs[Network.eyebrow_decomposer.outputs_key] = self.cached_eyebrow_decomposer_output + output = self.get_output(Branch.all_outputs.name, state) + if new_batch_0: + self.cached_batch_0 = state.batch[0] + self.cached_eyebrow_decomposer_output = state.outputs[Network.eyebrow_decomposer.outputs_key] + return output + + return func + + def compute_output(self, key: str, state: ComputationState) -> Any: + if key == Network.eyebrow_decomposer.outputs_key: + input_image = state.batch[0][:, :, 64:192, 64 + 128:192 + 128] + return state.modules[Network.eyebrow_decomposer.name].forward(input_image) + elif key == Network.eyebrow_morphing_combiner.outputs_key: + eyebrow_decomposer_output = self.get_output(Network.eyebrow_decomposer.outputs_key, state) + background_layer = eyebrow_decomposer_output[EyebrowDecomposer00.BACKGROUND_LAYER_INDEX] + eyebrow_layer = eyebrow_decomposer_output[EyebrowDecomposer00.EYEBROW_LAYER_INDEX] + eyebrow_pose = state.batch[1][:, :NUM_EYEBROW_PARAMS] + return state.modules[Network.eyebrow_morphing_combiner.name].forward( + background_layer, + eyebrow_layer, + eyebrow_pose) + elif key == Network.face_morpher.outputs_key: + eyebrow_morphing_combiner_output = self.get_output( + Network.eyebrow_morphing_combiner.outputs_key, state) + eyebrow_morphed_image = eyebrow_morphing_combiner_output[self.eyebrow_morphed_image_index] + input_image = state.batch[0][:, :, 32:32 + 192, (32 + 128):(32 + 192 + 128)].clone() + input_image[:, :, 32:32 + 128, 32:32 + 128] = eyebrow_morphed_image + face_pose = state.batch[1][:, NUM_EYEBROW_PARAMS:NUM_EYEBROW_PARAMS + NUM_FACE_PARAMS] + return state.modules[Network.face_morpher.name].forward(input_image, face_pose) + elif key == Branch.all_outputs.name: + face_morpher_output = self.get_output(Network.face_morpher.outputs_key, state) + eyebrow_morphing_combiner_output = self.get_output(Network.eyebrow_morphing_combiner.outputs_key, state) + eyebrow_decomposer_output = self.get_output(Network.eyebrow_decomposer.outputs_key, state) + output = face_morpher_output \ + + eyebrow_morphing_combiner_output \ + + eyebrow_decomposer_output + return output + else: + raise RuntimeError("Unsupported key: " + key) + + +def load_eyebrow_decomposer(file_name: str): + factory = EyebrowDecomposer00Factory( + EyebrowDecomposer00Args( + image_size=128, + image_channels=4, + start_channels=64, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)))) + print("Loading the eyebrow decomposer ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + print("DONE!!!") + return module + + +def load_eyebrow_morphing_combiner(file_name: str): + factory = EyebrowMorphingCombiner00Factory( + EyebrowMorphingCombiner00Args( + image_size=128, + image_channels=4, + start_channels=64, + num_pose_params=12, + bottleneck_image_size=16, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=True)))) + print("Loading the eyebrow morphing conbiner ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + print("DONE!!!") + return module + + +def load_face_morpher(file_name: str): + factory = FaceMorpher08Factory( + FaceMorpher08Args( + image_size=192, + image_channels=4, + num_expression_params=27, + start_channels=64, + bottleneck_image_size=24, + num_bottleneck_blocks=6, + max_channels=512, + block_args=BlockArgs( + initialization_method='he', + use_spectral_norm=False, + normalization_layer_factory=InstanceNorm2dFactory(), + nonlinearity_factory=ReLUFactory(inplace=False)), + output_iris_mouth_grid_change=True)) + print("Loading the face morpher ... ", end="") + module = factory.create() + module.load_state_dict(torch_load(file_name)) + print("DONE!!!") + return module + + +def apply_color_change(alpha, color_change, image: Tensor) -> Tensor: + return color_change * alpha + image * (1 - alpha) + + +def create_poser( + device: torch.device, + module_file_names: Optional[Dict[str, str]] = None, + eyebrow_morphed_image_index: int = EyebrowMorphingCombiner00.EYEBROW_IMAGE_NO_COMBINE_ALPHA_INDEX, + default_output_index: int = 0) -> GeneralPoser02: + if module_file_names is None: + module_file_names = {} + if Network.eyebrow_decomposer.name not in module_file_names: + file_name = "data/tha4/eyebrow_decomposer.pt" + module_file_names[Network.eyebrow_decomposer.name] = file_name + if Network.eyebrow_morphing_combiner.name not in module_file_names: + file_name = "data/tha4/eyebrow_morphing_combiner.pt" + module_file_names[Network.eyebrow_morphing_combiner.name] = file_name + if Network.face_morpher.name not in module_file_names: + file_name = "data/tha4/face_morpher.pt" + module_file_names[Network.face_morpher.name] = file_name + + loaders = { + Network.eyebrow_decomposer.name: + lambda: load_eyebrow_decomposer(module_file_names[Network.eyebrow_decomposer.name]), + Network.eyebrow_morphing_combiner.name: + lambda: load_eyebrow_morphing_combiner(module_file_names[Network.eyebrow_morphing_combiner.name]), + Network.face_morpher.name: + lambda: load_face_morpher(module_file_names[Network.face_morpher.name]), + } + return GeneralPoser02( + image_size=512, + module_loaders=loaders, + pose_parameters=get_pose_parameters().get_pose_parameter_groups(), + output_list_func=FiveStepPoserComputationProtocol(eyebrow_morphed_image_index).compute_func(), + subrect=None, + device=device, + output_length=5 + 5 + 8, + default_output_index=default_output_index) diff --git a/tha4/src/tha4/poser/modes/mode_14.py b/tha4/src/tha4/poser/modes/mode_14.py new file mode 100644 index 0000000..3fcf610 --- /dev/null +++ b/tha4/src/tha4/poser/modes/mode_14.py @@ -0,0 +1,162 @@ +from dataclasses import dataclass +from typing import List, Optional, Dict, Any + +import torch +from tha4.shion.core.cached_computation import CachedComputationProtocol, ComputationState +from tha4.shion.core.load_save import torch_load +from tha4.nn.siren.face_morpher.siren_face_morpher_00 import SirenFaceMorpher00Args, SirenFaceMorpher00 +from tha4.nn.siren.morpher.siren_morpher_03 import SirenMorpher03, SirenMorpher03Args, SirenMorpherLevelArgs +from tha4.nn.siren.vanilla.siren import SirenArgs +from tha4.poser.general_poser_02 import GeneralPoser02 +from tha4.poser.modes.pose_parameters import get_pose_parameters +from torch import Tensor + +KEY_FACE_MORPHER = "face_morpher" +KEY_BODY_MORPHER = "body_morpher" + + +@dataclass +class Keys: + face_morpher: str = KEY_FACE_MORPHER + face_morpher_output: str = "face_morpher_output" + + face_morpher_input_image: str = "face_morpher_input_image" + face_morpher_input_pose: str = "face_morpher_input_pose" + + body_morpher_input_image: str = "body_morpher_input_image" + + body_morpher: str = KEY_BODY_MORPHER + body_morpher_output: str = "body_morpher_output" + + all_outputs: str = "all_outputs" + + +@dataclass +class Indices: + original_image: int = 0 + original_pose: int = 1 + + +class TwoStepPoserComputationProtocol(CachedComputationProtocol): + def __init__(self, keys: Optional[Keys] = None, indices: Optional[Indices] = None): + super().__init__() + + if keys is None: + keys = Keys() + if indices is None: + indices = Indices() + + self.keys = keys + self.indices = indices + + def compute_func(self): + def func(state: ComputationState) -> List[Tensor]: + return self.get_output(self.keys.all_outputs, state) + + return func + + def compute_output(self, key: str, state: ComputationState) -> Any: + if key == self.keys.face_morpher_input_image: + image = state.batch[self.indices.original_image] + center_x = 256 + center_y = 128 + 16 + return image[:, :, center_y - 64:center_y + 64, center_x - 64:center_x + 64] + elif key == self.keys.face_morpher_input_pose: + pose = state.batch[self.indices.original_pose] + return pose[:, 0:39] + elif key == self.keys.face_morpher_output: + module = state.modules[self.keys.face_morpher] + pose = self.get_output(self.keys.face_morpher_input_pose, state) + with torch.no_grad(): + return module.forward(pose) + elif key == self.keys.body_morpher_input_image: + image = state.batch[self.indices.original_image].clone() + center_x = 256 + center_y = 128 + 16 + face_morphed_image = self.get_output(self.keys.face_morpher_output, state) + image[:, :, center_y - 64:center_y + 64, center_x - 64:center_x + 64] = face_morphed_image + return image + elif key == self.keys.body_morpher_output: + image = self.get_output(self.keys.body_morpher_input_image, state) + pose = state.batch[self.indices.original_pose] + body_morpher = state.modules[self.keys.body_morpher] + with torch.no_grad(): + return body_morpher.forward(image, pose) + elif key == self.keys.all_outputs: + body_morpher_output = self.get_output(self.keys.body_morpher_output, state) + face_morpher_output = self.get_output(self.keys.face_morpher_output, state) + return body_morpher_output + [face_morpher_output] + else: + raise RuntimeError("Unsupported key: " + key) + + +def load_face_morpher(file_name: Optional[str] = None): + module = SirenFaceMorpher00( + SirenFaceMorpher00Args( + image_size=128, + image_channels=4, + pose_size=39, + siren_args=SirenArgs( + in_channels=39 + 2, + out_channels=4, + intermediate_channels=128, + num_sine_layers=8))) + if file_name is not None: + module.load_state_dict(torch_load(file_name)) + return module + + +def load_body_morpher(file_name: Optional[str] = None): + module = SirenMorpher03( + SirenMorpher03Args( + image_size=512, + image_channels=4, + pose_size=45, + level_args=[ + SirenMorpherLevelArgs( + image_size=128, + intermediate_channels=360, + num_sine_layers=3), + SirenMorpherLevelArgs( + image_size=256, + intermediate_channels=180, + num_sine_layers=3), + SirenMorpherLevelArgs( + image_size=512, + intermediate_channels=90, + num_sine_layers=3), + ])) + if file_name is not None: + module.load_state_dict(torch_load(file_name)) + return module + + +def create_poser( + device: torch.device, + module_file_names: Optional[Dict[str, str]] = None, + default_output_index: int = 0) -> GeneralPoser02: + if module_file_names is None: + module_file_names = {} + if KEY_FACE_MORPHER not in module_file_names: + file_name = "data/character_models/lambda_00/face_morpher.pt" + module_file_names[KEY_FACE_MORPHER] = file_name + if KEY_BODY_MORPHER not in module_file_names: + file_name = "data/character_models/lambda_00/body_morpher.pt" + module_file_names[KEY_BODY_MORPHER] = file_name + + loaders = { + KEY_FACE_MORPHER: + lambda: load_face_morpher(module_file_names[KEY_FACE_MORPHER]), + KEY_BODY_MORPHER: + lambda: load_body_morpher(module_file_names[KEY_BODY_MORPHER]), + } + + return GeneralPoser02( + image_size=512, + module_loaders=loaders, + pose_parameters=get_pose_parameters().get_pose_parameter_groups(), + output_list_func=TwoStepPoserComputationProtocol().compute_func(), + subrect=None, + device=device, + output_length=5 + 1, + default_output_index=default_output_index) diff --git a/tha4/src/tha4/poser/modes/pose_parameters.py b/tha4/src/tha4/poser/modes/pose_parameters.py new file mode 100644 index 0000000..8f38aa2 --- /dev/null +++ b/tha4/src/tha4/poser/modes/pose_parameters.py @@ -0,0 +1,36 @@ +from tha4.poser.poser import PoseParameters, PoseParameterCategory + + +def get_pose_parameters(): + return PoseParameters.Builder() \ + .add_parameter_group("eyebrow_troubled", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_angry", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_lowered", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_raised", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_happy", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eyebrow_serious", PoseParameterCategory.EYEBROW, arity=2) \ + .add_parameter_group("eye_wink", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_happy_wink", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_surprised", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_relaxed", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_unimpressed", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("eye_raised_lower_eyelid", PoseParameterCategory.EYE, arity=2) \ + .add_parameter_group("iris_small", PoseParameterCategory.IRIS_MORPH, arity=2) \ + .add_parameter_group("mouth_aaa", PoseParameterCategory.MOUTH, arity=1, default_value=1.0) \ + .add_parameter_group("mouth_iii", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_uuu", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_eee", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_ooo", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_delta", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("mouth_lowered_corner", PoseParameterCategory.MOUTH, arity=2) \ + .add_parameter_group("mouth_raised_corner", PoseParameterCategory.MOUTH, arity=2) \ + .add_parameter_group("mouth_smirk", PoseParameterCategory.MOUTH, arity=1) \ + .add_parameter_group("iris_rotation_x", PoseParameterCategory.IRIS_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("iris_rotation_y", PoseParameterCategory.IRIS_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("head_x", PoseParameterCategory.FACE_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("head_y", PoseParameterCategory.FACE_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("neck_z", PoseParameterCategory.FACE_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("body_y", PoseParameterCategory.BODY_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("body_z", PoseParameterCategory.BODY_ROTATION, arity=1, range=(-1.0, 1.0)) \ + .add_parameter_group("breathing", PoseParameterCategory.BREATHING, arity=1, range=(0.0, 1.0)) \ + .build() \ No newline at end of file diff --git a/tha4/src/tha4/poser/poser.py b/tha4/src/tha4/poser/poser.py new file mode 100644 index 0000000..eb712ba --- /dev/null +++ b/tha4/src/tha4/poser/poser.py @@ -0,0 +1,162 @@ +from abc import ABC, abstractmethod +from enum import Enum +from typing import Tuple, List, Optional + +import torch +from torch import Tensor + + +class PoseParameterCategory(Enum): + EYEBROW = 1 + EYE = 2 + IRIS_MORPH = 3 + IRIS_ROTATION = 4 + MOUTH = 5 + FACE_ROTATION = 6 + BODY_ROTATION = 7 + BREATHING = 8 + + +class PoseParameterGroup: + def __init__(self, + group_name: str, + parameter_index: int, + category: PoseParameterCategory, + arity: int = 1, + discrete: bool = False, + default_value: float = 0.0, + range: Optional[Tuple[float, float]] = None): + assert arity == 1 or arity == 2 + if range is None: + range = (0.0, 1.0) + if arity == 1: + parameter_names = [group_name] + else: + parameter_names = [group_name + "_left", group_name + "_right"] + assert len(parameter_names) == arity + + self.parameter_names = parameter_names + self.range = range + self.default_value = default_value + self.discrete = discrete + self.arity = arity + self.category = category + self.parameter_index = parameter_index + self.group_name = group_name + + def get_arity(self) -> int: + return self.arity + + def get_group_name(self) -> str: + return self.group_name + + def get_parameter_names(self) -> List[str]: + return self.parameter_names + + def is_discrete(self) -> bool: + return self.discrete + + def get_range(self) -> Tuple[float, float]: + return self.range + + def get_default_value(self): + return self.default_value + + def get_parameter_index(self): + return self.parameter_index + + def get_category(self) -> PoseParameterCategory: + return self.category + + +class PoseParameters: + def __init__(self, pose_parameter_groups: List[PoseParameterGroup]): + self.pose_parameter_groups = pose_parameter_groups + + def get_parameter_index(self, name: str) -> int: + index = 0 + for parameter_group in self.pose_parameter_groups: + for param_name in parameter_group.parameter_names: + if name == param_name: + return index + index += 1 + raise RuntimeError("Cannot find parameter with name %s" % name) + + def get_parameter_name(self, index: int) -> str: + assert index >= 0 and index < self.get_parameter_count() + + for group in self.pose_parameter_groups: + if index < group.get_arity(): + return group.get_parameter_names()[index] + index -= group.arity + + raise RuntimeError("Something is wrong here!!!") + + def get_pose_parameter_groups(self): + return self.pose_parameter_groups + + def get_parameter_count(self): + count = 0 + for group in self.pose_parameter_groups: + count += group.arity + return count + + class Builder: + def __init__(self): + self.index = 0 + self.pose_parameter_groups = [] + + def add_parameter_group(self, + group_name: str, + category: PoseParameterCategory, + arity: int = 1, + discrete: bool = False, + default_value: float = 0.0, + range: Optional[Tuple[float, float]] = None): + self.pose_parameter_groups.append( + PoseParameterGroup( + group_name, + self.index, + category, + arity, + discrete, + default_value, + range)) + self.index += arity + return self + + def build(self) -> 'PoseParameters': + return PoseParameters(self.pose_parameter_groups) + + +class Poser(ABC): + @abstractmethod + def get_image_size(self) -> int: + pass + + @abstractmethod + def get_output_length(self) -> int: + pass + + @abstractmethod + def get_pose_parameter_groups(self) -> List[PoseParameterGroup]: + pass + + @abstractmethod + def get_num_parameters(self) -> int: + pass + + @abstractmethod + def pose(self, image: Tensor, pose: Tensor, output_index: int = 0) -> Tensor: + pass + + @abstractmethod + def get_posing_outputs(self, image: Tensor, pose: Tensor) -> List[Tensor]: + pass + + def get_dtype(self) -> torch.dtype: + return torch.float + + @abstractmethod + def to(self, device: torch.device): + pass \ No newline at end of file diff --git a/tha4/src/tha4/pytasuku/__init__.py b/tha4/src/tha4/pytasuku/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/pytasuku/indexed/__init__.py b/tha4/src/tha4/pytasuku/indexed/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/pytasuku/indexed/all_tasks.py b/tha4/src/tha4/pytasuku/indexed/all_tasks.py new file mode 100644 index 0000000..11ea938 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/all_tasks.py @@ -0,0 +1,26 @@ +from typing import Iterable + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_tasks import IndexedTasks +from tha4.pytasuku.indexed.no_index_command_tasks import NoIndexCommandTasks + + +class AllTasks(NoIndexCommandTasks): + def __init__( + self, + workspace: Workspace, prefix: str, + tasks: Iterable[IndexedTasks], + command_name: str = "all", + define_tasks_immediately: bool = True): + super().__init__(workspace, prefix, command_name, define_tasks_immediately) + self.tasks = [t for t in tasks] + if define_tasks_immediately: + self.define_tasks() + + def execute_run_command(self): + for task in self.tasks: + self.workspace.run(task.run_command) + + def execute_clean_command(self): + for task in self.tasks: + self.workspace.run(task.clean_command) \ No newline at end of file diff --git a/tha4/src/tha4/pytasuku/indexed/bundled_indexed_file_tasks.py b/tha4/src/tha4/pytasuku/indexed/bundled_indexed_file_tasks.py new file mode 100644 index 0000000..2ac9ab8 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/bundled_indexed_file_tasks.py @@ -0,0 +1,31 @@ +import abc +from typing import Iterable, List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_tasks import IndexedTasks +from tha4.pytasuku.workspace import do_nothing + + +class BundledIndexedTasks: + __metaclass__ = abc.ABCMeta + + @property + @abc.abstractmethod + def indexed_tasks_command_names(self) -> Iterable[str]: + pass + + @abc.abstractmethod + def get_indexed_tasks(self, command_name) -> IndexedTasks: + pass + + +def define_all_tasks_from_list(workspace: Workspace, prefix: str, tasks: List[BundledIndexedTasks]): + for command_name in tasks[0].indexed_tasks_command_names: + workspace.create_command_task( + prefix + "/" + command_name, + [x.get_indexed_tasks(command_name).run_command for x in tasks], + do_nothing) + workspace.create_command_task( + prefix + "/" + command_name + "_clean", + [x.get_indexed_tasks(command_name).clean_command for x in tasks], + do_nothing) diff --git a/tha4/src/tha4/pytasuku/indexed/indexed_file_tasks.py b/tha4/src/tha4/pytasuku/indexed/indexed_file_tasks.py new file mode 100644 index 0000000..3b53cfd --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/indexed_file_tasks.py @@ -0,0 +1,19 @@ +import abc +from typing import List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_tasks import IndexedTasks + + +class IndexedFileTasks(IndexedTasks, abc.ABC): + def __init__(self, workspace: Workspace, prefix: str): + super().__init__(workspace, prefix) + + @property + @abc.abstractmethod + def file_list(self) -> List[str]: + pass + + @abc.abstractmethod + def get_file_name(self, *indices: int) -> str: + pass diff --git a/tha4/src/tha4/pytasuku/indexed/indexed_tasks.py b/tha4/src/tha4/pytasuku/indexed/indexed_tasks.py new file mode 100644 index 0000000..75e6676 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/indexed_tasks.py @@ -0,0 +1,34 @@ +import abc +from typing import List + +from tha4.pytasuku.workspace import Workspace + + +class IndexedTasks(abc.ABC): + def __init__(self, workspace: Workspace, prefix: str): + self.prefix = prefix + self.workspace = workspace + + @property + @abc.abstractmethod + def run_command(self) -> str: + pass + + @property + @abc.abstractmethod + def clean_command(self) -> str: + pass + + @property + @abc.abstractmethod + def shape(self) -> List[int]: + pass + + @property + @abc.abstractmethod + def arity(self) -> int: + pass + + @abc.abstractmethod + def define_tasks(self): + pass diff --git a/tha4/src/tha4/pytasuku/indexed/no_index_command_tasks.py b/tha4/src/tha4/pytasuku/indexed/no_index_command_tasks.py new file mode 100644 index 0000000..6f8c400 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/no_index_command_tasks.py @@ -0,0 +1,41 @@ +import abc +from typing import List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_tasks import IndexedTasks + + +class NoIndexCommandTasks(IndexedTasks, abc.ABC): + def __init__(self, workspace: Workspace, prefix: str, command_name: str, define_tasks_immediately: bool = True): + super().__init__(workspace, prefix) + self.command_name = command_name + if define_tasks_immediately: + self.define_tasks() + + @property + def run_command(self): + return self.prefix + "/" + self.command_name + + @property + def clean_command(self): + return self.prefix + "/" + self.command_name + "_clean" + + @property + def arity(self) -> int: + return 0 + + @property + def shape(self) -> List[int]: + return [] + + @abc.abstractmethod + def execute_run_command(self): + pass + + @abc.abstractmethod + def execute_clean_command(self): + pass + + def define_tasks(self): + self.workspace.create_command_task(self.run_command, [], self.execute_run_command) + self.workspace.create_command_task(self.clean_command, [], self.execute_clean_command) diff --git a/tha4/src/tha4/pytasuku/indexed/no_index_file_tasks.py b/tha4/src/tha4/pytasuku/indexed/no_index_file_tasks.py new file mode 100644 index 0000000..70be842 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/no_index_file_tasks.py @@ -0,0 +1,56 @@ +import abc +from typing import List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_file_tasks import IndexedFileTasks +from tha4.pytasuku.indexed.util import delete_file + + +class NoIndexFileTasks(IndexedFileTasks, abc.ABC): + def __init__(self, workspace: Workspace, prefix: str, command_name: str, define_tasks_immediately: bool = True): + super().__init__(workspace, prefix) + self.command_name = command_name + if define_tasks_immediately: + self.define_tasks() + + @property + @abc.abstractmethod + def file_name(self): + pass + + @abc.abstractmethod + def create_file_task(self): + pass + + def get_file_name(self, *indices: int) -> str: + if len(indices) > 0: + raise IndexError("NoIndexFileTasks has arity 0, but get_file_name is called with an index.") + return self.file_name + + @property + def run_command(self): + return self.prefix + "/" + self.command_name + + @property + def clean_command(self): + return self.prefix + "/" + self.command_name + "_clean" + + @property + def arity(self) -> int: + return 0 + + @property + def shape(self) -> List[int]: + return [] + + @property + def file_list(self) -> List[str]: + return [self.file_name] + + def clean(self): + delete_file(self.file_name) + + def define_tasks(self): + self.create_file_task() + self.workspace.create_command_task(self.run_command, [self.file_name]) + self.workspace.create_command_task(self.clean_command, [], lambda: self.clean()) \ No newline at end of file diff --git a/tha4/src/tha4/pytasuku/indexed/one_index_file_tasks.py b/tha4/src/tha4/pytasuku/indexed/one_index_file_tasks.py new file mode 100644 index 0000000..ca5305c --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/one_index_file_tasks.py @@ -0,0 +1,66 @@ +import abc + +from typing import List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_file_tasks import IndexedFileTasks +from tha4.pytasuku.indexed.util import delete_file + + +class OneIndexFileTasks(IndexedFileTasks, abc.ABC): + def __init__(self, workspace: Workspace, prefix: str, command_name: str, count: int, + define_tasks_immediately: bool = True): + super().__init__(workspace, prefix) + self.command_name = command_name + self.count = count + self.file_list_ = [] + if define_tasks_immediately: + self.define_tasks() + + @property + def run_command(self) -> str: + return self.prefix + "/" + self.command_name + + @property + def clean_command(self) -> str: + return self.prefix + "/" + self.command_name + "_clean" + + @property + def shape(self) -> List[int]: + return [self.count] + + @property + def arity(self) -> int: + return 1 + + @abc.abstractmethod + def file_name(self, index): + pass + + @abc.abstractmethod + def create_file_tasks(self, index): + pass + + def get_file_name(self, *indices: int) -> str: + if len(indices) != 1: + raise IndexError("OneIndexFileTasks has arity 1, but " + "get_file_name does not get the appropriate number of arguments.") + return self.file_name(indices[0]) + + @property + def file_list(self): + if len(self.file_list_) == 0: + for i in range(self.count): + self.file_list_.append(self.file_name(i)) + return self.file_list_ + + def clean(self): + for file in self.file_list: + delete_file(file) + + def define_tasks(self): + for index in range(self.count): + self.create_file_tasks(index) + dependencies = self.file_list + self.workspace.create_command_task(self.run_command, dependencies) + self.workspace.create_command_task(self.clean_command, [], lambda: self.clean()) diff --git a/tha4/src/tha4/pytasuku/indexed/simple_no_index_file_tasks.py b/tha4/src/tha4/pytasuku/indexed/simple_no_index_file_tasks.py new file mode 100644 index 0000000..d4d1ed5 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/simple_no_index_file_tasks.py @@ -0,0 +1,28 @@ +from typing import Callable, List, Optional + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.no_index_file_tasks import NoIndexFileTasks + + +class SimpleNoIndexFileTasks(NoIndexFileTasks): + def __init__(self, + workspace: Workspace, + prefix: str, + command_name: str, + file_name: str, + run_func: Callable[[], None], + dependencies: Optional[List[str]] = None): + super().__init__(workspace, prefix, command_name, define_tasks_immediately=False) + if dependencies is None: + dependencies = [] + self.run_func = run_func + self._file_name = file_name + self.dependencies = dependencies + self.define_tasks() + + @property + def file_name(self): + return self._file_name + + def create_file_task(self): + self.workspace.create_file_task(self.file_name, self.dependencies, self.run_func) diff --git a/tha4/src/tha4/pytasuku/indexed/two_indices_file_tasks.py b/tha4/src/tha4/pytasuku/indexed/two_indices_file_tasks.py new file mode 100644 index 0000000..d3d36e0 --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/two_indices_file_tasks.py @@ -0,0 +1,67 @@ +import abc +from typing import List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.indexed_file_tasks import IndexedFileTasks +from tha4.pytasuku.indexed.util import delete_file + + +class TwoIndicesFileTasks(IndexedFileTasks, abc.ABC): + def __init__(self, workspace: Workspace, prefix: str, command_name: str, + count0: int, count1: int, define_tasks_immediately: bool = True): + super().__init__(workspace, prefix) + self.count1 = count1 + self.count0 = count0 + self.command_name = command_name + self.file_list_ = [] + if define_tasks_immediately: + self.define_tasks() + + @property + def run_command(self) -> str: + return self.prefix + "/" + self.command_name + + @property + def clean_command(self) -> str: + return self.prefix + "/" + self.command_name + "_clean" + + @property + def shape(self) -> List[int]: + return [self.count0, self.count1] + + @property + def arity(self) -> int: + return 2 + + @abc.abstractmethod + def file_name(self, index0: int, index1: int) -> str: + pass + + @property + def file_list(self) -> List[str]: + if len(self.file_list_) == 0: + for i in range(self.count0): + for j in range(self.count1): + self.file_list_.append(self.file_name(i, j)) + return self.file_list_ + + @abc.abstractmethod + def create_file_tasks(self, index0: int, index1: int): + pass + + def get_file_name(self, *indices: int) -> str: + if len(indices) != 2: + raise IndexError("TwoIndicesFileTasks.get_file_name require two indices, " + + "but not exactly 2 indices were provide") + return self.file_name(indices[0], indices[1]) + + def clean(self): + for file in self.file_list: + delete_file(file) + + def define_tasks(self): + for index0 in range(self.count0): + for index1 in range(self.count1): + self.create_file_tasks(index0, index1) + self.workspace.create_command_task(self.run_command, self.file_list) + self.workspace.create_command_task(self.clean_command, [], lambda: self.clean()) diff --git a/tha4/src/tha4/pytasuku/indexed/util.py b/tha4/src/tha4/pytasuku/indexed/util.py new file mode 100644 index 0000000..c1d9b4b --- /dev/null +++ b/tha4/src/tha4/pytasuku/indexed/util.py @@ -0,0 +1,71 @@ +import os +from typing import Iterable, Dict, Callable, List + +from tha4.pytasuku.workspace import Workspace +from tha4.pytasuku.indexed.all_tasks import AllTasks +from tha4.pytasuku.indexed.indexed_tasks import IndexedTasks + + +def delete_file(file_name): + if os.path.exists(file_name): + os.remove(file_name) + print("[delete] " + file_name) + else: + print("[not exist] " + file_name) + + +def all_tasks_from_named_tasks_map( + workspace: Workspace, + prefix: str, + tasks: Iterable[Dict[str, IndexedTasks]], + define_all_tasks: bool = True) \ + -> Dict[str, IndexedTasks]: + subtasks = [x for x in tasks] + name_to_subtask_list = {} + for a_subtasks in subtasks: + for name in a_subtasks: + if not define_all_tasks and name == "all": + continue + if name not in name_to_subtask_list: + name_to_subtask_list[name] = [] + name_to_subtask_list[name].append(a_subtasks[name]) + output = {} + for name in name_to_subtask_list: + output[name] = AllTasks(workspace, prefix, name_to_subtask_list[name], name) + return output + + +def create_tasks_hierarchy_helper( + workspace: Workspace, + prefix: str, + tasks_func: Callable[[Workspace, str, List[str]], Dict[str, IndexedTasks]], + branches: List[List[str]], + path: List[str]): + if len(branches) == 0: + return tasks_func(workspace, prefix, path) + else: + tasks = {} + for branch in branches[0]: + output_tasks = create_tasks_hierarchy_helper( + workspace, + f"{prefix}/{branch}", + tasks_func, + branches[1:], + path + [branch]) + if output_tasks is not None: + tasks[branch] = output_tasks + return all_tasks_from_named_tasks_map(workspace, prefix, tasks.values()) + + +def create_task_hierarchy( + workspace: Workspace, + prefix: str, + tasks_func: Callable[[Workspace, str, List[str]], Dict[str, IndexedTasks]], + branches: List[List[str]]) -> Dict[str, IndexedTasks]: + return create_tasks_hierarchy_helper(workspace, prefix, tasks_func, branches, []) + + +def write_done_file(file_name: str): + os.makedirs(os.path.dirname(file_name), exist_ok=True) + with open(file_name, "wt") as fout: + fout.write("DONE!!!") \ No newline at end of file diff --git a/tha4/src/tha4/pytasuku/task.py b/tha4/src/tha4/pytasuku/task.py new file mode 100644 index 0000000..f6909f9 --- /dev/null +++ b/tha4/src/tha4/pytasuku/task.py @@ -0,0 +1,100 @@ +import os +import logging +from typing import List + + +class Task: + def __init__(self, workspace: 'Workspace', name: str, dependencies: List[str]): + self._workspace = workspace + self._name = name + self._dependencies = dependencies + self._workspace.add_task(self) + + def run(self): + pass + + @property + def can_run(self) -> bool: + return True + + @property + def needs_to_be_run(self) -> bool: + return False + + @property + def name(self) -> str: + return self._name + + @property + def dependencies(self) -> List[str]: + return self._dependencies + + @property + def workspace(self) -> 'Workspace': + return self._workspace + + @property + def timestamp(self) -> float: + return float("inf") + + +class CommandTask(Task): + def __init__(self, workspace, name, dependencies): + super().__init__(workspace, name, dependencies) + + @property + def needs_to_be_run(self): + return True + + +class PlaceholderTask(Task): + def __init__(self, workspace, name): + super().__init__(workspace, name, []) + + @property + def can_run(self): + return False + + def run(self): + raise Exception("A placeholder task cannot be run! (%s)" % self.name) + + @property + def needs_to_be_run(self): + return not os.path.isfile(self.name) + + @property + def timestamp(self) -> float: + if not os.path.isfile(self.name): + return float("inf") + else: + return os.path.getmtime(self.name) + + +class FileTask(Task): + def __init__(self, workspace, name, dependencies): + super().__init__(workspace, name, dependencies) + + @property + def timestamp(self): + return os.path.getmtime(self.name) + + @property + def needs_to_be_run(self): + if not os.path.isfile(self.name): + logging.info("Task %s will be run because the corresponding file does not exist." % self.name) + return True + for dep in self.dependencies: + if self.workspace.needs_to_run(dep): + logging.info("Task %s will be run because dependency %s also needs to be run." % (self.name, dep)) + return True + else: + self_timestamp = self.timestamp + dep_task = self.workspace.get_task(dep) + if dep_task.timestamp > self_timestamp: + if isinstance(dep_task, FileTask) or isinstance(dep_task, PlaceholderTask): + logging.info("Task %s needs to be run because task %s has later timestamp." % + (self.name, dep)) + elif isinstance(dep_task, CommandTask): + logging.info("Task %s needs to be run because task %s is a command." % (self.name, dep)) + return True + return False diff --git a/tha4/src/tha4/pytasuku/task_selector_ui.py b/tha4/src/tha4/pytasuku/task_selector_ui.py new file mode 100644 index 0000000..94ba8b9 --- /dev/null +++ b/tha4/src/tha4/pytasuku/task_selector_ui.py @@ -0,0 +1,113 @@ +from tkinter import Tk, BOTH, Button, RIGHT, Scrollbar +from tkinter.ttk import Frame, Treeview + +from tha4.pytasuku.workspace import Workspace, PlaceholderTask + + +class TaskSelectorUi(Frame): + def __init__(self, root, workspace: Workspace): + super().__init__() + self.root = root + self.workspace = workspace + self.master.title("Tasks") + self.master.geometry("256x512") + + treeview_frame = Frame(self) + treeview_frame.pack(fill=BOTH, expand=True) + + self.treeview = Treeview(treeview_frame) + self.treeview["columns"] = ("task_name") + self.treeview.column("#0", width=256, minwidth=256) + self.treeview.heading("#0", text="Tree") + self.treeview.heading("task_name", text="Task Name") + + treeview_vertical_scroll = Scrollbar(treeview_frame, + orient='vertical', + command=self.treeview.yview) + self.treeview.configure(yscrollcommand=treeview_vertical_scroll.set) + treeview_vertical_scroll.pack(side=RIGHT, fill='y') + self.treeview.pack(fill=BOTH, expand=True) + + treeview_horizontal_scroll = Scrollbar(treeview_frame, + orient='horizontal', + command=self.treeview.xview) + self.treeview.configure(xscrollcommand=treeview_horizontal_scroll.set) + treeview_horizontal_scroll.pack(fill='x') + + self.add_tree_nodes() + + self.execute_button = Button(self, text="Execute!", command=self.run_selected_task) + self.execute_button.pack(side=RIGHT, padx=5, pady=5) + + self.pack(fill=BOTH, expand=True) + + self.selected_task_name = None + + def add_tree_nodes(self): + nodes = {} + + for task in self.workspace._tasks.values(): + if isinstance(task, PlaceholderTask): + continue + comps = task.name.split('/') + for i in range(1, len(comps)): + assert len(comps) > 0 + prefix = "" + index = 0 + for comp in comps: + index = index + 1 + parent = prefix + if prefix == "" and comp == "": + prefix = "/" + elif prefix == "": + prefix = prefix + comp + elif prefix == "/": + prefix = prefix + comp + else: + prefix = prefix + "/" + comp + if prefix in nodes: + continue + if index == len(comps): + data = prefix + else: + data = "" + if prefix == "/": + comp = "/" + nodes[prefix] = { + "name": str(prefix), + "display_name": comp, + "parent": parent, + "data": data + } + + sorted_node_names = sorted(nodes.keys()) + node_index = {} + for name in sorted_node_names: + node = nodes[name] + if node["parent"] == "": + id = self.treeview.insert("", "end", text=node["display_name"], values=node["data"], ) + else: + parent = node_index[node["parent"]] + id = self.treeview.insert(parent, "end", text=node["display_name"], values=node["data"], ) + node_index[node["name"]] = id + + def run_selected_task(self): + selection = self.treeview.selection() + item = self.treeview.item(selection) + if item['values'] == "": + return + task_name = item["values"][0] + self.selected_task_name = task_name + self.root.destroy() + + +def run_task_selector_ui(workspace: Workspace): + root = Tk() + task_selector_ui = TaskSelectorUi(root, workspace=workspace) + root.mainloop() + + task_name = task_selector_ui.selected_task_name + if task_name is not None: + print("Running", task_name, "...") + with workspace.session(): + workspace.run(task_name) diff --git a/tha4/src/tha4/pytasuku/util.py b/tha4/src/tha4/pytasuku/util.py new file mode 100644 index 0000000..24da79b --- /dev/null +++ b/tha4/src/tha4/pytasuku/util.py @@ -0,0 +1,15 @@ +import os.path +from typing import List +import logging + +from tha4.pytasuku.workspace import Workspace + + +def create_delete_all_task(workspace: Workspace, name: str, files: List[str]): + def delete_all(): + for file in files: + if os.path.exists(file): + logging.info("Removing %s ..." % file) + os.remove(file) + + workspace.create_command_task(name, [], delete_all) diff --git a/tha4/src/tha4/pytasuku/workspace.py b/tha4/src/tha4/pytasuku/workspace.py new file mode 100644 index 0000000..856ac0e --- /dev/null +++ b/tha4/src/tha4/pytasuku/workspace.py @@ -0,0 +1,168 @@ +from contextlib import contextmanager +from enum import Enum +from typing import List + +from tha4.pytasuku.task import Task, CommandTask, FileTask, PlaceholderTask + + +class WorkspaceState(Enum): + OUT_OF_SESSION = 1 + IN_SESSION = 2 + + +class NodeState(Enum): + IN_STACK = 1 + VISITED = 2 + + +class FuncCommandTask(CommandTask): + def __init__(self, workspace, name, dependencies, func): + super().__init__(workspace, name, dependencies) + self._func = func + + def run(self): + self._func() + + +class FuncFileTask(FileTask): + def __init__(self, workspace, name, dependencies, func): + super().__init__(workspace, name, dependencies) + self._func = func + + def run(self): + self._func() + + +def do_nothing(): + pass + + +class Workspace: + def __init__(self): + self._tasks = dict() + self._name_to_done = None + self._state = WorkspaceState.OUT_OF_SESSION + self._modified = False + + @property + def modified(self) -> bool: + return self._modified + + @property + def state(self) -> WorkspaceState: + return self._state + + @property + def in_session(self) -> bool: + return self._state == WorkspaceState.IN_SESSION + + def task_exists(self, name: str) -> bool: + return name in self._tasks + + def task_exists_and_not_placeholder(self, name: str) -> bool: + return self.task_exists(name) and not isinstance(self.get_task(name), PlaceholderTask) + + def get_task(self, name: str) -> Task: + return self._tasks[name] + + def add_task(self, task): + if self.in_session: + raise RuntimeError("New tasks can only be created when the workspace is out of session.") + if isinstance(task, PlaceholderTask): + if not self.task_exists(task.name): + self._tasks[task.name] = task + self._modified = True + else: + self._tasks[task.name] = task + for dep in task.dependencies: + PlaceholderTask(self, dep) + self._modified = True + + def start_session(self): + if self.in_session: + raise RuntimeError("A session can only be started when the workspace is out of session.") + if self.modified: + self.check_cycle() + self._state = WorkspaceState.IN_SESSION + self._name_to_done = dict() + self._modified = False + + def end_session(self): + if not self.in_session: + raise RuntimeError("A session can only be ended when the workspace is in session.") + self._state = WorkspaceState.OUT_OF_SESSION + self._name_to_done = None + + @contextmanager + def session(self): + try: + self.start_session() + yield + finally: + self.end_session() + + def check_cycle(self): + node_states = dict() + for name in self._tasks: + if name not in node_states: + self.dfs(name, node_states) + + def dfs(self, name, node_states): + node_states[name] = NodeState.IN_STACK + task = self.get_task(name) + for dep in task.dependencies: + if dep not in node_states: + self.dfs(dep, node_states) + else: + state = node_states[dep] + if state == NodeState.IN_STACK: + raise RuntimeError("Dicovered cyclic dependency!") + node_states[name] = NodeState.VISITED + + def run(self, name): + if not self.in_session: + raise RuntimeError("A task can only be run when the workspace is in session.") + if not self.task_exists(name): + raise RuntimeError("Task %s does not exists" % name) + self.run_helper(name) + + def run_helper(self, name): + task = self.get_task(name) + for dep in task.dependencies: + if self.needs_to_run(dep): + self.run_helper(dep) + if self.needs_to_run(name): + task.run() + self._name_to_done[name] = True + + def needs_to_run(self, name): + if not self.in_session: + raise RuntimeError("You can only check whether a task needs to run when the workspace is in session.") + if name in self._name_to_done: + return not self._name_to_done[name] + task = self.get_task(name) + need_to_run_value = task.needs_to_be_run + self._name_to_done[name] = not need_to_run_value + return need_to_run_value + + def create_command_task(self, name, dependencies, func=do_nothing): + return FuncCommandTask(self, name, dependencies, func) + + def create_file_task(self, name, dependencies, func): + return FuncFileTask(self, name, dependencies, func) + + +def command_task(workspace: Workspace, name: str, dependencies: List[str]): + def func(f): + workspace.create_command_task(name, dependencies, f) + return f + + return func + + +def file_task(workspace: Workspace, name: str, dependencies: List[str]): + def func(f): + workspace.create_file_task(name, dependencies, f) + return f + + return func diff --git a/tha4/src/tha4/sampleoutput/__init__.py b/tha4/src/tha4/sampleoutput/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/sampleoutput/general_sample_output_protocol.py b/tha4/src/tha4/sampleoutput/general_sample_output_protocol.py new file mode 100644 index 0000000..4bf8c2f --- /dev/null +++ b/tha4/src/tha4/sampleoutput/general_sample_output_protocol.py @@ -0,0 +1,142 @@ +import os +from enum import Enum +from typing import List, Dict + +import PIL.Image +import numpy +import torch +from tha4.shion.base.dataset.util import get_indexed_batch +from tha4.shion.base.image_util import pytorch_rgb_to_numpy_image, pytorch_rgba_to_numpy_image +from tha4.shion.core.cached_computation import TensorCachedComputationFunc, ComputationState +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.image_util import grid_change_to_numpy_image +from torch.nn import Module +from torch.nn.functional import interpolate +from torch.utils.data import Dataset + + +class ImageType(Enum): + COLOR = 0 + ALPHA = 1 + GRID_CHANGE = 2 + SIGMOID_LOGIT = 3 + + +class SampleImageSpec: + def __init__(self, value_func: TensorCachedComputationFunc, image_type: ImageType): + self.value_func = value_func + self.image_type = image_type + + +class SampleImageSaver: + def __init__(self, + cell_size: int, + image_channels: int, + sample_image_specs: List[SampleImageSpec]): + super().__init__() + self.sample_image_specs = sample_image_specs + self.cell_size = cell_size + self.image_channels = image_channels + + def save_sample_output_data(self, + state: ComputationState, + prefix: str, + examples_seen_so_far: int): + num_cols = len(self.sample_image_specs) + num_rows = state.batch[0].shape[0] + output_image = numpy.zeros([self.cell_size * num_rows, self.cell_size * num_cols, self.image_channels]) + + for col in range(num_cols): + spec = self.sample_image_specs[col] + images = spec.value_func(state) + start_col = col * self.cell_size + + for image_index in range(num_rows): + image = images[image_index].clone().detach() + row = image_index + start_row = row * self.cell_size + if spec.image_type == ImageType.COLOR: + c, h, w = image.shape + green_screen = torch.ones(3, h, w, device=image.device) * -1.0 + green_screen[1, :, :] = 1.0 + alpha = (image[3:4, :, :] + 1.0) * 0.5 + image[0:3, :, :] = image[0:3, :, :] * alpha + green_screen * (1 - alpha) + image[3:4, :, :] = 1.0 + image = image.cpu() + elif spec.image_type == ImageType.GRID_CHANGE: + image = image.cpu() + elif spec.image_type == ImageType.SIGMOID_LOGIT: + image = torch.sigmoid(image) + image = image.repeat(self.image_channels, 1, 1) + image = image * 2.0 - 1.0 + image = image.cpu() + elif spec.image_type == ImageType.ALPHA: + if image.shape[0] == 1: + image = image.repeat(self.image_channels, 1, 1) + image = image * 2.0 - 1.0 + image = image.cpu() + else: + raise RuntimeError(f"Unsupported image type: {spec.image_type}") + + output_image[start_row:start_row + self.cell_size, start_col:start_col + self.cell_size, :] \ + = self.convert_to_numpy_image(image) + + file_name = "%s/sample_output_%010d.png" % (prefix, examples_seen_so_far) + os.makedirs(os.path.dirname(file_name), exist_ok=True) + if self.image_channels == 3: + mode = 'RGB' + else: + mode = 'RGBA' + pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(output_image * 255.0)), mode=mode) + pil_image.save(file_name) + print("Saved %s" % file_name) + + def convert_to_numpy_image(self, image: torch.Tensor): + image_size = image.shape[-1] + if self.cell_size != image_size: + image = interpolate(image.unsqueeze(0), size=self.cell_size).squeeze(0) + if image.shape[0] == 2: + return grid_change_to_numpy_image(image, num_channels=self.image_channels) + elif self.image_channels == 3: + return pytorch_rgb_to_numpy_image(image) + else: + return pytorch_rgba_to_numpy_image(image) + + +class GeneralSampleOutputProtocol(SampleOutputProtocol): + def __init__(self, + sample_image_specs: List[SampleImageSpec], + num_images: int = 8, + cell_size: int = 256, + image_channels: int = 4, + examples_per_sample_output: int = 5000, + random_seed: int = 1203040687): + super().__init__() + self.num_images = num_images + self.random_seed = random_seed + self.examples_per_sample_output = examples_per_sample_output + self.sample_image_saver = SampleImageSaver(cell_size, image_channels, sample_image_specs) + + def get_examples_per_sample_output(self) -> int: + return self.examples_per_sample_output + + def get_random_seed(self) -> int: + return self.random_seed + + def get_sample_output_data(self, validation_dataset: Dataset, device: torch.device) -> dict: + example_indices = torch.randint(0, len(validation_dataset), (self.num_images,)) + example_indices = [example_indices[i].item() for i in range(self.num_images)] + batch = get_indexed_batch(validation_dataset, example_indices, device) + return {'batch': batch} + + def save_sample_output_data(self, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + sample_output_data: dict, prefix: str, + examples_seen_so_far: int, + device: torch.device): + for key in modules: + modules[key].train(False) + batch = sample_output_data['batch'] + state = ComputationState(modules, accumulated_modules, batch, {}) + self.sample_image_saver.save_sample_output_data(state, prefix, examples_seen_so_far) diff --git a/tha4/src/tha4/sampleoutput/poser_sampler_output_protocol.py b/tha4/src/tha4/sampleoutput/poser_sampler_output_protocol.py new file mode 100644 index 0000000..aba1fe9 --- /dev/null +++ b/tha4/src/tha4/sampleoutput/poser_sampler_output_protocol.py @@ -0,0 +1,73 @@ +from typing import Optional, List, Dict + +import torch +from torch.nn import Module +from torch.utils.data import Dataset + +from tha4.shion.base.dataset.util import get_indexed_batch +from tha4.shion.core.cached_computation import CachedComputationFunc, ComputationState +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.sampleoutput.sample_image_creator import SampleImageSpec, ImageSource, ImageType, SampleImageSaver + + +class PoserSampleOutputProtocol(SampleOutputProtocol): + def __init__(self, + output_list_func: Optional[CachedComputationFunc] = None, + num_images: int = 8, + image_size: int = 256, + cell_size: int = 256, + image_channels: int = 4, + examples_per_sample_output: int = 5000, + sample_image_specs: Optional[List[SampleImageSpec]] = None, + random_seed: int = 1203040687): + super().__init__() + self.num_images = num_images + self.random_seed = random_seed + self.examples_per_sample_output = examples_per_sample_output + self.output_list_func = output_list_func + + if sample_image_specs is None: + sample_image_specs = [ + SampleImageSpec(ImageSource.BATCH, 0, ImageType.COLOR), + SampleImageSpec(ImageSource.BATCH, 2, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 0, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 1, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 2, ImageType.ALPHA), + SampleImageSpec(ImageSource.BATCH, 3, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 3, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 4, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 5, ImageType.ALPHA), + SampleImageSpec(ImageSource.OUTPUT, 6, ImageType.COLOR), + SampleImageSpec(ImageSource.BATCH, 4, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 7, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 8, ImageType.COLOR), + SampleImageSpec(ImageSource.OUTPUT, 9, ImageType.ALPHA), + SampleImageSpec(ImageSource.OUTPUT, 10, ImageType.COLOR), + ] + + self.sample_image_saver = SampleImageSaver(image_size, cell_size, image_channels, sample_image_specs) + + def get_examples_per_sample_output(self) -> int: + return self.examples_per_sample_output + + def get_random_seed(self) -> int: + return self.random_seed + + def get_sample_output_data(self, validation_dataset: Dataset, device: torch.device) -> dict: + example_indices = torch.randint(0, len(validation_dataset), (self.num_images,)) + example_indices = [example_indices[i].item() for i in range(self.num_images)] + batch = get_indexed_batch(validation_dataset, example_indices, device) + return {'batch': batch} + + def save_sample_output_data(self, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + sample_output_data: dict, prefix: str, + examples_seen_so_far: int, + device: torch.device): + for key in modules: + modules[key].train(False) + batch = sample_output_data['batch'] + with torch.no_grad(): + outputs = self.output_list_func(ComputationState(modules, accumulated_modules, batch)) + self.sample_image_saver.save_sample_output_data(batch, outputs, prefix, examples_seen_so_far) diff --git a/tha4/src/tha4/sampleoutput/sample_image_creator.py b/tha4/src/tha4/sampleoutput/sample_image_creator.py new file mode 100644 index 0000000..f9504d8 --- /dev/null +++ b/tha4/src/tha4/sampleoutput/sample_image_creator.py @@ -0,0 +1,148 @@ +import math +import os +from enum import Enum +from typing import List + +import numpy +import torch +from matplotlib import cm +from torch import Tensor +from torch.nn.functional import interpolate + +from tha4.shion.base.image_util import save_numpy_image + + +class ImageSource(Enum): + BATCH = 0 + OUTPUT = 1 + + +class ImageType(Enum): + COLOR = 0 + ALPHA = 1 + GRID_CHANGE = 2 + SIGMOID_LOGIT = 3 + + +class SampleImageSpec: + def __init__(self, image_source: ImageSource, index: int, image_type: ImageType): + self.image_type = image_type + self.index = index + self.image_source = image_source + + +def torch_rgb_to_numpy_image(torch_image: Tensor, min_pixel_value=-1.0, max_pixel_value=1.0): + assert torch_image.dim() == 3 + assert torch_image.shape[0] == 3 + height = torch_image.shape[1] + width = torch_image.shape[2] + + reshaped_image = torch_image.numpy().reshape(3, height * width).transpose().reshape(height, width, 3) + numpy_image = (reshaped_image - min_pixel_value) / (max_pixel_value - min_pixel_value) + return numpy_image + + +def torch_rgba_to_numpy_image(torch_image: Tensor, min_pixel_value=-1.0, max_pixel_value=1.0): + assert torch_image.dim() == 3 + assert torch_image.shape[0] == 4 + height = torch_image.shape[1] + width = torch_image.shape[2] + + reshaped_image = torch_image.numpy().reshape(4, height * width).transpose().reshape(height, width, 4) + numpy_image = (reshaped_image - min_pixel_value) / (max_pixel_value - min_pixel_value) + numpy_image = numpy.clip(numpy_image, 0.0, 1.0) + return numpy_image + + +def torch_grid_change_to_numpy_image(torch_image, num_channels=3): + height = torch_image.shape[1] + width = torch_image.shape[2] + size_image = (torch_image[0, :, :] ** 2 + torch_image[1, :, :] ** 2).sqrt().view(height, width, 1).numpy() + hsv = cm.get_cmap('hsv') + angle_image = hsv(((torch.atan2( + torch_image[0, :, :].view(height * width), + torch_image[1, :, :].view(height * width)).view(height, width) + math.pi) / (2 * math.pi)).numpy()) * 3 + numpy_image = size_image * angle_image[:, :, 0:3] + if num_channels == 3: + return numpy_image + elif num_channels == 4: + return numpy.concatenate([numpy_image, numpy.ones_like(size_image)], axis=2) + else: + raise RuntimeError("Unsupported num_channels: " + str(num_channels)) + + +class SampleImageSaver: + def __init__(self, + image_size: int, + cell_size: int, + image_channels: int, + sample_image_specs: List[SampleImageSpec]): + super().__init__() + self.sample_image_specs = sample_image_specs + self.cell_size = cell_size + self.image_channels = image_channels + self.image_size = image_size + + def save_sample_output_image(self, batch: List[Tensor], outputs: List[Tensor], file_name: str): + num_cols = len(self.sample_image_specs) + + num_rows = batch[0].shape[0] + output_image = numpy.zeros([self.cell_size * num_rows, self.cell_size * num_cols, self.image_channels]) + + for image_index in range(num_rows): + row = image_index + start_row = row * self.cell_size + + for col in range(num_cols): + spec = self.sample_image_specs[col] + start_col = col * self.cell_size + + if spec.image_source == ImageSource.BATCH: + image = batch[spec.index][image_index].clone() + else: + image = outputs[spec.index][image_index].clone() + + if spec.image_type == ImageType.COLOR: + c, h, w = image.shape + green_screen = torch.ones(3, h, w, device=image.device) * -1.0 + green_screen[1, :, :] = 1.0 + alpha = (image[3:4, :, :] + 1.0) * 0.5 + image[0:3, :, :] = image[0:3, :, :] * alpha + green_screen * (1 - alpha) + image[3:4, :, :] = 1.0 + image = image.detach().cpu() + elif spec.image_type == ImageType.GRID_CHANGE: + image = image.detach().cpu() + elif spec.image_type == ImageType.SIGMOID_LOGIT: + image = torch.sigmoid(image) + image = image.repeat(self.image_channels, 1, 1) + image = image * 2.0 - 1.0 + image = image.detach().cpu() + else: + if image.shape[0] == 1: + image = image.repeat(self.image_channels, 1, 1) + image = image * 2.0 - 1.0 + image = image.detach().cpu() + + output_image[start_row:start_row + self.cell_size, start_col:start_col + self.cell_size, :] \ + = self.convert_to_numpy_image(image) + + os.makedirs(os.path.dirname(file_name), exist_ok=True) + save_numpy_image(output_image, file_name, save_straight_alpha=True) + + def save_sample_output_data(self, + batch: List[Tensor], + outputs: List[Tensor], + prefix: str, + examples_seen_so_far: int): + file_name = "%s/sample_output_%010d.png" % (prefix, examples_seen_so_far) + self.save_sample_output_image(batch, outputs, file_name) + + def convert_to_numpy_image(self, image: torch.Tensor): + if self.cell_size != self.image_size: + image = interpolate(image.unsqueeze(0), size=self.cell_size).squeeze(0) + if image.shape[0] == 2: + return torch_grid_change_to_numpy_image(image, num_channels=self.image_channels) + elif self.image_channels == 3: + return torch_rgb_to_numpy_image(image) + else: + return torch_rgba_to_numpy_image(image) diff --git a/tha4/src/tha4/shion/__init__.py b/tha4/src/tha4/shion/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/base/__init__.py b/tha4/src/tha4/shion/base/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/base/dataset/__init__.py b/tha4/src/tha4/shion/base/dataset/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/base/dataset/lazy_dataset.py b/tha4/src/tha4/shion/base/dataset/lazy_dataset.py new file mode 100644 index 0000000..2678878 --- /dev/null +++ b/tha4/src/tha4/shion/base/dataset/lazy_dataset.py @@ -0,0 +1,20 @@ +from typing import Callable + +from torch.utils.data import Dataset + + +class LazyDataset(Dataset): + def __init__(self, source_func: Callable[[], Dataset]): + self.source_func = source_func + self.source = None + + def get_source(self): + if self.source is None: + self.source = self.source_func() + return self.source + + def __len__(self): + return len(self.get_source()) + + def __getitem__(self, item): + return self.get_source()[item] diff --git a/tha4/src/tha4/shion/base/dataset/lazy_tensor_dataset.py b/tha4/src/tha4/shion/base/dataset/lazy_tensor_dataset.py new file mode 100644 index 0000000..f330f69 --- /dev/null +++ b/tha4/src/tha4/shion/base/dataset/lazy_tensor_dataset.py @@ -0,0 +1,33 @@ +import torch +from torch.utils.data import Dataset, TensorDataset + +from tha4.shion.core.load_save import torch_load + + +class LazyTensorDataset(Dataset): + def __init__(self, file_name: str): + self.file_name = file_name + self.dataset = None + + def get_dataset(self): + if self.dataset is None: + data = torch_load(self.file_name) + if isinstance(data, torch.Tensor): + self.dataset = TensorDataset(data) + elif isinstance(data, tuple): + self.dataset = TensorDataset(*data) + elif isinstance(data, list): + self.dataset = TensorDataset(*data) + else: + raise RuntimeError("Unsupported data type: " + type(data)) + return self.dataset + + def __len__(self): + dataset = self.get_dataset() + return len(dataset) + + def __getitem__(self, item): + dataset = self.get_dataset() + return dataset.__getitem__(item) + + diff --git a/tha4/src/tha4/shion/base/dataset/png_in_dir_dataset.py b/tha4/src/tha4/shion/base/dataset/png_in_dir_dataset.py new file mode 100644 index 0000000..9d11ef0 --- /dev/null +++ b/tha4/src/tha4/shion/base/dataset/png_in_dir_dataset.py @@ -0,0 +1,53 @@ +import os + +from torch.nn import functional +from torch.utils.data import Dataset +from os import listdir +from os.path import isfile + +from tha4.shion.base.image_util import extract_pytorch_image_from_filelike + + +class PngInDirDataset(Dataset): + def __init__(self, dir: str, + downscale_kernel_size: int = 1, + has_alpha=False, + scale=2.0, + offset=-1.0, + premultiply_alpha=True, + perfrom_srb_to_linear=True): + super().__init__() + self.perfrom_srb_to_linear = perfrom_srb_to_linear + self.premultiply_alpha = premultiply_alpha + self.offset = offset + self.scale = scale + self.has_alpha = has_alpha + self.downscale_kernel_size = downscale_kernel_size + self.dir = dir + self.file_names = None + + def get_file_names(self): + if self.file_names is None: + self.file_names = [os.path.join(self.dir, x) for x in listdir(self.dir)] + self.file_names = [x for x in self.file_names if isfile(x) and x[-4:].lower() == ".png"] + self.file_names = sorted(self.file_names) + return self.file_names + + def __len__(self): + file_names = self.get_file_names() + return len(file_names) + + def __getitem__(self, item): + file_names = self.get_file_names() + file_name = file_names[item] + image = extract_pytorch_image_from_filelike( + file_name, + scale=self.scale, + offset=self.offset, + premultiply_alpha=self.has_alpha and self.premultiply_alpha, + perform_srgb_to_linear=self.perfrom_srb_to_linear) + if self.downscale_kernel_size == 1: + return [image] + else: + image = functional.avg_pool2d(image.unsqueeze(0), kernel_size=self.downscale_kernel_size).squeeze(0) + return [image] diff --git a/tha4/src/tha4/shion/base/dataset/util.py b/tha4/src/tha4/shion/base/dataset/util.py new file mode 100644 index 0000000..77c6be3 --- /dev/null +++ b/tha4/src/tha4/shion/base/dataset/util.py @@ -0,0 +1,29 @@ +from typing import List + +import torch +from torch.utils.data import Dataset + + +def get_indexed_batch(dataset: Dataset, example_indices: List[int], device: torch.device): + if len(example_indices) == 0: + return [] + examples = [] + for index in range(len(example_indices)): + example_index = example_indices[index] + raw_example = dataset[example_index] + example = [] + for x in raw_example: + if isinstance(x, torch.Tensor): + y = x.to(device).unsqueeze(0) + elif isinstance(x, float) or isinstance(x, int): + y = torch.tensor([[x]], device=device) + else: + raise RuntimeError(f"get_indexed_batch: Data of type {type(x)} is not supported.") + example.append(y) + examples.append(example) + k = len(examples[0]) + transposed = [[] for i in range(k)] + for example in examples: + for i in range(k): + transposed[i].append(example[i]) + return [torch.cat(x, dim=0) for x in transposed] diff --git a/tha4/src/tha4/shion/base/dataset/xformed_dataset.py b/tha4/src/tha4/shion/base/dataset/xformed_dataset.py new file mode 100644 index 0000000..169b134 --- /dev/null +++ b/tha4/src/tha4/shion/base/dataset/xformed_dataset.py @@ -0,0 +1,15 @@ +from typing import Any, Callable + +from torch.utils.data import Dataset + + +class XformedDataset(Dataset): + def __init__(self, source: Dataset, xform_func: Callable[[Any], Any]): + self.xform_func = xform_func + self.source = source + + def __len__(self): + return len(self.source) + + def __getitem__(self, item): + return self.xform_func(self.source[item]) diff --git a/tha4/src/tha4/shion/base/image_util.py b/tha4/src/tha4/shion/base/image_util.py new file mode 100644 index 0000000..c89a962 --- /dev/null +++ b/tha4/src/tha4/shion/base/image_util.py @@ -0,0 +1,243 @@ +import os + +import PIL.Image +import numpy +import torch +from matplotlib import pyplot +from torch import Tensor + + +def numpy_srgb_to_linear(x): + x = numpy.clip(x, 0.0, 1.0) + return numpy.where(x <= 0.04045, x / 12.92, ((x + 0.055) / 1.055) ** 2.4) + + +def numpy_linear_to_srgb(x): + x = numpy.clip(x, 0.0, 1.0) + return numpy.where(x <= 0.003130804953560372, x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) + + +def numpy_alpha_devide(rgb, a, epsilon=1e-5): + aaa = numpy.repeat(a, 3, axis=2) + aaa_prime = aaa + numpy.where(numpy.abs(aaa) < epsilon, epsilon, 0.0) + return numpy.where(numpy.abs(aaa) < epsilon, 0.0, rgb / aaa_prime) + + +def torch_srgb_to_linear(x: torch.Tensor): + x = torch.clip(x, 0.0, 1.0) + return torch.where(torch.le(x, 0.04045), x / 12.92, ((x + 0.055) / 1.055) ** 2.4) + + +def torch_linear_to_srgb(x): + x = torch.clip(x, 0.0, 1.0) + return torch.where(torch.le(x, 0.003130804953560372), x * 12.92, 1.055 * (x ** (1.0 / 2.4)) - 0.055) + + +def numpy_image_linear_to_srgb(image): + assert image.shape[2] == 3 or image.shape[2] == 4 + if image.shape[2] == 3: + return numpy_linear_to_srgb(image) + else: + height, width, _ = image.shape + rgb_image = numpy_linear_to_srgb(image[:, :, 0:3]) + a_image = image[:, :, 3:4] + return numpy.concatenate((rgb_image, a_image), axis=2) + + +def numpy_image_srgb_to_linear(image): + assert image.shape[2] == 3 or image.shape[2] == 4 + if image.shape[2] == 3: + return numpy_srgb_to_linear(image) + else: + height, width, _ = image.shape + rgb_image = numpy_srgb_to_linear(image[:, :, 0:3]) + a_image = image[:, :, 3:4] + return numpy.concatenate((rgb_image, a_image), axis=2) + + +def pytorch_rgb_to_numpy_image(torch_image: Tensor, min_pixel_value=-1.0, max_pixel_value=1.0): + assert torch_image.dim() == 3 + assert torch_image.shape[0] == 3 + height = torch_image.shape[1] + width = torch_image.shape[2] + + reshaped_image = torch_image.numpy().reshape(3, height * width).transpose().reshape(height, width, 3) + numpy_image = (reshaped_image - min_pixel_value) / (max_pixel_value - min_pixel_value) + return numpy_linear_to_srgb(numpy_image) + + +def pytorch_rgba_to_numpy_image_greenscreen(torch_image: Tensor, + min_pixel_value=-1.0, + max_pixel_value=1.0, + include_alpha=False): + height = torch_image.shape[1] + width = torch_image.shape[2] + + numpy_image = (torch_image.numpy().reshape(4, height * width).transpose().reshape(height, width, + 4) - min_pixel_value) \ + / (max_pixel_value - min_pixel_value) + rgb_image = numpy_linear_to_srgb(numpy_image[:, :, 0:3]) + a_image = numpy_image[:, :, 3] + rgb_image[:, :, 0:3] = rgb_image[:, :, 0:3] * a_image.reshape(a_image.shape[0], a_image.shape[1], 1) + rgb_image[:, :, 1] = rgb_image[:, :, 1] + (1 - a_image) + + if not include_alpha: + return rgb_image + else: + return numpy.concatenate((rgb_image, numpy.ones_like(numpy_image[:, :, 3:4])), axis=2) + + +def pytorch_rgba_to_numpy_image( + torch_image: Tensor, + min_pixel_value=-1.0, + max_pixel_value=1.0, + perform_linear_to_srb: bool = True): + assert torch_image.dim() == 3 + assert torch_image.shape[0] == 4 + height = torch_image.shape[1] + width = torch_image.shape[2] + + reshaped_image = torch_image.numpy().reshape(4, height * width).transpose().reshape(height, width, 4) + numpy_image = (reshaped_image - min_pixel_value) / (max_pixel_value - min_pixel_value) + if perform_linear_to_srb: + rgb_image = numpy_linear_to_srgb(numpy_image[:, :, 0:3]) + else: + rgb_image = numpy.clip(numpy_image[:, :, 0:3], 0.0, 1.0) + a_image = numpy.clip(numpy_image[:, :, 3], 0.0, 1.0) + rgba_image = numpy.concatenate((rgb_image, a_image.reshape(height, width, 1)), axis=2) + return rgba_image + + +def pil_image_has_transparency(pil_image): + if pil_image.info.get("transparency", None) is not None: + return True + if pil_image.mode == "P": + transparent = pil_image.info.get("transparency", -1) + for _, index in pil_image.getcolors(): + if index == transparent: + return True + elif pil_image.mode == "RGBA": + extrema = pil_image.getextrema() + if extrema[3][0] < 255: + return True + + return False + + +def extract_numpy_image_from_PIL_image(pil_image, scale=2.0, offset=-1.0, + premultiply_alpha=True, + perform_srgb_to_linear=True): + has_alpha = pil_image_has_transparency(pil_image) + if has_alpha and pil_image.mode != 'RGBA': + pil_image = pil_image.convert("RGBA") + if not has_alpha and pil_image.mode != 'RGB': + pil_image = pil_image.convert("RGB") + if has_alpha: + num_channel = 4 + else: + num_channel = 3 + image_width = pil_image.width + image_height = pil_image.height + + raw_image = numpy.asarray(pil_image, dtype=numpy.float32) + image = (raw_image / 255.0).reshape(image_height, image_width, num_channel) + if perform_srgb_to_linear: + image[:, :, 0:3] = numpy_srgb_to_linear(image[:, :, 0:3]) + # Premultiply alpha + if has_alpha and premultiply_alpha: + image[:, :, 0:3] = image[:, :, 0:3] * image[:, :, 3:4] + return image * scale + offset + + +def extract_numpy_image_from_PIL_image_with_pytorch_layout(pil_image, scale=2.0, offset=-1.0, + premultiply_alpha=True, + perform_srgb_to_linear=True): + numpy_image = extract_numpy_image_from_PIL_image( + pil_image, scale, offset, premultiply_alpha, perform_srgb_to_linear) + image_height, image_width, num_channel = numpy_image.shape + image = numpy_image \ + .reshape(image_height * image_width, num_channel) \ + .transpose() \ + .reshape(num_channel, image_height, image_width) + return image + + +def extract_numpy_image_from_filelike_with_pytorch_layout(file, scale=2.0, offset=-1.0, premultiply_alpha=True): + try: + pil_image = PIL.Image.open(file) + except Exception as e: + raise RuntimeError(file) + return extract_numpy_image_from_PIL_image_with_pytorch_layout(pil_image, scale, offset, premultiply_alpha) + + +def extract_numpy_image_from_filelike(file, scale=1.0, offset=0.0, + premultiply_alpha=True, + perform_srgb_to_linear: bool = True): + try: + pil_image = PIL.Image.open(file) + except Exception as e: + raise RuntimeError(file) + return extract_numpy_image_from_PIL_image(pil_image, scale, offset, premultiply_alpha, perform_srgb_to_linear) + + +def extract_pytorch_image_from_filelike(file, scale=2.0, offset=-1.0, premultiply_alpha=True, + perform_srgb_to_linear=True): + try: + pil_image = PIL.Image.open(file) + except Exception as e: + raise RuntimeError(file) + image = extract_numpy_image_from_PIL_image_with_pytorch_layout(pil_image, scale, offset, premultiply_alpha, + perform_srgb_to_linear) + return torch.from_numpy(image).float() + + +def extract_pytorch_image_from_PIL_image(pil_image, scale=2.0, offset=-1.0, premultiply_alpha=True, + perform_srgb_to_linear=True): + image = extract_numpy_image_from_PIL_image_with_pytorch_layout( + pil_image, scale, offset, premultiply_alpha, perform_srgb_to_linear) + return torch.from_numpy(image).float() + + +def convert_pytorch_image_to_zero_to_one_numpy_image( + torch_image: torch.Tensor, + scale: float = 2.0, + offset: float = -1.0): + torch_image = (torch_image - offset) / scale + torch_image = torch.permute(torch_image, (1, 2, 0)) + numpy_image = torch_image.cpu().numpy() + return numpy_image + + +def convert_zero_to_one_numpy_image_to_PIL_image( + numpy_image, use_straight_alpha=True, perform_linear_to_srgb=True): + if numpy_image.shape[2] == 4: + rgb_image = numpy_image[:, :, 0:3] + a_image = numpy.clip(numpy_image[:, :, 3:4], 0.0, 1.0) + if use_straight_alpha: + rgb_image = numpy_alpha_devide(rgb_image, a_image) + if perform_linear_to_srgb: + rgb_image = numpy_linear_to_srgb(rgb_image) + else: + rgb_image = numpy.clip(rgb_image, 0.0, 1.0) + new_numpy_image = numpy.concatenate((rgb_image, a_image), axis=2) + pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(new_numpy_image * 255.0)), mode='RGBA') + else: + if perform_linear_to_srgb: + numpy_image = numpy_linear_to_srgb(numpy_image) + else: + numpy_image = numpy.clip(numpy_image, 0.0, 1.0) + pil_image = PIL.Image.fromarray(numpy.uint8(numpy.rint(numpy_image * 255.0)), mode='RGB') + return pil_image + + +def save_numpy_image(numpy_image, file_name: str, save_straight_alpha=True, perform_linear_to_srgb=True): + pil_image = convert_zero_to_one_numpy_image_to_PIL_image(numpy_image, save_straight_alpha, perform_linear_to_srgb) + os.makedirs(os.path.dirname(file_name), exist_ok=True) + pil_image.save(file_name) + + +def resize_PIL_image(pil_image, size=(256, 256)): + w, h = pil_image.size + d = min(w, h) + r = ((w - d) // 2, (h - d) // 2, (w + d) // 2, (h + d) // 2) + return pil_image.resize(size, resample=PIL.Image.LANCZOS, box=r) \ No newline at end of file diff --git a/tha4/src/tha4/shion/base/loss/__init__.py b/tha4/src/tha4/shion/base/loss/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/base/loss/computed_scale_loss.py b/tha4/src/tha4/shion/base/loss/computed_scale_loss.py new file mode 100644 index 0000000..18cb115 --- /dev/null +++ b/tha4/src/tha4/shion/base/loss/computed_scale_loss.py @@ -0,0 +1,22 @@ +from typing import Optional, Callable + +from tha4.shion.core.cached_computation import TensorCachedComputationFunc, ComputationState +from tha4.shion.core.loss import Loss + + +class ComputedScaleLoss(Loss): + def __init__(self, + scale_func: TensorCachedComputationFunc, + loss: Loss, + weight: float = 1.0): + self.weight = weight + self.loss = loss + self.scale_func = scale_func + + def compute(self, state: ComputationState, log_func: Optional[Callable[[str, float], None]] = None): + loss = self.loss.compute(state) + scale = self.scale_func(state) + loss = self.weight * scale * loss + if log_func is not None: + log_func("loss", loss.item()) + return loss diff --git a/tha4/src/tha4/shion/base/loss/computed_scaled_l2_loss.py b/tha4/src/tha4/shion/base/loss/computed_scaled_l2_loss.py new file mode 100644 index 0000000..b592b86 --- /dev/null +++ b/tha4/src/tha4/shion/base/loss/computed_scaled_l2_loss.py @@ -0,0 +1,29 @@ +from typing import Callable, Optional + +from tha4.shion.core.cached_computation import TensorCachedComputationFunc, ComputationState +from tha4.shion.core.loss import Loss + + +class ComputedScaledL2Loss(Loss): + def __init__(self, + expected_func: TensorCachedComputationFunc, + actual_func: TensorCachedComputationFunc, + element_scale_func: TensorCachedComputationFunc, + weight: float = 1.0): + self.element_scale_func = element_scale_func + self.actual_func = actual_func + self.expected_func = expected_func + self.weight = weight + + def compute( + self, + state: ComputationState, + log_func: Optional[Callable[[str, float], None]] = None): + element_scale = self.element_scale_func(state) + expected = self.expected_func(state) + actual = self.actual_func(state) + diff = (expected - actual) * element_scale + loss = self.weight * (diff ** 2).mean() + if log_func is not None: + log_func("loss", loss.item()) + return loss diff --git a/tha4/src/tha4/shion/base/loss/l1_loss.py b/tha4/src/tha4/shion/base/loss/l1_loss.py new file mode 100644 index 0000000..1ed7314 --- /dev/null +++ b/tha4/src/tha4/shion/base/loss/l1_loss.py @@ -0,0 +1,67 @@ +from typing import Callable, Optional + +import torch + +from tha4.shion.core.cached_computation import TensorCachedComputationFunc, ComputationState +from tha4.shion.core.loss import Loss + + +class L1Loss(Loss): + def __init__(self, + expected_func: TensorCachedComputationFunc, + actual_func: TensorCachedComputationFunc, + weight: float = 1.0): + self.actual_func = actual_func + self.expected_func = expected_func + self.weight = weight + + def compute(self, state: ComputationState, log_func: Optional[Callable[[str, float], None]] = None): + expected = self.expected_func(state) + actual = self.actual_func(state) + loss = self.weight * (expected - actual).abs().mean() + if log_func is not None: + log_func("loss", loss.item()) + return loss + + +class ListL1Loss(Loss): + def __init__(self, + expected_func: TensorCachedComputationFunc, + actual_func: TensorCachedComputationFunc, + weight: float = 1.0): + self.actual_func = actual_func + self.expected_func = expected_func + self.weight = weight + + def compute(self, state: ComputationState, log_func: Optional[Callable[[str, float], None]] = None): + expected = self.expected_func(state) + actual = self.actual_func(state) + assert len(expected) == len(actual) + loss = torch.zeros(1, device=expected[0].device) + for i in range(len(expected)): + loss += (expected[i] - actual[i]).abs().mean() + loss = self.weight * loss + if log_func is not None: + log_func("loss", loss.item()) + return loss + + +class MaskedL1Loss(Loss): + def __init__(self, + expected_func: TensorCachedComputationFunc, + actual_func: TensorCachedComputationFunc, + mask_func: TensorCachedComputationFunc, + weight: float = 1.0): + self.mask_func = mask_func + self.actual_func = actual_func + self.expected_func = expected_func + self.weight = weight + + def compute(self, state: ComputationState, log_func: Optional[Callable[[str, float], None]] = None): + mask = self.mask_func(state) + expected = self.expected_func(state) + actual = self.actual_func(state) + loss = self.weight * ((expected - actual) * mask).abs().mean() + if log_func is not None: + log_func("loss", loss.item()) + return loss \ No newline at end of file diff --git a/tha4/src/tha4/shion/base/loss/l2_loss.py b/tha4/src/tha4/shion/base/loss/l2_loss.py new file mode 100644 index 0000000..437ac36 --- /dev/null +++ b/tha4/src/tha4/shion/base/loss/l2_loss.py @@ -0,0 +1,25 @@ +from typing import Callable, Optional + +from tha4.shion.core.cached_computation import TensorCachedComputationFunc, ComputationState +from tha4.shion.core.loss import Loss + + +class L2Loss(Loss): + def __init__(self, + expected_func: TensorCachedComputationFunc, + actual_func: TensorCachedComputationFunc, + weight: float = 1.0): + self.actual_func = actual_func + self.expected_func = expected_func + self.weight = weight + + def compute( + self, + state: ComputationState, + log_func: Optional[Callable[[str, float], None]] = None): + expected = self.expected_func(state) + actual = self.actual_func(state) + loss = self.weight * ((expected - actual) ** 2).mean() + if log_func is not None: + log_func("loss", loss.item()) + return loss diff --git a/tha4/src/tha4/shion/base/loss/sum_loss.py b/tha4/src/tha4/shion/base/loss/sum_loss.py new file mode 100644 index 0000000..6dd604a --- /dev/null +++ b/tha4/src/tha4/shion/base/loss/sum_loss.py @@ -0,0 +1,32 @@ +from typing import List, Tuple, Callable, Optional + +import torch +from torch import Tensor + +from tha4.shion.core.cached_computation import ComputationState +from tha4.shion.core.loss import Loss + + +class SumLoss(Loss): + def __init__(self, losses: List[Tuple[str, Loss]]): + self.losses = losses + + def compute(self, + state: ComputationState, + log_func: Optional[Callable[[str, float], None]] = None) -> Tensor: + device = state.batch[0].device + loss_value = torch.zeros(1, device=device) + for loss_spec in self.losses: + loss_name = loss_spec[0] + loss = loss_spec[1] + if log_func is not None: + def loss_log_func(name, value): + log_func(loss_name + "_" + name, value) + else: + loss_log_func = None + loss_value = loss_value + loss.compute(state, loss_log_func) + + if log_func is not None: + log_func("loss", loss_value.item()) + + return loss_value diff --git a/tha4/src/tha4/shion/base/loss/time_dependently_weighted_loss.py b/tha4/src/tha4/shion/base/loss/time_dependently_weighted_loss.py new file mode 100644 index 0000000..b6c6da2 --- /dev/null +++ b/tha4/src/tha4/shion/base/loss/time_dependently_weighted_loss.py @@ -0,0 +1,29 @@ +from typing import Callable, Optional + +from torch import Tensor + +from tha4.shion.core.cached_computation import ComputationState, CachedComputationFunc +from tha4.shion.core.loss import Loss + + +class TimeDependentlyWeightedLoss(Loss): + def __init__(self, + base_loss: Loss, + examples_seen_so_far_func: CachedComputationFunc, + weight_func: Callable[[int], float]): + self.weight_func = weight_func + self.examples_seen_so_far_func = examples_seen_so_far_func + self.base_loss = base_loss + + def compute(self, + state: ComputationState, + log_func: Optional[Callable[[str, float], None]] = None) -> Tensor: + base_value = self.base_loss.compute(state) + examples_seen_so_far = self.examples_seen_so_far_func(state) + weight = self.weight_func(examples_seen_so_far) + loss_value = base_value * weight + + if log_func is not None: + log_func("loss", loss_value.item()) + + return loss_value diff --git a/tha4/src/tha4/shion/base/module_accumulators.py b/tha4/src/tha4/shion/base/module_accumulators.py new file mode 100644 index 0000000..7af98fe --- /dev/null +++ b/tha4/src/tha4/shion/base/module_accumulators.py @@ -0,0 +1,29 @@ +from typing import Optional + +import torch +from torch.nn import Module + +from tha4.shion.core.module_accumulator import ModuleAccumulator + + +# Code from https://github.com/rosinality/style-based-gan-pytorch/blob/8437a8bbd106ad4a4691b798ce35d30b5111990b/train.py +def accumulate_modules(new_module: Module, accumulated_module: Module, beta=0.99): + with torch.no_grad(): + new_module_params = dict(new_module.named_parameters()) + accumulated_module_params = dict(accumulated_module.named_parameters()) + for key in new_module_params.keys(): + accumulated_module_params[key].mul_(beta).add_(new_module_params[key] * (1 - beta)) + + new_module_buffers = dict(new_module.named_buffers()) + accumulated_module_buffers = dict(accumulated_module.named_buffers()) + for key in new_module_buffers.keys(): + accumulated_module_buffers[key].copy_(new_module_buffers[key]) + + +class DecayAccumulator(ModuleAccumulator): + def __init__(self, decay: float = 0.999): + self.decay = decay + + def accumulate(self, module: Module, output: Module, examples_seen_so_far: Optional[int] = None) -> Module: + accumulate_modules(module, output, self.decay) + return output diff --git a/tha4/src/tha4/shion/base/optimizer_factories.py b/tha4/src/tha4/shion/base/optimizer_factories.py new file mode 100644 index 0000000..d2fcbcc --- /dev/null +++ b/tha4/src/tha4/shion/base/optimizer_factories.py @@ -0,0 +1,46 @@ +from typing import Tuple, Iterable + +from torch.nn import Parameter +from torch.optim import Optimizer, Adam, AdamW, SparseAdam, RMSprop + +from tha4.shion.core.optimizer_factory import OptimizerFactory + + +class AdamOptimizerFactory(OptimizerFactory): + def __init__(self, betas: Tuple[float, float] = (0.9, 0.999), epsilon: float = 1e-8, weight_decay: float = 0.0): + super().__init__() + self.weight_decay = weight_decay + self.betas = betas + self.epsilon = epsilon + + def create(self, parameters: Iterable[Parameter]) -> Optimizer: + return Adam(parameters, betas=self.betas, eps=self.epsilon, weight_decay=self.weight_decay) + + +class AdamWOptimizerFactory(OptimizerFactory): + def __init__(self, betas: Tuple[float, float] = (0.9, 0.999), epsilon: float = 1e-8, weight_decay: float = 0.01): + super().__init__() + self.weight_decay = weight_decay + self.betas = betas + self.epsilon = epsilon + + def create(self, parameters: Iterable[Parameter]) -> Optimizer: + return AdamW(parameters, betas=self.betas, eps=self.epsilon, weight_decay=self.weight_decay) + + +class SparseAdamOptimizerFactory(OptimizerFactory): + def __init__(self, betas: Tuple[float, float] = (0.9, 0.999), epsilon: float = 1e-8): + super().__init__() + self.betas = betas + self.epsilon = epsilon + + def create(self, parameters: Iterable[Parameter]) -> Optimizer: + return SparseAdam(list(parameters), betas=self.betas, eps=self.epsilon) + + +class RMSpropOptimizerFactory(OptimizerFactory): + def __init__(self): + super().__init__() + + def create(self, parameters: Iterable[Parameter]) -> Optimizer: + return RMSprop(parameters) diff --git a/tha4/src/tha4/shion/base/protocol/single_network_from_batch_input_computation_protocol.py b/tha4/src/tha4/shion/base/protocol/single_network_from_batch_input_computation_protocol.py new file mode 100644 index 0000000..d52fee1 --- /dev/null +++ b/tha4/src/tha4/shion/base/protocol/single_network_from_batch_input_computation_protocol.py @@ -0,0 +1,29 @@ +from typing import Optional, Any, List + +from tha4.shion.core.cached_computation import CachedComputationProtocol, ComputationState + +KEY_NETWORK = "network" +KEY_NETWORK_OUTPUT = "network_output" + + +class SingleNetworkBatchInputComputationProtocol(CachedComputationProtocol): + def __init__(self, + key_network: str = KEY_NETWORK, + key_network_output: str = KEY_NETWORK_OUTPUT, + input_index_to_batch_index: Optional[List[int]] = None): + if input_index_to_batch_index is None: + input_index_to_batch_index = [0] + + self.input_index_to_batch_index = input_index_to_batch_index + self.key_network_output = key_network_output + self.key_network = key_network + + def compute_output(self, key: str, state: ComputationState) -> Any: + if key == self.key_network_output: + inputs = [] + for batch_index in self.input_index_to_batch_index: + inputs.append(state.batch[batch_index]) + network = state.modules[self.key_network] + return network.forward(*inputs) + else: + raise RuntimeError("Computing output for key " + key + " is not supported!") diff --git a/tha4/src/tha4/shion/base/training/__init__.py b/tha4/src/tha4/shion/base/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/base/training/single_network.py b/tha4/src/tha4/shion/base/training/single_network.py new file mode 100644 index 0000000..c0b334b --- /dev/null +++ b/tha4/src/tha4/shion/base/training/single_network.py @@ -0,0 +1,108 @@ +import time +from typing import List, Dict, Callable, Any, Optional + +import torch +from torch.nn import Module +from torch.nn.utils import clip_grad_norm_ +from torch.optim.optimizer import Optimizer + +from tha4.shion.core.cached_computation import ComputationState +from tha4.shion.core.loss import Loss +from tha4.shion.core.optimizer_factory import OptimizerFactory +from tha4.shion.core.training.training_protocol import TrainingProtocol +from tha4.shion.core.training.validation_protocol import ValidationProtocol + +KEY_NETWORK = "network" + + +class SingleNetworkTrainingProtocol(TrainingProtocol): + def __init__(self, + check_point_examples: List[int], + batch_size: int, + learning_rate: Callable[[int], Dict[str, float]], + optimizer_factories: Dict[str, OptimizerFactory], + module_key: str = KEY_NETWORK, + random_seed: int = 39549059840, + max_grad_norm: Optional[float] = None): + super().__init__() + self.max_grad_norm = max_grad_norm + self.module_key = module_key + self.optimizer_factories = optimizer_factories + self.learning_rate = learning_rate + self.batch_size = batch_size + self.random_seed = random_seed + self.check_point_examples = check_point_examples + + def get_optimizer_factories(self) -> Dict[str, OptimizerFactory]: + return self.optimizer_factories + + def get_checkpoint_examples(self) -> List[int]: + return self.check_point_examples + + def get_random_seed(self) -> int: + return self.random_seed + + def get_batch_size(self) -> int: + return self.batch_size + + def get_learning_rate(self, examples_seen_so_far: int) -> Dict[str, float]: + return self.learning_rate(examples_seen_so_far) + + def run_training_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer], + losses: Dict[str, Loss], + create_log_func: Optional[Callable[[str, int], Callable[[str, float], None]]], + device: torch.device): + module = modules[self.module_key] + module.train(True) + optimizers[self.module_key].zero_grad(set_to_none=True) + if create_log_func is not None: + log_func = create_log_func("training_" + self.module_key, examples_seen_so_far) + else: + log_func = None + losses[self.module_key].compute( + ComputationState(modules, accumulated_modules, batch), + log_func).backward() + if self.max_grad_norm is not None: + clip_grad_norm_(module.parameters(), self.max_grad_norm) + optimizers[self.module_key].step() + + +class SingleNetworkValidationProtocol(ValidationProtocol): + def __init__( + self, + example_per_validation_iteration: int, + batch_size: int, + module_key: str = KEY_NETWORK): + super().__init__() + self.module_key = module_key + self.batch_size = batch_size + self.example_per_validation_iteration = example_per_validation_iteration + + def get_batch_size(self, ) -> int: + return self.batch_size + + def get_examples_per_validation_iteration(self) -> int: + return self.example_per_validation_iteration + + def run_validation_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + losses: Dict[str, Loss], + create_log_func: Callable[[str, int], Callable[[str, float], None]], + device: torch.device): + module = modules[self.module_key] + module.train(False) + with torch.no_grad(): + log_func = create_log_func("validation_" + self.module_key, examples_seen_so_far) + losses[self.module_key].compute( + ComputationState(modules, accumulated_modules, batch), + log_func) diff --git a/tha4/src/tha4/shion/base/training/single_network_with_minibatch.py b/tha4/src/tha4/shion/base/training/single_network_with_minibatch.py new file mode 100644 index 0000000..e767b9d --- /dev/null +++ b/tha4/src/tha4/shion/base/training/single_network_with_minibatch.py @@ -0,0 +1,87 @@ +import time +from typing import List, Dict, Callable, Any, Optional + +import torch +from torch.nn import Module +from torch.nn.utils import clip_grad_norm_ +from torch.optim.optimizer import Optimizer + +from tha4.shion.core.cached_computation import ComputationState +from tha4.shion.core.loss import Loss +from tha4.shion.core.optimizer_factory import OptimizerFactory +from tha4.shion.core.training.training_protocol import TrainingProtocol +from tha4.shion.core.training.validation_protocol import ValidationProtocol + +KEY_NETWORK = "network" + + +class SingleNetworkWithMinibatchTrainingProtocol(TrainingProtocol): + def __init__(self, + check_point_examples: List[int], + batch_size: int, + minibatch_size: int, + learning_rate: Callable[[int], Dict[str, float]], + optimizer_factories: Dict[str, OptimizerFactory], + module_key: str = KEY_NETWORK, + random_seed: int = 39549059840, + max_grad_norm: Optional[float] = None): + super().__init__() + assert batch_size % minibatch_size == 0 + self.minibatch_size = minibatch_size + self.max_grad_norm = max_grad_norm + self.module_key = module_key + self.optimizer_factories = optimizer_factories + self.learning_rate = learning_rate + self.batch_size = batch_size + self.random_seed = random_seed + self.check_point_examples = check_point_examples + + def get_optimizer_factories(self) -> Dict[str, OptimizerFactory]: + return self.optimizer_factories + + def get_checkpoint_examples(self) -> List[int]: + return self.check_point_examples + + def get_random_seed(self) -> int: + return self.random_seed + + def get_batch_size(self) -> int: + return self.batch_size + + def get_learning_rate(self, examples_seen_so_far: int) -> Dict[str, float]: + return self.learning_rate(examples_seen_so_far) + + def run_training_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer], + losses: Dict[str, Loss], + create_log_func: Optional[Callable[[str, int], Callable[[str, float], None]]], + device: torch.device): + module = modules[self.module_key] + module.train(True) + optimizers[self.module_key].zero_grad(set_to_none=True) + if create_log_func is not None: + log_func = create_log_func("training_" + self.module_key, examples_seen_so_far) + else: + log_func = None + + num_minibatch = self.batch_size // self.minibatch_size + for minibatch_index in range(num_minibatch): + minibatch = [] + for item in batch: + minibatch.append( + item[minibatch_index * self.minibatch_size:(minibatch_index + 1) * self.minibatch_size]) + loss = losses[self.module_key].compute( + ComputationState(modules, accumulated_modules, minibatch), + log_func if minibatch_index == 0 else None) + loss = loss / num_minibatch + loss.backward() + + if self.max_grad_norm is not None: + clip_grad_norm_(module.parameters(), self.max_grad_norm) + + optimizers[self.module_key].step() diff --git a/tha4/src/tha4/shion/base/training/two_networks_training_protocol.py b/tha4/src/tha4/shion/base/training/two_networks_training_protocol.py new file mode 100644 index 0000000..8acf495 --- /dev/null +++ b/tha4/src/tha4/shion/base/training/two_networks_training_protocol.py @@ -0,0 +1,106 @@ +from typing import List, Dict, Callable, Any, Optional + +import torch +from torch.nn import Module +from torch.nn.utils import clip_grad_norm_ +from torch.optim.optimizer import Optimizer + +from tha4.shion.core.cached_computation import ComputationState +from tha4.shion.core.loss import Loss +from tha4.shion.core.optimizer_factory import OptimizerFactory +from tha4.shion.core.training.training_protocol import TrainingProtocol + + +class TwoNetworksWithMinibatchTrainingProtocol(TrainingProtocol): + def __init__(self, + check_point_examples: List[int], + batch_size: int, + learning_rate: Callable[[int], Dict[str, float]], + optimizer_factories: Dict[str, OptimizerFactory], + key_network_0: str, + key_network_1: str, + train_network_0: bool = False, + random_seed: int = 39549059840, + max_grad_norm: Optional[float] = None, + minibatch_size: Optional[int] = None): + super().__init__() + if minibatch_size is None: + minibatch_size = batch_size + assert batch_size % minibatch_size == 0 + self.train_network_0 = train_network_0 + self.key_network_1 = key_network_1 + self.key_network_0 = key_network_0 + self.minibatch_size = minibatch_size + self.max_grad_norm = max_grad_norm + self.optimizer_factories = optimizer_factories + self.learning_rate = learning_rate + self.batch_size = batch_size + self.random_seed = random_seed + self.check_point_examples = check_point_examples + + def get_optimizer_factories(self) -> Dict[str, OptimizerFactory]: + return self.optimizer_factories + + def get_checkpoint_examples(self) -> List[int]: + return self.check_point_examples + + def get_random_seed(self) -> int: + return self.random_seed + + def get_batch_size(self) -> int: + return self.batch_size + + def get_learning_rate(self, examples_seen_so_far: int) -> Dict[str, float]: + return self.learning_rate(examples_seen_so_far) + + def run_training_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer], + losses: Dict[str, Loss], + create_log_func: Optional[Callable[[str, int], Callable[[str, float], None]]], + device: torch.device): + network_0 = modules[self.key_network_0] + network_0.train(self.train_network_0) + + network_1 = modules[self.key_network_1] + network_1.train(True) + + if self.train_network_0: + optimizers[self.key_network_0].zero_grad(set_to_none=True) + optimizers[self.key_network_1].zero_grad(set_to_none=True) + + if create_log_func is not None: + network_0_log_func = create_log_func("training_" + self.key_network_0, examples_seen_so_far) + network_1_log_func = create_log_func("training_" + self.key_network_1, examples_seen_so_far) + else: + network_0_log_func = None + network_1_log_func = None + + num_minibatch = self.batch_size // self.minibatch_size + for minibatch_index in range(num_minibatch): + minibatch = [] + for item in batch: + minibatch.append( + item[minibatch_index * self.minibatch_size:(minibatch_index + 1) * self.minibatch_size]) + loss = losses[self.key_network_1].compute( + ComputationState(modules, accumulated_modules, minibatch), + network_1_log_func if minibatch_index == 0 else None) + if self.train_network_0 and self.key_network_0 in losses: + loss = loss + losses[self.key_network_0].compute( + ComputationState(modules, accumulated_modules, minibatch), + network_0_log_func if minibatch_index == 0 else None) + loss = loss / num_minibatch + loss.backward() + + if self.max_grad_norm is not None: + clip_grad_norm_(network_1.parameters(), self.max_grad_norm) + if self.train_network_0: + clip_grad_norm_(network_0.parameters(), self.max_grad_norm) + + optimizers[self.key_network_1].step() + if self.train_network_0: + optimizers[self.key_network_0].step() diff --git a/tha4/src/tha4/shion/core/__init__.py b/tha4/src/tha4/shion/core/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/core/cached_computation.py b/tha4/src/tha4/shion/core/cached_computation.py new file mode 100644 index 0000000..eac0815 --- /dev/null +++ b/tha4/src/tha4/shion/core/cached_computation.py @@ -0,0 +1,112 @@ +from abc import ABC, abstractmethod +from typing import Callable, Dict, Any, Optional + +import torch +from torch import Tensor +from torch.nn import Module + + +class ComputationState: + def __init__(self, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + batch: Any, + outputs: Optional[Dict[str, Any]] = None): + if outputs is None: + outputs = {} + self.outputs = outputs + self.batch = batch + self.accumulated_modules = accumulated_modules + self.modules = modules + + +CachedComputationFunc = Callable[[ComputationState], Any] +TensorCachedComputationFunc = Callable[[ComputationState], Tensor] + + +def create_get_item_func(func: CachedComputationFunc, index): + def _f(state: ComputationState): + output = func(state) + return output[index] + + return _f + + +def create_batch_element_func(index: int) -> TensorCachedComputationFunc: + def _f(state: ComputationState) -> Tensor: + return state.batch[index] + + return _f + + +class CachedComputationProtocol(ABC): + def get_output(self, key: str, state: ComputationState) -> Any: + if key in state.outputs: + return state.outputs[key] + else: + output = self.compute_output(key, state) + state.outputs[key] = output + return state.outputs[key] + + @abstractmethod + def compute_output(self, key: str, state: ComputationState) -> Any: + pass + + def get_output_func(self, key: str) -> CachedComputationFunc: + def func(state: ComputationState): + return self.get_output(key, state) + + return func + + +ComposableCachedComputationStep = Callable[[CachedComputationProtocol, ComputationState], Any] + + +class ComposableCachedComputationProtocol(CachedComputationProtocol): + def __init__(self, computation_steps: Optional[Dict[str, ComposableCachedComputationStep]] = None): + if computation_steps is None: + computation_steps = {} + self.computation_steps = computation_steps + + def compute_output(self, key: str, state: ComputationState) -> Any: + if key in self.computation_steps: + return self.computation_steps[key](self, state) + else: + raise RuntimeError("Computing output for key " + key + " is not supported!") + + +def batch_indexing_func(index: int): + def _f(protocol: CachedComputationProtocol, state: ComputationState): + return state.batch[index] + + return _f + + +def proxy_func(key: str): + def _f(protocol: CachedComputationProtocol, state: ComputationState): + return protocol.get_output(key, state) + + return _f + + +def output_array_indexing_func(key: str, index: int): + def _f(protocol: CachedComputationProtocol, state: ComputationState): + return protocol.get_output(key, state)[index] + + return _f + + +def add_step(step_dict: Dict[str, ComposableCachedComputationStep], name: str): + def _f(func): + step_dict[name] = func + return func + + return _f + + +def zeros_like_func(key: str): + def _f(protocol: CachedComputationProtocol, state: ComputationState): + prototype = protocol.get_output(key, state) + return torch.zeros_like(prototype) + + return _f diff --git a/tha4/src/tha4/shion/core/load_save.py b/tha4/src/tha4/shion/core/load_save.py new file mode 100644 index 0000000..d9f7e72 --- /dev/null +++ b/tha4/src/tha4/shion/core/load_save.py @@ -0,0 +1,14 @@ +import os + +import torch + + +def torch_save(content, file_name): + os.makedirs(os.path.dirname(file_name), exist_ok=True) + with open(file_name, 'wb') as f: + torch.save(content, f) + + +def torch_load(file_name): + with open(file_name, 'rb') as f: + return torch.load(f, map_location=lambda storage, loc: storage) diff --git a/tha4/src/tha4/shion/core/loss.py b/tha4/src/tha4/shion/core/loss.py new file mode 100644 index 0000000..d301051 --- /dev/null +++ b/tha4/src/tha4/shion/core/loss.py @@ -0,0 +1,15 @@ +from abc import ABC, abstractmethod +from typing import Callable, Optional + +from torch import Tensor + +from tha4.shion.core.cached_computation import ComputationState + + +class Loss(ABC): + @abstractmethod + def compute( + self, + state: ComputationState, + log_func: Optional[Callable[[str, float], None]] = None) -> Tensor: + pass diff --git a/tha4/src/tha4/shion/core/module_accumulator.py b/tha4/src/tha4/shion/core/module_accumulator.py new file mode 100644 index 0000000..e291268 --- /dev/null +++ b/tha4/src/tha4/shion/core/module_accumulator.py @@ -0,0 +1,10 @@ +from abc import ABC, abstractmethod +from typing import Optional + +from torch.nn import Module + + +class ModuleAccumulator(ABC): + @abstractmethod + def accumulate(self, module: Module, output: Module, examples_seen_so_far: Optional[int] = None) -> Module: + pass diff --git a/tha4/src/tha4/shion/core/module_factory.py b/tha4/src/tha4/shion/core/module_factory.py new file mode 100644 index 0000000..f6231c3 --- /dev/null +++ b/tha4/src/tha4/shion/core/module_factory.py @@ -0,0 +1,9 @@ +from abc import ABC, abstractmethod + +from torch.nn import Module + + +class ModuleFactory(ABC): + @abstractmethod + def create(self) -> Module: + pass \ No newline at end of file diff --git a/tha4/src/tha4/shion/core/optimizer_factory.py b/tha4/src/tha4/shion/core/optimizer_factory.py new file mode 100644 index 0000000..d5738b4 --- /dev/null +++ b/tha4/src/tha4/shion/core/optimizer_factory.py @@ -0,0 +1,10 @@ +from abc import ABC, abstractmethod +from typing import Iterable + +from torch.nn import Parameter + + +class OptimizerFactory(ABC): + @abstractmethod + def create(self, parameters: Iterable[Parameter]): + pass diff --git a/tha4/src/tha4/shion/core/training/__init__.py b/tha4/src/tha4/shion/core/training/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/core/training/distrib/__init__.py b/tha4/src/tha4/shion/core/training/distrib/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/core/training/distrib/device_mapper.py b/tha4/src/tha4/shion/core/training/distrib/device_mapper.py new file mode 100644 index 0000000..29673a7 --- /dev/null +++ b/tha4/src/tha4/shion/core/training/distrib/device_mapper.py @@ -0,0 +1,17 @@ +from typing import Dict + +import torch + + +class SimpleCudaDeviceMapper: + def __call__(self, rank, local_rank): + return torch.device("cuda", local_rank) + + +class UserSpecifiedLocalRankToDeviceMapper: + def __init__(self, device_map: Dict[int, torch.device]): + self.device_map = device_map + + def __call__(self, rank, local_rank): + assert local_rank in self.device_map + return self.device_map[local_rank] diff --git a/tha4/src/tha4/shion/core/training/distrib/distributed_trainer.py b/tha4/src/tha4/shion/core/training/distrib/distributed_trainer.py new file mode 100644 index 0000000..988ac27 --- /dev/null +++ b/tha4/src/tha4/shion/core/training/distrib/distributed_trainer.py @@ -0,0 +1,419 @@ +import argparse +import logging +import os.path +import time +from datetime import datetime +from typing import Dict, Optional, Callable, Any + +import torch +import torch.distributed +from torch.nn.parallel import DistributedDataParallel +from torch.utils.data import Dataset, DataLoader, DistributedSampler +from torch.utils.tensorboard import SummaryWriter + +from tha4.shion.core.load_save import torch_save, torch_load +from tha4.shion.core.loss import Loss +from tha4.shion.core.module_accumulator import ModuleAccumulator +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.core.training.distrib.device_mapper import SimpleCudaDeviceMapper +from tha4.shion.core.training.distrib.distributed_training_states import DistributedTrainingState +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.shion.core.training.training_protocol import TrainingProtocol +from tha4.shion.core.training.util import set_learning_rate, create_log_func, get_least_greater_multiple +from tha4.shion.core.training.validation_protocol import ValidationProtocol + +KEY_CHECKPOINT = 'checkpoint' +KEY_SNAPSHOT = 'snapshot' +KEY_VALIDATION = 'validation' +KEY_SAMPLE_OUTPUT = 'sample_output' + + +class DistributedTrainer: + def __init__(self, + prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + losses: Dict[str, Loss], + training_dataset: Dataset, + validation_dataset: Optional[Dataset], + training_protocol: TrainingProtocol, + validation_protocol: Optional[ValidationProtocol], + sample_output_protocol: Optional[SampleOutputProtocol], + pretrained_module_file_names: Dict[str, str], + example_per_snapshot: int, + num_data_loader_workers: int = 8, + distrib_backend: str = 'gloo'): + self.distrib_backend = distrib_backend + self.num_data_loader_workers = num_data_loader_workers + self.accumulators = accumulators + self.sample_output_protocol = sample_output_protocol + self.example_per_snapshot = example_per_snapshot + self.pretrained_module_file_names = pretrained_module_file_names + self.losses = losses + self.validation_protocol = validation_protocol + self.training_protocol = training_protocol + self.module_factories = module_factories + self.prefix = prefix + self.training_dataset = training_dataset + self.validation_dataset = validation_dataset + + self.checkpoint_examples = self.training_protocol.get_checkpoint_examples() + assert len(self.checkpoint_examples) >= 1 + assert self.checkpoint_examples[0] > 0 + self.checkpoint_examples = [0] + self.checkpoint_examples + + self.module_names = self.module_factories.keys() + assert len(self.module_names) > 0 + + self.training_data_loader = None + self.training_data_loader_iter = None + self.training_data_loader_batch_size = None + self.training_data_sampler = None + + self.validation_data_loader = None + self.validation_data_loader_iter = None + self.validation_data_loader_batch_size = None + + self.sample_output_data = None + self.summary_writer = None + self.log_dir = None + self.training_state = None + + def get_sample_output_data_file_name(self): + return self.prefix + "/sample_output_data.pt" + + def save_sample_output_data(self, rank: int, device: torch.device): + if rank != 0: + return + if os.path.exists(self.get_sample_output_data_file_name()): + return + if self.sample_output_protocol is not None: + torch.manual_seed(self.sample_output_protocol.get_random_seed()) + sample_output_data = self.sample_output_protocol.get_sample_output_data(self.validation_dataset, device) + torch_save(sample_output_data, self.get_sample_output_data_file_name()) + else: + torch_save({}, self.get_sample_output_data_file_name()) + + def load_sample_output_data(self, rank: int, device: torch.device): + if rank != 0: + return None + else: + self.save_sample_output_data(rank, device) + return torch_load(self.get_sample_output_data_file_name()) + + def get_snapshot_prefix(self) -> str: + return self.prefix + "/snapshot" + + def can_load_training_state(self, prefix: str, world_size: int) -> bool: + return DistributedTrainingState.can_load( + prefix, + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + world_size) + + def load_training_state(self, prefix, rank: int, local_rank: int, device: torch.device) -> DistributedTrainingState: + return DistributedTrainingState.load( + prefix, + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + rank, + local_rank, + device) + + @staticmethod + def checkpoint_prefix(prefix: str, checkpoint_index: int) -> str: + return "%s/checkpoint/%04d" % (prefix, checkpoint_index) + + def get_checkpoint_prefix(self, checkpoint_index) -> str: + return DistributedTrainer.checkpoint_prefix(self.prefix, checkpoint_index) + + def get_initial_training_state(self, rank: int, local_rank: int, device: torch.device) -> DistributedTrainingState: + training_state = DistributedTrainingState.new( + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + self.training_protocol.get_random_seed(), + rank, + local_rank, + device, + self.pretrained_module_file_names) + logging.info("Created a new initial training state.") + return training_state + + def load_previous_training_state(self, + target_checkpoint_examples: int, + world_size: int, + rank: int, + local_rank: int, + device: torch.device) -> DistributedTrainingState: + if self.can_load_training_state(self.get_snapshot_prefix(), world_size): + examples_seen_so_far = DistributedTrainingState.get_examples_seen_so_far(self.get_snapshot_prefix()) + diff = examples_seen_so_far - target_checkpoint_examples + if diff < self.training_protocol.get_batch_size(): + return self.load_training_state(self.get_snapshot_prefix(), rank, local_rank, device) + num_checkpoints = len(self.checkpoint_examples) + for checkpoint_index in range(num_checkpoints - 1, -1, -1): + if self.can_load_training_state(self.get_checkpoint_prefix(checkpoint_index), world_size): + examples_seen_so_far = DistributedTrainingState.get_examples_seen_so_far( + self.get_checkpoint_prefix(checkpoint_index)) + diff = examples_seen_so_far - target_checkpoint_examples + if diff < self.training_protocol.get_batch_size(): + return self.load_training_state( + self.get_checkpoint_prefix(checkpoint_index), rank, local_rank, device) + + training_state = self.get_initial_training_state(rank, local_rank, device) + training_state.save(self.get_checkpoint_prefix(0), rank, lambda: self.barrier(local_rank)) + training_state = self.load_training_state(self.get_checkpoint_prefix(0), rank, local_rank, device) + return training_state + + def get_log_dir(self): + if self.log_dir is None: + now = datetime.now() + self.log_dir = self.prefix + "/log/" + now.strftime("%Y_%m_%d__%H_%M_%S") + return self.log_dir + + def get_summary_writer(self, rank: int) -> Optional[SummaryWriter]: + if rank != 0: + return None + if self.summary_writer is None: + self.summary_writer = SummaryWriter(log_dir=self.get_log_dir()) + return self.summary_writer + + def get_effective_training_epoch_size(self, world_size: int): + batch_size = self.training_protocol.get_batch_size() + N = len(self.training_dataset) + N = (N // world_size) * world_size + N = (N // batch_size) * batch_size + return N + + def get_training_epoch_index(self, examples_seen_so_far: int, world_size: int): + epoch_size = self.get_effective_training_epoch_size(world_size) + batch_size = self.training_protocol.get_batch_size() + return (examples_seen_so_far + batch_size * world_size) // epoch_size + + def get_next_training_batch(self, examples_seen_so_far: int, world_size: int, device: torch.device): + batch_size = self.training_protocol.get_batch_size() + dataset = self.training_dataset + if self.training_data_loader is None: + self.training_data_sampler = DistributedSampler( + dataset, + shuffle=True, + drop_last=True) + self.training_data_loader = DataLoader( + dataset, + batch_size=batch_size, + sampler=self.training_data_sampler, + shuffle=False, + num_workers=self.num_data_loader_workers, + drop_last=True) + if self.training_data_loader_iter is None: + epoch_index = self.get_training_epoch_index(examples_seen_so_far, world_size) + logging.info(f"Started a new epoch: index = {epoch_index}, examples_seen_so_far = {examples_seen_so_far}") + self.training_data_sampler.set_epoch(epoch_index) + self.training_data_loader_iter = iter(self.training_data_loader) + try: + batch = next(self.training_data_loader_iter) + except StopIteration: + epoch_index = self.get_training_epoch_index(examples_seen_so_far, world_size) + logging.info(f"Started a new epoch: index = {epoch_index}, examples_seen_so_far = {examples_seen_so_far}") + self.training_data_sampler.set_epoch(epoch_index) + self.training_data_loader_iter = iter(self.training_data_loader) + batch = next(self.training_data_loader_iter) + return [x.to(device) for x in batch] + + def get_next_checkpoint_num_examples(self, examples_seen_so_far) -> int: + next_index = next( + (i for i in range(len(self.checkpoint_examples)) if self.checkpoint_examples[i] > examples_seen_so_far), + -1) + return self.checkpoint_examples[next_index] + + def get_next_snapshot_num_examples(self, examples_seen_so_far) -> int: + return get_least_greater_multiple(examples_seen_so_far, self.example_per_snapshot) + + def get_next_validation_num_examples(self, examples_seen_so_far) -> int: + if self.validation_protocol is None: + return -1 + return get_least_greater_multiple(examples_seen_so_far, + self.validation_protocol.get_examples_per_validation_iteration()) + + def get_next_sample_output_num_examples(self, examples_seen_so_far) -> int: + if self.sample_output_protocol is None: + return -1 + return get_least_greater_multiple(examples_seen_so_far, + self.sample_output_protocol.get_examples_per_sample_output()) + + def get_next_num_examples(self, examples_seen_so_far) -> Dict[str, int]: + return { + KEY_CHECKPOINT: self.get_next_checkpoint_num_examples(examples_seen_so_far), + KEY_SNAPSHOT: self.get_next_snapshot_num_examples(examples_seen_so_far), + KEY_VALIDATION: self.get_next_validation_num_examples(examples_seen_so_far), + KEY_SAMPLE_OUTPUT: self.get_next_sample_output_num_examples(examples_seen_so_far) + } + + def get_next_validation_batch(self, device: torch.device): + if self.validation_dataset is None: + return None + if self.validation_data_loader is None: + self.validation_data_loader = DataLoader( + self.validation_dataset, + batch_size=self.validation_protocol.get_batch_size(), + shuffle=True, + num_workers=1, + drop_last=True) + if self.validation_data_loader_iter is None: + self.validation_data_loader_iter = iter(self.validation_data_loader) + try: + batch = next(self.validation_data_loader_iter) + except StopIteration: + self.validation_data_loader_iter = iter(self.validation_data_loader) + batch = next(self.validation_data_loader_iter) + return [x.to(device) for x in batch] + + def get_checkpoint_index_to_save(self, examples_seen_so_far: int) -> int: + checkpoint_index = 0 + for i in range(len(self.checkpoint_examples)): + if self.checkpoint_examples[i] <= examples_seen_so_far: + checkpoint_index = i + return checkpoint_index + + def barrier(self, local_rank: int): + if self.distrib_backend == 'nccl': + torch.distributed.barrier(device_ids=[local_rank]) + else: + torch.distributed.barrier() + + def train(self, + world_size: int, + rank: int, + local_rank: int, + target_checkpoint_examples: Optional[int] = None, + device_mapper: Optional[Callable[[int, int], torch.device]] = None): + if target_checkpoint_examples is None: + target_checkpoint_examples = self.checkpoint_examples[-1] + + if device_mapper is None: + device_mapper = SimpleCudaDeviceMapper() + device = device_mapper(rank, local_rank) + + sample_output_data = self.load_sample_output_data(rank, device) + training_state = self.load_previous_training_state( + target_checkpoint_examples, world_size, rank, local_rank, device) + summary_writer = self.get_summary_writer(rank) + if summary_writer is not None: + log_func_factory = lambda name, num: create_log_func(summary_writer, name, num) + else: + log_func_factory = None + last_time = time.time() + + while training_state.examples_seen_so_far < target_checkpoint_examples: + # Set the learning rate + learning_rate_by_module_name = self.training_protocol.get_learning_rate(training_state.examples_seen_so_far) + for module_name in self.module_factories.keys(): + if module_name not in learning_rate_by_module_name or module_name not in training_state.optimizers: + continue + lr = learning_rate_by_module_name[module_name] + set_learning_rate(training_state.optimizers[module_name], lr) + if summary_writer is not None: + summary_writer.add_scalar( + module_name + "_learning_rate", lr, training_state.examples_seen_so_far) + + # One training iteration + training_batch = self.get_next_training_batch(training_state.examples_seen_so_far, world_size, device) + self.training_protocol.run_training_iteration( + training_batch, + training_state.examples_seen_so_far, + training_state.modules, + training_state.accumulated_modules, + training_state.optimizers, + self.losses, + log_func_factory, + device) + + # Accumulate model data + for module_name in self.accumulators: + new_module = training_state.modules[module_name] + if isinstance(new_module, DistributedDataParallel): + new_module = new_module.module + buffer_module = training_state.accumulated_modules[module_name] + self.accumulators[module_name].accumulate( + new_module, buffer_module, examples_seen_so_far=training_state.examples_seen_so_far) + + # Advance the number of examples seen so far + next_num_examples = self.get_next_num_examples(training_state.examples_seen_so_far) + training_state.examples_seen_so_far += self.training_protocol.get_batch_size() * world_size + + # Validation iteration + if self.validation_protocol is not None \ + and training_state.examples_seen_so_far >= next_num_examples[KEY_VALIDATION] \ + and rank == 0: + validation_batch = self.get_next_validation_batch(device) + self.validation_protocol.run_validation_iteration( + validation_batch, + training_state.examples_seen_so_far, + training_state.modules, + training_state.accumulated_modules, + self.losses, + log_func_factory, + device) + + # Save sample output + if self.sample_output_protocol is not None \ + and training_state.examples_seen_so_far >= next_num_examples[KEY_SAMPLE_OUTPUT]: + if rank == 0: + self.sample_output_protocol.save_sample_output_data( + training_state.modules, + training_state.accumulated_modules, + sample_output_data, + self.prefix + "/sample_outputs", + training_state.examples_seen_so_far, + device) + self.barrier(local_rank) + + # Save checkpoint + if training_state.examples_seen_so_far >= next_num_examples[KEY_CHECKPOINT]: + checkpoint_index = self.get_checkpoint_index_to_save(training_state.examples_seen_so_far) + training_state.save( + self.get_checkpoint_prefix(checkpoint_index), rank, lambda: self.barrier(local_rank)) + if next_num_examples[KEY_CHECKPOINT] != next_num_examples[KEY_SNAPSHOT]: + training_state.save(self.get_snapshot_prefix(), rank, lambda: self.barrier(local_rank)) + + # Save snapshot + if training_state.examples_seen_so_far >= next_num_examples[KEY_SNAPSHOT]: + training_state.save(self.get_snapshot_prefix(), rank, lambda: self.barrier(local_rank)) + + now = time.time() + if now - last_time > 10: + logging.info("Showed %d training examples." % training_state.examples_seen_so_far) + last_time = now + + @staticmethod + def get_default_arg_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser(description='Training script.') + parser.add_argument("--target_checkpoint_examples", type=int) + return parser + + @staticmethod + def run_with_args(trainer_factory: Callable[[int, str], 'DistributedTrainer'], + args, + backend: str = 'gloo', + device_mapper: Optional[Callable[[int, int], torch.device]] = None): + world_size = int(os.environ['WORLD_SIZE']) + rank = int(os.environ['RANK']) + local_rank = int(os.environ['LOCAL_RANK']) + + torch.distributed.init_process_group(backend) + trainer = trainer_factory(world_size, backend) + trainer.train(world_size, rank, local_rank, args.target_checkpoint_examples, device_mapper) + + @staticmethod + def run(trainer_factory: Callable[[int, str], 'DistributedTrainer'], + backend: str = 'gloo', + device_mapper: Optional[Callable[[int, int], torch.device]] = None, + args: Optional[Any] = None): + if args is None: + parser = DistributedTrainer.get_default_arg_parser() + args = parser.parse_args() + + DistributedTrainer.run_with_args(trainer_factory, args, backend, device_mapper) diff --git a/tha4/src/tha4/shion/core/training/distrib/distributed_training_states.py b/tha4/src/tha4/shion/core/training/distrib/distributed_training_states.py new file mode 100644 index 0000000..8eb1444 --- /dev/null +++ b/tha4/src/tha4/shion/core/training/distrib/distributed_training_states.py @@ -0,0 +1,234 @@ +import copy +import logging +import os +from typing import Dict, Optional, Callable + +import torch +from torch.nn import Module +from torch.nn.parallel import DistributedDataParallel +from torch.optim.optimizer import Optimizer + +from tha4.shion.core.load_save import torch_save, torch_load +from tha4.shion.core.module_accumulator import ModuleAccumulator +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.core.optimizer_factory import OptimizerFactory +from tha4.shion.core.training.util import optimizer_to_device + + +class DistributedTrainingState: + def __init__(self, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer]): + self.accumulated_modules = accumulated_modules + self.optimizers = optimizers + self.modules = modules + self.examples_seen_so_far = examples_seen_so_far + + @staticmethod + def get_examples_seen_so_far_file_name(prefix) -> str: + return prefix + "/examples_seen_so_far.txt" + + @staticmethod + def get_module_file_name(prefix, module_name) -> str: + return "%s/module_%s.pt" % (prefix, module_name) + + @staticmethod + def get_accumulated_module_file_name(prefix, module_name) -> str: + return "%s/accumulated_%s.pt" % (prefix, module_name) + + @staticmethod + def get_optimizer_file_name(prefix, module_name) -> str: + return "%s/optimizer_%s.pt" % (prefix, module_name) + + @staticmethod + def get_rng_state_file_name(prefix, rank: int): + return "%s/rng_state_%08d.pt" % (prefix, rank) + + def mkdir(self, prefix: str): + os.makedirs(prefix, exist_ok=True) + + def save_data(self, prefix: str, rank: int): + assert os.path.exists(prefix) + + torch_save(torch.get_rng_state(), DistributedTrainingState.get_rng_state_file_name(prefix, rank)) + logging.info("Saved %s" % DistributedTrainingState.get_rng_state_file_name(prefix, rank)) + + if rank == 0: + logging.info("Saving training state to %s" % prefix) + with open(DistributedTrainingState.get_examples_seen_so_far_file_name(prefix), "wt") as fout: + fout.write("%d\n" % self.examples_seen_so_far) + logging.info("Saved %s" % DistributedTrainingState.get_examples_seen_so_far_file_name(prefix)) + for module_name in self.modules: + file_name = DistributedTrainingState.get_module_file_name(prefix, module_name) + module = self.modules[module_name] + if isinstance(module, DistributedDataParallel): + state_dict = module.module.state_dict() + else: + state_dict = module.state_dict() + torch_save(state_dict, file_name) + logging.info("Saved %s" % file_name) + for module_name in self.accumulated_modules: + file_name = DistributedTrainingState.get_accumulated_module_file_name(prefix, module_name) + torch_save(self.accumulated_modules[module_name].state_dict(), file_name) + logging.info("Saved %s" % file_name) + for module_name in self.optimizers: + file_name = DistributedTrainingState.get_optimizer_file_name(prefix, module_name) + torch_save(self.optimizers[module_name].state_dict(), file_name) + logging.info("Saved %s" % file_name) + + logging.info("Done saving training state to %s" % prefix) + + def save(self, prefix: str, rank: int, barrier_func: Callable[[], None]): + if rank == 0: + self.mkdir(prefix) + barrier_func() + self.save_data(prefix, rank) + barrier_func() + + @staticmethod + def get_examples_seen_so_far(prefix: str) -> int: + with open(DistributedTrainingState.get_examples_seen_so_far_file_name(prefix)) as fin: + lines = fin.readlines() + return int(lines[0]) + + @staticmethod + def load( + prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + optimizer_factories: Dict[str, OptimizerFactory], + rank: int, + local_rank: int, + device: torch.device) -> 'DistributedTrainingState': + logging.info("Loading training state from %s" % prefix) + + with open(DistributedTrainingState.get_examples_seen_so_far_file_name(prefix)) as fin: + lines = fin.readlines() + examples_seen_so_far = int(lines[0]) + logging.info("Loaded %s" % DistributedTrainingState.get_examples_seen_so_far_file_name(prefix)) + + modules = { + module_name: factory.create() + for (module_name, factory) in module_factories.items() + } + for module_name in modules: + file_name = DistributedTrainingState.get_module_file_name(prefix, module_name) + module = modules[module_name] + state_dict = torch_load(file_name) + module.load_state_dict(state_dict) + module.to(device) + modules[module_name] = DistributedDataParallel( + module, + device_ids=[device.index], + output_device=device.index) + logging.info("Loaded %s" % file_name) + + accumulated_modules = {} + for module_name in accumulators: + module_factory = module_factories[module_name] + module = module_factory.create() + file_name = DistributedTrainingState.get_accumulated_module_file_name(prefix, module_name) + module.load_state_dict(torch_load(file_name)) + module.to(device) + accumulated_modules[module_name] = module + logging.info("Loaded %s" % file_name) + + optimizers = {} + for module_name in optimizer_factories: + optimizer = optimizer_factories[module_name].create(modules[module_name].parameters()) + file_name = DistributedTrainingState.get_optimizer_file_name(prefix, module_name) + optimizer.load_state_dict(torch_load(file_name)) + optimizer_to_device(optimizer, device) + optimizers[module_name] = optimizer + logging.info("Loaded %s" % file_name) + + torch.set_rng_state(torch_load(DistributedTrainingState.get_rng_state_file_name(prefix, rank))) + logging.info("Loaded %s" % DistributedTrainingState.get_rng_state_file_name(prefix, rank)) + + logging.info("Done loading training state from %s" % prefix) + + return DistributedTrainingState(examples_seen_so_far, modules, accumulated_modules, optimizers) + + @staticmethod + def new(module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + optimizer_factories: Dict[str, OptimizerFactory], + random_seed: int, + rank: int, + local_rank: int, + device: torch.device, + pretrained_module_file_names: Optional[Dict[str, str]] = None) -> 'DistributedTrainingState': + examples_seen_so_far = 0 + + modules = { + module_name: factory.create() + for (module_name, factory) in module_factories.items() + } + for module_name in modules: + modules[module_name].to(device) + if pretrained_module_file_names is not None: + for module_name in modules: + if module_name in pretrained_module_file_names: + file_name = pretrained_module_file_names[module_name] + modules[module_name].load_state_dict(torch_load(file_name)) + logging.info("Loaded initial state from %s ..." % file_name) + + accumulated_modules = {} + for module_name in accumulators: + accumulated_modules[module_name] = copy.deepcopy(modules[module_name]) + + for module_name in modules: + module = modules[module_name] + modules[module_name] = DistributedDataParallel( + module, + device_ids=[device.index], + output_device=device.index) + + optimizers = {} + for module_name in optimizer_factories: + module = modules[module_name] + optimizer = optimizer_factories[module_name].create(module.parameters()) + optimizer_to_device(optimizer, device) + optimizers[module_name] = optimizer + + torch.manual_seed(random_seed + rank) + + return DistributedTrainingState(examples_seen_so_far, modules, accumulated_modules, optimizers) + + @staticmethod + def can_load(prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + optimizer_factories: Dict[str, OptimizerFactory], + world_size: int) -> bool: + logging.info(f"Checking directory {prefix}") + if not os.path.isdir(prefix): + logging.info(f"Cannot load files in {prefix} because it is not a directory") + return False + examples_seen_so_far_file_name = DistributedTrainingState.get_examples_seen_so_far_file_name(prefix) + if not os.path.isfile(examples_seen_so_far_file_name): + logging.info(f"Cannot load files in {prefix} because {examples_seen_so_far_file_name} is not a file.") + return False + for module_name in module_factories.keys(): + file_name = DistributedTrainingState.get_module_file_name(prefix, module_name) + if not os.path.isfile(file_name): + logging.info(f"Cannot load files in {prefix} because {file_name} is not a file.") + return False + for module_name in accumulators: + file_name = DistributedTrainingState.get_accumulated_module_file_name(prefix, module_name) + if not os.path.isfile(file_name): + logging.info(f"Cannot load files in {prefix} because {file_name} is not a file.") + return False + for module_name in optimizer_factories: + file_name = DistributedTrainingState.get_optimizer_file_name(prefix, module_name) + if not os.path.isfile(file_name): + logging.info(f"Cannot load files in {prefix} because {file_name} is not a file.") + return False + for rank in range(world_size): + file_name = DistributedTrainingState.get_rng_state_file_name(prefix, rank) + if not os.path.isfile(file_name): + logging.info(f"Cannot load files in {prefix} because {file_name} is not a file.") + return False + return True diff --git a/tha4/src/tha4/shion/core/training/distrib/distributed_training_tasks.py b/tha4/src/tha4/shion/core/training/distrib/distributed_training_tasks.py new file mode 100644 index 0000000..6b65cfc --- /dev/null +++ b/tha4/src/tha4/shion/core/training/distrib/distributed_training_tasks.py @@ -0,0 +1,139 @@ +import logging +import os +import sys +from typing import Callable, List, Optional + +from tha4.pytasuku.workspace import Workspace +from tha4.shion.core.training.distrib.distributed_trainer import DistributedTrainer +from tha4.shion.core.training.distrib.distributed_training_states import DistributedTrainingState + + +def get_torchrun_executable(): + return os.path.dirname(sys.executable) + os.path.sep + "torchrun" + + +def run_distributed_training_script( + training_script_file_name: str, + num_nodes: int, + node_rank: int, + num_proc_per_node: int, + master_addr: int = "127.0.0.1", + master_port: int = 8888): + command = f"{get_torchrun_executable()} " \ + f"--nproc_per_node={num_proc_per_node} " \ + f"--nnodes={num_nodes} " \ + f"--node_rank={node_rank} " \ + f"--master_addr={master_addr} " \ + f"--master_port={master_port} " \ + f"{training_script_file_name}" + logging.info(f"Executing -- {command}") + os.system(command) + + +class RdzvConfig: + def __init__(self, id: int, port: int): + self.port = port + self.id = id + + +def run_standalone_distributed_training_script( + training_script_file_name: str, + num_proc_per_node: int, + target_checkpoint_examples: Optional[int] = None, + rdzv_config: Optional[RdzvConfig] = None): + command = f"{get_torchrun_executable()} " \ + f"--nnodes=1 " \ + f"--nproc_per_node={num_proc_per_node} " + if rdzv_config is not None: + command += f"--rdzv_endpoint=localhost:{rdzv_config.port} " + command += "--rdzv_backend=c10d " + command += f"--rdzv_id={rdzv_config.id} " + else: + command += "--standalone " + command += f"{training_script_file_name} " + if target_checkpoint_examples is not None: + command += f"--target_checkpoint_examples {target_checkpoint_examples} " + logging.info(f"Executing -- {command}") + os.system(command) + + +def define_distributed_training_tasks( + workspace: Workspace, + prefix: str, + training_script_file_name: str, + num_nodes: int, + num_proc_per_node: int, + master_addr: int = "127.0.0.1", + master_port: int = 8888): + def run_training_script_func(rank: int): + def _f(): + run_distributed_training_script( + training_script_file_name, + num_nodes, + rank, + num_proc_per_node, master_addr, + master_port) + + return _f + + for i in range(num_nodes): + workspace.create_command_task(f"{prefix}/train_node_%06d" % i, [], run_training_script_func(i)) + + +def define_standalone_distributed_training_tasks( + workspace: Workspace, + distributed_trainer_func: Callable[[int], DistributedTrainer], + training_script_file_name: str, + num_proc_per_node: int, + dependencies: Optional[List[str]] = None, + rdzv_config: Optional[RdzvConfig] = None): + trainer = distributed_trainer_func(1) + checkpoint_examples = trainer.training_protocol.get_checkpoint_examples() + assert len(checkpoint_examples) >= 1 + assert checkpoint_examples[0] > 0 + checkpoint_examples = [0] + checkpoint_examples + + if dependencies is None: + dependencies = [] + module_file_dependencies = dependencies[:] + for module_name in trainer.pretrained_module_file_names: + module_file_dependencies.append(trainer.pretrained_module_file_names[module_name]) + + def create_train_func(target_checkpoint_examples: int): + return lambda: run_standalone_distributed_training_script( + training_script_file_name, + num_proc_per_node, + target_checkpoint_examples, + rdzv_config=rdzv_config) + + train_tasks = [] + for checkpoint_index in range(0, len(checkpoint_examples)): + for module_name in trainer.module_names: + module_file_name = DistributedTrainingState.get_module_file_name( + trainer.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + module_file_name, + module_file_dependencies, + create_train_func(trainer.checkpoint_examples[checkpoint_index])) + for module_name in trainer.accumulators: + accumulated_module_file_name = DistributedTrainingState.get_accumulated_module_file_name( + trainer.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + accumulated_module_file_name, + module_file_dependencies, + create_train_func(checkpoint_examples[checkpoint_index])) + workspace.create_command_task( + trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standalone", + module_file_dependencies, + create_train_func(checkpoint_examples[checkpoint_index])) + train_tasks.append(trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standlone") + workspace.create_file_task( + trainer.prefix + "/train_standalone", + module_file_dependencies, + create_train_func(checkpoint_examples[-1])) + + +if __name__ == "__main__": + print(os.path.dirname(sys.executable) + os.path.sep + "torchrun") diff --git a/tha4/src/tha4/shion/core/training/sample_output_protocol.py b/tha4/src/tha4/shion/core/training/sample_output_protocol.py new file mode 100644 index 0000000..cc7705d --- /dev/null +++ b/tha4/src/tha4/shion/core/training/sample_output_protocol.py @@ -0,0 +1,43 @@ +from abc import ABC, abstractmethod +from typing import Dict, Any + +import torch +from torch.nn import Module +from torch.utils.data import Dataset + + +class SampleOutputProtocol(ABC): + @abstractmethod + def get_examples_per_sample_output(self) -> int: + pass + + @abstractmethod + def get_random_seed(self) -> int: + pass + + @abstractmethod + def get_sample_output_data(self, validation_dataset: Dataset, device: torch.device) -> Any: + pass + + @abstractmethod + def save_sample_output_data( + self, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + sample_output_data: Any, + prefix: str, + examples_seen_so_far: int, + device: torch.device): + pass + + +class AbstractSampleOutputProtocol(SampleOutputProtocol, ABC): + def __init__(self, examples_per_sample_output: int, random_seed: int): + self.random_seed = random_seed + self.examples_per_sample_output = examples_per_sample_output + + def get_examples_per_sample_output(self) -> int: + return self.examples_per_sample_output + + def get_random_seed(self) -> int: + return self.random_seed diff --git a/tha4/src/tha4/shion/core/training/single/__init__.py b/tha4/src/tha4/shion/core/training/single/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/core/training/single/training_states.py b/tha4/src/tha4/shion/core/training/single/training_states.py new file mode 100644 index 0000000..79c4434 --- /dev/null +++ b/tha4/src/tha4/shion/core/training/single/training_states.py @@ -0,0 +1,182 @@ +import copy +import logging +import os +from typing import Dict, Optional + +import torch +from torch.nn import Module +from torch.optim import Optimizer + +from tha4.shion.core.load_save import torch_save, torch_load +from tha4.shion.core.module_accumulator import ModuleAccumulator +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.core.optimizer_factory import OptimizerFactory +from tha4.shion.core.training.util import optimizer_to_device + + +class TrainingState: + def __init__(self, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer]): + self.accumulated_modules = accumulated_modules + self.optimizers = optimizers + self.modules = modules + self.examples_seen_so_far = examples_seen_so_far + + @staticmethod + def get_examples_seen_so_far_file_name(prefix) -> str: + return prefix + "/examples_seen_so_far.txt" + + @staticmethod + def get_module_file_name(prefix, module_name) -> str: + return "%s/module_%s.pt" % (prefix, module_name) + + @staticmethod + def get_accumulated_module_file_name(prefix, module_name) -> str: + return "%s/accumulated_%s.pt" % (prefix, module_name) + + @staticmethod + def get_optimizer_file_name(prefix, module_name) -> str: + return "%s/optimizer_%s.pt" % (prefix, module_name) + + @staticmethod + def get_rng_state_file_name(prefix): + return "%s/rng_state.pt" % prefix + + def save(self, prefix): + logging.info("Saving training state to %s" % prefix) + os.makedirs(prefix, exist_ok=True) + with open(TrainingState.get_examples_seen_so_far_file_name(prefix), "wt") as fout: + fout.write("%d\n" % self.examples_seen_so_far) + logging.info("Saved %s" % TrainingState.get_examples_seen_so_far_file_name(prefix)) + for module_name in self.modules: + file_name = TrainingState.get_module_file_name(prefix, module_name) + torch_save(self.modules[module_name].state_dict(), file_name) + logging.info("Saved %s" % file_name) + for module_name in self.accumulated_modules: + file_name = TrainingState.get_accumulated_module_file_name(prefix, module_name) + torch_save(self.accumulated_modules[module_name].state_dict(), file_name) + logging.info("Saved %s" % file_name) + for module_name in self.optimizers: + file_name = TrainingState.get_optimizer_file_name(prefix, module_name) + torch_save(self.optimizers[module_name].state_dict(), file_name) + logging.info("Saved %s" % file_name) + torch_save(torch.get_rng_state(), TrainingState.get_rng_state_file_name(prefix)) + logging.info("Saved %s" % TrainingState.get_rng_state_file_name(prefix)) + logging.info("Done saving training state to %s" % prefix) + + @staticmethod + def get_examples_seen_so_far(prefix: str) -> int: + with open(TrainingState.get_examples_seen_so_far_file_name(prefix)) as fin: + lines = fin.readlines() + return int(lines[0]) + + @staticmethod + def load(prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + optimizer_factories: Dict[str, OptimizerFactory], + device: torch.device) -> 'TrainingState': + logging.info("Loading training state from %s" % prefix) + + with open(TrainingState.get_examples_seen_so_far_file_name(prefix)) as fin: + lines = fin.readlines() + examples_seen_so_far = int(lines[0]) + logging.info("Loaded %s" % TrainingState.get_examples_seen_so_far_file_name(prefix)) + + modules = { + module_name: factory.create() + for (module_name, factory) in module_factories.items() + } + for module_name in modules: + file_name = TrainingState.get_module_file_name(prefix, module_name) + modules[module_name].load_state_dict(torch_load(file_name)) + modules[module_name].to(device) + logging.info("Loaded %s" % file_name) + + accumulated_modules = {} + for module_name in accumulators: + module_factory = module_factories[module_name] + module = module_factory.create() + file_name = TrainingState.get_accumulated_module_file_name(prefix, module_name) + module.load_state_dict(torch_load(file_name)) + module.to(device) + accumulated_modules[module_name] = module + logging.info("Loaded %s" % file_name) + + optimizers = {} + for module_name in optimizer_factories: + optimizer = optimizer_factories[module_name].create(modules[module_name].parameters()) + file_name = TrainingState.get_optimizer_file_name(prefix, module_name) + optimizer.load_state_dict(torch_load(file_name)) + optimizer_to_device(optimizer, device) + optimizers[module_name] = optimizer + logging.info("Loaded %s" % file_name) + + torch.set_rng_state(torch_load(TrainingState.get_rng_state_file_name(prefix))) + logging.info("Loaded %s" % TrainingState.get_rng_state_file_name(prefix)) + + logging.info("Done loading training state from %s" % prefix) + + return TrainingState(examples_seen_so_far, modules, accumulated_modules, optimizers) + + @staticmethod + def new(module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + optimizer_factories: Dict[str, OptimizerFactory], + random_seed: int, + device: torch.device, + pretrained_module_file_names: Optional[Dict[str, str]] = None) -> 'TrainingState': + examples_seen_so_far = 0 + + modules = { + module_name: factory.create() + for (module_name, factory) in module_factories.items() + } + for module_name in modules: + modules[module_name].to(device) + if pretrained_module_file_names is not None: + for module_name in modules: + if module_name in pretrained_module_file_names: + file_name = pretrained_module_file_names[module_name] + modules[module_name].load_state_dict(torch_load(file_name)) + logging.info("Loaded initial state from %s ..." % file_name) + + accumulated_modules = {} + for module_name in accumulators: + accumulated_modules[module_name] = copy.deepcopy(modules[module_name]) + + optimizers = {} + for module_name in optimizer_factories: + module = modules[module_name] + optimizer = optimizer_factories[module_name].create(module.parameters()) + optimizer_to_device(optimizer, device) + optimizers[module_name] = optimizer + + torch.manual_seed(random_seed) + + return TrainingState(examples_seen_so_far, modules, accumulated_modules, optimizers) + + @staticmethod + def can_load(prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + optimizer_factories: Dict[str, OptimizerFactory]) -> bool: + if not os.path.isdir(prefix): + return False + if not os.path.isfile(TrainingState.get_examples_seen_so_far_file_name(prefix)): + return False + for module_name in module_factories.keys(): + if not os.path.isfile(TrainingState.get_module_file_name(prefix, module_name)): + return False + for module_name in accumulators: + if not os.path.isfile(TrainingState.get_accumulated_module_file_name(prefix, module_name)): + return False + for module_name in optimizer_factories: + if not os.path.isfile(TrainingState.get_optimizer_file_name(prefix, module_name)): + return False + if not os.path.isfile(TrainingState.get_rng_state_file_name(prefix)): + return False + return True diff --git a/tha4/src/tha4/shion/core/training/single/training_tasks.py b/tha4/src/tha4/shion/core/training/single/training_tasks.py new file mode 100644 index 0000000..78a9c4e --- /dev/null +++ b/tha4/src/tha4/shion/core/training/single/training_tasks.py @@ -0,0 +1,357 @@ +import logging +import time +from datetime import datetime +from typing import Optional, Dict, List + +import torch +from torch.utils.data import Dataset, DataLoader +from torch.utils.tensorboard import SummaryWriter + +from tha4.pytasuku.workspace import Workspace +from tha4.shion.core.load_save import torch_save, torch_load +from tha4.shion.core.loss import Loss +from tha4.shion.core.module_accumulator import ModuleAccumulator +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.shion.core.training.single.training_states import TrainingState +from tha4.shion.core.training.training_protocol import TrainingProtocol +from tha4.shion.core.training.util import get_least_greater_multiple, create_log_func, set_learning_rate +from tha4.shion.core.training.validation_protocol import ValidationProtocol + +KEY_CHECKPOINT = 'checkpoint' +KEY_SNAPSHOT = 'snapshot' +KEY_VALIDATION = 'validation' +KEY_SAMPLE_OUTPUT = 'sample_output' + + +class TrainingTasks: + def __init__( + self, + workspace: Workspace, + prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + losses: Dict[str, Loss], + training_dataset: Dataset, + validation_dataset: Optional[Dataset], + training_protocol: TrainingProtocol, + validation_protocol: Optional[ValidationProtocol], + sample_output_protocol: Optional[SampleOutputProtocol], + pretrained_module_file_names: Dict[str, str], + example_per_snapshot: int, + device: torch.device, + num_data_loader_workers: int = 8, + dependencies: Optional[List[str]] = None): + super().__init__() + self.num_data_loader_workers = num_data_loader_workers + self.accumulators = accumulators + self.device = device + self.sample_output_protocol = sample_output_protocol + self.example_per_snapshot = example_per_snapshot + self.pretrained_module_file_names = pretrained_module_file_names + self.losses = losses + self.validation_protocol = validation_protocol + self.training_protocol = training_protocol + self.module_factories = module_factories + self.prefix = prefix + self.training_dataset = training_dataset + self.validation_dataset = validation_dataset + + self.checkpoint_examples = self.training_protocol.get_checkpoint_examples() + assert len(self.checkpoint_examples) >= 1 + assert self.checkpoint_examples[0] > 0 + self.checkpoint_examples = [0] + self.checkpoint_examples + + self.module_names = self.module_factories.keys() + assert len(self.module_names) > 0 + + self.training_data_loader = None + self.training_data_loader_iter = None + self.training_data_loader_batch_size = None + self.validation_data_loader = None + self.validation_data_loader_iter = None + self.validation_data_loader_batch_size = None + self.sample_output_data = None + self.summary_writer = None + self.log_dir = None + self.training_state = None + + if dependencies is None: + dependencies = [] + self.sample_output_data_task = workspace.create_file_task( + self.get_sample_output_data_file_name(), dependencies, self.save_sample_output_data) + + module_file_dependencies = [self.sample_output_data_task.name] + for module_name in pretrained_module_file_names: + module_file_dependencies.append(self.pretrained_module_file_names[module_name]) + + def create_train_func(target_examples: int): + return lambda: self.train(target_examples) + + train_tasks = [] + for checkpoint_index in range(1, len(self.checkpoint_examples)): + for module_name in self.module_names: + module_file_name = TrainingState.get_module_file_name( + self.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + module_file_name, + module_file_dependencies, + create_train_func(self.checkpoint_examples[checkpoint_index])) + for module_name in self.accumulators: + accumulated_module_file_name = TrainingState.get_accumulated_module_file_name( + self.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + accumulated_module_file_name, + module_file_dependencies, + create_train_func(self.checkpoint_examples[checkpoint_index])) + workspace.create_command_task( + self.get_checkpoint_prefix(checkpoint_index) + "/train", + module_file_dependencies, + create_train_func(self.checkpoint_examples[checkpoint_index])) + train_tasks.append(self.get_checkpoint_prefix(checkpoint_index) + "/train") + + self.train_task = workspace.create_file_task( + self.get_train_command_name(), + module_file_dependencies, + create_train_func(self.checkpoint_examples[-1])) + + def get_sample_output_data_file_name(self): + return self.prefix + "/sample_output_data.pt" + + def save_sample_output_data(self): + if self.sample_output_protocol is not None: + torch.manual_seed(self.sample_output_protocol.get_random_seed()) + sample_output_data = self.sample_output_protocol.get_sample_output_data(self.validation_dataset, + self.device) + torch_save(sample_output_data, self.get_sample_output_data_file_name()) + else: + torch_save({}, self.get_sample_output_data_file_name()) + + def get_module_file_name(self, checkpoint_index, module_name): + return TrainingState.get_module_file_name(self.get_checkpoint_prefix(checkpoint_index), module_name) + + def get_last_module_file_name(self, module_name): + return self.get_module_file_name(len(self.checkpoint_examples) - 1, module_name) + + def get_log_dir(self): + if self.log_dir is None: + now = datetime.now() + self.log_dir = self.prefix + "/log/" + now.strftime("%Y_%m_%d__%H_%M_%S") + return self.log_dir + + def get_summary_writer(self) -> SummaryWriter: + if self.summary_writer is None: + self.summary_writer = SummaryWriter(log_dir=self.get_log_dir()) + return self.summary_writer + + def get_train_command_name(self) -> str: + return self.prefix + "/train" + + def get_snapshot_prefix(self) -> str: + return self.prefix + "/snapshot" + + def get_checkpoint_prefix(self, checkpoint_index) -> str: + return "%s/checkpoint/%04d" % (self.prefix, checkpoint_index) + + def can_load_training_state(self, prefix) -> bool: + return TrainingState.can_load( + prefix, + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories()) + + def load_training_state(self, prefix) -> TrainingState: + return TrainingState.load( + prefix, + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + self.device) + + def get_initial_training_state(self) -> TrainingState: + training_state = TrainingState.new( + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + self.training_protocol.get_random_seed(), + self.device, + self.pretrained_module_file_names) + logging.info("Created a new initial training state.") + return training_state + + def load_previous_training_state(self, target_checkpoint_examples: int) -> TrainingState: + if self.can_load_training_state(self.get_snapshot_prefix()): + examples_seen_so_far = TrainingState.get_examples_seen_so_far(self.get_snapshot_prefix()) + diff = examples_seen_so_far - target_checkpoint_examples + if diff < self.training_protocol.get_batch_size(): + return self.load_training_state(self.get_snapshot_prefix()) + num_checkpoints = len(self.checkpoint_examples) + for checkpoint_index in range(num_checkpoints - 1, -1, -1): + if self.can_load_training_state(self.get_checkpoint_prefix(checkpoint_index)): + examples_seen_so_far = TrainingState.get_examples_seen_so_far( + self.get_checkpoint_prefix(checkpoint_index)) + diff = examples_seen_so_far - target_checkpoint_examples + if diff < self.training_protocol.get_batch_size(): + return self.load_training_state(self.get_checkpoint_prefix(checkpoint_index)) + return self.get_initial_training_state() + + def get_next_checkpoint_num_examples(self, examples_seen_so_far) -> int: + next_index = next( + (i for i in range(len(self.checkpoint_examples)) if self.checkpoint_examples[i] > examples_seen_so_far), + -1) + return self.checkpoint_examples[next_index] + + def get_next_snapshot_num_examples(self, examples_seen_so_far) -> int: + return get_least_greater_multiple(examples_seen_so_far, self.example_per_snapshot) + + def get_next_validation_num_examples(self, examples_seen_so_far) -> int: + if self.validation_protocol is None: + return -1 + return get_least_greater_multiple(examples_seen_so_far, + self.validation_protocol.get_examples_per_validation_iteration()) + + def get_next_sample_output_num_examples(self, examples_seen_so_far) -> int: + if self.sample_output_protocol is None: + return -1 + return get_least_greater_multiple(examples_seen_so_far, + self.sample_output_protocol.get_examples_per_sample_output()) + + def get_next_num_examples(self, examples_seen_so_far) -> Dict[str, int]: + return { + KEY_CHECKPOINT: self.get_next_checkpoint_num_examples(examples_seen_so_far), + KEY_SNAPSHOT: self.get_next_snapshot_num_examples(examples_seen_so_far), + KEY_VALIDATION: self.get_next_validation_num_examples(examples_seen_so_far), + KEY_SAMPLE_OUTPUT: self.get_next_sample_output_num_examples(examples_seen_so_far) + } + + def get_checkpoint_index_to_save(self, examples_seen_so_far: int) -> int: + checkpoint_index = 0 + for i in range(len(self.checkpoint_examples)): + if self.checkpoint_examples[i] <= examples_seen_so_far: + checkpoint_index = i + return checkpoint_index + + def get_next_training_batch(self): + if self.training_data_loader is None: + self.training_data_loader = DataLoader( + self.training_dataset, + batch_size=self.training_protocol.get_batch_size(), + shuffle=True, + num_workers=self.num_data_loader_workers, + drop_last=True) + if self.training_data_loader_iter is None: + self.training_data_loader_iter = iter(self.training_data_loader) + try: + batch = next(self.training_data_loader_iter) + except StopIteration: + self.training_data_loader_iter = iter(self.training_data_loader) + batch = next(self.training_data_loader_iter) + return [x.to(self.device) for x in batch] + + def get_next_validation_batch(self): + if self.validation_dataset is None: + return None + if self.validation_data_loader is None: + self.validation_data_loader = DataLoader( + self.validation_dataset, + batch_size=self.validation_protocol.get_batch_size(), + shuffle=True, + num_workers=self.num_data_loader_workers, + drop_last=True) + if self.validation_data_loader_iter is None: + self.validation_data_loader_iter = iter(self.validation_data_loader) + try: + batch = next(self.validation_data_loader_iter) + except StopIteration: + self.validation_data_loader_iter = iter(self.validation_data_loader) + batch = next(self.validation_data_loader_iter) + return [x.to(self.device) for x in batch] + + def get_checkpoint_index(self, target_checkpoint_examples: int): + return self.checkpoint_examples.index(target_checkpoint_examples) + + def train(self, target_checkpoint_examples: Optional[int] = None): + if target_checkpoint_examples is None: + target_checkpoint_examples = self.checkpoint_examples[-1] + + sample_output_data = torch_load(self.get_sample_output_data_file_name()) + logging.info("Loaded sampled output data from %s", self.get_sample_output_data_file_name()) + training_state = self.load_previous_training_state(target_checkpoint_examples) + summary_writer = self.get_summary_writer() + last_time = time.time() + + while training_state.examples_seen_so_far < target_checkpoint_examples: + # One training iteration + learning_rate = self.training_protocol.get_learning_rate(training_state.examples_seen_so_far) + for module_name in self.module_factories.keys(): + if module_name not in learning_rate or module_name not in training_state.optimizers: + continue + lr = learning_rate[module_name] + set_learning_rate(training_state.optimizers[module_name], lr) + self.get_summary_writer().add_scalar( + module_name + "_learning_rate", lr, training_state.examples_seen_so_far) + training_batch = self.get_next_training_batch() + self.training_protocol.run_training_iteration( + training_batch, + training_state.examples_seen_so_far, + training_state.modules, + training_state.accumulated_modules, + training_state.optimizers, + self.losses, + lambda name, num: create_log_func(summary_writer, name, num), + self.device) + + # Accumulate model data + for module_name in self.accumulators: + new_module = training_state.modules[module_name] + buffer_module = training_state.accumulated_modules[module_name] + self.accumulators[module_name].accumulate( + new_module, + buffer_module, + training_state.examples_seen_so_far) + + # Advance the number of examples seen so far + next_num_examples = self.get_next_num_examples(training_state.examples_seen_so_far) + training_state.examples_seen_so_far += self.training_protocol.get_batch_size() + + # Validation iteration + if self.validation_protocol is not None \ + and training_state.examples_seen_so_far >= next_num_examples[KEY_VALIDATION]: + validation_batch = self.get_next_validation_batch() + self.validation_protocol.run_validation_iteration( + validation_batch, + training_state.examples_seen_so_far, + training_state.modules, + training_state.accumulated_modules, + self.losses, + lambda name, num: create_log_func(summary_writer, name, num), + self.device) + + # Save sample output + if self.sample_output_protocol is not None \ + and training_state.examples_seen_so_far >= next_num_examples[KEY_SAMPLE_OUTPUT]: + self.sample_output_protocol.save_sample_output_data( + training_state.modules, + training_state.accumulated_modules, + sample_output_data, + self.prefix + "/sample_outputs", + training_state.examples_seen_so_far, + self.device) + + # Save checkpoint + if training_state.examples_seen_so_far >= next_num_examples[KEY_CHECKPOINT]: + checkpoint_index = self.get_checkpoint_index_to_save(training_state.examples_seen_so_far) + training_state.save(self.get_checkpoint_prefix(checkpoint_index)) + if next_num_examples[KEY_CHECKPOINT] != next_num_examples[KEY_SNAPSHOT]: + training_state.save(self.get_snapshot_prefix()) + + # Save snapshot + if training_state.examples_seen_so_far >= next_num_examples[KEY_SNAPSHOT]: + training_state.save(self.get_snapshot_prefix()) + + now = time.time() + if now - last_time > 10: + logging.info("Showed %d training examples." % training_state.examples_seen_so_far) + last_time = now diff --git a/tha4/src/tha4/shion/core/training/swarm/__init__.py b/tha4/src/tha4/shion/core/training/swarm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/core/training/swarm/swarm_training_tasks.py b/tha4/src/tha4/shion/core/training/swarm/swarm_training_tasks.py new file mode 100644 index 0000000..06360ef --- /dev/null +++ b/tha4/src/tha4/shion/core/training/swarm/swarm_training_tasks.py @@ -0,0 +1,62 @@ +from typing import Callable, Optional, List + +from tha4.pytasuku.workspace import Workspace +from tha4.shion.core.training.distrib.distributed_training_tasks import RdzvConfig, \ + run_standalone_distributed_training_script +from tha4.shion.core.training.single.training_states import TrainingState +from tha4.shion.core.training.swarm.swarm_unit_trainer import SwarmUnitTrainer + + +def define_standalone_swarm_training_tasks( + workspace: Workspace, + swarm_unit_trainer_func: Callable[[], SwarmUnitTrainer], + training_script_file_name: str, + num_proc_per_node: int, + dependencies: Optional[List[str]] = None, + rdzv_config: Optional[RdzvConfig] = None): + trainer = swarm_unit_trainer_func() + checkpoint_examples = trainer.training_protocol.get_checkpoint_examples() + assert len(checkpoint_examples) >= 1 + assert checkpoint_examples[0] > 0 + checkpoint_examples = [0] + checkpoint_examples + + if dependencies is None: + dependencies = [] + module_file_dependencies = dependencies[:] + for module_name in trainer.pretrained_module_file_names: + module_file_dependencies.append(trainer.pretrained_module_file_names[module_name]) + + def create_train_func(target_checkpoint_examples: int): + return lambda: run_standalone_distributed_training_script( + training_script_file_name, + num_proc_per_node, + target_checkpoint_examples, + rdzv_config=rdzv_config) + + train_tasks = [] + for checkpoint_index in range(0, len(checkpoint_examples)): + for module_name in trainer.module_names: + module_file_name = TrainingState.get_module_file_name( + trainer.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + module_file_name, + module_file_dependencies, + create_train_func(trainer.checkpoint_examples[checkpoint_index])) + for module_name in trainer.accumulators: + accumulated_module_file_name = TrainingState.get_accumulated_module_file_name( + trainer.get_checkpoint_prefix(checkpoint_index), + module_name) + workspace.create_file_task( + accumulated_module_file_name, + module_file_dependencies, + create_train_func(checkpoint_examples[checkpoint_index])) + workspace.create_command_task( + trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standalone", + module_file_dependencies, + create_train_func(checkpoint_examples[checkpoint_index])) + train_tasks.append(trainer.get_checkpoint_prefix(checkpoint_index) + "/train_standlone") + workspace.create_file_task( + trainer.prefix + "/train_standalone", + module_file_dependencies, + create_train_func(checkpoint_examples[-1])) diff --git a/tha4/src/tha4/shion/core/training/swarm/swarm_unit_trainer.py b/tha4/src/tha4/shion/core/training/swarm/swarm_unit_trainer.py new file mode 100644 index 0000000..e5ebc5a --- /dev/null +++ b/tha4/src/tha4/shion/core/training/swarm/swarm_unit_trainer.py @@ -0,0 +1,345 @@ +import argparse +import logging +import os +import time +from datetime import datetime +from typing import Dict, Optional, Callable +import torch.distributed + +import torch +from torch.utils.data import Dataset, DataLoader +from torch.utils.tensorboard import SummaryWriter + +from tha4.shion.core.load_save import torch_save, torch_load +from tha4.shion.core.loss import Loss +from tha4.shion.core.module_accumulator import ModuleAccumulator +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.core.training.distrib.device_mapper import SimpleCudaDeviceMapper +from tha4.shion.core.training.sample_output_protocol import SampleOutputProtocol +from tha4.shion.core.training.single.training_states import TrainingState +from tha4.shion.core.training.single.training_tasks import KEY_CHECKPOINT, KEY_SNAPSHOT, KEY_VALIDATION, KEY_SAMPLE_OUTPUT +from tha4.shion.core.training.training_protocol import TrainingProtocol +from tha4.shion.core.training.util import get_least_greater_multiple, create_log_func, set_learning_rate +from tha4.shion.core.training.validation_protocol import ValidationProtocol + + +class SwarmUnitTrainer: + def __init__(self, + prefix: str, + module_factories: Dict[str, ModuleFactory], + accumulators: Dict[str, ModuleAccumulator], + losses: Dict[str, Loss], + training_dataset: Dataset, + validation_dataset: Optional[Dataset], + training_protocol: TrainingProtocol, + validation_protocol: Optional[ValidationProtocol], + sample_output_protocol: Optional[SampleOutputProtocol], + pretrained_module_file_names: Dict[str, str], + example_per_snapshot: int, + num_data_loader_workers: int = 8): + self.num_data_loader_workers = num_data_loader_workers + self.accumulators = accumulators + self.sample_output_protocol = sample_output_protocol + self.example_per_snapshot = example_per_snapshot + self.pretrained_module_file_names = pretrained_module_file_names + self.losses = losses + self.validation_protocol = validation_protocol + self.training_protocol = training_protocol + self.module_factories = module_factories + self.prefix = prefix + self.training_dataset = training_dataset + self.validation_dataset = validation_dataset + + self.checkpoint_examples = self.training_protocol.get_checkpoint_examples() + assert len(self.checkpoint_examples) >= 1 + assert self.checkpoint_examples[0] > 0 + self.checkpoint_examples = [0] + self.checkpoint_examples + + self.module_names = self.module_factories.keys() + assert len(self.module_names) > 0 + + self.training_data_loader = None + self.training_data_loader_iter = None + self.training_data_loader_batch_size = None + self.training_data_sampler = None + + self.validation_data_loader = None + self.validation_data_loader_iter = None + self.validation_data_loader_batch_size = None + + self.sample_output_data = None + self.summary_writer = None + self.log_dir = None + self.training_state = None + + def get_sample_output_data_file_name(self): + return self.prefix + "/sample_output_data.pt" + + def save_sample_output_data(self, device: torch.device): + if os.path.exists(self.get_sample_output_data_file_name()): + return + if self.sample_output_protocol is not None: + torch.manual_seed(self.sample_output_protocol.get_random_seed()) + sample_output_data = self.sample_output_protocol.get_sample_output_data(self.validation_dataset, device) + torch_save(sample_output_data, self.get_sample_output_data_file_name()) + else: + torch_save({}, self.get_sample_output_data_file_name()) + + def load_sample_output_data(self, device: torch.device): + self.save_sample_output_data(device) + return torch_load(self.get_sample_output_data_file_name()) + + def get_snapshot_prefix(self) -> str: + return self.prefix + "/snapshot" + + def can_load_training_state(self, prefix: str) -> bool: + return TrainingState.can_load( + prefix, + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories()) + + def load_training_state(self, prefix, device: torch.device) -> TrainingState: + return TrainingState.load( + prefix, + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + device) + + @staticmethod + def checkpoint_prefix(prefix: str, checkpoint_index: int) -> str: + return "%s/checkpoint/%04d" % (prefix, checkpoint_index) + + def get_checkpoint_prefix(self, checkpoint_index) -> str: + return SwarmUnitTrainer.checkpoint_prefix(self.prefix, checkpoint_index) + + def get_initial_training_state(self, device: torch.device) -> TrainingState: + training_state = TrainingState.new( + self.module_factories, + self.accumulators, + self.training_protocol.get_optimizer_factories(), + self.training_protocol.get_random_seed(), + device, + self.pretrained_module_file_names) + logging.info("Created a new initial training state.") + return training_state + + def load_previous_training_state(self, + target_checkpoint_examples: int, + device: torch.device) -> TrainingState: + if self.can_load_training_state(self.get_snapshot_prefix()): + examples_seen_so_far = TrainingState.get_examples_seen_so_far(self.get_snapshot_prefix()) + diff = examples_seen_so_far - target_checkpoint_examples + if diff < self.training_protocol.get_batch_size(): + return self.load_training_state(self.get_snapshot_prefix(), device) + num_checkpoints = len(self.checkpoint_examples) + for checkpoint_index in range(num_checkpoints - 1, -1, -1): + if self.can_load_training_state(self.get_checkpoint_prefix(checkpoint_index)): + examples_seen_so_far = TrainingState.get_examples_seen_so_far( + self.get_checkpoint_prefix(checkpoint_index)) + diff = examples_seen_so_far - target_checkpoint_examples + if diff < self.training_protocol.get_batch_size(): + return self.load_training_state( + self.get_checkpoint_prefix(checkpoint_index), device) + + training_state = self.get_initial_training_state(device) + training_state.save(self.get_checkpoint_prefix(0)) + training_state = self.load_training_state(self.get_checkpoint_prefix(0), device) + return training_state + + def get_log_dir(self): + if self.log_dir is None: + now = datetime.now() + self.log_dir = self.prefix + "/log/" + now.strftime("%Y_%m_%d__%H_%M_%S") + return self.log_dir + + def get_summary_writer(self) -> Optional[SummaryWriter]: + if self.summary_writer is None: + self.summary_writer = SummaryWriter(log_dir=self.get_log_dir()) + return self.summary_writer + + def get_next_training_batch(self, device: torch.device): + if self.training_data_loader is None: + self.training_data_loader = DataLoader( + self.training_dataset, + batch_size=self.training_protocol.get_batch_size(), + shuffle=True, + num_workers=self.num_data_loader_workers, + drop_last=True) + if self.training_data_loader_iter is None: + self.training_data_loader_iter = iter(self.training_data_loader) + try: + batch = next(self.training_data_loader_iter) + except StopIteration: + self.training_data_loader_iter = iter(self.training_data_loader) + batch = next(self.training_data_loader_iter) + return [x.to(device) for x in batch] + + def get_next_checkpoint_num_examples(self, examples_seen_so_far) -> int: + next_index = next( + (i for i in range(len(self.checkpoint_examples)) if self.checkpoint_examples[i] > examples_seen_so_far), + -1) + return self.checkpoint_examples[next_index] + + def get_next_snapshot_num_examples(self, examples_seen_so_far) -> int: + return get_least_greater_multiple(examples_seen_so_far, self.example_per_snapshot) + + def get_next_validation_num_examples(self, examples_seen_so_far) -> int: + if self.validation_protocol is None: + return -1 + return get_least_greater_multiple(examples_seen_so_far, + self.validation_protocol.get_examples_per_validation_iteration()) + + def get_next_sample_output_num_examples(self, examples_seen_so_far) -> int: + if self.sample_output_protocol is None: + return -1 + return get_least_greater_multiple(examples_seen_so_far, + self.sample_output_protocol.get_examples_per_sample_output()) + + def get_next_num_examples(self, examples_seen_so_far) -> Dict[str, int]: + return { + KEY_CHECKPOINT: self.get_next_checkpoint_num_examples(examples_seen_so_far), + KEY_SNAPSHOT: self.get_next_snapshot_num_examples(examples_seen_so_far), + KEY_VALIDATION: self.get_next_validation_num_examples(examples_seen_so_far), + KEY_SAMPLE_OUTPUT: self.get_next_sample_output_num_examples(examples_seen_so_far) + } + + def get_next_validation_batch(self, device: torch.device): + if self.validation_dataset is None: + return None + if self.validation_data_loader is None: + self.validation_data_loader = DataLoader( + self.validation_dataset, + batch_size=self.validation_protocol.get_batch_size(), + shuffle=True, + num_workers=1, + drop_last=True) + if self.validation_data_loader_iter is None: + self.validation_data_loader_iter = iter(self.validation_data_loader) + try: + batch = next(self.validation_data_loader_iter) + except StopIteration: + self.validation_data_loader_iter = iter(self.validation_data_loader) + batch = next(self.validation_data_loader_iter) + return [x.to(device) for x in batch] + + def get_checkpoint_index_to_save(self, examples_seen_so_far: int) -> int: + checkpoint_index = 0 + for i in range(len(self.checkpoint_examples)): + if self.checkpoint_examples[i] <= examples_seen_so_far: + checkpoint_index = i + return checkpoint_index + + def train(self, + rank: int, + local_rank: int, + target_checkpoint_examples: Optional[int] = None, + device_mapper: Optional[Callable[[int, int], torch.device]] = None): + if target_checkpoint_examples is None: + target_checkpoint_examples = self.checkpoint_examples[-1] + + if device_mapper is None: + device_mapper = SimpleCudaDeviceMapper() + device = device_mapper(rank, local_rank) + + sample_output_data = self.load_sample_output_data(device) + training_state = self.load_previous_training_state( + target_checkpoint_examples, device) + summary_writer = self.get_summary_writer() + if summary_writer is not None: + log_func_factory = lambda name, num: create_log_func(summary_writer, name, num) + else: + log_func_factory = None + last_time = time.time() + + while training_state.examples_seen_so_far < target_checkpoint_examples: + # Set the learning rate + learning_rate_by_module_name = self.training_protocol.get_learning_rate(training_state.examples_seen_so_far) + for module_name in self.module_factories.keys(): + if module_name not in learning_rate_by_module_name or module_name not in training_state.optimizers: + continue + lr = learning_rate_by_module_name[module_name] + set_learning_rate(training_state.optimizers[module_name], lr) + if summary_writer is not None: + summary_writer.add_scalar( + module_name + "_learning_rate", lr, training_state.examples_seen_so_far) + + # One training iteration + training_batch = self.get_next_training_batch(device) + self.training_protocol.run_training_iteration( + training_batch, + training_state.examples_seen_so_far, + training_state.modules, + training_state.accumulated_modules, + training_state.optimizers, + self.losses, + log_func_factory, + device) + + # Accumulate model data + for module_name in self.accumulators: + new_module = training_state.modules[module_name] + buffer_module = training_state.accumulated_modules[module_name] + self.accumulators[module_name].accumulate( + new_module, buffer_module, examples_seen_so_far=training_state.examples_seen_so_far) + + # Advance the number of examples seen so far + next_num_examples = self.get_next_num_examples(training_state.examples_seen_so_far) + training_state.examples_seen_so_far += self.training_protocol.get_batch_size() + + # Validation iteration + if self.validation_protocol is not None \ + and training_state.examples_seen_so_far >= next_num_examples[KEY_VALIDATION]: + validation_batch = self.get_next_validation_batch(device) + self.validation_protocol.run_validation_iteration( + validation_batch, + training_state.examples_seen_so_far, + training_state.modules, + training_state.accumulated_modules, + self.losses, + log_func_factory, + device) + + # Save sample output + if self.sample_output_protocol is not None \ + and training_state.examples_seen_so_far >= next_num_examples[KEY_SAMPLE_OUTPUT]: + self.sample_output_protocol.save_sample_output_data( + training_state.modules, + training_state.accumulated_modules, + sample_output_data, + self.prefix + "/sample_outputs", + training_state.examples_seen_so_far, + device) + + # Save checkpoint + if training_state.examples_seen_so_far >= next_num_examples[KEY_CHECKPOINT]: + checkpoint_index = self.get_checkpoint_index_to_save(training_state.examples_seen_so_far) + training_state.save(self.get_checkpoint_prefix(checkpoint_index)) + if next_num_examples[KEY_CHECKPOINT] != next_num_examples[KEY_SNAPSHOT]: + training_state.save(self.get_snapshot_prefix()) + + # Save snapshot + if training_state.examples_seen_so_far >= next_num_examples[KEY_SNAPSHOT]: + training_state.save(self.get_snapshot_prefix()) + + now = time.time() + if now - last_time > 10: + logging.info("[Rank %d] Showed %d training examples." % (rank, training_state.examples_seen_so_far)) + last_time = now + + @staticmethod + def run(trainer_factory: Dict[int, Callable[[], 'SwarmUnitTrainer']], + backend: str = 'gloo', + device_mapper: Optional[Callable[[int, int], torch.device]] = None): + parser = argparse.ArgumentParser(description='Training script.') + parser.add_argument("--target_checkpoint_examples", type=int) + args = parser.parse_args() + + rank = int(os.environ['RANK']) + local_rank = int(os.environ['LOCAL_RANK']) + + torch.distributed.init_process_group(backend) + if rank in trainer_factory: + trainer = trainer_factory[rank]() + trainer.train(rank, local_rank, args.target_checkpoint_examples, device_mapper) diff --git a/tha4/src/tha4/shion/core/training/training_protocol.py b/tha4/src/tha4/shion/core/training/training_protocol.py new file mode 100644 index 0000000..da61864 --- /dev/null +++ b/tha4/src/tha4/shion/core/training/training_protocol.py @@ -0,0 +1,73 @@ +from abc import ABC, abstractmethod +from typing import Dict, List, Callable, Any, Optional + +import torch +from torch.nn import Module +from torch.optim.optimizer import Optimizer + +from tha4.shion.core.loss import Loss +from tha4.shion.core.optimizer_factory import OptimizerFactory + + +class TrainingProtocol(ABC): + @abstractmethod + def get_optimizer_factories(self) -> Dict[str, OptimizerFactory]: + pass + + @abstractmethod + def get_checkpoint_examples(self) -> List[int]: + pass + + @abstractmethod + def get_random_seed(self) -> int: + pass + + @abstractmethod + def get_batch_size(self) -> int: + pass + + @abstractmethod + def get_learning_rate(self, examples_seen_so_far: int) -> Dict[str, float]: + pass + + @abstractmethod + def run_training_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + optimizers: Dict[str, Optimizer], + losses: Dict[str, Loss], + create_log_func: Optional[Callable[[str, int], Callable[[str, float], None]]], + device: torch.device): + pass + + +class AbstractTrainingProtocol(TrainingProtocol, ABC): + def __init__(self, + check_point_examples: List[int], + batch_size: int, + learning_rate: Callable[[int], Dict[str, float]], + optimizer_factories: Dict[str, OptimizerFactory], + random_seed: int): + self.random_seed = random_seed + self.optimizer_factories = optimizer_factories + self.learning_rate = learning_rate + self.batch_size = batch_size + self.check_point_examples = check_point_examples + + def get_optimizer_factories(self) -> Dict[str, OptimizerFactory]: + return self.optimizer_factories + + def get_checkpoint_examples(self) -> List[int]: + return self.check_point_examples + + def get_random_seed(self) -> int: + return self.random_seed + + def get_batch_size(self) -> int: + return self.batch_size + + def get_learning_rate(self, examples_seen_so_far: int) -> Dict[str, float]: + return self.learning_rate(examples_seen_so_far) diff --git a/tha4/src/tha4/shion/core/training/util.py b/tha4/src/tha4/shion/core/training/util.py new file mode 100644 index 0000000..862c0ac --- /dev/null +++ b/tha4/src/tha4/shion/core/training/util.py @@ -0,0 +1,41 @@ +from typing import Callable + +import torch +from torch.nn import Module +from torch.optim import Optimizer + + +def optimizer_to_device(optim: Optimizer, device: torch.device): + for state in optim.state.values(): + for k, v in state.items(): + if isinstance(v, torch.Tensor): + state[k] = v.to(device) + + +def zero_module(module: Module): + parameters = dict(module.named_parameters()) + for k in parameters.keys(): + parameters[k].data.zero_() + + +def get_least_greater_multiple(x: int, m: int) -> int: + """ + :param x: a non-negative integer + :param m: a positive integer + :return: the next multiple of m that is greater than x + """ + assert x >= 0 + assert m > 0 + return (x // m + 1) * m + + +def create_log_func(summary_writer, prefix: str, examples_seen_so_far: int) -> Callable[[str, float], None]: + def log_func(tag: str, value: float): + summary_writer.add_scalar(prefix + "_" + tag, value, examples_seen_so_far) + + return log_func + + +def set_learning_rate(module, lr): + for param_group in module.param_groups: + param_group['lr'] = lr diff --git a/tha4/src/tha4/shion/core/training/validation_protocol.py b/tha4/src/tha4/shion/core/training/validation_protocol.py new file mode 100644 index 0000000..806b2da --- /dev/null +++ b/tha4/src/tha4/shion/core/training/validation_protocol.py @@ -0,0 +1,43 @@ +from abc import ABC, abstractmethod +from typing import Dict, Callable, Any + +import torch +from torch.nn import Module + +from tha4.shion.core.loss import Loss + + +class ValidationProtocol(ABC): + @abstractmethod + def get_batch_size(self) -> int: + pass + + @abstractmethod + def get_examples_per_validation_iteration(self) -> int: + pass + + @abstractmethod + def run_validation_iteration( + self, + batch: Any, + examples_seen_so_far: int, + modules: Dict[str, Module], + accumulated_modules: Dict[str, Module], + losses: Dict[str, Loss], + create_log_func: Callable[[str, int], Callable[[str, float], None]], + device: torch.device): + pass + + +class AbstractValidationProtocol(ValidationProtocol, ABC): + def __init__(self, + example_per_validation_iteration: int, + batch_size: int): + self.batch_size = batch_size + self.example_per_validation_iteration = example_per_validation_iteration + + def get_batch_size(self) -> int: + return self.batch_size + + def get_examples_per_validation_iteration(self) -> int: + return self.example_per_validation_iteration diff --git a/tha4/src/tha4/shion/nn00/__init__.py b/tha4/src/tha4/shion/nn00/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tha4/src/tha4/shion/nn00/block_args.py b/tha4/src/tha4/shion/nn00/block_args.py new file mode 100644 index 0000000..f1c39f8 --- /dev/null +++ b/tha4/src/tha4/shion/nn00/block_args.py @@ -0,0 +1,22 @@ +from typing import Optional + +from torch.nn import Module, Sequential + +from tha4.shion.core.module_factory import ModuleFactory +from tha4.shion.nn00.linear_module_args import LinearModuleArgs +from tha4.shion.nn00.nonlinearity_factories import resolve_nonlinearity_factory +from tha4.shion.nn00.normalization_layer_factories import resolve_normalization_layer_factory +from tha4.shion.nn00.normalization_layer_factory import NormalizationLayerFactory + + +class BlockArgs: + def __init__( + self, + linear_module_args: Optional[LinearModuleArgs] = None, + normalization_layer_factory: Optional[NormalizationLayerFactory] = None, + nonlinearity_factory: Optional[ModuleFactory] = None): + if linear_module_args is None: + linear_module_args = LinearModuleArgs() + self.nonlinearity_factory = resolve_nonlinearity_factory(nonlinearity_factory) + self.normalization_layer_factory = resolve_normalization_layer_factory(normalization_layer_factory) + self.linear_module_args = linear_module_args \ No newline at end of file diff --git a/tha4/src/tha4/shion/nn00/conv.py b/tha4/src/tha4/shion/nn00/conv.py new file mode 100644 index 0000000..c1718cc --- /dev/null +++ b/tha4/src/tha4/shion/nn00/conv.py @@ -0,0 +1,102 @@ +from typing import Optional, Union, Callable + +from torch.nn import Conv2d, Module, Sequential, ConvTranspose2d + +from tha4.shion.nn00.block_args import BlockArgs +from tha4.shion.nn00.linear_module_args import LinearModuleArgs, wrap_linear_module + + +def create_conv7( + in_channels: int, + out_channels: int, + bias: bool = False, + linear_module_args: Optional[LinearModuleArgs] = None) -> Module: + return wrap_linear_module( + Conv2d(in_channels, out_channels, kernel_size=7, stride=1, padding=3, bias=bias), + linear_module_args) + + +def create_conv3(in_channels: int, + out_channels: int, + bias: bool = False, + linear_module_args: Optional[LinearModuleArgs] = None) -> Module: + return wrap_linear_module( + Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=bias), + linear_module_args) + + +def create_conv1( + in_channels: int, out_channels: int, + bias: bool = False, + linear_module_args: Optional[LinearModuleArgs] = None) -> Module: + return wrap_linear_module( + Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0, bias=bias), + linear_module_args) + + +def create_conv7_block( + in_channels: int, + out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return Sequential( + create_conv7( + in_channels, + out_channels, + bias=False, + linear_module_args=block_args.linear_module_args), + block_args.normalization_layer_factory.create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) + + +def create_conv3_block( + in_channels: int, + out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return Sequential( + create_conv7( + in_channels, + out_channels, + bias=False, + linear_module_args=block_args.linear_module_args), + block_args.normalization_layer_factory.create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) + + +def create_downsample_block( + in_channels: int, + out_channels: int, + is_output_1x1: bool = False, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + if is_output_1x1: + return Sequential( + wrap_linear_module( + Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + block_args.linear_module_args), + block_args.nonlinearity_factory.create()) + else: + return Sequential( + wrap_linear_module( + Conv2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + block_args.linear_module_args), + block_args.normalization_layer_factory.create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) + + +def create_upsample_block( + in_channels: int, + out_channels: int, + block_args: Optional[BlockArgs] = None) -> Module: + if block_args is None: + block_args = BlockArgs() + return Sequential( + wrap_linear_module( + ConvTranspose2d(in_channels, out_channels, kernel_size=4, stride=2, padding=1, bias=False), + linear_module_args=block_args.linear_module_args), + block_args.normalization_layer_factory.create(out_channels, affine=True), + block_args.nonlinearity_factory.create()) diff --git a/tha4/src/tha4/shion/nn00/initialization_funcs.py b/tha4/src/tha4/shion/nn00/initialization_funcs.py new file mode 100644 index 0000000..fbf73ae --- /dev/null +++ b/tha4/src/tha4/shion/nn00/initialization_funcs.py @@ -0,0 +1,58 @@ +from typing import Callable, Optional + +import torch +from torch import zero_ +from torch.nn import Module +from torch.nn.init import kaiming_normal_, xavier_normal_, normal_ + + +class HeInitialization: + def __init__(self, a: int = 0, mode: str = 'fan_in', nonlinearity: str = 'leaky_relu'): + self.nonlinearity = nonlinearity + self.mode = mode + self.a = a + + def __call__(self, module: Module) -> Module: + with torch.no_grad(): + kaiming_normal_(module.weight, a=self.a, mode=self.mode, nonlinearity=self.nonlinearity) + return module + + +class NormalInitialization: + def __init__(self, mean: float = 0.0, std: float = 1.0): + self.std = std + self.mean = mean + + def __call__(self, module: Module) -> Module: + with torch.no_grad(): + normal_(module.weight, self.mean, self.std) + return module + + +class XavierInitialization: + def __init__(self, gain: float = 1.0): + self.gain = gain + + def __call__(self, module: Module) -> Module: + with torch.no_grad(): + xavier_normal_(module.weight, self.gain) + return module + + +class ZeroInitialization: + def __call__(self, module: Module) -> Module: + with torch.no_grad: + zero_(module.weight) + return module + + +class NoInitialization: + def __call__(self, module: Module) -> Module: + return module + + +def resolve_initialization_func(initialization: Optional[Callable[[Module], Module]]): + if initialization is None: + return NoInitialization() + else: + return initialization diff --git a/tha4/src/tha4/shion/nn00/linear_module_args.py b/tha4/src/tha4/shion/nn00/linear_module_args.py new file mode 100644 index 0000000..b8bf494 --- /dev/null +++ b/tha4/src/tha4/shion/nn00/linear_module_args.py @@ -0,0 +1,30 @@ +from typing import Optional, Callable + +from torch.nn import Module +from torch.nn.utils import spectral_norm + +from tha4.shion.nn00.initialization_funcs import resolve_initialization_func + + +class LinearModuleArgs: + def __init__( + self, + initialization_func: Optional[Callable[[Module], Module]] = None, + use_spectral_norm: bool = False): + self.use_spectral_norm = use_spectral_norm + self.initialization_func = resolve_initialization_func(initialization_func) + + def wrap_linear_module(self, module: Module) -> Module: + module = self.initialization_func(module) + if self.use_spectral_norm: + module = spectral_norm(module) + return module + + +def wrap_linear_module(module: Module, linear_module_args: Optional[LinearModuleArgs] = None): + if linear_module_args is None: + linear_module_args = LinearModuleArgs() + module = linear_module_args.initialization_func(module) + if linear_module_args.use_spectral_norm: + module = spectral_norm(module) + return module diff --git a/tha4/src/tha4/shion/nn00/nonlinearity_factories.py b/tha4/src/tha4/shion/nn00/nonlinearity_factories.py new file mode 100644 index 0000000..f88dcc8 --- /dev/null +++ b/tha4/src/tha4/shion/nn00/nonlinearity_factories.py @@ -0,0 +1,87 @@ +from typing import Optional + +import torch +from torch import Tensor +from torch.nn import Module, ReLU, LeakyReLU, ELU, ReLU6, Hardswish, SiLU, Tanh, Sigmoid + +from tha4.shion.core.module_factory import ModuleFactory + + +class ReLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return ReLU(self.inplace) + + +class LeakyReLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False, negative_slope: float = 1e-2): + self.negative_slope = negative_slope + self.inplace = inplace + + def create(self) -> Module: + return LeakyReLU(inplace=self.inplace, negative_slope=self.negative_slope) + + +class ELUFactory(ModuleFactory): + def __init__(self, inplace: bool = False, alpha: float = 1.0): + self.alpha = alpha + self.inplace = inplace + + def create(self) -> Module: + return ELU(inplace=self.inplace, alpha=self.alpha) + + +class ReLU6Factory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return ReLU6(inplace=self.inplace) + + +class SiLUFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return SiLU(inplace=self.inplace) + + +class HardswishFactory(ModuleFactory): + def __init__(self, inplace: bool = False): + self.inplace = inplace + + def create(self) -> Module: + return Hardswish(inplace=self.inplace) + + +class TanhFactory(ModuleFactory): + def create(self) -> Module: + return Tanh() + + +class SigmoidFactory(ModuleFactory): + def create(self) -> Module: + return Sigmoid() + + +class Swish(Module): + def __init__(self): + super().__init__() + + def forward(self, x: Tensor): + return x * torch.sigmoid(x) + + +class SwishFactory(ModuleFactory): + def create(self) -> Module: + return Swish() + + +def resolve_nonlinearity_factory(nonlinearity_factory: Optional[ModuleFactory]) -> ModuleFactory: + if nonlinearity_factory is None: + return ReLUFactory(inplace=False) + else: + return nonlinearity_factory \ No newline at end of file diff --git a/tha4/src/tha4/shion/nn00/normalization_layer_factories.py b/tha4/src/tha4/shion/nn00/normalization_layer_factories.py new file mode 100644 index 0000000..61ead7c --- /dev/null +++ b/tha4/src/tha4/shion/nn00/normalization_layer_factories.py @@ -0,0 +1,111 @@ +from typing import Optional + +import torch +from torch.nn import Module, Parameter, BatchNorm2d, InstanceNorm2d, GroupNorm +from torch.nn.functional import layer_norm +from torch.nn.init import normal_, constant_ + +from tha4.shion.nn00.normalization_layer_factory import NormalizationLayerFactory +from tha4.shion.nn00.pass_through import PassThrough + + +class Bias2d(Module): + def __init__(self, num_features: int): + super().__init__() + self.num_features = num_features + self.bias = Parameter(torch.zeros(1, num_features, 1, 1)) + + def forward(self, x): + return x + self.bias + + +class NoNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + if affine: + return Bias2d(num_features) + else: + return PassThrough() + + +class BatchNorm2dFactory(NormalizationLayerFactory): + def __init__(self, + weight_mean: Optional[float] = None, + weight_std: Optional[float] = None, + bias: Optional[float] = None): + super().__init__() + self.bias = bias + self.weight_std = weight_std + self.weight_mean = weight_mean + + def get_weight_mean(self): + if self.weight_mean is None: + return 1.0 + else: + return self.weight_mean + + def get_weight_std(self): + if self.weight_std is None: + return 0.02 + else: + return self.weight_std + + def create(self, num_features: int, affine: bool = True) -> Module: + module = BatchNorm2d(num_features=num_features, affine=affine) + if affine: + if self.weight_mean is not None or self.weight_std is not None: + normal_(module.weight, self.get_weight_mean(), self.get_weight_std()) + if self.bias is not None: + constant_(module.bias, self.bias) + return module + + +class InstanceNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return InstanceNorm2d(num_features=num_features, affine=affine) + + +class LayerNorm2d(Module): + def __init__(self, channels: int, affine: bool = True): + super(LayerNorm2d, self).__init__() + self.channels = channels + self.affine = affine + + if self.affine: + self.weight = Parameter(torch.ones(1, channels, 1, 1)) + self.bias = Parameter(torch.zeros(1, channels, 1, 1)) + + def forward(self, x): + shape = x.size()[1:] + y = layer_norm(x, shape) * self.weight + self.bias + return y + + +class LayerNorm2dFactory(NormalizationLayerFactory): + def __init__(self): + super().__init__() + + def create(self, num_features: int, affine: bool = True) -> Module: + return LayerNorm2d(channels=num_features, affine=affine) + + +class GroupNormFactory(NormalizationLayerFactory): + def __init__(self, num_groups: int, eps=1e-6): + super().__init__() + self.eps = eps + self.num_groups = num_groups + + def create(self, num_features: int, affine: bool = True) -> Module: + return GroupNorm(num_channels=num_features, num_groups=self.num_groups, eps=self.eps, affine=affine) + + +def resolve_normalization_layer_factory(factory: Optional['NormalizationLayerFactory']) -> 'NormalizationLayerFactory': + if factory is None: + return InstanceNorm2dFactory() + else: + return factory diff --git a/tha4/src/tha4/shion/nn00/normalization_layer_factory.py b/tha4/src/tha4/shion/nn00/normalization_layer_factory.py new file mode 100644 index 0000000..5ff61cd --- /dev/null +++ b/tha4/src/tha4/shion/nn00/normalization_layer_factory.py @@ -0,0 +1,12 @@ +from abc import ABC, abstractmethod + +from torch.nn import Module + + +class NormalizationLayerFactory(ABC): + def __init__(self): + super().__init__() + + @abstractmethod + def create(self, num_features: int, affine: bool = True) -> Module: + pass diff --git a/tha4/src/tha4/shion/nn00/pass_through.py b/tha4/src/tha4/shion/nn00/pass_through.py new file mode 100644 index 0000000..c64d678 --- /dev/null +++ b/tha4/src/tha4/shion/nn00/pass_through.py @@ -0,0 +1,9 @@ +from torch.nn import Module + + +class PassThrough(Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x \ No newline at end of file diff --git a/tha4/src/tha4/shion/nn00/resnet_block.py b/tha4/src/tha4/shion/nn00/resnet_block.py new file mode 100644 index 0000000..7989100 --- /dev/null +++ b/tha4/src/tha4/shion/nn00/resnet_block.py @@ -0,0 +1,55 @@ +from typing import Optional + +import torch +from torch.nn import Module, Sequential, Parameter + +from tha4.shion.nn00.block_args import BlockArgs +from tha4.shion.nn00.conv import create_conv1, create_conv3 + + +class ResnetBlock(Module): + def __init__(self, + num_channels: int, + is1x1: bool = False, + use_scale_parameter: bool = False, + block_args: Optional[BlockArgs] = None): + super().__init__() + if block_args is None: + block_args = BlockArgs() + self.use_scale_parameter = use_scale_parameter + if self.use_scale_parameter: + self.scale = Parameter(torch.zeros(1)) + if is1x1: + self.resnet_path = Sequential( + create_conv1( + num_channels, + num_channels, + bias=True, + linear_module_args=block_args.linear_module_args), + block_args.nonlinearity_factory.create(), + create_conv1( + num_channels, + num_channels, + bias=True, + linear_module_args=block_args.linear_module_args)) + else: + self.resnet_path = Sequential( + create_conv3( + num_channels, + num_channels, + bias=False, + linear_module_args=block_args.linear_module_args), + block_args.normalization_layer_factory.create(num_channels, affine=True), + block_args.nonlinearity_factory.create(), + create_conv3( + num_channels, + num_channels, + bias=False, + linear_module_args=block_args.linear_module_args), + block_args.normalization_layer_factory.create(num_channels, affine=True)) + + def forward(self, x): + if self.use_scale_parameter: + return x + self.scale * self.resnet_path(x) + else: + return x + self.resnet_path(x) diff --git a/tha4/test_tha4_student.py b/tha4/test_tha4_student.py new file mode 100644 index 0000000..f353120 --- /dev/null +++ b/tha4/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()) diff --git a/tha4/tha4_adapter.py b/tha4/tha4_adapter.py new file mode 100644 index 0000000..73c6432 --- /dev/null +++ b/tha4/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/tha4_convert_onnx_to_trt.py b/tha4/tha4_convert_onnx_to_trt.py new file mode 100644 index 0000000..4f538a2 --- /dev/null +++ b/tha4/tha4_convert_onnx_to_trt.py @@ -0,0 +1,76 @@ +""" +将 ONNX 模型转换为 fp32 和 fp16 精度的 TensorRT 模型 +并保存到对应的 fp32 和 fp16 目录 +""" +import os +import sys +from pathlib import Path + +# 添加 ezvtuber-rt 到路径 +sys.path.insert(0, str(Path(__file__).parent / 'ezvtuber-rt')) + +from ezvtb_rt.trt_utils import build_engine, save_engine + +def convert_onnx_to_trt(): + # 基础目录 + base_dir = Path('data/models/tha4') + onnx_dir = base_dir / 'fp32' + + # 检查 ONNX 文件目录是否存在 + if not onnx_dir.exists(): + print(f"错误: ONNX 模型目录不存在: {onnx_dir}") + return + + # 指定要转换的文件名列表 + target_names = ['body_morpher.onnx', 'combiner.onnx', 'decomposer.onnx', + 'morpher.onnx', 'upscaler.onnx'] + + # 获取指定的 ONNX 文件 + onnx_files = [onnx_dir / name for name in target_names if (onnx_dir / name).exists()] + + if not onnx_files: + print(f"在 {onnx_dir} 中未找到指定的 ONNX 模型文件") + print(f"需要的文件: {target_names}") + return + + print(f"找到 {len(onnx_files)} 个 ONNX 模型文件") + print("开始转换为 fp32 和 fp16 精度的 TensorRT 模型...\n") + + # 定义要转换的精度类型 + precisions = ['fp32', 'fp16'] + + for precision in precisions: + print(f"\n{'='*60}") + print(f"开始转换 {precision.upper()} 精度模型") + print(f"{'='*60}\n") + + # 设置输出目录 + output_dir = base_dir / precision + output_dir.mkdir(parents=True, exist_ok=True) + + for onnx_file in onnx_files: + trt_filename = onnx_file.stem + '.trt' + trt_path = output_dir / trt_filename + + # 检查是否已存在转换好的 TRT 文件 + if trt_path.exists(): + print(f"⊙ 跳过 ({precision}): {onnx_file.name} (TRT文件已存在)") + continue + + print(f"正在转换 ({precision}): {onnx_file.name}") + + try: + # 构建并保存 TensorRT 引擎 + engine = build_engine(str(onnx_file), precision) + save_engine(engine, str(trt_path)) + print(f"✓ 成功转换 ({precision}): {onnx_file.name} -> {trt_filename}\n") + except Exception as e: + print(f"✗ 转换失败 ({precision}): {onnx_file.name}") + print(f" 错误信息: {str(e)}\n") + + print("\n" + "="*60) + print("全部转换完成！") + print("="*60) + +if __name__ == '__main__': + convert_onnx_to_trt() diff --git a/tha4/tha4_core.py b/tha4/tha4_core.py new file mode 100644 index 0000000..6d2d33d --- /dev/null +++ b/tha4/tha4_core.py @@ -0,0 +1,504 @@ +""" +THA4 Core - PyTorch THA4 model with optional TensorRT RIFE/SR acceleration +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 + +# Add ezvtuber-rt path for TensorRT modules +dir_path = os.path.dirname(os.path.realpath(__file__)) +ezvtb_path = os.path.join(dir_path, 'ezvtuber-rt') +if os.path.exists(ezvtb_path) and ezvtb_path not in sys.path: + sys.path.append(ezvtb_path) + + +class THA4Core: + """ + THA4 Core class - Pure PyTorch implementation (No TensorRT) + + This class only uses PyTorch for THA4 model inference. + It does NOT support real frame interpolation or super resolution. + If interpolation_scale > 1, frames will be simply duplicated. + + For real RIFE interpolation and SR with TensorRT, use THA4CoreWithAccel. + + Uses CharacterModel to load YAML + PNG + PT files properly. + The PNG image from YAML is used as the source image. + + Expected interface: + - setImage(img: np.ndarray) - NOT USED, image comes from YAML + - inference(pose: np.ndarray) -> List[np.ndarray] - run inference + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1): + """ + Initialize THA4 Core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + interpolation_scale: number of output frames to generate + (THA4 doesn't support real interpolation, + will duplicate frames) + """ + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + self.use_eyebrow = use_eyebrow + self.interpolation_scale = interpolation_scale + + # Create THA4 wrapper (loads YAML + PNG + PT) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + + # Use character image from YAML (already loaded by wrapper) + self.input_image_tensor = self.wrapper.character_image + + # For compatibility (no caching in THA4 core) + self.cacher = None + + print(f"THA4 Core initialized on {self.device}") + print("Using character image from model (ignoring setImage calls)") + if interpolation_scale > 1: + print(f"Note: Pure PyTorch mode doesn't support real " + f"interpolation, will duplicate frames {interpolation_scale}x") + print(" Use THA4CoreWithAccel with TensorRT for " + "real RIFE interpolation") + + def setImage(self, img: np.ndarray): + """ + Set input character image - IGNORED for THA4 + + THA4 uses the character image from the YAML/PNG files loaded + during initialization. This method exists only for interface + compatibility with CoreTRT/CoreORT. + + Args: + img: numpy array in BGRA format (ignored) + """ + # THA4 uses character image from YAML, ignore external image + print("THA4: setImage() called but ignored (using image from YAML)") + + def inference(self, pose: np.ndarray) -> list: + """ + Run inference + + Args: + pose: numpy array of shape (1, 45) + [0:12] - eyebrow parameters + [12:39] - face parameters + [39:45] - head pose parameters + + Returns: + List of output images as numpy arrays [H, W, 4] + in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set. Call setImage() first.") + + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] # [1, 12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] # [1, 27] + head_pose = pose[:, 39:45] # [1, 6] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + # Call wrapper (it handles the pose conversion internally) + # We need to create dummy compressed versions for + # interface compatibility + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # THA4 outputs in linear color space, convert to sRGB + # Output is [batch, 4, H, W] in range [-1, 1] + output_image = output_tensor[0].float() # [4, 512, 512] + + # Clip to [-1, 1] and convert to [0, 1] + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + + # Convert linear RGB to sRGB (keep alpha as-is) + output_image = convert_linear_to_srgb(output_image) + + # Convert to [H, W, C] numpy + output_np = output_image.permute(1, 2, 0).cpu().numpy() + + # Convert to uint8 + output_np = (output_np * 255.0).astype(np.uint8) + + # Convert RGBA to BGRA for output + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Return as list with proper number of frames + # Pure PyTorch mode: duplicate frames (no real interpolation) + # Use THA4CoreWithAccel with TensorRT for real RIFE interpolation + return [output_bgra] * self.interpolation_scale + + +class THA4CoreWithAccel: + """ + THA4 Core with TensorRT acceleration for RIFE and SR + + Architecture: + - THA4 main model: PyTorch (PT files) - always PyTorch + - RIFE interpolation: TensorRT (optional) - REAL frame interpolation + - SR super-resolution: TensorRT (optional) - REAL upscaling + - Cacher: RAM cache (optional) + + This class provides REAL interpolation and super resolution using + TensorRT, unlike THA4Core which only duplicates frames. + + Follows the same design as CoreTRT for THA3 models. + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + use_tensorrt=False, use_interpolation=False, + interpolation_scale=2, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Initialize THA4 Core with optional TensorRT acceleration + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + use_tensorrt: enable TensorRT for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_scale: interpolation scale (2/3/4) + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x (1/2/3) + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + """ + self.device_id = device_id + self.use_eyebrow = use_eyebrow + self.use_tensorrt = use_tensorrt + self.use_interpolation = use_interpolation + self.use_sr = use_sr + + # Initialize PyTorch device + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + + # Initialize THA4 wrapper (PyTorch) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + self.input_image_tensor = self.wrapper.character_image + + # Initialize TensorRT components if requested + self.rife = None + self.sr = None + self.cacher = None + self.scale = 1 + self.instream = None + self.output_memory = None + + if use_tensorrt and (use_interpolation or use_sr): + # Initialize CUDA context for TensorRT + self._init_tensorrt(device_id) + + # Build model paths + rife_model_path = '' + if use_interpolation: + rife_model_path = os.path.join( + '.', 'data', 'models', 'rife_512', + f'x{interpolation_scale}', + 'fp16' if interpolation_half else 'fp32') + + sr_model_path = '' + if use_sr: + if sr_x4: + base_path = os.path.join('.', 'data', 'models', + 'Real-ESRGAN') + sr_model_path = os.path.join( + base_path, + 'exported_256_fp16' if sr_half else 'exported_256') + else: # x2 + base_path = os.path.join('.', 'data', 'models', + 'waifu2x_upconv', + 'fp16' if sr_half else 'fp32', + 'upconv_7', 'art') + sr_model_path = os.path.join( + base_path, f'noise{sr_noise}_scale2x') + + # Initialize TensorRT modules + self._init_trt_modules(rife_model_path, sr_model_path) + + # Initialize cacher + if cacher_ram_size > 0.0: + from ezvtb_rt.cache import Cacher + self.cacher = Cacher(cacher_ram_size, cacher_quality) + + print("THA4 Core initialized:") + print(f" Device: {self.device}") + print(f" TensorRT: {use_tensorrt}") + interp_info = f"scale: {interpolation_scale}" \ + if use_interpolation else "N/A" + print(f" Interpolation: {use_interpolation} ({interp_info})") + if use_interpolation: + print(" (Real RIFE interpolation via TensorRT)") + print(f" Super Resolution: {use_sr}") + + def _init_tensorrt(self, device_id): + """Initialize TensorRT and CUDA context""" + try: + from ezvtb_rt.trt_utils import cudaSetDevice + cudaSetDevice(device_id) + os.environ['CUDA_DEVICE'] = str(device_id) + import pycuda.autoinit + device_name = pycuda.autoinit.device.name() + print(f'TensorRT initialized on device: {device_name}') + except Exception as e: + print(f'Warning: Failed to initialize TensorRT: {e}') + self.use_tensorrt = False + + def _init_trt_modules(self, rife_model_path, sr_model_path): + """Initialize RIFE and SR TensorRT modules""" + try: + import pycuda.driver as cuda + + # Create CUDA stream + self.instream = cuda.Stream() + + # Create output memory compatible with RIFE input + # [1, 512, 512, 4] + output_shape = (1, 512, 512, 4) + self.output_memory = self._create_host_device_mem( + output_shape, np.uint8) + + # Initialize RIFE + if self.use_interpolation and rife_model_path: + from ezvtb_rt.rife import RIFE + + self.rife = RIFE(rife_model_path, self.instream, + self.output_memory) + self.scale = self.rife.scale + print(f"RIFE TensorRT module loaded: {rife_model_path}") + + # Initialize SR + if self.use_sr and sr_model_path: + from ezvtb_rt.sr import SR + + if self.rife is not None: + instream = self.rife.instream + mems = [self.rife.memories['framegen_'+str(i)] + for i in range(self.rife.scale)] + else: + instream = self.instream + mems = [self.output_memory] + + self.sr = SR(sr_model_path, instream, mems) + print(f"SR TensorRT module loaded: {sr_model_path}") + + except Exception as e: + print(f'Warning: Failed to initialize TensorRT modules: {e}') + self.rife = None + self.sr = None + + def _create_host_device_mem(self, shape, dtype): + """Create HostDeviceMem for TensorRT modules""" + import pycuda.driver as cuda + from ezvtb_rt.trt_utils import HostDeviceMem + + host_mem = cuda.pagelocked_empty(shape, dtype) + device_mem = cuda.mem_alloc(host_mem.nbytes) + return HostDeviceMem(host_mem, device_mem) + + def setImage(self, img: np.ndarray): + """Set input character image - IGNORED for THA4""" + # THA4 uses character image from YAML, ignore external image + pass + + def inference(self, pose: np.ndarray) -> list: + """ + Run full inference pipeline: THA4 -> RIFE (TRT) -> SR (TRT) + + Args: + pose: numpy array of shape (1, 45) + + Returns: + List of output images [H, W, 4] in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set.") + + pose = pose.astype(np.float32) + + # Cache management + need_cache_write = 0 + res_carrier = None + + # Step 1: THA4 inference (PyTorch) with optional caching + if self.cacher is None: + output_bgra = self._run_tha4_inference(pose) + res_carrier = [output_bgra] + else: + hs = hash(str(pose)) + cached = self.cacher.read(hs) + + if cached is not None: # Cache hit + output_bgra = cached + res_carrier = [cached] + + # If using TensorRT, upload to GPU + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + else: # Cache miss + output_bgra = self._run_tha4_inference(pose) + need_cache_write = hs + res_carrier = [output_bgra] + + # Step 2: RIFE interpolation (TensorRT) + if self.rife is not None: + self.rife.inference() + res_carrier = self.rife + + # Step 3: SR super resolution (TensorRT) + if self.sr is not None: + self.sr.inference() + res_carrier = self.sr + + # Step 4: Fetch results + if need_cache_write != 0: + self.cacher.write(need_cache_write, output_bgra) + + # Get final results + if type(res_carrier) is not list: + res_carrier = res_carrier.fetchRes() + + return res_carrier + + def _run_tha4_inference(self, pose: np.ndarray) -> np.ndarray: + """Run THA4 PyTorch inference and return BGRA numpy array""" + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] + head_pose = pose[:, 39:45] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + # Run THA4 inference + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # Post-process output + output_image = output_tensor[0].float() + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + output_image = convert_linear_to_srgb(output_image) + output_np = output_image.permute(1, 2, 0).cpu().numpy() + output_np = (output_np * 255.0).astype(np.uint8) + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Upload to GPU memory if TensorRT is active + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + + return output_bgra + + +def get_tha4_core(device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1, use_tensorrt=False, + use_interpolation=False, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Factory function to create THA4 core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow control + yaml_path: path to character_model.yaml + interpolation_scale: number of frames for interpolation + use_tensorrt: enable TensorRT acceleration for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + + Returns: + THA4Core or THA4CoreWithAccel instance + """ + # Use accelerated version if TensorRT is requested + if use_tensorrt and (use_interpolation or use_sr): + return THA4CoreWithAccel( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + use_tensorrt=use_tensorrt, + use_interpolation=use_interpolation, + interpolation_scale=interpolation_scale, + interpolation_half=interpolation_half, + use_sr=use_sr, + sr_x4=sr_x4, + sr_half=sr_half, + sr_noise=sr_noise, + cacher_quality=cacher_quality, + cacher_ram_size=cacher_ram_size + ) + else: + # Use simple PyTorch-only version + return THA4Core( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + interpolation_scale=interpolation_scale if use_interpolation + else 1 + ) diff --git a/tha4/tha4_core_old.py b/tha4/tha4_core_old.py new file mode 100644 index 0000000..5add4d5 --- /dev/null +++ b/tha4/tha4_core_old.py @@ -0,0 +1,498 @@ +""" +THA4 Core - ONNX-based THA4 model with optional TensorRT/RIFE/SR acceleration +Provides interface compatible with CoreTRT/CoreORT for THA4 models +Uses teacher models (5 models) from onnx_model_tha4 folder +""" +import numpy as np +import sys +import os + +# Add ezvtuber-rt path for TensorRT/RIFE/SR modules +dir_path = os.path.dirname(os.path.realpath(__file__)) +ezvtb_path = os.path.join(dir_path, 'ezvtuber-rt') +if os.path.exists(ezvtb_path) and ezvtb_path not in sys.path: + sys.path.append(ezvtb_path) + + +class THA4CoreTRT: + """ + THA4 Core class - Pure PyTorch implementation (No TensorRT) + + This class only uses PyTorch for THA4 model inference. + It does NOT support real frame interpolation or super resolution. + If interpolation_scale > 1, frames will be simply duplicated. + + For real RIFE interpolation and SR with TensorRT, use THA4CoreWithAccel. + + Uses CharacterModel to load YAML + PNG + PT files properly. + The PNG image from YAML is used as the source image. + + Expected interface: + - setImage(img: np.ndarray) - NOT USED, image comes from YAML + - inference(pose: np.ndarray) -> List[np.ndarray] - run inference + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1): + """ + Initialize THA4 Core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + interpolation_scale: number of output frames to generate + (THA4 doesn't support real interpolation, + will duplicate frames) + """ + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + self.use_eyebrow = use_eyebrow + self.interpolation_scale = interpolation_scale + + # Create THA4 wrapper (loads YAML + PNG + PT) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + + # Use character image from YAML (already loaded by wrapper) + self.input_image_tensor = self.wrapper.character_image + + # For compatibility (no caching in THA4 core) + self.cacher = None + + print(f"THA4 Core initialized on {self.device}") + print("Using character image from model (ignoring setImage calls)") + if interpolation_scale > 1: + print(f"Note: Pure PyTorch mode doesn't support real " + f"interpolation, will duplicate frames {interpolation_scale}x") + print(" Use THA4CoreWithAccel with TensorRT for " + "real RIFE interpolation") + + def setImage(self, img: np.ndarray): + """ + Set input character image - IGNORED for THA4 + + THA4 uses the character image from the YAML/PNG files loaded + during initialization. This method exists only for interface + compatibility with CoreTRT/CoreORT. + + Args: + img: numpy array in BGRA format (ignored) + """ + # THA4 uses character image from YAML, ignore external image + print("THA4: setImage() called but ignored (using image from YAML)") + + def inference(self, pose: np.ndarray) -> list: + """ + Run inference + + Args: + pose: numpy array of shape (1, 45) + [0:12] - eyebrow parameters + [12:39] - face parameters + [39:45] - head pose parameters + + Returns: + List of output images as numpy arrays [H, W, 4] + in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set. Call setImage() first.") + + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] # [1, 12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] # [1, 27] + head_pose = pose[:, 39:45] # [1, 6] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + # Call wrapper (it handles the pose conversion internally) + # We need to create dummy compressed versions for + # interface compatibility + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # THA4 outputs in linear color space, convert to sRGB + # Output is [batch, 4, H, W] in range [-1, 1] + output_image = output_tensor[0].float() # [4, 512, 512] + + # Clip to [-1, 1] and convert to [0, 1] + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + + # Convert linear RGB to sRGB (keep alpha as-is) + output_image = convert_linear_to_srgb(output_image) + + # Convert to [H, W, C] numpy + output_np = output_image.permute(1, 2, 0).cpu().numpy() + + # Convert to uint8 + output_np = (output_np * 255.0).astype(np.uint8) + + # Convert RGBA to BGRA for output + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Return as list with proper number of frames + # Pure PyTorch mode: duplicate frames (no real interpolation) + # Use THA4CoreWithAccel with TensorRT for real RIFE interpolation + return [output_bgra] * self.interpolation_scale + + +class THA4CoreWithAccel: + """ + THA4 Core with TensorRT acceleration for RIFE and SR + + Architecture: + - THA4 main model: PyTorch (PT files) - always PyTorch + - RIFE interpolation: TensorRT (optional) - REAL frame interpolation + - SR super-resolution: TensorRT (optional) - REAL upscaling + - Cacher: RAM cache (optional) + + This class provides REAL interpolation and super resolution using + TensorRT, unlike THA4Core which only duplicates frames. + + Follows the same design as CoreTRT for THA3 models. + """ + + def __init__(self, device_id=0, use_eyebrow=True, yaml_path=None, + use_tensorrt=False, use_interpolation=False, + interpolation_scale=2, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Initialize THA4 Core with optional TensorRT acceleration + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow parameters + yaml_path: path to character_model.yaml + use_tensorrt: enable TensorRT for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_scale: interpolation scale (2/3/4) + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x (1/2/3) + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + """ + self.device_id = device_id + self.use_eyebrow = use_eyebrow + self.use_tensorrt = use_tensorrt + self.use_interpolation = use_interpolation + self.use_sr = use_sr + + # Initialize PyTorch device + device_str = f'cuda:{device_id}' if torch.cuda.is_available() \ + else 'cpu' + self.device = torch.device(device_str) + + # Initialize THA4 wrapper (PyTorch) + self.wrapper = THA4Wrapper(device=self.device, yaml_path=yaml_path) + self.input_image_tensor = self.wrapper.character_image + + # Initialize TensorRT components if requested + self.rife = None + self.sr = None + self.cacher = None + self.scale = 1 + self.instream = None + self.output_memory = None + + if use_tensorrt and (use_interpolation or use_sr): + # Initialize CUDA context for TensorRT + self._init_tensorrt(device_id) + + # Build model paths + rife_model_path = '' + if use_interpolation: + rife_model_path = os.path.join( + '.', 'data', 'models', 'rife_512', + f'x{interpolation_scale}', + 'fp16' if interpolation_half else 'fp32') + + sr_model_path = '' + if use_sr: + if sr_x4: + base_path = os.path.join('.', 'data', 'models', + 'Real-ESRGAN') + sr_model_path = os.path.join( + base_path, + 'exported_256_fp16' if sr_half else 'exported_256') + else: # x2 + base_path = os.path.join('.', 'data', 'models', + 'waifu2x_upconv', + 'fp16' if sr_half else 'fp32', + 'upconv_7', 'art') + sr_model_path = os.path.join( + base_path, f'noise{sr_noise}_scale2x') + + # Initialize TensorRT modules + self._init_trt_modules(rife_model_path, sr_model_path) + + # Initialize cacher + if cacher_ram_size > 0.0: + from ezvtb_rt.cache import Cacher + self.cacher = Cacher(cacher_ram_size, cacher_quality) + + print("THA4 Core initialized:") + print(f" Device: {self.device}") + print(f" TensorRT: {use_tensorrt}") + interp_info = f"scale: {interpolation_scale}" \ + if use_interpolation else "N/A" + print(f" Interpolation: {use_interpolation} ({interp_info})") + if use_interpolation: + print(" (Real RIFE interpolation via TensorRT)") + print(f" Super Resolution: {use_sr}") + + def _init_tensorrt(self, device_id): + """Initialize TensorRT and CUDA context""" + try: + from ezvtb_rt.trt_utils import cudaSetDevice + cudaSetDevice(device_id) + os.environ['CUDA_DEVICE'] = str(device_id) + import pycuda.autoinit + device_name = pycuda.autoinit.device.name() + print(f'TensorRT initialized on device: {device_name}') + except Exception as e: + print(f'Warning: Failed to initialize TensorRT: {e}') + self.use_tensorrt = False + + def _init_trt_modules(self, rife_model_path, sr_model_path): + """Initialize RIFE and SR TensorRT modules""" + try: + import pycuda.driver as cuda + + # Create CUDA stream + self.instream = cuda.Stream() + + # Create output memory compatible with RIFE input + # [1, 512, 512, 4] + output_shape = (1, 512, 512, 4) + self.output_memory = self._create_host_device_mem( + output_shape, np.uint8) + + # Initialize RIFE + if self.use_interpolation and rife_model_path: + from ezvtb_rt.rife import RIFE + + self.rife = RIFE(rife_model_path, self.instream, + self.output_memory) + self.scale = self.rife.scale + print(f"RIFE TensorRT module loaded: {rife_model_path}") + + # Initialize SR + if self.use_sr and sr_model_path: + from ezvtb_rt.sr import SR + + if self.rife is not None: + instream = self.rife.instream + mems = [self.rife.memories['framegen_'+str(i)] + for i in range(self.rife.scale)] + else: + instream = self.instream + mems = [self.output_memory] + + self.sr = SR(sr_model_path, instream, mems) + print(f"SR TensorRT module loaded: {sr_model_path}") + + except Exception as e: + print(f'Warning: Failed to initialize TensorRT modules: {e}') + self.rife = None + self.sr = None + + def _create_host_device_mem(self, shape, dtype): + """Create HostDeviceMem for TensorRT modules""" + import pycuda.driver as cuda + from ezvtb_rt.trt_utils import HostDeviceMem + + host_mem = cuda.pagelocked_empty(shape, dtype) + device_mem = cuda.mem_alloc(host_mem.nbytes) + return HostDeviceMem(host_mem, device_mem) + + def setImage(self, img: np.ndarray): + """Set input character image - IGNORED for THA4""" + # THA4 uses character image from YAML, ignore external image + pass + + def inference(self, pose: np.ndarray) -> list: + """ + Run full inference pipeline: THA4 -> RIFE (TRT) -> SR (TRT) + + Args: + pose: numpy array of shape (1, 45) + + Returns: + List of output images [H, W, 4] in BGRA format, uint8 + """ + if self.input_image_tensor is None: + raise RuntimeError("Image not set.") + + pose = pose.astype(np.float32) + + # Cache management + need_cache_write = 0 + res_carrier = None + + # Step 1: THA4 inference (PyTorch) with optional caching + if self.cacher is None: + output_bgra = self._run_tha4_inference(pose) + res_carrier = [output_bgra] + else: + hs = hash(str(pose)) + cached = self.cacher.read(hs) + + if cached is not None: # Cache hit + output_bgra = cached + res_carrier = [cached] + + # If using TensorRT, upload to GPU + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + else: # Cache miss + output_bgra = self._run_tha4_inference(pose) + need_cache_write = hs + res_carrier = [output_bgra] + + # Step 2: RIFE interpolation (TensorRT) + if self.rife is not None: + self.rife.inference() + res_carrier = self.rife + + # Step 3: SR super resolution (TensorRT) + if self.sr is not None: + self.sr.inference() + res_carrier = self.sr + + # Step 4: Fetch results + if need_cache_write != 0: + self.cacher.write(need_cache_write, output_bgra) + + # Get final results + if type(res_carrier) is not list: + res_carrier = res_carrier.fetchRes() + + return res_carrier + + def _run_tha4_inference(self, pose: np.ndarray) -> np.ndarray: + """Run THA4 PyTorch inference and return BGRA numpy array""" + # Extract pose components + if self.use_eyebrow: + eyebrow_pose = pose[:, :12] + else: + eyebrow_pose = np.zeros((1, 12), dtype=np.float32) + + face_pose = pose[:, 12:39] + head_pose = pose[:, 39:45] + + # Convert to torch tensors + eyebrow_tensor = torch.from_numpy(eyebrow_pose).to(self.device) + face_tensor = torch.from_numpy(face_pose).to(self.device) + head_tensor = torch.from_numpy(head_pose).to(self.device) + + eyebrow_c = eyebrow_pose[0].tolist() + face_c = face_pose[0].tolist() + + # Run THA4 inference + with torch.no_grad(): + output_tensor = self.wrapper.forward( + self.input_image_tensor, + face_tensor, + head_tensor, + eyebrow_tensor, + face_c, + eyebrow_c, + ratio=None + ) + + # Post-process output + output_image = output_tensor[0].float() + output_image = torch.clamp((output_image + 1.0) / 2.0, 0.0, 1.0) + output_image = convert_linear_to_srgb(output_image) + output_np = output_image.permute(1, 2, 0).cpu().numpy() + output_np = (output_np * 255.0).astype(np.uint8) + output_bgra = cv2.cvtColor(output_np, cv2.COLOR_RGBA2BGRA) + + # Upload to GPU memory if TensorRT is active + if self.use_tensorrt and \ + (self.rife is not None or self.sr is not None): + np.copyto(self.output_memory.host, + output_bgra.reshape(1, 512, 512, 4)) + self.output_memory.htod(self.instream) + + return output_bgra + + +def get_tha4_core(device_id=0, use_eyebrow=True, yaml_path=None, + interpolation_scale=1, use_tensorrt=False, + use_interpolation=False, interpolation_half=True, + use_sr=False, sr_x4=True, sr_half=True, sr_noise=1, + cacher_quality=85, cacher_ram_size=2.0): + """ + Factory function to create THA4 core + + Args: + device_id: GPU device ID + use_eyebrow: whether to use eyebrow control + yaml_path: path to character_model.yaml + interpolation_scale: number of frames for interpolation + use_tensorrt: enable TensorRT acceleration for RIFE/SR + use_interpolation: enable RIFE frame interpolation + interpolation_half: use FP16 for RIFE + use_sr: enable super resolution + sr_x4: use 4x SR (else 2x) + sr_half: use FP16 for SR + sr_noise: noise level for waifu2x + cacher_quality: cache compression quality + cacher_ram_size: RAM cache size in GB + + Returns: + THA4Core or THA4CoreWithAccel instance + """ + # Use accelerated version if TensorRT is requested + if use_tensorrt and (use_interpolation or use_sr): + return THA4CoreWithAccel( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + use_tensorrt=use_tensorrt, + use_interpolation=use_interpolation, + interpolation_scale=interpolation_scale, + interpolation_half=interpolation_half, + use_sr=use_sr, + sr_x4=sr_x4, + sr_half=sr_half, + sr_noise=sr_noise, + cacher_quality=cacher_quality, + cacher_ram_size=cacher_ram_size + ) + else: + # Use simple PyTorch-only version + return THA4Core( + device_id=device_id, + use_eyebrow=use_eyebrow, + yaml_path=yaml_path, + interpolation_scale=interpolation_scale if use_interpolation + else 1 + ) diff --git a/tha4/tha4_student_convert_onnx_to_trt.py b/tha4/tha4_student_convert_onnx_to_trt.py new file mode 100644 index 0000000..c46a21b --- /dev/null +++ b/tha4/tha4_student_convert_onnx_to_trt.py @@ -0,0 +1,285 @@ +""" +将学生模型 ONNX 转换为 fp32 和 fp16 精度的 TensorRT 模型 +学生模型（mode_14）只有 2 个网络：face_morpher 和 body_morpher +""" +import os +import sys +from pathlib import Path + +# 添加 ezvtuber-rt 到路径 +sys.path.insert(0, str(Path(__file__).parent / 'ezvtuber-rt')) + +from ezvtb_rt.trt_utils import save_engine +import tensorrt as trt + + +# TensorRT Logger +TRT_LOGGER = trt.Logger(trt.Logger.INFO) + + +def build_engine_with_dynamic_shape(onnx_file_path: str, precision: str): + """ + 构建支持动态输入形状的 TensorRT 引擎 + 专门用于 body_morpher，它有动态的图像输入 + """ + builder = trt.Builder(TRT_LOGGER) + network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + config = builder.create_builder_config() + parser = trt.OnnxParser(network, TRT_LOGGER) + + # Parse model file + TRT_LOGGER.log(TRT_LOGGER.INFO, f'Loading ONNX file from path {onnx_file_path}...') + with open(onnx_file_path, 'rb') as model: + TRT_LOGGER.log(TRT_LOGGER.INFO, 'Beginning ONNX file parsing') + parse_res = parser.parse(model.read()) + if not parse_res: + for error in range(parser.num_errors): + TRT_LOGGER.log(TRT_LOGGER.ERROR, parser.get_error(error)) + raise ValueError('Failed to parse the ONNX file.') + + TRT_LOGGER.log(TRT_LOGGER.INFO, 'Completed parsing of ONNX file') + TRT_LOGGER.log(TRT_LOGGER.INFO, f'Input number: {network.num_inputs}') + TRT_LOGGER.log(TRT_LOGGER.INFO, f'Output number: {network.num_outputs}') + + # 检查是否有动态输入 + has_dynamic_shape = False + dynamic_inputs = [] + + for i in range(network.num_inputs): + input_tensor = network.get_input(i) + input_shape = input_tensor.shape + input_name = input_tensor.name + TRT_LOGGER.log(TRT_LOGGER.INFO, f'Input {i} ({input_name}): shape={input_shape}') + # 检查是否有 -1（动态维度） + if any(dim == -1 for dim in input_shape): + has_dynamic_shape = True + dynamic_inputs.append((input_name, input_shape)) + + # 设置优化配置文件（如果有动态输入） + if has_dynamic_shape: + TRT_LOGGER.log(TRT_LOGGER.INFO, 'Detected dynamic shapes, adding optimization profile...') + profile = builder.create_optimization_profile() + + for input_name, input_shape in dynamic_inputs: + input_shape_list = list(input_shape) + + # 为 input_image 设置动态形状配置 + if 'image' in input_name.lower(): + # 图像输入格式为 [H, W, C] + min_shape = [512, 512, 4] + opt_shape = [512, 512, 4] + max_shape = [512, 512, 4] + + TRT_LOGGER.log(TRT_LOGGER.INFO, + f'Setting shape for {input_name}: {min_shape}') + profile.set_shape(input_name, min_shape, opt_shape, max_shape) + + else: + # 对于 pose 和其他输入，固定为 batch_size=1 + # 这避免了条件层（If层）中的秩不匹配问题 + min_shape = [1 if dim == -1 else dim for dim in input_shape_list] + opt_shape = [1 if dim == -1 else dim for dim in input_shape_list] + max_shape = [1 if dim == -1 else dim for dim in input_shape_list] + + TRT_LOGGER.log(TRT_LOGGER.INFO, + f'Setting shape for {input_name}: shape={min_shape}') + profile.set_shape(input_name, min_shape, opt_shape, max_shape) + + config.add_optimization_profile(profile) + + # 设置内存限制 + def GiB(val): + return val * 1 << 30 + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, GiB(2)) # 2G for complex models + + # 设置精度 + if precision == 'fp32': + config.set_flag(trt.BuilderFlag.TF32) + elif precision == 'fp16': + config.set_flag(trt.BuilderFlag.FP16) + else: + raise ValueError('precision must be one of fp32 or fp16') + + # Build engine + TRT_LOGGER.log(TRT_LOGGER.INFO, f'Building an engine from file {onnx_file_path}; this may take a while...') + serialized_engine = builder.build_serialized_network(network, config) + TRT_LOGGER.log(TRT_LOGGER.INFO, 'Completed creating Engine') + + return serialized_engine + + +def convert_student_onnx_to_trt(model_name: str, precision: str = 'fp32'): + """ + 转换学生模型的 ONNX 文件为 TensorRT 引擎 + + Args: + model_name: 模型名称（例如：kanori_tha4） + precision: 精度类型，'fp32' 或 'fp16' + """ + # 基础目录 + base_dir = Path('data/models/custom_tha4_models') / model_name + onnx_dir = base_dir + + # 检查 ONNX 文件目录是否存在 + if not onnx_dir.exists(): + print(f"错误: ONNX 模型目录不存在: {onnx_dir}") + print(f"请先运行: python onnx_export_tha4_student.py {model_name} --{precision}") + return False + + # 学生模型只需要这两个文件 + target_names = ['face_morpher.onnx', 'body_morpher.onnx'] + + # 获取 ONNX 文件 + onnx_files = [onnx_dir / name for name in target_names if (onnx_dir / name).exists()] + + if len(onnx_files) != 2: + print(f"错误: 在 {onnx_dir} 中未找到所有必需的 ONNX 模型文件") + print(f"需要的文件: {target_names}") + print(f"找到的文件: {[f.name for f in onnx_files]}") + return False + + print(f"\n{'='*70}") + print(f"学生模型 TensorRT 转换工具") + print(f"{'='*70}") + print(f"模型名称: {model_name}") + print(f"精度类型: {precision.upper()}") + print(f"找到 {len(onnx_files)} 个 ONNX 模型文件") + print(f"{'='*70}\n") + + # 设置输出目录（与 ONNX 文件同目录） + output_dir = onnx_dir + success_count = 0 + + for onnx_file in onnx_files: + trt_filename = onnx_file.stem + '.trt' + trt_path = output_dir / trt_filename + + # 检查是否已存在转换好的 TRT 文件 + if trt_path.exists(): + user_input = input(f"⊙ TRT文件已存在: {trt_filename}\n 是否覆盖? (y/N): ") + if user_input.lower() != 'y': + print(f" 跳过: {onnx_file.name}\n") + success_count += 1 + continue + else: + print(f" 将覆盖现有文件\n") + + print(f"{'─'*70}") + print(f"正在转换: {onnx_file.name}") + print(f"{'─'*70}") + + try: + # 构建 TensorRT 引擎 + print(f" → 加载 ONNX 文件...") + + # 使用支持动态输入的构建函数（两个模型都可能有动态输入） + print(f" → 构建 TensorRT 引擎（支持动态输入）...") + engine = build_engine_with_dynamic_shape(str(onnx_file), precision) + + if engine is None: + raise RuntimeError("引擎构建返回 None") + + print(f" → 保存 TensorRT 引擎...") + save_engine(engine, str(trt_path)) + + # 验证文件是否成功保存 + if trt_path.exists(): + file_size = trt_path.stat().st_size / (1024 * 1024) # MB + print(f" ✓ 成功转换: {onnx_file.name} -> {trt_filename}") + print(f" 文件大小: {file_size:.2f} MB") + success_count += 1 + else: + raise RuntimeError("TRT 文件保存失败") + + except Exception as e: + print(f" ✗ 转换失败: {onnx_file.name}") + print(f" 错误信息: {str(e)}") + + print() + + # 输出总结 + print(f"{'='*70}") + if success_count == len(onnx_files): + print(f"✓ 全部转换完成！({success_count}/{len(onnx_files)})") + print(f"\nTensorRT 模型已保存到:") + print(f" {output_dir}") + print(f"\n生成的文件:") + for name in target_names: + trt_name = name.replace('.onnx', '.trt') + trt_file = output_dir / trt_name + if trt_file.exists(): + file_size = trt_file.stat().st_size / (1024 * 1024) + print(f" - {trt_name} ({file_size:.2f} MB)") + return True + else: + print(f"⚠ 部分转换失败 ({success_count}/{len(onnx_files)})") + return False + + +def main(): + """主函数：处理命令行参数并执行转换""" + if len(sys.argv) < 2: + print("用法: python tha4_student_convert_onnx_to_trt.py <模型名称> [精度]") + print("\n参数:") + print(" 模型名称: 自定义模型名称（例如：kanori_tha4）") + print(" 精度: fp32 或 fp16（默认：fp32）") + print("\n示例:") + print(" python tha4_student_convert_onnx_to_trt.py kanori_tha4") + print(" python tha4_student_convert_onnx_to_trt.py kanori_tha4 fp32") + print(" python tha4_student_convert_onnx_to_trt.py kanori_tha4 fp16") + print("\n注意:") + print(" - 请确保已经使用 onnx_export_tha4_student.py 导出了 ONNX 模型") + print(" - TensorRT 转换需要 CUDA 支持") + print(" - 转换过程可能需要几分钟时间") + sys.exit(1) + + model_name = sys.argv[1] + precision = sys.argv[2].lower() if len(sys.argv) > 2 else 'fp32' + + # 验证精度参数 + if precision not in ['fp32', 'fp16']: + print(f"错误: 不支持的精度类型 '{precision}'") + print("支持的精度: fp32, fp16") + sys.exit(1) + + # 检查 TensorRT 是否可用 + try: + import tensorrt as trt + print(f"TensorRT 版本: {trt.__version__}") + except ImportError: + print("错误: 未安装 TensorRT") + print("请安装 TensorRT: pip install tensorrt") + sys.exit(1) + + # 检查 CUDA 是否可用 + try: + import pycuda.driver as cuda + cuda.init() + print(f"CUDA 设备数量: {cuda.Device.count()}") + if cuda.Device.count() > 0: + device = cuda.Device(0) + print(f"GPU: {device.name()}") + except Exception as e: + print(f"警告: CUDA 初始化失败 - {e}") + print("TensorRT 转换需要 CUDA 支持") + sys.exit(1) + + print() + + # 执行转换 + success = convert_student_onnx_to_trt(model_name, precision) + + if success: + print(f"\n{'='*70}") + print("转换成功！可以在启动器中选择使用 TensorRT 模式") + print(f"{'='*70}") + sys.exit(0) + else: + print(f"\n{'='*70}") + print("转换失败，请检查错误信息") + print(f"{'='*70}") + sys.exit(1) + + +if __name__ == '__main__': + main()