# AUTOGENERATED! DO NOT EDIT! File to edit: ../../src/pipeline/diffusion_pipeline.ipynb.

# %% auto 0
__all__ = ['DiffusionPipeline']

# %% ../../src/pipeline/diffusion_pipeline.ipynb 2
from ..imports import *
from ..scheduler.scheduler import Scheduler
from .pipeline import Pipeline
from ..config_loader import *
from ..models.config_model import Config_Model

from huggingface_hub import snapshot_download

# %% ../../src/pipeline/diffusion_pipeline.ipynb 3
class DiffusionPipeline(Pipeline):   
    """A `Pipeline` for diffusion models. Implements train and inference functions. Diffusion parameters are defined inside a `Scheduler` object."""
    non_blocking = False
    
    def __init__(self, 
                 scheduler: Scheduler,   
                 model: nn.Module,
                 text_encoder: nn.Module,
                 device: torch.device,
                 enable_guidance_train = True,
                 guidance_train_p = 0.1,
                 cached_text_enc = True
                ):    
        super().__init__(model, device)
        self.scheduler = scheduler
        self.scheduler.to_device(device)    
        
        self.text_encoder = text_encoder
        self.text_encoder.eval()
        
        self.enable_guidance_train = enable_guidance_train           
        self.guidance_train_p      = guidance_train_p
           
        self.cached_text_enc = cached_text_enc        
        self.empty_token     = self.text_encoder.empty_token
        
        if cached_text_enc:                       
            def cached_empty_token_fn(c):
                if   c.dim() == 1: return self.text_encoder.cached_empty_token_index  # yields then a list of ints       
                elif c.dim() == 2: return self.empty_token.expand(c.shape)            # tokenized input      
                else: raise NotImplementedError("")
            
            self.empty_token_fn = cached_empty_token_fn    
                        
        else:
            self.empty_token_fn = lambda c: self.empty_token.expand(c.shape) # for own clip  

    #------------------------------------
         
    add_config = {}
        
    def params_config(self, save_path: str):         
        params_config = {}
                
        params_config["scheduler"]    = self.scheduler.get_config()
        params_config["model"]        = self.model.get_config(save_path=save_path+"model.pt")
        params_config["text_encoder"] = self.text_encoder.get_config(save_path=save_path+"text_encoder.pt")
        
        params_config["device"]                = str(self.device)
        params_config["enable_guidance_train"] = self.enable_guidance_train
        params_config["guidance_train_p"]      = self.guidance_train_p
        params_config["cached_text_enc"]       = self.cached_text_enc
        params_config["add_config"]            = self.add_config
        
        return params_config
    
    def store_pipeline(self, config_path: str, save_path: str):
        super().store_pipeline(config_path, save_path)
        config = self.get_config(save_path)
        save_dict_yaml(config, config_path+"config.yaml")
               
        #only store weights of these submodels
        self.model.store_model(config_path=None, save_path=save_path+"model.pt")
        self.text_encoder.store_model(config_path=None, save_path=save_path+"text_encoder.pt")
    
    @staticmethod
    def from_config_file(config_path, device: torch.device):    
        config = load_config(config_path+"config.yaml")   
        config = config_to_dict(config)

        if exists(device):
            config["params"]["device"]                        = device
            config["params"]["scheduler"]["params"]["device"] = device
        
        config["params"]["scheduler"] = instantiate_from_config(config["params"]["scheduler"])
 
        model_path                = config_path+"model.pt" if config["params"]["model"]["save_path"] is None else config["params"]["model"]["save_path"]
        config["params"]["model"] = Config_Model.from_config(config["params"]["model"], device, model_path)

        config["params"]["text_encoder"] = Config_Model.from_config(config["params"]["text_encoder"], device, config["params"]["text_encoder"]["save_path"])   
        add_config = config["params"].pop("add_config", None)

        pipeline = instantiate_from_config(config)
        
        if exists(pipeline.add_config):
            pipeline.gate_pool  = [gate for gate in add_config["dataset"]["params"]["gate_pool"]] 
            pipeline.add_config = add_config
        
        return pipeline


    @classmethod
    def from_pretrained(cls, repo_id: str, device: torch.device, **kwargs):  
        """Load a model pipeline directly from Huggingface."""
        model_path = snapshot_download(repo_id=repo_id, repo_type="model", allow_patterns=["*.pt", "*.yaml", "*.safetensors"], **kwargs) 
        pipeline   = cls.from_config_file(model_path+"/", device)  
        return pipeline
    
    #------------------------------------
    # Inference functions
    
    @torch.no_grad()
    def __call__(self, latents=None, c=None, seed=None, timesteps=None, no_bar=False, enable_guidance=True, g=7.5):        
        if exists(seed):      torch.manual_seed(seed)
        if exists(timesteps): self.scheduler.set_timesteps(self.timesteps)
            
        latents = latents.to(self.device)                  
        x0      = self.denoising(latents, c=c, no_bar=no_bar, enable_guidance=enable_guidance, g=g)  
        
        return x0
    
    @torch.no_grad()
    def latent_filling(self, org_latents: torch.Tensor, mask: torch.Tensor, c=None, enable_guidance=True, g=7.5, 
                       t_start_index=0, no_bar=False, return_predicted_x0=False, **kwargs):
        """mask: area with ones is going to be filled"""
        if   mask.dim() == 4: assert list(org_latents.shape) == list(mask.shape)     # diff mask per sample and channel
        elif mask.dim() == 3: assert list(org_latents.shape[1:]) == list(mask.shape) # diff mask per channel
        elif mask.dim() == 2: assert list(org_latents.shape[2:]) == list(mask.shape) # all same mask
        else:                 raise RuntimeError("")
    
        self.model.eval()
        self.text_encoder.eval()       
        self.scheduler.to_device(self.device)
                 
        c_emb = self.prepare_c_emb(c, enable_guidance, **kwargs)
    
        org_latents = org_latents.to(self.device, non_blocking=self.non_blocking)

        #---
        # ch to prep latent, init_image
        
        noise   = torch.randn(org_latents.shape, device=self.device)   
        
        if t_start_index > 0: latents = self.scheduler.add_noise(org_latents, noise, self.scheduler.timesteps[t_start_index])
        else:                 latents = noise
        
        #---

        timesteps = self.scheduler.timesteps[t_start_index:]
        
        if return_predicted_x0: predicted_x0 = list()
     
        for i, t in enumerate(tqdm(timesteps, disable=no_bar)):
            timesteps = (torch.ones((1)) * t).type(torch.int64).to(self.device, non_blocking=self.non_blocking)        

            #----------------------------------
            #gen the fill in part
            latents, x0 = self.denoising_step(latents, timesteps, c_emb=c_emb, enable_guidance=enable_guidance, g=g)
            
            #----------------------------------
            #create the part we don't fill 

            noisy_org_latents = org_latents
            
            if i < len(timesteps) - 1:            #else we are finished with denoising
                noise_timestep = timesteps[i + 1]
                # noise = torch.randn(org_latents.shape, device=self.device)     
                noisy_org_latents = self.scheduler.add_noise(noisy_org_latents, noise, noise_timestep) 
                
            #combine
            latents = (1.0-mask) * noisy_org_latents + mask * latents  

            #----------------------------------   
            if return_predicted_x0: 
                x0      = (1.0-mask) *  org_latents + mask * fill_x0  
                predicted_x0.append(x0)
        
        if return_predicted_x0: return latents.cpu(), predicted_x0     
        return latents.cpu()
    
    #------------------------------------
    # Helper functions
    
    def get_guidance_condition(self, c, enable_guidance):
        if not exists(c): return c      
        c = c.to(self.device)                
        if enable_guidance:             
            u = self.empty_token_fn(c).to(self.device)        
            c = torch.cat([u, c])            
        c = c.type(torch.int64) 
        return c

    def prepare_c_emb(self, c, enable_guidance, **kwargs):
        c     = self.get_guidance_condition(c, enable_guidance) 
        c_emb = self.text_encoder(c, pool=False)
        return c_emb
        
    @torch.no_grad()
    def denoising(self, latents: torch.Tensor, c=None, enable_guidance=True, g=7.5, t_start_index=0, no_bar=False, return_predicted_x0=False, **kwargs):
        self.model.eval()
        self.text_encoder.eval()
        self.scheduler.to_device(self.device)
    
        c_emb = self.prepare_c_emb(c, enable_guidance, **kwargs)
        
        latents = latents.to(self.device, non_blocking=self.non_blocking)
        
        if return_predicted_x0: predicted_x0 = list()
        
        for i, t in enumerate(tqdm(self.scheduler.timesteps[t_start_index:], disable=no_bar)):
            timesteps = (torch.ones((1)) * t).type(torch.int64).to(self.device, non_blocking=self.non_blocking)
                   
            latents, x0 = self.denoising_step(latents, timesteps, c_emb=c_emb, enable_guidance=enable_guidance, g=g, **kwargs)
        
            if return_predicted_x0: predicted_x0.append(x0.cpu())
        
        if return_predicted_x0: return latents.cpu(), predicted_x0     
        return latents.cpu()
  
    # @torch.no_grad()
    def denoising_step(self, latents: torch.Tensor, ts: Union[int, torch.IntTensor], c_emb: torch.Tensor=None, enable_guidance=True, g=7.5, **kwargs):    
        if enable_guidance:
            x = torch.cat([latents] * 2)     #uses batch layer combine here
            
            if ts.numel() > 1: chunk_ts = torch.cat([ts] * 2)
            else:              chunk_ts = ts
                
            eps_u, eps_c = self.model(x, chunk_ts, c_emb).chunk(2) 
            
            eps = self.CFG(eps_u, eps_c, g)
                    
        else:
            eps = self.model(latents, ts, c_emb)  
                 
        x = self.scheduler.step(eps, ts, latents)      
        return x.prev_sample, x.pred_original_sample
     
    guidance_sample_mode = "rescaled" # one of: normal, fastai, rescaled
        
    def CFG(self, eps_u, eps_c, g):
        """Apply Classifier-free-guidance sampling"""
        dim = list(range(1, eps_u.dim())) # reduce all but batches

        if self.guidance_sample_mode == "normal":   # from https://arxiv.org/pdf/2207.12598.pdf, w=g+1    
            eps = eps_u + g * (eps_c-eps_u)               

        elif self.guidance_sample_mode == "fastai":  # from fastAi less 11
            eps = eps_u + g*(eps_c-eps_u) * torch.linalg.vector_norm(eps_u, dim=dim, keepdim=True) / torch.linalg.vector_norm(eps_c-eps_u, dim=dim, keepdim=True)    
            eps = eps * torch.linalg.vector_norm(eps_u, dim=dim, keepdim=True) / torch.linalg.vector_norm(eps, dim=dim, keepdim=True)

        elif self.guidance_sample_mode == "rescaled": # from https://arxiv.org/pdf/2305.08891.pd
            phi = 0.7

            eps_cfg    = eps_u + g * (eps_c-eps_u)              
            eps_scaled = eps_cfg * eps_c.std(dim=dim, keepdim=True) / eps_cfg.std(dim=dim, keepdim=True)
            eps        = phi * eps_scaled + (1.0-phi) * eps_cfg

        else: raise NotImplementedError("Use as guidance_sample_mode one of: 'normal', 'fastai' or 'rescaled'")

        return eps
        
    #------------------------------------
    # Training functions

    def train_on_epoch(self, data_loader: DataLoader, train=True):   
        self.scheduler.to_device(self.device, non_blocking=self.non_blocking)    
        super().train_on_epoch(data_loader, train)

    #@torch.autocast(device_type=device.type)
    def train_step(self, data, **kwargs): 
        latents, y = data                
        b, s, t = latents.shape          
        
        #start async memcpy
        latents = latents.to(self.device, non_blocking=self.non_blocking)  
        latents = self.model.embedd_clrs(latents)                   #this is only new tensor
               
        #do the cond embedding with CLIP                     
        y = y.to(self.device, non_blocking=self.non_blocking)  
        
        if self.enable_guidance_train: 
            rnd = torch.rand((b,), device=self.device)            
            rnd = (rnd > self.guidance_train_p).type(torch.int64)          # todo: change to bernoulli dist fn
            rnd = self.scheduler.unsqueeze_vector_to_shape(rnd, y.shape)   # e.g. [b, 1, 1]            
            y   = y * rnd + (1-rnd) * self.empty_token_fn(y)
               
        y_emb = self.text_encoder(y, pool=False)
                          
        #sample timesteps
        timesteps = torch.randint(low=0, high=self.scheduler.num_train_timesteps, size=(b,), device=self.device, dtype=torch.int64)

        #forward noising    
        noise = torch.randn(latents.shape, device=self.device)     
        noisy_latents = self.scheduler.add_noise(latents, noise, timesteps)

        #predict eps
        eps = self.model(noisy_latents, timesteps, y_emb)
            
        #comp mse
        loss = self.loss_fn(eps, noise)
        
        #log the loss
        return loss
