Upload 2 files

PackedAvatar.py

@@ -214,14 +214,238 @@ def mount_zip_payload(zip_bytes: bytes, zip_sha256: str, target_dir: Path, marke
 # ============================================================
 # AVATAR BANK RUNTIME
 # ============================================================
-
 class AvatarBankRuntime:
-    def __init__(self, payload: Dict[str, Any], defaults: Optional[Dict[str, Any]] = None):
+    def __init__(
+        self,
+        payload: Dict[str, Any],
+        defaults: Optional[Dict[str, Any]] = None,
+    ):
         self.index: Dict[str, Dict[str, Any]] = payload.get("index", {}) or {}
         self.embeddings: Dict[str, Dict[str, Any]] = payload.get("embeddings", {}) or {}
         self.previews: Dict[str, Any] = payload.get("previews", {}) or {}
         self.defaults = defaults or {}
 
+    @classmethod
+    def load(
+        cls,
+        path: Path,
+        defaults: Optional[Dict[str, Any]] = None,
+    ) -> "AvatarBankRuntime":
+        payload = torch.load(str(path), map_location="cpu", weights_only=False)
+
+        if not isinstance(payload, dict):
+            raise ValueError(f"Avatar bank file did not contain a dictionary: {path}")
+
+        return cls(payload, defaults=defaults)
+
+    def save(self, path: Union[str, Path]) -> None:
+        torch.save(
+            {
+                "index": self.index,
+                "embeddings": self.embeddings,
+                "previews": self.previews,
+            },
+            str(path),
+        )
+
+    # --------------------------------------------------------
+    # BASIC ACCESS
+    # --------------------------------------------------------
+
+    def __contains__(self, avatar_id: str) -> bool:
+        return avatar_id in self.index
+
+    def exists(self, avatar_id: str) -> bool:
+        return avatar_id in self.index
+
+    def available_ids(self) -> List[str]:
+        return list(self.index.keys())
+
+    def list_avatars(self) -> List[str]:
+        return self.available_ids()
+
+    def get_metadata(self, avatar_id: str) -> Dict[str, Any]:
+        if avatar_id not in self.index:
+            raise KeyError(f"Avatar not found: {avatar_id}")
+
+        return dict(self.index[avatar_id])
+
+    def get_avatar(self, avatar_id: str) -> Dict[str, Any]:
+        return self.build_avatar_condition(avatar_id)
+
+    def get_embedding_bundle(self, avatar_id: str) -> Dict[str, Any]:
+        if avatar_id not in self.embeddings:
+            raise KeyError(f"Avatar not found: {avatar_id}")
+
+        return self.embeddings[avatar_id]
+
+    # --------------------------------------------------------
+    # FUZZY SEARCH
+    # --------------------------------------------------------
+
+    def _fuzzy_match_single(
+        self,
+        query: str,
+        choices: set,
+        cutoff: float = 0.6,
+    ):
+        if not query:
+            return None
+
+        query_lower = query.lower()
+
+        choice_map = {
+            c.lower(): c
+            for c in choices
+        }
+
+        matches = get_close_matches(
+            query_lower,
+            list(choice_map.keys()),
+            n=1,
+            cutoff=cutoff,
+        )
+
+        return choice_map[matches[0]] if matches else None
+
+    def fuzzy_search_id(
+        self,
+        query_id: str,
+        n: int = 5,
+        cutoff: float = 0.5,
+    ) -> List[str]:
+        query_lower = query_id.lower()
+
+        id_map = {
+            aid.lower(): aid
+            for aid in self.index.keys()
+        }
+
+        matches = get_close_matches(
+            query_lower,
+            list(id_map.keys()),
+            n=n,
+            cutoff=cutoff,
+        )
+
+        return [id_map[m] for m in matches]
+
+    # --------------------------------------------------------
+    # QUERY
+    # --------------------------------------------------------
+
+    def query(
+        self,
+        gender=None,
+        style=None,
+        fuzzy=True,
+        cutoff=0.6,
+    ) -> List[str]:
+
+        known_genders = {
+            meta["gender"]
+            for meta in self.index.values()
+            if meta.get("gender")
+        }
+
+        known_styles = {
+            meta["style"]
+            for meta in self.index.values()
+            if meta.get("style")
+        }
+
+        target_gender = gender
+        target_style = style
+
+        if fuzzy:
+            if gender:
+                target_gender = (
+                    self._fuzzy_match_single(
+                        gender,
+                        known_genders,
+                        cutoff,
+                    )
+                    or gender
+                )
+
+            if style:
+                target_style = (
+                    self._fuzzy_match_single(
+                        style,
+                        known_styles,
+                        cutoff,
+                    )
+                    or style
+                )
+
+        results = []
+
+        for aid, meta in self.index.items():
+
+            if (
+                target_gender
+                and meta.get("gender") != target_gender
+            ):
+                continue
+
+            if (
+                target_style
+                and meta.get("style") != target_style
+            ):
+                continue
+
+            results.append(aid)
+
+        return results
+
+    # --------------------------------------------------------
+    # PREVIEWS
+    # --------------------------------------------------------
+
+    def get_preview(self, avatar_id: str):
+        if avatar_id not in self.previews:
+            raise KeyError(f"Avatar not found: {avatar_id}")
+
+        return decode_png_or_zstd_image(
+            self.previews[avatar_id]
+        )
+
+    def get_preview_numpy(
+        self,
+        avatar_id: str,
+    ) -> Optional[np.ndarray]:
+        return self._preview_to_numpy(avatar_id)
+
+    def _preview_to_numpy(
+        self,
+        avatar_id: str,
+    ) -> Optional[np.ndarray]:
+
+        blob = self.previews.get(avatar_id)
+
+        if blob is None:
+            return None
+
+        try:
+            img = decode_png_or_zstd_image(blob)
+            return pil_to_numpy_rgb(img)
+
+        except Exception:
+            return None
+
+    # --------------------------------------------------------
+    # MUTATION
+    # --------------------------------------------------------
+
+    def delete_avatar(
+        self,
+        avatar_id: str,
+    ) -> None:
+
+        self.index.pop(avatar_id, None)
+        self.embeddings.pop(avatar_id, None)
+        self.previews.pop(avatar_id, None)
+
     @classmethod
     def load(cls, path: Path, defaults: Optional[Dict[str, Any]] = None) -> "AvatarBankRuntime":
         payload = torch.load(str(path), map_location="cpu", weights_only=False)
@@ -1254,6 +1478,30 @@ class PackedAvatar:
         subprocess.run(cmd, cwd=str(repo), check=True)
         return out_video
 
+    def list_avatars(self):
+        return self.avatar_bank.list_avatars()
+
+    def query_avatars(self, *args, **kwargs):
+        return self.avatar_bank.query(*args, **kwargs)
+
+    def fuzzy_search_avatar(self, query, n=5, cutoff=0.5):
+        return self.avatar_bank.fuzzy_search_id(query, n, cutoff)
+
+    def get_avatar(self, avatar_id):
+        return self.avatar_bank.get_avatar(avatar_id)
+
+    def get_avatar_preview(self, avatar_id):
+        return self.avatar_bank.get_preview(avatar_id)
+
+    def get_avatar_metadata(self, avatar_id):
+        return self.avatar_bank.get_metadata(avatar_id)
+
+    def delete_avatar(self, avatar_id):
+        self.avatar_bank.delete_avatar(avatar_id)
+
+    def save_avatar_bank(self, path):
+        self.avatar_bank.save(path)
+
     def generate(
         self,
         source_image: Optional[str] = None,

README.md

@@ -530,6 +530,223 @@ If `avatar_condition` is provided, it supersedes `source_image` conditioning.
 ## 6. Optional post-processing
 
 * Wav2Lip GAN lip-sync enhancement
+# AvatarBank API
+
+PackedAvatar exposes the packed AvatarBank directly, allowing you to browse, search, inspect, and manage avatars at runtime.
+
+The AvatarBank is loaded automatically when the model is initialized.
+
+```python
+from PackedAvatar import PackedAvatar
+
+model = PackedAvatar("PackedAvatar.pt")
+```
+
+## List available avatars
+
+```python
+avatars = model.list_avatars()
+
+print(avatars)
+```
+
+Returns:
+
+```python
+[
+    "Aaron",
+    "Rebecca",
+    "Amy",
+    ...
+]
+```
+
+---
+
+## Search by name
+
+Perform a fuzzy search against avatar IDs.
+
+```python
+matches = model.fuzzy_search_avatar("rebeca")
+
+print(matches)
+```
+
+Example:
+
+```python
+["Rebecca"]
+```
+
+---
+
+## Query by gender and style
+
+Search the bank using metadata filters.
+
+```python
+female_anime = model.query_avatars(
+    gender="female",
+    style="anime"
+)
+
+print(female_anime)
+```
+
+Example:
+
+```python
+[
+    "Alison",
+    "Amber",
+    "Andrea",
+    ...
+]
+```
+
+Fuzzy matching is supported:
+
+```python
+model.query_avatars(
+    gender="femal",
+    style="anim"
+)
+```
+
+---
+
+## Retrieve avatar metadata
+
+```python
+metadata = model.get_avatar_metadata("Rebecca")
+
+print(metadata)
+```
+
+Example:
+
+```python
+{
+    "gender": "female",
+    "style": "anime"
+}
+```
+
+---
+
+## Retrieve avatar conditioning
+
+Access the full conditioning bundle used internally by PackedAvatar.
+
+```python
+avatar = model.get_avatar("Rebecca")
+```
+
+Returned fields may include:
+
+```python
+{
+    "avatar_id": "Rebecca",
+    "gender": "female",
+    "style": "anime",
+    "coeff_3dmm": ...,
+    "motion_3dmm": ...,
+    "full_3dmm": ...,
+    "crop_info": ...,
+    "crop_preview": ...
+}
+```
+
+This bundle can be passed directly as an `avatar_condition`.
+
+```python
+video = model.generate(
+    avatar_condition=avatar,
+    driven_audio="speech.wav"
+)
+```
+
+---
+
+## Retrieve avatar previews
+
+Preview images stored in the AvatarBank can be accessed directly.
+
+```python
+preview = model.get_avatar_preview("Rebecca")
+
+preview.show()
+```
+
+Returns a PIL image.
+
+---
+
+## Delete avatars
+
+Remove an avatar from the in-memory runtime.
+
+```python
+model.delete_avatar("Rebecca")
+```
+
+This only affects the current runtime session.
+
+To persist changes, save the bank.
+
+---
+
+## Save AvatarBank changes
+
+```python
+model.save_avatar_bank("ModifiedAvatarBank.pt")
+```
+
+This writes the current AvatarBank state to disk.
+
+---
+
+## Check whether an avatar exists
+
+```python
+exists = model.avatar_bank.exists("Rebecca")
+
+print(exists)
+```
+
+Returns:
+
+```python
+True
+```
+
+---
+
+## Direct AvatarBank access
+
+The underlying AvatarBank runtime is also exposed.
+
+```python
+bank = model.avatar_bank
+```
+
+Available methods include:
+
+```python
+bank.available_ids()
+bank.list_avatars()
+bank.exists(...)
+bank.query(...)
+bank.fuzzy_search_id(...)
+bank.get_avatar(...)
+bank.get_metadata(...)
+bank.get_preview(...)
+bank.delete_avatar(...)
+bank.save(...)
+```
+
+This provides full programmatic access to the packed avatar database.
 
 ---