refactor: add type annotations and improve type handling

src/visualizr/anitalker/choices.py

@@ -5,35 +5,35 @@
 
 class TrainMode(Enum):
     # manipulate mode = training the classifier
-    manipulate = "manipulate"
+    manipulate: str = "manipulate"
     # default training mode!
-    diffusion = "diffusion"
+    diffusion: str = "diffusion"
 
 
 class ModelType(Enum):
     """Kinds of the backbone models."""
 
     # unconditional ddpm
-    ddpm = "ddpm"
+    ddpm: str = "ddpm"
     # autoencoding ddpm cannot do unconditional generation
-    autoencoder = "autoencoder"
+    autoencoder: str = "autoencoder"
 
-    def has_autoenc(self):
+    def has_autoenc(self) -> bool:
         return self in [ModelType.autoencoder]
 
 
 class ModelName(Enum):
     """List of all supported model classes."""
 
-    beatgans_ddpm = "beatgans_ddpm"
-    beatgans_autoenc = "beatgans_autoenc"
+    beatgans_ddpm: str = "beatgans_ddpm"
+    beatgans_autoenc: str = "beatgans_autoenc"
 
 
 class ModelMeanType(Enum):
     """Which type of output the model predicts."""
 
     # the model predicts epsilon
-    eps = "eps"
+    eps: str = "eps"
 
 
 class ModelVarType(Enum):
@@ -45,29 +45,29 @@ class ModelVarType(Enum):
     """
 
     # posterior beta_t
-    fixed_small = "fixed_small"
+    fixed_small: str = "fixed_small"
     # beta_t
-    fixed_large = "fixed_large"
+    fixed_large: str = "fixed_large"
 
 
 class LossType(Enum):
     # use raw MSE loss and KL when learning variances
-    mse = "mse"
+    mse: str = "mse"
 
 
 class GenerativeType(Enum):
     """where how a sample is generated."""
 
-    ddpm = "ddpm"
-    ddim = "ddim"
+    ddpm: str = "ddpm"
+    ddim: str = "ddim"
 
 
 class Activation(Enum):
-    none = "none"
-    relu = "relu"
-    lrelu = "lrelu"
-    silu = "silu"
-    tanh = "tanh"
+    none: str = "none"
+    relu: str = "relu"
+    lrelu: str = "lrelu"
+    silu: str = "silu"
+    tanh: str = "tanh"
 
     def get_act(self) -> Identity | ReLU | LeakyReLU | SiLU | Tanh:
         match self:

src/visualizr/anitalker/config.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import Literal
+from typing import Literal, Self
 
 from visualizr.anitalker.choices import (
     Activation,
@@ -24,17 +24,12 @@
     ModelConfig,
 )
 from visualizr.anitalker.model.latentnet import LatentNetType, MLPSkipNetConfig
+from visualizr.app.types import InferenceType
 
 
 @dataclass
 class TrainConfig(BaseConfig):
-    infer_type: Literal[
-        "mfcc_full_control",
-        "mfcc_pose_only",
-        "hubert_pose_only",
-        "hubert_audio_only",
-        "hubert_full_control",
-    ] = None
+    infer_type: InferenceType = None
     # random seed
     seed: int = 0
     train_mode: TrainMode = TrainMode.diffusion
@@ -94,19 +89,22 @@ class TrainConfig(BaseConfig):
     T: int = 1_000
     # to be overridden
     name: str = ""
-    decoder_layers = None
-    motion_dim = None
+    decoder_layers: int = None
+    motion_dim: int = None
+    mfcc: bool = None
+    face_scale: bool = None
+    face_location: bool = None
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self.batch_size_eval = self.batch_size_eval or self.batch_size
         self.data_val_name = self.data_val_name or self.data_name
 
-    def scale_up_gpus(self, num_gpus, num_nodes=1):
+    def scale_up_gpus(self, num_gpus: int, num_nodes: int = 1) -> Self:
         self.batch_size *= num_gpus * num_nodes
         self.batch_size_eval *= num_gpus * num_nodes
         return self
 
-    def _make_diffusion_conf(self, t: int):
+    def _make_diffusion_conf(self, t: int) -> SpacedDiffusionBeatGansConfig:
         if self.diffusion_type != "beatgans":
             raise NotImplementedError
         # can use t < `self.t` for evaluation
@@ -132,21 +130,21 @@ def _make_diffusion_conf(self, t: int):
         )
 
     @property
-    def model_out_channels(self):
+    def model_out_channels(self) -> int:
         return 3
 
     def make_t_sampler(self) -> UniformSampler:
         if self.T_sampler != "uniform":
             raise NotImplementedError
         return UniformSampler(self.T)
 
-    def make_diffusion_conf(self):
+    def make_diffusion_conf(self) -> SpacedDiffusionBeatGansConfig:
         return self._make_diffusion_conf(self.T)
 
-    def make_eval_diffusion_conf(self):
+    def make_eval_diffusion_conf(self) -> SpacedDiffusionBeatGansConfig:
         return self._make_diffusion_conf(self.T_eval)
 
-    def make_model_conf(self):
+    def make_model_conf(self) -> BeatGANsAutoencConfig | BeatGANsUNetConfig:
         if self.model_name == ModelName.beatgans_ddpm:
             self.model_type = ModelType.ddpm
             self.model_conf = BeatGANsUNetConfig(
@@ -172,16 +170,9 @@ def make_model_conf(self):
                 resnet_two_cond=self.net_beatgans_resnet_two_cond,
                 resnet_use_zero_module=self.net_beatgans_resnet_use_zero_module,
             )
-        elif self.model_name in [
-            ModelName.beatgans_autoenc,
-        ]:
+        elif self.model_name == ModelName.beatgans_autoenc:
             cls = BeatGANsAutoencConfig
-            # supports both autoenc and vaeddpm
-            if self.model_name == ModelName.beatgans_autoenc:
-                self.model_type = ModelType.autoencoder
-            else:
-                raise NotImplementedError
-
+            self.model_type = ModelType.autoencoder
             if self.net_latent_net_type == LatentNetType.none:
                 latent_net_conf = None
             elif self.net_latent_net_type == LatentNetType.skip:

src/visualizr/anitalker/config_base.py

@@ -17,41 +17,40 @@ class BaseConfig:
     and serialize/deserialize configurations to/from JSON.
     """
 
