


'''simple docstring'''




from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..utils import maybe_allow_in_graph
from .activations import get_activation
from .attention_processor import Attention
from .embeddings import CombinedTimestepLabelEmbeddings

@maybe_allow_in_graph
class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      Any    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int=0.0    ,      UpperCAmelCase__       :      Optional[int] = None    ,      UpperCAmelCase__       :      str = "geglu"    ,      UpperCAmelCase__       :      Optional[int] = None    ,      UpperCAmelCase__       :      bool = False    ,      UpperCAmelCase__       :      bool = False    ,      UpperCAmelCase__       :      bool = False    ,      UpperCAmelCase__       :      bool = False    ,      UpperCAmelCase__       :      bool = True    ,      UpperCAmelCase__       :      str = "layer_norm"    ,      UpperCAmelCase__       :      bool = False    ,      )				->  List[Any]:
        super().__init__()
        lowerCAmelCase     =			only_cross_attention

        lowerCAmelCase     =			(num_embeds_ada_norm is not None) and norm_type == 'ada_norm_zero'
        lowerCAmelCase     =			(num_embeds_ada_norm is not None) and norm_type == 'ada_norm'

        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
           raise ValueError(
               F'''`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to'''
               F''' define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}.'''    )

        # Define 3 blocks. Each block has its own normalization layer.
        # 1. Self-Attn
        if self.use_ada_layer_norm:
           lowerCAmelCase     =			AdaLayerNorm(UpperCAmelCase__    ,      UpperCAmelCase__    )
        elif self.use_ada_layer_norm_zero:
           lowerCAmelCase     =			AdaLayerNormZero(UpperCAmelCase__    ,      UpperCAmelCase__    )
        else:
           lowerCAmelCase     =			nn.LayerNorm(UpperCAmelCase__    ,      elementwise_affine=UpperCAmelCase__    )
        lowerCAmelCase     =			Attention(
            query_dim=UpperCAmelCase__    ,      heads=UpperCAmelCase__    ,      dim_head=UpperCAmelCase__    ,      dropout=UpperCAmelCase__    ,      bias=UpperCAmelCase__    ,      cross_attention_dim=cross_attention_dim if only_cross_attention else None    ,      upcast_attention=UpperCAmelCase__    ,      )

        # 2. Cross-Attn
        if cross_attention_dim is not None or double_self_attention:
           # We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
           # I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
           # the second cross attention block.
           lowerCAmelCase     =			(
               AdaLayerNorm(UpperCAmelCase__    ,      UpperCAmelCase__    )
               if self.use_ada_layer_norm
               else nn.LayerNorm(UpperCAmelCase__    ,      elementwise_affine=UpperCAmelCase__    )
           )
           lowerCAmelCase     =			Attention(
               query_dim=UpperCAmelCase__    ,      cross_attention_dim=cross_attention_dim if not double_self_attention else None    ,      heads=UpperCAmelCase__    ,      dim_head=UpperCAmelCase__    ,      dropout=UpperCAmelCase__    ,      bias=UpperCAmelCase__    ,      upcast_attention=UpperCAmelCase__    ,      )  # is self-attn if encoder_hidden_states is none
        else:
           lowerCAmelCase     =			None
           lowerCAmelCase     =			None

        # 3. Feed-forward
        lowerCAmelCase     =			nn.LayerNorm(UpperCAmelCase__    ,      elementwise_affine=UpperCAmelCase__    )
        lowerCAmelCase     =			FeedForward(UpperCAmelCase__    ,      dropout=UpperCAmelCase__    ,      activation_fn=UpperCAmelCase__    ,      final_dropout=UpperCAmelCase__    )

        # let chunk size default to None
        lowerCAmelCase     =			None
        lowerCAmelCase     =			0





     def 							__UpperCAmelCase      ( self       :      Dict    ,      UpperCAmelCase__       :      Optional[int]    ,      UpperCAmelCase__       :      int    )				->  Optional[int]:
        # Sets chunk feed-forward
        lowerCAmelCase     =			chunk_size
        lowerCAmelCase     =			dim





     def 							__UpperCAmelCase      ( self       :      Any    ,      UpperCAmelCase__       :      torch.FloatTensor    ,      UpperCAmelCase__       :      Optional[torch.FloatTensor] = None    ,      UpperCAmelCase__       :      Optional[torch.FloatTensor] = None    ,      UpperCAmelCase__       :      Optional[torch.FloatTensor] = None    ,      UpperCAmelCase__       :      Optional[torch.LongTensor] = None    ,      UpperCAmelCase__       :      Dict[str, Any] = None    ,      UpperCAmelCase__       :      Optional[torch.LongTensor] = None    ,      )				->  Any:
        # Notice that normalization is always applied before the real computation in the following blocks.
        # 1. Self-Attention
        if self.use_ada_layer_norm:
           lowerCAmelCase     =			self.norma(UpperCAmelCase__    ,      UpperCAmelCase__    )
        elif self.use_ada_layer_norm_zero:
           lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     =			self.norma(
               UpperCAmelCase__    ,      UpperCAmelCase__    ,      UpperCAmelCase__    ,      hidden_dtype=hidden_states.dtype    )
        else:
           lowerCAmelCase     =			self.norma(UpperCAmelCase__    )

        lowerCAmelCase     =			cross_attention_kwargs if cross_attention_kwargs is not None else {}

        lowerCAmelCase     =			self.attna(
            UpperCAmelCase__    ,      encoder_hidden_states=encoder_hidden_states if self.only_cross_attention else None    ,      attention_mask=UpperCAmelCase__    ,      **UpperCAmelCase__    ,      )
        if self.use_ada_layer_norm_zero:
           lowerCAmelCase     =			gate_msa.unsqueeze(1    ) * attn_output
        lowerCAmelCase     =			attn_output + hidden_states

        # 2. Cross-Attention
        if self.attna is not None:
           lowerCAmelCase     =			(
               self.norma(UpperCAmelCase__    ,      UpperCAmelCase__    ) if self.use_ada_layer_norm else self.norma(UpperCAmelCase__    )
           )

           lowerCAmelCase     =			self.attna(
               UpperCAmelCase__    ,      encoder_hidden_states=UpperCAmelCase__    ,      attention_mask=UpperCAmelCase__    ,      **UpperCAmelCase__    ,      )
           lowerCAmelCase     =			attn_output + hidden_states

        # 3. Feed-forward
        lowerCAmelCase     =			self.norma(UpperCAmelCase__    )

        if self.use_ada_layer_norm_zero:
           lowerCAmelCase     =			norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]

        if self._chunk_size is not None:
           # "feed_forward_chunk_size" can be used to save memory
           if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
              raise ValueError(
                  F'''`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`.'''    )

           lowerCAmelCase     =			norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
           lowerCAmelCase     =			torch.cat(
               [self.ff(UpperCAmelCase__    ) for hid_slice in norm_hidden_states.chunk(UpperCAmelCase__    ,      dim=self._chunk_dim    )]    ,      dim=self._chunk_dim    ,      )
        else:
           lowerCAmelCase     =			self.ff(UpperCAmelCase__    )

        if self.use_ada_layer_norm_zero:
           lowerCAmelCase     =			gate_mlp.unsqueeze(1    ) * ff_output

        lowerCAmelCase     =			ff_output + hidden_states

        return hidden_states




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      Tuple    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      Optional[int] = None    ,      UpperCAmelCase__       :      int = 4    ,      UpperCAmelCase__       :      float = 0.0    ,      UpperCAmelCase__       :      str = "geglu"    ,      UpperCAmelCase__       :      bool = False    ,      )				->  str:
        super().__init__()
        lowerCAmelCase     =			int(dim * mult    )
        lowerCAmelCase     =			dim_out if dim_out is not None else dim

        if activation_fn == "gelu":
           lowerCAmelCase     =			GELU(UpperCAmelCase__    ,      UpperCAmelCase__    )
        if activation_fn == "gelu-approximate":
           lowerCAmelCase     =			GELU(UpperCAmelCase__    ,      UpperCAmelCase__    ,      approximate='tanh'    )
        elif activation_fn == "geglu":
           lowerCAmelCase     =			GEGLU(UpperCAmelCase__    ,      UpperCAmelCase__    )
        elif activation_fn == "geglu-approximate":
           lowerCAmelCase     =			ApproximateGELU(UpperCAmelCase__    ,      UpperCAmelCase__    )

        lowerCAmelCase     =			nn.ModuleList([]    )
        # project in
        self.net.append(UpperCAmelCase__    )
        # project dropout
        self.net.append(nn.Dropout(UpperCAmelCase__    )    )
        # project out
        self.net.append(nn.Linear(UpperCAmelCase__    ,      UpperCAmelCase__    )    )
        # FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
        if final_dropout:
           self.net.append(nn.Dropout(UpperCAmelCase__    )    )





     def 							__UpperCAmelCase      ( self       :      Dict    ,      UpperCAmelCase__       :      Tuple    )				->  str:
        for module in self.net:
           lowerCAmelCase     =			module(UpperCAmelCase__    )
        return hidden_states




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      List[str]    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      str = "none"    )				->  List[Any]:
        super().__init__()
        lowerCAmelCase     =			nn.Linear(UpperCAmelCase__    ,      UpperCAmelCase__    )
        lowerCAmelCase     =			approximate





     def 							__UpperCAmelCase      ( self       :      List[str]    ,      UpperCAmelCase__       :      Optional[int]    )				->  Any:
        if gate.device.type != "mps":
           return F.gelu(UpperCAmelCase__    ,      approximate=self.approximate    )
        # mps: gelu is not implemented for float16
        return F.gelu(gate.to(dtype=torch.floataa    )    ,      approximate=self.approximate    ).to(dtype=gate.dtype    )





     def 							__UpperCAmelCase      ( self       :      Tuple    ,      UpperCAmelCase__       :      Tuple    )				->  List[str]:
        lowerCAmelCase     =			self.proj(UpperCAmelCase__    )
        lowerCAmelCase     =			self.gelu(UpperCAmelCase__    )
        return hidden_states




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      Any    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    )				->  int:
        super().__init__()
        lowerCAmelCase     =			nn.Linear(UpperCAmelCase__    ,      dim_out * 2    )





     def 							__UpperCAmelCase      ( self       :      Union[str, Any]    ,      UpperCAmelCase__       :      Optional[Any]    )				->  Any:
        if gate.device.type != "mps":
           return F.gelu(UpperCAmelCase__    )
        # mps: gelu is not implemented for float16
        return F.gelu(gate.to(dtype=torch.floataa    )    ).to(dtype=gate.dtype    )





     def 							__UpperCAmelCase      ( self       :      List[Any]    ,      UpperCAmelCase__       :      List[Any]    )				->  Optional[int]:
        lowerCAmelCase     ,  lowerCAmelCase     =			self.proj(UpperCAmelCase__    ).chunk(2    ,      dim=-1    )
        return hidden_states * self.gelu(UpperCAmelCase__    )




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      Optional[Any]    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    )				->  Optional[Any]:
        super().__init__()
        lowerCAmelCase     =			nn.Linear(UpperCAmelCase__    ,      UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      int    ,      UpperCAmelCase__       :      List[str]    )				->  List[Any]:
        lowerCAmelCase     =			self.proj(UpperCAmelCase__    )
        return x * torch.sigmoid(1.702 * x    )




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      Dict    ,      UpperCAmelCase__       :      Tuple    ,      UpperCAmelCase__       :      Union[str, Any]    )				->  Optional[int]:
        super().__init__()
        lowerCAmelCase     =			nn.Embedding(UpperCAmelCase__    ,      UpperCAmelCase__    )
        lowerCAmelCase     =			nn.SiLU()
        lowerCAmelCase     =			nn.Linear(UpperCAmelCase__    ,      embedding_dim * 2    )
        lowerCAmelCase     =			nn.LayerNorm(UpperCAmelCase__    ,      elementwise_affine=UpperCAmelCase__    )





     def 							__UpperCAmelCase      ( self       :      int    ,      UpperCAmelCase__       :      Dict    ,      UpperCAmelCase__       :      Optional[int]    )				->  Dict:
        lowerCAmelCase     =			self.linear(self.silu(self.emb(UpperCAmelCase__    )    )    )
        lowerCAmelCase     ,  lowerCAmelCase     =			torch.chunk(UpperCAmelCase__    ,      2    )
        lowerCAmelCase     =			self.norm(UpperCAmelCase__    ) * (1 + scale) + shift
        return x




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      List[Any]    ,      UpperCAmelCase__       :      Union[str, Any]    ,      UpperCAmelCase__       :      int    )				->  Dict:
        super().__init__()

        lowerCAmelCase     =			CombinedTimestepLabelEmbeddings(UpperCAmelCase__    ,      UpperCAmelCase__    )

        lowerCAmelCase     =			nn.SiLU()
        lowerCAmelCase     =			nn.Linear(UpperCAmelCase__    ,      6 * embedding_dim    ,      bias=UpperCAmelCase__    )
        lowerCAmelCase     =			nn.LayerNorm(UpperCAmelCase__    ,      elementwise_affine=UpperCAmelCase__    ,      eps=1E-6    )





     def 							__UpperCAmelCase      ( self       :      Union[str, Any]    ,      UpperCAmelCase__       :      Tuple    ,      UpperCAmelCase__       :      Tuple    ,      UpperCAmelCase__       :      Optional[int]    ,      UpperCAmelCase__       :      Optional[int]=None    )				->  List[Any]:
        lowerCAmelCase     =			self.linear(self.silu(self.emb(UpperCAmelCase__    ,      UpperCAmelCase__    ,      hidden_dtype=UpperCAmelCase__    )    )    )
        lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     ,  lowerCAmelCase     =			emb.chunk(6    ,      dim=1    )
        lowerCAmelCase     =			self.norm(UpperCAmelCase__    ) * (1 + scale_msa[:, None]) + shift_msa[:, None]
        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp




