{{ message }}
Repository navigation
Expand file tree
/
Copy pathapp.py
More file actions
302 lines (266 loc) · 14.4 KB
/
Copy pathapp.py
File metadata and controls
302 lines (266 loc) · 14.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
"""A talking head that answers questions about your documents. One file, one command.
modal setup # once
modal secret create talking-head-llm LLM_API_KEY=sk-… # your LLM key
modal deploy app.py # -> a public https URL
Everything expensive happens ONCE, AT BUILD TIME, on Modal's infrastructure:
build : deps -> 4.4 GB of weights -> read your corpus -> prep the avatar -> idle loop
boot : load 2 GB of weights, warm the CUDA kernels, warm the voice
ask : LLM -> speech -> one forward pass per frame, streamed
`Image.run_function(..., gpu=...)` is the trick that makes this practical: it runs a
function ON A GPU DURING THE BUILD, and its filesystem writes persist into the image
layer. So face detection, the VAE latents, the blend masks and the idle loop are all
BAKED IN. A cold container never redoes them.
That matters more than it sounds. Building this on a rented GPU box means dragging a 3 GB
python environment, 4.4 GB of weights and a GPU prep step down whatever network that
particular host happens to have -- and hosts vary by three orders of magnitude. One wrote
124 KB/s. Modal builds on its own infrastructure and caches the layer, which deletes the
failure mode instead of mitigating it. It also means you never need Docker locally.
"""
import os
import modal
APP_DIR = "/app"
MUSETALK = "https://github.com/TMElyralab/MuseTalk.git"
# MuseTalk peaks at ~3.1 GB of VRAM, so this is not a big-GPU problem. An A10 is plenty
# and costs about half of an A100. Kokoro (82M) rides along on the same card for free.
GPU = os.environ.get("GPU", "A10")
# Keep a container pinned warm. Costs the full hourly rate around the clock -- see the
# cold-start note further down. 0 = scale to zero, which is what you want by default.
KEEP_WARM = int(os.environ.get("KEEP_WARM", "0"))
app = modal.App("talking-head")
def _build_corpus() -> None:
"""Read corpus/ into ONE text file, baked into the image. Runs at build time, on CPU.
Handles the formats people actually have: PDF, Word, HTML, Markdown, text. Extraction
is deliberately dumb -- get the words out, drop the chrome.
"""
import pathlib
import sys
src = pathlib.Path(f"{APP_DIR}/corpus")
out = pathlib.Path(f"{APP_DIR}/corpus.txt")
parts, skipped = [], []
for path in sorted(src.rglob("*")):
if not path.is_file() or path.name.startswith("."):
continue
suffix = path.suffix.lower()
try:
if suffix in {".txt", ".md", ".markdown", ".rst", ".csv"}:
text = path.read_text(encoding="utf-8", errors="replace")
elif suffix in {".html", ".htm"}:
from bs4 import BeautifulSoup
soup = BeautifulSoup(path.read_text(encoding="utf-8", errors="replace"), "html.parser")
# Nav, headers, footers and cookie banners are on EVERY page, so they are
# the most repeated text in the corpus. Drop them or they crowd out the
# actual content.
for junk in soup(["script", "style", "nav", "header", "footer", "aside"]):
junk.decompose()
text = soup.get_text(" ", strip=True)
elif suffix == ".pdf":
from pypdf import PdfReader
text = "\n".join((p.extract_text() or "") for p in PdfReader(str(path)).pages)
elif suffix == ".docx":
import docx
text = "\n".join(p.text for p in docx.Document(str(path)).paragraphs)
else:
skipped.append(f"{path.name} (unsupported type '{suffix}')")
continue
except Exception as exc:
skipped.append(f"{path.name} ({type(exc).__name__}: {exc})")
continue
text = text.strip()
if not text:
# A PDF that yields no text is almost always a SCAN -- an image of a page with
# no text layer. Say so, because "it's in the folder" and "she can read it"
# are different states and the gap is invisible.
skipped.append(f"{path.name} (no text extracted -- a scanned PDF? needs OCR)")
continue
parts.append(f"### {path.relative_to(src)}\n\n{text}")
body = "\n\n".join(parts)
out.write_text(body, encoding="utf-8")
# ~4 chars/token is close enough to catch the case that matters.
tokens = len(body) // 4
print(f"\ncorpus: {len(parts)} document(s), ~{tokens:,} tokens")
for s in skipped:
print(f" SKIPPED: {s}")
if not parts:
# Fail the BUILD. The alternative is a container that deploys perfectly and
# answers every question from thin air -- green everywhere, product broken, and
# the symptom is a confident wrong answer rather than an error.
print("\nNOTHING WAS INGESTED. Put documents in corpus/ (.md .txt .html .pdf .docx).",
file=sys.stderr)
raise SystemExit(1)
# The whole corpus goes in the prompt on every question -- no vector store, nothing to
# keep in sync (see service/brain.py). That stops working when it stops fitting.
if tokens > 100_000:
print(f"\nWARNING: ~{tokens:,} tokens is a lot to put in every prompt. It will be "
f"slow and expensive, and past your model's context window it will simply "
f"fail. Trim corpus/, or move to a retrieval setup.", file=sys.stderr)
def _prepare_avatar() -> None:
"""Runs ON A GPU, AT BUILD TIME. Everything it writes persists into the image.
Two expensive, fragile things happen here once, forever, instead of on every cold
container: the avatar prep (landmarks, VAE latents, blend masks) and the idle loop
(which needs the UNet, so it cannot be done in prep).
"""
import os
import subprocess
import sys
if not os.path.isfile(f"{APP_DIR}/assets/source.mp4"):
raise SystemExit(
"\n NO FACE.\n\n"
" assets/source.mp4 does not exist. This repo deliberately ships without one:\n"
" a face is a person's likeness, and you need the right to use the one you pick.\n\n"
" Record or generate a ~10 second clip, CHECK IT, and drop it in:\n\n"
" python tools/validate-avatar-clip.py my-clip.mp4\n"
" cp my-clip.mp4 assets/source.mp4\n\n"
" See docs/YOUR-AVATAR.md. The clip MUST contain a blink and a moment of\n"
" stillness, and both failures are silent if you skip the validator.\n"
)
os.chdir(APP_DIR) # MuseTalk resolves ./models/... relative to CWD
sys.path.insert(0, APP_DIR)
subprocess.run(
[sys.executable, "-m", "service.prepare",
"--avatar-id", "avatar", "--video", f"{APP_DIR}/assets/source.mp4"],
check=True, cwd=APP_DIR,
)
import torch
from service.avatar import Avatar
from service.models import load_models
from service import server as srv
models = load_models(torch.device("cuda:0"), fp16=True)
avatar = Avatar("avatar", os.environ["AVATAR_ROOT"], models, batch_size=8)
srv.AVATAR_ROOT = os.environ["AVATAR_ROOT"]
srv._ensure_idle_loop(avatar)
print("build: avatar + idle loop baked into the image")
image = (
# CUDA 11.8 / torch 2.0.1 is not nostalgia. `mim install mmcv==2.0.1` resolves a
# PREBUILT wheel keyed on exactly cu118/torch2.0. Move off it and you are in a
# 10-minute source build of mmcv, or worse.
modal.Image.from_registry(
"nvidia/cuda:11.8.0-cudnn8-devel-ubuntu22.04", add_python="3.10"
)
.apt_install("git", "ffmpeg", "libgl1", "libglib2.0-0", "espeak-ng")
.pip_install(
"torch==2.0.1", "torchvision==0.15.2", "torchaudio==2.0.2",
extra_index_url="https://download.pytorch.org/whl/cu118",
)
.pip_install(
"diffusers==0.30.2", "accelerate==0.28.0", "transformers==4.39.2",
"huggingface_hub[cli]==0.30.2",
"numpy==1.23.5", # hard pin: mmcv/xtcocotools are numpy-1.x ABI
"opencv-python-headless==4.9.0.80",
"soundfile==0.12.1", "librosa==0.11.0", "einops==0.8.1",
"omegaconf", "ffmpeg-python", "imageio[ffmpeg]", "requests",
"fastapi==0.111.0", "uvicorn[standard]==0.30.1", "python-multipart==0.0.9",
# The voice. Apache-2.0 weights and all -- see service/tts.py on why that matters
# more than the quality does. `click` is a spacy dependency that kokoro's G2P
# needs and does not declare.
"kokoro>=0.9.4", "click",
# Corpus extraction, build-time only. All pure python: they cannot disturb the
# pins above, which is the entire reason for choosing them.
"pypdf", "python-docx", "beautifulsoup4",
)
.run_commands(
# xtcocotools (an mmpose dependency) has no wheel here and compiles from source.
# pip does that in an ISOLATED build env -- which, on Modal's `add_python`
# runtime, has no `pip` in it, so it dies with a baffling "No module named 'pip'"
# that has nothing to do with mmpose. Give it its build deps and turn isolation off.
"pip install --no-cache-dir cython setuptools wheel",
"pip install --no-cache-dir --no-build-isolation xtcocotools",
"pip install --no-cache-dir -U openmim",
"mim install mmengine",
"mim install mmcv==2.0.1",
"mim install mmdet==3.1.0",
"mim install mmpose==1.1.0",
# RE-ASSERT THE PINS AFTER mim, AND PROVE IT BY IMPORTING.
#
# `mim install mmengine` silently drags in numpy 2.x AND opencv-python 5.x over the
# top of the pins above. NEITHER breaks the build. BOTH break at import, much
# later: numpy 2 against the numpy-1.x ABI that mmcv and xtcocotools are compiled
# against, and opencv-python clobbering opencv-python-headless's cv2. An image
# that builds perfectly would die on its first request.
#
# A GREEN BUILD IS NOT EVIDENCE OF A WORKING ENVIRONMENT. Import it, in the build.
"pip install --no-cache-dir numpy==1.23.5",
"pip uninstall -y opencv-python opencv-contrib-python || true",
"pip install --no-cache-dir --force-reinstall --no-deps opencv-python-headless==4.9.0.80",
"python -c \""
"import numpy, cv2, mmcv, mmpose, kokoro; from xtcocotools import coco; "
"assert numpy.__version__.startswith('1.23'), numpy.__version__; "
"assert __import__('torch').__version__.startswith('2.0.1'); "
"print('pins OK -', numpy.__version__, cv2.__version__, mmcv.__version__)\"",
)
# Shallow: we never need MuseTalk's history, and a full clone is what wedged a slow
# host for 25 minutes.
.run_commands(f"git clone --depth 1 {MUSETALK} {APP_DIR}")
# s3fd is otherwise fetched AT RUNTIME from a personal website on first prep. A first
# request that reaches out to someone's academic homepage is a failure waiting to happen.
.run_commands(
"python -c \""
"from huggingface_hub import hf_hub_download; import shutil; "
"p = hf_hub_download('ByteDance/LatentSync', 'auxiliary/s3fd-619a316812.pth'); "
f"shutil.copy(p, '{APP_DIR}/musetalk/utils/face_detection/detection/sfd/s3fd.pth')\""
)
.add_local_dir("service", f"{APP_DIR}/service", copy=True)
.add_local_dir("ui", f"{APP_DIR}/ui", copy=True)
.add_local_dir("assets", f"{APP_DIR}/assets", copy=True)
.add_local_dir("corpus", f"{APP_DIR}/corpus", copy=True)
.add_local_file("fetch_weights.sh", f"{APP_DIR}/fetch_weights.sh", copy=True)
# 4.4 GB, pulled over MODAL's network at build time -- not over a rented box's.
.run_commands(f"cd {APP_DIR} && MODEL_ROOT={APP_DIR}/models bash fetch_weights.sh")
.env({
"MODEL_ROOT": f"{APP_DIR}/models",
"AVATAR_ROOT": f"{APP_DIR}/avatars",
"CORPUS_PATH": f"{APP_DIR}/corpus.txt",
"PYTHONPATH": APP_DIR,
"DEFAULT_AVATAR": "avatar",
"BATCH_SIZE": "8",
"BLEND_THREADS": "8",
# Thread pinning is not tuning -- it is the fix for a HANG. Unpinned, cv2 and BLAS
# spawn a thread per core; on a 128-core host the mask stage wedged at 0 frames
# for 30 minutes, at 1300% CPU, with no error.
"OMP_NUM_THREADS": "4",
"OPENBLAS_NUM_THREADS": "4",
"MKL_NUM_THREADS": "4",
})
.run_function(_build_corpus) # your documents -> one text file
.run_function(_prepare_avatar, gpu=GPU) # ← the whole reason for choosing Modal
)
# ── COLD START IS THE ONE THING TO PLAN FOR ──────────────────────────────────────────
#
# Scale-to-zero is why this costs almost nothing when nobody is using it. It is also why
# the FIRST question after an idle period is slow: the container must boot and load 2 GB
# of weights, and the first render pays for cold CUDA kernels on top.
#
# cold : first word ~35s (container boot + model load dominates)
# warm : first word ~2s
#
# The page pre-warms on load, which usually hides this -- by the time someone has typed a
# question, the GPU is up. But DO NOT let a stakeholder's first click be the cold one.
#
# Before a demo, pin a container warm:
#
# KEEP_WARM=1 modal deploy app.py
#
# That bills the full GPU rate around the clock (~$1/hr), so turn it off afterwards by
# redeploying without it. For a 30-minute demo it is pennies and it removes every cold
# start from the room.
# ─────────────────────────────────────────────────────────────────────────────────────
@app.cls(
image=image,
gpu=GPU,
min_containers=KEEP_WARM, # 0 = scale to zero (cheap); 1 = always warm (for a demo)
scaledown_window=300,
timeout=600,
secrets=[modal.Secret.from_name("talking-head-llm")], # LLM_API_KEY, and optionally
# LLM_BASE_URL / LLM_MODEL
)
class Renderer:
@modal.enter()
def load(self):
"""Runs once per container, BEFORE it is marked warm. Not per request."""
import os
os.chdir(APP_DIR) # MuseTalk's hardcoded ./models/... paths
from service import server
server._startup() # idempotent; the ASGI lifespan then no-ops
@modal.asgi_app()
def web(self):
from service.server import app as fastapi_app
return fastapi_app
You can’t perform that action at this time.