-    def inherit(self, another):
+    def inherit(self, another) -> None:
         """Inherit common keys from a given config."""
-        common_keys = set(self.__dict__.keys()) & set(another.__dict__.keys())
+        common_keys: set[str] = set(self.__dict__.keys()) & set(another.__dict__.keys())
         for k in common_keys:
             setattr(self, k, getattr(another, k))
 
-    def propagate(self):
+    def propagate(self) -> None:
         """Push down the configuration to all members."""
         for _, v in self.__dict__.items():
             if isinstance(v, BaseConfig):
                 v.inherit(self)
                 v.propagate()
 
-    def from_dict(self, config_dict, strict=False):
+    def from_dict(self, config_dict: dict, strict: bool = False) -> None:
         """
         Populate configuration attributes from a dictionary.
 
         Optionally, enforcing strict key checking.
         """
         for k, v in config_dict.items():
             if not hasattr(self, k):
+                _msg: str = f"loading extra '{k}'"
                 if strict:
-                    raise ValueError(f"loading extra '{k}'")
-                _msg = f"loading extra '{k}'"
-                logger.info(_msg)
-                Info(_msg)
+                    raise ValueError(_msg)
+                logger.warning(_msg)
                 continue
             if isinstance(self.__dict__[k], BaseConfig):
                 self.__dict__[k].from_dict(v)
             else:
                 self.__dict__[k] = v
 
-    def as_dict_jsonable(self):
+    def as_dict_jsonable(self) -> dict:
         """Convert the configuration to a JSON-serializable dictionary."""
-        conf = {}
+        conf: dict = {}
         for k, v in self.__dict__.items():
             if isinstance(v, BaseConfig):
                 conf[k] = v.as_dict_jsonable()
@@ -61,10 +60,10 @@ def as_dict_jsonable(self):
 
 
 @lru_cache
-def jsonable(x: Any) -> bool:
+def jsonable(x) -> bool:
     """Check if the object x is JSON serializable."""
     try:
         dumps(x)
-        return True
     except TypeError:
         return False
+    return True