class 		UpperCAmelCase_      (  nn.Module		):





     def __init__( self       :      Union[str, Any]    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      int    ,      UpperCAmelCase__       :      Optional[str] = None    ,      UpperCAmelCase__       :      float = 1E-5    )				->  Optional[Any]:
        super().__init__()
        lowerCAmelCase     =			num_groups
        lowerCAmelCase     =			eps

        if act_fn is None:
           lowerCAmelCase     =			None
        else:
           lowerCAmelCase     =			get_activation(UpperCAmelCase__    )

        lowerCAmelCase     =			nn.Linear(UpperCAmelCase__    ,      out_dim * 2    )





     def 							__UpperCAmelCase      ( self       :      str    ,      UpperCAmelCase__       :      List[Any]    ,      UpperCAmelCase__       :      Any    )				->  Tuple:
        if self.act:
           lowerCAmelCase     =			self.act(UpperCAmelCase__    )
        lowerCAmelCase     =			self.linear(UpperCAmelCase__    )
        lowerCAmelCase     =			emb[:, :, None, None]
        lowerCAmelCase     ,  lowerCAmelCase     =			emb.chunk(2    ,      dim=1    )

        lowerCAmelCase     =			F.group_norm(UpperCAmelCase__    ,      self.num_groups    ,      eps=self.eps    )
        lowerCAmelCase     =			x * (1 + scale) + shift
        return x










