





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       __A					(					lowerCAmelCase    ,       lowerCAmelCase    ,       unittest.TestCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		StableDiffusionXLImgaImgPipeline
			lowerCAmelCase_							  =		TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"""height""", """width"""}
			lowerCAmelCase_							  =		PipelineTesterMixin.required_optional_params - {"""latents"""}
			lowerCAmelCase_							  =		TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
			lowerCAmelCase_							  =		IMAGE_TO_IMAGE_IMAGE_PARAMS
			lowerCAmelCase_							  =		IMAGE_TO_IMAGE_IMAGE_PARAMS





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				torch.manual_seed(0							)
				lowerCamelCase__              =	UNetaDConditionModel(
				    block_out_channels=(3_2, 6_4) ,						layers_per_block=2 ,						sample_size=3_2 ,						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=8_0 ,						cross_attention_dim=6_4 ,						)
				lowerCamelCase__              =	EulerDiscreteScheduler(
				    beta_start=0.0_0085 ,						beta_end=0.012 ,						steps_offset=1 ,						beta_schedule='''scaled_linear''' ,						timestep_spacing='''leading''' ,						)
				torch.manual_seed(0							)
				lowerCamelCase__              =	AutoencoderKL(
				    block_out_channels=[3_2, 6_4] ,						in_channels=3 ,						out_channels=3 ,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D'''] ,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D'''] ,						latent_channels=4 ,						sample_size=1_2_8 ,						)
				torch.manual_seed(0							)
				lowerCamelCase__              =	CLIPTextConfig(
				    bos_token_id=0 ,						eos_token_id=2 ,						hidden_size=3_2 ,						intermediate_size=3_7 ,						layer_norm_eps=1E-05 ,						num_attention_heads=4 ,						num_hidden_layers=5 ,						pad_token_id=1 ,						vocab_size=1_0_0_0 ,						hidden_act='''gelu''' ,						projection_dim=3_2 ,						)
				lowerCamelCase__              =	CLIPTextModel(__lowerCAmelCase							)
				lowerCamelCase__              =	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' ,						local_files_only=__lowerCAmelCase							)

				lowerCamelCase__              =	CLIPTextModelWithProjection(__lowerCAmelCase							)
				lowerCamelCase__              =	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip''' ,						local_files_only=__lowerCAmelCase							)

				lowerCamelCase__              =	{
				    '''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 ,						__lowerCAmelCase ,						__lowerCAmelCase=0							):



				'''simple docstring'''







				lowerCamelCase__              =	floats_tensor((1, 3, 3_2, 3_2) ,						rng=random.Random(__lowerCAmelCase							)							).to(__lowerCAmelCase							)
				lowerCamelCase__              =	image / 2 + 0.5
				if str(__lowerCAmelCase							).startswith('''mps'''							):
					lowerCamelCase__              =	torch.manual_seed(__lowerCAmelCase							)
				else:
					lowerCamelCase__              =	torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
				lowerCamelCase__              =	{
				    '''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							):



				'''simple docstring'''







				lowerCamelCase__              =	'''cpu'''  # ensure determinism for the device-dependent torch.Generator
				lowerCamelCase__              =	self.get_dummy_components()
				lowerCamelCase__              =	StableDiffusionXLImgaImgPipeline(**__lowerCAmelCase							)
				lowerCamelCase__              =	sd_pipe.to(__lowerCAmelCase							)
				sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__              =	self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__              =	sd_pipe(**__lowerCAmelCase							).images
				lowerCamelCase__              =	image[0, -3:, -3:, -1]

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

				lowerCamelCase__              =	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							):



				'''simple docstring'''







				super().test_attention_slicing_forward_pass(expected_max_diff=3E-3							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				super().test_inference_batch_single_identical(expected_max_diff=3E-3							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				pass





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.get_dummy_components()
				lowerCamelCase__              =	StableDiffusionXLImgaImgPipeline(**__lowerCAmelCase							)
				lowerCamelCase__              =	sd_pipe.to(__lowerCAmelCase							)
				lowerCamelCase__              =	sd_pipe.to(__lowerCAmelCase							)
				sd_pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				# forward without prompt embeds
				lowerCamelCase__              =	self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__              =	3 * ['''this is a negative prompt''']
				lowerCamelCase__              =	negative_prompt
				lowerCamelCase__              =	3 * [inputs['''prompt''']]

				lowerCamelCase__              =	sd_pipe(**__lowerCAmelCase							)
				lowerCamelCase__              =	output.images[0, -3:, -3:, -1]

				# forward with prompt embeds
				lowerCamelCase__              =	self.get_dummy_inputs(__lowerCAmelCase							)
				lowerCamelCase__              =	3 * ['''this is a negative prompt''']
				lowerCamelCase__              =	3 * [inputs.pop('''prompt'''							)]

				(
				    (
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			
				)              =	sd_pipe.encode_prompt(__lowerCAmelCase ,						negative_prompt=__lowerCAmelCase							)

				lowerCamelCase__              =	sd_pipe(
				    **__lowerCAmelCase ,						prompt_embeds=__lowerCAmelCase ,						negative_prompt_embeds=__lowerCAmelCase ,						pooled_prompt_embeds=__lowerCAmelCase ,						negative_pooled_prompt_embeds=__lowerCAmelCase ,						)
				lowerCamelCase__              =	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       __A					(					unittest.TestCase				):

			'''simple docstring'''







			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				super().tearDown()
				gc.collect()
				torch.cuda.empty_cache()





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase="cpu" ,						__lowerCAmelCase=torch.floataa ,						__lowerCAmelCase=0							):



				'''simple docstring'''







				lowerCamelCase__              =	torch.Generator(device=__lowerCAmelCase							).manual_seed(__lowerCAmelCase							)
				lowerCamelCase__              =	np.random.RandomState(__lowerCAmelCase							).standard_normal((1, 4, 6_4, 6_4)							)
				lowerCamelCase__              =	torch.from_numpy(__lowerCAmelCase							).to(device=__lowerCAmelCase ,						dtype=__lowerCAmelCase							)
				lowerCamelCase__              =	{
				    '''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							):



				'''simple docstring'''







				lowerCamelCase__              =	DiffusionPipeline.from_pretrained('''stabilityai/stable-diffusion-2-base'''							)
				pipe.to(__lowerCAmelCase							)
				pipe.set_progress_bar_config(disable=__lowerCAmelCase							)

				lowerCamelCase__              =	self.get_inputs(__lowerCAmelCase							)
				lowerCamelCase__              =	pipe(**__lowerCAmelCase							).images
				lowerCamelCase__              =	image[0, -3:, -3:, -1].flatten()

				assert image.shape == (1, 5_1_2, 5_1_2, 3)
				lowerCamelCase__              =	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










def lowerCAmelCase__(__snake_case							)       ->   int:







	'''simple docstring'''







	if not grid or not grid[0]:
		raise TypeError('''The grid does not contain the appropriate information'''							)

	for cell_n in range(1 ,len(grid[0]							)							):
		grid[0][cell_n] += grid[0][cell_n - 1]
	lowerCamelCase__              =	grid[0]

	for row_n in range(1 ,len(__snake_case							)							):
		lowerCamelCase__              =	grid[row_n]
		lowerCamelCase__              =	fill_row(__snake_case ,__snake_case							)
		lowerCamelCase__              =	grid[row_n]

	return grid[-1][-1]





def lowerCAmelCase__(__snake_case ,__snake_case							)       ->   list:







	'''simple docstring'''







	current_row[0] += row_above[0]
	for cell_n in range(1 ,len(__snake_case							)							):
		current_row[cell_n] += min(current_row[cell_n - 1] ,row_above[cell_n]							)

	return current_row


if __name__ == "__main__":
			import doctest

			doctest.testmod()



