codec.py

import torch
from torch.nn.attention.flex_attention import flex_attention as flex,create_block_mask

class Conv(torch.nn.Module):                                                 # conv  padding: left = (k-1)*d+1-s, right = (((l - k + left) / s +1 ).ceil() -1) * s + k - left -l
    def __init__(self,i,o,k,s=1,d=1,g=1,snake=True,T=False):                 # s=1 -> right=0, s!=1 -> k=2*s,d=1 -> left=s -> right = ceil(l / s) * s - l
        super().__init__()                                                   # convT padding: right = k - s, left = 0, x -> x[..., left : x.shape[-1] - right] 
        self.alpha=torch.nn.Parameter(torch.ones(1,i,1,1)) if snake else None
        self.conv=torch.nn.ConvTranspose2d(i,o,(k,1),(s,1)) if T else torch.nn.Conv2d(i,o,(k,1),(s,1),dilation=d,groups=g)
        self.T,self.left,self.right,self.leftT,self.rightT=T,(k-1)*d+1-s,0 if s==1 else s,0,k-s

    def forward(self,x):
        if self.alpha!=None: x=x+(x*self.alpha).sin().square()*(1.0/(self.alpha+1e-9))
        if self.T: x=self.conv(x); return x[:,:,:x.shape[-2]-self.rightT,:]
        padding=(self.left,0 if self.right==0 or x.shape[-2]%self.right==0 else self.right-x.shape[-2]%self.right)
        return self.conv(torch.nn.functional.pad(x,(0,0)+padding,"constant",0))

class ResidualUnit(torch.nn.Module):
    def __init__(self,C,d):
        super().__init__()
        self.conv1,self.conv2=Conv(C,C,7,1,d,snake=True),Conv(C,C,1,1,1,snake=True)

    def forward(self,x):
        return x+self.conv2(self.conv1(x))

class ConvNeXtBlock(torch.nn.Module):
    def __init__(self,C):
        super().__init__()
        self.gamma,self.dwconv,self.norm=torch.nn.Parameter(1e-6*torch.ones((C))),Conv(C,C,7,g=C,snake=False),torch.nn.LayerNorm((C,),eps=1e-6)
        self.up,self.down=[torch.nn.Linear(n*C,m*C) for n,m in [[1,4],[4,1]]]
    
    def forward(self,x):
        x_n=self.norm(self.dwconv(x).permute(0,3,2,1))
        return x+(self.down(torch.nn.functional.gelu(self.up(x_n)))*self.gamma).permute(0,3,2,1)