'''simple docstring'''




import multiprocessing
import os
from typing import BinaryIO, Optional, Union

import fsspec

from .. import Dataset, Features, NamedSplit, config
from ..formatting import query_table
from ..packaged_modules.json.json import Json
from ..utils import logging
from ..utils.typing import NestedDataStructureLike, PathLike
from .abc import AbstractDatasetReader







class 	_lowerCAmelCase							(			__UpperCAmelCase ):


	def __init__(self       ,      lowercase       ,      lowercase = None       ,      lowercase = None       ,      lowercase = None       ,      lowercase = False       ,      lowercase = False       ,      lowercase = None       ,      lowercase = None       ,      **lowercase       ,      ):
				super().__init__(
				    lowercase       ,      split=lowercase       ,      features=lowercase       ,      cache_dir=lowercase       ,      keep_in_memory=lowercase       ,      streaming=lowercase       ,      num_proc=lowercase       ,      **lowercase       ,      )
				A_	:						Any	=     field
				A_	:						Optional[int]	=     path_or_paths if isinstance(lowercase       ,      lowercase							) else {self.split: path_or_paths}
				A_	:						str	=     Json(
				    cache_dir=lowercase       ,      data_files=lowercase       ,      features=lowercase       ,      field=lowercase       ,      **lowercase       ,      )


	def     _a	(self							):
				# Build iterable dataset
				if self.streaming:
							A_	:						Optional[int]	=     self.builder.as_streaming_dataset(split=self.split							)
				# Build regular (map-style) dataset
				else:
							A_	:						Optional[Any]	=     None
							A_	:						Optional[Any]	=     None
							A_	:						Tuple	=     None
							A_	:						Optional[int]	=     None

							self.builder.download_and_prepare(
							    download_config=lowercase       ,      download_mode=lowercase       ,      verification_mode=lowercase       ,      base_path=lowercase       ,      num_proc=self.num_proc       ,      )
							A_	:						Union[str, Any]	=     self.builder.as_dataset(
							    split=self.split       ,      verification_mode=lowercase       ,      in_memory=self.keep_in_memory							)
				return dataset







