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

# %% auto 0
__all__ = ['DiffusionPipeline_attnPadded', 'DiffusionPipeline_Compilation']

# %% ../../src/pipeline/diffusion_pipeline_special.ipynb 2
from ..imports import *
from .diffusion_pipeline import DiffusionPipeline

# %% ../../src/pipeline/diffusion_pipeline_special.ipynb 3
class DiffusionPipeline_attnPadded(DiffusionPipeline): 
    """A special `DiffusionPipeline` with attention masking."""
    def train_step(self, data, **kwargs): 
        latents, y, key_padding_mask_list = data                
        b, s, t = latents.shape          
        
        #start async memcpy
        loss_mask = (key_padding_mask_list[0].to(self.device, non_blocking=self.non_blocking)>-1.0).float().unsqueeze(1)
      
        shaped_mask = []
        for key_padding_mask in key_padding_mask_list:
            key_padding_mask = key_padding_mask.to(self.device, non_blocking=self.non_blocking)    
            key_padding_mask = key_padding_mask.reshape((b, -1)) #from [b, s, t] to [b, -1] aka [N, L]           
            shaped_mask.append(key_padding_mask)         
        
        latents = latents.to(self.device, non_blocking=self.non_blocking)  
        latents = self.model.embedd_clrs(latents)                   #this is only new tensor
        self.scheduler.to_device(self.device, non_blocking=self.non_blocking) 
            
        #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, key_padding_mask=shaped_mask)
            
        #comp mse
        loss = self.loss_fn(eps*loss_mask, noise*loss_mask)
        # loss = self.loss_fn(eps, noise)
        
        #log the loss
        return loss

# %% ../../src/pipeline/diffusion_pipeline_special.ipynb 4
class DiffusionPipeline_Compilation(DiffusionPipeline):   
    """A special `DiffusionPipeline` that accounts for unitary conditions, i.e. compilation."""
    
    #------------------------------------
    
    @torch.no_grad()
    def __call__(self, latents, c, U, g, no_bar=False):
        
        latents = latents.to(self.device)
        c       = c.to(self.device)
        U       = U.to(self.device)
        
        return self.denoising(latents, c=c, U=U, enable_guidance=True, g=g, no_bar=no_bar)

    #------------------------------------
    
    def get_guidance_U(self, U, enable_guidance):
        if not exists(U): return U      
        U = U.to(self.device)                
        if enable_guidance:             
            u = torch.zeros_like(U, device=self.device)       
            U = torch.cat([u, U])            
        U = U.type(torch.float32)
        return U
    
    @torch.no_grad()
    def denoising(self, latents, c, U, enable_guidance, g, no_bar=False, return_predicted_x0=False):       
        U = self.get_guidance_U(U, enable_guidance)  
        # self.unitary_encoder.eval()   
        return super().denoising(latents, c, enable_guidance, g, no_bar=no_bar, return_predicted_x0=return_predicted_x0, U=U)

    def denoising_step(self, latents: torch.Tensor, ts: Union[int, torch.IntTensor], c_emb: torch.Tensor=None, enable_guidance=False, g=7.5, U: torch.Tensor=None):    
        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, U=U).chunk(2) 
            
            eps = self.CFG(eps_u, eps_c, g)
                    
        else:
            eps = self.model(latents, ts, c_em, U=U)  
                 
        x = self.scheduler.step(eps, ts, latents)      
        return x.prev_sample, x.pred_original_sample
    
    #------------------------------------

    def train_step(self, data, **kwargs): 
        latents, y, U = 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)  
        U = U.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_y = self.scheduler.unsqueeze_vector_to_shape(rnd, y.shape)   # e.g. [b, 1, 1]            
            y   = y * rnd_y + (1-rnd_y) * self.empty_token_fn(y)
            
            U   = U * self.scheduler.unsqueeze_vector_to_shape(rnd, U.shape)
            
        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, U=U)
            
        #comp mse
        loss = self.loss_fn(eps, noise)
        
        #log the loss
        return loss
