



'''simple docstring'''






import gc
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    PNDMScheduler,
    StableDiffusionLDMaDPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import nightly, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_IMAGE_PARAMS, TEXT_TO_IMAGE_PARAMS


enable_full_determinism()

class        SCREAMING_SNAKE_CASE(							unittest.TestCase    ):







	"""simple docstring"""






	lowerCamelCase__    				= StableDiffusionLDMaDPipeline
	lowerCamelCase__    				= TEXT_TO_IMAGE_PARAMS
	lowerCamelCase__    				= TEXT_TO_IMAGE_BATCH_PARAMS
	lowerCamelCase__    				= TEXT_TO_IMAGE_IMAGE_PARAMS







	def   A			(		self    :  int	)    ->	Any:
								torch.manual_seed(0	)
								UpperCAmelCase				: Optional[int]											=	UNetaDConditionModel(
								    block_out_channels=(32, 64)	,    layers_per_block=2	,    sample_size=32	,    in_channels=4	,    out_channels=4	,    down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')	,    up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')	,    cross_attention_dim=32	,    )
								UpperCAmelCase				: Tuple											=	DDIMScheduler(
								    beta_start=0.0_00_85	,    beta_end=0.0_12	,    beta_schedule='''scaled_linear'''	,    clip_sample=__snake_case	,    set_alpha_to_one=__snake_case	,    )
								torch.manual_seed(0	)
								UpperCAmelCase				: Dict											=	AutoencoderKL(
								    block_out_channels=[32, 64]	,    in_channels=6	,    out_channels=6	,    down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']	,    up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']	,    latent_channels=4	,    )
								torch.manual_seed(0	)
								UpperCAmelCase				: Dict											=	CLIPTextConfig(
								    bos_token_id=0	,    eos_token_id=2	,    hidden_size=32	,    intermediate_size=37	,    layer_norm_eps=1E-05	,    num_attention_heads=4	,    num_hidden_layers=5	,    pad_token_id=1	,    vocab_size=1000	,    )
								UpperCAmelCase				: Any											=	CLIPTextModel(__snake_case	)
								UpperCAmelCase				: List[Any]											=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''	)

								UpperCAmelCase				: Union[str, Any]											=	{
								    '''unet''': unet,
								    '''scheduler''': scheduler,
								    '''vae''': vae,
								    '''text_encoder''': text_encoder,
								    '''tokenizer''': tokenizer,
								    '''safety_checker''': None,
								    '''feature_extractor''': None,
								}
								return components







	def   A			(		self    :  List[str]	,    __snake_case    :  str	,    __snake_case    :  Dict=0	)    ->	Dict:
								if str(__snake_case	).startswith('''mps'''	):
															UpperCAmelCase				: Optional[Any]											=	torch.manual_seed(__snake_case	)
								else:
															UpperCAmelCase				: Optional[int]											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: str											=	{
								    '''prompt''': '''A painting of a squirrel eating a burger''',
								    '''generator''': generator,
								    '''num_inference_steps''': 2,
								    '''guidance_scale''': 6.0,
								    '''output_type''': '''numpy''',
								}
								return inputs







	def   A			(		self    :  Union[str, Any]	)    ->	str:
								UpperCAmelCase				: Dict											=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator

								UpperCAmelCase				: int											=	self.get_dummy_components()
								UpperCAmelCase				: Optional[int]											=	StableDiffusionLDMaDPipeline(**__snake_case	)
								UpperCAmelCase				: List[str]											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Any											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: str											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase				: Union[str, Any]											=	output.rgb, output.depth

								UpperCAmelCase				: List[Any]											=	rgb[0, -3:, -3:, -1]
								UpperCAmelCase				: List[str]											=	depth[0, -3:, -1]

								assert rgb.shape == (1, 64, 64, 3)
								assert depth.shape == (1, 64, 64)

								UpperCAmelCase				: Optional[Any]											=	np.array(
								    [0.37_33_81_76, 0.7_02_47, 0.74_20_31_93, 0.51_64_36_04, 0.58_25_67_93, 0.60_93_21_36, 0.4_18_10_95, 0.48_35_58_77, 0.46_53_52_62]	)
								UpperCAmelCase				: Optional[Any]											=	np.array([103.46727, 85.81_20_04, 87.84_92_36]	)

								assert np.abs(image_slice_rgb.flatten() - expected_slice_rgb	).max() < 1E-2
								assert np.abs(image_slice_depth.flatten() - expected_slice_depth	).max() < 1E-2







	def   A			(		self    :  Optional[int]	)    ->	Tuple:
								UpperCAmelCase				: Tuple											=	self.get_dummy_components()
								UpperCAmelCase				: Any											=	StableDiffusionLDMaDPipeline(**__snake_case	)
								UpperCAmelCase				: Any											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: str											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: str											=	3 * [inputs['''prompt''']]

								# forward
								UpperCAmelCase				: Tuple											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase				: List[Any]											=	output.rgb, output.depth
								UpperCAmelCase				: Optional[Any]											=	rgb_slice_a[0, -3:, -3:, -1]
								UpperCAmelCase				: Optional[Any]											=	depth_slice_a[0, -3:, -1]

								UpperCAmelCase				: Tuple											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: Union[str, Any]											=	3 * [inputs.pop('''prompt'''	)]

								UpperCAmelCase				: Optional[Any]											=	ldmad_pipe.tokenizer(
								    __snake_case	,    padding='''max_length'''	,    max_length=ldmad_pipe.tokenizer.model_max_length	,    truncation=__snake_case	,    return_tensors='''pt'''	,    )
								UpperCAmelCase				: Union[str, Any]											=	text_inputs['''input_ids'''].to(__snake_case	)

								UpperCAmelCase				: Optional[int]											=	ldmad_pipe.text_encoder(__snake_case	)[0]

								UpperCAmelCase				: Optional[int]											=	prompt_embeds

								# forward
								UpperCAmelCase				: int											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase				: int											=	output.rgb, output.depth
								UpperCAmelCase				: Union[str, Any]											=	rgb_slice_a[0, -3:, -3:, -1]
								UpperCAmelCase				: int											=	depth_slice_a[0, -3:, -1]

								assert np.abs(rgb_slice_a.flatten() - rgb_slice_a.flatten()	).max() < 1E-4
								assert np.abs(depth_slice_a.flatten() - depth_slice_a.flatten()	).max() < 1E-4







	def   A			(		self    :  Optional[Any]	)    ->	Union[str, Any]:
								UpperCAmelCase				: List[Any]											=	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
								UpperCAmelCase				: Optional[Any]											=	self.get_dummy_components()
								UpperCAmelCase				: Union[str, Any]											=	PNDMScheduler(skip_prk_steps=__snake_case	)
								UpperCAmelCase				: Tuple											=	StableDiffusionLDMaDPipeline(**__snake_case	)
								UpperCAmelCase				: Any											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: int											=	self.get_dummy_inputs(__snake_case	)
								UpperCAmelCase				: Optional[Any]											=	'''french fries'''
								UpperCAmelCase				: Dict											=	ldmad_pipe(**__snake_case	,    negative_prompt=__snake_case	)

								UpperCAmelCase				: Tuple											=	output.rgb, output.depth
								UpperCAmelCase				: List[str]											=	rgb[0, -3:, -3:, -1]
								UpperCAmelCase				: List[Any]											=	depth[0, -3:, -1]

								assert rgb.shape == (1, 64, 64, 3)
								assert depth.shape == (1, 64, 64)

								UpperCAmelCase				: Any											=	np.array(
								    [0.3_70_44, 0.71_81_15_03, 0.7_22_32_51, 0.48_60_36_75, 0.5_63_83_91, 0.6_36_49_48, 0.42_83_37_04, 0.4_90_13_15, 0.47_92_62_17]	)
								UpperCAmelCase				: str											=	np.array([107.84738, 84.6_28_02, 89.96_21_35]	)
								assert np.abs(rgb_slice.flatten() - expected_slice_rgb	).max() < 1E-2
								assert np.abs(depth_slice.flatten() - expected_slice_depth	).max() < 1E-2

