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

# %% auto 0
__all__ = []

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

# %% ../nbs/Model inference tests-Copy2.ipynb 2
import torch
#import pylab as plt
from fastprogress import progress_bar
import time

#from torch.profiler import profile, record_function, ProfilerActivity

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

# %% ../nbs/Model inference tests-Copy2.ipynb 7
pipe = Pipeline(t2s_ref='t2s-small-en+pl.model', s2a_ref='s2a-q4-tiny-en+pl.model', use_kv_cache=False)

for emb in [pipe.t2s.embeddings.embedding, pipe.t2s.embeddings.embedding]:
    emb.convert_for_eval()
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)
#     l.attn.convert_for_eval()
#     l.cross_attn.convert_for_eval()
# for l in pipe.s2a.encoder:
#     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.generate_next = torch.compile(pipe.s2a.generate_next, mode="reduce-overhead", fullgraph=True)

# %% ../nbs/Model inference tests-Copy2.ipynb 10
#pipe.vocoder.decode_to_notebook(pipe.generate_atoks("Welcome to Burger Shack!"))
