from deepaudiox.modules.backbones import BACKBONES
from ptflops import get_model_complexity_info

import torch

x = torch.randn(1, 10 * 16_000)

backbone = BACKBONES["beats"]()
backbone.sample_rate = 16_000

y = backbone.extract_features(x)

with torch.no_grad():
    flops, params = get_model_complexity_info(
        backbone,
        (y.shape[1], y.shape[2], y.shape[3]),
        as_strings=True,
        print_per_layer_stat=True,
        verbose=True,
    )
    print(f"FLOPs: {flops}, Params: {params}")
