# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/05a_inference_anime_heads.ipynb (unless otherwise specified).

__all__ = []

# Cell
import torch
import matplotlib.pyplot as plt
from fastcore.all import *

from .torch_utils import Normalizer
from .data_anime_heads import Tokenizer
from .model_anime_heads import ExportedModel
from .constants import AnimeHeads

# Internal Cell
normalizer = Normalizer()
tokenizer = Tokenizer()

# Internal Cell
def decode_img(img):
    "img: (bs, 3, _, _), returns: (bs, _, _, 3)"
    imgs = normalizer.decode(img).permute(0, 2, 3, 1).cpu()
    return imgs

# Internal Cell
def get_attn_w(model):
    "return: (bs, seq_len, w, h)"
    attn_w = model.g_net.gs[-1].conv[0].attn_block.attn_w # (bs, w, h, seq_len)
    return attn_w.permute(0, 3, 1, 2) # (bs, seq_len, w, h)

# Cell
@patch
@torch.no_grad()
def predict(self: ExportedModel, cap):
    ''' cap: 'white hair yellow eyes'
        returns: img: (64, 64, 3), attn_w: (2, 64, 64) '''
    cap = torch.tensor([tokenizer.encode(cap)])
    self.eval()
    img = self(cap)
    img = decode_img(img)
    attn_w = get_attn_w(self)
    return img[0], attn_w[0]

# Internal Cell
def show_pred(img, attn_w, cap):
    ''' img: (64, 64, 3), attn_w: (2, 64, 64), cap: 'white hair yellow eyes' '''
    tags = [' '.join(cap.split()[:2]), ' '.join(cap.split()[2:])]

    fig, ax = plt.subplots(nrows=1, ncols=3, figsize=(12, 4))
    axs = ax.flatten()
    axs[0].imshow(img)
    axs[1].set_title(tags[0])
    axs[1].imshow(attn_w[0])
    axs[2].set_title(tags[1])
    axs[2].imshow(attn_w[1])

# Cell
@patch
def pred_and_show(self: ExportedModel, cap):
    "cap: 'white hair yellow eyes' "
    img, attn_w = self.predict(cap)
    show_pred(img, attn_w, cap)