Delete modelA.py

modelA.py

@@ -1,1047 +0,0 @@
-import os
-import math
-import warnings
-import logging
-from itertools import chain
-import torch
-import torch.nn.functional as F
-from torch import nn, Tensor
-from tensordict import TensorDict
-from typing import Optional, Dict, Union, List, Tuple
-import numpy as np
-from functools import partial
-from datetime import datetime
-from tensordict import TensorDict
-from transformers.trainer_seq2seq import Seq2SeqTrainer
-from transformers.training_args_seq2seq import Seq2SeqTrainingArguments
-from echoutils import *
-
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-dtype = torch.float32
-warnings.filterwarnings("ignore")
-logging.basicConfig(level=logging.ERROR)
-
-class rotary(nn.Module):
-    def __init__(self, dims, head, max_ctx=1500, radii=False, debug: List[str] = [], use_pbias=False, axial=False, spec_shape=None):
-
-        super(rotary, self).__init__()
-        self.use_pbias = use_pbias
-        self.dims = dims
-        self.head = head
-        self.head_dim = dims // head
-        self.radii = radii
-        self.debug = debug
-        self.counter = 0
-        self.last_theta = None
-        self.axial = axial
-
-        self.bias = nn.Parameter(torch.zeros(max_ctx, dims // 2), requires_grad=True if use_pbias else False)
-        theta = (torch.tensor(10000, device=device, dtype=dtype))
-        self.theta = nn.Parameter(theta, requires_grad=True)    
-        self.theta_values = []
-
-        if axial and spec_shape is not None:
-            time_frames, freq_bins = spec_shape
-            self.time_frames = time_frames
-            self.freq_bins = freq_bins
-            
-            time_theta = 50.0
-            time_freqs = 1.0 / (time_theta ** (torch.arange(0, dims, 4)[:(dims // 4)].float() / dims))
-            self.register_buffer('time_freqs', time_freqs)
-            
-            freq_theta = 100.0
-            freq_freqs = 1.0 / (freq_theta ** (torch.arange(0, dims, 4)[:(dims // 4)].float() / dims))
-            self.register_buffer('freq_freqs', freq_freqs)
-
-    def pitch_bias(self, f0):
-        if f0 is None:
-            return None
-        f0_flat = f0.squeeze().float()
-        f0_norm = (f0_flat - f0_flat.mean()) / (f0_flat.std() + 1e-8)
-        f0_sim = torch.exp(-torch.cdist(f0_norm.unsqueeze(1), 
-                                    f0_norm.unsqueeze(1)))
-        return f0_sim.unsqueeze(0).unsqueeze(0)
-
-    def theta_freqs(self, theta):
-        if theta.dim() == 0:
-            theta = theta.unsqueeze(0)
-        freq = (theta.unsqueeze(-1) / 220.0) * 700 * (
-            torch.pow(10, torch.linspace(0, 2595 * torch.log10(torch.tensor(1 + 8000/700)), 
-                    self.head_dim // 2, device=theta.device, dtype=theta.dtype) / 2595) - 1) / 1000
-        return freq
-
-    def _apply_radii(self, freqs, f0, ctx):
-        if self.radii and f0 is not None:
-            radius = f0.to(device, dtype)
-            L = radius.shape[0]
-            if L != ctx:
-                F = L / ctx
-                idx = torch.arange(ctx, device=f0.device)
-                idx = (idx * F).long().clamp(0, L - 1)
-                radius = radius[idx]
-                return torch.polar(radius.unsqueeze(-1), freqs), radius
-            else:
-                return torch.polar(radius.unsqueeze(-1), freqs), radius
-        else:
-            return torch.polar(torch.ones_like(freqs), freqs), None
-
-    def check_f0(self, f0, f0t, ctx):
-        if f0 is not None and f0.shape[1] == ctx:
-            return f0
-        elif f0t is not None and f0t.shape[1] == ctx:
-            return f0t
-        else:
-            return None         
-
-    def axial_freqs(self, ctx):
-        if not self.axial:
-            return None
-        time_frames = self.time_frames
-        freq_bins = self.freq_bins
-    
-        t = torch.arange(ctx, device=device, dtype=dtype)
-        t_x = (t % time_frames).float()
-        t_y = torch.div(t, time_frames, rounding_mode='floor').float()
-        freqs_x = torch.outer(t_x, self.time_freqs)
-        freqs_y = torch.outer(t_y, self.freq_freqs)
-        freqs_cis_x = torch.polar(torch.ones_like(freqs_x), freqs_x)
-        freqs_cis_y = torch.polar(torch.ones_like(freqs_y), freqs_y)
-        return torch.cat([freqs_cis_x, freqs_cis_y], dim=-1)
-
-    def forward(self, x=None, en=None, f=None, layer=None) -> Tensor:
-        ctx=x
-        f0 = en.get("f0") if en is not None else None 
-        f0t = en.get("f0t") if en is not None else None 
-
-        f0 = self.check_f0(f0, f0t, ctx)
-        if f0 is not None:
-            if f0.dim() == 2:
-                f0 = f0.squeeze(0) 
-            theta = f0 + self.theta  
-        else:
-            theta = self.theta 
-        freqs = self.theta_freqs(theta)
-        t = torch.arange(ctx, device=device, dtype=dtype)
-        freqs = t[:, None] * freqs
-        freqs, radius = self._apply_radii(freqs, f0, ctx)
-
-        if self.axial and f == "spectrogram":
-            freqs_2d = self.axial_freqs(ctx)
-            if freqs_2d is not None:
-                return freqs_2d.unsqueeze(0)
-
-        if "radius" in self.debug and self.counter == 10:
-            print(f"  [{layer}] [Radius] {radius.shape if radius is not None else None} {radius.mean() if radius is not None else None} [Theta] {theta.mean() if theta is not None else None} [f0] {f0.shape if f0 is not None else None} [Freqs] {freqs.shape} {freqs.mean():.2f} [ctx] {ctx}")
-        self.counter += 1
-        return freqs.unsqueeze(0)
-
-    @staticmethod
-    def apply_rotary(x, freqs):
-        x1 = x[..., :freqs.shape[-1]*2]
-        x2 = x[..., freqs.shape[-1]*2:]
-        orig_shape = x1.shape
-        if x1.ndim == 2:
-            x1 = x1.unsqueeze(0)
-        x1 = x1.float().reshape(*x1.shape[:-1], -1, 2).contiguous()
-        x1 = torch.view_as_complex(x1) * freqs
-        x1 = torch.view_as_real(x1).flatten(-2)
-        x1 = x1.view(orig_shape)
-        return torch.cat([x1.type_as(x), x2], dim=-1)
-
-class MultiheadA(nn.Module):
-
-    rbf = False
-    def __init__(self, dims: int, head: int, rotary_emb: bool = True, 
-                 zero_val: float = 1e-7, minz: float = 1e-8, maxz: float = 1e-6, debug: List[str] = [], optim_attn=False, use_pbias=False):
-        super(MultiheadA, self).__init__()
-
-        self.dims = dims
-        self.head = head
-        self.head_dim = dims // head
-        self.debug = debug
-        self.counter = 0
-        self.use_pbias = use_pbias
-
-        self.q = nn.Linear(dims, dims).to(device, dtype)
-        self.k = nn.Linear(dims, dims, bias=False).to(device, dtype)
-        self.v = nn.Linear(dims, dims).to(device, dtype)
-        self.o = nn.Linear(dims, dims).to(device, dtype)
-
-        self.pad_token = 0
-        self.rotary_emb = rotary_emb
-        self.minz = minz
-        self.maxz = maxz
-        self.zero_val = zero_val
-        self.optim_attn = optim_attn        
-        self.fzero = nn.Parameter(torch.tensor(zero_val, device=device, dtype=dtype), requires_grad=False)
-        
-        if rotary_emb:
-            self.rope = rotary(
-                dims=dims,
-                head=head,
-                debug=debug,
-                radii=False,
-                )
-        else:
-            self.rope = None
-
-    def cos_sim(self, q: Tensor, k: Tensor, v: Tensor, mask) -> Tensor:
-        q_norm = torch.nn.functional.normalize(q, dim=-1, eps=1e-12)
-        k_norm = torch.nn.functional.normalize(k, dim=-1, eps=1e-12)
-        qk_cosine = torch.matmul(q_norm, k_norm.transpose(-1, -2))
-        qk_cosine = qk_cosine + mask
-        weights = F.softmax(qk_cosine, dim=-1)
-        out = torch.matmul(weights, v)
-        return out
-
-    def rbf_scores(self, q, k, rbf_sigma=1.0, rbf_ratio=0.0):
-        scale = (self.dims // self.head) ** -0.25
-        dot_scores = torch.matmul(q, k.transpose(-1, -2)) * scale
-        if rbf_ratio <= 0.0:
-            return dot_scores
-        q_norm = q.pow(2).sum(dim=-1, keepdim=True)
-        k_norm = k.pow(2).sum(dim=-1, keepdim=True)
-        qk = torch.matmul(q, k.transpose(-1, -2))
-        dist_sq = q_norm + k_norm.transpose(-1, -2) - 2 * qk
-        rbf_scores = torch.exp(-dist_sq / (2 * rbf_sigma**2))
-        return (1 - rbf_ratio) * dot_scores + rbf_ratio * rbf_scores
-          
-    def forward(self, x: Tensor, xa = None, mask = None, en= None, layer = None, f=None) -> tuple:
-
-        x = x.to(device, dtype)
-        if xa is not None:
-            xa = xa.to(device, dtype)
-        scale = (self.dims // self.head) ** -0.25
-        
-        z = default(xa, x).to(device, dtype)
-        q = self.q(x)
-        k = self.k(z)
-        v = self.v(z)
-
-        if self.rotary_emb:   
-            q = q.view(*q.shape[:2], self.head, -1).permute(0, 2, 1, 3)
-            k = k.view(*k.shape[:2], self.head, -1).permute(0, 2, 1, 3)
-            v = v.view(*v.shape[:2], self.head, -1).permute(0, 2, 1, 3)
-            q2 = q.shape[2]
-            k2 = k.shape[2]
-
-            q = self.rope.apply_rotary(q, (self.rope(x=q2, en=en, f=f, layer=layer)))
-            k = self.rope.apply_rotary(k, (self.rope(x=k2, en=en, f=f, layer=layer)))
-        else:
-            q = q.view(*q.shape[:2], self.head, -1).permute(0, 2, 1, 3)
-            k = k.view(*k.shape[:2], self.head, -1).permute(0, 2, 1, 3)
-            v = v.view(*v.shape[:2], self.head, -1).permute(0, 2, 1, 3)
-        
-        qk = (q * scale) @ (k * scale).transpose(-1, -2)
-
-        if self.rbf:
-            qk = self.rbf_scores(q * scale, k * scale, rbf_sigma=1.0, rbf_ratio=0.3)
-        if self.use_pbias:
-            pbias = self.rope.pitch_bias(f0 = en.get("f0", None) if en is not None else None) 
-            if pbias is not None:
-                qk = qk + pbias[:,:,:q2,:q2]
-
-        token_ids = k[:, :, :, 0]
-        zscale = torch.ones_like(token_ids)
-        fzero = torch.clamp(F.softplus(self.fzero), self.minz, self.maxz)
-        zscale[token_ids.float() == self.pad_token] = fzero
-        
-        if mask is not None:
-            if mask.dim() == 4:
-                mask = mask[0, 0]
-            mask = mask[:q2, :k2] if xa is not None else mask[:q2, :q2]
-            qk = qk + mask * zscale.unsqueeze(-2).expand(qk.shape)
-
-        qk = qk * zscale.unsqueeze(-2)
-        w = F.softmax(qk, dim=-1).to(q.dtype)
-        wv = (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2)
-
-        if "multihead" in self.debug and self.counter % 100 == 0:
-            print(f"MHA: q={q.shape}, k={k.shape}, v={v.shape} - {qk.shape}, wv shape: {wv.shape}")
-        self.counter += 1        
-        return self.o(wv), qk
-
-    @staticmethod
-    def split(X: Tensor) -> (Tensor, Tensor):
-        half_dim = X.shape[-1] // 2
-        return X[..., :half_dim], X[..., half_dim:]
-
-class t_gate(nn.Module):
-    def __init__(self, dims, num_types=4, enabled=True):
-        super().__init__()
-        self.enabled = enabled
-        self.gate_projections = nn.ModuleList([
-            nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-            for _ in range(num_types)])
-        self.type_classifier = nn.Sequential(
-            Linear(dims, num_types),
-            nn.Softmax(dim=-1))
-    def forward(self, x):
-        if not self.enabled:
-            return None
-        type_probs = self.type_classifier(x)
-        gates = torch.stack([gate(x) for gate in self.gate_projections], dim=-1)
-        comb_gate = torch.sum(gates * type_probs.unsqueeze(2), dim=-1)
-        return comb_gate
-
-class m_gate(nn.Module):
-    def __init__(self, dims, mem_size=64, enabled=True):
-        super().__init__()
-        self.enabled = enabled
-        if enabled:
-            self.m_key = nn.Parameter(torch.randn(mem_size, dims))
-            self.m_val = nn.Parameter(torch.randn(mem_size, 1))
-            self.gate_proj = nn.Sequential(Linear(dims, dims//2), nn.SiLU(), Linear(dims//2, 1))
-            
-    def forward(self, x):
-        if not self.enabled:
-            return None
-        d_gate = torch.sigmoid(self.gate_proj(x))
-        attention = torch.matmul(x, self.m_key.transpose(0, 1))
-        attention = F.softmax(attention / math.sqrt(x.shape[-1]), dim=-1)
-        m_gate = torch.matmul(attention, self.m_val)
-        m_gate = torch.sigmoid(m_gate)
-        return 0.5 * (d_gate + m_gate)
-
-class c_gate(nn.Module):
-    def __init__(self, dims, enabled=True):
-        super().__init__()
-        self.enabled = enabled
-        if enabled:
-            self.s_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-            self.w_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-            self.p_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-            self.e_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-            self.ph_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-            self.integ = Linear(dims*5, dims)
-        
-    def forward(self, x, features):
-        if not self.enabled:
-            return None
-        s_feat = features.get("spectrogram", x)
-        w_feat = features.get("waveform", x)
-        p_feat = features.get("pitch", x)
-        e_feat = features.get("envelope", x)
-        ph_feat = features.get("phase", x)
-        s = self.s_gate(x) * s_feat
-        w = self.w_gate(x) * w_feat
-        p = self.p_gate(x) * p_feat
-        e = self.e_gate(x) * e_feat
-        ph = self.ph_gate(x) * ph_feat
-        comb = torch.cat([s, w, p, e, ph], dim=-1)
-        return self.integ(comb)
-
-class mlp_gate(nn.Module):
-    def __init__(self, dims, head, enabled=True, one_shot=True):
-        super().__init__()
-        self.enabled = enabled
-        if enabled:
-            self.gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-
-    def forward(self, x, xa=None, f=None):
-        if not self.enabled:
-            return None
-        return self.gate(x)
-
-class Residual(nn.Module):
-    _seen = set()  
-    def __init__(self, ctx, dims, head, act, debug: List[str] = [], 
-                 tgate=True, mgate=False, cgate=False, mem_size=512, features=None, one_shot=False):
-        super().__init__()
-        
-        self.dims = dims
-        self.head = head
-        self.ctx = ctx
-        self.head_dim = dims // head
-        self.features = features
-        self.debug = debug
-        self.counter = 0
-        self.dropout = 0.01
-        self.one_shot = one_shot
-
-        self.blend = nn.Parameter(torch.tensor(0.5)) 
-        act_fn = get_activation(act)
-        self.attn = MultiheadA(dims, head, rotary_emb=True, debug=debug)
-        self.curiosity = curiosity(dims, head)
-
-        if not any([tgate, mgate, cgate]):
-            self.mlp_gate = nn.Sequential(Linear(dims, 1), nn.Sigmoid())
-        else:
-            self.mlp_gate = None
-        
-        mlp = dims * 4
-        self.mlp = nn.Sequential(Linear(dims, mlp), act_fn, Linear(mlp, dims))
-        
-        self.t_gate = t_gate(dims=dims, num_types=4*2, enabled=tgate)
-        self.m_gate = m_gate(dims=dims, mem_size=mem_size, enabled=mgate)
-        self.c_gate = c_gate(dims=dims, enabled=cgate)
-        self.mlp_gate = mlp_gate(dims=dims, head=head, enabled=not any([tgate, mgate, cgate]), one_shot=True)
-        
-        self.lna = RMSNorm(dims)
-        self.lnb = RMSNorm(dims)
-        self.lnc = RMSNorm(dims)
-
-    def forward(self, x, xa=None, mask=None, en=None, layer=None, f=None) -> Tensor:
- 
-        b = torch.sigmoid(self.blend)
-        ax = x + self.attn(self.lna(x), xa=xa, mask=mask, en=en, layer=layer, f=f)[0]
-        bx = b * ax + (1 - b) * x
-        cx = self.lnb(bx)
-        dx = self.mlp(cx)
-        ex = self.t_gate(cx) if not None else self.default(self.m_gate(cx), self.mlp_gate(cx))
-        fx = x + ex + dx
-        gx = self.lnc(fx)
-        return gx
-            
-class OneShot(nn.Module):
-    def __init__(self, dims: int, head: int, scale: float = 0.3):
-        super().__init__()
-        self.head  = head
-        self.hdim  = dims // head
-        self.scale = scale                      
-        self.q_proj = Linear(dims, dims)
-        self.k_proj = Linear(dims, dims)
-
-    def forward(self, x: Tensor, guide: Tensor, f=None) -> Tensor | None:
-        B, Q, _ = x.shape
-        K       = guide.size(1)
-        q = self.q_proj(x ).view(B, Q, self.head, self.hdim).transpose(1,2)
-        k = self.k_proj(guide).view(B, K, self.head, self.hdim).transpose(1,2)
-        bias = (q @ k.transpose(-1, -2)) * self.scale / math.sqrt(self.hdim)
-        return bias
-
-class curiosity(nn.Module):
-    def __init__(self, d, h, bias=True):
-        super().__init__()
-        self.h  = h
-        self.dh = d // h
-        self.qkv = nn.Linear(d, d * 3, bias=bias)
-        self.qkv_aux = nn.Linear(d, d * 3, bias=bias)
-        self.o  = nn.Linear(d, d, bias=bias)
-        self.g  = nn.Parameter(torch.zeros(h))
-
-    def split(self, x):
-        b, t, _ = x.shape
-        return x.view(b, t, self.h, self.dh).transpose(1, 2)
-
-    def merge(self, x):
-        b, h, t, dh = x.shape
-        return x.transpose(1, 2).contiguous().view(b, t, h * dh)
-
-    def forward(self, x, xa, mask=None):
-        q, k, v   = self.qkv(x).chunk(3, -1)
-        qa, ka, va = self.qkv_aux(xa).chunk(3, -1)
-        q, k, v   = map(self.split, (q, k, v))
-        qa, ka, va = map(self.split, (qa, ka, va))
-        dots      = (q @ k.transpose(-2, -1)) / self.dh**0.5
-        dots_aux  = (q @ ka.transpose(-2, -1)) / self.dh**0.5
-        if mask is not None: dots = dots.masked_fill(mask, -9e15)
-        p   = dots.softmax(-1)
-        pa  = dots_aux.softmax(-1)
-        h_main = p  @ v
-        h_aux  = pa @ va
-        g = torch.sigmoid(self.g).view(1, -1, 1, 1)
-        out = self.merge(h_main * (1 - g) + h_aux * g)
-        return self.o(out)
-
-class PositionalEncoding(nn.Module):
-    def __init__(self, dims, ctx):
-        super(PositionalEncoding, self).__init__()
-        self.dims = dims
-        self.ctx = ctx
-        self.pe = self.get_positional_encoding(max_ctx=ctx)
-
-    def get_positional_encoding(self, max_ctx):
-        pe = torch.zeros(max_ctx, self.dims)
-        position = torch.arange(0, max_ctx, dtype=torch.float32).unsqueeze(1)
-        div_term = torch.exp(
-            torch.arange(0, self.dims, 2, dtype=torch.float32)
-            * (-math.log(10000.0) / self.dims)
-        )
-        pe[:, 0::2] = torch.sin(position * div_term)
-        pe[:, 1::2] = torch.cos(position * div_term)
-        pe = pe.unsqueeze(0)
-        return pe.to(device)
-
-    def forward(self, x):
-        ctx = x.size(1)
-        pe = self.pe[:, :ctx, :]
-        x = x * math.sqrt(self.dims)
-        x = x + pe
-        return x
-
-class FEncoder(nn.Module):
-    def __init__(self, mels, dims, head, layer, kernel_size, act, stride=1, use_rope=False, spec_shape=None, debug=[]):
-        super().__init__()
-        
-        self.head = head
-        self.head_dim = dims // head  
-        self.dropout = 0.01 
-        self.use_rope = use_rope
-        self.dims = dims
-        self.debug = debug
-        act_fn = get_activation(act)
-        self.attend_pitch = False
-
-        if self.attend_pitch:
-            self.q, self.k, self.v, self.o, self.scale = qkv_init(dims, head)
-            self.mlp = nn.Sequential(
-                nn.Linear(dims, dims),
-                nn.ReLU(),
-                nn.Linear(dims, dims),
-            )
-        else:
-            self.q, self.k, self.v, self.o, self.scale = None, None, None, None, None
-            self.mlp = None
-
-        self.encoder = nn.Sequential(
-            Conv1d(mels, dims, kernel_size=3, stride=1, padding=1), act_fn,
-            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1), act_fn,
-            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1, groups=dims), act_fn)
-
-        if use_rope:
-            if spec_shape is not None:
-                self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape)
-        else:
-            self.rope = None
-            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
-        self.norm = RMSNorm(dims)
-
-    def apply_rope_to_features(self, x, en=None, f=None, layer="audio"):
-        batch, ctx, dims = x.shape
-        x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
-        freqs = self.rope(ctx, en=en, f=f, layer=layer)
-        x = self.rope.apply_rotary(x, freqs)
-        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
-
-        return x
-
-    def forward(self, x: Tensor, en=None, f=None, layer = None):
-        x = self.encoder(x).permute(0, 2, 1)
-        if self.use_rope:
-            x = self.apply_rope_to_features(x, en=en, f=f, layer=layer)
-        else:
-            x = x + self.positional(x.shape[1], x.shape[-1], 10000).to(device, dtype)
-
-        if self.mlp is not None:
-            x = self.mlp(x)
-
-        if self.attend_pitch:
-            xa = en["input_ids"]
-            if xa is not None:
-                q, k, v = create_qkv(self.q, self.k, self.v, x=xa, xa=x, head=self.head)
-                out, _ = calculate_attention(q, k, v, mask=None, temperature=1.0, is_causal=True)
-                out = self.o(out)
-                x = x + out
-
-        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
-        x = self.norm(x)
-        return x
-
-class WEncoder(nn.Module):
-    def __init__(self, input_dims, dims, head, layer, kernel_size, act, use_rope=False, debug=[], spec_shape=None):
-        super().__init__()
-        
-        self.head = head
-        self.head_dim = dims // head
-        self.dropout = 0.01
-        self.use_rope = use_rope
-        self.dims = dims
-        self.debug = debug
-        act_fn = get_activation(act)
-        self.target_length = None
-        self.encoder = nn.Sequential(
-            Conv1d(input_dims, dims//4, kernel_size=15, stride=4, padding=7), act_fn,
-            Conv1d(dims//4, dims//2, kernel_size=7, stride=2, padding=3), act_fn,
-            Conv1d(dims//2, dims, kernel_size=5, stride=2, padding=2), act_fn)
-            
-        if use_rope:
-            if spec_shape is not None:
-                self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape)
-        else:
-            self.rope = None
-            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
-        self.norm = RMSNorm(dims)
-
-    def apply_rope_to_features(self, x, en=None, f=None, layer="audio"):
-        batch, ctx, dims = x.shape
-        x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
-        freqs = self.rope(ctx, en=en, f=f, layer=layer)
-        x = self.rope.apply_rotary(x, freqs)
-        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
-        return x
-        
-    def forward(self, x: Tensor, en= None, f=None, layer = None):
-        x = self.encoder(x).permute(0, 2, 1)
-        if self.target_length and x.shape[1] != self.target_length:
-            x = F.adaptive_avg_pool1d(x.transpose(1, 2), self.target_length).transpose(1, 2)
-        if self.use_rope:
-            x = self.apply_rope_to_features(x, en=en, f=f, layer=layer)
-        else:
-            x = x + self.positional(x.shape[1], x.shape[-1], 10000).to(device, dtype)
-        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
-
-        x = self.ln(x)
-        print(f"X: {x.shape} {f}") if "encoder" in self.debug else None
-        return self.norm(x)
-
-class PEncoder(nn.Module):
-    def __init__(self, input_dims, dims, head, layer, kernel_size, act, use_rope=True, debug=[], one_shot=False, spec_shape=None):
-        super().__init__()
-        
-        self.head = head
-        self.head_dim = dims // head
-        self.dims = dims
-        self.dropout = 0.01
-        self.use_rope = use_rope
-        self.debug = debug
-        act_fn = get_activation(act)
-        
-        self.encoder = nn.Sequential(
-            Conv1d(input_dims, dims, kernel_size=7, stride=1, padding=3), act_fn,
-            Conv1d(dims, dims, kernel_size=5, stride=1, padding=2), act_fn,
-            Conv1d(dims, dims, kernel_size=3, stride=1, padding=1, groups=dims), act_fn)
-
-        if use_rope:
-                self.rope = rotary(dims=dims, head=head, radii=False, debug=[], use_pbias=False, axial=False, spec_shape=spec_shape)
-        else:
-            self.rope = None
-            self.positional = lambda length, dims, max_tscale: sinusoids(length, dims, max_tscale)
-        
-        self.norm = RMSNorm(dims)
-        
-    def rope_to_feature(self, x, en=None, f="pitch", layer="PEncoder"):
-        batch, ctx, dims = x.shape
-        x = x.view(batch, ctx, self.head, self.head_dim).permute(0, 2, 1, 3)
-        freqs = self.rope(ctx, en=en, f=f, layer=layer)
-        x = self.rope.apply_rotary(x, freqs)
-        x = x.permute(0, 2, 1, 3).contiguous().view(batch, ctx, dims)
-        return x
-        
-    def forward(self, x: Tensor, en= None, f="pitch", layer="PEncoder"):
-
-        if x.dim() == 2:
-            x = x.unsqueeze(0)
-        
-        x = self.encoder(x).permute(0, 2, 1)
-        if self.use_rope:
-            x = self.rope_to_feature(x, en=en, f=f, layer=layer)
-        else:
-            x = x + self.positional(x.shape[1], x.shape[-1], 10000).to(device, dtype)
-        x = nn.functional.dropout(x, p=self.dropout, training=self.training)
-        x = self.norm(x)
-        print(f"X: {x.shape} {f}") if "PEncoder" in self.debug else None
-        return x
-
-class theBridge(nn.Module):
-    def __init__(self, vocab: int, mels: int, ctx: int, dims: int, head: int, layer: int, 
-                debug: List[str], features: List[str], act: str = "gelu"): 
-        super(theBridge, self).__init__()
-    
-        tgate = True
-        mgate = False
-        cgate = False
-
-        self.debug = debug
-        self.counter = 0
-        self.dropout = 0.01 
-        self.features = features
-        self.do_blend = "no_blend" not in self.debug
-        self.sequential = "sequential" in self.debug
-        self.layer = layer
-
-        self.token = nn.Embedding(vocab, dims, device=device, dtype=dtype)
-        self.positional = nn.Parameter(torch.empty(ctx, dims, device=device, dtype=dtype), requires_grad=True)
-        self.blend = nn.Parameter(torch.tensor(0.5, device=device, dtype=dtype), requires_grad=True)
-        self.norm = RMSNorm(dims)
-        self.sinusoid_pos = lambda length, dims, max_tscale: sinusoids(length, dims, 10000)
-        self.rotary = rotary(dims=dims,  head=head, debug=debug, radii=False)
-
-        with torch.no_grad():
-            self.token.weight[0].zero_()
-        
-        act_fn = get_activation(act)
-        if features == ["spectrogram", "waveform", "pitch"]:
-            cgate=True
-        else:
-            cgate = False
-    
-        self.blockA = nn.ModuleDict()
-        self.blockA["waveform"] = nn.ModuleList(
-            [WEncoder(input_dims=1, dims=dims, head=head, layer=layer, kernel_size=11, act=act_fn)] + 
-            [Residual(ctx=ctx, dims=dims, head=head, act=act_fn, tgate=tgate, mgate=mgate, cgate=cgate, debug=debug, features=features) 
-            for _ in range(layer)] if "waveform" in features else None) 
-
-        for feature_type in ["spectrogram", "aperiodic", "harmonic"]:
-            if feature_type in features:
-                self.blockA[feature_type] = nn.ModuleList(
-                    [FEncoder(mels=mels, dims=dims, head=head, layer=layer, kernel_size=3, act=act_fn)] + 
-                    [Residual(ctx=ctx, dims=dims, head=head, act=act_fn, tgate=tgate, mgate=mgate, cgate=cgate, debug=debug, features=features) for _ in range(layer)] if feature_type in features else None)
-            else:
-                self.blockA[feature_type] = None
-
-        for feature_type in ["pitch", "phase"]:
-            if feature_type in features:
-                self.blockA[feature_type] = nn.ModuleList(
-                    [PEncoder(input_dims=1, dims=dims, head=head, layer=layer, kernel_size=9, act=act_fn)] + 
-                    [Residual(ctx=ctx, dims=dims, head=head, act=act_fn, tgate=tgate, mgate=mgate, cgate=cgate, debug=debug, features=features) for _ in range(layer)] if feature_type in features else None)
-            else:
-                self.blockA[feature_type] = None
-
-        self.blockB = nn.ModuleList([
-            Residual(ctx=ctx, dims=dims, head=head, act=act_fn, tgate=tgate, mgate=mgate, cgate=cgate, debug=debug, features=features)
-            for _ in range(layer)])
-
-        self.modal = nn.ModuleList([
-            Residual(ctx=ctx, dims=dims, head=head, act=act_fn, tgate=tgate, mgate=mgate, cgate=cgate, debug=debug, features=features)
-            for _ in range(layer)]) 
-
-        mask = torch.tril(torch.ones(ctx, ctx), diagonal=0) 
-        self.register_buffer("mask", mask, persistent=False)
-
-        self.norm = RMSNorm(dims)
-
-    def forward(self, x, xa, en, f, sequential=False) -> Tensor:
-        mask = self.mask[:x.shape[1], :x.shape[1]] 
-        x = self.token(x.long()) + self.positional[:x.shape[1]]
-
-        out = {}
-        out["input_ids"] = x
-        out.update(en)
-
-        for b in chain(self.blockA[f] or []):
-            xa = b(x=xa, en=out, f=f, layer="en")
-
-        for b in chain(self.blockB or []):
-            x = b(x=x, xa=None, mask=mask, en=out, f=f, layer="dec")
-            y = b(x, xa=xa, mask=None, en=out, f=f, layer="cross")
-            if sequential:
-                x = y
-            else:
-                a = torch.sigmoid(self.blend)
-                x = a * y + (1 - a) * x 
-        for b in self.modal:
-            xc = b(x=torch.cat([x, xa], dim=1), xa=None, mask=None, en=out, f=f, layer="modal")    
-            xm = b(x=xc[:, :x.shape[1]], xa=xc[:, x.shape[1]:], mask=None, en=out, f=f, layer="modal")
-            if sequential:
-                x = xm
-            else:
-                a = torch.sigmoid(self.blend)
-                x = a * x + (1 - a) * xm
-
-        if self.counter < 1 and "encoder" in self.debug:      
-                shapes = {k: v.shape for k, v in en.items()}
-                print(f"Step {self.counter}: mode: {list(en.keys()) }: shapes: {shapes}")
-        self.counter += 1
-
-        x = self.norm(x)
-        x = x @ torch.transpose(self.token.weight.to(dtype), 0, 1).float()
-
-        return x
-   
-class Echo(nn.Module):
-    def __init__(self, param: Dimensions):
-        super().__init__()
-        self.param = param
-        
-        self.processor = theBridge(
-            vocab=param.vocab,
-            mels=param.mels,
-            ctx=param.ctx,
-            dims=param.dims,
-            head=param.head,
-            layer=param.layer,
-            features=param.features,
-            act=param.act,
-            debug=param.debug,
-            )       
-        
-    def forward(self,
-        labels=None,
-        input_ids=None,
-        waveform: Optional[torch.Tensor]=None,
-        spectrogram: Optional[torch.Tensor]=None,
-        pitch: Optional[torch.Tensor]=None,
-        f0: Optional[torch.Tensor]=None,
-        f0t: Optional[torch.Tensor]=None,
-        harmonic: Optional[torch.Tensor]=None,
-        aperiodic: Optional[torch.Tensor]=None,
-        phase: Optional[torch.Tensor]=None,
-        ) -> Dict[str, Optional[torch.Tensor]]:
-
-        en= TensorDict(batch_size=[1], device=self.device, dtype=self.dtype)
-
-        en= {}
-        if f0 is not None:
-            en["f0"] = f0
-        if f0t is not None:
-            en["f0t"] = f0t
-        if harmonic is not None:
-            en["harmonic"] = harmonic
-        if aperiodic is not None:
-            en["aperiodic"] = aperiodic
-        if phase is not None:
-            en["phase"] = phase
-        if pitch is not None:
-            en["pitch"] = pitch
-        if waveform is not None:
-            en["waveform"] = waveform
-        if spectrogram is not None:
-            en["spectrogram"] = spectrogram
-
-        x = input_ids
-        for f, xa in en.items():
-
-            logits = self.processor(x, xa, en, f)
-
-        loss = None
-        if labels is not None:
-            loss = F.cross_entropy(
-                logits.view(-1, logits.shape[-1]), labels.view(-1), ignore_index=0)
-        
-        return {"logits": logits, "loss": loss} 
-
-    @property
-    def device(self):
-        return next(self.parameters()).device
-    @property
-    def dtype(self):
-        return next(self.parameters()).dtype
-
-    def _init_weights(self, module):
-        std = 0.02
-        self.init_counts = {
-            "Linear": 0, "Conv1d": 0, "LayerNorm": 0, "RMSNorm": 0,
-            "Conv2d": 0, "theBridge": 0, "Echo": 0, 
-            "Residual": 0, "MultiheadA": 0, 
-            "MultiheadC": 0, "MultiheadD": 0, "FEncoder": 0,
-            "WEncoder": 0, "PEncoder": 0}
-
-        for name, module in self.named_modules():
-            if isinstance(module, RMSNorm):
-                nn.init.ones_(module.weight)
-                self.init_counts["RMSNorm"] += 1
-            elif isinstance(module, nn.Linear):
-                if module.weight is not None:
-                    nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-                self.init_counts["Linear"] += 1
-            elif isinstance(module, Conv1d):
-                nn.init.normal_(module.weight, mean=0.0, std=std)
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-                self.init_counts["Conv1d"] += 1
-            elif isinstance(module, Conv2d):
-                nn.init.normal_(module.weight, mean=0.0, std=std)
-                if module.bias is not None:
-                    nn.init.zeros_(module.bias)
-                self.init_counts["Conv2d"] += 1
-            elif isinstance(module, MultiheadA):
-                self.init_counts["MultiheadA"] += 1
-            elif isinstance(module, Residual):
-                self.init_counts["Residual"] += 1
-            elif isinstance(module, PEncoder):
-                self.init_counts["PEncoder"] += 1
-            elif isinstance(module, FEncoder):
-                self.init_counts["FEncoder"] += 1
-            elif isinstance(module, WEncoder):
-                self.init_counts["WEncoder"] += 1
-            elif isinstance(module, theBridge):
-                self.init_counts["theBridge"] += 1
-            elif isinstance(module, Echo):
-                self.init_counts["Echo"] += 1
-
-    def init_weights(self):
-        print("Initializing model weights...")
-        self.apply(self._init_weights)
-        print("Initialization summary:")
-        for module_type, count in self.init_counts.items():
-            if count > 0:
-                print(f"{module_type}: {count}")
-
-    def generate(self, input_ids=None, spectrogram=None, waveform=None, pitch=None, f0=None, 
-        envelope=None, phase=None, tokenizer=None, max_length=128, min_length=1, device=None, **kwargs):
-        if device is None:
-            device = self.device
-        pad_token_id = getattr(tokenizer, "pad_token_id", 0)
-        bos_token_id = getattr(tokenizer, "bos_token_id", 1)
-        eos_token_id = getattr(tokenizer, "eos_token_id", 2)
-        batch_size = 1
-        for x in [spectrogram, waveform, pitch, f0, envelope, phase]:
-            if x is not None:
-                batch_size = x.shape[0]
-                break
-        ids = torch.full((batch_size, 1), bos_token_id, dtype=torch.long, device=device)
-        feature = {}
-        if spectrogram is not None:
-            feature["spectrogram"] = spectrogram
-        if waveform is not None:
-            feature["waveform"] = waveform
-        if pitch is not None:
-            feature["pitch"] = pitch
-        if envelope is not None:
-            feature["envelope"] = envelope
-        if phase is not None:
-            feature["phase"] = phase
-        if f0 is not None:
-            feature["f0"] = f0
-        
-        for i in range(max_length - 1):
-            with torch.no_grad():
-                feature["input_ids"] = ids
-                logits = self.SpeechTransformer(feature)
-            next_token_logits = logits[:, -1, :]
-            if i < min_length:
-                next_token_logits[:, eos_token_id] = 0
-            next_tokens = torch.argmax(next_token_logits, dim=-1, keepdim=True)
-            ids = torch.cat([ids, next_tokens], dim=1)
-            if (next_tokens == eos_token_id).all() and i >= min_length:
-                break
-        return ids
-
-    @property
-    def config(self):
-        class Config:
-            pad_token_id = getattr(self.param, "pad_token_id", 0)
-            bos_token_id = getattr(self.param, "bos_token_id", 1)
-            eos_token_id = getattr(self.param, "eos_token_id", 2)
-            def to_json_string(self):
-                import json
-                return json.dumps({
-                    "pad_token_id": self.pad_token_id,
-                    "bos_token_id": self.bos_token_id,
-                    "eos_token_id": self.eos_token_id,
-                })
-        return Config()
-
-def main():
-    token = ""
-    log_dir = os.path.join('./output/logs', datetime.now().strftime('%m-%d_%H_%M_%S'))
-    os.makedirs(log_dir, exist_ok=True)
-    tokenizer = setup_tokenizer("./")
-
-    sanity_check = False
-    streaming = False
-    load_saved = False
-    save_dataset = False
-    cache_dir = None
-    extract_args = None    
-
-    extract_args = {
-        "waveform": False,
-        "spec": True,
-        "f0": False,
-        "f0t": False,
-        "pitch": True,
-        "harmonics": False,
-        "aperiodics": False,
-        "phase_mod": False,
-        "crepe": False,        
-        "sample_rate": 16000,
-        "hop_length": 256,
-        "mode": "mean",
-        "debug": False,
-    }
-
-    param = Dimensions(
-        vocab=40000,
-        mels=128,
-        ctx=2048,
-        dims=512,
-        head=4,
-        layer=4,
-        act="swish",
-        debug={"encoder"},
-        features = ["spectrogram", "pitch"],
-        )
-
-    train_dataset, test_dataset = prepare_datasets(tokenizer, token, sanity_check=sanity_check, sample_rate=16000, streaming=streaming,
-        load_saved=load_saved, save_dataset=save_dataset, cache_dir=cache_dir, extract_args=extract_args, max_ctx=param.ctx)
-
-    model = Echo(param).to('cuda')
-    print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
-    print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")
-    
-    from functools import partial
-    metrics_fn = partial(compute_metrics, 
-    print_pred=True, 
-    num_samples=1, 
-    tokenizer=tokenizer, model=model)
-
-    if sanity_check:
-        training_args = Seq2SeqTrainingArguments(
-            output_dir=log_dir,
-            per_device_train_batch_size=1,
-            per_device_eval_batch_size=1,
-            max_steps=10,
-            eval_steps=5,
-            save_steps=0,
-            warmup_steps=0,
-            logging_steps=1,
-            logging_dir=log_dir,
-            eval_strategy="steps",
-            save_strategy="no",
-            logging_strategy="no",
-            report_to=["tensorboard"],
-            push_to_hub=False,
-            save_total_limit=1,
-            label_names=["labels"],
-            save_safetensors=False,
-            eval_on_start=True,
-            batch_eval_metrics=False,
-            disable_tqdm=False,
-            include_tokens_per_second=True,
-            include_num_input_tokens_seen=True,
-            learning_rate=1e-7,
-            weight_decay=0.01,
-        )
-    else:
-        training_args = Seq2SeqTrainingArguments(
-            output_dir=log_dir,
-            per_device_train_batch_size=1,
-            per_device_eval_batch_size=1,
-            max_steps=1000,
-            eval_steps=100,
-            save_steps=1000,
-            warmup_steps=100,
-            logging_steps=10,
-            logging_dir=log_dir,
-            logging_strategy="steps",
-            eval_strategy="steps",
-            save_strategy="no",
-            report_to=["tensorboard"],
-            push_to_hub=False,
-            save_total_limit=1,
-            label_names=["labels"],
-            save_safetensors=False,
-            eval_on_start=True,
-            batch_eval_metrics=False,
-            disable_tqdm=False,
-            include_tokens_per_second=True,
-            include_num_input_tokens_seen=True,
-            learning_rate=0.00025,
-            weight_decay=0.025,
-        )
-
-    optimizer = torch.optim.AdamW(model.parameters(), lr=training_args.learning_rate, eps=1e-8, weight_decay=training_args.weight_decay, betas=(0.9, 0.999), 
-    amsgrad=False, foreach=False, fused=False, capturable=False, differentiable=False, maximize=False)
-    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=training_args.max_steps, eta_min=1e-9, last_epoch=-1)
-
-    trainer = Seq2SeqTrainer(
-        args=training_args,
-        model=model,
-        train_dataset=train_dataset,
-        eval_dataset=test_dataset,
-        data_collator=DataCollator(tokenizer=tokenizer),
-        preprocess_logits_for_metrics=preprocess_logits_for_metrics,
-        compute_metrics=metrics_fn,
-        optimizers=(optimizer, scheduler)
-    )
-
-    model.init_weights()
-    trainer.train()
-if __name__ == "__main__":
-
-    main()
-