@slow
@require_torch_gpu
class        SCREAMING_SNAKE_CASE(							unittest.TestCase    ):







	"""simple docstring"""






	def   A			(		self    :  str	)    ->	Dict:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def   A			(		self    :  List[Any]	,    __snake_case    :  int	,    __snake_case    :  Tuple="cpu"	,    __snake_case    :  Tuple=torch.floataa	,    __snake_case    :  List[Any]=0	)    ->	Union[str, Any]:
								UpperCAmelCase				: List[str]											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: List[str]											=	np.random.RandomState(__snake_case	).standard_normal((1, 4, 64, 64)	)
								UpperCAmelCase				: Optional[Any]											=	torch.from_numpy(__snake_case	).to(device=__snake_case	,    dtype=__snake_case	)
								UpperCAmelCase				: List[str]											=	{
								    '''prompt''': '''a photograph of an astronaut riding a horse''',
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 3,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs







	def   A			(		self    :  str	)    ->	str:
								UpperCAmelCase				: Union[str, Any]											=	StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''	)
								UpperCAmelCase				: Tuple											=	ldmad_pipe.to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: Union[str, Any]											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: Optional[int]											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase				: List[str]											=	output.rgb, output.depth
								UpperCAmelCase				: Union[str, Any]											=	rgb[0, -3:, -3:, -1].flatten()
								UpperCAmelCase				: Any											=	rgb[0, -3:, -1].flatten()

								assert rgb.shape == (1, 512, 512, 3)
								assert depth.shape == (1, 512, 512)

								UpperCAmelCase				: Optional[Any]											=	np.array(
								    [0.53_80_54_65, 0.56_70_73_05, 0.5_48_65_15, 0.57_01_22_36, 0.5_81_45_11, 0.56_25_34_87, 0.54_84_30_14, 0.55_09_22_63, 0.6_45_97_06]	)
								UpperCAmelCase				: str											=	np.array(
								    [0.9_26_37_81, 0.6_67_86_72, 0.5_48_65_15, 0.92_20_21_45, 0.67_83_11_35, 0.56_25_34_87, 0.9_24_16_94, 0.7_55_14_78, 0.6_45_97_06]	)
								assert np.abs(rgb_slice - expected_slice_rgb	).max() < 3E-3
								assert np.abs(depth_slice - expected_slice_depth	).max() < 3E-3

