commit 1972d794d9d06e854e02f49bfcf1fc9a27dfe76e
Author: Isi <86603298+Isi-dev@users.noreply.github.com>
Date: Thu Aug 1 21:09:40 2024 +0100
First commit
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..53200a9
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,18 @@
+*.pkl
+*.pt
+*.mov
+*.pth
+*.mov
+*.npz
+*.npy
+*.boj
+*.onnx
+*.tar
+*.bin
+cache*
+.DS_Store
+*DS_Store
+outputs/
+**/__pycache__
+***/__pycache__
+*/__pycache__
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..6755761
--- /dev/null
+++ b/README.md
@@ -0,0 +1,104 @@
+
+
+
+
+
+
+#This is my ComfyUi-windows implementation for the image animation project -> UniAnimate: Taming Unified Video Diffusion Models for Consistent Human Image Animation
+
+[🎨 Source Project Page](https://unianimate.github.io/)
+
+
+
+
+## Getting Started
+
+The ComfyUI nodes created are "Align & Generate poses for UniAnimate" & "Animate image with UniAnimate"
+
+I used a ComfyUI_windows_portable to test the nodes in a Windows 10 OS with 16GB RAM & 12GB VRAM Nvidia Graphics Card
+
+Download or clone this repository and place it in ComfyUI_windows_portable\ComfyUI\custom_nodes\
+
+You will need python>=3.9 in your ComfyUI Environment.
+I tested the project with the following pytorch versions which you can install as follows
+
+```
+conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.8 -c pytorch -c nvidia
+
+```
+
+Or
+
+```
+conda install pytorch==2.3.1 torchvision==0.18.1 torchaudio==2.3.1 pytorch-cuda=11.8 -c pytorch -c nvidia
+
+```
+
+If not installed, then:
+pip install opencv-python
+pip install pytorch_lightning
+pip install lightning_utilities #if not installed
+pip install lightning_fabric #if not installed
+pip install torchmetrics
+pip install xFormers = 0.0.20 or copy torch 2.0.1 and supporting libraries and xFormer from A1111 and place in your Environment\Lib\site-packages (or)
+pip3 install -U xformers --index-url https://download.pytorch.org/whl/cu118 (for pytorch==2.3.1)
+pip install oss2
+pip install einops
+pip install args
+pip install modelscope
+
+
+Download the required models (Around 14GB) after installing modelscope :
+
+```
+python modeldownloader.py
+
+```
+
+After downloading all the models, move them manually from 'checkpoints/iic/unianimate/' to the 'checkpoints' directory
+Or move them via your command line interface:
+
+```
+python mv ./checkpoints/iic/unianimate/* ./checkpoints/
+
+```
+
+All the models should be in the '\Path-to-UniAnimate\checkpoints' folder as follows:
+
+
+```
+./checkpoints/
+|---- dw-ll_ucoco_384.onnx
+|---- open_clip_pytorch_model.bin
+|---- unianimate_16f_32f_non_ema_223000.pth
+|---- v2-1_512-ema-pruned.ckpt
+â””---- yolox_l.onnx
+
+```
+
+You can now upload the workflow in your '\Path-to-UniAnimate\' folder which is titled 'basicUniAnimateWorkflow.json', install missing custom nodes with the ComfyUI Manager if necessary, upload a picture & video (You can use those in the 'assets' folder), and run!
+
+
+**✔ Some tips**:
+
+- > In the 'Align & Generate poses for UniAnimate' node, the first frame in the target pose sequence is used to calculate the scale coefficient of the alignment. Therefore, if the first frame in the target pose sequence contains the entire face and pose (hand and foot), it can help obtain more accurate estimation and better video generation results.
+
+- > Generating 32 frames of video with a resolution of [512, 768] usually takes about 7 minutes.
+
+
+
+- > To run the "Animate image with UniAnimate" node, **~12G** GPU memory will be used. If your GPU is smaller than this, you can change the `max_frames: 32` to other values, e.g., 24, 16, and 8.
+
+You can also generate a video first, and then upload the last frame of the video as a pic to generate the next frames.
+
+
+@article{
+ project={ComfyUi-windows implementation for the image animation project -> UniAnimate: Taming Unified Video Diffusion Models for Consistent Human Image Animation},
+ developer={Isimemen Omoifo Jnr},
+ year={2024}
+}
+
+
+## Disclaimer
+
+I explicitly disclaim any responsibility for user-generated content. Users are solely liable for their actions while using these nodes and the generative model. I and the source project contributors have no legal affiliation with, nor accountability for, users' behaviors. It is imperative to use these nodes and the generative model responsibly, adhering to both ethical and legal standards.
diff --git a/__init__.py b/__init__.py
new file mode 100644
index 0000000..c24e4d6
--- /dev/null
+++ b/__init__.py
@@ -0,0 +1,3 @@
+from .uniAnimate_Inference import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS']
\ No newline at end of file
diff --git a/basicUniAnimateWorkflow.json b/basicUniAnimateWorkflow.json
new file mode 100644
index 0000000..7e19d87
--- /dev/null
+++ b/basicUniAnimateWorkflow.json
@@ -0,0 +1,427 @@
+{
+ "last_node_id": 16,
+ "last_link_id": 18,
+ "nodes": [
+ {
+ "id": 11,
+ "type": "VHS_LoadVideo",
+ "pos": [
+ 580,
+ 465
+ ],
+ "size": [
+ 235.1999969482422,
+ 620.5687812408271
+ ],
+ "flags": {},
+ "order": 0,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "meta_batch",
+ "type": "VHS_BatchManager",
+ "link": null
+ }
+ ],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 11
+ ],
+ "shape": 3,
+ "slot_index": 0
+ },
+ {
+ "name": "frame_count",
+ "type": "INT",
+ "links": null,
+ "shape": 3
+ },
+ {
+ "name": "audio",
+ "type": "VHS_AUDIO",
+ "links": null,
+ "shape": 3
+ },
+ {
+ "name": "video_info",
+ "type": "VHS_VIDEOINFO",
+ "links": null,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "VHS_LoadVideo"
+ },
+ "widgets_values": {
+ "video": "0001-0032.mp4",
+ "force_rate": 0,
+ "force_size": "Disabled",
+ "custom_width": 512,
+ "custom_height": 512,
+ "frame_load_cap": 0,
+ "skip_first_frames": 0,
+ "select_every_nth": 1,
+ "choose video to upload": "image",
+ "videopreview": {
+ "hidden": false,
+ "paused": false,
+ "params": {
+ "frame_load_cap": 0,
+ "skip_first_frames": 0,
+ "force_rate": 0,
+ "filename": "0001-0032.mp4",
+ "type": "input",
+ "format": "video/mp4",
+ "select_every_nth": 1
+ }
+ }
+ }
+ },
+ {
+ "id": 13,
+ "type": "PreviewImage",
+ "pos": [
+ 900,
+ 207
+ ],
+ "size": [
+ 210,
+ 246
+ ],
+ "flags": {},
+ "order": 3,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 12
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "PreviewImage"
+ }
+ },
+ {
+ "id": 10,
+ "type": "LoadImage",
+ "pos": [
+ 538,
+ 95
+ ],
+ "size": [
+ 315,
+ 314.00000762939453
+ ],
+ "flags": {},
+ "order": 1,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 10,
+ 15
+ ],
+ "shape": 3,
+ "slot_index": 0
+ },
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": null,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "LoadImage"
+ },
+ "widgets_values": [
+ "untitled.png",
+ "image"
+ ]
+ },
+ {
+ "id": 12,
+ "type": "Gen_align_pose",
+ "pos": [
+ 880,
+ 493
+ ],
+ "size": {
+ "0": 267,
+ "1": 46
+ },
+ "flags": {},
+ "order": 2,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "reference_image",
+ "type": "IMAGE",
+ "link": 10
+ },
+ {
+ "name": "video",
+ "type": "IMAGE",
+ "link": 11
+ }
+ ],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 12,
+ 16
+ ],
+ "shape": 3,
+ "slot_index": 0
+ },
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 14,
+ 17
+ ],
+ "shape": 3,
+ "slot_index": 1
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "Gen_align_pose"
+ }
+ },
+ {
+ "id": 14,
+ "type": "PreviewImage",
+ "pos": [
+ 904,
+ 595
+ ],
+ "size": [
+ 210,
+ 246
+ ],
+ "flags": {},
+ "order": 4,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 14
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "PreviewImage"
+ }
+ },
+ {
+ "id": 15,
+ "type": "UniAnimateImage",
+ "pos": [
+ 1178,
+ 470
+ ],
+ "size": {
+ "0": 315,
+ "1": 194
+ },
+ "flags": {},
+ "order": 5,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "reference_image",
+ "type": "IMAGE",
+ "link": 15
+ },
+ {
+ "name": "ref_pose",
+ "type": "IMAGE",
+ "link": 16
+ },
+ {
+ "name": "pose_sequence",
+ "type": "IMAGE",
+ "link": 17
+ }
+ ],
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 18
+ ],
+ "shape": 3,
+ "slot_index": 0
+ },
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": null,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "UniAnimateImage"
+ },
+ "widgets_values": [
+ 30,
+ true,
+ 1,
+ 32,
+ 512
+ ]
+ },
+ {
+ "id": 16,
+ "type": "VHS_VideoCombine",
+ "pos": [
+ 1522,
+ 225
+ ],
+ "size": [
+ 315,
+ 746.5
+ ],
+ "flags": {},
+ "order": 6,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "images",
+ "type": "IMAGE",
+ "link": 18
+ },
+ {
+ "name": "audio",
+ "type": "VHS_AUDIO",
+ "link": null
+ },
+ {
+ "name": "meta_batch",
+ "type": "VHS_BatchManager",
+ "link": null
+ }
+ ],
+ "outputs": [
+ {
+ "name": "Filenames",
+ "type": "VHS_FILENAMES",
+ "links": null,
+ "shape": 3
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "VHS_VideoCombine"
+ },
+ "widgets_values": {
+ "frame_rate": 16,
+ "loop_count": 0,
+ "filename_prefix": "UniAnimate/vid",
+ "format": "video/h264-mp4",
+ "pix_fmt": "yuv420p",
+ "crf": 19,
+ "save_metadata": true,
+ "pingpong": false,
+ "save_output": true,
+ "videopreview": {
+ "hidden": false,
+ "paused": false,
+ "params": {
+ "filename": "vid_00001.mp4",
+ "subfolder": "UniAnimate",
+ "type": "output",
+ "format": "video/h264-mp4"
+ }
+ }
+ }
+ }
+ ],
+ "links": [
+ [
+ 10,
+ 10,
+ 0,
+ 12,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 11,
+ 11,
+ 0,
+ 12,
+ 1,
+ "IMAGE"
+ ],
+ [
+ 12,
+ 12,
+ 0,
+ 13,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 14,
+ 12,
+ 1,
+ 14,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 15,
+ 10,
+ 0,
+ 15,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 16,
+ 12,
+ 0,
+ 15,
+ 1,
+ "IMAGE"
+ ],
+ [
+ 17,
+ 12,
+ 1,
+ 15,
+ 2,
+ "IMAGE"
+ ],
+ [
+ 18,
+ 15,
+ 0,
+ 16,
+ 0,
+ "IMAGE"
+ ]
+ ],
+ "groups": [],
+ "config": {},
+ "extra": {
+ "ds": {
+ "scale": 1,
+ "offset": [
+ -94.68478436382884,
+ 8.13263156125987
+ ]
+ }
+ },
+ "version": 0.4
+}
\ No newline at end of file
diff --git a/environment.yaml b/environment.yaml
new file mode 100644
index 0000000..a60b160
--- /dev/null
+++ b/environment.yaml
@@ -0,0 +1,236 @@
+name: /mnt/user/miniconda3/envs/dtrans
+channels:
+ - http://mirrors.aliyun.com/anaconda/pkgs/main
+ - defaults
+dependencies:
+ - _libgcc_mutex=0.1=main
+ - _openmp_mutex=5.1=1_gnu
+ - ca-certificates=2023.12.12=h06a4308_0
+ - ld_impl_linux-64=2.38=h1181459_1
+ - libffi=3.4.4=h6a678d5_0
+ - libgcc-ng=11.2.0=h1234567_1
+ - libgomp=11.2.0=h1234567_1
+ - libstdcxx-ng=11.2.0=h1234567_1
+ - ncurses=6.4=h6a678d5_0
+ - openssl=3.0.12=h7f8727e_0
+ - pip=23.3.1=py39h06a4308_0
+ - python=3.9.18=h955ad1f_0
+ - readline=8.2=h5eee18b_0
+ - setuptools=68.2.2=py39h06a4308_0
+ - sqlite=3.41.2=h5eee18b_0
+ - tk=8.6.12=h1ccaba5_0
+ - wheel=0.41.2=py39h06a4308_0
+ - xz=5.4.5=h5eee18b_0
+ - zlib=1.2.13=h5eee18b_0
+ - pip:
+ - aiofiles==23.2.1
+ - aiohttp==3.9.1
+ - aiosignal==1.3.1
+ - aliyun-python-sdk-core==2.14.0
+ - aliyun-python-sdk-kms==2.16.2
+ - altair==5.2.0
+ - annotated-types==0.6.0
+ - antlr4-python3-runtime==4.9.3
+ - anyio==4.2.0
+ - argparse==1.4.0
+ - asttokens==2.4.1
+ - async-timeout==4.0.3
+ - attrs==23.2.0
+ - automat==22.10.0
+ - beartype==0.16.4
+ - blessed==1.20.0
+ - buildtools==1.0.6
+ - causal-conv1d==1.1.3.post1
+ - certifi==2023.11.17
+ - cffi==1.16.0
+ - chardet==5.2.0
+ - charset-normalizer==3.3.2
+ - clean-fid==0.1.35
+ - click==8.1.7
+ - clip==1.0
+ - cmake==3.28.1
+ - colorama==0.4.6
+ - coloredlogs==15.0.1
+ - constantly==23.10.4
+ - contourpy==1.2.0
+ - crcmod==1.7
+ - cryptography==41.0.7
+ - cycler==0.12.1
+ - decorator==5.1.1
+ - decord==0.6.0
+ - diffusers==0.26.3
+ - docopt==0.6.2
+ - easydict==1.11
+ - einops==0.7.0
+ - exceptiongroup==1.2.0
+ - executing==2.0.1
+ - fairscale==0.4.13
+ - fastapi==0.109.0
+ - ffmpeg==1.4
+ - ffmpy==0.3.1
+ - filelock==3.13.1
+ - flatbuffers==24.3.25
+ - fonttools==4.47.2
+ - frozenlist==1.4.1
+ - fsspec==2023.12.2
+ - ftfy==6.1.3
+ - furl==2.1.3
+ - gpustat==1.1.1
+ - gradio==4.14.0
+ - gradio-client==0.8.0
+ - greenlet==3.0.3
+ - h11==0.14.0
+ - httpcore==1.0.2
+ - httpx==0.26.0
+ - huggingface-hub==0.20.2
+ - humanfriendly==10.0
+ - hyperlink==21.0.0
+ - idna==3.6
+ - imageio==2.33.1
+ - imageio-ffmpeg==0.4.9
+ - importlib-metadata==7.0.1
+ - importlib-resources==6.1.1
+ - incremental==22.10.0
+ - ipdb==0.13.13
+ - ipython==8.18.1
+ - jedi==0.19.1
+ - jinja2==3.1.3
+ - jmespath==0.10.0
+ - joblib==1.3.2
+ - jsonschema==4.21.0
+ - jsonschema-specifications==2023.12.1
+ - kiwisolver==1.4.5
+ - kornia==0.7.1
+ - lazy-loader==0.3
+ - lightning-utilities==0.10.0
+ - lit==17.0.6
+ - lpips==0.1.4
+ - mamba-ssm==1.1.4
+ - markdown-it-py==3.0.0
+ - markupsafe==2.1.3
+ - matplotlib==3.8.2
+ - matplotlib-inline==0.1.6
+ - mdurl==0.1.2
+ - motion-vector-extractor==1.0.6
+ - mpmath==1.3.0
+ - multidict==6.0.4
+ - mypy-extensions==1.0.0
+ - networkx==3.2.1
+ - ninja==1.11.1.1
+ - numpy==1.26.3
+ - nvidia-cublas-cu11==11.10.3.66
+ - nvidia-cublas-cu12==12.1.3.1
+ - nvidia-cuda-cupti-cu11==11.7.101
+ - nvidia-cuda-cupti-cu12==12.1.105
+ - nvidia-cuda-nvrtc-cu11==11.7.99
+ - nvidia-cuda-nvrtc-cu12==12.1.105
+ - nvidia-cuda-runtime-cu11==11.7.99
+ - nvidia-cuda-runtime-cu12==12.1.105
+ - nvidia-cudnn-cu11==8.5.0.96
+ - nvidia-cudnn-cu12==8.9.2.26
+ - nvidia-cufft-cu11==10.9.0.58
+ - nvidia-cufft-cu12==11.0.2.54
+ - nvidia-curand-cu11==10.2.10.91
+ - nvidia-curand-cu12==10.3.2.106
+ - nvidia-cusolver-cu11==11.4.0.1
+ - nvidia-cusolver-cu12==11.4.5.107
+ - nvidia-cusparse-cu11==11.7.4.91
+ - nvidia-cusparse-cu12==12.1.0.106
+ - nvidia-ml-py==12.535.133
+ - nvidia-nccl-cu11==2.14.3
+ - nvidia-nccl-cu12==2.19.3
+ - nvidia-nvjitlink-cu12==12.3.101
+ - nvidia-nvtx-cu11==11.7.91
+ - nvidia-nvtx-cu12==12.1.105
+ - omegaconf==2.3.0
+ - onnxruntime==1.18.0
+ - open-clip-torch==2.24.0
+ - opencv-python==4.5.3.56
+ - opencv-python-headless==4.9.0.80
+ - orderedmultidict==1.0.1
+ - orjson==3.9.10
+ - oss2==2.18.4
+ - packaging==23.2
+ - pandas==2.1.4
+ - parso==0.8.3
+ - pexpect==4.9.0
+ - pillow==10.2.0
+ - piq==0.8.0
+ - pkgconfig==1.5.5
+ - prompt-toolkit==3.0.43
+ - protobuf==4.25.2
+ - psutil==5.9.8
+ - ptflops==0.7.2.2
+ - ptyprocess==0.7.0
+ - pure-eval==0.2.2
+ - pycparser==2.21
+ - pycryptodome==3.20.0
+ - pydantic==2.5.3
+ - pydantic-core==2.14.6
+ - pydub==0.25.1
+ - pygments==2.17.2
+ - pynvml==11.5.0
+ - pyparsing==3.1.1
+ - pyre-extensions==0.0.29
+ - python-dateutil==2.8.2
+ - python-multipart==0.0.6
+ - pytorch-lightning==2.1.3
+ - pytz==2023.3.post1
+ - pyyaml==6.0.1
+ - redo==2.0.4
+ - referencing==0.32.1
+ - regex==2023.12.25
+ - requests==2.31.0
+ - rich==13.7.0
+ - rotary-embedding-torch==0.5.3
+ - rpds-py==0.17.1
+ - ruff==0.2.0
+ - safetensors==0.4.1
+ - scikit-image==0.22.0
+ - scikit-learn==1.4.0
+ - scipy==1.11.4
+ - semantic-version==2.10.0
+ - sentencepiece==0.1.99
+ - shellingham==1.5.4
+ - simplejson==3.19.2
+ - six==1.16.0
+ - sk-video==1.1.10
+ - sniffio==1.3.0
+ - sqlalchemy==2.0.27
+ - stack-data==0.6.3
+ - starlette==0.35.1
+ - sympy==1.12
+ - thop==0.1.1-2209072238
+ - threadpoolctl==3.2.0
+ - tifffile==2023.12.9
+ - timm==0.9.12
+ - tokenizers==0.15.0
+ - tomli==2.0.1
+ - tomlkit==0.12.0
+ - toolz==0.12.0
+ - torch==2.0.1+cu118
+ - torchaudio==2.0.2+cu118
+ - torchdiffeq==0.2.3
+ - torchmetrics==1.3.0.post0
+ - torchsde==0.2.6
+ - torchvision==0.15.2+cu118
+ - tqdm==4.66.1
+ - traitlets==5.14.1
+ - trampoline==0.1.2
+ - transformers==4.36.2
+ - triton==2.0.0
+ - twisted==23.10.0
+ - typer==0.9.0
+ - typing-extensions==4.9.0
+ - typing-inspect==0.9.0
+ - tzdata==2023.4
+ - urllib3==2.1.0
+ - uvicorn==0.26.0
+ - wcwidth==0.2.13
+ - websockets==11.0.3
+ - xformers==0.0.20
+ - yarl==1.9.4
+ - zipp==3.17.0
+ - zope-interface==6.2
+ - onnxruntime-gpu==1.13.1
+prefix: /mnt/user/miniconda3/envs/dtrans
diff --git a/modeldownloader.py b/modeldownloader.py
new file mode 100644
index 0000000..3916220
--- /dev/null
+++ b/modeldownloader.py
@@ -0,0 +1,2 @@
+from modelscope.hub.snapshot_download import snapshot_download
+model_dir = snapshot_download('iic/unianimate', cache_dir='checkpoints/')
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..57b8664
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,208 @@
+aiofiles==23.2.1
+aiohttp==3.9.1
+aiosignal==1.3.1
+aliyun-python-sdk-core==2.14.0
+aliyun-python-sdk-kms==2.16.2
+altair==5.2.0
+annotated-types==0.6.0
+antlr4-python3-runtime==4.9.3
+anyio==4.2.0
+asttokens==2.4.1
+async-timeout==4.0.3
+attrs==23.2.0
+Automat==22.10.0
+beartype==0.16.4
+blessed==1.20.0
+buildtools==1.0.6
+# causal-conv1d==1.1.3.post1
+certifi==2023.11.17
+cffi==1.16.0
+chardet==5.2.0
+charset-normalizer==3.3.2
+clean-fid==0.1.35
+click==8.1.7
+# clip==1.0
+cmake==3.28.1
+colorama==0.4.6
+coloredlogs==15.0.1
+constantly==23.10.4
+contourpy==1.2.0
+crcmod==1.7
+cryptography==41.0.7
+cycler==0.12.1
+decorator==5.1.1
+decord==0.6.0
+diffusers==0.26.3
+docopt==0.6.2
+easydict==1.11
+einops==0.7.0
+exceptiongroup==1.2.0
+executing==2.0.1
+fairscale==0.4.13
+fastapi==0.109.0
+ffmpeg==1.4
+ffmpy==0.3.1
+filelock==3.13.1
+flatbuffers==24.3.25
+fonttools==4.47.2
+frozenlist==1.4.1
+fsspec==2023.12.2
+ftfy==6.1.3
+furl==2.1.3
+gpustat==1.1.1
+gradio==4.14.0
+gradio_client==0.8.0
+greenlet==3.0.3
+h11==0.14.0
+httpcore==1.0.2
+httpx==0.26.0
+huggingface-hub==0.20.2
+humanfriendly==10.0
+hyperlink==21.0.0
+idna==3.6
+imageio==2.33.1
+imageio-ffmpeg==0.4.9
+importlib-metadata==7.0.1
+importlib-resources==6.1.1
+incremental==22.10.0
+ipdb==0.13.13
+ipython==8.18.1
+jedi==0.19.1
+Jinja2==3.1.3
+jmespath==0.10.0
+joblib==1.3.2
+jsonschema==4.21.0
+jsonschema-specifications==2023.12.1
+kiwisolver==1.4.5
+kornia==0.7.1
+lazy_loader==0.3
+lightning-utilities==0.10.0
+lit==17.0.6
+lpips==0.1.4
+markdown-it-py==3.0.0
+MarkupSafe==2.1.3
+matplotlib==3.8.2
+matplotlib-inline==0.1.6
+mdurl==0.1.2
+# motion-vector-extractor==1.0.6
+mpmath==1.3.0
+multidict==6.0.4
+mypy-extensions==1.0.0
+networkx==3.2.1
+ninja==1.11.1.1
+numpy==1.26.3
+nvidia-cublas-cu11==11.10.3.66
+nvidia-cublas-cu12==12.1.3.1
+nvidia-cuda-cupti-cu11==11.7.101
+nvidia-cuda-cupti-cu12==12.1.105
+nvidia-cuda-nvrtc-cu11==11.7.99
+nvidia-cuda-nvrtc-cu12==12.1.105
+nvidia-cuda-runtime-cu11==11.7.99
+nvidia-cuda-runtime-cu12==12.1.105
+nvidia-cudnn-cu11==8.5.0.96
+nvidia-cudnn-cu12==8.9.2.26
+nvidia-cufft-cu11==10.9.0.58
+nvidia-cufft-cu12==11.0.2.54
+nvidia-curand-cu11==10.2.10.91
+nvidia-curand-cu12==10.3.2.106
+nvidia-cusolver-cu11==11.4.0.1
+nvidia-cusolver-cu12==11.4.5.107
+nvidia-cusparse-cu11==11.7.4.91
+nvidia-cusparse-cu12==12.1.0.106
+nvidia-ml-py==12.535.133
+nvidia-nccl-cu11==2.14.3
+nvidia-nccl-cu12==2.19.3
+nvidia-nvjitlink-cu12==12.3.101
+nvidia-nvtx-cu11==11.7.91
+nvidia-nvtx-cu12==12.1.105
+omegaconf==2.3.0
+onnxruntime==1.18.0
+open-clip-torch==2.24.0
+opencv-python==4.5.3.56
+opencv-python-headless==4.9.0.80
+orderedmultidict==1.0.1
+orjson==3.9.10
+oss2==2.18.4
+# packaging==23.2
+pandas==2.1.4
+parso==0.8.3
+pexpect==4.9.0
+pillow==10.2.0
+piq==0.8.0
+pkgconfig==1.5.5
+prompt-toolkit==3.0.43
+protobuf==4.25.2
+psutil==5.9.8
+ptflops==0.7.2.2
+ptyprocess==0.7.0
+pure-eval==0.2.2
+pycparser==2.21
+pycryptodome==3.20.0
+pydantic==2.5.3
+pydantic_core==2.14.6
+pydub==0.25.1
+Pygments==2.17.2
+pynvml==11.5.0
+pyparsing==3.1.1
+pyre-extensions==0.0.29
+python-dateutil==2.8.2
+python-multipart==0.0.6
+pytorch-lightning==2.1.3
+pytz==2023.3.post1
+PyYAML==6.0.1
+redo==2.0.4
+referencing==0.32.1
+regex==2023.12.25
+requests==2.31.0
+rich==13.7.0
+rotary-embedding-torch==0.5.3
+rpds-py==0.17.1
+ruff==0.2.0
+safetensors==0.4.1
+scikit-image==0.22.0
+scikit-learn==1.4.0
+scipy==1.11.4
+semantic-version==2.10.0
+sentencepiece==0.1.99
+shellingham==1.5.4
+simplejson==3.19.2
+six==1.16.0
+sk-video==1.1.10
+sniffio==1.3.0
+SQLAlchemy==2.0.27
+stack-data==0.6.3
+starlette==0.35.1
+sympy==1.12
+thop==0.1.1.post2209072238
+threadpoolctl==3.2.0
+tifffile==2023.12.9
+timm==0.9.12
+tokenizers==0.15.0
+tomli==2.0.1
+tomlkit==0.12.0
+toolz==0.12.0
+# torch==2.0.1+cu118
+# torchaudio==2.0.2+cu118
+torchdiffeq==0.2.3
+torchmetrics==1.3.0.post0
+torchsde==0.2.6
+# torchvision==0.15.2+cu118
+tqdm==4.66.1
+traitlets==5.14.1
+trampoline==0.1.2
+transformers==4.36.2
+triton==2.0.0
+Twisted==23.10.0
+typer==0.9.0
+typing-inspect==0.9.0
+typing_extensions==4.9.0
+tzdata==2023.4
+urllib3==2.1.0
+uvicorn==0.26.0
+wcwidth==0.2.13
+websockets==11.0.3
+xformers==0.0.20
+yarl==1.9.4
+zipp==3.17.0
+zope.interface==6.2
+onnxruntime-gpu==1.13.1
diff --git a/run_align_pose.py b/run_align_pose.py
new file mode 100644
index 0000000..b0d0c9a
--- /dev/null
+++ b/run_align_pose.py
@@ -0,0 +1,709 @@
+# Openpose
+# Original from CMU https://github.com/CMU-Perceptual-Computing-Lab/openpose
+# 2nd Edited by https://github.com/Hzzone/pytorch-openpose
+# 3rd Edited by ControlNet
+# 4th Edited by ControlNet (added face and correct hands)
+
+import os
+os.environ["KMP_DUPLICATE_LIB_OK"]="TRUE"
+import cv2
+import torch
+import numpy as np
+import json
+import copy
+import torch
+import random
+import argparse
+import shutil
+import tempfile
+import subprocess
+import numpy as np
+import math
+
+import torch.multiprocessing as mp
+import torch.distributed as dist
+# import pickle
+import logging
+# from io import BytesIO
+# import oss2 as oss
+# import os.path as osp
+
+import sys
+from .dwpose import util
+from .dwpose.wholebody import Wholebody
+
+
+def smoothing_factor(t_e, cutoff):
+ r = 2 * math.pi * cutoff * t_e
+ return r / (r + 1)
+
+
+def exponential_smoothing(a, x, x_prev):
+ return a * x + (1 - a) * x_prev
+
+
+class OneEuroFilter:
+ def __init__(self, t0, x0, dx0=0.0, min_cutoff=1.0, beta=0.0,
+ d_cutoff=1.0):
+ """Initialize the one euro filter."""
+ # The parameters.
+ self.min_cutoff = float(min_cutoff)
+ self.beta = float(beta)
+ self.d_cutoff = float(d_cutoff)
+ # Previous values.
+ self.x_prev = x0
+ self.dx_prev = float(dx0)
+ self.t_prev = float(t0)
+
+ def __call__(self, t, x):
+ """Compute the filtered signal."""
+ t_e = t - self.t_prev
+
+ # The filtered derivative of the signal.
+ a_d = smoothing_factor(t_e, self.d_cutoff)
+ dx = (x - self.x_prev) / t_e
+ dx_hat = exponential_smoothing(a_d, dx, self.dx_prev)
+
+ # The filtered signal.
+ cutoff = self.min_cutoff + self.beta * abs(dx_hat)
+ a = smoothing_factor(t_e, cutoff)
+ x_hat = exponential_smoothing(a, x, self.x_prev)
+
+ # Memorize the previous values.
+ self.x_prev = x_hat
+ self.dx_prev = dx_hat
+ self.t_prev = t
+
+ return x_hat
+
+
+def get_logger(name="essmc2"):
+ logger = logging.getLogger(name)
+ logger.propagate = False
+ if len(logger.handlers) == 0:
+ std_handler = logging.StreamHandler(sys.stdout)
+ formatter = logging.Formatter(
+ '%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+ std_handler.setFormatter(formatter)
+ std_handler.setLevel(logging.INFO)
+ logger.setLevel(logging.INFO)
+ logger.addHandler(std_handler)
+ return logger
+
+class DWposeDetector:
+ def __init__(self):
+
+ self.pose_estimation = Wholebody()
+
+ def __call__(self, oriImg):
+ oriImg = oriImg.copy()
+ # print(f'The shape of the image should be in HWC format but it is currently: {oriImg.shape} ')
+ H, W, C = oriImg.shape
+ with torch.no_grad():
+ candidate, subset = self.pose_estimation(oriImg)
+ candidate = candidate[0][np.newaxis, :, :]
+ subset = subset[0][np.newaxis, :]
+ nums, keys, locs = candidate.shape
+ candidate[..., 0] /= float(W)
+ candidate[..., 1] /= float(H)
+ body = candidate[:,:18].copy()
+ body = body.reshape(nums*18, locs)
+ score = subset[:,:18].copy()
+
+ for i in range(len(score)):
+ for j in range(len(score[i])):
+ if score[i][j] > 0.3:
+ score[i][j] = int(18*i+j)
+ else:
+ score[i][j] = -1
+
+ un_visible = subset<0.3
+ candidate[un_visible] = -1
+
+ bodyfoot_score = subset[:,:24].copy()
+ for i in range(len(bodyfoot_score)):
+ for j in range(len(bodyfoot_score[i])):
+ if bodyfoot_score[i][j] > 0.3:
+ bodyfoot_score[i][j] = int(18*i+j)
+ else:
+ bodyfoot_score[i][j] = -1
+ if -1 not in bodyfoot_score[:,18] and -1 not in bodyfoot_score[:,19]:
+ bodyfoot_score[:,18] = np.array([18.])
+ else:
+ bodyfoot_score[:,18] = np.array([-1.])
+ if -1 not in bodyfoot_score[:,21] and -1 not in bodyfoot_score[:,22]:
+ bodyfoot_score[:,19] = np.array([19.])
+ else:
+ bodyfoot_score[:,19] = np.array([-1.])
+ bodyfoot_score = bodyfoot_score[:, :20]
+
+ bodyfoot = candidate[:,:24].copy()
+
+ for i in range(nums):
+ if -1 not in bodyfoot[i][18] and -1 not in bodyfoot[i][19]:
+ bodyfoot[i][18] = (bodyfoot[i][18]+bodyfoot[i][19])/2
+ else:
+ bodyfoot[i][18] = np.array([-1., -1.])
+ if -1 not in bodyfoot[i][21] and -1 not in bodyfoot[i][22]:
+ bodyfoot[i][19] = (bodyfoot[i][21]+bodyfoot[i][22])/2
+ else:
+ bodyfoot[i][19] = np.array([-1., -1.])
+
+ bodyfoot = bodyfoot[:,:20,:]
+ bodyfoot = bodyfoot.reshape(nums*20, locs)
+
+ foot = candidate[:,18:24]
+
+ faces = candidate[:,24:92]
+
+ hands = candidate[:,92:113]
+ hands = np.vstack([hands, candidate[:,113:]])
+
+ # bodies = dict(candidate=body, subset=score)
+ bodies = dict(candidate=bodyfoot, subset=bodyfoot_score)
+ pose = dict(bodies=bodies, hands=hands, faces=faces)
+
+ # return draw_pose(pose, H, W)
+ return pose
+
+def draw_pose(pose, H, W):
+ bodies = pose['bodies']
+ faces = pose['faces']
+ hands = pose['hands']
+ candidate = bodies['candidate']
+ subset = bodies['subset']
+ canvas = np.zeros(shape=(H, W, 3), dtype=np.uint8)
+
+ canvas = util.draw_body_and_foot(canvas, candidate, subset)
+
+ canvas = util.draw_handpose(canvas, hands)
+
+ canvas_without_face = copy.deepcopy(canvas)
+
+ canvas = util.draw_facepose(canvas, faces)
+
+ return canvas_without_face, canvas
+
+def dw_func(_id, frame, dwpose_model, dwpose_woface_folder='tmp_dwpose_wo_face', dwpose_withface_folder='tmp_dwpose_with_face'):
+
+ # frame = cv2.imread(frame_name, cv2.IMREAD_COLOR)
+ pose = dwpose_model(frame)
+
+ return pose
+
+
+def mp_main(reference_image, video):
+
+ logger.info(f"There are {video.size(0)} frames for extracting poses")
+
+ logger.info('LOAD: DW Pose Model')
+ dwpose_model = DWposeDetector()
+
+ results_vis = []
+ num_frames = video.size(0)
+
+
+ for i in range(num_frames):
+ logger.info(f"Processing frame {i + 1}/{num_frames}")
+ frame = video[i].permute(0, 1, 2).cpu().numpy()*255 # Convert to HWC format and numpy array
+ frame = np.flip(frame, axis=2)
+ pose = dw_func(i, frame, dwpose_model)
+ results_vis.append(pose)
+
+ logger.info(f'All frames have been processed.')
+ print(len(results_vis))
+
+ # Process the reference image
+ ref_frame = reference_image.squeeze(0).cpu().numpy()*255 # Convert to HWC format and numpy array
+ ref_frame = np.flip(ref_frame, axis=2)
+ pose_ref = dw_func(-1, ref_frame, dwpose_model)
+ # print(f'The content of the image is currently: {pose_ref} ')
+
+ bodies = results_vis[0]['bodies']
+ faces = results_vis[0]['faces']
+ hands = results_vis[0]['hands']
+ candidate = bodies['candidate']
+
+ ref_bodies = pose_ref['bodies']
+ ref_faces = pose_ref['faces']
+ ref_hands = pose_ref['hands']
+ ref_candidate = ref_bodies['candidate']
+
+
+ ref_2_x = ref_candidate[2][0]
+ ref_2_y = ref_candidate[2][1]
+ ref_5_x = ref_candidate[5][0]
+ ref_5_y = ref_candidate[5][1]
+ ref_8_x = ref_candidate[8][0]
+ ref_8_y = ref_candidate[8][1]
+ ref_11_x = ref_candidate[11][0]
+ ref_11_y = ref_candidate[11][1]
+ ref_center1 = 0.5*(ref_candidate[2]+ref_candidate[5])
+ ref_center2 = 0.5*(ref_candidate[8]+ref_candidate[11])
+
+ zero_2_x = candidate[2][0]
+ zero_2_y = candidate[2][1]
+ zero_5_x = candidate[5][0]
+ zero_5_y = candidate[5][1]
+ zero_8_x = candidate[8][0]
+ zero_8_y = candidate[8][1]
+ zero_11_x = candidate[11][0]
+ zero_11_y = candidate[11][1]
+ zero_center1 = 0.5*(candidate[2]+candidate[5])
+ zero_center2 = 0.5*(candidate[8]+candidate[11])
+
+ x_ratio, y_ratio = 1, 1
+ if (zero_5_x-zero_2_x) > 0 :
+ x_ratio = (ref_5_x-ref_2_x)/(zero_5_x-zero_2_x)
+ if (zero_center2[1]-zero_center1[1]) > 0 :
+ y_ratio = (ref_center2[1]-ref_center1[1])/(zero_center2[1]-zero_center1[1])
+
+ results_vis[0]['bodies']['candidate'][:,0] *= x_ratio
+ results_vis[0]['bodies']['candidate'][:,1] *= y_ratio
+ results_vis[0]['faces'][:,:,0] *= x_ratio
+ results_vis[0]['faces'][:,:,1] *= y_ratio
+ results_vis[0]['hands'][:,:,0] *= x_ratio
+ results_vis[0]['hands'][:,:,1] *= y_ratio
+
+ ########neck########
+ neck_ratio = 1
+ l_neck_ref = ((ref_candidate[0][0] - ref_candidate[1][0]) ** 2 + (ref_candidate[0][1] - ref_candidate[1][1]) ** 2) ** 0.5
+ l_neck_0 = ((candidate[0][0] - candidate[1][0]) ** 2 + (candidate[0][1] - candidate[1][1]) ** 2) ** 0.5
+ if l_neck_0 != 0:
+ neck_ratio = l_neck_ref / l_neck_0
+
+ x_offset_neck = (candidate[1][0]-candidate[0][0])*(1.-neck_ratio)
+ y_offset_neck = (candidate[1][1]-candidate[0][1])*(1.-neck_ratio)
+
+ results_vis[0]['bodies']['candidate'][0,0] += x_offset_neck
+ results_vis[0]['bodies']['candidate'][0,1] += y_offset_neck
+ results_vis[0]['bodies']['candidate'][14,0] += x_offset_neck
+ results_vis[0]['bodies']['candidate'][14,1] += y_offset_neck
+ results_vis[0]['bodies']['candidate'][15,0] += x_offset_neck
+ results_vis[0]['bodies']['candidate'][15,1] += y_offset_neck
+ results_vis[0]['bodies']['candidate'][16,0] += x_offset_neck
+ results_vis[0]['bodies']['candidate'][16,1] += y_offset_neck
+ results_vis[0]['bodies']['candidate'][17,0] += x_offset_neck
+ results_vis[0]['bodies']['candidate'][17,1] += y_offset_neck
+
+ ########shoulder2########
+ shoulder2_ratio = 1
+ l_shoulder2_ref = ((ref_candidate[2][0] - ref_candidate[1][0]) ** 2 + (ref_candidate[2][1] - ref_candidate[1][1]) ** 2) ** 0.5
+ l_shoulder2_0 = ((candidate[2][0] - candidate[1][0]) ** 2 + (candidate[2][1] - candidate[1][1]) ** 2) ** 0.5
+ if l_shoulder2_0 != 0:
+ shoulder2_ratio = l_shoulder2_ref / l_shoulder2_0
+
+ x_offset_shoulder2 = (candidate[1][0]-candidate[2][0])*(1.-shoulder2_ratio)
+ y_offset_shoulder2 = (candidate[1][1]-candidate[2][1])*(1.-shoulder2_ratio)
+
+ results_vis[0]['bodies']['candidate'][2,0] += x_offset_shoulder2
+ results_vis[0]['bodies']['candidate'][2,1] += y_offset_shoulder2
+ results_vis[0]['bodies']['candidate'][3,0] += x_offset_shoulder2
+ results_vis[0]['bodies']['candidate'][3,1] += y_offset_shoulder2
+ results_vis[0]['bodies']['candidate'][4,0] += x_offset_shoulder2
+ results_vis[0]['bodies']['candidate'][4,1] += y_offset_shoulder2
+ results_vis[0]['hands'][1,:,0] += x_offset_shoulder2
+ results_vis[0]['hands'][1,:,1] += y_offset_shoulder2
+
+ ########shoulder5########
+ shoulder5_ratio = 1
+ l_shoulder5_ref = ((ref_candidate[5][0] - ref_candidate[1][0]) ** 2 + (ref_candidate[5][1] - ref_candidate[1][1]) ** 2) ** 0.5
+ l_shoulder5_0 = ((candidate[5][0] - candidate[1][0]) ** 2 + (candidate[5][1] - candidate[1][1]) ** 2) ** 0.5
+ if l_shoulder5_0 != 0:
+ shoulder5_ratio = l_shoulder5_ref / l_shoulder5_0
+
+ x_offset_shoulder5 = (candidate[1][0]-candidate[5][0])*(1.-shoulder5_ratio)
+ y_offset_shoulder5 = (candidate[1][1]-candidate[5][1])*(1.-shoulder5_ratio)
+
+ results_vis[0]['bodies']['candidate'][5,0] += x_offset_shoulder5
+ results_vis[0]['bodies']['candidate'][5,1] += y_offset_shoulder5
+ results_vis[0]['bodies']['candidate'][6,0] += x_offset_shoulder5
+ results_vis[0]['bodies']['candidate'][6,1] += y_offset_shoulder5
+ results_vis[0]['bodies']['candidate'][7,0] += x_offset_shoulder5
+ results_vis[0]['bodies']['candidate'][7,1] += y_offset_shoulder5
+ results_vis[0]['hands'][0,:,0] += x_offset_shoulder5
+ results_vis[0]['hands'][0,:,1] += y_offset_shoulder5
+
+ ########arm3########
+ arm3_ratio = 1
+ l_arm3_ref = ((ref_candidate[3][0] - ref_candidate[2][0]) ** 2 + (ref_candidate[3][1] - ref_candidate[2][1]) ** 2) ** 0.5
+ l_arm3_0 = ((candidate[3][0] - candidate[2][0]) ** 2 + (candidate[3][1] - candidate[2][1]) ** 2) ** 0.5
+ if l_arm3_0 != 0:
+ arm3_ratio = l_arm3_ref / l_arm3_0
+
+ x_offset_arm3 = (candidate[2][0]-candidate[3][0])*(1.-arm3_ratio)
+ y_offset_arm3 = (candidate[2][1]-candidate[3][1])*(1.-arm3_ratio)
+
+ results_vis[0]['bodies']['candidate'][3,0] += x_offset_arm3
+ results_vis[0]['bodies']['candidate'][3,1] += y_offset_arm3
+ results_vis[0]['bodies']['candidate'][4,0] += x_offset_arm3
+ results_vis[0]['bodies']['candidate'][4,1] += y_offset_arm3
+ results_vis[0]['hands'][1,:,0] += x_offset_arm3
+ results_vis[0]['hands'][1,:,1] += y_offset_arm3
+
+ ########arm4########
+ arm4_ratio = 1
+ l_arm4_ref = ((ref_candidate[4][0] - ref_candidate[3][0]) ** 2 + (ref_candidate[4][1] - ref_candidate[3][1]) ** 2) ** 0.5
+ l_arm4_0 = ((candidate[4][0] - candidate[3][0]) ** 2 + (candidate[4][1] - candidate[3][1]) ** 2) ** 0.5
+ if l_arm4_0 != 0:
+ arm4_ratio = l_arm4_ref / l_arm4_0
+
+ x_offset_arm4 = (candidate[3][0]-candidate[4][0])*(1.-arm4_ratio)
+ y_offset_arm4 = (candidate[3][1]-candidate[4][1])*(1.-arm4_ratio)
+
+ results_vis[0]['bodies']['candidate'][4,0] += x_offset_arm4
+ results_vis[0]['bodies']['candidate'][4,1] += y_offset_arm4
+ results_vis[0]['hands'][1,:,0] += x_offset_arm4
+ results_vis[0]['hands'][1,:,1] += y_offset_arm4
+
+ ########arm6########
+ arm6_ratio = 1
+ l_arm6_ref = ((ref_candidate[6][0] - ref_candidate[5][0]) ** 2 + (ref_candidate[6][1] - ref_candidate[5][1]) ** 2) ** 0.5
+ l_arm6_0 = ((candidate[6][0] - candidate[5][0]) ** 2 + (candidate[6][1] - candidate[5][1]) ** 2) ** 0.5
+ if l_arm6_0 != 0:
+ arm6_ratio = l_arm6_ref / l_arm6_0
+
+ x_offset_arm6 = (candidate[5][0]-candidate[6][0])*(1.-arm6_ratio)
+ y_offset_arm6 = (candidate[5][1]-candidate[6][1])*(1.-arm6_ratio)
+
+ results_vis[0]['bodies']['candidate'][6,0] += x_offset_arm6
+ results_vis[0]['bodies']['candidate'][6,1] += y_offset_arm6
+ results_vis[0]['bodies']['candidate'][7,0] += x_offset_arm6
+ results_vis[0]['bodies']['candidate'][7,1] += y_offset_arm6
+ results_vis[0]['hands'][0,:,0] += x_offset_arm6
+ results_vis[0]['hands'][0,:,1] += y_offset_arm6
+
+ ########arm7########
+ arm7_ratio = 1
+ l_arm7_ref = ((ref_candidate[7][0] - ref_candidate[6][0]) ** 2 + (ref_candidate[7][1] - ref_candidate[6][1]) ** 2) ** 0.5
+ l_arm7_0 = ((candidate[7][0] - candidate[6][0]) ** 2 + (candidate[7][1] - candidate[6][1]) ** 2) ** 0.5
+ if l_arm7_0 != 0:
+ arm7_ratio = l_arm7_ref / l_arm7_0
+
+ x_offset_arm7 = (candidate[6][0]-candidate[7][0])*(1.-arm7_ratio)
+ y_offset_arm7 = (candidate[6][1]-candidate[7][1])*(1.-arm7_ratio)
+
+ results_vis[0]['bodies']['candidate'][7,0] += x_offset_arm7
+ results_vis[0]['bodies']['candidate'][7,1] += y_offset_arm7
+ results_vis[0]['hands'][0,:,0] += x_offset_arm7
+ results_vis[0]['hands'][0,:,1] += y_offset_arm7
+
+ ########head14########
+ l_head14_ref = ((ref_candidate[14][0] - ref_candidate[0][0]) ** 2 + (ref_candidate[14][1] - ref_candidate[0][1]) ** 2) ** 0.5
+ l_head14_0 = ((candidate[14][0] - candidate[0][0]) ** 2 + (candidate[14][1] - candidate[0][1]) ** 2) ** 0.5
+
+ head14_ratio = l_head14_ref / l_head14_0
+
+ x_offset_head14 = (candidate[0][0]-candidate[14][0])*(1.-head14_ratio)
+ y_offset_head14 = (candidate[0][1]-candidate[14][1])*(1.-head14_ratio)
+
+ results_vis[0]['bodies']['candidate'][14,0] += x_offset_head14
+ results_vis[0]['bodies']['candidate'][14,1] += y_offset_head14
+ results_vis[0]['bodies']['candidate'][16,0] += x_offset_head14
+ results_vis[0]['bodies']['candidate'][16,1] += y_offset_head14
+
+ ########head15########
+ l_head15_ref = ((ref_candidate[15][0] - ref_candidate[0][0]) ** 2 + (ref_candidate[15][1] - ref_candidate[0][1]) ** 2) ** 0.5
+ l_head15_0 = ((candidate[15][0] - candidate[0][0]) ** 2 + (candidate[15][1] - candidate[0][1]) ** 2) ** 0.5
+
+ head15_ratio = l_head15_ref / l_head15_0
+
+ x_offset_head15 = (candidate[0][0]-candidate[15][0])*(1.-head15_ratio)
+ y_offset_head15 = (candidate[0][1]-candidate[15][1])*(1.-head15_ratio)
+
+ results_vis[0]['bodies']['candidate'][15,0] += x_offset_head15
+ results_vis[0]['bodies']['candidate'][15,1] += y_offset_head15
+ results_vis[0]['bodies']['candidate'][17,0] += x_offset_head15
+ results_vis[0]['bodies']['candidate'][17,1] += y_offset_head15
+
+ ########head16########
+ l_head16_ref = ((ref_candidate[16][0] - ref_candidate[14][0]) ** 2 + (ref_candidate[16][1] - ref_candidate[14][1]) ** 2) ** 0.5
+ l_head16_0 = ((candidate[16][0] - candidate[14][0]) ** 2 + (candidate[16][1] - candidate[14][1]) ** 2) ** 0.5
+
+ head16_ratio = l_head16_ref / l_head16_0
+
+ x_offset_head16 = (candidate[14][0]-candidate[16][0])*(1.-head16_ratio)
+ y_offset_head16 = (candidate[14][1]-candidate[16][1])*(1.-head16_ratio)
+
+ results_vis[0]['bodies']['candidate'][16,0] += x_offset_head16
+ results_vis[0]['bodies']['candidate'][16,1] += y_offset_head16
+
+ ########head17########
+ l_head17_ref = ((ref_candidate[17][0] - ref_candidate[15][0]) ** 2 + (ref_candidate[17][1] - ref_candidate[15][1]) ** 2) ** 0.5
+ l_head17_0 = ((candidate[17][0] - candidate[15][0]) ** 2 + (candidate[17][1] - candidate[15][1]) ** 2) ** 0.5
+
+ head17_ratio = l_head17_ref / l_head17_0
+
+ x_offset_head17 = (candidate[15][0]-candidate[17][0])*(1.-head17_ratio)
+ y_offset_head17 = (candidate[15][1]-candidate[17][1])*(1.-head17_ratio)
+
+ results_vis[0]['bodies']['candidate'][17,0] += x_offset_head17
+ results_vis[0]['bodies']['candidate'][17,1] += y_offset_head17
+
+ ########MovingAverage########
+
+ ########left leg########
+ ll1_ratio = 1
+ l_ll1_ref = ((ref_candidate[8][0] - ref_candidate[9][0]) ** 2 + (ref_candidate[8][1] - ref_candidate[9][1]) ** 2) ** 0.5
+ l_ll1_0 = ((candidate[8][0] - candidate[9][0]) ** 2 + (candidate[8][1] - candidate[9][1]) ** 2) ** 0.5
+ if l_ll1_0 != 0 :
+ ll1_ratio = l_ll1_ref / l_ll1_0
+
+ x_offset_ll1 = (candidate[9][0]-candidate[8][0])*(ll1_ratio-1.)
+ y_offset_ll1 = (candidate[9][1]-candidate[8][1])*(ll1_ratio-1.)
+
+ results_vis[0]['bodies']['candidate'][9,0] += x_offset_ll1
+ results_vis[0]['bodies']['candidate'][9,1] += y_offset_ll1
+ results_vis[0]['bodies']['candidate'][10,0] += x_offset_ll1
+ results_vis[0]['bodies']['candidate'][10,1] += y_offset_ll1
+ results_vis[0]['bodies']['candidate'][19,0] += x_offset_ll1
+ results_vis[0]['bodies']['candidate'][19,1] += y_offset_ll1
+
+ l_ll2_ref = ((ref_candidate[9][0] - ref_candidate[10][0]) ** 2 + (ref_candidate[9][1] - ref_candidate[10][1]) ** 2) ** 0.5
+ l_ll2_0 = ((candidate[9][0] - candidate[10][0]) ** 2 + (candidate[9][1] - candidate[10][1]) ** 2) ** 0.5
+ ll2_ratio = l_ll2_ref / l_ll2_0
+
+ x_offset_ll2 = (candidate[10][0]-candidate[9][0])*(ll2_ratio-1.)
+ y_offset_ll2 = (candidate[10][1]-candidate[9][1])*(ll2_ratio-1.)
+
+ results_vis[0]['bodies']['candidate'][10,0] += x_offset_ll2
+ results_vis[0]['bodies']['candidate'][10,1] += y_offset_ll2
+ results_vis[0]['bodies']['candidate'][19,0] += x_offset_ll2
+ results_vis[0]['bodies']['candidate'][19,1] += y_offset_ll2
+
+ ########right leg########
+ rl1_ratio = 1
+ l_rl1_ref = ((ref_candidate[11][0] - ref_candidate[12][0]) ** 2 + (ref_candidate[11][1] - ref_candidate[12][1]) ** 2) ** 0.5
+ l_rl1_0 = ((candidate[11][0] - candidate[12][0]) ** 2 + (candidate[11][1] - candidate[12][1]) ** 2) ** 0.5
+ if l_rl1_0 != 0:
+ rl1_ratio = l_rl1_ref / l_rl1_0
+
+ x_offset_rl1 = (candidate[12][0]-candidate[11][0])*(rl1_ratio-1.)
+ y_offset_rl1 = (candidate[12][1]-candidate[11][1])*(rl1_ratio-1.)
+
+ results_vis[0]['bodies']['candidate'][12,0] += x_offset_rl1
+ results_vis[0]['bodies']['candidate'][12,1] += y_offset_rl1
+ results_vis[0]['bodies']['candidate'][13,0] += x_offset_rl1
+ results_vis[0]['bodies']['candidate'][13,1] += y_offset_rl1
+ results_vis[0]['bodies']['candidate'][18,0] += x_offset_rl1
+ results_vis[0]['bodies']['candidate'][18,1] += y_offset_rl1
+
+ l_rl2_ref = ((ref_candidate[12][0] - ref_candidate[13][0]) ** 2 + (ref_candidate[12][1] - ref_candidate[13][1]) ** 2) ** 0.5
+ l_rl2_0 = ((candidate[12][0] - candidate[13][0]) ** 2 + (candidate[12][1] - candidate[13][1]) ** 2) ** 0.5
+ rl2_ratio = l_rl2_ref / l_rl2_0
+
+ x_offset_rl2 = (candidate[13][0]-candidate[12][0])*(rl2_ratio-1.)
+ y_offset_rl2 = (candidate[13][1]-candidate[12][1])*(rl2_ratio-1.)
+
+ results_vis[0]['bodies']['candidate'][13,0] += x_offset_rl2
+ results_vis[0]['bodies']['candidate'][13,1] += y_offset_rl2
+ results_vis[0]['bodies']['candidate'][18,0] += x_offset_rl2
+ results_vis[0]['bodies']['candidate'][18,1] += y_offset_rl2
+
+ offset = ref_candidate[1] - results_vis[0]['bodies']['candidate'][1]
+
+ results_vis[0]['bodies']['candidate'] += offset[np.newaxis, :]
+ results_vis[0]['faces'] += offset[np.newaxis, np.newaxis, :]
+ results_vis[0]['hands'] += offset[np.newaxis, np.newaxis, :]
+
+ for i in range(1, len(results_vis)):
+ results_vis[i]['bodies']['candidate'][:,0] *= x_ratio
+ results_vis[i]['bodies']['candidate'][:,1] *= y_ratio
+ results_vis[i]['faces'][:,:,0] *= x_ratio
+ results_vis[i]['faces'][:,:,1] *= y_ratio
+ results_vis[i]['hands'][:,:,0] *= x_ratio
+ results_vis[i]['hands'][:,:,1] *= y_ratio
+
+ ########neck########
+ x_offset_neck = (results_vis[i]['bodies']['candidate'][1][0]-results_vis[i]['bodies']['candidate'][0][0])*(1.-neck_ratio)
+ y_offset_neck = (results_vis[i]['bodies']['candidate'][1][1]-results_vis[i]['bodies']['candidate'][0][1])*(1.-neck_ratio)
+
+ results_vis[i]['bodies']['candidate'][0,0] += x_offset_neck
+ results_vis[i]['bodies']['candidate'][0,1] += y_offset_neck
+ results_vis[i]['bodies']['candidate'][14,0] += x_offset_neck
+ results_vis[i]['bodies']['candidate'][14,1] += y_offset_neck
+ results_vis[i]['bodies']['candidate'][15,0] += x_offset_neck
+ results_vis[i]['bodies']['candidate'][15,1] += y_offset_neck
+ results_vis[i]['bodies']['candidate'][16,0] += x_offset_neck
+ results_vis[i]['bodies']['candidate'][16,1] += y_offset_neck
+ results_vis[i]['bodies']['candidate'][17,0] += x_offset_neck
+ results_vis[i]['bodies']['candidate'][17,1] += y_offset_neck
+
+ ########shoulder2########
+
+
+ x_offset_shoulder2 = (results_vis[i]['bodies']['candidate'][1][0]-results_vis[i]['bodies']['candidate'][2][0])*(1.-shoulder2_ratio)
+ y_offset_shoulder2 = (results_vis[i]['bodies']['candidate'][1][1]-results_vis[i]['bodies']['candidate'][2][1])*(1.-shoulder2_ratio)
+
+ results_vis[i]['bodies']['candidate'][2,0] += x_offset_shoulder2
+ results_vis[i]['bodies']['candidate'][2,1] += y_offset_shoulder2
+ results_vis[i]['bodies']['candidate'][3,0] += x_offset_shoulder2
+ results_vis[i]['bodies']['candidate'][3,1] += y_offset_shoulder2
+ results_vis[i]['bodies']['candidate'][4,0] += x_offset_shoulder2
+ results_vis[i]['bodies']['candidate'][4,1] += y_offset_shoulder2
+ results_vis[i]['hands'][1,:,0] += x_offset_shoulder2
+ results_vis[i]['hands'][1,:,1] += y_offset_shoulder2
+
+ ########shoulder5########
+
+ x_offset_shoulder5 = (results_vis[i]['bodies']['candidate'][1][0]-results_vis[i]['bodies']['candidate'][5][0])*(1.-shoulder5_ratio)
+ y_offset_shoulder5 = (results_vis[i]['bodies']['candidate'][1][1]-results_vis[i]['bodies']['candidate'][5][1])*(1.-shoulder5_ratio)
+
+ results_vis[i]['bodies']['candidate'][5,0] += x_offset_shoulder5
+ results_vis[i]['bodies']['candidate'][5,1] += y_offset_shoulder5
+ results_vis[i]['bodies']['candidate'][6,0] += x_offset_shoulder5
+ results_vis[i]['bodies']['candidate'][6,1] += y_offset_shoulder5
+ results_vis[i]['bodies']['candidate'][7,0] += x_offset_shoulder5
+ results_vis[i]['bodies']['candidate'][7,1] += y_offset_shoulder5
+ results_vis[i]['hands'][0,:,0] += x_offset_shoulder5
+ results_vis[i]['hands'][0,:,1] += y_offset_shoulder5
+
+ ########arm3########
+
+ x_offset_arm3 = (results_vis[i]['bodies']['candidate'][2][0]-results_vis[i]['bodies']['candidate'][3][0])*(1.-arm3_ratio)
+ y_offset_arm3 = (results_vis[i]['bodies']['candidate'][2][1]-results_vis[i]['bodies']['candidate'][3][1])*(1.-arm3_ratio)
+
+ results_vis[i]['bodies']['candidate'][3,0] += x_offset_arm3
+ results_vis[i]['bodies']['candidate'][3,1] += y_offset_arm3
+ results_vis[i]['bodies']['candidate'][4,0] += x_offset_arm3
+ results_vis[i]['bodies']['candidate'][4,1] += y_offset_arm3
+ results_vis[i]['hands'][1,:,0] += x_offset_arm3
+ results_vis[i]['hands'][1,:,1] += y_offset_arm3
+
+ ########arm4########
+
+ x_offset_arm4 = (results_vis[i]['bodies']['candidate'][3][0]-results_vis[i]['bodies']['candidate'][4][0])*(1.-arm4_ratio)
+ y_offset_arm4 = (results_vis[i]['bodies']['candidate'][3][1]-results_vis[i]['bodies']['candidate'][4][1])*(1.-arm4_ratio)
+
+ results_vis[i]['bodies']['candidate'][4,0] += x_offset_arm4
+ results_vis[i]['bodies']['candidate'][4,1] += y_offset_arm4
+ results_vis[i]['hands'][1,:,0] += x_offset_arm4
+ results_vis[i]['hands'][1,:,1] += y_offset_arm4
+
+ ########arm6########
+
+ x_offset_arm6 = (results_vis[i]['bodies']['candidate'][5][0]-results_vis[i]['bodies']['candidate'][6][0])*(1.-arm6_ratio)
+ y_offset_arm6 = (results_vis[i]['bodies']['candidate'][5][1]-results_vis[i]['bodies']['candidate'][6][1])*(1.-arm6_ratio)
+
+ results_vis[i]['bodies']['candidate'][6,0] += x_offset_arm6
+ results_vis[i]['bodies']['candidate'][6,1] += y_offset_arm6
+ results_vis[i]['bodies']['candidate'][7,0] += x_offset_arm6
+ results_vis[i]['bodies']['candidate'][7,1] += y_offset_arm6
+ results_vis[i]['hands'][0,:,0] += x_offset_arm6
+ results_vis[i]['hands'][0,:,1] += y_offset_arm6
+
+ ########arm7########
+
+ x_offset_arm7 = (results_vis[i]['bodies']['candidate'][6][0]-results_vis[i]['bodies']['candidate'][7][0])*(1.-arm7_ratio)
+ y_offset_arm7 = (results_vis[i]['bodies']['candidate'][6][1]-results_vis[i]['bodies']['candidate'][7][1])*(1.-arm7_ratio)
+
+ results_vis[i]['bodies']['candidate'][7,0] += x_offset_arm7
+ results_vis[i]['bodies']['candidate'][7,1] += y_offset_arm7
+ results_vis[i]['hands'][0,:,0] += x_offset_arm7
+ results_vis[i]['hands'][0,:,1] += y_offset_arm7
+
+ ########head14########
+
+ x_offset_head14 = (results_vis[i]['bodies']['candidate'][0][0]-results_vis[i]['bodies']['candidate'][14][0])*(1.-head14_ratio)
+ y_offset_head14 = (results_vis[i]['bodies']['candidate'][0][1]-results_vis[i]['bodies']['candidate'][14][1])*(1.-head14_ratio)
+
+ results_vis[i]['bodies']['candidate'][14,0] += x_offset_head14
+ results_vis[i]['bodies']['candidate'][14,1] += y_offset_head14
+ results_vis[i]['bodies']['candidate'][16,0] += x_offset_head14
+ results_vis[i]['bodies']['candidate'][16,1] += y_offset_head14
+
+ ########head15########
+
+ x_offset_head15 = (results_vis[i]['bodies']['candidate'][0][0]-results_vis[i]['bodies']['candidate'][15][0])*(1.-head15_ratio)
+ y_offset_head15 = (results_vis[i]['bodies']['candidate'][0][1]-results_vis[i]['bodies']['candidate'][15][1])*(1.-head15_ratio)
+
+ results_vis[i]['bodies']['candidate'][15,0] += x_offset_head15
+ results_vis[i]['bodies']['candidate'][15,1] += y_offset_head15
+ results_vis[i]['bodies']['candidate'][17,0] += x_offset_head15
+ results_vis[i]['bodies']['candidate'][17,1] += y_offset_head15
+
+ ########head16########
+
+ x_offset_head16 = (results_vis[i]['bodies']['candidate'][14][0]-results_vis[i]['bodies']['candidate'][16][0])*(1.-head16_ratio)
+ y_offset_head16 = (results_vis[i]['bodies']['candidate'][14][1]-results_vis[i]['bodies']['candidate'][16][1])*(1.-head16_ratio)
+
+ results_vis[i]['bodies']['candidate'][16,0] += x_offset_head16
+ results_vis[i]['bodies']['candidate'][16,1] += y_offset_head16
+
+ ########head17########
+ x_offset_head17 = (results_vis[i]['bodies']['candidate'][15][0]-results_vis[i]['bodies']['candidate'][17][0])*(1.-head17_ratio)
+ y_offset_head17 = (results_vis[i]['bodies']['candidate'][15][1]-results_vis[i]['bodies']['candidate'][17][1])*(1.-head17_ratio)
+
+ results_vis[i]['bodies']['candidate'][17,0] += x_offset_head17
+ results_vis[i]['bodies']['candidate'][17,1] += y_offset_head17
+
+ # ########MovingAverage########
+
+ ########left leg########
+ x_offset_ll1 = (results_vis[i]['bodies']['candidate'][9][0]-results_vis[i]['bodies']['candidate'][8][0])*(ll1_ratio-1.)
+ y_offset_ll1 = (results_vis[i]['bodies']['candidate'][9][1]-results_vis[i]['bodies']['candidate'][8][1])*(ll1_ratio-1.)
+
+ results_vis[i]['bodies']['candidate'][9,0] += x_offset_ll1
+ results_vis[i]['bodies']['candidate'][9,1] += y_offset_ll1
+ results_vis[i]['bodies']['candidate'][10,0] += x_offset_ll1
+ results_vis[i]['bodies']['candidate'][10,1] += y_offset_ll1
+ results_vis[i]['bodies']['candidate'][19,0] += x_offset_ll1
+ results_vis[i]['bodies']['candidate'][19,1] += y_offset_ll1
+
+
+
+ x_offset_ll2 = (results_vis[i]['bodies']['candidate'][10][0]-results_vis[i]['bodies']['candidate'][9][0])*(ll2_ratio-1.)
+ y_offset_ll2 = (results_vis[i]['bodies']['candidate'][10][1]-results_vis[i]['bodies']['candidate'][9][1])*(ll2_ratio-1.)
+
+ results_vis[i]['bodies']['candidate'][10,0] += x_offset_ll2
+ results_vis[i]['bodies']['candidate'][10,1] += y_offset_ll2
+ results_vis[i]['bodies']['candidate'][19,0] += x_offset_ll2
+ results_vis[i]['bodies']['candidate'][19,1] += y_offset_ll2
+
+ ########right leg########
+
+ x_offset_rl1 = (results_vis[i]['bodies']['candidate'][12][0]-results_vis[i]['bodies']['candidate'][11][0])*(rl1_ratio-1.)
+ y_offset_rl1 = (results_vis[i]['bodies']['candidate'][12][1]-results_vis[i]['bodies']['candidate'][11][1])*(rl1_ratio-1.)
+
+ results_vis[i]['bodies']['candidate'][12,0] += x_offset_rl1
+ results_vis[i]['bodies']['candidate'][12,1] += y_offset_rl1
+ results_vis[i]['bodies']['candidate'][13,0] += x_offset_rl1
+ results_vis[i]['bodies']['candidate'][13,1] += y_offset_rl1
+ results_vis[i]['bodies']['candidate'][18,0] += x_offset_rl1
+ results_vis[i]['bodies']['candidate'][18,1] += y_offset_rl1
+
+
+ x_offset_rl2 = (results_vis[i]['bodies']['candidate'][13][0]-results_vis[i]['bodies']['candidate'][12][0])*(rl2_ratio-1.)
+ y_offset_rl2 = (results_vis[i]['bodies']['candidate'][13][1]-results_vis[i]['bodies']['candidate'][12][1])*(rl2_ratio-1.)
+
+ results_vis[i]['bodies']['candidate'][13,0] += x_offset_rl2
+ results_vis[i]['bodies']['candidate'][13,1] += y_offset_rl2
+ results_vis[i]['bodies']['candidate'][18,0] += x_offset_rl2
+ results_vis[i]['bodies']['candidate'][18,1] += y_offset_rl2
+
+ results_vis[i]['bodies']['candidate'] += offset[np.newaxis, :]
+ results_vis[i]['faces'] += offset[np.newaxis, np.newaxis, :]
+ results_vis[i]['hands'] += offset[np.newaxis, np.newaxis, :]
+
+ # Prepare to return the dwpose images
+ dwpose_images = []
+
+ for i in range(len(results_vis)):
+ dwpose_woface, _ = draw_pose(results_vis[i], H=768, W=512)
+ dwpose_tensor = torch.from_numpy(dwpose_woface).permute(2, 0, 1).unsqueeze(0).float() # Convert to tensor and CHW format
+ dwpose_tensor = dwpose_tensor.permute(0, 2, 3, 1)
+ dwpose_images.append(dwpose_tensor)
+
+ dwpose_images = torch.cat(dwpose_images, dim=0)
+ dwpose_woface_ref, _ = draw_pose(pose_ref, H=768, W=512)
+ dwpose_ref_tensor = torch.from_numpy(dwpose_woface_ref).permute(2, 0, 1).unsqueeze(0).float() # Convert to tensor and CHW format
+ dwpose_ref_tensor = dwpose_ref_tensor.permute(0, 2, 3, 1)
+
+ # print(f'The type of the pose from run_align_pose is currently of the form : {type(dwpose_ref_tensor)} ')
+ # print(f'The content of the pose from run_align_pose is currently: {dwpose_ref_tensor} ')
+
+ return dwpose_images, dwpose_ref_tensor
+
+
+logger = get_logger('dw pose extraction')
+
diff --git a/uniAnimate_Inference.py b/uniAnimate_Inference.py
new file mode 100644
index 0000000..db10817
--- /dev/null
+++ b/uniAnimate_Inference.py
@@ -0,0 +1,83 @@
+import torch
+
+from .utils.config import Config
+from .tools.inferences import inference_unianimate_entrance
+from . import run_align_pose
+
+# from tools import *
+
+
+
+class UniAnimateImage:
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "steps": ("INT", {"default": 30, "min": 25, "max": 50, "step": 1}),
+ "useFirstFrame": ("BOOLEAN", { "default": False }),
+ "reference_image": ("IMAGE",), # single image
+ "ref_pose": ("IMAGE",), # single image
+ "pose_sequence": ("IMAGE",), # Batch of pose images
+ "frame_interval": ("INT", {"default": 1, "min": 1, "max": 8, "step": 1}),
+ "max_frames": ("INT", {"default": 32, "min": 1, "max": 64, "step": 1}),
+ "resolution_x": ("INT", {"default": 512, "min": 512, "max": 768, "step": 256}),
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE", "MASK")
+ FUNCTION = "process"
+ CATEGORY = "image"
+
+ def process(self, steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution_x):
+ cfg_update = Config(load=True)
+ resolution_y = 768
+ if resolution_x == 768:
+ resolution_y = 1216
+ resolution = [resolution_x, resolution_y]
+ print("Ready for inference.")
+
+ # print(f"image is: {reference_image}")
+
+ frames = inference_unianimate_entrance(steps, useFirstFrame, reference_image, ref_pose, pose_sequence, frame_interval, max_frames, resolution, cfg_update=cfg_update.cfg_dict)
+ mask_template = torch.zeros((1, resolution_y, resolution_x), dtype=torch.float32)
+ masks = [mask_template.clone() for _ in range(len(pose_sequence))]
+ masks = torch.cat(masks, dim=0)
+ return (frames, masks)
+
+class Gen_align_pose:
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "reference_image": ("IMAGE",), # single image
+ "video": ("IMAGE",), # video
+ }
+ }
+
+ RETURN_TYPES = ("IMAGE", "IMAGE")
+ FUNCTION = "process"
+ CATEGORY = "image"
+
+ def process(self, reference_image, video):
+ if torch.cuda.is_available():
+ print(f"CUDA version: {torch.version.cuda}")
+ print(f"CUDNN version: {torch.backends.cudnn.version()}")
+ print(f"Device name: {torch.cuda.get_device_name(0)}")
+ else:
+ print("CUDA is not available")
+ poses, refPose = run_align_pose.mp_main(reference_image, video)
+ return (refPose, poses)
+
+
+
+NODE_CLASS_MAPPINGS = {
+ "UniAnimateImage" : UniAnimateImage,
+ "Gen_align_pose" : Gen_align_pose,
+
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+ "UniAnimateImage" :"Animate image with UniAnimate",
+ "Gen_align_pose" :"Align & Generate poses for UniAnimate",
+
+}
\ No newline at end of file