| """Self-contained PyTorch loader for ThomasThebaud/PROPS_K16.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| from torch.nn import functional as F |
|
|
|
|
| class ComposedGMMMDN(nn.Module): |
| """Map a 384-D description embedding to a GMM over 192-D speaker embeddings.""" |
|
|
| def __init__( |
| self, |
| input_dim: int = 384, |
| output_dim: int = 192, |
| effective_components: int = 15072, |
| hidden_dims=(1024, 2048, 1024), |
| dropout: float = 0.1, |
| min_sigma: float = 1e-4, |
| ): |
| super().__init__() |
| self.input_dim = input_dim |
| self.output_dim = output_dim |
| self.K = effective_components |
| self.min_sigma = min_sigma |
|
|
| layers = [] |
| previous = input_dim |
| for width in hidden_dims: |
| layers.extend( |
| [nn.Linear(previous, width), nn.LayerNorm(width), nn.GELU(), nn.Dropout(dropout)] |
| ) |
| previous = width |
| self.backbone = nn.Sequential(*layers) |
| self.pi_head = nn.Linear(previous, effective_components) |
| self.mu = nn.Parameter(torch.empty(effective_components, output_dim), requires_grad=False) |
| self.raw_sigma = nn.Parameter( |
| torch.empty(effective_components, output_dim), requires_grad=False |
| ) |
|
|
| |
| |
| requested_components = 16 |
| self.mu_head = nn.Linear(previous, requested_components * output_dim) |
| self.sigma_head = nn.Linear(previous, requested_components * output_dim) |
|
|
| def forward(self, description_embeddings: torch.Tensor): |
| h = self.backbone(description_embeddings) |
| pi_logits = self.pi_head(h) |
| sigma = F.softplus(self.raw_sigma) + self.min_sigma |
| return pi_logits, self.mu.unsqueeze(0).expand(h.shape[0], -1, -1), sigma.unsqueeze(0).expand( |
| h.shape[0], -1, -1 |
| ) |
|
|
| @torch.inference_mode() |
| def distribution(self, description_embeddings: torch.Tensor): |
| """Return normalized weights, component means, and diagonal standard deviations.""" |
| pi_logits, mu, sigma = self(description_embeddings) |
| return torch.softmax(pi_logits, dim=-1), mu, sigma |
|
|
| @torch.inference_mode() |
| def sample(self, description_embeddings: torch.Tensor, samples_per_prompt: int = 1): |
| """Draw speaker embeddings; output shape is [batch, samples_per_prompt, 192].""" |
| pi, mu, sigma = self.distribution(description_embeddings) |
| component = torch.multinomial(pi, samples_per_prompt, replacement=True) |
| gather_index = component.unsqueeze(-1).expand(-1, -1, self.output_dim) |
| chosen_mu = mu.gather(1, gather_index) |
| chosen_sigma = sigma.gather(1, gather_index) |
| return chosen_mu + chosen_sigma * torch.randn_like(chosen_mu) |
|
|
| @classmethod |
| def from_pretrained(cls, repo_or_path="ThomasThebaud/PROPS_K16", device="cpu"): |
| path = Path(repo_or_path) |
| if path.is_dir(): |
| config_path = path / "config.json" |
| else: |
| from huggingface_hub import hf_hub_download |
|
|
| config_path = Path(hf_hub_download(repo_or_path, "config.json")) |
| config = json.loads(config_path.read_text()) |
| if path.is_dir(): |
| checkpoint_path = path / config["checkpoint_file"] |
| else: |
| checkpoint_path = Path(hf_hub_download(repo_or_path, config["checkpoint_file"])) |
|
|
| model = cls( |
| input_dim=config["input_dim"], |
| output_dim=config["output_dim"], |
| effective_components=config["effective_components"], |
| hidden_dims=tuple(config["hidden_dims"]), |
| dropout=config["dropout"], |
| min_sigma=config["min_sigma"], |
| ) |
| try: |
| checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True) |
| except TypeError: |
| checkpoint = torch.load(checkpoint_path, map_location="cpu") |
| model.load_state_dict(checkpoint["model_state_dict"], strict=True) |
| return model.to(device).eval() |
|
|
