
"""simple docstring"""




import gc
import random
import unittest

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

from diffusers import (
    AutoencoderKL,
    DiffusionPipeline,
    EulerDiscreteScheduler,
    StableDiffusionXLImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.utils import floats_tensor, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()

class     _UpperCamelCase			(   UpperCamelCase__       ,    UpperCamelCase__       ,    unittest.TestCase					):






			'''simple docstring'''

			__UpperCamelCase :  Tuple    =	StableDiffusionXLImgaImgPipeline
			__UpperCamelCase :  str    =	TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
			__UpperCamelCase :  int    =	PipelineTesterMixin.required_optional_params - {"""latents"""}
			__UpperCamelCase :  List[str]    =	TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			__UpperCamelCase :  str    =	IMAGE_TO_IMAGE_IMAGE_PARAMS
			__UpperCamelCase :  List[str]    =	IMAGE_TO_IMAGE_IMAGE_PARAMS

			def  lowerCAmelCase__		(    self	:    Any ):
							torch.manual_seed(0 )
							UpperCamelCase_:							str		=  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""")  ,		attention_head_dim=(2, 4)  ,		use_linear_projection=__lowerCamelCase  ,		addition_embed_type="""text_time"""  ,		addition_time_embed_dim=8  ,		transformer_layers_per_block=(1, 2)  ,		projection_class_embeddings_input_dim=80  ,		cross_attention_dim=64  ,		)
							UpperCamelCase_:							Dict		=  EulerDiscreteScheduler(
							    beta_start=0.0_0085  ,		beta_end=0.012  ,		steps_offset=1  ,		beta_schedule="""scaled_linear"""  ,		timestep_spacing="""leading"""  ,		)
							torch.manual_seed(0 )
							UpperCamelCase_:							Union[str, Any]		=  AutoencoderKL(
							    block_out_channels=[32, 64]  ,		in_channels=3  ,		out_channels=3  ,		down_block_types=["""DownEncoderBlock2D""", """DownEncoderBlock2D"""]  ,		up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D"""]  ,		latent_channels=4  ,		sample_size=128  ,		)
							torch.manual_seed(0 )
							UpperCamelCase_:							str		=  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  ,		hidden_act="""gelu"""  ,		projection_dim=32  ,		)
							UpperCamelCase_:							Tuple		=  CLIPTextModel(__lowerCamelCase )
							UpperCamelCase_:							str		=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  ,		local_files_only=__lowerCamelCase )

							UpperCamelCase_:							List[str]		=  CLIPTextModelWithProjection(__lowerCamelCase )
							UpperCamelCase_:							Union[str, Any]		=  CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  ,		local_files_only=__lowerCamelCase )

							UpperCamelCase_:							int		=  {
							    """unet""": unet,
							    """scheduler""": scheduler,
							    """vae""": vae,
							    """text_encoder""": text_encoder,
							    """tokenizer""": tokenizer,
							    """text_encoder_2""": text_encoder_a,
							    """tokenizer_2""": tokenizer_a,
							    # "safety_checker": None,
							    # "feature_extractor": None,
							}
							return components

			def  lowerCAmelCase__		(    self	:    Any  ,		snake_case_	:    int  ,		snake_case_	:    Tuple=0 ):
							UpperCamelCase_:							str		=  floats_tensor((1, 3, 32, 32)  ,		rng=random.Random(__lowerCamelCase ) ).to(__lowerCamelCase )
							UpperCamelCase_:							Dict		=  image / 2 + 0.5
							if str(__lowerCamelCase ).startswith("""mps""" ):
											UpperCamelCase_:							Tuple		=  torch.manual_seed(__lowerCamelCase )
							else:
											UpperCamelCase_:							Tuple		=  torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
							UpperCamelCase_:							Union[str, Any]		=  {
							    """prompt""": """A painting of a squirrel eating a burger""",
							    """image""": image,
							    """generator""": generator,
							    """num_inference_steps""": 2,
							    """guidance_scale""": 5.0,
							    """output_type""": """numpy""",
							    """strength""": 0.75,
							}
							return inputs

			def  lowerCAmelCase__		(    self	:    Optional[int] ):
							UpperCamelCase_:							Optional[Any]		=  """cpu"""  # ensure determinism for the device-dependent torch.Generator
							UpperCamelCase_:							str		=  self.get_dummy_components()
							UpperCamelCase_:							Any		=  StableDiffusionXLImgaImgPipeline(**__lowerCamelCase )
							UpperCamelCase_:							Union[str, Any]		=  sd_pipe.to(__lowerCamelCase )
							sd_pipe.set_progress_bar_config(disable=__lowerCamelCase )

							UpperCamelCase_:							List[str]		=  self.get_dummy_inputs(__lowerCamelCase )
							UpperCamelCase_:							List[Any]		=  sd_pipe(**__lowerCamelCase ).images
							UpperCamelCase_:							int		=  image[0, -3:, -3:, -1]

							assert image.shape == (1, 32, 32, 3)

							UpperCamelCase_:							Any		=  np.array([0.4656, 0.4840, 0.4439, 0.6698, 0.5574, 0.4524, 0.5799, 0.5943, 0.5165] )

							assert np.abs(image_slice.flatten() - expected_slice ).max() < 1e-2

			def  lowerCAmelCase__		(    self	:    Dict ):
							super().test_attention_slicing_forward_pass(expected_max_diff=3e-3 )

			def  lowerCAmelCase__		(    self	:    int ):
							super().test_inference_batch_single_identical(expected_max_diff=3e-3 )

			def  lowerCAmelCase__		(    self	:    List[str] ):
							pass




			def  lowerCAmelCase__		(    self	:    int ):
							UpperCamelCase_:							List[Any]		=  self.get_dummy_components()
							UpperCamelCase_:							Tuple		=  StableDiffusionXLImgaImgPipeline(**__lowerCamelCase )
							UpperCamelCase_:							int		=  sd_pipe.to(__lowerCamelCase )
							UpperCamelCase_:							Tuple		=  sd_pipe.to(__lowerCamelCase )
							sd_pipe.set_progress_bar_config(disable=__lowerCamelCase )

							# forward without prompt embeds
							UpperCamelCase_:							Dict		=  self.get_dummy_inputs(__lowerCamelCase )
							UpperCamelCase_:							Union[str, Any]		=  3 * ["""this is a negative prompt"""]
							UpperCamelCase_:							int		=  negative_prompt
							UpperCamelCase_:							int		=  3 * [inputs["""prompt"""]]

							UpperCamelCase_:							List[Any]		=  sd_pipe(**__lowerCamelCase )
							UpperCamelCase_:							List[str]		=  output.images[0, -3:, -3:, -1]

							# forward with prompt embeds
							UpperCamelCase_:							List[Any]		=  self.get_dummy_inputs(__lowerCamelCase )
							UpperCamelCase_:							Optional[Any]		=  3 * ["""this is a negative prompt"""]
							UpperCamelCase_:							str		=  3 * [inputs.pop("""prompt""" )]

							(
							    (
							    UpperCamelCase_     
							),				(
							    UpperCamelCase_     
							),				(
							    UpperCamelCase_     
							),				(
							    UpperCamelCase_     
							),				     
							):							Tuple		=  sd_pipe.encode_prompt(__lowerCamelCase  ,		negative_prompt=__lowerCamelCase )

							UpperCamelCase_:							Dict		=  sd_pipe(
							    **__lowerCamelCase  ,		prompt_embeds=__lowerCamelCase  ,		negative_prompt_embeds=__lowerCamelCase  ,		pooled_prompt_embeds=__lowerCamelCase  ,		negative_pooled_prompt_embeds=__lowerCamelCase  ,		)
							UpperCamelCase_:							int		=  output.images[0, -3:, -3:, -1]

							# make sure that it's equal
							assert np.abs(image_slice_a.flatten() - image_slice_a.flatten() ).max() < 1e-4




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






			'''simple docstring'''

			def  lowerCAmelCase__		(    self	:    Union[str, Any] ):
							super().tearDown()
							gc.collect()
							torch.cuda.empty_cache()

			def  lowerCAmelCase__		(    self	:    List[str]  ,		snake_case_	:    List[str]  ,		snake_case_	:    Tuple="cpu"  ,		snake_case_	:    Optional[Any]=torch.floataa  ,		snake_case_	:    Optional[int]=0 ):
							UpperCamelCase_:							Any		=  torch.Generator(device=__lowerCamelCase ).manual_seed(__lowerCamelCase )
							UpperCamelCase_:							Any		=  np.random.RandomState(__lowerCamelCase ).standard_normal((1, 4, 64, 64) )
							UpperCamelCase_:							Optional[int]		=  torch.from_numpy(__lowerCamelCase ).to(device=__lowerCamelCase  ,		dtype=__lowerCamelCase )
							UpperCamelCase_:							Dict		=  {
							    """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  lowerCAmelCase__		(    self	:    List[Any] ):
							UpperCamelCase_:							Optional[Any]		=  DiffusionPipeline.from_pretrained("""stabilityai/stable-diffusion-2-base""" )
							pipe.to(__lowerCamelCase )
							pipe.set_progress_bar_config(disable=__lowerCamelCase )

							UpperCamelCase_:							Optional[int]		=  self.get_inputs(__lowerCamelCase )
							UpperCamelCase_:							Optional[Any]		=  pipe(**__lowerCamelCase ).images
							UpperCamelCase_:							Optional[Any]		=  image[0, -3:, -3:, -1].flatten()

							assert image.shape == (1, 512, 512, 3)
							UpperCamelCase_:							List[str]		=  np.array([0.4_9493, 0.4_7896, 0.4_0798, 0.5_4214, 0.5_3212, 0.4_8202, 0.4_7656, 0.4_6329, 0.4_8506] )
							assert np.abs(image_slice - expected_slice ).max() < 7e-3








import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


lowerCamelCase_  :	Any 				=					5_00_00
lowerCamelCase_  :	Any 				=					50_00

lowerCamelCase_				,    lowerCamelCase_  :	Tuple 				=					os.path.split(__file__)
lowerCamelCase_  :	int 				=					os.path.join(RESULTS_BASEPATH, """results""", RESULTS_FILENAME.replace(""".py""", """.json"""))







@get_duration
def     A__   (  lowerCamelCase       ,      lowerCamelCase  )					->		Tuple:
				for i in range(lowerCamelCase  ):
								UpperCamelCase_:							Dict		=  dataset[i]







@get_duration
def     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase  )					->		Tuple:
				for i in range(0       ,      len(lowerCamelCase  )       ,      lowerCamelCase  ):
								UpperCamelCase_:							List[Any]		=  dataset[i : i + batch_size]







@get_duration
def     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase  )					->		Any:
				with dataset.formatted_as(type=lowerCamelCase  ):
								for i in range(lowerCamelCase  ):
												UpperCamelCase_:							List[str]		=  dataset[i]







@get_duration
def     A__   (  lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase       ,      lowerCamelCase  )					->		Optional[int]:
				with dataset.formatted_as(type=lowerCamelCase  ):
								for i in range(0       ,      lowerCamelCase       ,      lowerCamelCase  ):
												UpperCamelCase_:							Union[str, Any]		=  dataset[i : i + batch_size]







def     A__   (  )					->		Tuple:
				UpperCamelCase_:							int		=  {"""num examples""": SPEED_TEST_N_EXAMPLES}
				UpperCamelCase_:							Union[str, Any]		=  [
				    (read, {"""length""": SMALL_TEST}),
				    (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
				    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
				    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1_00}),
				    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10_00}),
				    (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
				    (read_formatted, {"""type""": """pandas""", """length""": SMALL_TEST}),
				    (read_formatted, {"""type""": """torch""", """length""": SMALL_TEST}),
				    (read_formatted, {"""type""": """tensorflow""", """length""": SMALL_TEST}),
				    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
				    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10_00}),
				]

				UpperCamelCase_:							Tuple		=  [
				    (read, {"""length""": SMALL_TEST}),
				    (read, {"""length""": SPEED_TEST_N_EXAMPLES}),
				    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10}),
				    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 1_00}),
				    (read_batch, {"""length""": SPEED_TEST_N_EXAMPLES, """batch_size""": 10_00}),
				    (read_formatted, {"""type""": """numpy""", """length""": SMALL_TEST}),
				    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10}),
				    (read_formatted_batch, {"""type""": """numpy""", """length""": SMALL_TEST, """batch_size""": 10_00}),
				]
				with tempfile.TemporaryDirectory() as tmp_dir:
								print("""generating dataset"""  )
								UpperCamelCase_:							int		=  datasets.Features(
								    {"""list""": datasets.Sequence(datasets.Value("""float32"""  )  ), """numbers""": datasets.Value("""float32"""  )}  )
								UpperCamelCase_:							Optional[int]		=  generate_example_dataset(
								    os.path.join(lowerCamelCase       ,      """dataset.arrow"""  )       ,      lowerCamelCase       ,      num_examples=lowerCamelCase       ,      seq_shapes={"""list""": (1_00,)}       ,      )
								print("""first set of iterations"""  )
								for func, kwargs in functions:
												print(func.__name__       ,      str(lowerCamelCase  )  )
												UpperCamelCase_:							List[Any]		=  func(lowerCamelCase       ,      **lowerCamelCase  )

								print("""shuffling dataset"""  )
								UpperCamelCase_:							Dict		=  dataset.shuffle()
								print("""Second set of iterations (after shuffling"""  )
								for func, kwargs in functions_shuffled:
												print("""shuffled """       ,      func.__name__       ,      str(lowerCamelCase  )  )
												UpperCamelCase_:							List[Any]		=  func(
												    lowerCamelCase       ,      **lowerCamelCase  )

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


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






