diff --git a/02.构建TensorRT模型（N卡开启TensorRT使用）.bat b/02.构建TensorRT模型（N卡开启TensorRT使用）.bat index 68a0f18..0315d6f 100644 --- a/02.构建TensorRT模型（N卡开启TensorRT使用）.bat +++ b/02.构建TensorRT模型（N卡开启TensorRT使用）.bat @@ -6,4 +6,6 @@ call conda env list call conda activate ezvtb_rt_venv call conda env list -python ezvtb_rt_interface.py \ No newline at end of file +python ezvtb_rt_interface.py + +pause \ No newline at end of file diff --git a/env_conda.yaml b/env_conda.yaml deleted file mode 100644 index 8041af1..0000000 --- a/env_conda.yaml +++ /dev/null @@ -1,139 +0,0 @@ -name: talking-head-anime-2-demo -channels: - - pytorch - - conda-forge - - defaults -dependencies: - - argon2-cffi=20.1.0=py38h294d835_2 - - async_generator=1.10=py_0 - - attrs=20.3.0=pyhd3deb0d_0 - - backcall=0.2.0=pyh9f0ad1d_0 - - backports=1.0=py_2 - - backports.functools_lru_cache=1.6.1=py_0 - - blas=1.0=mkl - - bleach=3.3.0=pyh44b312d_0 - - brotlipy=0.7.0=py38h2bbff1b_1003 - - ca-certificates=2022.3.29=haa95532_1 - - certifi=2021.10.8=py38haa95532_2 - - cffi=1.14.5=py38hd8c33c5_0 - - charset-normalizer=2.0.4=pyhd3eb1b0_0 - - colorama=0.4.4=pyh9f0ad1d_0 - - cryptography=36.0.0=py38h21b164f_0 - - cudatoolkit=11.3.1=h59b6b97_2 - - cycler=0.10.0=py38_0 - - decorator=4.4.2=py_0 - - defusedxml=0.6.0=py_0 - - entrypoints=0.3=pyhd8ed1ab_1003 - - freetype=2.10.4=hd328e21_0 - - icc_rt=2019.0.0=h0cc432a_1 - - icu=58.2=ha925a31_3 - - idna=3.3=pyhd3eb1b0_0 - - importlib-metadata=3.4.0=py38haa244fe_0 - - importlib_metadata=3.4.0=hd8ed1ab_0 - - intel-openmp=2020.2=254 - - ipykernel=5.5.0=py38hc5df569_1 - - ipython=7.20.0=py38hc5df569_2 - - ipython_genutils=0.2.0=py_1 - - ipywidgets=7.6.3=pyhd3deb0d_0 - - jedi=0.18.0=py38haa244fe_2 - - jinja2=2.11.3=pyh44b312d_0 - - jpeg=9b=hb83a4c4_2 - - jsonschema=3.2.0=py_2 - - jupyter_client=6.1.11=pyhd8ed1ab_1 - - jupyter_core=4.7.1=py38haa244fe_0 - - jupyterlab_pygments=0.1.2=pyh9f0ad1d_0 - - jupyterlab_widgets=1.0.0=pyhd8ed1ab_1 - - kiwisolver=1.3.1=py38hd77b12b_0 - - libpng=1.6.37=h2a8f88b_0 - - libsodium=1.0.18=h8d14728_1 - - libtiff=4.1.0=h56a325e_1 - - libuv=1.40.0=he774522_0 - - lz4-c=1.9.3=h2bbff1b_0 - - m2w64-gcc-libgfortran=5.3.0=6 - - m2w64-gcc-libs=5.3.0=7 - - m2w64-gcc-libs-core=5.3.0=7 - - m2w64-gmp=6.1.0=2 - - m2w64-libwinpthread-git=5.0.0.4634.697f757=2 - - markupsafe=1.1.1=py38h294d835_3 - - matplotlib=3.3.4=py38haa95532_0 - - matplotlib-base=3.3.4=py38h49ac443_0 - - mistune=0.8.4=py38h294d835_1003 - - mkl=2020.2=256 - - mkl-service=2.3.0=py38h196d8e1_0 - - mkl_fft=1.2.1=py38h46781fe_0 - - mkl_random=1.1.1=py38h47e9c7a_0 - - msys2-conda-epoch=20160418=1 - - nbclient=0.5.2=pyhd8ed1ab_0 - - nbconvert=6.0.7=py38haa244fe_3 - - nbformat=5.1.2=pyhd8ed1ab_1 - - nest-asyncio=1.4.3=pyhd8ed1ab_0 - - ninja=1.10.2=py38h6d14046_0 - - notebook=6.2.0=py38haa244fe_0 - - numpy=1.19.2=py38hadc3359_0 - - numpy-base=1.19.2=py38ha3acd2a_0 - - olefile=0.46=py_0 - - openssl=1.1.1n=h2bbff1b_0 - - packaging=20.9=pyh44b312d_0 - - pandoc=2.11.4=h8ffe710_0 - - pandocfilters=1.4.2=py_1 - - parso=0.8.1=pyhd8ed1ab_0 - - pickleshare=0.7.5=py_1003 - - pillow=8.1.0=py38h4fa10fc_0 - - pip=21.0.1=py38haa95532_0 - - prometheus_client=0.9.0=pyhd3deb0d_0 - - prompt-toolkit=3.0.16=pyha770c72_0 - - pycparser=2.20=pyh9f0ad1d_2 - - pygments=2.8.0=pyhd8ed1ab_0 - - pyopenssl=22.0.0=pyhd3eb1b0_0 - - pyparsing=2.4.7=pyhd3eb1b0_0 - - pyqt=5.9.2=py38ha925a31_4 - - pyrsistent=0.17.3=py38h294d835_2 - - pysocks=1.7.1=py38haa95532_0 - - python=3.8.5=h5fd99cc_1 - - python-dateutil=2.8.1=pyhd3eb1b0_0 - - python_abi=3.8=1_cp38 - - pytorch=1.11.0=py3.8_cuda11.3_cudnn8_0 - - pytorch-mutex=1.0=cuda - - pywin32=300=py38h294d835_0 - - pywinpty=0.5.7=py38h32f6830_1 - - pyzmq=22.0.3=py38h7a0e47e_0 - - qt=5.9.7=vc14h73c81de_0 - - requests=2.27.1=pyhd3eb1b0_0 - - scipy=1.6.0=py38h14eb087_0 - - send2trash=1.5.0=py_0 - - setuptools=52.0.0=py38haa95532_0 - - sip=4.19.13=py38ha925a31_0 - - six=1.15.0=py38haa95532_0 - - sqlite=3.33.0=h2a8f88b_0 - - terminado=0.9.2=py38haa244fe_0 - - testpath=0.4.4=py_0 - - tk=8.6.10=he774522_0 - - torchaudio=0.11.0=py38_cu113 - - torchvision=0.12.0=py38_cu113 - - tornado=6.1=py38h2bbff1b_0 - - traitlets=5.0.5=py_0 - - typing_extensions=3.7.4.3=pyha847dfd_0 - - urllib3=1.26.8=pyhd3eb1b0_0 - - vc=14.2=h21ff451_1 - - vs2015_runtime=14.27.29016=h5e58377_2 - - wcwidth=0.2.5=pyh9f0ad1d_2 - - webencodings=0.5.1=py_1 - - wheel=0.36.2=pyhd3eb1b0_0 - - widgetsnbextension=3.5.1=py38haa244fe_4 - - win_inet_pton=1.1.0=py38haa95532_0 - - wincertstore=0.2=py38_0 - - winpty=0.4.3=4 - - xz=5.2.5=h62dcd97_0 - - zeromq=4.3.3=h0e60522_3 - - zipp=3.4.0=py_0 - - zlib=1.2.11=h62dcd97_4 - - zstd=1.4.5=h04227a9_0 - - pip: - - absl-py==1.0.0 - - mediapipe==0.8.9.1 - - opencv-contrib-python==4.5.5.64 - - opencv-python==4.5.5.64 - - protobuf==3.20.1 - - pyvirtualcam==0.9.1 - - pynput==1.7.6 - - pyanime4k==2.5.2 diff --git a/ezvtb_interface.md b/ezvtb_interface.md deleted file mode 100644 index 988ee9f..0000000 --- a/ezvtb_interface.md +++ /dev/null @@ -1,87 +0,0 @@ -# Interface -## Basic -Use the get_core() function below to get a `core` object, -you only need to interact with the `core` object with two methods: -`core.setImage()` and `core.inference()`. For example: - -``` -core = get_core() #create a core -core.setImage(img) #pass image to model -while True: #the eventloop - ret = core.inference(pose_params) #go forward for models - #do what ever you want with returned images -``` - -## Data interface -The core object is interacted through numpy interface -### Image data type -All images input& output using `np.uint8` datatype -### Image data layout -All images inputs&outputs are in `(Hight, Width, Channel)` data layout. For example, the input image that passed to `core.setImage()` must be in a shape of `(512, 512, 4)` -### Image channels -Inputs and output images are having alpha channel therefore 4 channels. The channel format must be BGRA. -### pose_params -Pose params are passed every time an inference going forward. It should be `np.float32` data in a shape of `(1, 45)` in a way works the same as Tha3 project. -### Output images -The returned object of `core.inference()` is a `List[np.ndarray]` object, which returns a list of one or multiple image results according to if interpolation is enabled, we will talk about what and how an interpolation happens soon. The output images in the list subjects to the above datatype&layout&channel constrains. - -The only thing you need to keep in mind is that, always consume(by transforming, or sending to other process, or do a copy to save) the result images before calling next `core.inference()`. The next return of inference will use those `np.ndarray`'s memory inplace. - -## Concepts -The `core` object hides all details but still you need to know the following concepts to correctly configurate `get_core()` function. -There are four conponents in a core : `tha`, `cacher`, `rife`, and `sr`(undergoing integration). -`tha` component is mandatory, all other components are optional - -### tha -`tha` component is where THA3 algorithm happens. It has following options. -1. `tha` component can use either *seperable* or *standard* convolution graph structure. This controls by `model_seperable` option in `get_core`, default to `True`. -2. By setting `model_half` to `True` to use half percision for `tha` component and vice versa. Default to `True`. -3. `model_cache` controls if `tha` model try to cache the intermediate results. Exchange RAM/VRAM space for time and GPU power. It is recommended and defaulted to `True` -4. Model caching intermediate result can happen on VRAM or RAM, controled by `model_vram_cache`. Default and recommended to `True`. -5. `model_cache_size` controls size of above cache in a unit of gigabytes, default to `1` as 1 gigabytes cache. - -### cacher -In the implementation, `cacher` component is optionally appended to the end of `tha` component output to cache image results of `tha` output. It has following options. - -1. `use_cacher` controls if cacher is enabled, default and recommended to `True`. -2. `cacher_on_database` controls if cacher is used on RAM or hard drive database. Default to `False` on RAM. -3. Since we are using **turbojpeg** to compress and save images, we should sepecify the compression ratio of JPEG algorithm, which is `cacher_quality` param in integer from 0 to 100. Smaller ratio == smaller compressed size == more cached items == higher cache hit rates == lower GPU computation comsumption. And vice versa. From my observation, I can not see great difference when `cacher_quality >= 85`. And if `sr` component is enabled, `cacher_quality >= 70` still works good. When `cacher_quality == 100`, the core will give up **turbojpeg** and directly storing in non-compressed format. Default to `85`. -4. `cacher_ram_size` specifies cache size on ram in a unit of gigabytes. Default to `2`. -5. `cacher_db_path` to specify path to database file path if `cacher_on_database == True` . - -### rife -In the core, we use RIFE algorithm to interpolate THA3 output. When this is enabled, the `core.inference()` will runs a bit slower and return multiple images at a time instead of just one. Please reference to above Interface section to understand how to handle outputs. -1. `use_interpolation` controls if you use `rife` component in core, default to `True`. -2. `interpolation_scale` is the scaling factor of `rife` interpolation. It can only be integer `2|3|4`. You will get same numbers of output images everything running `core.inference()`. For example, if `interpolation_scale==3` then `ret = core.inference()`, ret is a list of length 3, the first two images in the list are generated by `rife` and the last one is the orginal `tha` output. Default to `2`. -3. `interpolation_half` controls if using half percision `rife` model. Default to `True` - - -## Device setting -You can also specify some device setting in `get_core()` for multiple GPU and non-Nvidia GPU. -### device_id -`device_id` specifies device id integer, default to `0`. -### use_tensorrt -`use_tensorrt` specifies if TensorRT is used for inference. When `use_tensorrt == False`, DirectML will be used to support non-Nvidia GPUs. default to `True`. - -## More detailed example -suppose we want to work on GPU1, use full percision seperable float, and half percison 4x interpolation, a 8Gb RAM cacher. -``` -from ezvtb_interface import get_core -import cv2 -import numpy as np -core = get_core( - device_id = 1, - model_seperable = True, - model_half = True, - interpolation_scale = 4, - interpolation_half = True, - cacher_ram_size = 8 - ) -img = cv2.imread('img1.png', cv2.IMREAD_UNCHANGED) -core.setImage(img) -while True: #Eventloop here - pose_param = pipe.get(blocking=True) - ret = core.inference(pose_param) - for ret_img in ret: - cv2.imwrite('result.png',ret_img) #imwrite just for example -``` \ No newline at end of file diff --git a/ezvtb_rt_interface.py b/ezvtb_rt_interface.py index 5d9df7f..ebceb42 100644 --- a/ezvtb_rt_interface.py +++ b/ezvtb_rt_interface.py @@ -114,8 +114,13 @@ def get_core( if __name__ == '__main__': init_ezvtb_rt() - from ezvtb_rt.trt_utils import check_build_all_models + 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 4baea94..d602899 160000 --- a/ezvtuber-rt +++ b/ezvtuber-rt @@ -1 +1 @@ -Subproject commit 4baea942fc512e849f277b6a63f016fa66f2d445 +Subproject commit d602899f2dde6561d6d5dfbbfa908117c10900b3 diff --git a/main.py b/main.py index 03d12ec..6302a0c 100644 --- a/main.py +++ b/main.py @@ -512,11 +512,8 @@ class ModelClientProcess(Process): 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 + if args.use_tensorrt: + self.gpu_cache_hit_ratio.value = self.model.tha.morpher_cacher.hits / (self.model.tha.morpher_cacher.hits + self.model.tha.morpher_cacher.miss) @torch.no_grad() @@ -903,15 +900,16 @@ def main(): 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 + if args.extend_movement or args.bongo: + rm = cv2.getRotationMatrix2D((args.model_output_size / 2, args.model_output_size / 2), rotate_angle, k_scale) + rm[0, 2] += dx + rm[1, 2] += dy - postprocessed_image = cv2.warpAffine( - postprocessed_image, - rm, - (args.model_output_size, args.model_output_size)) + 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)