Deploy ARTalk ZeroGPU Space

.gitattributes

@@ -43,3 +43,6 @@ assets/GAGAvatar/examples/6.jpg filter=lfs diff=lfs merge=lfs -text
 assets/GAGAvatar/examples/7.jpg filter=lfs diff=lfs merge=lfs -text
 assets/GAGAvatar/examples/9.jpg filter=lfs diff=lfs merge=lfs -text
 assets/GAGAvatar/gagavatar_logo.png filter=lfs diff=lfs merge=lfs -text
+wheels/*.whl filter=lfs diff=lfs merge=lfs -text
+assets/*.pt filter=lfs diff=lfs merge=lfs -text
+app/**/*.pt filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.pyc
+.space_checkout/
+render_results/
+.env

README.md

@@ -4,11 +4,20 @@ emoji: 🎭
 colorFrom: blue
 colorTo: purple
 sdk: gradio
-sdk_version: 5.9.1
+sdk_version: 4.44.0
 app_file: app.py
 pinned: false
+license: mit
+short_description: Speech-driven 3D head animation for YouTube automation
+models:
+  - xg-chu/ARTalk
+  - xg-chu/GAGAvatar
+hardware: zero-a10g
+suggested_hardware: zero-a10g
 ---
 
-# ARTalk
+# ARTalk on ZeroGPU
 
