(
self,
repo_id: Optional[str] = None,
config: Union[Dict, str, None] = None,
model: Optional[str] = None,
disable_complex: bool = False
)
| 30 | } |
| 31 | |
| 32 | def __init__( |
| 33 | self, |
| 34 | repo_id: Optional[str] = None, |
| 35 | config: Union[Dict, str, None] = None, |
| 36 | model: Optional[str] = None, |
| 37 | disable_complex: bool = False |
| 38 | ): |
| 39 | super().__init__() |
| 40 | if repo_id is None: |
| 41 | repo_id = 'hexgrad/Kokoro-82M' |
| 42 | print(f"WARNING: Defaulting repo_id to {repo_id}. Pass repo_id='{repo_id}' to suppress this warning.") |
| 43 | self.repo_id = repo_id |
| 44 | if not isinstance(config, dict): |
| 45 | if not config: |
| 46 | logger.debug("No config provided, downloading from HF") |
| 47 | config = hf_hub_download(repo_id=repo_id, filename='config.json') |
| 48 | with open(config, 'r', encoding='utf-8') as r: |
| 49 | config = json.load(r) |
| 50 | logger.debug(f"Loaded config: {config}") |
| 51 | self.vocab = config['vocab'] |
| 52 | self.bert = CustomAlbert(AlbertConfig(vocab_size=config['n_token'], **config['plbert'])) |
| 53 | self.bert_encoder = torch.nn.Linear(self.bert.config.hidden_size, config['hidden_dim']) |
| 54 | self.context_length = self.bert.config.max_position_embeddings |
| 55 | self.predictor = ProsodyPredictor( |
| 56 | style_dim=config['style_dim'], d_hid=config['hidden_dim'], |
| 57 | nlayers=config['n_layer'], max_dur=config['max_dur'], dropout=config['dropout'] |
| 58 | ) |
| 59 | self.text_encoder = TextEncoder( |
| 60 | channels=config['hidden_dim'], kernel_size=config['text_encoder_kernel_size'], |
| 61 | depth=config['n_layer'], n_symbols=config['n_token'] |
| 62 | ) |
| 63 | self.decoder = Decoder( |
| 64 | dim_in=config['hidden_dim'], style_dim=config['style_dim'], |
| 65 | dim_out=config['n_mels'], disable_complex=disable_complex, **config['istftnet'] |
| 66 | ) |
| 67 | if not model: |
| 68 | try: |
| 69 | model = hf_hub_download(repo_id=repo_id, filename=KModel.MODEL_NAMES[repo_id]) |
| 70 | except: |
| 71 | model = os.path.join(repo_id, 'kokoro-v1_0.pth') |
| 72 | for key, state_dict in torch.load(model, map_location='cpu', weights_only=True).items(): |
| 73 | assert hasattr(self, key), key |
| 74 | try: |
| 75 | getattr(self, key).load_state_dict(state_dict) |
| 76 | except: |
| 77 | logger.debug(f"Did not load {key} from state_dict") |
| 78 | state_dict = {k[7:]: v for k, v in state_dict.items()} |
| 79 | getattr(self, key).load_state_dict(state_dict, strict=False) |
| 80 | |
| 81 | @property |
| 82 | def device(self): |
no test coverage detected