




'''simple docstring'''







from dataclasses import dataclass
from typing import Optional, Tuple, Union

import torch
import torch.nn as nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .embeddings import GaussianFourierProjection, TimestepEmbedding, Timesteps
from .modeling_utils import ModelMixin
from .unet_ad_blocks import UNetMidBlockaD, get_down_block, get_up_block


@dataclass
class      _a    (      __a					):
       __a							:	torch.FloatTensor






class      _a    (      __a				,				__a					):
       @register_to_config
       def __init__(      self		:      List[str]  ,							lowercase		:      Optional[Union[int, Tuple[int, int]]] = None  ,							lowercase		:      int = 3  ,							lowercase		:      int = 3  ,							lowercase		:      bool = False  ,							lowercase		:      str = "positional"  ,							lowercase		:      int = 0  ,							lowercase		:      bool = True  ,							lowercase		:      Tuple[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")  ,							lowercase		:      Tuple[str] = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")  ,							lowercase		:      Tuple[int] = (224, 448, 672, 896)  ,							lowercase		:      int = 2  ,							lowercase		:      float = 1  ,							lowercase		:      int = 1  ,							lowercase		:      str = "conv"  ,							lowercase		:      str = "conv"  ,							lowercase		:      str = "silu"  ,							lowercase		:      Optional[int] = 8  ,							lowercase		:      int = 32  ,							lowercase		:      float = 1E-5  ,							lowercase		:      str = "default"  ,							lowercase		:      bool = True  ,							lowercase		:      Optional[str] = None  ,							lowercase		:      Optional[int] = None  ,							):



           '''simple docstring'''

           super().__init__()

           UpperCAmelCase     			=						sample_size
           UpperCAmelCase     			=						block_out_channels[0] * 4

           # Check inputs
           if len(lowercase							) != len(lowercase							):
               raise ValueError(
                   f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."							)

           if len(lowercase							) != len(lowercase							):
               raise ValueError(
                   f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."							)

           # input
           UpperCAmelCase     			=						nn.Convad(lowercase  ,							block_out_channels[0]  ,							kernel_size=3  ,							padding=(1, 1)							)

           # time
           if time_embedding_type == "fourier":
               UpperCAmelCase     			=						GaussianFourierProjection(embedding_size=block_out_channels[0]  ,							scale=16							)
               UpperCAmelCase     			=						2 * block_out_channels[0]
           elif time_embedding_type == "positional":
               UpperCAmelCase     			=						Timesteps(block_out_channels[0]  ,							lowercase  ,							lowercase							)
               UpperCAmelCase     			=						block_out_channels[0]

           UpperCAmelCase     			=						TimestepEmbedding(lowercase  ,							lowercase							)

           # class embedding
           if class_embed_type is None and num_class_embeds is not None:
               UpperCAmelCase     			=						nn.Embedding(lowercase  ,							lowercase							)
           elif class_embed_type == "timestep":
               UpperCAmelCase     			=						TimestepEmbedding(lowercase  ,							lowercase							)
           elif class_embed_type == "identity":
               UpperCAmelCase     			=						nn.Identity(lowercase  ,							lowercase							)
           else:
               UpperCAmelCase     			=						None

           UpperCAmelCase     			=						nn.ModuleList([]							)
           UpperCAmelCase     			=						None
           UpperCAmelCase     			=						nn.ModuleList([]							)

           # down
           UpperCAmelCase     			=						block_out_channels[0]
           for i, down_block_type in enumerate(lowercase							):
               UpperCAmelCase     			=						output_channel
               UpperCAmelCase     			=						block_out_channels[i]
               UpperCAmelCase     			=						i == len(lowercase							) - 1

               UpperCAmelCase     			=						get_down_block(
                   lowercase  ,							num_layers=lowercase  ,							in_channels=lowercase  ,							out_channels=lowercase  ,							temb_channels=lowercase  ,							add_downsample=not is_final_block  ,							resnet_eps=lowercase  ,							resnet_act_fn=lowercase  ,							resnet_groups=lowercase  ,							attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel  ,							downsample_padding=lowercase  ,							resnet_time_scale_shift=lowercase  ,							downsample_type=lowercase  ,							)
               self.down_blocks.append(lowercase							)

           # mid
           UpperCAmelCase     			=						UNetMidBlockaD(
               in_channels=block_out_channels[-1]  ,							temb_channels=lowercase  ,							resnet_eps=lowercase  ,							resnet_act_fn=lowercase  ,							output_scale_factor=lowercase  ,							resnet_time_scale_shift=lowercase  ,							attention_head_dim=attention_head_dim if attention_head_dim is not None else block_out_channels[-1]  ,							resnet_groups=lowercase  ,							add_attention=lowercase  ,							)

           # up
           UpperCAmelCase     			=						list(reversed(lowercase							)							)
           UpperCAmelCase     			=						reversed_block_out_channels[0]
           for i, up_block_type in enumerate(lowercase							):
               UpperCAmelCase     			=						output_channel
               UpperCAmelCase     			=						reversed_block_out_channels[i]
               UpperCAmelCase     			=						reversed_block_out_channels[min(i + 1  ,							len(lowercase							) - 1							)]

               UpperCAmelCase     			=						i == len(lowercase							) - 1

               UpperCAmelCase     			=						get_up_block(
                   lowercase  ,							num_layers=layers_per_block + 1  ,							in_channels=lowercase  ,							out_channels=lowercase  ,							prev_output_channel=lowercase  ,							temb_channels=lowercase  ,							add_upsample=not is_final_block  ,							resnet_eps=lowercase  ,							resnet_act_fn=lowercase  ,							resnet_groups=lowercase  ,							attention_head_dim=attention_head_dim if attention_head_dim is not None else output_channel  ,							resnet_time_scale_shift=lowercase  ,							upsample_type=lowercase  ,							)
               self.up_blocks.append(lowercase							)
               UpperCAmelCase     			=						output_channel

           # out
           UpperCAmelCase     			=						norm_num_groups if norm_num_groups is not None else min(block_out_channels[0] // 4  ,							32							)
           UpperCAmelCase     			=						nn.GroupNorm(num_channels=block_out_channels[0]  ,							num_groups=lowercase  ,							eps=lowercase							)
           UpperCAmelCase     			=						nn.SiLU()
           UpperCAmelCase     			=						nn.Convad(block_out_channels[0]  ,							lowercase  ,							kernel_size=3  ,							padding=1							)







       def        A (      self		:      Union[str, Any]  ,							lowercase		:      torch.FloatTensor  ,							lowercase		:      Union[torch.Tensor, float, int]  ,							lowercase		:      Optional[torch.Tensor] = None  ,							lowercase		:      bool = True  ,							):



           '''simple docstring'''

           if self.config.center_input_sample:
               UpperCAmelCase     			=						2 * sample - 1.0

           # 1. time
           UpperCAmelCase     			=						timestep
           if not torch.is_tensor(lowercase							):
               UpperCAmelCase     			=						torch.tensor([timesteps]  ,							dtype=torch.long  ,							device=sample.device							)
           elif torch.is_tensor(lowercase							) and len(timesteps.shape							) == 0:
               UpperCAmelCase     			=						timesteps[None].to(sample.device							)

           # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
           UpperCAmelCase     			=						timesteps * torch.ones(sample.shape[0]  ,							dtype=timesteps.dtype  ,							device=timesteps.device							)

           UpperCAmelCase     			=						self.time_proj(lowercase							)

           # timesteps does not contain any weights and will always return f32 tensors
           # but time_embedding might actually be running in fp16. so we need to cast here.
           # there might be better ways to encapsulate this.
           UpperCAmelCase     			=						t_emb.to(dtype=self.dtype							)
           UpperCAmelCase     			=						self.time_embedding(lowercase							)

           if self.class_embedding is not None:
               if class_labels is None:
                   raise ValueError('''class_labels should be provided when doing class conditioning'''							)

               if self.config.class_embed_type == "timestep":
                   UpperCAmelCase     			=						self.time_proj(lowercase							)

               UpperCAmelCase     			=						self.class_embedding(lowercase							).to(dtype=self.dtype							)
               UpperCAmelCase     			=						emb + class_emb

           # 2. pre-process
           UpperCAmelCase     			=						sample
           UpperCAmelCase     			=						self.conv_in(lowercase							)

           # 3. down
           UpperCAmelCase     			=						(sample,)
           for downsample_block in self.down_blocks:
               if hasattr(lowercase  ,							'''skip_conv'''							):
                   UpperCAmelCase			,	UpperCAmelCase			,	UpperCAmelCase     			=						downsample_block(
                       hidden_states=lowercase  ,							temb=lowercase  ,							skip_sample=lowercase							)
               else:
                   UpperCAmelCase			,	UpperCAmelCase     			=						downsample_block(hidden_states=lowercase  ,							temb=lowercase							)

               down_block_res_samples += res_samples

           # 4. mid
           UpperCAmelCase     			=						self.mid_block(lowercase  ,							lowercase							)

           # 5. up
           UpperCAmelCase     			=						None
           for upsample_block in self.up_blocks:
               UpperCAmelCase     			=						down_block_res_samples[-len(upsample_block.resnets							) :]
               UpperCAmelCase     			=						down_block_res_samples[: -len(upsample_block.resnets							)]

               if hasattr(lowercase  ,							'''skip_conv'''							):
                   UpperCAmelCase			,	UpperCAmelCase     			=						upsample_block(lowercase  ,							lowercase  ,							lowercase  ,							lowercase							)
               else:
                   UpperCAmelCase     			=						upsample_block(lowercase  ,							lowercase  ,							lowercase							)

        # 6. post-process
           UpperCAmelCase     			=						self.conv_norm_out(lowercase							)
           UpperCAmelCase     			=						self.conv_act(lowercase							)
           UpperCAmelCase     			=						self.conv_out(lowercase							)

           if skip_sample is not None:
               sample += skip_sample

           if self.config.time_embedding_type == "fourier":
               UpperCAmelCase     			=						timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]							)))							)
               UpperCAmelCase     			=						sample / timesteps

           if not return_dict:
               return (sample,)

           return UNetaDOutput(sample=lowercase							)