-Speech-driven 3D head animation. Powered by HuggingFace ZeroGPU.
+Speech → Talking head mp4. Powered by [ARTalk](https://github.com/xg-chu/ARTalk) + GAGAvatar + ZeroGPU.
+
+See [C:/DK/AT/deploy/hf-space/README.md] for deployment details.

app.py

@@ -88,14 +88,18 @@ def _duration_for(audio_path, shape_id, style_id, clip_length):
 
 
 @spaces.GPU(duration=_duration_for)
-def generate(audio_path: str, shape_id: str, style_id: str, clip_length: int):
+def generate(audio_path: str, shape_id: str, style_id: str, clip_length: int, tracked_pt_file=None):
     """Generate talking-head mp4 from audio + shape + style.
 
     Inputs:
-      audio_path: uploaded audio file path (Gradio writes to temp)
-      shape_id:   "mesh" or pre-registered character name
-      style_id:   one of STYLE_MOTIONS (or "default")
-      clip_length: max frames at 25fps (750 = 30s)
+      audio_path:      uploaded audio file path (Gradio writes to temp)
+      shape_id:        "mesh", pre-registered character, or custom avatar id
+                       (matching the key in an uploaded tracked_pt_file)
+      style_id:        one of STYLE_MOTIONS (or "default")
+      clip_length:     max frames at 25fps (750 = 30s)
+      tracked_pt_file: (optional) user-supplied tracked.pt produced by the
+                       character-tracking Colab notebook. Registers new
+                       avatar ids at runtime.
 
     Returns:
       mp4 file path (Gradio streams to client)
@@ -121,11 +125,23 @@ def generate(audio_path: str, shape_id: str, style_id: str, clip_length: int):
     else:
         ENGINE.style_motion = None
 
-    # 3. Shape fallback
-    if shape_id not in AVAILABLE_SHAPES:
+    # 3. Optional custom avatar registration
+    custom_ids = []
+    if tracked_pt_file:
+        custom_path = tracked_pt_file if isinstance(tracked_pt_file, str) else getattr(tracked_pt_file, "name", None)
+        if custom_path and os.path.exists(custom_path):
+            try:
+                custom_ids = ENGINE.add_custom_avatars(custom_path)
+                print(f"[generate] registered custom avatars: {custom_ids}")
+            except Exception as e:
+                raise gr.Error(f"Custom tracked.pt load failed: {e}")
+
+    # 4. Shape fallback — include freshly-registered custom ids
+    valid_shapes = set(AVAILABLE_SHAPES) | set(custom_ids) | set(ENGINE.GAGAvatar.all_gagavatar_id.keys())
+    if shape_id not in valid_shapes:
         shape_id = "mesh"
 
-    # 4. Inference
+    # 5. Inference
     pred_motions = ENGINE.inference(audio, clip_length=int(clip_length))
 
     # 5. Render → mp4
@@ -145,9 +161,10 @@ demo = gr.Interface(
     fn=generate,
     inputs=[
         gr.Audio(type="filepath", label="Audio (wav/mp3/m4a)"),
-        gr.Textbox(value="mesh", label="shape_id (use 'mesh' or pre-registered avatar name)"),
+        gr.Textbox(value="mesh", label="shape_id (use 'mesh', a built-in, or a custom avatar id from tracked_pt_file)"),
         gr.Textbox(value="natural_0", label="style_id (e.g. natural_0, happy_0, curious_0)"),
         gr.Number(value=750, label="clip_length (max frames @ 25fps, 750=30s)", precision=0),
+        gr.File(label="tracked_pt_file (optional — custom character from Colab notebook)", file_types=[".pt"]),
     ],
     outputs=gr.Video(label="Generated talking head"),
     title="ARTalk — Speech-Driven 3D Head Animation",
@@ -208,6 +225,16 @@ _API_INFO = {
                     "component": "Number",
                     "example_input": 750,
                 },
+                {
+                    "label": "tracked_pt_file",
+                    "parameter_name": "tracked_pt_file",
+                    "parameter_has_default": True,
+                    "parameter_default": None,
+                    "type": {"type": "string"},
+                    "python_type": {"type": "Dict(path: str, url: str | None, size: int | None, orig_name: str | None, mime_type: str | None, is_stream: bool, meta: Dict(_type: Literal['gradio.FileData'])) | None", "description": ""},
+                    "component": "File",
+                    "example_input": None,
+                },
             ],
             "returns": [
                 {

inference_engine.py

@@ -156,6 +156,54 @@ class ARTAvatarInferEngine:
                 n_shape=300, n_exp=100, scale=5.0, no_lmks=True
             ).to(device)
 
+    def add_custom_avatars(self, tracked_pt_path):
+        """Merge a user-supplied tracked.pt into GAGAvatar.all_gagavatar_id.
+
+        Accepted formats:
+          1. dict-of-dicts: {"alice": {image, shapecode, transform_matrix, ...}}
+             → all entries merged as-is.
+          2. single entry dict: {image, shapecode, transform_matrix, ...}
+             → registered under key derived from filename (without 'tracked_' prefix and '.pt').
+
+        Returns list of newly registered avatar_ids.
+        """
+        if not hasattr(self, "GAGAvatar"):
+            raise RuntimeError("GAGAvatar not loaded (load_gaga=False)")
+        loaded = torch.load(tracked_pt_path, map_location="cpu", weights_only=False)
+        if not isinstance(loaded, dict):
+            raise ValueError(f"tracked.pt must be a dict, got {type(loaded).__name__}")
+
+        required = {"image", "shapecode", "transform_matrix"}
+        added = []
+
+        # Detect format: dict-of-dicts vs single entry
+        values_are_dicts = all(isinstance(v, dict) for v in loaded.values()) if loaded else False
+        looks_like_single = required.issubset(loaded.keys())
+
+        if looks_like_single and not values_are_dicts:
+            stem = Path(tracked_pt_path).stem
+            if stem.startswith("tracked_"):
+                stem = stem[len("tracked_"):]
+            avatar_id = stem or "custom"
+            entries = {avatar_id: loaded}
+        elif values_are_dicts:
+            entries = loaded
+        else:
+            raise ValueError(
+                "tracked.pt must be either a single-avatar dict with keys "
+                "{image, shapecode, transform_matrix} or a dict-of-dicts mapping "
+                "avatar_id -> {image, shapecode, transform_matrix, ...}"
+            )
+
+        for avatar_id, entry in entries.items():
+            missing = required - set(entry.keys())
+            if missing:
+                raise ValueError(f"avatar '{avatar_id}' missing keys: {missing}")
+            self.GAGAvatar.all_gagavatar_id[avatar_id] = entry
+            added.append(avatar_id)
+
+        return added
+
     def set_style_motion(self, style_motion):
         if isinstance(style_motion, str):
             style_motion = torch.load(