github-actions[bot] commited on
Commit
cf6e6cb
·
0 Parent(s):

Deploy DoppelGen compiled C-extension binary distribution to Hugging Face Space

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +7 -0
  2. .github/scripts/compile_binaries.py +144 -0
  3. .github/workflows/deploy_hf.yml +77 -0
  4. .gitignore +31 -0
  5. README.md +151 -0
  6. actora/.dockerignore +8 -0
  7. actora/.gitignore +12 -0
  8. actora/README.md +253 -0
  9. actora/actora/__init__.py +0 -0
  10. actora/actora/api/__init__.py +1 -0
  11. actora/actora/api/api.cpython-312-x86_64-linux-gnu.so +3 -0
  12. actora/actora/api/main.cpython-312-x86_64-linux-gnu.so +3 -0
  13. actora/actora/core/__init__.py +1 -0
  14. actora/actora/core/arcanegan.cpython-312-x86_64-linux-gnu.so +3 -0
  15. actora/actora/core/dctnet.cpython-312-x86_64-linux-gnu.so +3 -0
  16. actora/actora/core/filters.cpython-312-x86_64-linux-gnu.so +3 -0
  17. actora/actora/core/matting.cpython-312-x86_64-linux-gnu.so +3 -0
  18. actora/actora/core/neural_accelerator.cpython-312-x86_64-linux-gnu.so +3 -0
  19. actora/actora/core/postprocess.cpython-312-x86_64-linux-gnu.so +3 -0
  20. actora/actora/models/__init__.py +1 -0
  21. actora/actora/models/base.cpython-312-x86_64-linux-gnu.so +3 -0
  22. actora/actora/models/dreamtalk.cpython-312-x86_64-linux-gnu.so +3 -0
  23. actora/actora/models/faster_liveportrait.cpython-312-x86_64-linux-gnu.so +3 -0
  24. actora/actora/models/fomm.cpython-312-x86_64-linux-gnu.so +3 -0
  25. actora/actora/models/generator.cpython-312-x86_64-linux-gnu.so +3 -0
  26. actora/actora/models/models.cpython-312-x86_64-linux-gnu.so +3 -0
  27. actora/actora/models/tpsmm.cpython-312-x86_64-linux-gnu.so +3 -0
  28. actora/actora/schema/__init__.py +0 -0
  29. actora/actora/test/create_test_inputs.py +53 -0
  30. actora/actora/test/test.py +35 -0
  31. actora/actora/utils/__init__.py +0 -0
  32. actora/actora/utils/utils.cpython-312-x86_64-linux-gnu.so +3 -0
  33. actora/actora/utils/warnings_patch.cpython-312-x86_64-linux-gnu.so +3 -0
  34. actora/docker/Dockerfile +14 -0
  35. actora/docker/docker-compose.yml +13 -0
  36. actora/docker/entrypoint.sh +3 -0
  37. actora/docker/requirements.txt +27 -0
  38. actora/preload_models.py +496 -0
  39. actora/run_actora.sh +28 -0
  40. actora/third_party/dreamtalk_src/LICENSE +21 -0
  41. actora/third_party/dreamtalk_src/README.md +99 -0
  42. actora/third_party/dreamtalk_src/configs/__init__.py +0 -0
  43. actora/third_party/dreamtalk_src/configs/default.py +91 -0
  44. actora/third_party/dreamtalk_src/core/__init__.py +0 -0
  45. actora/third_party/dreamtalk_src/core/networks/__init__.py +14 -0
  46. actora/third_party/dreamtalk_src/core/networks/diffusion_net.py +340 -0
  47. actora/third_party/dreamtalk_src/core/networks/diffusion_util.py +131 -0
  48. actora/third_party/dreamtalk_src/core/networks/disentangle_decoder.py +240 -0
  49. actora/third_party/dreamtalk_src/core/networks/dynamic_conv.py +156 -0
  50. 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
+ ![teaser](media/teaser.gif "teaser")
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)