Aditya02/IndicF5

1import sys2import os3 4sys.path.append(os.getcwd())5 6from f5_tts.model import CFM, DiT7 8import torch9import thop10 11 12""" ~155M """13# transformer =     UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4)14# transformer =     UNetT(dim = 768, depth = 20, heads = 12, ff_mult = 4, text_dim = 512, conv_layers = 4)15# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2)16# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4)17# transformer =       DiT(dim = 768, depth = 18, heads = 12, ff_mult = 2, text_dim = 512, conv_layers = 4, long_skip_connection = True)18# transformer =     MMDiT(dim = 512, depth = 16, heads = 16, ff_mult = 2)19 20""" ~335M """21# FLOPs: 622.1 G, Params: 333.2 M22# transformer =     UNetT(dim = 1024, depth = 24, heads = 16, ff_mult = 4)23# FLOPs: 363.4 G, Params: 335.8 M24transformer = DiT(dim=1024, depth=22, heads=16, ff_mult=2, text_dim=512, conv_layers=4)25 26 27model = CFM(transformer=transformer)28target_sample_rate = 2400029n_mel_channels = 10030hop_length = 25631duration = 2032frame_length = int(duration * target_sample_rate / hop_length)33text_length = 15034 35flops, params = thop.profile(36    model, inputs=(torch.randn(1, frame_length, n_mel_channels), torch.zeros(1, text_length, dtype=torch.long))37)38print(f"FLOPs: {flops / 1e9} G")39print(f"Params: {params / 1e6} M")40