Spaces:
Running on Zero
Running on Zero
github-actions[bot] commited on
Commit ·
cf6e6cb
0
Parent(s):
Deploy DoppelGen compiled C-extension binary distribution to Hugging Face Space
Browse filesThis view is limited to 50 files because it contains too many changes. See raw diff
- .gitattributes +7 -0
- .github/scripts/compile_binaries.py +144 -0
- .github/workflows/deploy_hf.yml +77 -0
- .gitignore +31 -0
- README.md +151 -0
- actora/.dockerignore +8 -0
- actora/.gitignore +12 -0
- actora/README.md +253 -0
- actora/actora/__init__.py +0 -0
- actora/actora/api/__init__.py +1 -0
- actora/actora/api/api.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/api/main.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/core/__init__.py +1 -0
- actora/actora/core/arcanegan.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/core/dctnet.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/core/filters.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/core/matting.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/core/neural_accelerator.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/core/postprocess.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/__init__.py +1 -0
- actora/actora/models/base.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/dreamtalk.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/faster_liveportrait.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/fomm.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/generator.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/models.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/models/tpsmm.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/schema/__init__.py +0 -0
- actora/actora/test/create_test_inputs.py +53 -0
- actora/actora/test/test.py +35 -0
- actora/actora/utils/__init__.py +0 -0
- actora/actora/utils/utils.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/actora/utils/warnings_patch.cpython-312-x86_64-linux-gnu.so +3 -0
- actora/docker/Dockerfile +14 -0
- actora/docker/docker-compose.yml +13 -0
- actora/docker/entrypoint.sh +3 -0
- actora/docker/requirements.txt +27 -0
- actora/preload_models.py +496 -0
- actora/run_actora.sh +28 -0
- actora/third_party/dreamtalk_src/LICENSE +21 -0
- actora/third_party/dreamtalk_src/README.md +99 -0
- actora/third_party/dreamtalk_src/configs/__init__.py +0 -0
- actora/third_party/dreamtalk_src/configs/default.py +91 -0
- actora/third_party/dreamtalk_src/core/__init__.py +0 -0
- actora/third_party/dreamtalk_src/core/networks/__init__.py +14 -0
- actora/third_party/dreamtalk_src/core/networks/diffusion_net.py +340 -0
- actora/third_party/dreamtalk_src/core/networks/diffusion_util.py +131 -0
- actora/third_party/dreamtalk_src/core/networks/disentangle_decoder.py +240 -0
- actora/third_party/dreamtalk_src/core/networks/dynamic_conv.py +156 -0
- actora/third_party/dreamtalk_src/core/networks/dynamic_fc_decoder.py +178 -0
.gitattributes
ADDED
|
@@ -0,0 +1,7 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.so filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.mp4 filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.webm filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
doppelgen/data/jobs/test/input/** filter=lfs diff=lfs merge=lfs -text
|
.github/scripts/compile_binaries.py
ADDED
|
@@ -0,0 +1,144 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/usr/bin/env python3
|
| 2 |
+
"""
|
| 3 |
+
DoppelGen Hardened Native Compiler Script
|
| 4 |
+
Compiles core Python packages into Cython native C-extensions (.so)
|
| 5 |
+
with fallback to bytecode (.pyc) for dynamic modules, and strips raw .py source files.
|
| 6 |
+
"""
|
| 7 |
+
|
| 8 |
+
import os
|
| 9 |
+
import sys
|
| 10 |
+
import shutil
|
| 11 |
+
import py_compile
|
| 12 |
+
import subprocess
|
| 13 |
+
from setuptools import setup, Extension
|
| 14 |
+
from Cython.Build import cythonize
|
| 15 |
+
|
| 16 |
+
def main():
|
| 17 |
+
print("=" * 60)
|
| 18 |
+
print("🚀 Starting DoppelGen Cython Native Binary Compilation (.so)...")
|
| 19 |
+
print("=" * 60)
|
| 20 |
+
|
| 21 |
+
base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
| 22 |
+
os.chdir(base_dir)
|
| 23 |
+
|
| 24 |
+
target_packages = ["doppelgen", "actora", "voxa", "sonora", "scenea", "captiona", "ocula"]
|
| 25 |
+
|
| 26 |
+
compiled_so_count = 0
|
| 27 |
+
compiled_pyc_count = 0
|
| 28 |
+
deleted_count = 0
|
| 29 |
+
|
| 30 |
+
for pkg in target_packages:
|
| 31 |
+
pkg_dir = os.path.join(base_dir, pkg)
|
| 32 |
+
if not os.path.exists(pkg_dir):
|
| 33 |
+
continue
|
| 34 |
+
|
| 35 |
+
os.chdir(pkg_dir)
|
| 36 |
+
|
| 37 |
+
py_files_in_pkg = []
|
| 38 |
+
for root, dirs, files in os.walk(pkg_dir):
|
| 39 |
+
if "third_party" in root or "hf_cache" in root or ".cache" in root or "checkpoints" in root:
|
| 40 |
+
continue
|
| 41 |
+
if os.path.basename(root) in ["runners", "test", "tests", "docker", "build", "__pycache__"]:
|
| 42 |
+
continue
|
| 43 |
+
for file in files:
|
| 44 |
+
if file.endswith(".py") and not file.startswith("."):
|
| 45 |
+
if file in ["__init__.py", "preload_models.py", "app.py", "setup.py"]:
|
| 46 |
+
continue
|
| 47 |
+
full_path = os.path.join(root, file)
|
| 48 |
+
rel_to_pkg = os.path.relpath(full_path, pkg_dir)
|
| 49 |
+
module_name = rel_to_pkg[:-3].replace(os.sep, ".")
|
| 50 |
+
py_files_in_pkg.append((full_path, rel_to_pkg, module_name))
|
| 51 |
+
|
| 52 |
+
print(f"📦 Compiling {len(py_files_in_pkg)} modules in {pkg}...")
|
| 53 |
+
|
| 54 |
+
for full_path, rel_to_pkg, module_name in py_files_in_pkg:
|
| 55 |
+
ext = Extension(
|
| 56 |
+
module_name,
|
| 57 |
+
sources=[rel_to_pkg],
|
| 58 |
+
extra_compile_args=["-O3", "-fPIC", "-Wno-unused-variable"],
|
| 59 |
+
)
|
| 60 |
+
|
| 61 |
+
success = False
|
| 62 |
+
try:
|
| 63 |
+
setup(
|
| 64 |
+
ext_modules=cythonize(
|
| 65 |
+
[ext],
|
| 66 |
+
compiler_directives={
|
| 67 |
+
'language_level': "3",
|
| 68 |
+
'always_allow_keywords': True,
|
| 69 |
+
'embedsignature': False,
|
| 70 |
+
'annotation_typing': False,
|
| 71 |
+
},
|
| 72 |
+
quiet=True,
|
| 73 |
+
),
|
| 74 |
+
script_args=["build_ext", "--inplace"],
|
| 75 |
+
)
|
| 76 |
+
dir_name = os.path.dirname(rel_to_pkg)
|
| 77 |
+
base_name = os.path.basename(rel_to_pkg)[:-3]
|
| 78 |
+
search_dir = os.path.join(pkg_dir, dir_name) if dir_name else pkg_dir
|
| 79 |
+
matching_so = [f for f in os.listdir(search_dir) if f.startswith(base_name) and f.endswith(".so")]
|
| 80 |
+
if matching_so:
|
| 81 |
+
success = True
|
| 82 |
+
compiled_so_count += 1
|
| 83 |
+
except Exception as e:
|
| 84 |
+
print(f"⚠️ Cython compilation fallback to .pyc for {rel_to_pkg}: {e}")
|
| 85 |
+
|
| 86 |
+
if not success:
|
| 87 |
+
pyc_path = full_path + "c"
|
| 88 |
+
try:
|
| 89 |
+
py_compile.compile(full_path, cfile=pyc_path, doraise=True, optimize=2)
|
| 90 |
+
compiled_pyc_count += 1
|
| 91 |
+
except Exception as pe:
|
| 92 |
+
print(f"❌ Bytecode compilation failed for {rel_to_pkg}: {pe}")
|
| 93 |
+
|
| 94 |
+
for full_path, rel_to_pkg, module_name in py_files_in_pkg:
|
| 95 |
+
dir_name = os.path.dirname(rel_to_pkg)
|
| 96 |
+
base_name = os.path.basename(rel_to_pkg)[:-3]
|
| 97 |
+
search_dir = os.path.join(pkg_dir, dir_name) if dir_name else pkg_dir
|
| 98 |
+
matching_so = [f for f in os.listdir(search_dir) if f.startswith(base_name) and f.endswith(".so")] if os.path.exists(search_dir) else []
|
| 99 |
+
if matching_so:
|
| 100 |
+
if os.path.exists(full_path):
|
| 101 |
+
os.remove(full_path)
|
| 102 |
+
deleted_count += 1
|
| 103 |
+
else:
|
| 104 |
+
print(f"ℹ️ Preserved {rel_to_pkg} as .py source (no .so generated)", flush=True)
|
| 105 |
+
|
| 106 |
+
os.chdir(base_dir)
|
| 107 |
+
|
| 108 |
+
# Strip debugging symbols from generated .so files
|
| 109 |
+
for root, dirs, files in os.walk(base_dir):
|
| 110 |
+
if ".git" in root or "venv" in root or ".venv" in root:
|
| 111 |
+
continue
|
| 112 |
+
for file in files:
|
| 113 |
+
if file.endswith(".so"):
|
| 114 |
+
so_path = os.path.join(root, file)
|
| 115 |
+
try:
|
| 116 |
+
subprocess.run(["strip", "--strip-debug", so_path], check=False)
|
| 117 |
+
except Exception:
|
| 118 |
+
pass
|
| 119 |
+
|
| 120 |
+
print(f"🔒 Compiled {compiled_so_count} Cython .so C-extensions (stripped symbols).")
|
| 121 |
+
print(f"⚡ Compiled {compiled_pyc_count} modules to optimized .pyc bytecode.")
|
| 122 |
+
|
| 123 |
+
# Clean up temporary Cython .c, .cpp, and build directories
|
| 124 |
+
for root, dirs, files in os.walk(base_dir):
|
| 125 |
+
for file in files:
|
| 126 |
+
if file.endswith(".c") or file.endswith(".cpp"):
|
| 127 |
+
if "third_party" not in root and not file.startswith("c_"):
|
| 128 |
+
c_path = os.path.join(root, file)
|
| 129 |
+
try:
|
| 130 |
+
os.remove(c_path)
|
| 131 |
+
except OSError:
|
| 132 |
+
pass
|
| 133 |
+
|
| 134 |
+
build_dir = os.path.join(root, "build")
|
| 135 |
+
if os.path.exists(build_dir):
|
| 136 |
+
shutil.rmtree(build_dir, ignore_errors=True)
|
| 137 |
+
|
| 138 |
+
print(f"🧹 Removed {deleted_count} original source .py files (replaced with compiled binaries).")
|
| 139 |
+
print("=" * 60)
|
| 140 |
+
print("✨ DoppelGen Hardened Binary Build Complete!")
|
| 141 |
+
print("=" * 60)
|
| 142 |
+
|
| 143 |
+
if __name__ == "__main__":
|
| 144 |
+
main()
|
.github/workflows/deploy_hf.yml
ADDED
|
@@ -0,0 +1,77 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
name: Sync to Hugging Face Spaces
|
| 2 |
+
|
| 3 |
+
on:
|
| 4 |
+
push:
|
| 5 |
+
branches: [ main ]
|
| 6 |
+
workflow_dispatch:
|
| 7 |
+
|
| 8 |
+
# Single-flight deploys: each run force-pushes an orphan branch with fresh
|
| 9 |
+
# LFS objects; interleaved runs corrupt HF repo state (exit-128 build failures).
|
| 10 |
+
concurrency:
|
| 11 |
+
group: deploy-hf
|
| 12 |
+
cancel-in-progress: true
|
| 13 |
+
|
| 14 |
+
jobs:
|
| 15 |
+
sync-to-hf:
|
| 16 |
+
runs-on: ubuntu-latest
|
| 17 |
+
steps:
|
| 18 |
+
- name: Checkout Repository
|
| 19 |
+
uses: actions/checkout@v4
|
| 20 |
+
with:
|
| 21 |
+
fetch-depth: 1
|
| 22 |
+
|
| 23 |
+
- name: Set up Python 3.12 & GCC Build Toolchain
|
| 24 |
+
uses: actions/setup-python@v5
|
| 25 |
+
with:
|
| 26 |
+
python-version: "3.12"
|
| 27 |
+
|
| 28 |
+
- name: Install Compiler Dependencies & Git LFS
|
| 29 |
+
run: |
|
| 30 |
+
sudo apt-get update && sudo apt-get install -y git-lfs
|
| 31 |
+
git lfs install
|
| 32 |
+
python -m pip install --upgrade pip
|
| 33 |
+
pip install Cython setuptools wheel
|
| 34 |
+
|
| 35 |
+
- name: Compile Python Core Packages to Native .so C-Binaries
|
| 36 |
+
run: |
|
| 37 |
+
python .github/scripts/compile_binaries.py
|
| 38 |
+
|
| 39 |
+
- name: Push Compiled Binaries to Hugging Face Space via Git LFS
|
| 40 |
+
env:
|
| 41 |
+
HF_TOKEN: ${{ secrets.HF_TOKEN }}
|
| 42 |
+
run: |
|
| 43 |
+
git config user.name "github-actions[bot]"
|
| 44 |
+
git config user.email "github-actions[bot]@users.noreply.github.com"
|
| 45 |
+
|
| 46 |
+
# Create a clean orphan branch without historical git blobs
|
| 47 |
+
git checkout --orphan hf-deploy
|
| 48 |
+
|
| 49 |
+
# Setup Git LFS for binary tracking
|
| 50 |
+
git lfs install
|
| 51 |
+
git lfs track "*.so"
|
| 52 |
+
git lfs track "*.mp4"
|
| 53 |
+
git lfs track "*.webm"
|
| 54 |
+
# Default avatar/background fixtures must ride LFS — HF rejects plain binary blobs
|
| 55 |
+
git lfs track "doppelgen/data/jobs/test/input/**"
|
| 56 |
+
git add .gitattributes
|
| 57 |
+
|
| 58 |
+
# Add all files including compiled .so binaries
|
| 59 |
+
git add -A
|
| 60 |
+
|
| 61 |
+
# Remove heavy binary files to comply with HF storage policy
|
| 62 |
+
# (Preserve ocula_visuals — tracked via Git LFS, and the small default
|
| 63 |
+
# avatar/background fixtures Actora needs at runtime. DreamTalk's
|
| 64 |
+
# style-clip/pose .mat assets are downloaded by preload_models.py
|
| 65 |
+
# at boot, so they are intentionally stripped here.)
|
| 66 |
+
find . -type f \( \
|
| 67 |
+
-name "*.png" -o -name "*.jpg" -o -name "*.jpeg" -o \
|
| 68 |
+
-name "*.wav" -o -name "*.mp3" -o -name "*.mp4" -o -name "*.m4a" -o -name "*.webm" -o -name "*.gif" -o \
|
| 69 |
+
-name "*.onnx" -o -name "*.pth" -o -name "*.pt" -o -name "*.tar" -o -name "*.npy" -o -name "*.mat" -o \
|
| 70 |
+
-name "*.pkl" -o -name "*.zip" -o -name "*.data" -o -name "*.bin" \
|
| 71 |
+
\) ! -path "*/ocula_visuals/*" ! -path "*doppelgen/data/jobs/test/input/*" -exec git rm -f --cached {} + || true
|
| 72 |
+
|
| 73 |
+
git commit -m "Deploy DoppelGen compiled C-extension binary distribution to Hugging Face Space"
|
| 74 |
+
|
| 75 |
+
git remote add hf https://Hazeezadebayo:$HF_TOKEN@huggingface.co/spaces/Hazeezadebayo/doppelgen
|
| 76 |
+
git push --force hf hf-deploy:main
|
| 77 |
+
|
.gitignore
ADDED
|
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Environments
|
| 2 |
+
.env
|
| 3 |
+
|
| 4 |
+
# Model weights (top-level only, NOT the Python model code sub-packages)
|
| 5 |
+
/*/models/
|
| 6 |
+
|
| 7 |
+
# Python
|
| 8 |
+
__pycache__/
|
| 9 |
+
*.pyc
|
| 10 |
+
*.pyo
|
| 11 |
+
*.egg-info/
|
| 12 |
+
|
| 13 |
+
# Hugging Face cache
|
| 14 |
+
.hf_cache/
|
| 15 |
+
**/.hf_cache/
|
| 16 |
+
|
| 17 |
+
# OS
|
| 18 |
+
.DS_Store
|
| 19 |
+
|
| 20 |
+
# Third-party model checkpoints (large weight files)
|
| 21 |
+
**/checkpoints/
|
| 22 |
+
|
| 23 |
+
# Generated outputs
|
| 24 |
+
**/test/output/
|
| 25 |
+
jobs/
|
| 26 |
+
*.db
|
| 27 |
+
|
| 28 |
+
# Runtime logs / reports
|
| 29 |
+
**/output_log.md
|
| 30 |
+
**/project_report.md
|
| 31 |
+
walkthrough.md
|
README.md
ADDED
|
@@ -0,0 +1,151 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: DoppelGen
|
| 3 |
+
emoji: 🎬
|
| 4 |
+
colorFrom: indigo
|
| 5 |
+
colorTo: purple
|
| 6 |
+
sdk: gradio
|
| 7 |
+
sdk_version: 6.24.0
|
| 8 |
+
python_version: "3.12"
|
| 9 |
+
app_file: app.py
|
| 10 |
+
pinned: false
|
| 11 |
+
---
|
| 12 |
+
# Talkinghead Orchestrator
|
| 13 |
+
|
| 14 |
+
**Flora** serves as the centralized master director and orchestration service for the **Creatorium** suite (Voxa, Sonora, Scenea, Actora, Captiona, Ocula, Tempora). Powered by **Google ADK (`google-adk`)** and **Gemini 3.5 Flash**, it unifies deep learning models, computer vision tools, audio processing, NLP sub-services, and social publishing into a single containerized cloud runtime.
|
| 15 |
+
|
| 16 |
+
## Google Cloud Infrastructure & ADK Multi-Agent Architecture
|
| 17 |
+
|
| 18 |
+
```mermaid
|
| 19 |
+
graph TD
|
| 20 |
+
A[User / Single-Page Studio UI] -->|HTTPS POST| B[Google Cloud Run - Flora API Container]
|
| 21 |
+
B -->|Master Orchestrator| C[Flora ADK Director Agent - Gemini 3.5 Flash]
|
| 22 |
+
|
| 23 |
+
C -->|Persists Job State| D[(Google Cloud Firestore DB)]
|
| 24 |
+
C -->|Stores Media Assets| E[(Google Cloud Storage GCS Bucket)]
|
| 25 |
+
C -->|Streams Reasoning Spans| F[OpenTelemetry Telemetry Dashboard]
|
| 26 |
+
|
| 27 |
+
C -->|1. Voxa TTS| G[Speech Waveform & Transcript]
|
| 28 |
+
C -->|2. Sonora Music| H[Ambient Soundtracks]
|
| 29 |
+
C -->|3. Scenea B-roll| I[Context B-rolls]
|
| 30 |
+
C -->|4. Actora LipSync| J[LipSync 1080p Video]
|
| 31 |
+
C -->|5. Captiona Subtitles| K[Dynamic Styled Subtitles]
|
| 32 |
+
C -->|6. Tempora Publisher| L[YouTube / Twitter / Playwright CDP Social Media]
|
| 33 |
+
```
|
| 34 |
+
|
| 35 |
+
## Quick Start — Deploy to Google Cloud Run
|
| 36 |
+
|
| 37 |
+
Flora is pre-configured for 1-click cloud deployment on **Google Cloud Run** using **Firestore** and **Google Cloud Storage (GCS)**:
|
| 38 |
+
|
| 39 |
+
```bash
|
| 40 |
+
# Set GCP Project Environment Variables
|
| 41 |
+
export GCP_PROJECT_ID="your-gcp-project-id"
|
| 42 |
+
export GEMINI_API_KEY="your-gemini-api-key"
|
| 43 |
+
|
| 44 |
+
# Deploy to Google Cloud Run via master runner
|
| 45 |
+
./run_flora.sh deploy
|
| 46 |
+
```
|
| 47 |
+
|
| 48 |
+
## High-Level Architectural Flow
|
| 49 |
+
|
| 50 |
+
```mermaid
|
| 51 |
+
graph TD
|
| 52 |
+
A[Frontend UI] -->|POST Form Data / File paths| B[Flora API]
|
| 53 |
+
B -->|1. Voxa Runner| C[TTS Waveform & Transcript]
|
| 54 |
+
B -->|2. Sonora Runner| D[Ambient Soundtracks]
|
| 55 |
+
B -->|3. Scenea Runner| E[Context B-rolls]
|
| 56 |
+
B -->|4. Actora Runner| F[LipSync Talking Head]
|
| 57 |
+
B -->|5. Captiona Runner| G[Composed Subtitled Video]
|
| 58 |
+
G -->|Success Payload| A
|
| 59 |
+
B -.->|GET Status Poll| A
|
| 60 |
+
```
|
| 61 |
+
|
| 62 |
+
## The Creatorium Pipeline
|
| 63 |
+
|
| 64 |
+
Flora features a unified Web UI (`flora/flora/web`) that serves as a single-page pipeline, sequentially triggering the independent nodes below:
|
| 65 |
+
|
| 66 |
+
1. **Sonora**
|
| 67 |
+
|
| 68 |
+
- **Input:** `[speech.txt] + [optional speech.wav]`
|
| 69 |
+
- **Process:** Performs semantic similarity analysis on the speech text (and speech.wav if provided) to identify the ideal background context, then anlyze pulled high-quality, royalty-free audio tracks from Mixkit, Freesound, and OpenGameArt for alignment.
|
| 70 |
+
- **Output:** Background ambient audio track `background.wav`.
|
| 71 |
+
2. **Voxa**
|
| 72 |
+
|
| 73 |
+
- **Input:** `[speech.txt] + [sample_audio.wav]`
|
| 74 |
+
- **Process:** Utilizes advanced ASR and TTS models to clone the provided voice and narrate the speech text.
|
| 75 |
+
- **Output:** `timestamped_transcript.txt` and `speech.wav` (the cloned narration).
|
| 76 |
+
3. **Scenea**
|
| 77 |
+
|
| 78 |
+
- **Input:** `[timestamped_transcript.txt]`
|
| 79 |
+
- **Process:** Analyzes the transcript to determine which segments require visual enhancement. It fetches relevant, concise B-roll videos (from Pexels or custom sets) tailored perfectly to those specific speech segments (e.g., generating 2 B-rolls if the user specifies a limit of 2).
|
| 80 |
+
- **Output:** B-roll video assets tightly bound to their transcript timestamps.
|
| 81 |
+
4. **Actora**
|
| 82 |
+
|
| 83 |
+
- **Input:** `[background.jpg] + [me.jpg] + [driving_video.mp4] + [speech.wav]`
|
| 84 |
+
- **Process:** Fuses the assets together, applying human-like mannerisms extracted from the driving video to the static image of "me". The lip-syncing is perfectly matched to `speech.wav`.
|
| 85 |
+
- **Output:** `talkinghead.mp4` (A complete, high-fidelity talking head video).
|
| 86 |
+
5. **Captiona**
|
| 87 |
+
|
| 88 |
+
- **Input:** `[talkinghead.mp4] + [timestamped_transcript.txt]`
|
| 89 |
+
- **Process:** Overlays dynamic, styled text captions onto the video, perfectly synchronized with the speech and configured to the user's stylistic preferences.
|
| 90 |
+
- **Output:** `talkinghead_captioned.mp4` (The final, ready-to-publish video).
|
| 91 |
+
|
| 92 |
+
# Pipeline Parallelization Plan
|
| 93 |
+
|
| 94 |
+
We will optimize the execution speed of the Flora orchestration pipeline by running independent tasks concurrently.
|
| 95 |
+
|
| 96 |
+
## Parallel Execution Architecture
|
| 97 |
+
|
| 98 |
+
Currently, the pipeline runs sequentially:
|
| 99 |
+
|
| 100 |
+
```
|
| 101 |
+
Voxa (TTS) -> Sonora (Ambient Audio) -> Scenea (B-rolls) -> Actora (Talking Head) -> Captiona (Subtitles)
|
| 102 |
+
```
|
| 103 |
+
|
| 104 |
+
However, after Voxa runs and produces the `speech.wav` and `speech_transcript.txt` files, the subsequent stages have no data dependencies on each other:
|
| 105 |
+
|
| 106 |
+
* **Sonora** only depends on `speech_transcript.txt`.
|
| 107 |
+
* **Scenea** only depends on `speech_transcript.txt`.
|
| 108 |
+
* **Actora** only depends on `speech.wav`.
|
| 109 |
+
|
| 110 |
+
Thus, we can execute Sonora, Scenea, and Actora concurrently using Python's `concurrent.futures.ThreadPoolExecutor`.
|
| 111 |
+
|
| 112 |
+
```mermaid
|
| 113 |
+
graph TD
|
| 114 |
+
A[Voxa TTS] --> B[Sonora Ambient Audio]
|
| 115 |
+
A --> C[Scenea B-rolls]
|
| 116 |
+
A --> D[Actora Talking Head]
|
| 117 |
+
B --> E[Captiona Subtitles]
|
| 118 |
+
C --> E
|
| 119 |
+
D --> E
|
| 120 |
+
```
|
| 121 |
+
|
| 122 |
+
## Impact on Execution Time
|
| 123 |
+
|
| 124 |
+
The total execution time will drop from:
|
| 125 |
+
`Time(Voxa) + Time(Sonora) + Time(Scenea) + Time(Actora) + Time(Captiona)`
|
| 126 |
+
to:
|
| 127 |
+
`Time(Voxa) + max(Time(Sonora), Time(Scenea), Time(Actora)) + Time(Captiona)`
|
| 128 |
+
|
| 129 |
+
With the host's 32-core CPU capacity, running these three processes simultaneously will not bottleneck local resources, leading to a substantial performance improvement.
|
| 130 |
+
|
| 131 |
+
## Review
|
| 132 |
+
|
| 133 |
+
> [!IMPORTANT]
|
| 134 |
+
> Because subprocesses are run concurrently, stdout and stderr logs will write to the container console in an interleaved manner. However, each sub-process will still run as an isolated execution thread and write to its own independent logs if needed. We will update `PIPELINE_STATUS` to show active progress for all running components (e.g. "Generating Video & Fetching Assets...").
|
| 135 |
+
|
| 136 |
+
## Similar apps:
|
| 137 |
+
|
| 138 |
+
Flora expects a strict Input/Output contract to seamlessly pass data between the nodes:
|
| 139 |
+
|
| 140 |
+
- **Talking head gen:**
|
| 141 |
+
- `https://www.veed.io/tools/text-to-speech-avatar/talking-head-video`
|
| 142 |
+
- `https://www.synthesia.io/tools/talking-head-video-maker`
|
| 143 |
+
- `https://toki.ai/ai-talking-avatar`
|
| 144 |
+
- `/https://captions.ai/solutions/talking-head-videos`
|
| 145 |
+
- **Video understanding:**
|
| 146 |
+
- `/https://huggingface.co/openai/clip-vit-base-patch32/tree/main`
|
| 147 |
+
- `/https://huggingface.co/google/siglip-base-patch16-224/tree/main`
|
| 148 |
+
- `/https://huggingface.co/microsoft/xclip-base-patch32/tree/main`
|
| 149 |
+
- `/https://huggingface.co/apple/MobileCLIP2-S3/tree/main`
|
| 150 |
+
|
| 151 |
+
By treating `flora` as the orchestrator, the entire Creatorium ecosystem operates as a cohesive, highly-optimized production engine.
|
actora/.dockerignore
ADDED
|
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
models/
|
| 2 |
+
.git/
|
| 3 |
+
test/output/
|
| 4 |
+
**/__pycache__/
|
| 5 |
+
*.pyc
|
| 6 |
+
*.pyo
|
| 7 |
+
*.pyd
|
| 8 |
+
.env
|
actora/.gitignore
ADDED
|
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
.env
|
| 2 |
+
__pycache__/
|
| 3 |
+
*.pyc
|
| 4 |
+
.hf_cache/
|
| 5 |
+
# Ignore generated output files but keep the output folder structure
|
| 6 |
+
b_roll_rag/data/output/*
|
| 7 |
+
!b_roll_rag/data/output/.gitkeep
|
| 8 |
+
|
| 9 |
+
# Ignore local AI agent logs and living documents
|
| 10 |
+
output_log.md
|
| 11 |
+
project_report.md
|
| 12 |
+
git_workflow.md
|
actora/README.md
ADDED
|
@@ -0,0 +1,253 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# ⚡ Edge Talking Head Pipeline (Raspberry Pi Optimized)
|
| 2 |
+
|
| 3 |
+
An ultra-lightweight, hardware-optimized content creation pipeline designed to run on low-resource edge devices (e.g., Raspberry Pi 4/5). It takes a portrait image, removes the background, composites it onto a beautiful room background, and uses recorded audio to generate a lip-synced talking head video.
|
| 4 |
+
|
| 5 |
+
---
|
| 6 |
+
|
| 7 |
+
## Architectural Flow
|
| 8 |
+
```mermaid
|
| 9 |
+
graph TD
|
| 10 |
+
A[Image / Audio / Driving Video] --> B[EdgePipeline: api.py]
|
| 11 |
+
B --> C{Talking Engine Model?}
|
| 12 |
+
C -->|FasterLivePortrait / FOMM / TPSMM| D[Adaptive Keyframe Sampling: skip_pct=0.80]
|
| 13 |
+
C -->|DreamTalk / Wav2Lip| E[Subsampled Frame Generation: render_every_n=2]
|
| 14 |
+
D --> F[Linear Interpolation & Composite]
|
| 15 |
+
E --> F
|
| 16 |
+
F --> G[Talking Head Output Video]
|
| 17 |
+
```
|
| 18 |
+
|
| 19 |
+
---
|
| 20 |
+
|
| 21 |
+
## 📂 Project Directory Structure
|
| 22 |
+
|
| 23 |
+
```
|
| 24 |
+
actora/
|
| 25 |
+
├── docker/
|
| 26 |
+
│ ├── Dockerfile
|
| 27 |
+
│ ├── docker-compose.yml
|
| 28 |
+
│ ├── entrypoint.sh
|
| 29 |
+
│ └── requirements.txt
|
| 30 |
+
├── models/ # Auto-created; stores cached ONNX and PyTorch weights
|
| 31 |
+
├── actora_core/ # Core library package
|
| 32 |
+
│ ├── __init__.py
|
| 33 |
+
│ ├── api.py # Pipeline orchestration API
|
| 34 |
+
│ ├── base.py # Abstract base class definitions
|
| 35 |
+
│ ├── dreamtalk.py # DreamTalk Audio-driven Expressive Face Generator
|
| 36 |
+
│ ├── dreamtalk_src/ # DreamTalk cloned source repository (auto-cloned)
|
| 37 |
+
│ ├── generator.py # Lip-sync generation interfaces (Wav2Lip)
|
| 38 |
+
│ ├── matting.py # Background removal interfaces (MODNet / MediaPipe)
|
| 39 |
+
│ ├── models.py # Swappable model registry factory
|
| 40 |
+
│ ├── postprocess.py # Post-processing super-resolution engine (GRL-GAN)
|
| 41 |
+
│ └── utils.py # Mathematical operations and video encoding helpers
|
| 42 |
+
├── test/
|
| 43 |
+
│ ├── input/ # Place inputs here (me_1.png, background_1.png, audio_1.wav)
|
| 44 |
+
│ ├── output/ # Generated outputs will be stored here
|
| 45 |
+
│ ├── test.py # Integration testing suite
|
| 46 |
+
│ ├── test_cartoonify.py # AnimeGAN cartoonification verification test script
|
| 47 |
+
│ └── verify_integrity.py # Verification utility checking dimensions and frame stats
|
| 48 |
+
├── preload_models.py # Pre-caches models (checkpoints, Wav2Vec2, ONNX models)
|
| 49 |
+
├── run_actora.sh # Pipeline container manager wrapper script
|
| 50 |
+
├── project_report.md # Continuous technical and architectural progress report
|
| 51 |
+
└── README.md # This documentation file
|
| 52 |
+
```
|
| 53 |
+
|
| 54 |
+
---
|
| 55 |
+
|
| 56 |
+
## 🏗️ Architectural Flow
|
| 57 |
+
|
| 58 |
+
```
|
| 59 |
+
[ INPUT TYPE ]
|
| 60 |
+
│
|
| 61 |
+
┌──────────────────────────┴──────────────────────────┐
|
| 62 |
+
▼ ▼
|
| 63 |
+
[ IMAGE INPUT ] [ VIDEO INPUT ]
|
| 64 |
+
│ │
|
| 65 |
+
1. Extract alpha matte once 1. Read video frames
|
| 66 |
+
2. Composite foreground on background once 2. Perform soft-blending
|
| 67 |
+
3. Detect & crop face region (Haar Cascade) frame-by-frame (EMA smoothed)
|
| 68 |
+
4. Run Expressive Head Generation: 3. Save composited video
|
| 69 |
+
a. Generate expressive facial animation sequence
|
| 70 |
+
using DreamTalk (Diffusion 3DMM + PIRenderer)
|
| 71 |
+
b. Run optional Wav2Lip ONNX for secondary sync
|
| 72 |
+
c. Paste crop sequence back onto canvas
|
| 73 |
+
5. Mux audio and video stream (FFmpeg)
|
| 74 |
+
```
|
| 75 |
+
|
| 76 |
+
---
|
| 77 |
+
|
| 78 |
+
## Technical & Algorithmic Breakdown
|
| 79 |
+
|
| 80 |
+
1. **Background Matting (Segmentation)**:
|
| 81 |
+
|
| 82 |
+
- Uses optimized ONNX runtimes.
|
| 83 |
+
- Available models: `modnet` (~25 MB, high accuracy) and `mediapipe` (<3 MB, ultra-lightweight and optimized for mobile CPU execution).
|
| 84 |
+
- Operation: Extract a single-channel alpha matte $A \in [0, 1]^{H \times W}$, then perform alpha compositing:
|
| 85 |
+
$$
|
| 86 |
+
I_{out} = I_{fg} \odot A + I_{bg} \odot (1 - A)
|
| 87 |
+
$$
|
| 88 |
+
2. **Talking Head Generation (Lip Sync & Face Warping)**:
|
| 89 |
+
|
| 90 |
+
- **DreamTalk (Diffusion + PIRenderer)**: The default high-fidelity audio-driven talking face generator. It uses a diffusion-based denoising network (DiffusionNet) to generate 3DMM facial motion coefficients from Wav2Vec2 audio embeddings, and a neural PIRenderer to render the generated coefficients onto a static face crop. This produces highly expressive facial animation (gaze shifts, natural expression) from a single photo and an audio file.
|
| 91 |
+
- **Wav2Lip ONNX + BiomechanicalJoint**: An alternative audio-driven fallback generator. It uses a lightweight Wav2Lip ONNX (~145 MB) model to sync lip movements. To prevent the "lifeless/frozen head" effect typical of Wav2Lip, a physics spring system (`BiomechanicalJoint`) extracts prosodic audio impulses and applies procedural vertical head bobbing to simulate natural human mannerisms.
|
| 92 |
+
- **FOMM (First-Order Motion Model)**: A video-driven generator. It uses two Qualcomm-exported ONNX models (`detector.onnx` to extract 10 facial keypoints + Jacobians, and `generator.onnx` to warp the source face). Motion is transferred from a driving video using relative keypoint normalization, preserving the source person's identity while applying the driver's head pose, eye blinks, and expressions.
|
| 93 |
+
|
| 94 |
+
---
|
| 95 |
+
|
| 96 |
+
## Architectural Flow
|
| 97 |
+
|
| 98 |
+
```
|
| 99 |
+
[ INPUT CHECK ]
|
| 100 |
+
│
|
| 101 |
+
├──► IF IMAGE + DRIVING VIDEO (video-driven, FOMM):
|
| 102 |
+
│ 1. Extract alpha matte once (MODNet/MediaPipe ONNX)
|
| 103 |
+
│ 2. Composite foreground once onto static background (OpenCV) -> Phase 1 static composite
|
| 104 |
+
│ 3. Save static composite to first_phase_overlay.png
|
| 105 |
+
│ 4. Detect and crop face region from the Phase 1 static composite
|
| 106 |
+
│ 5. Run FOMM: detect source keypoints once; for each driving frame,
|
| 107 |
+
│ detect driving keypoints, normalize via relative motion transfer,
|
| 108 |
+
│ run generator to produce animated 256x256 face crop.
|
| 109 |
+
│ 6. Paste generated face crop sequence back onto Phase 1 static composite
|
| 110 |
+
│ 7. Mux final video and input audio (FFmpeg)
|
| 111 |
+
│
|
| 112 |
+
├──► IF IMAGE (audio-driven, DreamTalk):
|
| 113 |
+
│ 1. Extract alpha matte once (MODNet/MediaPipe ONNX)
|
| 114 |
+
│ 2. Composite foreground once onto static background (OpenCV) -> Phase 1 static composite
|
| 115 |
+
│ 3. Save static composite to first_phase_overlay.png
|
| 116 |
+
│ 4. Detect and crop face region from the Phase 1 static composite
|
| 117 |
+
│ 5. Run Selected Generator on the face crop with audio:
|
| 118 |
+
│ - DreamTalk: Extract Wav2Vec2 audio embeddings, denoise motion coefficients
|
| 119 |
+
│ via DiffusionNet, render face crops via PIRenderer.
|
| 120 |
+
│ - Wav2Lip: Run Wav2Lip directly on face crop.
|
| 121 |
+
│ 6. Paste generated/synchronized face crop sequence directly back onto the Phase 1 static composite
|
| 122 |
+
│ 7. Mux final video and input audio (FFmpeg)
|
| 123 |
+
│
|
| 124 |
+
└──► IF VIDEO:
|
| 125 |
+
1. Frame-by-frame background matting (MediaPipe ONNX)
|
| 126 |
+
2. Soft-blend video frames onto chosen background
|
| 127 |
+
```
|
| 128 |
+
|
| 129 |
+
---
|
| 130 |
+
|
| 131 |
+
## 🛠️ Execution & Installation Guide
|
| 132 |
+
|
| 133 |
+
### 1. Prerequisites
|
| 134 |
+
|
| 135 |
+
Ensure your edge device has `docker`, `docker-compose`, `python3`, and `pip` installed.
|
| 136 |
+
|
| 137 |
+
### 2. Model Weight Caching
|
| 138 |
+
|
| 139 |
+
Run the helper script on the host system to pre-download the optimized model weights, clone the DreamTalk repo, and cache Wav2Vec2 transformers weights:
|
| 140 |
+
|
| 141 |
+
```bash
|
| 142 |
+
pip install huggingface_hub transformers
|
| 143 |
+
python3 preload_models.py
|
| 144 |
+
```
|
| 145 |
+
|
| 146 |
+
### 3. Container Lifecycle Commands
|
| 147 |
+
|
| 148 |
+
Control execution using the centralized wrapper script:
|
| 149 |
+
|
| 150 |
+
```bash
|
| 151 |
+
# Build the Docker image
|
| 152 |
+
./run_actora.sh build
|
| 153 |
+
|
| 154 |
+
# Launch the container in the background
|
| 155 |
+
./run_actora.sh up
|
| 156 |
+
|
| 157 |
+
# Run test execution (processes inputs in test/input/)
|
| 158 |
+
./run_actora.sh test
|
| 159 |
+
|
| 160 |
+
# Run cartoonification test (processes me_6.jpg through ArcaneGAN and DCT-Net filters)
|
| 161 |
+
docker exec -t actora_edge python3 /app/test/test_cartoonify.py
|
| 162 |
+
|
| 163 |
+
# Stop container and clean up output folders
|
| 164 |
+
./run_actora.sh down
|
| 165 |
+
./run_actora.sh clean
|
| 166 |
+
```
|
| 167 |
+
|
| 168 |
+
---
|
| 169 |
+
|
| 170 |
+
## ⚙️ Configuration & Customization API
|
| 171 |
+
|
| 172 |
+
The `EdgePipeline` class exposes customization options to control aspect ratio, grading effects, and natural head/body movements:
|
| 173 |
+
|
| 174 |
+
### 1. Aspect Ratio Canvas Formatting
|
| 175 |
+
|
| 176 |
+
Define the target output shape of the composition using the `aspect_ratio` parameter:
|
| 177 |
+
|
| 178 |
+
* `horizontal` (Default: `1024x576`, 16:9 widescreen format)
|
| 179 |
+
* `vertical` (`576x1024`, 9:16 portrait format for social reels/shorts)
|
| 180 |
+
* `square` (`768x768`, 1:1 format)
|
| 181 |
+
|
| 182 |
+
The system center-crops the background without stretching and auto-scales the foreground subject to occupy ~85% of the frame height, anchoring them at the bottom-center.
|
| 183 |
+
|
| 184 |
+
### 2. Swappable Generator Engines
|
| 185 |
+
|
| 186 |
+
Choose the animation quality and execution speed using the `generator_type` parameter in `test/test.py`:
|
| 187 |
+
|
| 188 |
+
* `dreamtalk` (Default): Expressive, diffusion-based 3DMM talking face generator with natural facial mannerisms.
|
| 189 |
+
* `wav2lip`: Standard audio-driven lip sync with static head posture.
|
| 190 |
+
|
| 191 |
+
### 3. Cinematic Blending & Effects
|
| 192 |
+
|
| 193 |
+
Harmonize the composited layers using the `blend_effect` parameter:
|
| 194 |
+
|
| 195 |
+
* `none`: standard compositing.
|
| 196 |
+
* `cinematic_warm`: warm sunlight grading with enhanced contrast.
|
| 197 |
+
* `b_and_w`: classic high-contrast silver-halide film simulation.
|
| 198 |
+
* `bokeh_blur`: Gaussian blur applied to the background *prior* to compositing, separating the subject with depth-of-field.
|
| 199 |
+
* `arcanegan`: stylized Arcane-like cartoon look using the JIT-compiled ArcaneGANv0.4 model.
|
| 200 |
+
* `dctnet_artstyle`: Artstyle cartoon look using DCT-Net.
|
| 201 |
+
* `dctnet_3d`: 3D cartoon style using DCT-Net.
|
| 202 |
+
* `dctnet_anime`: Anime cartoon style using DCT-Net.
|
| 203 |
+
* **Ambient Color Matcher**: Automatically computes the average color of the background and casts 6% illumination onto the subject for lighting integration.
|
| 204 |
+
|
| 205 |
+
### 4. Post-Processing & Super-Resolution
|
| 206 |
+
|
| 207 |
+
Improve facial clarity and eliminate neural morphing/blur artifacts around the talking head's boundary using the `--postprocess` parameter:
|
| 208 |
+
|
| 209 |
+
* `none` (Default): Bypasses post-processing.
|
| 210 |
+
* `grl_gan`: Applies GRL-GAN ONNX super-resolution upscaling (4x) and downscaling back to 256x256 using Lanczos4, resulting in highly detailed facial features.
|
| 211 |
+
|
| 212 |
+
---
|
| 213 |
+
|
| 214 |
+
## 💡 Edge Optimizations & Real-Time Plan
|
| 215 |
+
|
| 216 |
+
To run this pipeline in real-time on a Raspberry Pi:
|
| 217 |
+
|
| 218 |
+
1. **MediaPipe Segmentation**: Use the `mediapipe` model key (<3 MB) for matting to reduce compute overhead compared to MODNet (~25 MB).
|
| 219 |
+
2. **Single-Pass Matting**: For static image inputs, perform background matting and compositing *once* instead of frame-by-frame.
|
| 220 |
+
3. **Face Region Localization**: Only run the generator model on a cropped $256 \times 256$ face boundary, then paste it back. This avoids distorting the high-resolution background and significantly reduces processing time.
|
| 221 |
+
4. **Memory / CPU Optimizations**: Load models on CPU with PyTorch CPU-optimized wheels. Monkeypatch `.cuda()` to return CPU tensors to avoid runtime failures without bloated CUDA dependencies.
|
| 222 |
+
|
| 223 |
+
[github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model](https://github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model)
|
| 224 |
+
|
| 225 |
+
what our system does:
|
| 226 |
+
|
| 227 |
+
lipsync: takes a image of me, takes a background image that i desire. puts my segmented cutout on top of the desired background neetly and ensures perfect blending to make a new image. outputs this as this is the first success. then, takes my input audio and process the image to lipsync to the audio and then outputs a video. we have the option to apply visual effects/filters as well as postprocessing for higher quality.
|
| 228 |
+
|
| 229 |
+
videosync: similar to the above, same image technique and success yardstick. except these models take a video as a driving entity to animate the mannerism so that it is more convincing and then offers a chance to also lypsinc but not required sometimes since the driving video themselves might have contained mouth animations. thereby giving the illusion of speed. this system is therefore matched to our audio and as such outputted. like in the previous options for effects/filters and postprocessing into 4k exists.
|
| 230 |
+
|
| 231 |
+
lastly, we do:
|
| 232 |
+
|
| 233 |
+
scenesync: this is essentially us, instead of begining with an image like in the first 2 cases, we have a video of us and perhaps a audio as well, but we simply wanna change the background of our video and blend it appropriately so that it harmonizes well with the new background and audio. then we choose this option, it too has the effect/filters choise as well as post processing options.
|
| 234 |
+
|
| 235 |
+
here are the desireables:
|
| 236 |
+
|
| 237 |
+
1. our system is built for edge devices, and as such must priotise cpu only systems. that is, real time operations guarantees and speed comparable to big tech on tiny hardwares.
|
| 238 |
+
2. we do not want a system bloat. hence our models are largely under 300mb individually and as such we do our best to avoid redundancies. perfect object oriented senior level programming is expected. that also means something as simple as investigating which model "/home/azeez/ws/dev_env/py_code/projects/actora/models" no longer referenced or used within our pipeline and safely deleting it.
|
| 239 |
+
|
| 240 |
+
next, we wanted two things:
|
| 241 |
+
|
| 242 |
+
1. to tidy up our system and seperate what was thirdparty "/home/azeez/ws/dev_env/py_code/projects/actora/third_party" which would be pulled from the internet on building of our open source talking head project "/home/azeez/ws/dev_env/py_code/projects/actora" and the core files "/home/azeez/ws/dev_env/py_code/projects/actora/actora_core" for orchestration and pipeline.
|
| 243 |
+
2. Optimize our running/algorithmic speed so that processes no longer take an uncomfortable amount of time before getting completed.
|
| 244 |
+
|
| 245 |
+
in the process of trying to do the above, you have broken things and i require them fixed:
|
| 246 |
+
|
| 247 |
+
1. the faster_liveportrait no longer works. it no longer cuts out the person, puts them on a new background and animates them. in fact, in its case, we get a video of the desired background but no talking head. just a video of the desired background and the voice overlay. which is sad as we had successfully made this work.
|
| 248 |
+
2. tpsmm and fomm no longer high quality either and rather than a square around the face being animated since its required the detect face to match the driving video face stuff, it appears that we using a face video to animate an entire body hence it morphs in weird shapes and very irregular dimensions are experienced by the human. this was not the case before either.
|
| 249 |
+
|
| 250 |
+
please i need you to investigate our goals against what we have currently and ensure there is an alignment.
|
| 251 |
+
i need you to only make an implementation plan only after you have completely understood the codebase and the stakes.
|
| 252 |
+
i need you to evaluate the codebase for redundancies and inefficiencies and write a plan to address them and make the code usable to humans.
|
| 253 |
+
fix all errors.
|
actora/actora/__init__.py
ADDED
|
File without changes
|
actora/actora/api/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
from .api import EdgePipeline
|
actora/actora/api/api.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:5b86ab223099863e8abc86a088334c99d42fe000444b4f24fd887a8e3142d422
|
| 3 |
+
size 409528
|
actora/actora/api/main.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:1429a26fc2fd3697a372480954e919a895ff01b239502a097c044414a38e0f57
|
| 3 |
+
size 146648
|
actora/actora/core/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
from .postprocess import PostProcessEngine
|
actora/actora/core/arcanegan.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2ec6ee10b89ab087180e54099fb7419fb243bde4cb77e0ea9d2f5124e162f70c
|
| 3 |
+
size 153536
|
actora/actora/core/dctnet.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:0ec7a2b05d11720bd551c780c33fac85361ef96c5310e38cbf1fa8ce21338466
|
| 3 |
+
size 101560
|
actora/actora/core/filters.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:18d5be9c0a526deaae8a0ede8219179bdcbb8e8809f9cda37964dbe3444323ca
|
| 3 |
+
size 116976
|
actora/actora/core/matting.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:89f15c88a04dce67866e960f4e1c4d3b5840d15d54f5f388ce52f621c65cbb09
|
| 3 |
+
size 106056
|
actora/actora/core/neural_accelerator.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:dfcb044b5bf77bed64a62f9d0496ec6d044887c04451e2758dc37f7b5929e36e
|
| 3 |
+
size 136264
|
actora/actora/core/postprocess.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:9fa34853533a79334ff8293dd376752b1c309406dd706559b6caf6350e639536
|
| 3 |
+
size 132096
|
actora/actora/models/__init__.py
ADDED
|
@@ -0,0 +1 @@
|
|
|
|
|
|
|
| 1 |
+
from .models import ModelFactory
|
actora/actora/models/base.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:7243b4def147827eb493d29b3244b7c5206974aa3ddbb880e979c81bb111aa88
|
| 3 |
+
size 68120
|
actora/actora/models/dreamtalk.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:348549ec8077521cb570decb4f7f9065c0a77a6c3b6af22bdd47a9a028adafd1
|
| 3 |
+
size 318032
|
actora/actora/models/faster_liveportrait.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a5bacee2cc64b93262508610f76761c686cfd79a6be7713b15f1c8b6c611271b
|
| 3 |
+
size 247792
|
actora/actora/models/fomm.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:ee61a8ca7c7d1c05fe1ecb3c0567ca26f72b23bc251176bce3bd1b17bc702064
|
| 3 |
+
size 229384
|
actora/actora/models/generator.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:d5dba13ff15535b66906f1191640c569c17e5e02757e5a743904f21e53399197
|
| 3 |
+
size 193760
|
actora/actora/models/models.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:b45810a7353e02cd2cf23e8c610342db69a829f7a4c56733beacbbca93f6fd4e
|
| 3 |
+
size 67912
|
actora/actora/models/tpsmm.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:23b83d2b6f28a3bd6e6b0a4877b6cef65541abda6f6d2dbfefa03d55440e8edd
|
| 3 |
+
size 283848
|
actora/actora/schema/__init__.py
ADDED
|
File without changes
|
actora/actora/test/create_test_inputs.py
ADDED
|
@@ -0,0 +1,53 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import cv2
|
| 2 |
+
import numpy as np
|
| 3 |
+
import wave
|
| 4 |
+
import struct
|
| 5 |
+
import os
|
| 6 |
+
|
| 7 |
+
def create_inputs():
|
| 8 |
+
input_dir = "/app/test/input"
|
| 9 |
+
os.makedirs(input_dir, exist_ok=True)
|
| 10 |
+
|
| 11 |
+
# 1. Create me.jpg (foreground: simple drawing of a person/face box)
|
| 12 |
+
me_img = np.zeros((512, 512, 3), dtype=np.uint8) + 240 # Off-white background
|
| 13 |
+
# Draw head
|
| 14 |
+
cv2.circle(me_img, (256, 200), 100, (200, 150, 150), -1)
|
| 15 |
+
# Draw body
|
| 16 |
+
cv2.rectangle(me_img, (150, 300), (362, 512), (100, 100, 250), -1)
|
| 17 |
+
# Draw eyes
|
| 18 |
+
cv2.circle(me_img, (220, 180), 10, (50, 50, 50), -1)
|
| 19 |
+
cv2.circle(me_img, (292, 180), 10, (50, 50, 50), -1)
|
| 20 |
+
# Draw mouth
|
| 21 |
+
cv2.rectangle(me_img, (220, 240), (292, 260), (50, 50, 200), -1)
|
| 22 |
+
|
| 23 |
+
cv2.imwrite(os.path.join(input_dir, "me.jpg"), me_img)
|
| 24 |
+
print("Created me.jpg")
|
| 25 |
+
|
| 26 |
+
# 2. Create background.png (background space)
|
| 27 |
+
bg_img = np.zeros((512, 512, 3), dtype=np.uint8)
|
| 28 |
+
# Draw some grid lines to simulate a living room background
|
| 29 |
+
for y in range(0, 512, 64):
|
| 30 |
+
cv2.line(bg_img, (0, y), (512, y), (120, 120, 120), 2)
|
| 31 |
+
for x in range(0, 512, 64):
|
| 32 |
+
cv2.line(bg_img, (x, 0), (x, 512), (120, 120, 120), 2)
|
| 33 |
+
cv2.imwrite(os.path.join(input_dir, "background.png"), bg_img)
|
| 34 |
+
print("Created background.png")
|
| 35 |
+
|
| 36 |
+
# 3. Create audio.wav (3 seconds of a 440Hz sine wave)
|
| 37 |
+
sample_rate = 16000
|
| 38 |
+
duration = 3.0
|
| 39 |
+
num_samples = int(sample_rate * duration)
|
| 40 |
+
|
| 41 |
+
audio_file = os.path.join(input_dir, "audio.wav")
|
| 42 |
+
wav_file = wave.open(audio_file, 'w')
|
| 43 |
+
wav_file.setparams((1, 2, sample_rate, num_samples, 'NONE', 'not compressed'))
|
| 44 |
+
|
| 45 |
+
for i in range(num_samples):
|
| 46 |
+
value = int(32767.0 * np.sin(2.0 * np.pi * 440.0 * i / sample_rate))
|
| 47 |
+
data = struct.pack('<h', value)
|
| 48 |
+
wav_file.writeframesraw(data)
|
| 49 |
+
wav_file.close()
|
| 50 |
+
print("Created audio.wav")
|
| 51 |
+
|
| 52 |
+
if __name__ == "__main__":
|
| 53 |
+
create_inputs()
|
actora/actora/test/test.py
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import sys
|
| 2 |
+
import os
|
| 3 |
+
import time
|
| 4 |
+
# Ensure project root is in PYTHONPATH
|
| 5 |
+
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))
|
| 6 |
+
|
| 7 |
+
import actora.utils.warnings_patch
|
| 8 |
+
|
| 9 |
+
from actora.api.main import main
|
| 10 |
+
|
| 11 |
+
if __name__ == "__main__":
|
| 12 |
+
run_start = time.time()
|
| 13 |
+
|
| 14 |
+
# Test by explicitly passing desired arguments
|
| 15 |
+
test_args = [
|
| 16 |
+
"--mode", "LipSync", # 'VideoSync', 'LipSync', 'SceneSync'
|
| 17 |
+
"--img", "/app/actora/test/input/me_6.jpg",
|
| 18 |
+
"--vid", "/app/actora/test/input/vid_h1.mp4",
|
| 19 |
+
"--audio", "/app/actora/test/input/audio_1.wav",
|
| 20 |
+
"--bg", "/app/actora/test/input/background_n5.png",
|
| 21 |
+
"--effect", "arcanegan",
|
| 22 |
+
"--limit-frames", "150",
|
| 23 |
+
"--blend-amount", "0",
|
| 24 |
+
"--generator", "dreamtalk", #(fomm/tpsmm/faster_liveportrait = VideoSync), (wav2lip/dreamtalk = LipSync)
|
| 25 |
+
"--arcanegan-size", "480",
|
| 26 |
+
"--postprocess", "grl_gan", # "none,grl_gan,real_esrgan"
|
| 27 |
+
"--matting-rate", "100"
|
| 28 |
+
]
|
| 29 |
+
main(test_args)
|
| 30 |
+
|
| 31 |
+
print(f"\n [TIMER] Total wall-clock time: {time.time() - run_start:.1f}s")
|
| 32 |
+
|
| 33 |
+
# Filters: 'none', 'cinematic_warm', 'b_and_w', 'bokeh_blur',
|
| 34 |
+
# 'arcanegan', 'dctnet_artstyle', 'dctnet_3d', 'dctnet_anime'
|
| 35 |
+
|
actora/actora/utils/__init__.py
ADDED
|
File without changes
|
actora/actora/utils/utils.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:f0553c81dd4296bed1f5e87f1f3cc7aae5b481ca08b3c3af79320ad6f77fb96a
|
| 3 |
+
size 573664
|
actora/actora/utils/warnings_patch.cpython-312-x86_64-linux-gnu.so
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:2890d4b1a5b01c48ba7848874ac6a80ebf128213cded923f9a98ce122ac496c6
|
| 3 |
+
size 63192
|
actora/docker/Dockerfile
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
FROM python:3.11-slim
|
| 2 |
+
RUN apt-get update && apt-get install -y ffmpeg libsm6 libxext6 git && rm -rf /var/lib/apt/lists/*
|
| 3 |
+
WORKDIR /app
|
| 4 |
+
COPY docker/requirements.txt .
|
| 5 |
+
RUN pip install --no-cache-dir -r requirements.txt
|
| 6 |
+
|
| 7 |
+
COPY docker/entrypoint.sh /entrypoint.sh
|
| 8 |
+
RUN chmod +x /entrypoint.sh
|
| 9 |
+
|
| 10 |
+
# Set HF_HOME so huggingface_hub uses the persistent models cache directory at runtime
|
| 11 |
+
ENV HF_HOME=/app/models/.cache/huggingface
|
| 12 |
+
ENV MODELSCOPE_CACHE=/app/models/.cache/modelscope
|
| 13 |
+
|
| 14 |
+
ENTRYPOINT ["/entrypoint.sh"]
|
actora/docker/docker-compose.yml
ADDED
|
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
services:
|
| 2 |
+
edge_processor:
|
| 3 |
+
build:
|
| 4 |
+
context: ../
|
| 5 |
+
dockerfile: docker/Dockerfile
|
| 6 |
+
container_name: Actora
|
| 7 |
+
image: actora:latest
|
| 8 |
+
volumes:
|
| 9 |
+
- ../actora:/app/actora
|
| 10 |
+
- ../models:/app/models
|
| 11 |
+
- ../preload_models.py:/app/preload_models.py
|
| 12 |
+
- ../third_party:/app/third_party
|
| 13 |
+
command: tail -f /dev/null
|
actora/docker/entrypoint.sh
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
echo "Talking Head Edge Container Initialized."
|
| 3 |
+
exec "$@"
|
actora/docker/requirements.txt
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
--extra-index-url https://download.pytorch.org/whl/cpu
|
| 2 |
+
onnxruntime>=1.18.0
|
| 3 |
+
numpy<2.0
|
| 4 |
+
opencv-python-headless
|
| 5 |
+
librosa
|
| 6 |
+
soundfile
|
| 7 |
+
huggingface_hub
|
| 8 |
+
torch
|
| 9 |
+
torchvision
|
| 10 |
+
torchaudio
|
| 11 |
+
transformers
|
| 12 |
+
scipy
|
| 13 |
+
yacs
|
| 14 |
+
tensorflow-cpu==2.14.0
|
| 15 |
+
modelscope>=1.14.0
|
| 16 |
+
addict
|
| 17 |
+
datasets
|
| 18 |
+
oss2
|
| 19 |
+
yapf
|
| 20 |
+
simplejson
|
| 21 |
+
sortedcontainers
|
| 22 |
+
easydict
|
| 23 |
+
torchgeometry
|
| 24 |
+
omegaconf
|
| 25 |
+
munch
|
| 26 |
+
scikit-image
|
| 27 |
+
ffmpeg-python
|
actora/preload_models.py
ADDED
|
@@ -0,0 +1,496 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""
|
| 2 |
+
TL;DR: Preloads all model weights (ONNX, PyTorch checkpoints, and Transformers cache)
|
| 3 |
+
so the container starts fully self-contained with no runtime downloads required.
|
| 4 |
+
|
| 5 |
+
Models managed:
|
| 6 |
+
- modnet.onnx — Background matting (Hugging Face)
|
| 7 |
+
- wav2lip.onnx — Lip sync engine (Hugging Face)
|
| 8 |
+
- mediapipe.onnx — Light background matting (Hugging Face)
|
| 9 |
+
- models/fomm/ — First-Order Motion Model ONNX (Qualcomm S3)
|
| 10 |
+
- DreamTalk checkpoints— Denoising network + PIRenderer (Hugging Face)
|
| 11 |
+
- Wav2Vec2 — Transformers audio feature extractor (Hugging Face)
|
| 12 |
+
- FasterLivePortrait — 9 ONNX models for video-driven talking head (Hugging Face)
|
| 13 |
+
- models/animegan/ — AnimeGANv2 and AnimeGANv3 ONNX style-transfer models
|
| 14 |
+
|
| 15 |
+
Run this once (or on `./run_actora.sh up`) to populate the `models/` directory.
|
| 16 |
+
All checks are local-first: no network call is made if the file already exists.
|
| 17 |
+
"""
|
| 18 |
+
import os
|
| 19 |
+
import sys
|
| 20 |
+
import shutil
|
| 21 |
+
import subprocess
|
| 22 |
+
import urllib.request
|
| 23 |
+
import zipfile
|
| 24 |
+
import asyncio
|
| 25 |
+
import asyncio.base_events
|
| 26 |
+
|
| 27 |
+
def _patch_asyncio_del():
|
| 28 |
+
"""Monkeypatch BaseEventLoop.__del__ to suppress noisy 'ValueError: Invalid file descriptor: -1' during GC in child preload processes."""
|
| 29 |
+
_orig_del = asyncio.base_events.BaseEventLoop.__del__
|
| 30 |
+
def _patched_del(self):
|
| 31 |
+
try:
|
| 32 |
+
_orig_del(self)
|
| 33 |
+
except Exception as e:
|
| 34 |
+
if "Invalid file descriptor: -1" not in str(e):
|
| 35 |
+
raise
|
| 36 |
+
asyncio.base_events.BaseEventLoop.__del__ = _patched_del
|
| 37 |
+
|
| 38 |
+
_patch_asyncio_del()
|
| 39 |
+
|
| 40 |
+
def cleanup_event_loops():
|
| 41 |
+
try:
|
| 42 |
+
p = asyncio.get_event_loop_policy()
|
| 43 |
+
if hasattr(p, "_local") and getattr(p._local, "_loop", None) is not None:
|
| 44 |
+
lp = p._local._loop
|
| 45 |
+
if lp and not lp.is_running() and not lp.is_closed():
|
| 46 |
+
lp.close()
|
| 47 |
+
except Exception:
|
| 48 |
+
pass
|
| 49 |
+
|
| 50 |
+
from huggingface_hub import hf_hub_download
|
| 51 |
+
|
| 52 |
+
|
| 53 |
+
project_dir = os.path.dirname(os.path.abspath(__file__))
|
| 54 |
+
third_party_dir = os.environ.get("THIRD_PARTY_DIR", os.path.join(project_dir, "third_party"))
|
| 55 |
+
os.makedirs(third_party_dir, exist_ok=True)
|
| 56 |
+
|
| 57 |
+
|
| 58 |
+
# ---------------------------------------------------------------------------
|
| 59 |
+
# Hugging Face ONNX manifests (modnet, wav2lip, mediapipe)
|
| 60 |
+
# ---------------------------------------------------------------------------
|
| 61 |
+
MANIFEST = {
|
| 62 |
+
"modnet": {"repo": "Xenova/modnet", "file": "onnx/model.onnx", "target": "modnet.onnx"},
|
| 63 |
+
"wav2lip": {"repo": "bluefoxcreation/Wav2lip-Onnx", "file": "wav2lip.onnx", "target": "wav2lip.onnx"},
|
| 64 |
+
"mediapipe": {"repo": "onnx-community/mediapipe_selfie_segmentation", "file": "onnx/model.onnx", "target": "mediapipe.onnx"},
|
| 65 |
+
}
|
| 66 |
+
|
| 67 |
+
# ---------------------------------------------------------------------------
|
| 68 |
+
# FOMM — hosted on Qualcomm's public S3 bucket (not Hugging Face)
|
| 69 |
+
# ---------------------------------------------------------------------------
|
| 70 |
+
FOMM_URL = "https://qaihub-public-assets.s3.us-west-2.amazonaws.com/qai-hub-models/models/fomm/releases/v0.56.0/fomm-onnx-float.zip"
|
| 71 |
+
FOMM_SENTINEL = "fomm/generator.onnx" # Presence of this file means extraction is complete
|
| 72 |
+
|
| 73 |
+
# ---------------------------------------------------------------------------
|
| 74 |
+
# Thin-Plate Spline Motion Model (TPSMM)
|
| 75 |
+
# ---------------------------------------------------------------------------
|
| 76 |
+
# NOTE: The original vessl/thin-plate-spline-motion-model HF repo is gone (404).
|
| 77 |
+
# AlekseyKorshuk's HF Space mirrors the exact checkpoint (SHA-256 pinned below,
|
| 78 |
+
# verified bit-for-bit identical to the known-good local copy).
|
| 79 |
+
TPSMM_HF_REPO = "AlekseyKorshuk/thin-plate-spline-motion-model"
|
| 80 |
+
TPSMM_FILE = "checkpoints/vox.pth.tar"
|
| 81 |
+
TPSMM_SENTINEL = "tpsmm/vox.pth.tar"
|
| 82 |
+
TPSMM_SHA256 = "52ad8c848e2a1d91b621de96fea83faf57ce3b8c1c06424e317f4df1d3998204"
|
| 83 |
+
|
| 84 |
+
# ---------------------------------------------------------------------------
|
| 85 |
+
# ArcaneGAN Model Manifest
|
| 86 |
+
# ---------------------------------------------------------------------------
|
| 87 |
+
ARCANEGAN_HF_REPO = "akhaliq/ArcaneGANv0.4"
|
| 88 |
+
ARCANEGAN_FILE = "ArcaneGANv0.4.jit"
|
| 89 |
+
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
|
| 93 |
+
def create_progress_hook(name="Model"):
|
| 94 |
+
"""Creates a simple download progress reporter for urllib.request."""
|
| 95 |
+
def _hook(block_num, block_size, total_size):
|
| 96 |
+
if total_size > 0:
|
| 97 |
+
downloaded = block_num * block_size
|
| 98 |
+
pct = min(100.0, downloaded / total_size * 100.0)
|
| 99 |
+
print(f"\r Downloading {name}... {pct:.1f}%", end="", flush=True)
|
| 100 |
+
return _hook
|
| 101 |
+
|
| 102 |
+
|
| 103 |
+
def sync_cache():
|
| 104 |
+
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
| 105 |
+
os.makedirs(base_dir, exist_ok=True)
|
| 106 |
+
|
| 107 |
+
# ------------------------------------------------------------------
|
| 108 |
+
# 1. Download base ONNX models from Hugging Face
|
| 109 |
+
# ------------------------------------------------------------------
|
| 110 |
+
for key, info in MANIFEST.items():
|
| 111 |
+
target = os.path.join(base_dir, info["target"])
|
| 112 |
+
if os.path.exists(target):
|
| 113 |
+
print(f" [SKIP] {key} already present.")
|
| 114 |
+
continue
|
| 115 |
+
print(f" [SYNC] Downloading {key}...")
|
| 116 |
+
path = hf_hub_download(
|
| 117 |
+
repo_id=info["repo"],
|
| 118 |
+
filename=info["file"],
|
| 119 |
+
local_dir=base_dir,
|
| 120 |
+
local_dir_use_symlinks=False
|
| 121 |
+
)
|
| 122 |
+
if path != target and os.path.exists(path):
|
| 123 |
+
shutil.move(path, target)
|
| 124 |
+
|
| 125 |
+
# ------------------------------------------------------------------
|
| 126 |
+
# 2. Download First-Order Motion Model (FOMM) from Qualcomm S3
|
| 127 |
+
# Local-first: skip if sentinel file already exists.
|
| 128 |
+
# ------------------------------------------------------------------
|
| 129 |
+
fomm_sentinel = os.path.join(base_dir, FOMM_SENTINEL)
|
| 130 |
+
fomm_dir = os.path.join(base_dir, "fomm")
|
| 131 |
+
if os.path.exists(fomm_sentinel):
|
| 132 |
+
print(" [SKIP] FOMM models already present.")
|
| 133 |
+
else:
|
| 134 |
+
print(" [SYNC] Downloading FOMM ONNX package from Qualcomm S3...")
|
| 135 |
+
zip_path = os.path.join(base_dir, "_fomm_tmp.zip")
|
| 136 |
+
try:
|
| 137 |
+
urllib.request.urlretrieve(FOMM_URL, zip_path, reporthook=create_progress_hook("FOMM"))
|
| 138 |
+
print() # newline after progress
|
| 139 |
+
print(" [EXTRACT] Unpacking FOMM...")
|
| 140 |
+
os.makedirs(fomm_dir, exist_ok=True)
|
| 141 |
+
with zipfile.ZipFile(zip_path, "r") as zf:
|
| 142 |
+
for member in zf.infolist():
|
| 143 |
+
# Strip the top-level directory from the zip path
|
| 144 |
+
parts = member.filename.split("/", 1)
|
| 145 |
+
if len(parts) < 2 or not parts[1]:
|
| 146 |
+
continue # skip the root directory entry itself
|
| 147 |
+
dest = os.path.join(fomm_dir, parts[1])
|
| 148 |
+
if member.is_dir():
|
| 149 |
+
os.makedirs(dest, exist_ok=True)
|
| 150 |
+
else:
|
| 151 |
+
with zf.open(member) as src, open(dest, "wb") as out:
|
| 152 |
+
shutil.copyfileobj(src, out)
|
| 153 |
+
print(f" [OK] FOMM models extracted to: {fomm_dir}")
|
| 154 |
+
finally:
|
| 155 |
+
if os.path.exists(zip_path):
|
| 156 |
+
os.remove(zip_path)
|
| 157 |
+
|
| 158 |
+
# ------------------------------------------------------------------
|
| 159 |
+
# 3. Clone DreamTalk repository if not already present
|
| 160 |
+
# ------------------------------------------------------------------
|
| 161 |
+
dreamtalk_src = os.path.join(third_party_dir, "dreamtalk_src")
|
| 162 |
+
if not os.path.exists(dreamtalk_src):
|
| 163 |
+
print(" [SYNC] Cloning DreamTalk repository...")
|
| 164 |
+
subprocess.run(["git", "clone", "https://github.com/camenduru/dreamtalk.git", dreamtalk_src], check=True)
|
| 165 |
+
else:
|
| 166 |
+
print(" [SKIP] DreamTalk repository already cloned.")
|
| 167 |
+
|
| 168 |
+
# Mark vendored import roots as regular packages. Upstream has no __init__.py,
|
| 169 |
+
# so its 'core'/'configs'/'generators' are PEP 420 namespace packages and lose to
|
| 170 |
+
# identically-named regular packages (e.g. captiona/core, sonora/core) that other
|
| 171 |
+
# engines place on sys.path. A regular package at the front of sys.path always wins.
|
| 172 |
+
for _pkg_dir in ("core", "configs", "generators"):
|
| 173 |
+
_init = os.path.join(dreamtalk_src, _pkg_dir, "__init__.py")
|
| 174 |
+
if not os.path.exists(_init):
|
| 175 |
+
with open(_init, "w") as f:
|
| 176 |
+
f.write("# vendored package marker: prevents namespace-package collision\n")
|
| 177 |
+
|
| 178 |
+
# ------------------------------------------------------------------
|
| 179 |
+
# 4. Download DreamTalk PyTorch checkpoints
|
| 180 |
+
# ------------------------------------------------------------------
|
| 181 |
+
checkpoints_dir = os.path.join(dreamtalk_src, "checkpoints")
|
| 182 |
+
os.makedirs(checkpoints_dir, exist_ok=True)
|
| 183 |
+
|
| 184 |
+
dreamtalk_checkpoints = {
|
| 185 |
+
"denoising_network.pth": "damo/dreamtalk/checkpoints/denoising_network.pth",
|
| 186 |
+
"renderer.pt": "damo/dreamtalk/checkpoints/renderer.pt",
|
| 187 |
+
}
|
| 188 |
+
for filename, hf_path in dreamtalk_checkpoints.items():
|
| 189 |
+
target_path = os.path.join(checkpoints_dir, filename)
|
| 190 |
+
if os.path.exists(target_path):
|
| 191 |
+
print(f" [SKIP] DreamTalk checkpoint {filename} already present.")
|
| 192 |
+
continue
|
| 193 |
+
print(f" [SYNC] Downloading DreamTalk checkpoint: {filename}")
|
| 194 |
+
downloaded = hf_hub_download(
|
| 195 |
+
repo_id="impactframes/dreamtalk",
|
| 196 |
+
filename=hf_path,
|
| 197 |
+
local_dir=base_dir,
|
| 198 |
+
local_dir_use_symlinks=False
|
| 199 |
+
)
|
| 200 |
+
shutil.move(downloaded, target_path)
|
| 201 |
+
|
| 202 |
+
# ------------------------------------------------------------------
|
| 203 |
+
# 4b. Style-clip & pose .mat assets (required by DreamTalkEngine).
|
| 204 |
+
# Not part of the upstream git repo — distributed separately by the
|
| 205 |
+
# DreamTalk authors. Fetched from a verified public mirror; only the
|
| 206 |
+
# two files dreamtalk.py actually references are needed.
|
| 207 |
+
# (Uses the module-level `import urllib.request` — a function-local
|
| 208 |
+
# import here would shadow it for the whole scope and crash the
|
| 209 |
+
# earlier FOMM download with UnboundLocalError.)
|
| 210 |
+
# ------------------------------------------------------------------
|
| 211 |
+
_dt_asset_urls = {
|
| 212 |
+
"data/style_clip/3DMM/M030_front_neutral_level1_001.mat":
|
| 213 |
+
"https://storage.googleapis.com/falserverless/model_tests/dream_talk/style_clip/3DMM/M030_front_neutral_level1_001.mat",
|
| 214 |
+
"data/pose/RichardShelby_front_neutral_level1_001.mat":
|
| 215 |
+
"https://storage.googleapis.com/falserverless/model_tests/dream_talk/pose/RichardShelby_front_neutral_level1_001.mat",
|
| 216 |
+
}
|
| 217 |
+
for _rel, _url in _dt_asset_urls.items():
|
| 218 |
+
_dest = os.path.join(dreamtalk_src, _rel)
|
| 219 |
+
if os.path.exists(_dest):
|
| 220 |
+
print(f" [SKIP] DreamTalk asset {_rel} already present.")
|
| 221 |
+
continue
|
| 222 |
+
try:
|
| 223 |
+
print(f" [SYNC] Downloading DreamTalk asset: {_rel}")
|
| 224 |
+
os.makedirs(os.path.dirname(_dest), exist_ok=True)
|
| 225 |
+
urllib.request.urlretrieve(_url, _dest)
|
| 226 |
+
print(f" [OK] Saved to {_dest}")
|
| 227 |
+
except Exception as e:
|
| 228 |
+
print(f" [WARN] Could not download DreamTalk asset {_rel}: {e}")
|
| 229 |
+
|
| 230 |
+
# ------------------------------------------------------------------
|
| 231 |
+
# 5. Preload Wav2Vec2 audio feature extractor (required by DreamTalkEngine)
|
| 232 |
+
# Must land in the exact cache layout dreamtalk.py expects:
|
| 233 |
+
# <engine_root>/models/.cache/huggingface
|
| 234 |
+
# ------------------------------------------------------------------
|
| 235 |
+
print(" [SYNC] Preloading Wav2Vec2 audio feature extractor...")
|
| 236 |
+
try:
|
| 237 |
+
from transformers import Wav2Vec2Model, Wav2Vec2Processor
|
| 238 |
+
_wt_cache = os.path.join(base_dir, ".cache", "huggingface")
|
| 239 |
+
Wav2Vec2Processor.from_pretrained(
|
| 240 |
+
"jonatasgrosman/wav2vec2-large-xlsr-53-english", cache_dir=_wt_cache
|
| 241 |
+
)
|
| 242 |
+
Wav2Vec2Model.from_pretrained(
|
| 243 |
+
"jonatasgrosman/wav2vec2-large-xlsr-53-english", cache_dir=_wt_cache
|
| 244 |
+
)
|
| 245 |
+
print(" [OK] Wav2Vec2 cached.")
|
| 246 |
+
except Exception as e:
|
| 247 |
+
print(f" [WARN] Could not preload Wav2Vec2 (runtime will retry online): {e}")
|
| 248 |
+
|
| 249 |
+
print("\nPre-caching successfully completed.")
|
| 250 |
+
|
| 251 |
+
|
| 252 |
+
# Main execution logic moved to bottom
|
| 253 |
+
def sync_tpsmm():
|
| 254 |
+
"""
|
| 255 |
+
Downloads the TPSMM checkpoint (vox.pth.tar) into models/tpsmm/
|
| 256 |
+
and clones the Thin-Plate-Spline-Motion-Model source repo into
|
| 257 |
+
actora/tpsmm_src/.
|
| 258 |
+
"""
|
| 259 |
+
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
| 260 |
+
tpsmm_dir = os.path.join(base_dir, "tpsmm")
|
| 261 |
+
os.makedirs(tpsmm_dir, exist_ok=True)
|
| 262 |
+
|
| 263 |
+
sentinel = os.path.join(base_dir, TPSMM_SENTINEL)
|
| 264 |
+
if os.path.exists(sentinel):
|
| 265 |
+
print(" [SKIP] TPSMM model checkpoint already present.")
|
| 266 |
+
else:
|
| 267 |
+
print(f" [SYNC] Downloading TPSMM model checkpoint from {TPSMM_HF_REPO} (space)...")
|
| 268 |
+
try:
|
| 269 |
+
downloaded = hf_hub_download(
|
| 270 |
+
repo_id=TPSMM_HF_REPO,
|
| 271 |
+
filename=TPSMM_FILE,
|
| 272 |
+
repo_type="space",
|
| 273 |
+
local_dir=tpsmm_dir,
|
| 274 |
+
)
|
| 275 |
+
dest = os.path.join(tpsmm_dir, "vox.pth.tar")
|
| 276 |
+
if downloaded != dest and os.path.exists(downloaded):
|
| 277 |
+
shutil.move(downloaded, dest)
|
| 278 |
+
# Integrity check: reject corrupt/incomplete downloads
|
| 279 |
+
import hashlib
|
| 280 |
+
h = hashlib.sha256()
|
| 281 |
+
with open(dest, "rb") as f:
|
| 282 |
+
for chunk in iter(lambda: f.read(1 << 20), b""):
|
| 283 |
+
h.update(chunk)
|
| 284 |
+
if h.hexdigest() != TPSMM_SHA256:
|
| 285 |
+
os.remove(dest)
|
| 286 |
+
raise ValueError(f"SHA-256 mismatch: expected {TPSMM_SHA256}, got {h.hexdigest()}")
|
| 287 |
+
print(f" [OK] TPSMM model checkpoint verified and saved to: {dest}")
|
| 288 |
+
except Exception as e:
|
| 289 |
+
print(f" [WARN] Could not download TPSMM checkpoint ({TPSMM_HF_REPO}): {e}")
|
| 290 |
+
|
| 291 |
+
# Clone the Thin-Plate Spline Motion Model source repository
|
| 292 |
+
tpsmm_src = os.path.join(third_party_dir, "tpsmm_src")
|
| 293 |
+
if os.path.exists(tpsmm_src):
|
| 294 |
+
print(" [SKIP] TPSMM source repo already cloned.")
|
| 295 |
+
else:
|
| 296 |
+
print(" [SYNC] Cloning TPSMM source repo (shallow)...")
|
| 297 |
+
subprocess.run(
|
| 298 |
+
["git", "clone", "--depth", "1",
|
| 299 |
+
"https://github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model.git", tpsmm_src],
|
| 300 |
+
check=True
|
| 301 |
+
)
|
| 302 |
+
|
| 303 |
+
# Same namespace-package vaccine as DreamTalk: tpsmm.py imports bare 'modules.*'
|
| 304 |
+
_modules_init = os.path.join(tpsmm_src, "modules", "__init__.py")
|
| 305 |
+
if not os.path.exists(_modules_init):
|
| 306 |
+
with open(_modules_init, "w") as f:
|
| 307 |
+
f.write("# vendored package marker: prevents namespace-package collision\n")
|
| 308 |
+
|
| 309 |
+
print(" [OK] TPSMM setup complete.")
|
| 310 |
+
|
| 311 |
+
def sync_arcanegan():
|
| 312 |
+
"""
|
| 313 |
+
Downloads ArcaneGAN v0.4 JIT weights from Hugging Face.
|
| 314 |
+
"""
|
| 315 |
+
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
| 316 |
+
arcanegan_dir = os.path.join(base_dir, "arcanegan")
|
| 317 |
+
os.makedirs(arcanegan_dir, exist_ok=True)
|
| 318 |
+
target = os.path.join(arcanegan_dir, ARCANEGAN_FILE)
|
| 319 |
+
if os.path.exists(target):
|
| 320 |
+
print(f" [SKIP] arcanegan/{ARCANEGAN_FILE} already present.")
|
| 321 |
+
return
|
| 322 |
+
print(f" [SYNC] Downloading ArcaneGAN from {ARCANEGAN_HF_REPO}...")
|
| 323 |
+
downloaded = hf_hub_download(
|
| 324 |
+
repo_id=ARCANEGAN_HF_REPO,
|
| 325 |
+
filename=ARCANEGAN_FILE,
|
| 326 |
+
local_dir=arcanegan_dir,
|
| 327 |
+
local_dir_use_symlinks=False,
|
| 328 |
+
)
|
| 329 |
+
if downloaded != target and os.path.exists(downloaded):
|
| 330 |
+
shutil.move(downloaded, target)
|
| 331 |
+
print(f" [OK] Saved to {target}")
|
| 332 |
+
|
| 333 |
+
def sync_dctnet():
|
| 334 |
+
"""
|
| 335 |
+
Downloads DCT-Net models (Artstyle, 3D, Anime) from ModelScope.
|
| 336 |
+
"""
|
| 337 |
+
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
| 338 |
+
|
| 339 |
+
# Check if models are already downloaded to skip network access
|
| 340 |
+
models_to_check = [
|
| 341 |
+
"dctnet/artstyle/damo/cv_unet_person-image-cartoon-artstyle_compound-models",
|
| 342 |
+
"dctnet/3d/damo/cv_unet_person-image-cartoon-3d_compound-models",
|
| 343 |
+
"dctnet/anime/damo/cv_unet_person-image-cartoon_compound-models"
|
| 344 |
+
]
|
| 345 |
+
|
| 346 |
+
all_present = True
|
| 347 |
+
for path in models_to_check:
|
| 348 |
+
full_path = os.path.join(base_dir, path)
|
| 349 |
+
if not os.path.exists(full_path) or len(os.listdir(full_path)) == 0:
|
| 350 |
+
all_present = False
|
| 351 |
+
break
|
| 352 |
+
|
| 353 |
+
if all_present:
|
| 354 |
+
print(" [SKIP] DCT-Net models already present.")
|
| 355 |
+
return
|
| 356 |
+
|
| 357 |
+
try:
|
| 358 |
+
from modelscope.hub.snapshot_download import snapshot_download
|
| 359 |
+
print("\nPreloading DCT-Net Models (ModelScope)...")
|
| 360 |
+
# Artstyle
|
| 361 |
+
snapshot_download('damo/cv_unet_person-image-cartoon-artstyle_compound-models', cache_dir=os.path.join(base_dir, 'dctnet', 'artstyle'))
|
| 362 |
+
# 3D
|
| 363 |
+
snapshot_download('damo/cv_unet_person-image-cartoon-3d_compound-models', cache_dir=os.path.join(base_dir, 'dctnet', '3d'))
|
| 364 |
+
# Anime
|
| 365 |
+
snapshot_download('damo/cv_unet_person-image-cartoon_compound-models', cache_dir=os.path.join(base_dir, 'dctnet', 'anime'))
|
| 366 |
+
print(" [OK] DCT-Net Models downloaded.")
|
| 367 |
+
except ImportError as e:
|
| 368 |
+
print(f" [WARNING] Could not import modelscope. Skipping DCT-Net download: {e}")
|
| 369 |
+
|
| 370 |
+
|
| 371 |
+
def sync_faster_liveportrait():
|
| 372 |
+
"""
|
| 373 |
+
Downloads the FasterLivePortrait checkpoints and clones the repo.
|
| 374 |
+
"""
|
| 375 |
+
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
| 376 |
+
flip_dir = os.path.join(base_dir, "faster_liveportrait")
|
| 377 |
+
os.makedirs(flip_dir, exist_ok=True)
|
| 378 |
+
|
| 379 |
+
# Clone the FasterLivePortrait source repository
|
| 380 |
+
flip_src = os.path.join(third_party_dir, "faster_liveportrait_src")
|
| 381 |
+
if os.path.exists(flip_src):
|
| 382 |
+
print(" [SKIP] FasterLivePortrait source repo already cloned.")
|
| 383 |
+
else:
|
| 384 |
+
print(" [SYNC] Cloning FasterLivePortrait source repo (shallow)...")
|
| 385 |
+
subprocess.run(
|
| 386 |
+
["git", "clone", "--depth", "1",
|
| 387 |
+
"https://github.com/warmshao/FasterLivePortrait.git", flip_src],
|
| 388 |
+
check=True
|
| 389 |
+
)
|
| 390 |
+
|
| 391 |
+
# Download human ONNX weights
|
| 392 |
+
onnx_files = [
|
| 393 |
+
"appearance_feature_extractor.onnx",
|
| 394 |
+
"face_2dpose_106_static.onnx",
|
| 395 |
+
"landmark.onnx",
|
| 396 |
+
"motion_extractor.onnx",
|
| 397 |
+
"retinaface_det_static.onnx",
|
| 398 |
+
"stitching.onnx",
|
| 399 |
+
"stitching_eye.onnx",
|
| 400 |
+
"stitching_lip.onnx",
|
| 401 |
+
"warping_spade.onnx"
|
| 402 |
+
]
|
| 403 |
+
|
| 404 |
+
target_dir = os.path.join(flip_dir, "liveportrait_onnx")
|
| 405 |
+
os.makedirs(target_dir, exist_ok=True)
|
| 406 |
+
|
| 407 |
+
files_to_download = [f for f in onnx_files if not os.path.exists(os.path.join(target_dir, f))]
|
| 408 |
+
|
| 409 |
+
if not files_to_download:
|
| 410 |
+
print(" [SKIP] FasterLivePortrait ONNX weights already present.")
|
| 411 |
+
else:
|
| 412 |
+
print(f" [SYNC] Downloading {len(files_to_download)} missing FasterLivePortrait ONNX weights...")
|
| 413 |
+
for f in files_to_download:
|
| 414 |
+
print(f" [SYNC] Downloading {f}...")
|
| 415 |
+
downloaded = hf_hub_download(
|
| 416 |
+
repo_id="warmshao/FasterLivePortrait",
|
| 417 |
+
filename=f"liveportrait_onnx/{f}",
|
| 418 |
+
local_dir=flip_dir,
|
| 419 |
+
local_dir_use_symlinks=False,
|
| 420 |
+
)
|
| 421 |
+
dest = os.path.join(target_dir, f)
|
| 422 |
+
if downloaded != dest and os.path.exists(downloaded):
|
| 423 |
+
shutil.move(downloaded, dest)
|
| 424 |
+
|
| 425 |
+
print(" [OK] FasterLivePortrait setup complete.")
|
| 426 |
+
|
| 427 |
+
|
| 428 |
+
|
| 429 |
+
|
| 430 |
+
def sync_postprocess():
|
| 431 |
+
"""
|
| 432 |
+
Downloads post-processing super-resolution models (GRL-GAN and Real-ESRGAN).
|
| 433 |
+
"""
|
| 434 |
+
base_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "models")
|
| 435 |
+
postprocess_dir = os.path.join(base_dir, "postprocess")
|
| 436 |
+
os.makedirs(postprocess_dir, exist_ok=True)
|
| 437 |
+
|
| 438 |
+
# 1. GRL-GAN
|
| 439 |
+
grl_gan_target = os.path.join(postprocess_dir, "grl_gan.onnx")
|
| 440 |
+
if os.path.exists(grl_gan_target):
|
| 441 |
+
print(" [SKIP] postprocess/grl_gan.onnx already present.")
|
| 442 |
+
else:
|
| 443 |
+
print(" [SYNC] Downloading 4x_APISR_GRL_GAN_generator-onnx...")
|
| 444 |
+
downloaded = hf_hub_download(
|
| 445 |
+
repo_id="Xenova/4x_APISR_GRL_GAN_generator-onnx",
|
| 446 |
+
filename="onnx/model.onnx",
|
| 447 |
+
local_dir=postprocess_dir,
|
| 448 |
+
local_dir_use_symlinks=False,
|
| 449 |
+
)
|
| 450 |
+
if downloaded != grl_gan_target and os.path.exists(downloaded):
|
| 451 |
+
shutil.move(downloaded, grl_gan_target)
|
| 452 |
+
onnx_subdir = os.path.join(postprocess_dir, "onnx")
|
| 453 |
+
if os.path.exists(onnx_subdir):
|
| 454 |
+
shutil.rmtree(onnx_subdir)
|
| 455 |
+
print(f" [OK] Saved to {grl_gan_target}")
|
| 456 |
+
|
| 457 |
+
# 2. Real-ESRGAN
|
| 458 |
+
real_esrgan_url = "https://qaihub-public-assets.s3.us-west-2.amazonaws.com/qai-hub-models/models/real_esrgan_x4plus/releases/v0.57.0/real_esrgan_x4plus-onnx-float.zip"
|
| 459 |
+
real_esrgan_dir = os.path.join(postprocess_dir, "real_esrgan")
|
| 460 |
+
real_esrgan_sentinel = os.path.join(real_esrgan_dir, "real_esrgan_x4plus.onnx")
|
| 461 |
+
|
| 462 |
+
if os.path.exists(real_esrgan_sentinel):
|
| 463 |
+
print(" [SKIP] postprocess/real_esrgan already present.")
|
| 464 |
+
else:
|
| 465 |
+
print(" [SYNC] Downloading Real-ESRGAN ONNX package from Qualcomm S3...")
|
| 466 |
+
zip_path = os.path.join(postprocess_dir, "_realesrgan_tmp.zip")
|
| 467 |
+
try:
|
| 468 |
+
urllib.request.urlretrieve(real_esrgan_url, zip_path, reporthook=create_progress_hook("Real-ESRGAN"))
|
| 469 |
+
print()
|
| 470 |
+
print(" [EXTRACT] Unpacking Real-ESRGAN...")
|
| 471 |
+
os.makedirs(real_esrgan_dir, exist_ok=True)
|
| 472 |
+
with zipfile.ZipFile(zip_path, "r") as zf:
|
| 473 |
+
for member in zf.infolist():
|
| 474 |
+
parts = member.filename.split("/", 1)
|
| 475 |
+
if len(parts) < 2 or not parts[1]:
|
| 476 |
+
continue
|
| 477 |
+
dest = os.path.join(real_esrgan_dir, parts[1])
|
| 478 |
+
if member.is_dir():
|
| 479 |
+
os.makedirs(dest, exist_ok=True)
|
| 480 |
+
else:
|
| 481 |
+
with zf.open(member) as src, open(dest, "wb") as out:
|
| 482 |
+
shutil.copyfileobj(src, out)
|
| 483 |
+
print(f" [OK] Real-ESRGAN extracted to: {real_esrgan_dir}")
|
| 484 |
+
finally:
|
| 485 |
+
if os.path.exists(zip_path):
|
| 486 |
+
os.remove(zip_path)
|
| 487 |
+
|
| 488 |
+
|
| 489 |
+
if __name__ == "__main__":
|
| 490 |
+
sync_cache()
|
| 491 |
+
sync_tpsmm()
|
| 492 |
+
sync_faster_liveportrait()
|
| 493 |
+
sync_arcanegan()
|
| 494 |
+
sync_dctnet()
|
| 495 |
+
sync_postprocess()
|
| 496 |
+
|
actora/run_actora.sh
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
#!/bin/bash
|
| 2 |
+
# TL;DR: Docker lifecycle management script for the Talking Head edge pipeline.
|
| 3 |
+
|
| 4 |
+
#!/bin/bash
|
| 5 |
+
case "$1" in
|
| 6 |
+
build)
|
| 7 |
+
docker compose -f docker/docker-compose.yml build
|
| 8 |
+
;;
|
| 9 |
+
up)
|
| 10 |
+
docker compose -f docker/docker-compose.yml up -d
|
| 11 |
+
docker exec -t Actora python /app/preload_models.py
|
| 12 |
+
;;
|
| 13 |
+
down)
|
| 14 |
+
docker compose -f docker/docker-compose.yml down
|
| 15 |
+
;;
|
| 16 |
+
clean)
|
| 17 |
+
docker system prune -f
|
| 18 |
+
docker exec -t Actora chown -R $(id -u):$(id -g) /app 2>/dev/null || true
|
| 19 |
+
rm -rf actora/actora/test/output/*
|
| 20 |
+
;;
|
| 21 |
+
test)
|
| 22 |
+
docker exec -it Actora python /app/actora/test/test.py
|
| 23 |
+
docker exec -t Actora chown -R $(id -u):$(id -g) /app 2>/dev/null || true
|
| 24 |
+
;;
|
| 25 |
+
*)
|
| 26 |
+
echo "Usage: $0 {build|up|down|clean|test}"
|
| 27 |
+
exit 1
|
| 28 |
+
esac
|
actora/third_party/dreamtalk_src/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
MIT License
|
| 2 |
+
|
| 3 |
+
Copyright (c) 2023 Alibaba TongYi Vision Intelligence Lab
|
| 4 |
+
|
| 5 |
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
| 6 |
+
of this software and associated documentation files (the "Software"), to deal
|
| 7 |
+
in the Software without restriction, including without limitation the rights
|
| 8 |
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
| 9 |
+
copies of the Software, and to permit persons to whom the Software is
|
| 10 |
+
furnished to do so, subject to the following conditions:
|
| 11 |
+
|
| 12 |
+
The above copyright notice and this permission notice shall be included in all
|
| 13 |
+
copies or substantial portions of the Software.
|
| 14 |
+
|
| 15 |
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
| 16 |
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
| 17 |
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
| 18 |
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
| 19 |
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
| 20 |
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
| 21 |
+
SOFTWARE.
|
actora/third_party/dreamtalk_src/README.md
ADDED
|
@@ -0,0 +1,99 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
<h2 align="center">DreamTalk: When Expressive Talking Head Generation <br> Meets Diffusion Probabilistic Models</h2>
|
| 2 |
+
<p align="center">
|
| 3 |
+
<a href='https://dreamtalk-project.github.io/'><img src='https://img.shields.io/badge/Project-Page-Green'></a> <a href='https://arxiv.org/abs/2312.09767'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a> <a href='https://youtu.be/VF4vlE6ZqWQ'><img src='https://badges.aleen42.com/src/youtube.svg'></a>
|
| 4 |
+
</p>
|
| 5 |
+
|
| 6 |
+

|
| 7 |
+
|
| 8 |
+
DreamTalk is a diffusion-based audio-driven expressive talking head generation framework that can produce high-quality talking head videos across diverse speaking styles. DreamTalk exhibits robust performance with a diverse array of inputs, including songs, speech in multiple languages, noisy audio, and out-of-domain portraits.
|
| 9 |
+
|
| 10 |
+
## News
|
| 11 |
+
- __[2023.12]__ Release inference code and pretrained checkpoint.
|
| 12 |
+
|
| 13 |
+
## Installation
|
| 14 |
+
|
| 15 |
+
```
|
| 16 |
+
conda create -n dreamtalk python=3.7.0
|
| 17 |
+
conda activate dreamtalk
|
| 18 |
+
pip install -r requirements.txt
|
| 19 |
+
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.1 -c pytorch -c conda-forge
|
| 20 |
+
conda update ffmpeg
|
| 21 |
+
|
| 22 |
+
pip install urllib3==1.26.6
|
| 23 |
+
pip install transformers==4.28.1
|
| 24 |
+
pip install dlib
|
| 25 |
+
```
|
| 26 |
+
|
| 27 |
+
## Download Checkpoints
|
| 28 |
+
Download the checkpoint of the denoising network and the renderer:
|
| 29 |
+
* [HuggingFace](https://huggingface.co/damo-vilab/dreamtalk)
|
| 30 |
+
* [ModelScope](https://modelscope.cn/models/damo/dreamtalk/files) (in `checkpoints` folder)
|
| 31 |
+
|
| 32 |
+
|
| 33 |
+
Put the downloaded checkpoints into `checkpoints` folder.
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
## Inference
|
| 37 |
+
Run the script:
|
| 38 |
+
|
| 39 |
+
```
|
| 40 |
+
python inference_for_demo_video.py \
|
| 41 |
+
--wav_path data/audio/acknowledgement_english.m4a \
|
| 42 |
+
--style_clip_path data/style_clip/3DMM/M030_front_neutral_level1_001.mat \
|
| 43 |
+
--pose_path data/pose/RichardShelby_front_neutral_level1_001.mat \
|
| 44 |
+
--image_path data/src_img/uncropped/male_face.png \
|
| 45 |
+
--cfg_scale 1.0 \
|
| 46 |
+
--max_gen_len 30 \
|
| 47 |
+
--output_name acknowledgement_english@M030_front_neutral_level1_001@male_face
|
| 48 |
+
```
|
| 49 |
+
|
| 50 |
+
`wav_path` specifies the input audio. The input audio file extensions such as wav, mp3, m4a, and mp4 (video with sound) should all be compatible.
|
| 51 |
+
|
| 52 |
+
`style_clip_path` specifies the reference speaking style and `pose_path` specifies head pose. They are 3DMM paramenter sequences extracted from reference videos. You can follow [PIRenderer](https://github.com/RenYurui/PIRender) to extract 3DMM parameters from your own videos. Note that the video frame rate should be 25 FPS. Besides, videos used for head pose reference should be first cropped to $256\times256$ using scripts in [FOMM video preprocessing](https://github.com/AliaksandrSiarohin/video-preprocessing).
|
| 53 |
+
|
| 54 |
+
`image_path` specifies the input portrait. Its resolution should be larger than $256\times256$. Frontal portraits, with the face directly facing forward and not tilted to one side, usually achieve satisfactory results. The input portrait will be cropped to $256\times256$. If your portrait is already cropped to $256\times256$ and you want to disable cropping, use option `--disable_img_crop` like this:
|
| 55 |
+
|
| 56 |
+
```
|
| 57 |
+
python inference_for_demo_video.py \
|
| 58 |
+
--wav_path data/audio/acknowledgement_chinese.m4a \
|
| 59 |
+
--style_clip_path data/style_clip/3DMM/M030_front_surprised_level3_001.mat \
|
| 60 |
+
--pose_path data/pose/RichardShelby_front_neutral_level1_001.mat \
|
| 61 |
+
--image_path data/src_img/cropped/zp1.png \
|
| 62 |
+
--disable_img_crop \
|
| 63 |
+
--cfg_scale 1.0 \
|
| 64 |
+
--max_gen_len 30 \
|
| 65 |
+
--output_name acknowledgement_chinese@M030_front_surprised_level3_001@zp1
|
| 66 |
+
```
|
| 67 |
+
|
| 68 |
+
`cfg_scale` controls the scale of classifer-free guidance. It can adjust the intensity of speaking styles.
|
| 69 |
+
|
| 70 |
+
`max_gen_len` is the maximum video generation duration, measured in seconds. If the input audio exceeds this length, it will be truncated.
|
| 71 |
+
|
| 72 |
+
The generated video will be named `$(output_name).mp4` and put in the output_video folder. Intermediate results, including the cropped portrait, will be in the `tmp/$(output_name)` folder.
|
| 73 |
+
|
| 74 |
+
Sample inputs are presented in `data` folder. Due to copyright issues, we are unable to include the songs we have used in this folder.
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
## Acknowledgements
|
| 78 |
+
|
| 79 |
+
We extend our heartfelt thanks for the invaluable contributions made by preceding works to the development of DreamTalk. This includes, but is not limited to:
|
| 80 |
+
[PIRenderer](https://github.com/RenYurui/PIRender)
|
| 81 |
+
,[AVCT](https://github.com/FuxiVirtualHuman/AAAI22-one-shot-talking-face)
|
| 82 |
+
,[StyleTalk](https://github.com/FuxiVirtualHuman/styletalk)
|
| 83 |
+
,[Deep3DFaceRecon_pytorch](https://github.com/sicxu/Deep3DFaceRecon_pytorch)
|
| 84 |
+
,[Wav2vec2.0](https://huggingface.co/jonatasgrosman/wav2vec2-large-xlsr-53-english)
|
| 85 |
+
,[diffusion-point-cloud](https://github.com/luost26/diffusion-point-cloud)
|
| 86 |
+
,[FOMM video preprocessing](https://github.com/AliaksandrSiarohin/video-preprocessing). We are dedicated to advancing upon these foundational works with the utmost respect for their original contributions.
|
| 87 |
+
|
| 88 |
+
## Citation
|
| 89 |
+
If you find this codebase useful for your research, please use the following entry.
|
| 90 |
+
```BibTeX
|
| 91 |
+
@article{ma2023dreamtalk,
|
| 92 |
+
title={DreamTalk: When Expressive Talking Head Generation Meets Diffusion Probabilistic Models},
|
| 93 |
+
author={Ma, Yifeng and Zhang, Shiwei and Wang, Jiayu and Wang, Xiang and Zhang, Yingya and Deng, Zhidong},
|
| 94 |
+
journal={arXiv preprint arXiv:2312.09767},
|
| 95 |
+
year={2023}
|
| 96 |
+
}
|
| 97 |
+
```
|
| 98 |
+
|
| 99 |
+
|
actora/third_party/dreamtalk_src/configs/__init__.py
ADDED
|
File without changes
|
actora/third_party/dreamtalk_src/configs/default.py
ADDED
|
@@ -0,0 +1,91 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from yacs.config import CfgNode as CN
|
| 2 |
+
|
| 3 |
+
|
| 4 |
+
_C = CN()
|
| 5 |
+
_C.TAG = "style_id_emotion"
|
| 6 |
+
_C.DECODER_TYPE = "DisentangleDecoder"
|
| 7 |
+
_C.CONTENT_ENCODER_TYPE = "ContentW2VEncoder"
|
| 8 |
+
_C.STYLE_ENCODER_TYPE = "StyleEncoder"
|
| 9 |
+
|
| 10 |
+
_C.DIFFNET_TYPE = "DiffusionNet"
|
| 11 |
+
|
| 12 |
+
_C.WIN_SIZE = 5
|
| 13 |
+
_C.D_MODEL = 256
|
| 14 |
+
|
| 15 |
+
_C.DATASET = CN()
|
| 16 |
+
_C.DATASET.FACE3D_DIM = 64
|
| 17 |
+
_C.DATASET.NUM_FRAMES = 64
|
| 18 |
+
_C.DATASET.STYLE_MAX_LEN = 256
|
| 19 |
+
|
| 20 |
+
_C.TRAIN = CN()
|
| 21 |
+
_C.TRAIN.FACE3D_LATENT = CN()
|
| 22 |
+
_C.TRAIN.FACE3D_LATENT.TYPE = "face3d"
|
| 23 |
+
|
| 24 |
+
_C.DIFFUSION = CN()
|
| 25 |
+
_C.DIFFUSION.PREDICT_WHAT = "x0" # noise | x0
|
| 26 |
+
_C.DIFFUSION.SCHEDULE = CN()
|
| 27 |
+
_C.DIFFUSION.SCHEDULE.NUM_STEPS = 1000
|
| 28 |
+
_C.DIFFUSION.SCHEDULE.BETA_1 = 1e-4
|
| 29 |
+
_C.DIFFUSION.SCHEDULE.BETA_T = 0.02
|
| 30 |
+
_C.DIFFUSION.SCHEDULE.MODE = "linear"
|
| 31 |
+
|
| 32 |
+
_C.CONTENT_ENCODER = CN()
|
| 33 |
+
_C.CONTENT_ENCODER.d_model = _C.D_MODEL
|
| 34 |
+
_C.CONTENT_ENCODER.nhead = 8
|
| 35 |
+
_C.CONTENT_ENCODER.num_encoder_layers = 3
|
| 36 |
+
_C.CONTENT_ENCODER.dim_feedforward = 4 * _C.D_MODEL
|
| 37 |
+
_C.CONTENT_ENCODER.dropout = 0.1
|
| 38 |
+
_C.CONTENT_ENCODER.activation = "relu"
|
| 39 |
+
_C.CONTENT_ENCODER.normalize_before = False
|
| 40 |
+
_C.CONTENT_ENCODER.pos_embed_len = 2 * _C.WIN_SIZE + 1
|
| 41 |
+
|
| 42 |
+
_C.STYLE_ENCODER = CN()
|
| 43 |
+
_C.STYLE_ENCODER.d_model = _C.D_MODEL
|
| 44 |
+
_C.STYLE_ENCODER.nhead = 8
|
| 45 |
+
_C.STYLE_ENCODER.num_encoder_layers = 3
|
| 46 |
+
_C.STYLE_ENCODER.dim_feedforward = 4 * _C.D_MODEL
|
| 47 |
+
_C.STYLE_ENCODER.dropout = 0.1
|
| 48 |
+
_C.STYLE_ENCODER.activation = "relu"
|
| 49 |
+
_C.STYLE_ENCODER.normalize_before = False
|
| 50 |
+
_C.STYLE_ENCODER.pos_embed_len = _C.DATASET.STYLE_MAX_LEN
|
| 51 |
+
_C.STYLE_ENCODER.aggregate_method = (
|
| 52 |
+
"self_attention_pooling" # average | self_attention_pooling
|
| 53 |
+
)
|
| 54 |
+
# _C.STYLE_ENCODER.input_dim = _C.DATASET.FACE3D_DIM
|
| 55 |
+
|
| 56 |
+
_C.DECODER = CN()
|
| 57 |
+
_C.DECODER.d_model = _C.D_MODEL
|
| 58 |
+
_C.DECODER.nhead = 8
|
| 59 |
+
_C.DECODER.num_decoder_layers = 3
|
| 60 |
+
_C.DECODER.dim_feedforward = 4 * _C.D_MODEL
|
| 61 |
+
_C.DECODER.dropout = 0.1
|
| 62 |
+
_C.DECODER.activation = "relu"
|
| 63 |
+
_C.DECODER.normalize_before = False
|
| 64 |
+
_C.DECODER.return_intermediate_dec = False
|
| 65 |
+
_C.DECODER.pos_embed_len = 2 * _C.WIN_SIZE + 1
|
| 66 |
+
_C.DECODER.network_type = "TransformerDecoder"
|
| 67 |
+
_C.DECODER.dynamic_K = None
|
| 68 |
+
_C.DECODER.dynamic_ratio = None
|
| 69 |
+
# _C.DECODER.output_dim = _C.DATASET.FACE3D_DIM
|
| 70 |
+
# LSFM basis:
|
| 71 |
+
# _C.DECODER.upper_face3d_indices = tuple(list(range(19)) + list(range(46, 51)))
|
| 72 |
+
# _C.DECODER.lower_face3d_indices = tuple(range(19, 46))
|
| 73 |
+
# BFM basis:
|
| 74 |
+
# fmt: off
|
| 75 |
+
_C.DECODER.upper_face3d_indices = [6, 8, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63]
|
| 76 |
+
# fmt: on
|
| 77 |
+
_C.DECODER.lower_face3d_indices = [0, 1, 2, 3, 4, 5, 7, 9, 10, 11, 12, 13, 14]
|
| 78 |
+
|
| 79 |
+
_C.CF_GUIDANCE = CN()
|
| 80 |
+
_C.CF_GUIDANCE.TRAINING = True
|
| 81 |
+
_C.CF_GUIDANCE.INFERENCE = True
|
| 82 |
+
_C.CF_GUIDANCE.NULL_PROB = 0.1
|
| 83 |
+
_C.CF_GUIDANCE.SCALE = 1.0
|
| 84 |
+
|
| 85 |
+
_C.INFERENCE = CN()
|
| 86 |
+
_C.INFERENCE.CHECKPOINT = "checkpoints/denoising_network.pth"
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def get_cfg_defaults():
|
| 90 |
+
"""Get a yacs CfgNode object with default values for my_project."""
|
| 91 |
+
return _C.clone()
|
actora/third_party/dreamtalk_src/core/__init__.py
ADDED
|
File without changes
|
actora/third_party/dreamtalk_src/core/networks/__init__.py
ADDED
|
@@ -0,0 +1,14 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
from core.networks.generator import (
|
| 2 |
+
StyleEncoder,
|
| 3 |
+
Decoder,
|
| 4 |
+
ContentW2VEncoder,
|
| 5 |
+
)
|
| 6 |
+
from core.networks.disentangle_decoder import DisentangleDecoder
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def get_network(name: str):
|
| 10 |
+
obj = globals().get(name)
|
| 11 |
+
if obj is None:
|
| 12 |
+
raise KeyError("Unknown Network: %s" % name)
|
| 13 |
+
else:
|
| 14 |
+
return obj
|
actora/third_party/dreamtalk_src/core/networks/diffusion_net.py
ADDED
|
@@ -0,0 +1,340 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import math
|
| 2 |
+
import torch
|
| 3 |
+
import torch.nn.functional as F
|
| 4 |
+
from torch.nn import Module
|
| 5 |
+
from core.networks.diffusion_util import VarianceSchedule
|
| 6 |
+
import numpy as np
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
def face3d_raw_to_norm(face3d_raw, exp_min, exp_max):
|
| 10 |
+
"""
|
| 11 |
+
|
| 12 |
+
Args:
|
| 13 |
+
face3d_raw (_type_): (B, L, C_face3d)
|
| 14 |
+
exp_min (_type_): (C_face3d)
|
| 15 |
+
exp_max (_type_): (C_face3d)
|
| 16 |
+
|
| 17 |
+
Returns:
|
| 18 |
+
_type_: (B, L, C_face3d) in [-1, 1]
|
| 19 |
+
"""
|
| 20 |
+
exp_min_expand = exp_min[None, None, :]
|
| 21 |
+
exp_max_expand = exp_max[None, None, :]
|
| 22 |
+
face3d_norm_01 = (face3d_raw - exp_min_expand) / (exp_max_expand - exp_min_expand)
|
| 23 |
+
face3d_norm = face3d_norm_01 * 2 - 1
|
| 24 |
+
return face3d_norm
|
| 25 |
+
|
| 26 |
+
|
| 27 |
+
def face3d_norm_to_raw(face3d_norm, exp_min, exp_max):
|
| 28 |
+
"""
|
| 29 |
+
|
| 30 |
+
Args:
|
| 31 |
+
face3d_norm (_type_): (B, L, C_face3d)
|
| 32 |
+
exp_min (_type_): (C_face3d)
|
| 33 |
+
exp_max (_type_): (C_face3d)
|
| 34 |
+
|
| 35 |
+
Returns:
|
| 36 |
+
_type_: (B, L, C_face3d)
|
| 37 |
+
"""
|
| 38 |
+
exp_min_expand = exp_min[None, None, :]
|
| 39 |
+
exp_max_expand = exp_max[None, None, :]
|
| 40 |
+
face3d_norm_01 = (face3d_norm + 1) / 2
|
| 41 |
+
face3d_raw = face3d_norm_01 * (exp_max_expand - exp_min_expand) + exp_min_expand
|
| 42 |
+
return face3d_raw
|
| 43 |
+
|
| 44 |
+
|
| 45 |
+
class DiffusionNet(Module):
|
| 46 |
+
def __init__(self, cfg, net, var_sched: VarianceSchedule):
|
| 47 |
+
super().__init__()
|
| 48 |
+
self.cfg = cfg
|
| 49 |
+
self.net = net
|
| 50 |
+
self.var_sched = var_sched
|
| 51 |
+
self.face3d_latent_type = self.cfg.TRAIN.FACE3D_LATENT.TYPE
|
| 52 |
+
self.predict_what = self.cfg.DIFFUSION.PREDICT_WHAT
|
| 53 |
+
|
| 54 |
+
if self.cfg.CF_GUIDANCE.TRAINING:
|
| 55 |
+
null_style_clip = torch.zeros(
|
| 56 |
+
self.cfg.DATASET.STYLE_MAX_LEN, self.cfg.DATASET.FACE3D_DIM
|
| 57 |
+
)
|
| 58 |
+
self.register_buffer("null_style_clip", null_style_clip)
|
| 59 |
+
|
| 60 |
+
null_pad_mask = torch.tensor([False] * self.cfg.DATASET.STYLE_MAX_LEN)
|
| 61 |
+
self.register_buffer("null_pad_mask", null_pad_mask)
|
| 62 |
+
|
| 63 |
+
def _face3d_to_latent(self, face3d):
|
| 64 |
+
latent = None
|
| 65 |
+
if self.face3d_latent_type == "face3d":
|
| 66 |
+
latent = face3d
|
| 67 |
+
elif self.face3d_latent_type == "normalized_face3d":
|
| 68 |
+
latent = face3d_raw_to_norm(
|
| 69 |
+
face3d, exp_min=self.exp_min, exp_max=self.exp_max
|
| 70 |
+
)
|
| 71 |
+
else:
|
| 72 |
+
raise ValueError(f"Invalid face3d latent type: {self.face3d_latent_type}")
|
| 73 |
+
return latent
|
| 74 |
+
|
| 75 |
+
def _latent_to_face3d(self, latent):
|
| 76 |
+
face3d = None
|
| 77 |
+
if self.face3d_latent_type == "face3d":
|
| 78 |
+
face3d = latent
|
| 79 |
+
elif self.face3d_latent_type == "normalized_face3d":
|
| 80 |
+
latent = torch.clamp(latent, min=-1, max=1)
|
| 81 |
+
face3d = face3d_norm_to_raw(
|
| 82 |
+
latent, exp_min=self.exp_min, exp_max=self.exp_max
|
| 83 |
+
)
|
| 84 |
+
else:
|
| 85 |
+
raise ValueError(f"Invalid face3d latent type: {self.face3d_latent_type}")
|
| 86 |
+
return face3d
|
| 87 |
+
|
| 88 |
+
def ddim_sample(
|
| 89 |
+
self,
|
| 90 |
+
audio,
|
| 91 |
+
style_clip,
|
| 92 |
+
style_pad_mask,
|
| 93 |
+
output_dim,
|
| 94 |
+
flexibility=0.0,
|
| 95 |
+
ret_traj=False,
|
| 96 |
+
use_cf_guidance=False,
|
| 97 |
+
cfg_scale=2.0,
|
| 98 |
+
ddim_num_step=50,
|
| 99 |
+
ready_style_code=None,
|
| 100 |
+
):
|
| 101 |
+
"""
|
| 102 |
+
|
| 103 |
+
Args:
|
| 104 |
+
audio (_type_): (B, L, W) or (B, L, W, C)
|
| 105 |
+
style_clip (_type_): (B, L_clipmax, C_face3d)
|
| 106 |
+
style_pad_mask : (B, L_clipmax)
|
| 107 |
+
pose_dim (_type_): int
|
| 108 |
+
flexibility (float, optional): _description_. Defaults to 0.0.
|
| 109 |
+
ret_traj (bool, optional): _description_. Defaults to False.
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
Returns:
|
| 113 |
+
_type_: (B, L, C_face)
|
| 114 |
+
"""
|
| 115 |
+
if self.predict_what != "x0":
|
| 116 |
+
raise NotImplementedError(self.predict_what)
|
| 117 |
+
|
| 118 |
+
if ready_style_code is not None and use_cf_guidance:
|
| 119 |
+
raise NotImplementedError("not implement cfg for ready style code")
|
| 120 |
+
|
| 121 |
+
c = self.var_sched.num_steps // ddim_num_step
|
| 122 |
+
time_steps = torch.tensor(
|
| 123 |
+
np.asarray(list(range(0, self.var_sched.num_steps, c))) + 1
|
| 124 |
+
)
|
| 125 |
+
assert len(time_steps) == ddim_num_step
|
| 126 |
+
prev_time_steps = torch.cat((torch.tensor([0]), time_steps[:-1]))
|
| 127 |
+
|
| 128 |
+
batch_size, output_len = audio.shape[:2]
|
| 129 |
+
# batch_size = context.size(0)
|
| 130 |
+
context = {
|
| 131 |
+
"audio": audio,
|
| 132 |
+
"style_clip": style_clip,
|
| 133 |
+
"style_pad_mask": style_pad_mask,
|
| 134 |
+
"ready_style_code": ready_style_code,
|
| 135 |
+
}
|
| 136 |
+
if use_cf_guidance:
|
| 137 |
+
uncond_style_clip = self.null_style_clip.unsqueeze(0).repeat(
|
| 138 |
+
batch_size, 1, 1
|
| 139 |
+
)
|
| 140 |
+
uncond_pad_mask = self.null_pad_mask.unsqueeze(0).repeat(batch_size, 1)
|
| 141 |
+
|
| 142 |
+
context_double = {
|
| 143 |
+
"audio": torch.cat([audio] * 2, dim=0),
|
| 144 |
+
"style_clip": torch.cat([style_clip, uncond_style_clip], dim=0),
|
| 145 |
+
"style_pad_mask": torch.cat([style_pad_mask, uncond_pad_mask], dim=0),
|
| 146 |
+
"ready_style_code": None
|
| 147 |
+
if ready_style_code is None
|
| 148 |
+
else torch.cat(
|
| 149 |
+
[
|
| 150 |
+
ready_style_code,
|
| 151 |
+
self.net.style_encoder(uncond_style_clip, uncond_pad_mask),
|
| 152 |
+
],
|
| 153 |
+
dim=0,
|
| 154 |
+
),
|
| 155 |
+
}
|
| 156 |
+
|
| 157 |
+
x_t = torch.randn([batch_size, output_len, output_dim]).to(audio.device)
|
| 158 |
+
|
| 159 |
+
for idx in list(range(ddim_num_step))[::-1]:
|
| 160 |
+
t = time_steps[idx]
|
| 161 |
+
t_prev = prev_time_steps[idx]
|
| 162 |
+
ddim_alpha = self.var_sched.alpha_bars[t]
|
| 163 |
+
ddim_alpha_prev = self.var_sched.alpha_bars[t_prev]
|
| 164 |
+
|
| 165 |
+
t_tensor = torch.tensor([t] * batch_size).to(audio.device).float()
|
| 166 |
+
if use_cf_guidance:
|
| 167 |
+
x_t_double = torch.cat([x_t] * 2, dim=0)
|
| 168 |
+
t_tensor_double = torch.cat([t_tensor] * 2, dim=0)
|
| 169 |
+
cond_output, uncond_output = self.net(
|
| 170 |
+
x_t_double, t=t_tensor_double, **context_double
|
| 171 |
+
).chunk(2)
|
| 172 |
+
diff_output = uncond_output + cfg_scale * (cond_output - uncond_output)
|
| 173 |
+
else:
|
| 174 |
+
diff_output = self.net(x_t, t=t_tensor, **context)
|
| 175 |
+
|
| 176 |
+
pred_x0 = diff_output
|
| 177 |
+
eps = (x_t - torch.sqrt(ddim_alpha) * pred_x0) / torch.sqrt(1 - ddim_alpha)
|
| 178 |
+
c1 = torch.sqrt(ddim_alpha_prev)
|
| 179 |
+
c2 = torch.sqrt(1 - ddim_alpha_prev)
|
| 180 |
+
|
| 181 |
+
x_t = c1 * pred_x0 + c2 * eps
|
| 182 |
+
|
| 183 |
+
latent_output = x_t
|
| 184 |
+
face3d_output = self._latent_to_face3d(latent_output)
|
| 185 |
+
return face3d_output
|
| 186 |
+
|
| 187 |
+
def sample(
|
| 188 |
+
self,
|
| 189 |
+
audio,
|
| 190 |
+
style_clip,
|
| 191 |
+
style_pad_mask,
|
| 192 |
+
output_dim,
|
| 193 |
+
flexibility=0.0,
|
| 194 |
+
ret_traj=False,
|
| 195 |
+
use_cf_guidance=False,
|
| 196 |
+
cfg_scale=2.0,
|
| 197 |
+
sample_method="ddpm",
|
| 198 |
+
ddim_num_step=50,
|
| 199 |
+
ready_style_code=None,
|
| 200 |
+
):
|
| 201 |
+
# sample_method = kwargs["sample_method"]
|
| 202 |
+
if sample_method == "ddpm":
|
| 203 |
+
if ready_style_code is not None:
|
| 204 |
+
raise NotImplementedError("ready style code in ddpm")
|
| 205 |
+
return self.ddpm_sample(
|
| 206 |
+
audio,
|
| 207 |
+
style_clip,
|
| 208 |
+
style_pad_mask,
|
| 209 |
+
output_dim,
|
| 210 |
+
flexibility=flexibility,
|
| 211 |
+
ret_traj=ret_traj,
|
| 212 |
+
use_cf_guidance=use_cf_guidance,
|
| 213 |
+
cfg_scale=cfg_scale,
|
| 214 |
+
)
|
| 215 |
+
elif sample_method == "ddim":
|
| 216 |
+
return self.ddim_sample(
|
| 217 |
+
audio,
|
| 218 |
+
style_clip,
|
| 219 |
+
style_pad_mask,
|
| 220 |
+
output_dim,
|
| 221 |
+
flexibility=flexibility,
|
| 222 |
+
ret_traj=ret_traj,
|
| 223 |
+
use_cf_guidance=use_cf_guidance,
|
| 224 |
+
cfg_scale=cfg_scale,
|
| 225 |
+
ddim_num_step=ddim_num_step,
|
| 226 |
+
ready_style_code=ready_style_code,
|
| 227 |
+
)
|
| 228 |
+
|
| 229 |
+
def ddpm_sample(
|
| 230 |
+
self,
|
| 231 |
+
audio,
|
| 232 |
+
style_clip,
|
| 233 |
+
style_pad_mask,
|
| 234 |
+
output_dim,
|
| 235 |
+
flexibility=0.0,
|
| 236 |
+
ret_traj=False,
|
| 237 |
+
use_cf_guidance=False,
|
| 238 |
+
cfg_scale=2.0,
|
| 239 |
+
):
|
| 240 |
+
"""
|
| 241 |
+
|
| 242 |
+
Args:
|
| 243 |
+
audio (_type_): (B, L, W) or (B, L, W, C)
|
| 244 |
+
style_clip (_type_): (B, L_clipmax, C_face3d)
|
| 245 |
+
style_pad_mask : (B, L_clipmax)
|
| 246 |
+
pose_dim (_type_): int
|
| 247 |
+
flexibility (float, optional): _description_. Defaults to 0.0.
|
| 248 |
+
ret_traj (bool, optional): _description_. Defaults to False.
|
| 249 |
+
|
| 250 |
+
|
| 251 |
+
Returns:
|
| 252 |
+
_type_: (B, L, C_face)
|
| 253 |
+
"""
|
| 254 |
+
batch_size, output_len = audio.shape[:2]
|
| 255 |
+
# batch_size = context.size(0)
|
| 256 |
+
context = {
|
| 257 |
+
"audio": audio,
|
| 258 |
+
"style_clip": style_clip,
|
| 259 |
+
"style_pad_mask": style_pad_mask,
|
| 260 |
+
}
|
| 261 |
+
if use_cf_guidance:
|
| 262 |
+
uncond_style_clip = self.null_style_clip.unsqueeze(0).repeat(
|
| 263 |
+
batch_size, 1, 1
|
| 264 |
+
)
|
| 265 |
+
uncond_pad_mask = self.null_pad_mask.unsqueeze(0).repeat(batch_size, 1)
|
| 266 |
+
context_double = {
|
| 267 |
+
"audio": torch.cat([audio] * 2, dim=0),
|
| 268 |
+
"style_clip": torch.cat([style_clip, uncond_style_clip], dim=0),
|
| 269 |
+
"style_pad_mask": torch.cat([style_pad_mask, uncond_pad_mask], dim=0),
|
| 270 |
+
}
|
| 271 |
+
|
| 272 |
+
x_T = torch.randn([batch_size, output_len, output_dim]).to(audio.device)
|
| 273 |
+
traj = {self.var_sched.num_steps: x_T}
|
| 274 |
+
for t in range(self.var_sched.num_steps, 0, -1):
|
| 275 |
+
alpha = self.var_sched.alphas[t]
|
| 276 |
+
alpha_bar = self.var_sched.alpha_bars[t]
|
| 277 |
+
alpha_bar_prev = self.var_sched.alpha_bars[t - 1]
|
| 278 |
+
sigma = self.var_sched.get_sigmas(t, flexibility)
|
| 279 |
+
|
| 280 |
+
z = torch.randn_like(x_T) if t > 1 else torch.zeros_like(x_T)
|
| 281 |
+
x_t = traj[t]
|
| 282 |
+
t_tensor = torch.tensor([t] * batch_size).to(audio.device).float()
|
| 283 |
+
if use_cf_guidance:
|
| 284 |
+
x_t_double = torch.cat([x_t] * 2, dim=0)
|
| 285 |
+
t_tensor_double = torch.cat([t_tensor] * 2, dim=0)
|
| 286 |
+
cond_output, uncond_output = self.net(
|
| 287 |
+
x_t_double, t=t_tensor_double, **context_double
|
| 288 |
+
).chunk(2)
|
| 289 |
+
diff_output = uncond_output + cfg_scale * (cond_output - uncond_output)
|
| 290 |
+
else:
|
| 291 |
+
diff_output = self.net(x_t, t=t_tensor, **context)
|
| 292 |
+
|
| 293 |
+
if self.predict_what == "noise":
|
| 294 |
+
c0 = 1.0 / torch.sqrt(alpha)
|
| 295 |
+
c1 = (1 - alpha) / torch.sqrt(1 - alpha_bar)
|
| 296 |
+
x_next = c0 * (x_t - c1 * diff_output) + sigma * z
|
| 297 |
+
elif self.predict_what == "x0":
|
| 298 |
+
d0 = torch.sqrt(alpha) * (1 - alpha_bar_prev) / (1 - alpha_bar)
|
| 299 |
+
d1 = torch.sqrt(alpha_bar_prev) * (1 - alpha) / (1 - alpha_bar)
|
| 300 |
+
x_next = d0 * x_t + d1 * diff_output + sigma * z
|
| 301 |
+
traj[t - 1] = x_next.detach()
|
| 302 |
+
traj[t] = traj[t].cpu()
|
| 303 |
+
if not ret_traj:
|
| 304 |
+
del traj[t]
|
| 305 |
+
|
| 306 |
+
if ret_traj:
|
| 307 |
+
raise NotImplementedError
|
| 308 |
+
return traj
|
| 309 |
+
else:
|
| 310 |
+
latent_output = traj[0]
|
| 311 |
+
face3d_output = self._latent_to_face3d(latent_output)
|
| 312 |
+
return face3d_output
|
| 313 |
+
|
| 314 |
+
|
| 315 |
+
if __name__ == "__main__":
|
| 316 |
+
from core.networks.diffusion_util import NoisePredictor, VarianceSchedule
|
| 317 |
+
|
| 318 |
+
diffnet = DiffusionNet(
|
| 319 |
+
net=NoisePredictor(),
|
| 320 |
+
var_sched=VarianceSchedule(
|
| 321 |
+
num_steps=500, beta_1=1e-4, beta_T=0.02, mode="linear"
|
| 322 |
+
),
|
| 323 |
+
)
|
| 324 |
+
|
| 325 |
+
import torch
|
| 326 |
+
|
| 327 |
+
gt_face3d = torch.randn(16, 64, 64)
|
| 328 |
+
audio = torch.randn(16, 64, 11)
|
| 329 |
+
style_clip = torch.randn(16, 256, 64)
|
| 330 |
+
style_pad_mask = torch.ones(16, 256)
|
| 331 |
+
|
| 332 |
+
context = {
|
| 333 |
+
"audio": audio,
|
| 334 |
+
"style_clip": style_clip,
|
| 335 |
+
"style_pad_mask": style_pad_mask,
|
| 336 |
+
}
|
| 337 |
+
|
| 338 |
+
loss = diffnet.get_loss(gt_face3d, context)
|
| 339 |
+
|
| 340 |
+
print("hello")
|
actora/third_party/dreamtalk_src/core/networks/diffusion_util.py
ADDED
|
@@ -0,0 +1,131 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import numpy as np
|
| 2 |
+
import torch
|
| 3 |
+
import torch.nn as nn
|
| 4 |
+
from torch.nn import Module
|
| 5 |
+
from core.networks import get_network
|
| 6 |
+
from core.utils import sinusoidal_embedding
|
| 7 |
+
|
| 8 |
+
|
| 9 |
+
class VarianceSchedule(Module):
|
| 10 |
+
def __init__(self, num_steps, beta_1, beta_T, mode="linear"):
|
| 11 |
+
super().__init__()
|
| 12 |
+
assert mode in ("linear",)
|
| 13 |
+
self.num_steps = num_steps
|
| 14 |
+
self.beta_1 = beta_1
|
| 15 |
+
self.beta_T = beta_T
|
| 16 |
+
self.mode = mode
|
| 17 |
+
|
| 18 |
+
if mode == "linear":
|
| 19 |
+
betas = torch.linspace(beta_1, beta_T, steps=num_steps)
|
| 20 |
+
|
| 21 |
+
betas = torch.cat([torch.zeros([1]), betas], dim=0) # Padding
|
| 22 |
+
|
| 23 |
+
alphas = 1 - betas
|
| 24 |
+
log_alphas = torch.log(alphas)
|
| 25 |
+
for i in range(1, log_alphas.size(0)): # 1 to T
|
| 26 |
+
log_alphas[i] += log_alphas[i - 1]
|
| 27 |
+
alpha_bars = log_alphas.exp()
|
| 28 |
+
|
| 29 |
+
sigmas_flex = torch.sqrt(betas)
|
| 30 |
+
sigmas_inflex = torch.zeros_like(sigmas_flex)
|
| 31 |
+
for i in range(1, sigmas_flex.size(0)):
|
| 32 |
+
sigmas_inflex[i] = ((1 - alpha_bars[i - 1]) / (1 - alpha_bars[i])) * betas[
|
| 33 |
+
i
|
| 34 |
+
]
|
| 35 |
+
sigmas_inflex = torch.sqrt(sigmas_inflex)
|
| 36 |
+
|
| 37 |
+
self.register_buffer("betas", betas)
|
| 38 |
+
self.register_buffer("alphas", alphas)
|
| 39 |
+
self.register_buffer("alpha_bars", alpha_bars)
|
| 40 |
+
self.register_buffer("sigmas_flex", sigmas_flex)
|
| 41 |
+
self.register_buffer("sigmas_inflex", sigmas_inflex)
|
| 42 |
+
|
| 43 |
+
def uniform_sample_t(self, batch_size):
|
| 44 |
+
ts = np.random.choice(np.arange(1, self.num_steps + 1), batch_size)
|
| 45 |
+
return ts.tolist()
|
| 46 |
+
|
| 47 |
+
def get_sigmas(self, t, flexibility):
|
| 48 |
+
assert 0 <= flexibility and flexibility <= 1
|
| 49 |
+
sigmas = self.sigmas_flex[t] * flexibility + self.sigmas_inflex[t] * (
|
| 50 |
+
1 - flexibility
|
| 51 |
+
)
|
| 52 |
+
return sigmas
|
| 53 |
+
|
| 54 |
+
|
| 55 |
+
class NoisePredictor(nn.Module):
|
| 56 |
+
def __init__(self, cfg):
|
| 57 |
+
super().__init__()
|
| 58 |
+
|
| 59 |
+
content_encoder_class = get_network(cfg.CONTENT_ENCODER_TYPE)
|
| 60 |
+
self.content_encoder = content_encoder_class(**cfg.CONTENT_ENCODER)
|
| 61 |
+
|
| 62 |
+
style_encoder_class = get_network(cfg.STYLE_ENCODER_TYPE)
|
| 63 |
+
cfg.defrost()
|
| 64 |
+
cfg.STYLE_ENCODER.input_dim = cfg.DATASET.FACE3D_DIM
|
| 65 |
+
cfg.freeze()
|
| 66 |
+
self.style_encoder = style_encoder_class(**cfg.STYLE_ENCODER)
|
| 67 |
+
|
| 68 |
+
decoder_class = get_network(cfg.DECODER_TYPE)
|
| 69 |
+
cfg.defrost()
|
| 70 |
+
cfg.DECODER.output_dim = cfg.DATASET.FACE3D_DIM
|
| 71 |
+
cfg.freeze()
|
| 72 |
+
self.decoder = decoder_class(**cfg.DECODER)
|
| 73 |
+
|
| 74 |
+
self.content_xt_to_decoder_input_wo_time = nn.Sequential(
|
| 75 |
+
nn.Linear(cfg.D_MODEL + cfg.DATASET.FACE3D_DIM, cfg.D_MODEL),
|
| 76 |
+
nn.ReLU(),
|
| 77 |
+
nn.Linear(cfg.D_MODEL, cfg.D_MODEL),
|
| 78 |
+
nn.ReLU(),
|
| 79 |
+
nn.Linear(cfg.D_MODEL, cfg.D_MODEL),
|
| 80 |
+
)
|
| 81 |
+
|
| 82 |
+
self.time_sinusoidal_dim = cfg.D_MODEL
|
| 83 |
+
self.time_embed_net = nn.Sequential(
|
| 84 |
+
nn.Linear(cfg.D_MODEL, cfg.D_MODEL),
|
| 85 |
+
nn.SiLU(),
|
| 86 |
+
nn.Linear(cfg.D_MODEL, cfg.D_MODEL),
|
| 87 |
+
)
|
| 88 |
+
|
| 89 |
+
def forward(self, x_t, t, audio, style_clip, style_pad_mask, ready_style_code=None):
|
| 90 |
+
"""_summary_
|
| 91 |
+
|
| 92 |
+
Args:
|
| 93 |
+
x_t (_type_): (B, L, C_face)
|
| 94 |
+
t (_type_): (B,) dtype:float32
|
| 95 |
+
audio (_type_): (B, L, W)
|
| 96 |
+
style_clip (_type_): (B, L_clipmax, C_face3d)
|
| 97 |
+
style_pad_mask : (B, L_clipmax)
|
| 98 |
+
ready_style_code: (B, C_model)
|
| 99 |
+
Returns:
|
| 100 |
+
e_theta : (B, L, C_face)
|
| 101 |
+
"""
|
| 102 |
+
W = audio.shape[2]
|
| 103 |
+
content = self.content_encoder(audio)
|
| 104 |
+
# (B, L, W, C_model)
|
| 105 |
+
x_t_expand = x_t.unsqueeze(2).repeat(1, 1, W, 1)
|
| 106 |
+
# (B, L, C_face) -> (B, L, W, C_face)
|
| 107 |
+
content_xt_concat = torch.cat((content, x_t_expand), dim=3)
|
| 108 |
+
# (B, L, W, C_model+C_face)
|
| 109 |
+
decoder_input_without_time = self.content_xt_to_decoder_input_wo_time(
|
| 110 |
+
content_xt_concat
|
| 111 |
+
)
|
| 112 |
+
# (B, L, W, C_model)
|
| 113 |
+
|
| 114 |
+
time_sinusoidal = sinusoidal_embedding(t, self.time_sinusoidal_dim)
|
| 115 |
+
# (B, C_embed)
|
| 116 |
+
time_embedding = self.time_embed_net(time_sinusoidal)
|
| 117 |
+
# (B, C_model)
|
| 118 |
+
B, C = time_embedding.shape
|
| 119 |
+
time_embed_expand = time_embedding.view(B, 1, 1, C)
|
| 120 |
+
decoder_input = decoder_input_without_time + time_embed_expand
|
| 121 |
+
# (B, L, W, C_model)
|
| 122 |
+
|
| 123 |
+
if ready_style_code is not None:
|
| 124 |
+
style_code = ready_style_code
|
| 125 |
+
else:
|
| 126 |
+
style_code = self.style_encoder(style_clip, style_pad_mask)
|
| 127 |
+
# (B, C_model)
|
| 128 |
+
|
| 129 |
+
e_theta = self.decoder(decoder_input, style_code)
|
| 130 |
+
# (B, L, C_face)
|
| 131 |
+
return e_theta
|
actora/third_party/dreamtalk_src/core/networks/disentangle_decoder.py
ADDED
|
@@ -0,0 +1,240 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch
|
| 2 |
+
from torch import nn
|
| 3 |
+
|
| 4 |
+
from .transformer import (
|
| 5 |
+
PositionalEncoding,
|
| 6 |
+
TransformerDecoderLayer,
|
| 7 |
+
TransformerDecoder,
|
| 8 |
+
)
|
| 9 |
+
from core.networks.dynamic_fc_decoder import DynamicFCDecoderLayer, DynamicFCDecoder
|
| 10 |
+
from core.utils import _reset_parameters
|
| 11 |
+
|
| 12 |
+
|
| 13 |
+
def get_decoder_network(
|
| 14 |
+
network_type,
|
| 15 |
+
d_model,
|
| 16 |
+
nhead,
|
| 17 |
+
dim_feedforward,
|
| 18 |
+
dropout,
|
| 19 |
+
activation,
|
| 20 |
+
normalize_before,
|
| 21 |
+
num_decoder_layers,
|
| 22 |
+
return_intermediate_dec,
|
| 23 |
+
dynamic_K,
|
| 24 |
+
dynamic_ratio,
|
| 25 |
+
):
|
| 26 |
+
decoder = None
|
| 27 |
+
if network_type == "TransformerDecoder":
|
| 28 |
+
decoder_layer = TransformerDecoderLayer(
|
| 29 |
+
d_model, nhead, dim_feedforward, dropout, activation, normalize_before
|
| 30 |
+
)
|
| 31 |
+
norm = nn.LayerNorm(d_model)
|
| 32 |
+
decoder = TransformerDecoder(
|
| 33 |
+
decoder_layer,
|
| 34 |
+
num_decoder_layers,
|
| 35 |
+
norm,
|
| 36 |
+
return_intermediate_dec,
|
| 37 |
+
)
|
| 38 |
+
elif network_type == "DynamicFCDecoder":
|
| 39 |
+
d_style = d_model
|
| 40 |
+
decoder_layer = DynamicFCDecoderLayer(
|
| 41 |
+
d_model,
|
| 42 |
+
nhead,
|
| 43 |
+
d_style,
|
| 44 |
+
dynamic_K,
|
| 45 |
+
dynamic_ratio,
|
| 46 |
+
dim_feedforward,
|
| 47 |
+
dropout,
|
| 48 |
+
activation,
|
| 49 |
+
normalize_before,
|
| 50 |
+
)
|
| 51 |
+
norm = nn.LayerNorm(d_model)
|
| 52 |
+
decoder = DynamicFCDecoder(
|
| 53 |
+
decoder_layer, num_decoder_layers, norm, return_intermediate_dec
|
| 54 |
+
)
|
| 55 |
+
elif network_type == "DynamicFCEncoder":
|
| 56 |
+
d_style = d_model
|
| 57 |
+
decoder_layer = DynamicFCEncoderLayer(
|
| 58 |
+
d_model,
|
| 59 |
+
nhead,
|
| 60 |
+
d_style,
|
| 61 |
+
dynamic_K,
|
| 62 |
+
dynamic_ratio,
|
| 63 |
+
dim_feedforward,
|
| 64 |
+
dropout,
|
| 65 |
+
activation,
|
| 66 |
+
normalize_before,
|
| 67 |
+
)
|
| 68 |
+
norm = nn.LayerNorm(d_model)
|
| 69 |
+
decoder = DynamicFCEncoder(decoder_layer, num_decoder_layers, norm)
|
| 70 |
+
|
| 71 |
+
else:
|
| 72 |
+
raise ValueError(f"Invalid network_type {network_type}")
|
| 73 |
+
|
| 74 |
+
return decoder
|
| 75 |
+
|
| 76 |
+
|
| 77 |
+
class DisentangleDecoder(nn.Module):
|
| 78 |
+
def __init__(
|
| 79 |
+
self,
|
| 80 |
+
d_model=512,
|
| 81 |
+
nhead=8,
|
| 82 |
+
num_decoder_layers=3,
|
| 83 |
+
dim_feedforward=2048,
|
| 84 |
+
dropout=0.1,
|
| 85 |
+
activation="relu",
|
| 86 |
+
normalize_before=False,
|
| 87 |
+
return_intermediate_dec=False,
|
| 88 |
+
pos_embed_len=80,
|
| 89 |
+
upper_face3d_indices=tuple(list(range(19)) + list(range(46, 51))),
|
| 90 |
+
lower_face3d_indices=tuple(range(19, 46)),
|
| 91 |
+
network_type="None",
|
| 92 |
+
dynamic_K=None,
|
| 93 |
+
dynamic_ratio=None,
|
| 94 |
+
**_,
|
| 95 |
+
) -> None:
|
| 96 |
+
super().__init__()
|
| 97 |
+
|
| 98 |
+
self.upper_face3d_indices = upper_face3d_indices
|
| 99 |
+
self.lower_face3d_indices = lower_face3d_indices
|
| 100 |
+
|
| 101 |
+
# upper_decoder_layer = TransformerDecoderLayer(
|
| 102 |
+
# d_model, nhead, dim_feedforward, dropout, activation, normalize_before
|
| 103 |
+
# )
|
| 104 |
+
# upper_decoder_norm = nn.LayerNorm(d_model)
|
| 105 |
+
# self.upper_decoder = TransformerDecoder(
|
| 106 |
+
# upper_decoder_layer,
|
| 107 |
+
# num_decoder_layers,
|
| 108 |
+
# upper_decoder_norm,
|
| 109 |
+
# return_intermediate=return_intermediate_dec,
|
| 110 |
+
# )
|
| 111 |
+
self.upper_decoder = get_decoder_network(
|
| 112 |
+
network_type,
|
| 113 |
+
d_model,
|
| 114 |
+
nhead,
|
| 115 |
+
dim_feedforward,
|
| 116 |
+
dropout,
|
| 117 |
+
activation,
|
| 118 |
+
normalize_before,
|
| 119 |
+
num_decoder_layers,
|
| 120 |
+
return_intermediate_dec,
|
| 121 |
+
dynamic_K,
|
| 122 |
+
dynamic_ratio,
|
| 123 |
+
)
|
| 124 |
+
_reset_parameters(self.upper_decoder)
|
| 125 |
+
|
| 126 |
+
# lower_decoder_layer = TransformerDecoderLayer(
|
| 127 |
+
# d_model, nhead, dim_feedforward, dropout, activation, normalize_before
|
| 128 |
+
# )
|
| 129 |
+
# lower_decoder_norm = nn.LayerNorm(d_model)
|
| 130 |
+
# self.lower_decoder = TransformerDecoder(
|
| 131 |
+
# lower_decoder_layer,
|
| 132 |
+
# num_decoder_layers,
|
| 133 |
+
# lower_decoder_norm,
|
| 134 |
+
# return_intermediate=return_intermediate_dec,
|
| 135 |
+
# )
|
| 136 |
+
self.lower_decoder = get_decoder_network(
|
| 137 |
+
network_type,
|
| 138 |
+
d_model,
|
| 139 |
+
nhead,
|
| 140 |
+
dim_feedforward,
|
| 141 |
+
dropout,
|
| 142 |
+
activation,
|
| 143 |
+
normalize_before,
|
| 144 |
+
num_decoder_layers,
|
| 145 |
+
return_intermediate_dec,
|
| 146 |
+
dynamic_K,
|
| 147 |
+
dynamic_ratio,
|
| 148 |
+
)
|
| 149 |
+
_reset_parameters(self.lower_decoder)
|
| 150 |
+
|
| 151 |
+
self.pos_embed = PositionalEncoding(d_model, pos_embed_len)
|
| 152 |
+
|
| 153 |
+
tail_hidden_dim = d_model // 2
|
| 154 |
+
self.upper_tail_fc = nn.Sequential(
|
| 155 |
+
nn.Linear(d_model, tail_hidden_dim),
|
| 156 |
+
nn.ReLU(),
|
| 157 |
+
nn.Linear(tail_hidden_dim, tail_hidden_dim),
|
| 158 |
+
nn.ReLU(),
|
| 159 |
+
nn.Linear(tail_hidden_dim, len(upper_face3d_indices)),
|
| 160 |
+
)
|
| 161 |
+
self.lower_tail_fc = nn.Sequential(
|
| 162 |
+
nn.Linear(d_model, tail_hidden_dim),
|
| 163 |
+
nn.ReLU(),
|
| 164 |
+
nn.Linear(tail_hidden_dim, tail_hidden_dim),
|
| 165 |
+
nn.ReLU(),
|
| 166 |
+
nn.Linear(tail_hidden_dim, len(lower_face3d_indices)),
|
| 167 |
+
)
|
| 168 |
+
|
| 169 |
+
def forward(self, content, style_code):
|
| 170 |
+
"""
|
| 171 |
+
|
| 172 |
+
Args:
|
| 173 |
+
content (_type_): (B, num_frames, window, C_dmodel)
|
| 174 |
+
style_code (_type_): (B, C_dmodel)
|
| 175 |
+
|
| 176 |
+
Returns:
|
| 177 |
+
face3d: (B, L_clip, C_3dmm)
|
| 178 |
+
"""
|
| 179 |
+
B, N, W, C = content.shape
|
| 180 |
+
style = style_code.reshape(B, 1, 1, C).expand(B, N, W, C)
|
| 181 |
+
style = style.permute(2, 0, 1, 3).reshape(W, B * N, C)
|
| 182 |
+
# (W, B*N, C)
|
| 183 |
+
|
| 184 |
+
content = content.permute(2, 0, 1, 3).reshape(W, B * N, C)
|
| 185 |
+
# (W, B*N, C)
|
| 186 |
+
tgt = torch.zeros_like(style)
|
| 187 |
+
pos_embed = self.pos_embed(W)
|
| 188 |
+
pos_embed = pos_embed.permute(1, 0, 2)
|
| 189 |
+
|
| 190 |
+
upper_face3d_feat = self.upper_decoder(
|
| 191 |
+
tgt, content, pos=pos_embed, query_pos=style
|
| 192 |
+
)[0]
|
| 193 |
+
# (W, B*N, C)
|
| 194 |
+
upper_face3d_feat = upper_face3d_feat.permute(1, 0, 2).reshape(B, N, W, C)[
|
| 195 |
+
:, :, W // 2, :
|
| 196 |
+
]
|
| 197 |
+
# (B, N, C)
|
| 198 |
+
upper_face3d = self.upper_tail_fc(upper_face3d_feat)
|
| 199 |
+
# (B, N, C_exp)
|
| 200 |
+
|
| 201 |
+
lower_face3d_feat = self.lower_decoder(
|
| 202 |
+
tgt, content, pos=pos_embed, query_pos=style
|
| 203 |
+
)[0]
|
| 204 |
+
lower_face3d_feat = lower_face3d_feat.permute(1, 0, 2).reshape(B, N, W, C)[
|
| 205 |
+
:, :, W // 2, :
|
| 206 |
+
]
|
| 207 |
+
lower_face3d = self.lower_tail_fc(lower_face3d_feat)
|
| 208 |
+
C_exp = len(self.upper_face3d_indices) + len(self.lower_face3d_indices)
|
| 209 |
+
face3d = torch.zeros(B, N, C_exp).to(upper_face3d)
|
| 210 |
+
face3d[:, :, self.upper_face3d_indices] = upper_face3d
|
| 211 |
+
face3d[:, :, self.lower_face3d_indices] = lower_face3d
|
| 212 |
+
return face3d
|
| 213 |
+
|
| 214 |
+
|
| 215 |
+
if __name__ == "__main__":
|
| 216 |
+
import sys
|
| 217 |
+
|
| 218 |
+
sys.path.append("/home/mayifeng/Research/styleTH")
|
| 219 |
+
|
| 220 |
+
from configs.default import get_cfg_defaults
|
| 221 |
+
|
| 222 |
+
cfg = get_cfg_defaults()
|
| 223 |
+
cfg.merge_from_file("configs/styleTH_unpair_lsfm_emotion.yaml")
|
| 224 |
+
cfg.freeze()
|
| 225 |
+
|
| 226 |
+
# content_encoder = ContentEncoder(**cfg.CONTENT_ENCODER)
|
| 227 |
+
|
| 228 |
+
# dummy_audio = torch.randint(0, 41, (5, 64, 11))
|
| 229 |
+
# dummy_content = content_encoder(dummy_audio)
|
| 230 |
+
|
| 231 |
+
# style_encoder = StyleEncoder(**cfg.STYLE_ENCODER)
|
| 232 |
+
# dummy_face3d_seq = torch.randn(5, 64, 64)
|
| 233 |
+
# dummy_style_code = style_encoder(dummy_face3d_seq)
|
| 234 |
+
|
| 235 |
+
decoder = DisentangleDecoder(**cfg.DECODER)
|
| 236 |
+
dummy_content = torch.randn(5, 64, 11, 256)
|
| 237 |
+
dummy_style = torch.randn(5, 256)
|
| 238 |
+
dummy_output = decoder(dummy_content, dummy_style)
|
| 239 |
+
|
| 240 |
+
print("hello")
|
actora/third_party/dreamtalk_src/core/networks/dynamic_conv.py
ADDED
|
@@ -0,0 +1,156 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import math
|
| 2 |
+
|
| 3 |
+
import torch
|
| 4 |
+
from torch import nn
|
| 5 |
+
from torch.nn import functional as F
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class Attention(nn.Module):
|
| 9 |
+
def __init__(self, cond_planes, ratio, K, temperature=30, init_weight=True):
|
| 10 |
+
super().__init__()
|
| 11 |
+
# self.avgpool = nn.AdaptiveAvgPool2d(1)
|
| 12 |
+
self.temprature = temperature
|
| 13 |
+
assert cond_planes > ratio
|
| 14 |
+
hidden_planes = cond_planes // ratio
|
| 15 |
+
self.net = nn.Sequential(
|
| 16 |
+
nn.Conv2d(cond_planes, hidden_planes, kernel_size=1, bias=False),
|
| 17 |
+
nn.ReLU(),
|
| 18 |
+
nn.Conv2d(hidden_planes, K, kernel_size=1, bias=False),
|
| 19 |
+
)
|
| 20 |
+
|
| 21 |
+
if init_weight:
|
| 22 |
+
self._initialize_weights()
|
| 23 |
+
|
| 24 |
+
def update_temprature(self):
|
| 25 |
+
if self.temprature > 1:
|
| 26 |
+
self.temprature -= 1
|
| 27 |
+
|
| 28 |
+
def _initialize_weights(self):
|
| 29 |
+
for m in self.modules():
|
| 30 |
+
if isinstance(m, nn.Conv2d):
|
| 31 |
+
nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu")
|
| 32 |
+
if m.bias is not None:
|
| 33 |
+
nn.init.constant_(m.bias, 0)
|
| 34 |
+
if isinstance(m, nn.BatchNorm2d):
|
| 35 |
+
nn.init.constant_(m.weight, 1)
|
| 36 |
+
nn.init.constant_(m.bias, 0)
|
| 37 |
+
|
| 38 |
+
def forward(self, cond):
|
| 39 |
+
"""
|
| 40 |
+
|
| 41 |
+
Args:
|
| 42 |
+
cond (_type_): (B, C_style)
|
| 43 |
+
|
| 44 |
+
Returns:
|
| 45 |
+
_type_: (B, K)
|
| 46 |
+
"""
|
| 47 |
+
|
| 48 |
+
# att = self.avgpool(cond) # bs,dim,1,1
|
| 49 |
+
att = cond.view(cond.shape[0], cond.shape[1], 1, 1)
|
| 50 |
+
att = self.net(att).view(cond.shape[0], -1) # bs,K
|
| 51 |
+
return F.softmax(att / self.temprature, -1)
|
| 52 |
+
|
| 53 |
+
|
| 54 |
+
class DynamicConv(nn.Module):
|
| 55 |
+
def __init__(
|
| 56 |
+
self,
|
| 57 |
+
in_planes,
|
| 58 |
+
out_planes,
|
| 59 |
+
cond_planes,
|
| 60 |
+
kernel_size,
|
| 61 |
+
stride,
|
| 62 |
+
padding=0,
|
| 63 |
+
dilation=1,
|
| 64 |
+
groups=1,
|
| 65 |
+
bias=True,
|
| 66 |
+
K=4,
|
| 67 |
+
temperature=30,
|
| 68 |
+
ratio=4,
|
| 69 |
+
init_weight=True,
|
| 70 |
+
):
|
| 71 |
+
super().__init__()
|
| 72 |
+
self.in_planes = in_planes
|
| 73 |
+
self.out_planes = out_planes
|
| 74 |
+
self.cond_planes = cond_planes
|
| 75 |
+
self.kernel_size = kernel_size
|
| 76 |
+
self.stride = stride
|
| 77 |
+
self.padding = padding
|
| 78 |
+
self.dilation = dilation
|
| 79 |
+
self.groups = groups
|
| 80 |
+
self.bias = bias
|
| 81 |
+
self.K = K
|
| 82 |
+
self.init_weight = init_weight
|
| 83 |
+
self.attention = Attention(
|
| 84 |
+
cond_planes=cond_planes, ratio=ratio, K=K, temperature=temperature, init_weight=init_weight
|
| 85 |
+
)
|
| 86 |
+
|
| 87 |
+
self.weight = nn.Parameter(
|
| 88 |
+
torch.randn(K, out_planes, in_planes // groups, kernel_size, kernel_size), requires_grad=True
|
| 89 |
+
)
|
| 90 |
+
if bias:
|
| 91 |
+
self.bias = nn.Parameter(torch.randn(K, out_planes), requires_grad=True)
|
| 92 |
+
else:
|
| 93 |
+
self.bias = None
|
| 94 |
+
|
| 95 |
+
if self.init_weight:
|
| 96 |
+
self._initialize_weights()
|
| 97 |
+
|
| 98 |
+
def _initialize_weights(self):
|
| 99 |
+
for i in range(self.K):
|
| 100 |
+
nn.init.kaiming_uniform_(self.weight[i], a=math.sqrt(5))
|
| 101 |
+
if self.bias is not None:
|
| 102 |
+
fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight[i])
|
| 103 |
+
if fan_in != 0:
|
| 104 |
+
bound = 1 / math.sqrt(fan_in)
|
| 105 |
+
nn.init.uniform_(self.bias, -bound, bound)
|
| 106 |
+
|
| 107 |
+
def forward(self, x, cond):
|
| 108 |
+
"""
|
| 109 |
+
|
| 110 |
+
Args:
|
| 111 |
+
x (_type_): (B, C_in, L, 1)
|
| 112 |
+
cond (_type_): (B, C_style)
|
| 113 |
+
|
| 114 |
+
Returns:
|
| 115 |
+
_type_: (B, C_out, L, 1)
|
| 116 |
+
"""
|
| 117 |
+
bs, in_planels, h, w = x.shape
|
| 118 |
+
softmax_att = self.attention(cond) # bs,K
|
| 119 |
+
x = x.view(1, -1, h, w)
|
| 120 |
+
weight = self.weight.view(self.K, -1) # K,-1
|
| 121 |
+
aggregate_weight = torch.mm(softmax_att, weight).view(
|
| 122 |
+
bs * self.out_planes, self.in_planes // self.groups, self.kernel_size, self.kernel_size
|
| 123 |
+
) # bs*out_p,in_p,k,k
|
| 124 |
+
|
| 125 |
+
if self.bias is not None:
|
| 126 |
+
bias = self.bias.view(self.K, -1) # K,out_p
|
| 127 |
+
aggregate_bias = torch.mm(softmax_att, bias).view(-1) # bs*out_p
|
| 128 |
+
output = F.conv2d(
|
| 129 |
+
x, # 1, bs*in_p, L, 1
|
| 130 |
+
weight=aggregate_weight,
|
| 131 |
+
bias=aggregate_bias,
|
| 132 |
+
stride=self.stride,
|
| 133 |
+
padding=self.padding,
|
| 134 |
+
groups=self.groups * bs,
|
| 135 |
+
dilation=self.dilation,
|
| 136 |
+
)
|
| 137 |
+
else:
|
| 138 |
+
output = F.conv2d(
|
| 139 |
+
x,
|
| 140 |
+
weight=aggregate_weight,
|
| 141 |
+
bias=None,
|
| 142 |
+
stride=self.stride,
|
| 143 |
+
padding=self.padding,
|
| 144 |
+
groups=self.groups * bs,
|
| 145 |
+
dilation=self.dilation,
|
| 146 |
+
)
|
| 147 |
+
|
| 148 |
+
output = output.view(bs, self.out_planes, h, w)
|
| 149 |
+
return output
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
if __name__ == "__main__":
|
| 153 |
+
input = torch.randn(3, 32, 64, 64)
|
| 154 |
+
m = DynamicConv(in_planes=32, out_planes=64, kernel_size=3, stride=1, padding=1, bias=True)
|
| 155 |
+
out = m(input)
|
| 156 |
+
print(out.shape)
|
actora/third_party/dreamtalk_src/core/networks/dynamic_fc_decoder.py
ADDED
|
@@ -0,0 +1,178 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import torch.nn as nn
|
| 2 |
+
import torch
|
| 3 |
+
|
| 4 |
+
from core.networks.transformer import _get_activation_fn, _get_clones
|
| 5 |
+
from core.networks.dynamic_linear import DynamicLinear
|
| 6 |
+
|
| 7 |
+
|
| 8 |
+
class DynamicFCDecoderLayer(nn.Module):
|
| 9 |
+
def __init__(
|
| 10 |
+
self,
|
| 11 |
+
d_model,
|
| 12 |
+
nhead,
|
| 13 |
+
d_style,
|
| 14 |
+
dynamic_K,
|
| 15 |
+
dynamic_ratio,
|
| 16 |
+
dim_feedforward=2048,
|
| 17 |
+
dropout=0.1,
|
| 18 |
+
activation="relu",
|
| 19 |
+
normalize_before=False,
|
| 20 |
+
):
|
| 21 |
+
super().__init__()
|
| 22 |
+
self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
|
| 23 |
+
self.multihead_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout)
|
| 24 |
+
# Implementation of Feedforward model
|
| 25 |
+
# self.linear1 = nn.Linear(d_model, dim_feedforward)
|
| 26 |
+
self.linear1 = DynamicLinear(d_model, dim_feedforward, d_style, K=dynamic_K, ratio=dynamic_ratio)
|
| 27 |
+
self.dropout = nn.Dropout(dropout)
|
| 28 |
+
self.linear2 = nn.Linear(dim_feedforward, d_model)
|
| 29 |
+
# self.linear2 = DynamicLinear(dim_feedforward, d_model, d_style, K=dynamic_K, ratio=dynamic_ratio)
|
| 30 |
+
|
| 31 |
+
self.norm1 = nn.LayerNorm(d_model)
|
| 32 |
+
self.norm2 = nn.LayerNorm(d_model)
|
| 33 |
+
self.norm3 = nn.LayerNorm(d_model)
|
| 34 |
+
self.dropout1 = nn.Dropout(dropout)
|
| 35 |
+
self.dropout2 = nn.Dropout(dropout)
|
| 36 |
+
self.dropout3 = nn.Dropout(dropout)
|
| 37 |
+
|
| 38 |
+
self.activation = _get_activation_fn(activation)
|
| 39 |
+
self.normalize_before = normalize_before
|
| 40 |
+
|
| 41 |
+
def with_pos_embed(self, tensor, pos):
|
| 42 |
+
return tensor if pos is None else tensor + pos
|
| 43 |
+
|
| 44 |
+
def forward_post(
|
| 45 |
+
self,
|
| 46 |
+
tgt,
|
| 47 |
+
memory,
|
| 48 |
+
style,
|
| 49 |
+
tgt_mask=None,
|
| 50 |
+
memory_mask=None,
|
| 51 |
+
tgt_key_padding_mask=None,
|
| 52 |
+
memory_key_padding_mask=None,
|
| 53 |
+
pos=None,
|
| 54 |
+
query_pos=None,
|
| 55 |
+
):
|
| 56 |
+
# q = k = self.with_pos_embed(tgt, query_pos)
|
| 57 |
+
tgt2 = self.self_attn(tgt, tgt, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
|
| 58 |
+
tgt = tgt + self.dropout1(tgt2)
|
| 59 |
+
tgt = self.norm1(tgt)
|
| 60 |
+
tgt2 = self.multihead_attn(
|
| 61 |
+
query=tgt, key=memory, value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask
|
| 62 |
+
)[0]
|
| 63 |
+
tgt = tgt + self.dropout2(tgt2)
|
| 64 |
+
tgt = self.norm2(tgt)
|
| 65 |
+
# tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt, style))), style)
|
| 66 |
+
tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt, style))))
|
| 67 |
+
tgt = tgt + self.dropout3(tgt2)
|
| 68 |
+
tgt = self.norm3(tgt)
|
| 69 |
+
return tgt
|
| 70 |
+
|
| 71 |
+
# def forward_pre(
|
| 72 |
+
# self,
|
| 73 |
+
# tgt,
|
| 74 |
+
# memory,
|
| 75 |
+
# tgt_mask=None,
|
| 76 |
+
# memory_mask=None,
|
| 77 |
+
# tgt_key_padding_mask=None,
|
| 78 |
+
# memory_key_padding_mask=None,
|
| 79 |
+
# pos=None,
|
| 80 |
+
# query_pos=None,
|
| 81 |
+
# ):
|
| 82 |
+
# tgt2 = self.norm1(tgt)
|
| 83 |
+
# # q = k = self.with_pos_embed(tgt2, query_pos)
|
| 84 |
+
# tgt2 = self.self_attn(tgt2, tgt2, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]
|
| 85 |
+
# tgt = tgt + self.dropout1(tgt2)
|
| 86 |
+
# tgt2 = self.norm2(tgt)
|
| 87 |
+
# tgt2 = self.multihead_attn(
|
| 88 |
+
# query=tgt2, key=memory, value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask
|
| 89 |
+
# )[0]
|
| 90 |
+
# tgt = tgt + self.dropout2(tgt2)
|
| 91 |
+
# tgt2 = self.norm3(tgt)
|
| 92 |
+
# tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2))))
|
| 93 |
+
# tgt = tgt + self.dropout3(tgt2)
|
| 94 |
+
# return tgt
|
| 95 |
+
|
| 96 |
+
def forward(
|
| 97 |
+
self,
|
| 98 |
+
tgt,
|
| 99 |
+
memory,
|
| 100 |
+
style,
|
| 101 |
+
tgt_mask=None,
|
| 102 |
+
memory_mask=None,
|
| 103 |
+
tgt_key_padding_mask=None,
|
| 104 |
+
memory_key_padding_mask=None,
|
| 105 |
+
pos=None,
|
| 106 |
+
query_pos=None,
|
| 107 |
+
):
|
| 108 |
+
if self.normalize_before:
|
| 109 |
+
raise NotImplementedError
|
| 110 |
+
# return self.forward_pre(
|
| 111 |
+
# tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos
|
| 112 |
+
# )
|
| 113 |
+
return self.forward_post(
|
| 114 |
+
tgt, memory, style, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos
|
| 115 |
+
)
|
| 116 |
+
|
| 117 |
+
|
| 118 |
+
class DynamicFCDecoder(nn.Module):
|
| 119 |
+
def __init__(self, decoder_layer, num_layers, norm=None, return_intermediate=False):
|
| 120 |
+
super().__init__()
|
| 121 |
+
self.layers = _get_clones(decoder_layer, num_layers)
|
| 122 |
+
self.num_layers = num_layers
|
| 123 |
+
self.norm = norm
|
| 124 |
+
self.return_intermediate = return_intermediate
|
| 125 |
+
|
| 126 |
+
def forward(
|
| 127 |
+
self,
|
| 128 |
+
tgt,
|
| 129 |
+
memory,
|
| 130 |
+
tgt_mask=None,
|
| 131 |
+
memory_mask=None,
|
| 132 |
+
tgt_key_padding_mask=None,
|
| 133 |
+
memory_key_padding_mask=None,
|
| 134 |
+
pos=None,
|
| 135 |
+
query_pos=None,
|
| 136 |
+
):
|
| 137 |
+
style = query_pos[0]
|
| 138 |
+
# (B*N, C)
|
| 139 |
+
output = tgt + pos + query_pos
|
| 140 |
+
|
| 141 |
+
intermediate = []
|
| 142 |
+
|
| 143 |
+
for layer in self.layers:
|
| 144 |
+
output = layer(
|
| 145 |
+
output,
|
| 146 |
+
memory,
|
| 147 |
+
style,
|
| 148 |
+
tgt_mask=tgt_mask,
|
| 149 |
+
memory_mask=memory_mask,
|
| 150 |
+
tgt_key_padding_mask=tgt_key_padding_mask,
|
| 151 |
+
memory_key_padding_mask=memory_key_padding_mask,
|
| 152 |
+
pos=pos,
|
| 153 |
+
query_pos=query_pos,
|
| 154 |
+
)
|
| 155 |
+
if self.return_intermediate:
|
| 156 |
+
intermediate.append(self.norm(output))
|
| 157 |
+
|
| 158 |
+
if self.norm is not None:
|
| 159 |
+
output = self.norm(output)
|
| 160 |
+
if self.return_intermediate:
|
| 161 |
+
intermediate.pop()
|
| 162 |
+
intermediate.append(output)
|
| 163 |
+
|
| 164 |
+
if self.return_intermediate:
|
| 165 |
+
return torch.stack(intermediate)
|
| 166 |
+
|
| 167 |
+
return output.unsqueeze(0)
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
if __name__ == "__main__":
|
| 171 |
+
query = torch.randn(11, 1024, 256)
|
| 172 |
+
content = torch.randn(11, 1024, 256)
|
| 173 |
+
style = torch.randn(1024, 256)
|
| 174 |
+
pos = torch.randn(11, 1, 256)
|
| 175 |
+
m = DynamicFCDecoderLayer(256, 4, 256, 4, 4, 1024)
|
| 176 |
+
|
| 177 |
+
out = m(query, content, style, pos=pos)
|
| 178 |
+
print(out.shape)
|