# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/Model inference tests-Copy1.ipynb.

# %% auto 0
__all__ = ['text', 'stoks', 'atoks']

# %% ../nbs/Model inference tests-Copy1.ipynb 1
import os
# os.environ['TORCH_LOGS'] = ""
# os.environ['TORCH_LOGS'] = "+guards,+perf_hints,+dynamo"
# os.environ['TORCH_LOGS'] = "+guards,+perf_hints"
# os.environ['TORCHDYNAMO_VERBOSE'] = '1'

# %% ../nbs/Model inference tests-Copy1.ipynb 2
import torch
from fastprogress import progress_bar
import time

from torch.profiler import profile, record_function, ProfilerActivity

# %% ../nbs/Model inference tests-Copy1.ipynb 4
from whisperspeech.pipeline import Pipeline
from whisperspeech import modules

# %% ../nbs/Model inference tests-Copy1.ipynb 11
pipe = Pipeline(s2a_ref='collabora/whisperspeech:s2a-q4-tiny-en+pl.model') #t2s_ref='t2s-small-en+pl.model', s2a_ref='s2a-q4-tiny-en+pl.model')

# for emb in pipe.s2a.embds.embeddings:
#     emb.convert_for_eval()

# for l in pipe.s2a.decoder.layers:
#     l.setup_kv_cache(1, 2250, 750)
    
# pipe.t2s.optimize()
# pipe.s2a.optimize()
    
# for l in pipe.s2a.decoder.layers:
#     l.attn.convert_for_eval()
#     l.cross_attn.convert_for_eval()
# for l in pipe.s2a.encoder:
#     l.attn.convert_for_eval()

# for l in pipe.t2s.decoder.layers:
#     l.attn.convert_for_eval()
#     l.cross_attn.convert_for_eval()
# for l in pipe.t2s.encoder.layers:
#     l.attn.convert_for_eval()
    
# torch._dynamo.reset()
# pipe.t2s.encoder = torch.compile(pipe.t2s.encoder, mode="reduce-overhead", fullgraph=True)
# pipe.t2s.decoder = torch.compile(pipe.t2s.decoder, mode="reduce-overhead", fullgraph=True)
# pipe.s2a.decoder = torch.compile(pipe.s2a.decoder, mode="reduce-overhead", fullgraph=True)
# pipe.s2a.prefill = torch.compile(pipe.s2a.prefill, mode="reduce-overhead", fullgraph=True)

# pipe.t2s.encoder = torch.compile(pipe.t2s.encoder, mode="reduce-overhead", fullgraph=True)
# pipe.s2a._encoder = torch.compile(pipe.s2a._encoder, fullgraph=True)

# pipe.s2a.generate_next = torch.compile(pipe.s2a.generate_next, mode="max-autotune", fullgraph=True)
# pipe.t2s.generate_next = torch.compile(pipe.t2s.generate_next, mode="max-autotune", fullgraph=True)

# %% ../nbs/Model inference tests-Copy1.ipynb 13
print("Warmup:")
text = """
Welcome to Burger Shack! What can I get for you today?
"""
stoks = pipe.t2s.generate(text, cps=14, lang='en', top_k=32, T=.9)
atoks = pipe.s2a.generate(stoks, pipe.default_speaker.unsqueeze(0).to(torch.float32), top_k=64, T=.7)
print()

# %% ../nbs/Model inference tests-Copy1.ipynb 14
def generate_test():
    text = """
    Welcome to Burger Shack. What can I get for you today?
    """
    start = time.time()
    with record_function("t2s"):
        stoks = pipe.t2s.generate(text, cps=14, lang='en', top_k=32, T=.9, show_progress_bar=False)
        stoks[0].item() # force a CUDA sync
    s2a_start = time.time()
    with record_function("s2a"):
        atoks = pipe.s2a.generate(stoks, pipe.default_speaker.to(torch.float32).unsqueeze(0), top_k=64, T=.7, show_progress_bar=False)
        atoks[0,0].item() # force a CUDA sync
    with record_function("decode"):
        pipe.vocoder.decode_to_notebook(atoks)
    print("Stats:", time.time() - start, time.time() - s2a_start, len(stoks)/25)

# %% ../nbs/Model inference tests-Copy1.ipynb 15
print("Optimized:")
generate_test()
print()

# %% ../nbs/Model inference tests-Copy1.ipynb 16
print("Optimized:")
generate_test()
print()

# %% ../nbs/Model inference tests-Copy1.ipynb 17
print("Optimized:")
generate_test()
print()
