




"""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__      (      lowercase_     ):
	A_							:				int						     =      4_2

class     A__      (      lowercase_     ,	lowercase_     ):




	@register_to_config
	def __init__(   self  :	List[str]      ,    lowercase  :	Optional[int] = None      ,    lowercase  :	Optional[int] = 3      ,    lowercase  :	List[str] = 3      ,    lowercase  :	Dict = False      ,    lowercase  :	Tuple = "positional"      ,    lowercase  :	List[Any] = 0      ,    lowercase  :	List[str] = True      ,    lowercase  :	List[str] = ("DownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D", "AttnDownBlock2D")      ,    lowercase  :	int = ("AttnUpBlock2D", "AttnUpBlock2D", "AttnUpBlock2D", "UpBlock2D")      ,    lowercase  :	str = (224, 448, 672, 896)      ,    lowercase  :	str = 2      ,    lowercase  :	Optional[int] = 1      ,    lowercase  :	Dict = 1      ,    lowercase  :	List[Any] = "conv"      ,    lowercase  :	Union[str, Any] = "conv"      ,    lowercase  :	Optional[int] = "silu"      ,    lowercase  :	List[str] = 8      ,    lowercase  :	List[str] = 32      ,    lowercase  :	int = 1e-5      ,    lowercase  :	int = "default"      ,    lowercase  :	Optional[Any] = True      ,    lowercase  :	Tuple = None      ,    lowercase  :	Union[str, Any] = None      ,    )					-> List[str]:
						super().__init__()

						__snake_case          =	sample_size
						__snake_case          =	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
						__snake_case          =	nn.Convad(lowercase      ,    block_out_channels[0]      ,    kernel_size=3      ,    padding=(1, 1)	)

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

						__snake_case          =	TimestepEmbedding(lowercase      ,    lowercase	)

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

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

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

											__snake_case          =	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
						__snake_case          =	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
						__snake_case          =	list(reversed(lowercase	)	)
						__snake_case          =	reversed_block_out_channels[0]
						for i, up_block_type in enumerate(lowercase	):
											__snake_case          =	output_channel
											__snake_case          =	reversed_block_out_channels[i]
											__snake_case          =	reversed_block_out_channels[min(i + 1      ,    len(lowercase	) - 1	)]

											__snake_case          =	i == len(lowercase	) - 1

											__snake_case          =	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	)
											__snake_case          =	output_channel

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




	def 							a_		(   self  :	Union[str, Any]      ,    lowercase  :	Tuple      ,    lowercase  :	Any      ,    lowercase  :	int = None      ,    lowercase  :	Tuple = True      ,    )					-> Union[UNetaDOutput, Tuple]:
						# 0. center input if necessary
						if self.config.center_input_sample:
											__snake_case          =	2 * sample - 1.0

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

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

						__snake_case          =	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.
						__snake_case          =	t_emb.to(dtype=self.dtype	)
						__snake_case          =	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":
																__snake_case          =	self.time_proj(lowercase	)

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

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

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

											down_block_res_samples += res_samples

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

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

											if hasattr(lowercase      ,    """skip_conv"""	):
																__snake_case          =	upsample_block(lowercase      ,    lowercase      ,    lowercase      ,    lowercase	)
											else:
																__snake_case          =	upsample_block(lowercase      ,    lowercase      ,    lowercase	)

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

						if skip_sample is not None:
											sample += skip_sample

						if self.config.time_embedding_type == "fourier":
											__snake_case          =	timesteps.reshape((sample.shape[0], *([1] * len(sample.shape[1:]	)))	)
											__snake_case          =	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


UpperCamelCase__ :    Optional[int] 					=							10_00_00_00_00_00
UpperCamelCase__ :    int 					=							1_00_00

UpperCamelCase__					,					UpperCamelCase__ :    List[str] 					=							os.path.split(__file__)
UpperCamelCase__ :    Any 					=							os.path.join(RESULTS_BASEPATH, "results", RESULTS_FILENAME.replace(".py", ".json"))




def       snake_case__	(       lowerCAmelCase   :		int				,					lowerCAmelCase   :		int		):
					__snake_case          =	pa.Table.from_pydict({"""col""": [0] * chunk_size}		)
					__snake_case          =	pa.concat_tables([table] * (num_examples // chunk_size)		)
					return datasets.Dataset(lowerCAmelCase				,					fingerprint="""table_100B"""		)
@dataclass
class     A__      :
	A_							:				int
	A_							:				int
	A_							:				int
	A_							:				int




	def 							a_		(   self  :	Optional[Any]	)					-> Any:
						__snake_case          =	np.random.default_rng(self.seed	)
						__snake_case          =	rng.integers(low=self.low      ,    high=self.high      ,    size=self.size	).tolist()




	def __iter__(   self  :	str	)					-> int:
						return iter(self._sampled_values	)




	def __len__(   self  :	Union[str, Any]	)					-> Any:
						return self.size




@get_duration
def       snake_case__	(       lowerCAmelCase   :		datasets.Dataset		):
					__snake_case          =	dataset[0]




@get_duration
def       snake_case__	(       lowerCAmelCase   :		datasets.Dataset		):
					__snake_case          =	dataset[-1]




@get_duration
def       snake_case__	(       lowerCAmelCase   :		datasets.Dataset		):
					__snake_case          =	dataset[range(len(lowerCAmelCase		) // 2				,					len(lowerCAmelCase		) // 2 + 1_024		)]




@get_duration
def       snake_case__	(       lowerCAmelCase   :		datasets.Dataset		):
					__snake_case          =	dataset[RandIter(0				,					len(lowerCAmelCase		)				,					1_024				,					seed=42		)]




def       snake_case__	(       ):
					__snake_case          =	{"""num examples""": SPEED_TEST_N_EXAMPLES}
					__snake_case          =	(get_first_row, get_last_row, get_batch_of_aaaa_rows, get_batch_of_aaaa_random_rows)
					print("""generating dataset"""		)
					__snake_case          =	generate_100B_dataset(num_examples=lowerCAmelCase				,					chunk_size=lowerCAmelCase		)
					print("""Functions"""		)
					for func in functions:
										print(func.__name__		)
										__snake_case          =	func(lowerCAmelCase		)

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


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


