'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
from ..attention import BasicTransformerBlock, TemporalBasicTransformerBlock
from ...configuration_utils import ConfigMixin, register_to_config
from ..embeddings import TimestepEmbedding, Timesteps
from ..modeling_utils import ModelMixin
from typing import Any, Dict, Optional
from dataclasses import dataclass
from ..resnet import AlphaBlender
from ...utils import BaseOutput
from torch import nn
import torch
@dataclass
class TransformerTemporalModelOutput(BaseOutput):
    """Args:"""
    sample: torch.Tensor
class TransformerTemporalModel(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(self, num_attention_heads: int=16, attention_head_dim: int=88, in_channels: Optional[int]=None, out_channels: Optional[int]=None, num_layers: int=1,
    dropout: float=0.0, norm_num_groups: int=32, cross_attention_dim: Optional[int]=None, attention_bias: bool=False, sample_size: Optional[int]=None,
    activation_fn: str='geglu', norm_elementwise_affine: bool=True, double_self_attention: bool=True, positional_embeddings: Optional[str]=None, num_positional_embeddings: Optional[int]=None):
        super().__init__()
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        inner_dim = num_attention_heads * attention_head_dim
        self.in_channels = in_channels
        self.norm = torch.nn.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-06, affine=True)
        self.proj_in = nn.Linear(in_channels, inner_dim)
        self.transformer_blocks = nn.ModuleList([BasicTransformerBlock(inner_dim, num_attention_heads, attention_head_dim, dropout=dropout, cross_attention_dim=cross_attention_dim,
        activation_fn=activation_fn, attention_bias=attention_bias, double_self_attention=double_self_attention, norm_elementwise_affine=norm_elementwise_affine,
        positional_embeddings=positional_embeddings, num_positional_embeddings=num_positional_embeddings) for d in range(num_layers)])
        self.proj_out = nn.Linear(inner_dim, in_channels)
    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.LongTensor]=None, timestep: Optional[torch.LongTensor]=None, class_labels: torch.LongTensor=None,
    num_frames: int=1, cross_attention_kwargs: Optional[Dict[str, Any]]=None, return_dict: bool=True) -> TransformerTemporalModelOutput:
        """Returns:"""
        batch_frames, channel, height, width = hidden_states.shape
        batch_size = batch_frames // num_frames
        residual = hidden_states
        hidden_states = hidden_states[None, :].reshape(batch_size, num_frames, channel, height, width)
        hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
        hidden_states = self.norm(hidden_states)
        hidden_states = hidden_states.permute(0, 3, 4, 2, 1).reshape(batch_size * height * width, num_frames, channel)
        hidden_states = self.proj_in(hidden_states)
        for block in self.transformer_blocks: hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states,
        timestep=timestep, cross_attention_kwargs=cross_attention_kwargs, class_labels=class_labels)
        hidden_states = self.proj_out(hidden_states)
        hidden_states = hidden_states[None, None, :].reshape(batch_size, height, width, num_frames, channel).permute(0, 3, 4, 1, 2).contiguous()
        hidden_states = hidden_states.reshape(batch_frames, channel, height, width)
        output = hidden_states + residual
        if not return_dict: return (output,)
        return TransformerTemporalModelOutput(sample=output)
class TransformerSpatioTemporalModel(nn.Module):
    def __init__(self, num_attention_heads: int=16, attention_head_dim: int=88, in_channels: int=320, out_channels: Optional[int]=None, num_layers: int=1, cross_attention_dim: Optional[int]=None):
        super().__init__()
        self.num_attention_heads = num_attention_heads
        self.attention_head_dim = attention_head_dim
        inner_dim = num_attention_heads * attention_head_dim
        self.inner_dim = inner_dim
        self.in_channels = in_channels
        self.norm = torch.nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-06)
        self.proj_in = nn.Linear(in_channels, inner_dim)
        self.transformer_blocks = nn.ModuleList([BasicTransformerBlock(inner_dim, num_attention_heads, attention_head_dim, cross_attention_dim=cross_attention_dim) for d in range(num_layers)])
        time_mix_inner_dim = inner_dim
        self.temporal_transformer_blocks = nn.ModuleList([TemporalBasicTransformerBlock(inner_dim, time_mix_inner_dim, num_attention_heads,
        attention_head_dim, cross_attention_dim=cross_attention_dim) for _ in range(num_layers)])
        time_embed_dim = in_channels * 4
        self.time_pos_embed = TimestepEmbedding(in_channels, time_embed_dim, out_dim=in_channels)
        self.time_proj = Timesteps(in_channels, True, 0)
        self.time_mixer = AlphaBlender(alpha=0.5, merge_strategy='learned_with_images')
        self.out_channels = in_channels if out_channels is None else out_channels
        self.proj_out = nn.Linear(inner_dim, in_channels)
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor]=None, image_only_indicator: Optional[torch.Tensor]=None, return_dict: bool=True):
        """Returns:"""
        batch_frames, _, height, width = hidden_states.shape
        num_frames = image_only_indicator.shape[-1]
        batch_size = batch_frames // num_frames
        time_context = encoder_hidden_states
        time_context_first_timestep = time_context[None, :].reshape(batch_size, num_frames, -1, time_context.shape[-1])[:, 0]
        time_context = time_context_first_timestep[:, None].broadcast_to(batch_size, height * width, time_context.shape[-2], time_context.shape[-1])
        time_context = time_context.reshape(batch_size * height * width, -1, time_context.shape[-1])
        residual = hidden_states
        hidden_states = self.norm(hidden_states)
        inner_dim = hidden_states.shape[1]
        hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(batch_frames, height * width, inner_dim)
        hidden_states = self.proj_in(hidden_states)
        num_frames_emb = torch.arange(num_frames, device=hidden_states.device)
        num_frames_emb = num_frames_emb.repeat(batch_size, 1)
        num_frames_emb = num_frames_emb.reshape(-1)
        t_emb = self.time_proj(num_frames_emb)
        t_emb = t_emb.to(dtype=hidden_states.dtype)
        emb = self.time_pos_embed(t_emb)
        emb = emb[:, None, :]
        for block, temporal_block in zip(self.transformer_blocks, self.temporal_transformer_blocks):
            if torch.is_grad_enabled() and self.gradient_checkpointing: hidden_states = torch.utils.checkpoint.checkpoint(block, hidden_states, None, encoder_hidden_states, None, use_reentrant=False)
            else: hidden_states = block(hidden_states, encoder_hidden_states=encoder_hidden_states)
            hidden_states_mix = hidden_states
            hidden_states_mix = hidden_states_mix + emb
            hidden_states_mix = temporal_block(hidden_states_mix, num_frames=num_frames, encoder_hidden_states=time_context)
            hidden_states = self.time_mixer(x_spatial=hidden_states, x_temporal=hidden_states_mix, image_only_indicator=image_only_indicator)
        hidden_states = self.proj_out(hidden_states)
        hidden_states = hidden_states.reshape(batch_frames, height, width, inner_dim).permute(0, 3, 1, 2).contiguous()
        output = hidden_states + residual
        if not return_dict: return (output,)
        return TransformerTemporalModelOutput(sample=output)
'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
