dalle/models/__init__.py

# ------------------------------------------------------------------------------------
# Minimal DALL-E
# Copyright (c) 2021 KakaoBrain. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------------------

import os
import torch
import torch.nn as nn
import pytorch_lightning as pl
from typing import Optional, Tuple, Union
from omegaconf import OmegaConf
from torch.cuda.amp import autocast
from torch.optim.lr_scheduler import CosineAnnealingLR, LambdaLR
from torch.nn import functional as F
from .stage1.vqgan import VQGAN
from .stage2.transformer import Transformer1d, iGPT
from .stage2.layers import Block
from .. import utils
from ..utils.config import get_base_config
from ..utils.sampling import sampling, sampling_igpt, get_positional_encoding, sampling_prefix, sampling_conditional
from ..utils.utils import save_image
from .tokenizer import build_tokenizer
import numpy as np
from .stage2.layers import CrossAttentionLayer

_MODELS = {
    'minDALL-E/1.3B': 'https://arena.kakaocdn.net/brainrepo/models/minDALL-E/57b008f02ceaa02b779c8b7463143315/1.3B.tar.gz'
}

class Dalle(pl.LightningModule):
    def __init__(self,
                 config: OmegaConf) -> None:
        super().__init__()
        self.tokenizer = None
        self.stage1 = VQGAN(n_embed=config.stage1.n_embed,
                            embed_dim=config.stage1.embed_dim,
                            hparams=config.stage1.hparams)
        self.stage2 = Transformer1d(vocab_size_txt=config.stage2.vocab_size_txt,
                                    vocab_size_img=config.stage2.vocab_size_img,
                                    hparams=config.stage2.hparams)
        self.config = config
        self.config_stage1 = config.stage1
        self.config_stage2 = config.stage2
        self.config_dataset = config.dataset

        # # make the parameters in stage 1 not trainable
        # self.stage1.eval()
        # for p in self.stage1.parameters():
        #     p.requires_grad = False

    @classmethod
    def from_pretrained(cls, args) -> Tuple[nn.Module, OmegaConf]:

        path = args.model_name_or_path
        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
        if args.do_train:
            config_base = get_base_config('finetuning')
            config_update = OmegaConf.merge(config_base, config_new)
            for key, val in vars(args).items():
                if key in config_update.optimizer.keys():
                    OmegaConf.update(config_update, "optimizer.%s" % key, val, merge=False)
                if key in config_update.experiment.keys():
                    OmegaConf.update(config_update, "experiment.%s" % key, val, merge=False)
        else:
            config_base = get_base_config('default')
            config_update = OmegaConf.merge(config_base, config_new)

        model = cls(config_update)
        model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
                                          context_length=model.config_dataset.context_length,
                                          lowercase=True,
                                          dropout=None)

        print("Loading models from checkpoint %s" % path)

        if hasattr(args, 'dalle_path') and args.dalle_path and args.dalle_path.endswith('.pth'):
            model.load_state_dict(torch.load(args.dalle_path)["model_state_dict"])
        else:
            model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
            model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))

        return model, config_update
        

    @torch.no_grad()
    def sampling(self,
                 prompt: Union[str, torch.LongTensor],
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True) -> torch.FloatTensor:
        self.stage1.eval()
        self.stage2.eval()

        if type(prompt) == str:
            tokens = self.tokenizer.encode(prompt)
            tokens = torch.LongTensor(tokens.ids)
        else:
            tokens = prompt
        tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        # Check if the encoding works as intended
        # print(self.tokenizer.decode_batch(tokens.tolist(), skip_special_tokens=True)[0])

        tokens = tokens.to(device)
        codes = sampling(self.stage2,
                         tokens,
                         top_k=top_k,
                         top_p=top_p,
                         softmax_temperature=softmax_temperature,
                         use_fp16=use_fp16)
        codes = codes.view(num_candidates, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels

    def forward(self,
                images: torch.FloatTensor,
                texts: Optional[torch.LongTensor],
                past=None
                ) -> tuple:
        B, C, H, W = images.shape
        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()
        pos_enc_tokens = get_positional_encoding(texts, mode='1d')
        codes = codes.clone().detach()
        pos_enc_code = get_positional_encoding(codes, mode='1d')
        # codes = codes.unsqueeze(-1)
        # pos_enc_code = pos_enc_code.unsqueeze(-1)
        logits_img, logits_txt = self.stage2(codes, texts, pos_enc_code, pos_enc_tokens, past)
        return logits_img, logits_txt, codes

    def training_step(self, batch, batch_idx):
        images, texts = batch
        logits_img, logits_txt, codes = self(images, texts)

        loss_img = F.cross_entropy(logits_img.view(-1, logits_img.shape[-1]), codes.view(-1))
        loss_txt = F.cross_entropy(logits_txt.view(-1, logits_txt.shape[-1]), texts[:, 1:].reshape(-1))
        self.log("train/loss_img", loss_img, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        self.log("train/loss_txt", loss_txt, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        return loss_img + loss_txt

    def validation_step(self, batch, batch_idx):
        images, texts = batch
        logits_img, logits_txt, codes = self(images, texts)
        # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)

        loss_img = F.cross_entropy(logits_img.view(-1, logits_img.shape[-1]), codes.view(-1))
        loss_txt = F.cross_entropy(logits_txt.view(-1, logits_txt.shape[-1]), texts[:, 1:].reshape(-1))
        self.log("val/loss_img", loss_img, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        self.log("val/loss_txt", loss_txt, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        return loss_img + loss_txt

    def configure_optimizers(self):
        assert self.config.optimizer.opt_type == 'adamW'
        # assert self.config.optimizer.sched_type == 'cosine'

        opt = torch.optim.AdamW(self.parameters(),
                                lr=self.config.optimizer.learning_rate,
                                betas=self.config.optimizer.betas,
                                weight_decay=self.config.optimizer.weight_decay)
        # sched = CosineAnnealingLR(opt,
        #                           T_max=self.config.optimizer.max_steps,
        #                           eta_min=self.config.optimizer.min_lr)

        def lr_lambda(current_step: int):
            return max(
                0.0, float(self.config.optimizer.max_steps - current_step) / float(max(1, self.config.optimizer.max_steps))
            )

        sched = LambdaLR(opt, lr_lambda)
        sched = {
            'scheduler': sched,
            'name': 'linear'
        }
        return [opt], [sched]

    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure,
                       on_tpu=False, using_native_amp=False, using_lbfgs=False):
        optimizer.step(closure=optimizer_closure)
        self.lr_schedulers().step()
        self.log("lr", self.lr_schedulers().get_last_lr()[0], on_step=True, on_epoch=False, prog_bar=True, logger=True)

    def on_epoch_start(self):
        self.stage1.eval()


class ImageGPT(pl.LightningModule):
    def __init__(self,
                 config: OmegaConf) -> None:
        super().__init__()
        self.stage1 = VQGAN(n_embed=config.stage1.n_embed,
                            embed_dim=config.stage1.embed_dim,
                            hparams=config.stage1.hparams)
        self.stage2 = iGPT(vocab_size_img=config.stage2.vocab_size_img,
                           use_cls_cond=config.stage2.use_cls_cond,
                           hparams=config.stage2.hparams)
        self.config = config
        self.use_cls_cond = config.stage2.use_cls_cond

        # make the parameters in stage 1 not trainable
        self.stage1.eval()
        for p in self.stage1.parameters():
            p.requires_grad = False

    @classmethod
    def from_pretrained(cls,
                        path_upstream: str,
                        path_downstream: str) -> Tuple[nn.Module, OmegaConf]:
        config_base = get_base_config(use_default=False)
        config_down = OmegaConf.load(path_downstream)
        config_down = OmegaConf.merge(config_base, config_down)

        model = cls(config_down)
        model.stage1.from_ckpt(os.path.join(path_upstream, 'stage1_last.ckpt'), strict=True)
        model.stage2.from_ckpt(os.path.join(path_upstream, 'stage2_last.ckpt'), strict=False)
        return model, config_down

    def sample(self,
               cls_idx: Optional[int] = None,
               top_k: int = 256,
               top_p: Optional[float] = None,
               softmax_temperature: float = 1.0,
               num_candidates: int = 16,
               device: str = 'cuda:0',
               use_fp16: bool = True,
               is_tqdm: bool = True) -> torch.FloatTensor:
        self.stage1.eval()
        self.stage2.eval()

        if cls_idx is None:
            sos = self.stage2.sos.repeat(num_candidates, 1, 1)
        else:
            sos = torch.LongTensor([cls_idx]).to(device=device)
            sos = sos.repeat(num_candidates)
            sos = self.stage2.sos(sos).unsqueeze(1)

        codes = sampling_igpt(self.stage2,
                              sos=sos,
                              top_k=top_k,
                              top_p=top_p,
                              softmax_temperature=softmax_temperature,
                              use_fp16=use_fp16,
                              is_tqdm=is_tqdm)
        codes = codes.view(num_candidates, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels

    def forward(self,
                images: torch.FloatTensor,
                labels: Optional[torch.LongTensor] = None) -> torch.FloatTensor:
        B, C, H, W = images.shape
        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()
        logits = self.stage2(codes, labels)
        return logits, codes

    def training_step(self, batch, batch_idx):
        images, labels = batch
        logits, codes = self(images, labels=labels if self.use_cls_cond else None)
        loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), codes.view(-1))
        self.log("train/loss", loss, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        return loss

    def validation_step(self, batch, batch_idx):
        images, labels = batch
        logits, codes = self(images, labels=labels if self.use_cls_cond else None)
        loss = F.cross_entropy(logits.view(-1, logits.shape[-1]), codes.view(-1))
        self.log("val/loss", loss, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        return loss

    def configure_optimizers(self):
        assert self.config.optimizer.opt_type == 'adamW'
        assert self.config.optimizer.sched_type == 'cosine'

        opt = torch.optim.AdamW(self.parameters(),
                                lr=self.config.optimizer.base_lr,
                                betas=self.config.optimizer.betas,
                                weight_decay=self.config.optimizer.weight_decay)
        sched = CosineAnnealingLR(opt,
                                  T_max=self.config.optimizer.max_steps,
                                  eta_min=self.config.optimizer.min_lr)
        sched = {
            'scheduler': sched,
            'name': 'cosine'
        }
        return [opt], [sched]

    def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure,
                       on_tpu=False, using_native_amp=False, using_lbfgs=False):
        optimizer.step(closure=optimizer_closure)
        self.lr_schedulers().step()
        self.log("lr", self.lr_schedulers().get_last_lr()[0], on_step=True, on_epoch=False, prog_bar=True, logger=True)

    def on_epoch_start(self):
        self.stage1.eval()


class PromptDalle(Dalle):
    """Classification Head for  transformer encoders"""
    def __init__(self, config):
        super().__init__(config)
        print('Initializing the PromptTuning model')

        self.config = config
        self.n_embd = config.stage2.hparams.embed_dim
        self.preseqlen = config.prompt.preseqlen
        self.prefix_dropout = config.prompt.prefix_dropout

        # DIFFERENT PARAMETRIZATION:

        print('[Full prompt-tuning Setting :) ]')
        self.input_tokens = torch.arange(self.preseqlen).long()
        self.wte = nn.Embedding(self.preseqlen, self.n_embd)
        self.control_trans = nn.Sequential(
            nn.Linear(self.n_embd, self.n_embd),
            nn.Tanh(),
            nn.Linear(self.n_embd, self.n_embd))
        self.get_prompt = self.get_prompt_p5
        self.dropout = nn.Dropout(self.prefix_dropout)

        ###### NUM PARAMS #########
        total_param = 0
        for name, param in self.named_parameters():
            # print(param.shape)
            total_param += param.numel()
        print('Total parameters is {}'.format(total_param))


    @classmethod
    def from_pretrained(cls, args) -> Tuple[nn.Module, OmegaConf]:

        # if not args.model_name_or_path:
        #     args.model_name_or_path = args.prefix_model_name_or_path

        path = args.prefix_model_name_or_path
        path = _MODELS[path] if path in _MODELS else path
        path = utils.realpath_url_or_path(path, root=os.path.expanduser("~/.cache/minDALL-E"))

        config_base = get_base_config('prompt_tuning')
        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
        config_update = OmegaConf.merge(config_base, config_new)

        for key, val in vars(args).items():
            if key in config_update.prompt.keys():
                OmegaConf.update(config_update, "prompt.%s" % key, val, merge=False)
            if key in config_update.optimizer.keys():
                OmegaConf.update(config_update, "optimizer.%s" % key, val, merge=False)
            if key in config_update.experiment.keys():
                OmegaConf.update(config_update, "experiment.%s" % key, val, merge=False)

        model = cls(config_update)
        model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
                                          context_length=model.config_dataset.context_length,
                                          lowercase=True,
                                          dropout=None)

        if args.model_name_or_path:
            print("Loading model from pretrained checkpoint %s" % args.model_name_or_path)
            # model.from_ckpt(args.model_name_or_path)
            try:
                model.load_state_dict(torch.load(args.model_name_or_path)['state_dict'])
            except KeyError:
                model.load_state_dict(torch.load(args.model_name_or_path)['model_state_dict'])

        else:
            print("Loading models from checkpoint %s" % path)
            model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
            model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))

        return model, config_update

    def get_prompt_p5(self, bsz=None, eval=False):
        input_tokens = self.input_tokens.unsqueeze(0).expand(bsz, -1).to(self.device)
        temp_control = self.wte(input_tokens)
        past_key_values = self.control_trans(temp_control) #bsz, seqlen, layer*emb
        if not eval:
            past_key_values = self.dropout(past_key_values)
        return past_key_values

    def forward(self,
                images: torch.FloatTensor,
                texts: Optional[torch.LongTensor],
                **kwargs,
        ):

        #{"input_ids": batch, "labels": labels, 'src_attn': src_attn, 'tgt_attn':tgt_attn, 'src':src}

        B, C, H, W = images.shape
        prompt = self.get_prompt(bsz=B)
        pos_enc_prompt = get_positional_encoding(self.input_tokens.unsqueeze(0).expand(B, -1).to(self.device), mode='1d')

        # if self.mode_para == 2 and src_attn is not None and tgt_attn is not None:
        #     attention_mask = torch.cat([src_attn, tgt_attn], dim=1)


        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()

        pos_enc_tokens = get_positional_encoding(texts, mode='1d')
        codes = codes.clone().detach()
        pos_enc_code = get_positional_encoding(codes, mode='1d')
        # codes = codes.unsqueeze(-1)
        # pos_enc_code = pos_enc_code.unsqueeze(-1)
        # print(images.shape, codes.shape, texts.shape)
        logits_img, logits_txt = self.stage2(codes, texts, pos_enc_code, pos_enc_tokens, prompt=prompt, pos_prompt=pos_enc_prompt)
        return logits_img, logits_txt, codes


    @torch.no_grad()
    def sampling(self,
                 tokens: torch.LongTensor,
                 prompt: torch.FloatTensor,
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True,
                 labels = None) -> torch.FloatTensor:
        self.stage1.eval()
        self.stage2.eval()

        # tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        tokens = tokens.to(device)
        pos_enc_prompt = get_positional_encoding(self.input_tokens.unsqueeze(0).expand(num_candidates, -1).to(self.device), mode='1d')

        codes = sampling(self.stage2,
                         tokens,
                         top_k=top_k,
                         top_p=top_p,
                         softmax_temperature=softmax_temperature,
                         use_fp16=use_fp16,
                         prompt=prompt,
                         pos_prompt=pos_enc_prompt)

        codes = codes.view(-1, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels


    @torch.no_grad()
    def predict_step(self, batch, batch_idx, return_images=False):
        orig_images, texts = batch

        # extra for checks
        logits_img, logits_txt, codes = self(orig_images, texts)
        pred = torch.argmax(logits_img.view(-1, logits_img.shape[-1]), dim=-1)
        bs = orig_images.shape[0]
        pred = pred.view(bs, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(pred) * 0.5 + 0.5, 0, 1).cpu().numpy()  # [B, 256, 256]
        pixels = np.transpose(pixels, (0, 2, 3, 1))

        # print(texts.shape, orig_images.shape)
        prompt = self.get_prompt(bsz=5, eval=True)

        images = []
        for i, t in enumerate(texts):
            pixels = self.sampling(t, prompt, top_k=16, num_candidates=5, labels=codes[i]).cpu().numpy()
            pixels = np.transpose(pixels, (0, 2, 3, 1))
            images.append(pixels)

        if return_images:
            return images
        else:
            save_image(orig_images, pixels, './out/images/pororo_prompt', batch_idx+10)
            save_image(orig_images, images, './out/images/pororo_prompt', batch_idx)


class PrefixTuningDalle(Dalle):
    """Classification Head for  transformer encoders"""
    def __init__(self, config):
        super().__init__(config)
        print('Initializing the PrefixTuning model')

        self.config = config

        self.match_n_layer = config.stage2.hparams.n_layers
        self.match_n_head = config.stage2.hparams.n_heads
        self.match_n_embd = config.stage2.hparams.embed_dim // config.stage2.hparams.n_heads
        self.n_embd = config.stage2.hparams.embed_dim

        self.optim_prefix = config.prefix.optim_prefix
        self.preseqlen = config.prefix.preseqlen
        self.prefix_dropout = config.prefix.prefix_dropout
        self.init_random = config.prefix.init_random
        self.hidden_dim_prefix = config.prefix.hidden_dim_prefix

        self.lowdata_token = config.prefix.lowdata_token
        self.init_shallow = config.prefix.init_shallow
        self.init_shallow_word = config.prefix.init_shallow_word
        self.mode_para = 0

        print('PrefixTuning')
        print('preseqlen is {}, optimizing the prefix directly'.format(self.preseqlen))

        # DIFFERENT PARAMETRIZATION:

        print('[Full prefix-tuning Setting :) ]')
        self.input_tokens = torch.arange(self.preseqlen).long()
        self.wte = nn.Embedding(self.preseqlen, self.n_embd)
        self.control_trans = nn.Sequential(
            nn.Linear(self.n_embd, self.hidden_dim_prefix),
            nn.Tanh(),
            nn.Linear(self.hidden_dim_prefix, self.match_n_layer * 2 * self.n_embd))
        self.get_prompt = self.get_prompt_p5
        self.dropout = nn.Dropout(self.prefix_dropout)

        ###### NUM PARAMS #########
        total_param = 0
        for name, param in self.named_parameters():
            # print(param.shape)
            total_param += param.numel()
        print('Total parameters is {}'.format(total_param))


    @classmethod
    def from_pretrained(cls, args) -> Tuple[nn.Module, OmegaConf]:

        # if not args.model_name_or_path:
        #     args.model_name_or_path = args.prefix_model_name_or_path

        path = args.prefix_model_name_or_path
        path = _MODELS[path] if path in _MODELS else path
        path = utils.realpath_url_or_path(path, root=os.path.expanduser("~/.cache/minDALL-E"))

        config_base = get_base_config('prefixtuning')
        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
        config_update = OmegaConf.merge(config_base, config_new)

        for key, val in vars(args).items():
            if key in config_update.prefix.keys():
                OmegaConf.update(config_update, "prefix.%s" % key, val, merge=False)
            if key in config_update.optimizer.keys():
                OmegaConf.update(config_update, "optimizer.%s" % key, val, merge=False)
            if key in config_update.experiment.keys():
                OmegaConf.update(config_update, "experiment.%s" % key, val, merge=False)

        model = cls(config_update)
        model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
                                          context_length=model.config_dataset.context_length,
                                          lowercase=True,
                                          dropout=None)

        if args.model_name_or_path:
            print("Loading model from pretrained checkpoint %s" % args.model_name_or_path)
            # model.from_ckpt(args.model_name_or_path)
            try:
                model.load_state_dict(torch.load(args.model_name_or_path)['state_dict'])
            except KeyError:
                model.load_state_dict(torch.load(args.model_name_or_path)['model_state_dict'])

        else:
            print("Loading models from checkpoint %s" % path)
            model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
            model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))

        return model, config_update

    def get_prompt_p5(self, bsz=None, eval=False):
        input_tokens = self.input_tokens.unsqueeze(0).expand(bsz, -1).to(self.device)
        temp_control = self.wte(input_tokens)
        past_key_values = self.control_trans(temp_control) #bsz, seqlen, layer*emb
        bsz, seqlen, _ = past_key_values.shape
        past_key_values = past_key_values.view(bsz, seqlen, self.match_n_layer * 2, self.match_n_head,
                                               self.match_n_embd)
        if not eval:
            past_key_values = self.dropout(past_key_values)
        # past_key_values = past_key_values.permute([2, 0, 3, 1, 4]).split(2)
        past_key_values = past_key_values.permute([2, 0, 3, 1, 4])
        # print(past_key_values.shape)
        return past_key_values.split(2)

    def forward(self,
                images: torch.FloatTensor,
                texts: Optional[torch.LongTensor],
                **kwargs,
        ):

        #{"input_ids": batch, "labels": labels, 'src_attn': src_attn, 'tgt_attn':tgt_attn, 'src':src}

        B, C, H, W = images.shape

        if self.mode_para == 2:
            past_key_values_prompt = self.get_prompt(bsz=B)
        else:
            past_key_values_prompt = self.get_prompt(bsz=B)

        # if self.mode_para == 2 and src_attn is not None and tgt_attn is not None:
        #     attention_mask = torch.cat([src_attn, tgt_attn], dim=1)


        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()

        pos_enc_tokens = get_positional_encoding(texts, mode='1d')
        codes = codes.clone().detach()
        pos_enc_code = get_positional_encoding(codes, mode='1d')
        # codes = codes.unsqueeze(-1)
        # pos_enc_code = pos_enc_code.unsqueeze(-1)
        # print(images.shape, codes.shape, texts.shape)
        logits_img, logits_txt = self.stage2(codes, texts, pos_enc_code, pos_enc_tokens, past_key_values_prompt)
        return logits_img, logits_txt, codes

    @torch.no_grad()
    def sampling(self,
                 tokens: torch.LongTensor,
                 past: torch.FloatTensor,
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True,
                 labels = None) -> torch.FloatTensor:
        self.stage1.eval()
        self.stage2.eval()

        if len(past.shape) == 6:
            n_layers, temp, bs, n_heads, seq_len, n_dim = past.shape
            past = past.view(n_layers, temp, bs*n_heads, seq_len, n_dim)

        tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        # Check if the encoding works as intended
        # print(self.tokenizer.decode_batch(tokens.tolist(), skip_special_tokens=True)[0])

        tokens = tokens.to(device)
        codes = sampling_prefix(self.stage2,
                                tokens,
                                past,
                                top_k=top_k,
                                top_p=top_p,
                                softmax_temperature=softmax_temperature,
                                use_fp16=use_fp16,
                                labels = None if labels is None else labels.view(-1))

        # codes = sampling(self.stage2,
        #                  tokens,
        #                  top_k=top_k,
        #                  top_p=top_p,
        #                  softmax_temperature=softmax_temperature,
        #                  use_fp16=use_fp16)

        codes = codes.view(num_candidates, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels

    def training_step(self, batch, batch_idx):
        images, texts = batch
        logits_img, logits_txt, codes = self(images, texts)

        loss_img = F.cross_entropy(logits_img.view(-1, logits_img.shape[-1]), codes.view(-1))
        loss_txt = F.cross_entropy(logits_txt.view(-1, logits_txt.shape[-1]), texts[:, 1:].reshape(-1))
        self.log("train/loss_img", loss_img, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        self.log("train/loss_txt", loss_txt, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        return loss_img + loss_txt

    def validation_step(self, batch, batch_idx):
        images, texts = batch
        logits_img, logits_txt, codes = self(images, texts)
        # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)

        loss_img = F.cross_entropy(logits_img.view(-1, logits_img.shape[-1]), codes.view(-1))
        loss_txt = F.cross_entropy(logits_txt.view(-1, logits_txt.shape[-1]), texts[:, 1:].reshape(-1))
        self.log("val/loss_img", loss_img, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        self.log("val/loss_txt", loss_txt, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        return loss_img + loss_txt

    @torch.no_grad()
    def predict_step(self, batch, batch_idx, return_images=False):
        orig_images, texts = batch

        # extra for checks
        logits_img, logits_txt, codes = self(orig_images, texts)
        pred = torch.argmax(logits_img.view(-1, logits_img.shape[-1]), dim=-1)
        bs = orig_images.shape[0]
        pred = pred.view(bs, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(pred) * 0.5 + 0.5, 0, 1).cpu().numpy()  # [B, 256, 256]
        pixels = np.transpose(pixels, (0, 2, 3, 1))


        # print(texts.shape, orig_images.shape)
        # concatenate the list of prompts (split by n_head) for better downstream processing
        past_key_values_prompt = self.get_prompt(bsz=5, eval=True)
        # print(past_key_values_prompt[0].shape, past_key_values_prompt[1].shape, len(past_key_values_prompt))
        past_key_values_prompt = torch.cat([x.unsqueeze(0) for x in past_key_values_prompt], dim=0)
        n_layers, temp, bs, n_heads, seq_len, n_dim = past_key_values_prompt.shape
        past_key_values_prompt = past_key_values_prompt.view(n_layers, temp, bs*n_heads, seq_len, n_dim)
        # print(past_key_values_prompt.shape)
        images = []
        for i, t in enumerate(texts):
            pixels = self.sampling(t, past_key_values_prompt, top_k=16, num_candidates=5, labels=codes[i]).cpu().numpy()
            pixels = np.transpose(pixels, (0, 2, 3, 1))
            images.append(pixels)
            # images.extend([p for p in pixels])
            # print([i.shape for i in images])


        if return_images:
            return images
        else:
            save_image(orig_images, pixels, './out/images/pororo_prefix', batch_idx+10)
            save_image(orig_images, images, './out/images/pororo_prefix', batch_idx)


class ConditionalDalle(Dalle):
    """Classification Head for  transformer encoders"""
    def __init__(self, config):
        super().__init__(config)
        print('Initializing the Conditional Dalle model')

        self.config = config

        print('Setting up Cross-attention Layers')
        self.init_cross_attention(list(range(2,42,3)), config.stage2.hparams)

        ###### NUM PARAMS #########
        total_param = 0
        for name, param in self.named_parameters():
            # print(param.shape)
            total_param += param.numel()
        print('Total parameters is {}'.format(total_param))

    @classmethod
    def from_pretrained(cls, args) -> Tuple[nn.Module, OmegaConf]:

        # if not args.model_name_or_path:
        #     args.model_name_or_path = args.prefix_model_name_or_path

        path = args.model_name_or_path
        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
        if args.do_train:
            config_base = get_base_config('finetuning')
            config_update = OmegaConf.merge(config_base, config_new)
            for key, val in vars(args).items():
                if key in config_update.optimizer.keys():
                    OmegaConf.update(config_update, "optimizer.%s" % key, val, merge=False)
                if key in config_update.experiment.keys():
                    OmegaConf.update(config_update, "experiment.%s" % key, val, merge=False)
        else:
            config_base = get_base_config('default')
            config_update = OmegaConf.merge(config_base, config_new)

        model = cls(config_update)
        model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
                                          context_length=model.config_dataset.context_length,
                                          lowercase=True,
                                          dropout=None)
        print(model.cross_attention_idxs)
        # print(next(model.cross_attention_layers[0].parameters()).is_cuda)

        if args.dalle_path:
            print("Loading model from pretrained checkpoint %s" % args.dalle_path)
            # model.from_ckpt(args.model_name_or_path)
            model.load_state_dict(torch.load(args.dalle_path)['model_state_dict'])
        else:
            print("Loading models from checkpoint %s" % path)
            model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
            model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))

        return model, config_update


    def init_cross_attention(self, cross_attention_layers, hparams):
        self.cross_attention_idxs = cross_attention_layers
        self.cross_attention_layers = [CrossAttentionLayer(ctx_len=hparams.ctx_len_img + hparams.ctx_len_txt,
                                         embed_dim=hparams.embed_dim,
                                         n_heads=hparams.n_heads,
                                         attn_bias=hparams.attn_bias,
                                         resid_pdrop=hparams.resid_pdrop,
                                         attn_pdrop=hparams.attn_pdrop) for i in cross_attention_layers]


    def forward(self,
                images: torch.FloatTensor,
                src_images: Optional[torch.FloatTensor],
                texts: Optional[torch.LongTensor],
                **kwargs,
        ):

        #{"input_ids": batch, "labels": labels, 'src_attn': src_attn, 'tgt_attn':tgt_attn, 'src':src}

        # print(images.shape, src_images.shape, texts.shape)
        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()
                src_codes = self.stage1.get_codes(src_images).detach()

        pos_enc_tokens = get_positional_encoding(texts, mode='1d')
        codes = codes.clone().detach()
        pos_enc_code = get_positional_encoding(codes, mode='1d')
        src_codes = src_codes.clone().detach()
        src_pos_enc_code = get_positional_encoding(src_codes, mode='1d')
        # codes = codes.unsqueeze(-1)
        # pos_enc_code = pos_enc_code.unsqueeze(-1)
        # print(images.shape, codes.shape, texts.shape)
        logits_img, logits_txt = self.stage2.forward_with_context(codes, texts,
                                                                  pos_enc_code, pos_enc_tokens, src_codes, src_pos_enc_code,
                                                                  self.cross_attention_idxs, self.cross_attention_layers)
        # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)
        return logits_img, logits_txt, codes

    @torch.no_grad()
    def sampling(self,
                 prompt: torch.LongTensor,
                 source: torch.FloatTensor,
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True) -> torch.FloatTensor:
        self.stage1.eval()
        self.stage2.eval()

        if type(prompt) == str:
            tokens = self.tokenizer.encode(prompt)
            tokens = torch.LongTensor(tokens.ids)
        else:
            tokens = prompt

        tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        # Check if the encoding works as intended
        # print(self.tokenizer.decode_batch(tokens.tolist(), skip_special_tokens=True)[0])

        tokens = tokens.to(device)
        source = source.to(device)

        with autocast(enabled=False):
            src_codes = self.stage1.get_codes(source).detach()
        src_codes = torch.repeat_interleave(src_codes, num_candidates, dim=0)

        codes = sampling_conditional(self.stage2,
                                     self.cross_attention_idxs,
                                     self.cross_attention_layers,
                                     tokens,
                                     src_codes,
                                     top_k=top_k,
                                     top_p=top_p,
                                     softmax_temperature=softmax_temperature,
                                     use_fp16=use_fp16)
        codes = codes.view(num_candidates, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels

    def training_step(self, batch, batch_idx):
        images, texts = batch
        logits_img, logits_txt, codes = self(images, texts)

        loss_img = F.cross_entropy(logits_img.view(-1, logits_img.shape[-1]), codes.view(-1))
        loss_txt = F.cross_entropy(logits_txt.view(-1, logits_txt.shape[-1]), texts[:, 1:].reshape(-1))
        self.log("train/loss_img", loss_img, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        self.log("train/loss_txt", loss_txt, on_step=True, on_epoch=True, prog_bar=False, logger=True)
        return loss_img + loss_txt

    def validation_step(self, batch, batch_idx):
        images, texts = batch
        logits_img, logits_txt, codes = self(images, texts)
        # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)

        loss_img = F.cross_entropy(logits_img.view(-1, logits_img.shape[-1]), codes.view(-1))
        loss_txt = F.cross_entropy(logits_txt.view(-1, logits_txt.shape[-1]), texts[:, 1:].reshape(-1))
        self.log("val/loss_img", loss_img, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        self.log("val/loss_txt", loss_txt, on_step=False, on_epoch=True, prog_bar=False, logger=True)
        return loss_img + loss_txt

    @torch.no_grad()
    def predict_step(self, batch, batch_idx):
        orig_images, texts = batch
        # concatenate the list of prompts (split by n_head) for better downstream processing
        past_key_values_prompt = self.get_prompt(bsz=5)
        past_key_values_prompt = torch.cat([x.unsqueeze(0) for x in past_key_values_prompt], dim=0)
        images = []
        for t in texts:
            pixels = self.sampling(t, past_key_values_prompt, top_k=64, num_candidates=5).cpu().numpy()
            pixels = np.transpose(pixels, (0, 2, 3, 1))
            images.append(pixels)
            # images.extend([p for p in pixels])
            # print([i.shape for i in images])

        save_image(orig_images, images, './out/images/', batch_idx)


class PromptConditionalDalle(Dalle):
    """Classification Head for  transformer encoders"""
    def __init__(self, config):
        super().__init__(config)
        print('Initializing the Conditional Dalle model')

        self.config = config

        print('Setting up Cross-attention Layers')
        self.init_cross_attention(list(range(2,42,3)), config.stage2.hparams)

        self.n_embd = config.stage2.hparams.embed_dim
        self.preseqlen = config.story.preseqlen
        self.prefix_dropout = config.story.prefix_dropout

        # DIFFERENT PARAMETRIZATION:

        print('[Full prompt-tuning Setting :) ]')
        self.input_tokens = torch.arange(self.preseqlen).long()
        self.wte = nn.Embedding(self.preseqlen, self.n_embd)
        self.control_trans = nn.Sequential(
            nn.Linear(self.n_embd, self.n_embd),
            nn.Tanh(),
            nn.Linear(self.n_embd, self.n_embd))
        self.get_prompt = self.get_prompt_p5
        self.dropout = nn.Dropout(self.prefix_dropout)

        ###### NUM PARAMS #########
        total_param = 0
        for name, param in self.named_parameters():
            # print(param.shape)
            total_param += param.numel()
        print('Total parameters is {}'.format(total_param))

    @classmethod
    def from_pretrained(cls, args) -> Tuple[nn.Module, OmegaConf]:

        # if not args.model_name_or_path:
        #     args.model_name_or_path = args.prefix_model_name_or_path

        path = args.prefix_model_name_or_path
        path = _MODELS[path] if path in _MODELS else path
        path = utils.realpath_url_or_path(path, root=os.path.expanduser("~/.cache/minDALL-E"))

        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
        if args.do_train:
            config_base = get_base_config('story')
            config_update = OmegaConf.merge(config_base, config_new)
            for key, val in vars(args).items():
                if key in config_update.story.keys():
                    OmegaConf.update(config_update, "story.%s" % key, val, merge=False)
                if key in config_update.optimizer.keys():
                    OmegaConf.update(config_update, "optimizer.%s" % key, val, merge=False)
                if key in config_update.experiment.keys():
                    OmegaConf.update(config_update, "experiment.%s" % key, val, merge=False)
        else:
            config_base = get_base_config('default')
            config_update = OmegaConf.merge(config_base, config_new)

        model = cls(config_update)
        model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
                                          context_length=model.config_dataset.context_length,
                                          lowercase=True,
                                          dropout=None)
        print(model.cross_attention_idxs)
        # print(next(model.cross_attention_layers[0].parameters()).is_cuda)

        if args.model_name_or_path:
            print("Loading model from pretrained checkpoint %s" % args.model_name_or_path)
            # model.from_ckpt(args.model_name_or_path)
            try:
                model.load_state_dict(torch.load(args.model_name_or_path)['state_dict'])
            except KeyError:
                model.load_state_dict(torch.load(args.model_name_or_path)['model_state_dict'])

        else:
            print("Loading models from checkpoint %s" % path)
            model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
            model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))

        return model, config_update


    def init_cross_attention(self, cross_attention_layers, hparams):
        self.cross_attention_idxs = cross_attention_layers
        self.cross_attention_layers = [CrossAttentionLayer(ctx_len=hparams.ctx_len_img + hparams.ctx_len_txt,
                                         embed_dim=hparams.embed_dim,
                                         n_heads=hparams.n_heads,
                                         attn_bias=hparams.attn_bias,
                                         resid_pdrop=hparams.resid_pdrop,
                                         attn_pdrop=hparams.attn_pdrop) for i in cross_attention_layers]

    def get_prompt_p5(self, bsz=None, eval=False):
        input_tokens = self.input_tokens.unsqueeze(0).expand(bsz, -1).to(self.device)
        temp_control = self.wte(input_tokens)
        past_key_values = self.control_trans(temp_control) #bsz, seqlen, layer*emb
        if not eval:
            past_key_values = self.dropout(past_key_values)
        return past_key_values

    def forward(self,
                images: torch.FloatTensor,
                src_images: Optional[torch.FloatTensor],
                texts: Optional[torch.LongTensor],
                **kwargs,
        ):

        #{"input_ids": batch, "labels": labels, 'src_attn': src_attn, 'tgt_attn':tgt_attn, 'src':src}

        # print(images.shape, src_images.shape, texts.shape)
        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()
                src_codes = self.stage1.get_codes(src_images).detach()

        B, C, H, W = images.shape
        prompt = self.get_prompt(bsz=B)
        pos_enc_prompt = get_positional_encoding(self.input_tokens.unsqueeze(0).expand(B, -1).to(self.device), mode='1d')

        pos_enc_tokens = get_positional_encoding(texts, mode='1d')
        codes = codes.clone().detach()
        pos_enc_code = get_positional_encoding(codes, mode='1d')
        src_codes = src_codes.clone().detach()
        src_pos_enc_code = get_positional_encoding(src_codes, mode='1d')
        # codes = codes.unsqueeze(-1)
        # pos_enc_code = pos_enc_code.unsqueeze(-1)
        # print(images.shape, codes.shape, texts.shape)
        logits_img, logits_txt = self.stage2.forward_with_context(codes, texts,
                                                                  pos_enc_code, pos_enc_tokens, src_codes, src_pos_enc_code,
                                                                  self.cross_attention_idxs, self.cross_attention_layers,
                                                                  prompt=prompt, pos_prompt=pos_enc_prompt)
        # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)
        return logits_img, logits_txt, codes

    @torch.no_grad()
    def sampling(self,
                 tokens: torch.LongTensor,
                 prompt: torch.LongTensor,
                 source: torch.FloatTensor,
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True,
                 labels=None) -> torch.FloatTensor:

        self.stage1.eval()
        self.stage2.eval()

        if type(tokens) == str:
            tokens = self.tokenizer.encode(prompt)
            tokens = torch.LongTensor(tokens.ids)
        else:
            pass

        tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        # Check if the encoding works as intended
        # print(self.tokenizer.decode_batch(tokens.tolist(), skip_special_tokens=True)[0])

        tokens = tokens.to(device)
        source = source.to(device)

        pos_enc_prompt = get_positional_encoding(self.input_tokens.unsqueeze(0).expand(num_candidates, -1).to(self.device), mode='1d')

        with autocast(enabled=False):
            src_codes = self.stage1.get_codes(source).detach()
        src_codes = torch.repeat_interleave(src_codes, num_candidates, dim=0)

        codes = sampling_conditional(self.stage2,
                                     self.cross_attention_idxs,
                                     self.cross_attention_layers,
                                     tokens,
                                     src_codes,
                                     top_k=top_k,
                                     top_p=top_p,
                                     softmax_temperature=softmax_temperature,
                                     use_fp16=use_fp16,
                                     prompt=prompt,
                                     pos_prompt=pos_enc_prompt)

        codes = codes.view(num_candidates, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels


    @torch.no_grad()
    def predict_step(self, batch, batch_idx, return_images=False):
        orig_images, texts = batch
        # concatenate the list of prompts (split by n_head) for better downstream processing

        # extra for checks
        logits_img, logits_txt, codes = self(orig_images, texts)
        pred = torch.argmax(logits_img.view(-1, logits_img.shape[-1]), dim=-1)
        bs = orig_images.shape[0]
        pred = pred.view(bs, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(pred) * 0.5 + 0.5, 0, 1).cpu().numpy()  # [B, 256, 256]
        pixels = np.transpose(pixels, (0, 2, 3, 1))

        prompt = self.get_prompt(bsz=5, eval=True)

        images = []
        for t in texts:
            pixels = self.sampling(t, prompt, top_k=64, num_candidates=5, labels=codes[i]).cpu().numpy()
            pixels = np.transpose(pixels, (0, 2, 3, 1))
            images.append(pixels)
            # images.extend([p for p in pixels])
            # print([i.shape for i in images])

        if return_images:
            return images
        else:
            save_image(orig_images, pixels, './out/images/pororo_story', batch_idx+10)
            save_image(orig_images, images, './out/images/pororo_story', batch_idx)


class StoryDalle(Dalle):
    """Base model with story block"""
    def __init__(self, config):
        super().__init__(config)
        print('Initializing the Conditional Dalle model')

        self.config = config

        self.story_linear = nn.Linear(config.story.sent_embed, config.stage2.hparams.embed_dim)
        self.story_block = Block(ctx_len=config.story.story_len,
                             embed_dim=config.stage2.hparams.embed_dim,
                             n_heads=config.stage2.hparams.n_heads,
                             mlp_bias=config.stage2.hparams.mlp_bias,
                             attn_bias=config.stage2.hparams.attn_bias,
                             resid_pdrop=config.stage2.hparams.resid_pdrop,
                             attn_pdrop=config.stage2.hparams.attn_pdrop,
                             gelu_use_approx=config.stage2.hparams.gelu_use_approx)

        if self.config.story.prompt:
            self.n_embd = config.stage2.hparams.embed_dim
            self.preseqlen = config.story.preseqlen
            self.prefix_dropout = config.story.prefix_dropout

            # DIFFERENT PARAMETRIZATION:

            print('[Full prompt-tuning Setting :) ]')
            self.input_tokens = torch.arange(self.preseqlen).long()
            self.wte = nn.Embedding(self.preseqlen, self.n_embd)
            self.control_trans = nn.Sequential(
                nn.Linear(self.n_embd, self.n_embd),
                nn.Tanh(),
                nn.Linear(self.n_embd, self.n_embd))
            self.get_prompt = self.get_prompt_p5
            self.dropout = nn.Dropout(self.prefix_dropout)

        if self.config.story.condition:
            print('Setting up Cross-attention Layers')
            self.init_cross_attention(list(range(2,42,3)), config.stage2.hparams)

        ###### NUM PARAMS #########
        total_param = 0
        for name, param in self.named_parameters():
            # print(param.shape)
            total_param += param.numel()
        print('Total parameters is {}'.format(total_param))

    @classmethod
    def from_pretrained(cls, args) -> Tuple[nn.Module, OmegaConf]:

        # if not args.model_name_or_path:
        #     args.model_name_or_path = args.prefix_model_name_or_path

        path = args.prefix_model_name_or_path
        path = _MODELS[path] if path in _MODELS else path
        path = utils.realpath_url_or_path(path, root=os.path.expanduser("~/.cache/minDALL-E"))

        config_new = OmegaConf.load(os.path.join(path, 'config.yaml'))
        # if args.do_train:
        config_base = get_base_config('story')
        config_update = OmegaConf.merge(config_base, config_new)
        for key, val in vars(args).items():
            if key in config_update.story.keys():
                OmegaConf.update(config_update, "story.%s" % key, val, merge=False)
            if key in config_update.optimizer.keys():
                OmegaConf.update(config_update, "optimizer.%s" % key, val, merge=False)
            if key in config_update.experiment.keys():
                OmegaConf.update(config_update, "experiment.%s" % key, val, merge=False)
        # else:
        #     config_base = get_base_config('story')
        #     config_update = OmegaConf.merge(config_base, config_new)
        # print(next(model.cross_attention_layers[0].parameters()).is_cuda)

        if args.model_name_or_path:
            if 'pororo' in args.model_name_or_path:
                config_update.stage2.vocab_size_txt = config_update.stage2.vocab_size_txt + 9
            elif 'flintstones' in args.model_name_or_path:
                config_update.stage2.vocab_size_txt = config_update.stage2.vocab_size_txt + 7
            model = cls(config_update)
            model_dir = os.path.dirname(args.model_name_or_path)
            print(model_dir)
            model.tokenizer = build_tokenizer(model_dir,
                                              context_length=model.config_dataset.context_length,
                                              lowercase=True,
                                              dropout=None)
            print("Loaded tokenizer from finetuned checkpoint")
            print(model.cross_attention_idxs)
            print("Loading model from pretrained checkpoint %s" % args.model_name_or_path)
            # model.from_ckpt(args.model_name_or_path)
            try:
                model.load_state_dict(torch.load(args.model_name_or_path)['state_dict'])
            except KeyError:
                model.load_state_dict(torch.load(args.model_name_or_path)['model_state_dict'])
        else:
            model = cls(config_update)
            print(model.cross_attention_idxs)
            print("Loading models from checkpoint %s" % path)
            model.stage1.from_ckpt(os.path.join(path, 'stage1_last.ckpt'))
            model.stage2.from_ckpt(os.path.join(path, 'stage2_last.ckpt'))

            model.tokenizer = build_tokenizer(os.path.join(path, 'tokenizer'),
                                              context_length=model.config_dataset.context_length,
                                              lowercase=True,
                                              dropout=None)


        return model, config_update


    def init_cross_attention(self, cross_attention_layers, hparams):
        self.cross_attention_idxs = cross_attention_layers
        self.cross_attention_layers = [CrossAttentionLayer(ctx_len=hparams.ctx_len_img + hparams.ctx_len_txt,
                                         embed_dim=hparams.embed_dim,
                                         n_heads=hparams.n_heads,
                                         attn_bias=hparams.attn_bias,
                                         resid_pdrop=hparams.resid_pdrop,
                                         attn_pdrop=hparams.attn_pdrop) for i in cross_attention_layers]

    def get_prompt_p5(self, bsz=None, eval=False):
        input_tokens = self.input_tokens.unsqueeze(0).expand(bsz, -1).to(self.device)
        temp_control = self.wte(input_tokens)
        past_key_values = self.control_trans(temp_control) #bsz, seqlen, layer*emb
        if not eval:
            past_key_values = self.dropout(past_key_values)
        return past_key_values

    def forward(self,
                images: torch.FloatTensor,
                src_images: Optional[torch.FloatTensor],
                texts: Optional[torch.LongTensor],
                sent_embeds: Optional[torch.FloatTensor],
                **kwargs,
        ):

        # print(images.shape, src_images.shape, texts.shape, sent_embeds.shape)

        B, L, C, H, W = images.shape
        images = images.view(B*L, C, H, W)
        src_images = src_images.unsqueeze(1).expand(-1, L, -1, -1, -1).reshape(B*L, C, H, W)
        sent_embeds = self.story_block(self.story_linear(sent_embeds)).view(B * L, -1).unsqueeze(1)
        texts = texts.view(B * L, -1)

        #{"input_ids": batch, "labels": labels, 'src_attn': src_attn, 'tgt_attn':tgt_attn, 'src':src}

        with torch.no_grad():
            with autocast(enabled=False):
                codes = self.stage1.get_codes(images).detach()
                src_codes = self.stage1.get_codes(src_images).detach()

        B, C, H, W = images.shape

        if self.config.story.prompt:
            prompt = self.get_prompt(bsz=B)
            prompt = torch.cat([prompt, sent_embeds], dim=1)
        else:
            prompt = sent_embeds

        # dim = 0 for full-model finetuning??
        pos_enc_prompt = get_positional_encoding(torch.arange(prompt.shape[1]).long().unsqueeze(0).expand(B, -1).to(self.device),
                                                 mode='1d')

        pos_enc_tokens = get_positional_encoding(texts, mode='1d')
        codes = codes.clone().detach()
        pos_enc_code = get_positional_encoding(codes, mode='1d')
        src_codes = src_codes.clone().detach()
        src_pos_enc_code = get_positional_encoding(src_codes, mode='1d')
        # codes = codes.unsqueeze(-1)
        # pos_enc_code = pos_enc_code.unsqueeze(-1)
        # print(images.shape, codes.shape, texts.shape)
        if self.config.story.condition:
            logits_img, logits_txt = self.stage2.forward_with_context(codes, texts,
                                                                      pos_enc_code, pos_enc_tokens, src_codes, src_pos_enc_code,
                                                                      self.cross_attention_idxs, self.cross_attention_layers,
                                                                      prompt=prompt, pos_prompt=pos_enc_prompt)
        else:
            logits_img, logits_txt = self.stage2(codes, texts, pos_enc_code, pos_enc_tokens, prompt=prompt,
                                                 pos_prompt=pos_enc_prompt)

        # print(logits_img.shape, logits_txt.shape, codes.shape, texts.shape)
        return logits_img, logits_txt, codes

    @torch.no_grad()
    def sampling(self,
                 tokens: torch.LongTensor,
                 source: torch.FloatTensor,
                 sent_embeds: torch.FloatTensor,
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True,
                 labels=None,
                 prompt = None) -> torch.FloatTensor:

        self.stage1.eval()
        self.stage2.eval()

        if type(tokens) == str:
            tokens = self.tokenizer.encode(tokens)
            tokens = torch.LongTensor(tokens.ids)

        # tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        # Check if the encoding works as intended
        # print(self.tokenizer.decode_batch(tokens.tolist(), skip_special_tokens=True)[0])

        tokens = tokens.to(device)
        source = source.to(device)

        # print(tokens.shape, sent_embeds.shape, prompt.shape)
        B, L, _ = sent_embeds.shape
        sent_embeds = self.story_block(self.story_linear(sent_embeds)).view(B * L, -1).unsqueeze(1)
        if prompt is not None:
            prompt = torch.cat([prompt, sent_embeds], dim=1)
        else:
            prompt = sent_embeds
        pos_enc_prompt = get_positional_encoding(torch.arange(prompt.shape[1]).long().unsqueeze(0).expand(B*L, -1).to(self.device), mode='1d')

        with autocast(enabled=False):
            src_codes = self.stage1.get_codes(source).detach()
        src_codes = torch.repeat_interleave(src_codes, self.config.story.story_len, dim=0)
        print(tokens.shape, src_codes.shape, prompt.shape)
        if self.config.story.condition:
            codes = sampling_conditional(self.stage2,
                                         self.cross_attention_idxs,
                                         self.cross_attention_layers,
                                         tokens,
                                         src_codes,
                                         top_k=top_k,
                                         top_p=top_p,
                                         softmax_temperature=softmax_temperature,
                                         use_fp16=use_fp16,
                                         prompt=prompt,
                                         pos_prompt=pos_enc_prompt)
        else:
            codes = sampling(self.stage2,
                             tokens,
                             top_k=top_k,
                             top_p=top_p,
                             softmax_temperature=softmax_temperature,
                             use_fp16=use_fp16,
                             prompt=prompt,
                             pos_prompt=pos_enc_prompt)

        codes = codes.view(self.config.story.story_len, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 256, 256]
        return pixels

    @torch.no_grad()
    def sampling_batch(self,
                 tokens: torch.LongTensor,
                 source: torch.FloatTensor,
                 sent_embeds: torch.FloatTensor,
                 top_k: int = 256,
                 top_p: Optional[float] = None,
                 softmax_temperature: float = 1.0,
                 num_candidates: int = 96,
                 device: str = 'cuda:0',
                 use_fp16: bool = True,
                 labels=None,
                 prompt=None, n_candidates=1) -> torch.FloatTensor:

        self.stage1.eval()
        self.stage2.eval()

        if type(tokens) == str:
            tokens = self.tokenizer.encode(tokens)
            tokens = torch.LongTensor(tokens.ids)

        # tokens = torch.repeat_interleave(tokens.unsqueeze(0), num_candidates, dim=0)

        # Check if the encoding works as intended
        # print(self.tokenizer.decode_batch(tokens.tolist(), skip_special_tokens=True)[0])

        tokens = tokens.to(device)
        source = source.to(device)

        # print(tokens.shape, sent_embeds.shape, prompt.shape)
        B, L, _ = sent_embeds.shape
        sent_embeds = self.story_block(self.story_linear(sent_embeds)).view(B * L, -1).unsqueeze(1)
        if prompt is not None:
            prompt = torch.cat([prompt, sent_embeds], dim=1)
        else:
            prompt = sent_embeds
        pos_enc_prompt = get_positional_encoding(
            torch.arange(prompt.shape[1]).long().unsqueeze(0).expand(B * L, -1).to(self.device), mode='1d')

        with autocast(enabled=False):
            src_codes = self.stage1.get_codes(source).detach()

        # repeat inputs to adjust to n_candidates and story length
        src_codes = torch.repeat_interleave(src_codes, self.config.story.story_len * n_candidates, dim=0)
        prompt = prompt.repeat(n_candidates, 1, 1)
        pos_enc_prompt = pos_enc_prompt.repeat(n_candidates, 1)
        tokens = tokens.repeat(n_candidates, 1)
        print(tokens.shape, src_codes.shape, prompt.shape, pos_enc_prompt.shape)
        if self.config.story.condition:
            codes = sampling_conditional(self.stage2,
                                         self.cross_attention_idxs,
                                         self.cross_attention_layers,
                                         tokens,
                                         src_codes,
                                         top_k=top_k,
                                         top_p=top_p,
                                         softmax_temperature=softmax_temperature,
                                         use_fp16=use_fp16,
                                         prompt=prompt,
                                         pos_prompt=pos_enc_prompt)
        else:
            codes = sampling(self.stage2,
                             tokens,
                             top_k=top_k,
                             top_p=top_p,
                             softmax_temperature=softmax_temperature,
                             use_fp16=use_fp16,
                             prompt=prompt,
                             pos_prompt=pos_enc_prompt)

        codes = codes.view(self.config.story.story_len * n_candidates, 16, 16)  # [B, 16, 16]
        print(codes.shape)
        pixels = torch.clamp(self.stage1.decode_code(codes) * 0.5 + 0.5, 0, 1)  # [B, 3, 256, 256]
        print(pixels.shape)
        return pixels.view(n_candidates, self.config.story.story_len, pixels.shape[-3], pixels.shape[-2], pixels.shape[-1])


    @torch.no_grad()
    def predict_step(self, batch, batch_idx, return_images=False):
        orig_images, texts = batch
        # concatenate the list of prompts (split by n_head) for better downstream processing

        # extra for checks
        logits_img, logits_txt, codes = self(orig_images, texts)
        pred = torch.argmax(logits_img.view(-1, logits_img.shape[-1]), dim=-1)
        bs = orig_images.shape[0]
        pred = pred.view(bs, 16, 16)  # [B, 16, 16]
        pixels = torch.clamp(self.stage1.decode_code(pred) * 0.5 + 0.5, 0, 1).cpu().numpy()  # [B, 256, 256]
        pixels = np.transpose(pixels, (0, 2, 3, 1))

        prompt = self.get_prompt(bsz=5, eval=True)

        images = []
        for t in texts:
            pixels = self.sampling(t, prompt, top_k=64, num_candidates=5, labels=codes[i]).cpu().numpy()
            pixels = np.transpose(pixels, (0, 2, 3, 1))
            images.append(pixels)
            # images.extend([p for p in pixels])
            # print([i.shape for i in images])

        if return_images:
            return images
        else:
            save_image(orig_images, pixels, './out/images/pororo_story', batch_idx+10)
            save_image(orig_images, images, './out/images/pororo_story', batch_idx)