Upload folder using huggingface_hub

config.json

@@ -0,0 +1,11 @@
+{
+    "_name_or_path": "oriyonay/myna-vertical",
+    "architectures": [
+        "Myna"
+    ],
+    "auto_map": {
+        "AutoConfig": "myna.MynaConfig",
+        "AutoModel": "myna.Myna"
+    },
+    "model_type": "myna"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5dde8aa186da0fc0bfdd2e5a7c1ea3e6c6af6cd4683089b65c4fed6af8335374
+size 85516720

myna.py

@@ -0,0 +1,339 @@
+'''
+Modified from the vit_pytorch library: https://github.com/lucidrains/vit-pytorch
+'''
+
+from einops import rearrange
+from einops.layers.torch import Rearrange
+import json
+import math
+from nnAudio.features.mel import MelSpectrogram
+import os
+import torch
+from torch import nn
+import torchaudio
+import torchaudio.transforms as T
+
+# for uploading to huggingface hub
+from huggingface_hub import HfApi, PyTorchModelHubMixin
+from transformers import PretrainedConfig, PreTrainedModel
+import shutil
+
+
+def pair(t):
+    return t if isinstance(t, tuple) else (t, t)
+
+
+def posemb_sincos_2d(h, w, dim, temperature: int = 10000, dtype = torch.float32):
+    y, x = torch.meshgrid(torch.arange(h), torch.arange(w), indexing="ij")
+    assert (dim % 4) == 0, "feature dimension must be multiple of 4 for sincos emb"
+    omega = torch.arange(dim // 4) / (dim // 4 - 1)
+    omega = 1.0 / (temperature ** omega)
+
+    y = y.flatten()[:, None] * omega[None, :]
+    x = x.flatten()[:, None] * omega[None, :]
+    pe = torch.cat((x.sin(), x.cos(), y.sin(), y.cos()), dim=1)
+    return pe.type(dtype)
+
+
+def load_model(model: nn.Module, checkpoint_path: str, device: str = 'cpu', ignore_layers: list = ['linear_head'], verbose: bool = False):
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+
+    filtered_state_dict = {
+        k: v for k, v in checkpoint.items() 
+        if not any(k.startswith(layer) for layer in ignore_layers)
+    }
+
+    model.load_state_dict(filtered_state_dict, strict=False)
+
+    if ignore_layers and verbose:
+        print(f'==> Loaded model from {checkpoint_path}, ignoring layers: {", ".join(ignore_layers)}')
+
+
+class FeedForward(nn.Module):
+    def __init__(self, dim, hidden_dim):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(dim),
+            nn.Linear(dim, hidden_dim),
+            nn.GELU(),
+            nn.Linear(hidden_dim, dim),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, heads = 8, dim_head = 64):
+        super().__init__()
+        inner_dim = dim_head *  heads
+        self.heads = heads
+        self.scale = dim_head ** -0.5
+        self.norm = nn.LayerNorm(dim)
+
+        self.attend = nn.Softmax(dim = -1)
+
+        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
+        self.to_out = nn.Linear(inner_dim, dim, bias = False)
+
+    def forward(self, x):
+        x = self.norm(x)
+
+        qkv = self.to_qkv(x).chunk(3, dim = -1)
+        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
+
+        dots = torch.matmul(q, k.transpose(-1, -2)) * self.scale
+
+        attn = self.attend(dots)
+
+        out = torch.matmul(attn, v)
+        out = rearrange(out, 'b h n d -> b n (h d)')
+        return self.to_out(out)
+
+
+class Transformer(nn.Module):
+    def __init__(self, dim, depth, heads, dim_head, mlp_dim):
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(nn.ModuleList([
+                Attention(dim, heads = heads, dim_head = dim_head),
+                FeedForward(dim, mlp_dim)
+            ]))
+    def forward(self, x):
+        for attn, ff in self.layers:
+            x = attn(x) + x
+            x = ff(x) + x
+        return self.norm(x)
+
+
+class MynaPreprocessor:
+    def __init__(self, target_sr: int = 16000, n_mels: int = 128):
+        self.target_sr = target_sr
+        self.n_mels = n_mels
+        self.mel_spec = MelSpectrogram(sr=target_sr, n_mels=n_mels, verbose=False)
+
+    def __call__(self, filename: str, n_frames: int = None):
+        # loads audio from file and returns a 3D tensor (B, n_mels, n_frames)
+        signal, sr = torchaudio.load(filename)
+        if signal.shape[0] > 1:
+            signal = signal.mean(dim=0, keepdim=True)
+        if sr != self.target_sr:
+            resampler = T.Resample(orig_freq=sr, new_freq=self.target_sr)
+            signal = resampler(signal)
+        ms = self.mel_spec(signal)
+
+        if n_frames:
+            ms = self._batch_spectrogram(ms, n_frames)
+
+        return ms
+    
+    def _batch_spectrogram(self, ms: torch.Tensor, n_frames: int):
+        # sanity check
+        assert ms.dim() == 3 and ms.shape[0] == 1
+
+        # discard excess frames
+        num_chunks = ms.shape[-1] // n_frames
+        ms = ms[:, :, :num_chunks * n_frames]
+
+        # split the tensor into chunks and stack them
+        chunks = torch.chunk(ms, num_chunks, dim=2)
+        batch = torch.stack(chunks)
+
+        return batch
+
+
+class MynaConfig(PretrainedConfig):
+    model_type = 'myna'
+    def __init__(
+        self, spec_size=(128, 4096), patch_size=16, dim=384, depth=12, 
+        heads=6, mlp_dim=1536, dim_head = 64, arch=None, additional_patch_size = None, 
+        hybrid_mode: bool = False, n_samples = 50000, sr = 16000, **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.spec_size = spec_size
+        self.patch_size = patch_size
+        self.dim = dim
+        self.depth = depth
+        self.heads = heads
+        self.mlp_dim = mlp_dim
+        self.dim_head = dim_head
+        self.arch = arch
+        self.additional_patch_size = additional_patch_size
+        self.hybrid_mode = hybrid_mode
+
+        self.n_samples = n_samples # number of samples for inference
+        self.sr = sr # for preprocessing
+        self.n_frames = self._get_n_frames(n_samples)
+
+        # load architecture if provided
+        if arch:
+            arch = self._get_arch(arch)
+            self.dim = arch['dim']
+            self.depth = arch['depth']
+            self.heads = arch['heads']
+            self.mlp_dim = arch['mlp_dim']
+
+    def _get_arch(self, arch: str):
+        if arch.lower() in ['vit-s-16', 'vit-s-32']:
+            # dim 384, depth 12, MLP 1536, 6 heads, 22M parameters
+            return {'dim': 384, 'depth': 12, 'mlp_dim': 1536, 'heads': 6}
+        if arch.lower() == 'vit-b-16':
+            # dim 768, depth 12, MLP 3072, 12 heads, 87M parameters
+            return {'dim': 768, 'depth': 12, 'mlp_dim': 3072, 'heads': 12}
+        if arch.lower() == 'vit-l-16':
+            # dim 1024, depth 24, MLP 4096, 16 heads, 303M parameters
+            return {'dim': 1024, 'depth': 24, 'mlp_dim': 4096, 'heads': 16}
+        
+        raise ValueError(f'Architecture {arch} not implemented')
+
+    def _get_n_frames(self, n_samples: int):
+        ''' How many frames is n_samples samples? '''
+        mel_spectrogram = MelSpectrogram(sr=self.sr, n_mels=self.spec_size[0], verbose=False)
+        patch_size_time = self.patch_size if isinstance(self.patch_size, int) else self.patch_size[1]
+        mel_frames = mel_spectrogram(torch.randn(1, 1, n_samples)).shape[-1]
+        mel_frames = math.floor(mel_frames / patch_size_time) * patch_size_time
+        return mel_frames
+
+
+class Myna(PreTrainedModel, PyTorchModelHubMixin):
+    config_class = MynaConfig
+    def __init__(self, config: MynaConfig):
+        super().__init__(config)
+
+        self.preprocessor = MynaPreprocessor()
+        self.hybrid_mode = config.hybrid_mode
+        spec_height, spec_width = pair(config.spec_size)
+        patch_height, patch_width = pair(config.patch_size)
+
+        assert spec_height % patch_height == 0 and spec_width % patch_width == 0, 'Spectrogram dimensions must be divisible by the patch size.'
+
+        self.additional_patch_size = config.additional_patch_size
+        if config.additional_patch_size:
+            patch_height_b, patch_width_b = pair(config.additional_patch_size)
+            patch_dim_b = patch_height_b * patch_width_b
+
+            self.to_patch_embedding_b, self.pos_embedding_b = self._make_embeddings(
+                patch_height_b, patch_width_b, patch_dim_b, config.dim, spec_height, spec_width
+            )
+
+        patch_dim = patch_height * patch_width
+
+        self.to_patch_embedding, self.pos_embedding = self._make_embeddings(
+            patch_height, patch_width, patch_dim, config.dim, spec_height, spec_width
+        )
+
+        self.transformer = Transformer(config.dim, config.depth, config.heads, config.dim_head, config.mlp_dim)
+
+        self.pool = 'mean'
+        self.to_latent = nn.Identity()
+
+        self.linear_head = nn.Identity()
+
+    def forward(self, spec, recurse=True):
+        if self.hybrid_mode and recurse:
+            a = self(spec, recurse=False)
+            self.toggle_embeddings()
+            b = self(spec, recurse=False)
+            self.toggle_embeddings()
+            return torch.cat((a, b), dim=-1)
+        
+        # if input shape is not 4d, make it 4d:
+        if spec.dim() == 2:
+            # unbatched: n_mels, n_frames
+            spec = spec.unsqueeze(0).unsqueeze(0)
+        elif spec.dim() == 3:
+            # batched but without channels: B, n_mels, n_frames
+            spec = spec.unsqueeze(1)
+        assert spec.dim() == 4
+
+        device = spec.device
+
+        x = self.to_patch_embedding(spec)
+        n_patches = x.shape[1] # x is of shape (B, n_patches, dim)
+        x += self.pos_embedding[:n_patches].to(device, dtype=x.dtype)
+
+        x = self.transformer(x)
+        x = x.mean(dim = 1)
+
+        x = self.to_latent(x)
+        return self.linear_head(x)
+    
+    def toggle_embeddings(self):
+        if not self.additional_patch_size: 
+            print('toggle_embeddings() called but no additional patch size provided! Ignoring call.')
+            return
+        self.to_patch_embedding, self.to_patch_embedding_b = self.to_patch_embedding_b, self.to_patch_embedding
+        self.pos_embedding, self.pos_embedding_b = self.pos_embedding_b, self.pos_embedding
+
+    def _make_embeddings(self, patch_height, patch_width, patch_dim, dim, image_height, image_width):
+        to_patch_embedding = nn.Sequential(
+            Rearrange("b c (h p1) (w p2) -> b (h w) (p1 p2 c)", p1 = patch_height, p2 = patch_width),
+            nn.LayerNorm(patch_dim),
+            nn.Linear(patch_dim, dim),
+            nn.LayerNorm(dim),
+        )
+
+        pos_embedding = posemb_sincos_2d(
+            h = image_height // patch_height,
+            w = image_width // patch_width,
+            dim = dim,
+        ) 
+
+        return to_patch_embedding, pos_embedding
+    
+    def from_file(self, filename: str, n_samples: int = None):
+        n_frames = self.config.n_frames
+        if n_samples and n_samples != self.config.n_samples:
+            n_frames = self.config._get_n_frames(n_samples)
+        spec = self.preprocessor(filename, n_frames)
+        return self(spec)
+
+    @property
+    def n_params(self):
+        return sum(p.numel() for p in self.parameters())
+    
+
+def save_model_and_push(model, repo_name, save_dir='myna-temp', to_hub=False):
+    model.save_pretrained(save_dir)
+    shutil.copy('myna.py', save_dir)
+
+    config = {
+        '_name_or_path': repo_name,
+        'architectures': ['Myna'],
+        'auto_map': {
+            'AutoConfig': 'myna.MynaConfig',
+            'AutoModel': 'myna.Myna'
+        },
+        'model_type': 'myna'
+    }
+    
+    with open(os.path.join(save_dir, 'config.json'), 'w') as f:
+        json.dump(config, f, indent=4)
+
+    print(f'Model saved locally to {save_dir}')
+
+    if to_hub:
+        api = HfApi()
+        api.create_repo(repo_name, exist_ok=True)
+        api.upload_folder(folder_path=save_dir, repo_id=repo_name)
+        print(f"Model pushed to: https://huggingface.co/{repo_name}")
+
+
+if __name__ == '__main__':
+    config = MynaConfig(
+        patch_size=(128, 2),
+        arch='vit-s-16',
+        additional_patch_size=None,
+        hybrid_mode=False
+    )
+    model = Myna(config)
+    load_model(model, 'checkpoints/myna-vertical.pth', verbose=True)
+    print(f'Model contains {model.n_params:,} parameters')
+
+    save_model_and_push(
+        model, 
+        repo_name='oriyonay/myna-vertical', 
+        save_dir='myna-vertical',
+        to_hub=True
+    )