Create quantizer_config.py

quantizer_config.py

@@ -0,0 +1,167 @@
+
+from transformers import PretrainedConfig
+from typing import List, Optional
+
+class QuantizerConfig(PretrainedConfig):
+    model_type = "prosody_quantizer"
+    
+    def __init__(
+        self,
+        # VQ parameters
+        l_bins: int = 320,
+        emb_width: int = 64,
+        mu: float = 0.99,
+        levels: int = 1,
+        
+        # Encoder parameters
+        encoder_input_emb_width: int = 3,
+        encoder_output_emb_width: int = 64,
+        encoder_levels: int = 1,
+        encoder_downs_t: List[int] = [4],
+        encoder_strides_t: List[int] = [2],
+        encoder_width: int = 32,
+        encoder_depth: int = 4,
+        encoder_m_conv: float = 1.0,
+        encoder_dilation_growth_rate: int = 3,
+        
+        # Decoder parameters
+        decoder_input_emb_width: int = 3,
+        decoder_output_emb_width: int = 64,
+        decoder_levels: int = 1,
+        decoder_downs_t: List[int] = [4],
+        decoder_strides_t: List[int] = [2],
+        decoder_width: int = 32,
+        decoder_depth: int = 4,
+        decoder_m_conv: float = 1.0,
+        decoder_dilation_growth_rate: int = 3,
+        
+        # Training parameters
+        lambda_commit: float = 0.02,
+        f0_normalize: bool = True,
+        intensity_normalize: bool = True,
+        multispkr: str = "single",
+        f0_feats: bool = False,
+        f0_median: bool = False,
+        
+        # Optional training hyperparameters
+        learning_rate: float = 0.0002,
+        adam_b1: float = 0.8,
+        adam_b2: float = 0.99,
+        lr_decay: float = 0.999,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        
+        # VQ parameters
+        self.l_bins = l_bins
+        self.emb_width = emb_width
+        self.mu = mu
+        self.levels = levels
+        
+        # Encoder parameters
+        self.encoder_input_emb_width = encoder_input_emb_width
+        self.encoder_output_emb_width = encoder_output_emb_width
+        self.encoder_levels = encoder_levels
+        self.encoder_downs_t = encoder_downs_t
+        self.encoder_strides_t = encoder_strides_t
+        self.encoder_width = encoder_width
+        self.encoder_depth = encoder_depth
+        self.encoder_m_conv = encoder_m_conv
+        self.encoder_dilation_growth_rate = encoder_dilation_growth_rate
+        
+        # Decoder parameters
+        self.decoder_input_emb_width = decoder_input_emb_width
+        self.decoder_output_emb_width = decoder_output_emb_width
+        self.decoder_levels = decoder_levels
+        self.decoder_downs_t = decoder_downs_t
+        self.decoder_strides_t = decoder_strides_t
+        self.decoder_width = decoder_width
+        self.decoder_depth = decoder_depth
+        self.decoder_m_conv = decoder_m_conv
+        self.decoder_dilation_growth_rate = decoder_dilation_growth_rate
+        
+        # Training parameters
+        self.lambda_commit = lambda_commit
+        self.f0_normalize = f0_normalize
+        self.intensity_normalize = intensity_normalize
+        self.multispkr = multispkr
+        self.f0_feats = f0_feats
+        self.f0_median = f0_median
+        
+        # Training hyperparameters
+        self.learning_rate = learning_rate
+        self.adam_b1 = adam_b1
+        self.adam_b2 = adam_b2
+        self.lr_decay = lr_decay
+    
+    @property
+    def f0_vq_params(self):
+        return {
+            "l_bins": self.l_bins,
+            "emb_width": self.emb_width,
+            "mu": self.mu,
+            "levels": self.levels
+        }
+    
+    @property
+    def f0_encoder_params(self):
+        return {
+            "input_emb_width": self.encoder_input_emb_width,
+            "output_emb_width": self.encoder_output_emb_width,
+            "levels": self.encoder_levels,
+            "downs_t": self.encoder_downs_t,
+            "strides_t": self.encoder_strides_t,
+            "width": self.encoder_width,
+            "depth": self.encoder_depth,
+            "m_conv": self.encoder_m_conv,
+            "dilation_growth_rate": self.encoder_dilation_growth_rate
+        }
+    
+    @property
+    def f0_decoder_params(self):
+        return {
+            "input_emb_width": self.decoder_input_emb_width,
+            "output_emb_width": self.decoder_output_emb_width,
+            "levels": self.decoder_levels,
+            "downs_t": self.decoder_downs_t,
+            "strides_t": self.decoder_strides_t,
+            "width": self.decoder_width,
+            "depth": self.decoder_depth,
+            "m_conv": self.decoder_m_conv,
+            "dilation_growth_rate": self.decoder_dilation_growth_rate
+        }
+    
+    @classmethod
+    def from_yaml(cls, yaml_path: str):
+        """Load config from yaml file"""
+        import yaml
+        with open(yaml_path, 'r') as f:
+            config = yaml.safe_load(f)
+        
+        # Convert yaml config to kwargs
+        kwargs = {
+            # VQ params
+            **{k: v for k, v in config['f0_vq_params'].items()},
+            
+            # Encoder params
+            **{f"encoder_{k}": v for k, v in config['f0_encoder_params'].items()},
+            
+            # Decoder params  
+            **{f"decoder_{k}": v for k, v in config['f0_decoder_params'].items()},
+            
+            # Training params
+            "lambda_commit": config.get('lambda_commit', 0.02),
+            "f0_normalize": config.get('f0_normalize', True),
+            "intensity_normalize": config.get('intensity_normalize', True),
+            "multispkr": config.get('multispkr', "single"),
+            "f0_feats": config.get('f0_feats', False),
+            "f0_median": config.get('f0_median', False),
+            
+            # Training hyperparams
+            "learning_rate": config.get('learning_rate', 0.0002),
+            "adam_b1": config.get('adam_b1', 0.8), 
+            "adam_b2": config.get('adam_b2', 0.99),
+            "lr_decay": config.get('lr_decay', 0.999),
+        }
+        
+        return cls(**kwargs)