class 	_lowerCAmelCase							:


	def __init__(self       ,      lowercase       ,      lowercase       ,      lowercase = None       ,      lowercase = None       ,      **lowercase       ,      ):
				if num_proc is not None and num_proc <= 0:
							raise ValueError(F'num_proc {num_proc} must be an integer > 0.'							)

				A_	:						Union[str, Any]	=     dataset
				A_	:						Optional[int]	=     path_or_buf
				A_	:						str	=     batch_size if batch_size else config.DEFAULT_MAX_BATCH_SIZE
				A_	:						List[str]	=     num_proc
				A_	:						Union[str, Any]	=     """utf-8"""
				A_	:						Dict	=     to_json_kwargs


	def     _a	(self							):
				A_	:						Optional[Any]	=     self.to_json_kwargs.pop("""path_or_buf"""       ,      lowercase							)
				A_	:						Tuple	=     self.to_json_kwargs.pop("""orient"""       ,      """records"""							)
				A_	:						Union[str, Any]	=     self.to_json_kwargs.pop("""lines"""       ,      True if orient == """records""" else False							)
				A_	:						str	=     self.to_json_kwargs.pop("""index"""       ,      False if orient in ["""split""", """table"""] else True							)
				A_	:						Dict	=     self.to_json_kwargs.pop("""compression"""       ,      lowercase							)

				if compression not in [None, "infer", "gzip", "bz2", "xz"]:
							raise NotImplementedError(F'`datasets` currently does not support {compression} compression'							)

				if isinstance(self.path_or_buf       ,      (str, bytes, os.PathLike)							):
							with fsspec.open(self.path_or_buf       ,      """wb"""       ,      compression=lowercase							) as buffer:
										A_	:						List[str]	=     self._write(file_obj=lowercase       ,      orient=lowercase       ,      lines=lowercase       ,      index=lowercase       ,      **self.to_json_kwargs							)
				else:
							if compression:
										raise NotImplementedError(
										    F'The compression parameter is not supported when writing to a buffer, but compression={compression}'
										    """ was passed. Please provide a local path instead."""							)
							A_	:						List[Any]	=     self._write(
							    file_obj=self.path_or_buf       ,      orient=lowercase       ,      lines=lowercase       ,      index=lowercase       ,      **self.to_json_kwargs							)
				return written


	def     _a	(self       ,      lowercase							):
				A_,    A_,    A_,    A_,    A_	:						Dict	=     args

				A_	:						Any	=     query_table(
				    table=self.dataset.data       ,      key=slice(lowercase       ,      offset + self.batch_size							)       ,      indices=self.dataset._indices       ,      )
				A_	:						Union[str, Any]	=     batch.to_pandas().to_json(
				    path_or_buf=lowercase       ,      orient=lowercase       ,      lines=lowercase       ,      index=lowercase       ,      **lowercase							)
				if not json_str.endswith("""\n"""							):
							json_str += "\n"
				return json_str.encode(self.encoding							)


	def     _a	(self       ,      lowercase       ,      lowercase       ,      lowercase       ,      lowercase       ,      **lowercase       ,      ):
				A_	:						Optional[Any]	=     0

				if self.num_proc is None or self.num_proc == 1:
							for offset in logging.tqdm(
							    range(0       ,      len(self.dataset							)       ,      self.batch_size							)       ,      unit="""ba"""       ,      disable=not logging.is_progress_bar_enabled()       ,      desc="""Creating json from Arrow format"""       ,      ):
										A_	:						Any	=     self._batch_json((offset, orient, lines, index, to_json_kwargs)							)
										written += file_obj.write(lowercase							)
				else:
							A_,    A_	:						List[Any]	=     len(self.dataset							), self.batch_size
							with multiprocessing.Pool(self.num_proc							) as pool:
										for json_str in logging.tqdm(
										    pool.imap(
										        self._batch_json       ,      [(offset, orient, lines, index, to_json_kwargs) for offset in range(0       ,      lowercase       ,      lowercase							)]       ,      )       ,      total=(num_rows // batch_size) + 1 if num_rows % batch_size else num_rows // batch_size       ,      unit="""ba"""       ,      disable=not logging.is_progress_bar_enabled()       ,      desc="""Creating json from Arrow format"""       ,      ):
													written += file_obj.write(lowercase							)

				return written