@nightly
@require_torch_gpu
class        SCREAMING_SNAKE_CASE(							unittest.TestCase    ):







	"""simple docstring"""






	def   A			(		self    :  Tuple	)    ->	int:
								super().tearDown()
								gc.collect()
								torch.cuda.empty_cache()







	def   A			(		self    :  str	,    __snake_case    :  Optional[Any]	,    __snake_case    :  Optional[Any]="cpu"	,    __snake_case    :  Dict=torch.floataa	,    __snake_case    :  int=0	)    ->	int:
								UpperCAmelCase				: Tuple											=	torch.Generator(device=__snake_case	).manual_seed(__snake_case	)
								UpperCAmelCase				: Dict											=	np.random.RandomState(__snake_case	).standard_normal((1, 4, 64, 64)	)
								UpperCAmelCase				: int											=	torch.from_numpy(__snake_case	).to(device=__snake_case	,    dtype=__snake_case	)
								UpperCAmelCase				: Optional[int]											=	{
								    '''prompt''': '''a photograph of an astronaut riding a horse''',
								    '''latents''': latents,
								    '''generator''': generator,
								    '''num_inference_steps''': 50,
								    '''guidance_scale''': 7.5,
								    '''output_type''': '''numpy''',
								}
								return inputs







	def   A			(		self    :  Tuple	)    ->	Any:
								UpperCAmelCase				: Any											=	StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d'''	).to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: List[str]											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: int											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase				: Dict											=	output.rgb, output.depth

								UpperCAmelCase				: List[Any]											=	0.49_55_86
								UpperCAmelCase				: List[Any]											=	0.33_79_55_15
								UpperCAmelCase				: int											=	112.48518
								UpperCAmelCase				: Tuple											=	98.48_97_46
								assert np.abs(expected_rgb_mean - rgb.mean()	) < 1E-3
								assert np.abs(expected_rgb_std - rgb.std()	) < 1E-3
								assert np.abs(expected_depth_mean - depth.mean()	) < 1E-3
								assert np.abs(expected_depth_std - depth.std()	) < 1E-3







	def   A			(		self    :  int	)    ->	Any:
								UpperCAmelCase				: Optional[Any]											=	StableDiffusionLDMaDPipeline.from_pretrained('''Intel/ldm3d-4c'''	).to(__snake_case	)
								ldmad_pipe.set_progress_bar_config(disable=__snake_case	)

								UpperCAmelCase				: List[str]											=	self.get_inputs(__snake_case	)
								UpperCAmelCase				: Dict											=	ldmad_pipe(**__snake_case	)
								UpperCAmelCase				: int											=	output.rgb, output.depth

								UpperCAmelCase				: List[Any]											=	0.4_19_41_27
								UpperCAmelCase				: str											=	0.35_37_55_86
								UpperCAmelCase				: Union[str, Any]											=	0.5_63_85_02
								UpperCAmelCase				: Optional[int]											=	0.34_68_61_03

								assert rgb.shape == (1, 512, 512, 3)
								assert depth.shape == (1, 512, 512, 1)
								assert np.abs(expected_rgb_mean - rgb.mean()	) < 1E-3
								assert np.abs(expected_rgb_std - rgb.std()	) < 1E-3
								assert np.abs(expected_depth_mean - depth.mean()	) < 1E-3
								assert np.abs(expected_depth_std - depth.std()	) < 1E-3










'''simple docstring'''






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


UpperCamelCase__:     List[Any]										=    logging.getLogger(__name__)


def 				snake_case_  (      )   ->     int:
							UpperCAmelCase				: int											=	argparse.ArgumentParser(
							    description='''Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.'''						)
							parser.add_argument(
							    '''--dataset_name'''     ,   type=_lowerCAmelCase     ,   default='''wikitext'''     ,   help='''Name of the training. Explore datasets at: hf.co/datasets.'''     ,   )
							parser.add_argument(
							    '''--dataset_config'''     ,   type=_lowerCAmelCase     ,   default='''wikitext-103-raw-v1'''     ,   help='''Configuration name of the dataset.'''						)
							parser.add_argument(
							    '''--tokenizer_name_or_path'''     ,   type=_lowerCAmelCase     ,   default='''sayakpaul/unigram-tokenizer-wikitext'''     ,   help='''Tokenizer identifier. Can be a local filepath or a Hub identifier.'''     ,   )
							parser.add_argument(
							    '''--shard_size'''     ,   type=_lowerCAmelCase     ,   default=1000     ,   help='''Number of entries to go in a single shard.'''     ,   )
							parser.add_argument('''--split'''     ,   type=_lowerCAmelCase     ,   default='''train'''     ,   choices=['''train''', '''test''', '''validation''']						)
							parser.add_argument(
							    '''--limit'''     ,   default=_lowerCAmelCase     ,   type=_lowerCAmelCase     ,   help='''Limit the number of shards (used for debugging).'''     ,   )
							parser.add_argument(
							    '''--max_length'''     ,   type=_lowerCAmelCase     ,   default=512     ,   help='''Maximum sequence length. For training on TPUs, it helps to have a maximum'''
							    ''' sequence length that is a multiple of 8.'''     ,   )
							parser.add_argument(
							    '''--output_dir'''     ,   default='''tf-tpu'''     ,   type=_lowerCAmelCase     ,   help='''Output directory where the TFRecord shards will be saved. If the'''
							    ''' path is appended with `gs://` (\'gs://tf-tpu\', for example) then the TFRecord'''
							    ''' shards will be directly saved to a Google Cloud Storage bucket.'''     ,   )

							UpperCAmelCase				: List[str]											=	parser.parse_args()
							return args