'''simple docstring'''







import json
import os
from dataclasses import dataclass

import numpy as np
import pyarrow as pa

import datasets
from utils import get_duration


A					       =10_00_00_00_00_00
A					       =1_00_00

A       , A					       =os.path.split(__file__)
A					       =os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))





def 		snake_case_				(_a      :				int       ,  _a      :				int	):
    UpperCAmelCase     			=						pa.Table.from_pydict({'''col''': [0] * chunk_size}	)
    UpperCAmelCase     			=						pa.concat_tables([table] * (num_examples // chunk_size)	)
    return datasets.Dataset(_a       ,  fingerprint='''table_100B'''	)


@dataclass
class      _a    :
       __a							:	int
       __a							:	int
       __a							:	int
       __a							:	int
       def        A (      self		:      List[str]							):



           '''simple docstring'''

           UpperCAmelCase     			=						np.random.default_rng(self.seed							)
           UpperCAmelCase     			=						rng.integers(low=self.low  ,							high=self.high  ,							size=self.size							).tolist()
       def __iter__(      self		:      Tuple							):



           '''simple docstring'''

           return iter(self._sampled_values							)







       def __len__(      self		:      Dict							):



           '''simple docstring'''

           return self.size






@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[0]





@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[-1]





@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[range(len(_a	) // 2       ,  len(_a	) // 2 + 1_0_2_4	)]





@get_duration
def 		snake_case_				(_a      :				datasets.Dataset	):
    UpperCAmelCase     			=						dataset[RandIter(0       ,  len(_a	)       ,  1_0_2_4       ,  seed=4_2	)]





def 		snake_case_				():
    UpperCAmelCase     			=						{'''num examples''': SPEED_TEST_N_EXAMPLES}
    UpperCAmelCase     			=						(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
    print('''generating dataset'''	)
    UpperCAmelCase     			=						generate_100B_dataset(num_examples=_a       ,  chunk_size=_a	)
    print('''Functions'''	)
    for func in functions:
        print(func.__name__	)
        UpperCAmelCase     			=						func(_a	)

    with open(_a       ,  '''wb'''	) as f:
        f.write(json.dumps(_a	).encode('''utf-8'''	)	)


if __name__ == "__main__":  # useful to run the profiler
   benchmark_table_100B()