def apply_rope(x,cache):
    x_type,x=x.dtype,x.to(torch.float32).unflatten(-1,[x.shape[-1]//2,2])
    return torch.stack([x[...,0]*cache[...,0]-x[...,1]*cache[...,1],x[...,1]*cache[...,0]+x[...,0]*cache[...,1]],-1).flatten(-2).to(x_type)

class TransformerLayer(torch.nn.Module):
    def __init__(self,C=1024):
        super().__init__()
        self.qkv,self.o,self.up,self.gate,self.down=[torch.nn.Linear(i,o,bias=False) for i,o in [[C,3*C],[C,C],[C,3*C],[C,3*C],[3*C,C]]]
        self.Anorm,self.Fnorm,self.Agamma,self.Fgamma=*[torch.nn.RMSNorm(i,eps=1e-5) for i in (C,C)],*[torch.nn.Parameter(1e-2*torch.ones(C)) for _ in range(2)]

    def forward(self,x,cache,mask):
        q,k,v=[i.unflatten(-1,[-1,cache.shape[-2]*2]).transpose(1,2) for i in self.qkv(self.Anorm(x)).split([x.shape[-1],x.shape[-1],x.shape[-1]],dim=-1)]
        q,k=[apply_rope(i,cache) for i in (q,k)]
        x=x+self.o(flex(q,k,v,block_mask=mask).transpose(1,2).contiguous().flatten(-2))*self.Agamma
        x_n=self.Fnorm(x); x=x+self.down(torch.nn.functional.silu(self.gate(x_n))*self.up(x_n))*self.Fgamma
        return x

class Transformer(torch.nn.Module):
    def __init__(self,L=4,C=1024):
        super().__init__()
        self.layers,self.norm=torch.nn.ModuleList([TransformerLayer(C) for _ in range(L)]),torch.nn.RMSNorm(C,eps=1e-5)

    def forward(self,x,cache,window_size):
        block_mask=create_block_mask(lambda b,h,q_idx,kv_idx: (q_idx>=kv_idx)&(q_idx<kv_idx+window_size),B=None,H=None,Q_LEN=x.shape[1],KV_LEN=x.shape[1])
        for layer in self.layers: x=layer(x,cache[:x.shape[1]][None,None],block_mask)                                                # max_len,c/2,2 -> T,c/2,2 -> 1,1,T,c/2,2 X B,A,T,c/2,2
        return self.norm(x)

class Quantizer(torch.nn.Module):
    def __init__(self,C=1024,bookC=8,n=10,layers=8,elayers=4):
        super().__init__()
        self.etransformer,self.eout=Transformer(elayers,C),Conv(C,C,3)
        self.books=torch.nn.ModuleList([m for V in [4096]+[1024]*(n-1) for m in [torch.nn.Linear(C,bookC),torch.nn.Linear(bookC,C),torch.nn.Embedding(V,bookC)]])
        self.downsample=torch.nn.Sequential(Conv(C,C,2,2,snake=False),ConvNeXtBlock(C),Conv(C,C,2,2,snake=False),ConvNeXtBlock(C))
        self.upsample=torch.nn.Sequential(Conv(C,C,2,2,snake=False,T=True),ConvNeXtBlock(C),Conv(C,C,2,2,snake=False,T=True),ConvNeXtBlock(C))
        self.pre_module,self.post_module=[Transformer(L=layers,C=C) for _ in range(2)]

    def encode(self,z,cache):
        z=self.eout(self.etransformer(z.squeeze(-1).transpose(1,2),cache,512).transpose(1,2).unsqueeze(-1))                          # B,C,T,1 -> B,T,C -> B,C,T,1
        z=self.pre_module(self.downsample(z).squeeze(-1).transpose(1,2),cache,128)                                                   # B,C,T,1 -> B,T,C
        codes=torch.zeros([z.shape[0],z.shape[1],10],dtype=torch.int32,device=z.device)
        for i in range(10):                                                                                                          # B,T,C   -> B,T,8 -> B*T,8
            e_norm,z_norm=[torch.nn.functional.normalize(w) for w in [self.books[i*3+2].weight,self.books[i*3](z).view(-1,self.books[2].embedding_dim)]]
            d=z_norm.pow(2).sum(dim=1,keepdim=True)-2*z_norm@(e_norm.T)+(e_norm).pow(2).sum(dim=1,keepdim=True).T                    # (z-e)²=z²+e²-2ze
            codes[...,i]=d.argmin(dim=1).view(z.shape[0],-1)                                                                         # B*T     -> B,T
            zq=self.books[i*3+1](self.books[i*3+2](codes[...,i])); z=z-zq                                                            # B,T     -> B,T,8 -> B,T,C
        return codes.transpose(1,2)                                                                                                  # B,T,N   -> B,N,T

    def decode(self,codes,cache):                                                                                                    # B,N,T
        z=self.books[1](self.books[2](codes[:,0,:]))                                                                                 # B,T     -> B,T,8 -> B,T,C
        for i in range(1,10): z=z+self.books[i*3+1](self.books[i*3+2](codes[:,i,:]))
        return self.upsample(self.post_module(z,cache,128).transpose(1,2).unsqueeze(-1))                                             # B,T,C   -> B,C,T,1

class Codec(torch.nn.Module):
    def __init__(self,max_len=640*4,strides=[8,8,4,2],dilations=[1,3,9]):                                                            # encoder 44100/2/4/8/8=86Hz, max_len=2560, 2560/86≈30 s
        super().__init__()                                                                                                           # rope cache: 2560*32*2*4/1e6=0.66 MB
        eblock=[Conv(1,64,7,snake=False)]+[m for s,C in zip(strides[::-1],[64,128,256,512]) for m in [ResidualUnit(C,d) for d in dilations]+[Conv(C,C*2,2*s,s)]]
        dblock=[Conv(1024,1536,7,snake=False)]+[m for s,C in zip(strides,[1536,768,384,192]) for m in [Conv(C,C//2,2*s,s,T=True)]+[ResidualUnit(C//2,d) for d in dilations]]+[Conv(96,1,7),torch.nn.Tanh()]
        self.encoder,self.quantizer,self.decoder=torch.nn.Sequential(*eblock),Quantizer(),torch.nn.Sequential(*dblock)               # quantizer 44100/2/4/8/8/2/2=21.5Hz
        self.cache=torch.nn.Buffer(torch.stack([f(torch.tensor([[j*10000**(-i/(32)) for i in range(32)] for j in range(max_len)])) for f in (torch.cos,torch.sin)],dim=-1),persistent=False)

    def encode(self,wave):
        z=self.encoder(wave.unsqueeze(-1))                                                                                           # B,C,T -> B,C,T,1
        return self.quantizer.encode(z,self.cache)

    def decode(self,codes):                                                                                                          # exactly 84 line after removing the blank lines
        y=self.quantizer.decode(codes,self.cache)                                                                                    # extra 1 line to load remapped‌ weights
        return self.decoder(y).squeeze(-1)                                                                                           # B,C,T,1 -> B,C,T

    def load_weights(self,weight_path="./s2-pro/codec.pth"):                                                                         # remap origianl weights
        pth=torch.load(weight_path,map_location="cpu")
        keys,values,ori_keys,i=[],[],list(pth.keys()),0
        while i<len(ori_keys):
            if ori_keys[i].endswith(".freqs_cis") or ori_keys[i].endswith(".causal_mask"): i+=1                                      # ignore rope cache and causal mask
            elif ori_keys[i].endswith(".bias") and (ori_keys[i+1].endswith(".original0") or ori_keys[i+1].endswith(".weight_g")):    # remove weight_norm
                b,g,v=ori_keys[i:i+3]; i+=3
                w=pth[g]*pth[v]/pth[v].view(pth[v].shape[0],-1).norm(2,dim=1,keepdim=True).view(-1,1,1)
                keys+=[g,b]; values+=[w.squeeze(-1),pth[b]] if g.endswith(".weight_g") else [w,pth[b]]                               # convert quantizer conv1d to linear
            elif ori_keys[i].endswith(".wqkv.weight"):
                keys+=[ori_keys[i+offset] for offset in [7,8,0,1,3,2,4,6,5]]; values+=[pth[k] for k in keys[-9:]]; i+=9              # reoder transformer weight
            else:
                keys.append(ori_keys[i]); values.append(pth[ori_keys[i]]); i+=1                                                      # conv1d to conv2d, <All keys matched successfully>
        self.load_state_dict({k:v.unsqueeze(-1) if k.endswith("conv.weight") or k.endswith(".alpha") else v for k,v in zip(self.state_dict().keys(),values)})

def benchmark(use_tensorrt=True):
    import soundfile,soxr
    codec=Codec().eval().to(device=0,dtype=torch.float16)
    codec.load_weights("./codec.pth")

    if use_tensorrt:
        import torch_tensorrt
        codec.quantizer.encode=torch.compile(codec.quantizer.encode,mode="max-autotune-no-cudagraphs")
        codec.quantizer.decode=torch.compile(codec.quantizer.decode,mode="max-autotune-no-cudagraphs")
        options={"optimization_level":5,"tiling_optimization_level":"full","cuda_graph_strategy":"whole_graph_capture"}              # a litter quicker ~1%
        codec.encoder=torch.compile(codec.encoder,fullgraph=True,dynamic=False,backend="tensorrt",options=options)
        codec.decoder=torch.compile(codec.decoder,fullgraph=True,dynamic=False,backend="tensorrt",options=options)
    else:
        codec.encode,codec.decode=[torch.compile(f) for f in [codec.encode,codec.decode]]

    # wget https://raw.githubusercontent.com/boson-ai/higgs-audio/main/examples/voice_prompts/mabaoguo.wav
    wave,sr=soundfile.read("./mabaoguo.wav",dtype="float32",always_2d=True)
    if wave.shape[1]>1: wave=wave.mean(axis=1,keepdims=True)
    if sr!=44100: wave=soxr.resample(wave,sr,44100) 
    wave=torch.as_tensor(wave).T.to(device=0,dtype=torch.float16)[None]

    # using buckets for compile mode to processing dynamic audio length, padding to bucket than clip to origianl
    # buckets=[324*i for i in range(3)]
    # padding=buckets[int((wave.shape[-1]/2048+128-1)/128)]*2048-wave.shape[-1]

    padding=0 if wave.shape[-1]%2048==0 else 2048-wave.shape[-1]%2048                                                                # frame_length=pord([2,4,8,8]+[2,2])=2048
    wave=torch.nn.functional.pad(wave,(0,padding))                                       

    for _ in range(5):                                                                                                               # warmup
        dummy=torch.randn_like(wave)
        with torch.inference_mode():
            codes=codec.encode(dummy)
            rwave=codec.decode(codes)

    encode_time,decode_time=[],[]
    duration=wave.shape[0]*wave.shape[2]/44100
    torch.cuda.synchronize()
    start_event,end_event=torch.cuda.Event(enable_timing=True),torch.cuda.Event(enable_timing=True)
    for step in range(10):
        with torch.inference_mode():
            start_event.record()
            codes=codec.encode(wave)
            end_event.record(); torch.cuda.synchronize(); time_cost=start_event.elapsed_time(end_event); encode_time.append(time_cost)

            start_event.record()
            rwave=codec.decode(codes)
            end_event.record(); torch.cuda.synchronize(); time_cost=start_event.elapsed_time(end_event); decode_time.append(time_cost)
            print(f"step: {step+1}  encode: {encode_time[-1]:.4f} ms {duration*1000/encode_time[-1]:.2f}X  decode: {decode_time[-1]:.4f} ms {duration*1000/decode_time[-1]:.2f}X")

    encode_ratio,decode_ratio=[duration*1000/sum(t)*len(t) for t in [encode_time,decode_time]]
    print(f"input: {wave.shape[0]} X {wave.shape[2]/44100:.2f} second audio, duration: {duration:.2f} second, shape: {list(wave.shape)}")
    print(f"encode Ratio: {encode_ratio:.2f}X, RTF: {1/encode_ratio:.6f}, decode Ratio: {decode_ratio:.2f}X, RTF: {1/decode_ratio:.6f}")
    soundfile.write("./offline_fp16.wav",rwave[:,:,:-padding][0].cpu().clamp(-1,1).mul(32768).to(torch.int16).T,samplerate=44100)

if __name__=="__main__":
    torch.backends.cudnn.benchmark=True
    benchmark()

''' benchmark on RTX 5090
step: 1  encode: 13.4393 ms 2062.95X  decode: 20.9561 ms 1322.98X
step: 2  encode: 13.5243 ms 2049.99x  decode: 21.0252 ms 1318.64X
step: 3  encode: 13.6576 ms 2029.97X  decode: 21.9429 ms 1263.49X
step: 4  encode: 13.8758 ms 1998.05x  decode: 21.8393 ms 1269.48X
step: 5  encode: 13.8591 ms 2000.46X  decode: 21.8003 ms 1271.76X
step: 6  encode: 13.8426 ms 2002.85x  decode: 21.7833 ms 1272.75X
step: 7  encode: 13.8718 ms 1998.64X  decode: 22.0168 ms 1259.25X
step: 8  encode: 13.8848 ms 1996.77x  decode: 22.0511 ms 1257.29X
step: 9  encode: 13.9095 ms 1993.22X  decode: 21.9796 ms 1261.38X
step: 10  encode: 13.9011 ms 1994.42X  decode: 21.9001 ms 1265.96X
input: 1 X 27.72 second audio, duration: 27.72 second, shape: [1, 1, 1222656]
encode Ratio: 2012.45X, RTF: 0.000497, decode Ratio: 1275.90X, RTF: 0.000784
'''