def 				snake_case_  (      _lowerCAmelCase						:		Union[str, Any]						)   ->     Optional[int]:
							def fn(_lowerCAmelCase						:		Tuple						):
														return tokenizer(examples['''text''']						)

							return fn


def 				snake_case_  (      _lowerCAmelCase						:		int						)   ->     Dict:
							UpperCAmelCase				: Optional[Any]											=	[]
							for i in range(len(tokenized_data['''input_ids''']						)						):
														UpperCAmelCase				: Optional[Any]											=	{
														    '''input_ids''': tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data['''input_ids'''][i]						)						),
														    '''attention_mask''': tf.train.Feature(
														        intaa_list=tf.train.IntaaList(value=tokenized_data['''attention_mask'''][i]						)						),
														}
														UpperCAmelCase				: Optional[Any]											=	tf.train.Features(feature=_lowerCAmelCase						)
														UpperCAmelCase				: Union[str, Any]											=	tf.train.Example(features=_lowerCAmelCase						)
														UpperCAmelCase				: Any											=	example.SerializeToString()
														records.append(_lowerCAmelCase						)
							return records


def 				snake_case_  (      _lowerCAmelCase						:		List[str]						)   ->     List[str]:
							UpperCAmelCase				: List[str]											=	datasets.load_dataset(args.dataset_name     ,   args.dataset_config     ,   split=args.split						)

							if args.limit is not None:
														UpperCAmelCase				: List[Any]											=	min(len(_lowerCAmelCase						)     ,   args.limit						)
														UpperCAmelCase				: Dict											=	dataset.select(range(_lowerCAmelCase						)						)
														print(f"""Limiting the dataset to {args.limit} entries."""						)

							UpperCAmelCase				: List[Any]											=	AutoTokenizer.from_pretrained(args.tokenizer_name_or_path						)

							# Handle output directory creation.
							# For serializing into a Google Cloud Storage Bucket, one needs to first
							# create a bucket.
							if "gs" not in args.output_dir:
														if not os.path.exists(args.output_dir						):
																					os.makedirs(args.output_dir						)
														UpperCAmelCase				: List[Any]											=	os.path.join(args.output_dir     ,   args.split						)
														if not os.path.exists(_lowerCAmelCase						):
																					os.makedirs(_lowerCAmelCase						)
							else:
														UpperCAmelCase				: Tuple											=	os.path.join(args.output_dir     ,   args.split						)

							# Tokenize the whole dataset at once.
							UpperCAmelCase				: Optional[Any]											=	tokenize_function(_lowerCAmelCase						)
							UpperCAmelCase				: int											=	dataset.map(_lowerCAmelCase     ,   batched=_lowerCAmelCase     ,   num_proc=4     ,   remove_columns=['''text''']						)

							# We need to concatenate all our texts together, and then split the result
							# into chunks of a fixed size, which we will call block_size. To do this, we
							# will use the map method again, with the option batched=True. When we use batched=True,
							# the function we pass to map() will be passed multiple inputs at once, allowing us
							# to group them into more or fewer examples than we had in the input.
							# This allows us to create our new fixed-length samples. The advantage of this
							# method is that we don't lose a whole lot of content from the dataset compared to the
							# case where we simply tokenize with a pre-defined max_length.

							def group_texts(_lowerCAmelCase						:		List[str]						):
														# Concatenate all texts.
														UpperCAmelCase				: Optional[Any]											=	{k: sum(examples[k]     ,   []						) for k in examples.keys()}
														UpperCAmelCase				: Optional[Any]											=	len(concatenated_examples[list(examples.keys()						)[0]]						)
														# We drop the small remainder, though you could add padding instead if the model supports it
														# In this, as in all things, we advise you to follow your heart 🫀
														UpperCAmelCase				: Any											=	(total_length // args.max_length) * args.max_length
														# Split by chunks of max_len.
														UpperCAmelCase				: List[str]											=	{
														    k: [t[i : i + args.max_length] for i in range(0     ,   _lowerCAmelCase     ,   args.max_length						)]
														    for k, t in concatenated_examples.items()
														}
														return result

							UpperCAmelCase				: str											=	dataset_tokenized.map(_lowerCAmelCase     ,   batched=_lowerCAmelCase     ,   batch_size=1000     ,   num_proc=4						)

							UpperCAmelCase				: Union[str, Any]											=	0
							UpperCAmelCase				: Dict											=	0
							for shard in range(0     ,   len(_lowerCAmelCase						)     ,   args.shard_size						):
														UpperCAmelCase				: Union[str, Any]											=	grouped_dataset[shard : shard + args.shard_size]
														UpperCAmelCase				: Optional[int]											=	len(dataset_snapshot['''input_ids''']						)
														UpperCAmelCase				: Dict											=	os.path.join(_lowerCAmelCase     ,   f"""dataset-{shard_count}-{records_containing}.tfrecord"""						)
														UpperCAmelCase				: List[Any]											=	get_serialized_examples(_lowerCAmelCase						)

														with tf.io.TFRecordWriter(_lowerCAmelCase						) as out_file:
																					for i in range(len(_lowerCAmelCase						)						):
																												UpperCAmelCase				: Any											=	serialized_examples[i]
																												out_file.write(_lowerCAmelCase						)
																					print('''Wrote file {} containing {} records'''.format(_lowerCAmelCase     ,   _lowerCAmelCase						)						)

														shard_count += 1
														total_records += records_containing

							with open(f"""split-{args.split}-records-count.txt"""     ,   '''w'''						) as f:
														print(f"""Total {args.split} records: {total_records}"""     ,   file=_lowerCAmelCase						)


if __name__ == "__main__":
						UpperCamelCase__:     List[Any]										=    parse_args()
						main(args)





