


'''simple docstring'''


import gc
import random
import unittest

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

import diffusers
from diffusers import (
    AutoencoderKL,
    EulerDiscreteScheduler,
    StableDiffusionLatentUpscalePipeline,
    StableDiffusionPipeline,
    UNetaDConditionModel,
)
from diffusers.schedulers import KarrasDiffusionSchedulers
from diffusers.utils import floats_tensor, load_image, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS, TEXT_GUIDED_IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()







def 			snake_case_   (  SCREAMING_SNAKE_CASE__     ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE :       str     		=      [tensor.shape for tensor in tensor_list]
				return all(shape == shapes[0] for shape in shapes[1:]     )





class     lowercase__  (						lowercase_     ,     lowercase_     ,     lowercase_     ,     unittest.TestCase					):





							'''simple docstring'''


							A_							:      Union[str, Any]      	=  StableDiffusionLatentUpscalePipeline
							A_							:      str      	=  TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {
							    '''height''',
							    '''width''',
							    '''cross_attention_kwargs''',
							    '''negative_prompt_embeds''',
							    '''prompt_embeds''',
							}
							A_							:      Dict      	=  PipelineTesterMixin.required_optional_params - {'''num_images_per_prompt'''}
							A_							:      Optional[int]      	=  TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							A_							:      Optional[int]      	=  frozenset(
							    []					)  # TO-DO: update image_params once pipeline is refactored with VaeImageProcessor.preprocess
							A_							:      Any      	=  frozenset([]					)

							A_							:      Any      	=  True






							@property
							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       str     		=      1
											_SCREAMING_SNAKE_CASE :       Dict     		=      4
											_SCREAMING_SNAKE_CASE :       str     		=      (16, 16)

											_SCREAMING_SNAKE_CASE :       int     		=      floats_tensor((batch_size, num_channels) + sizes    ,					rng=random.Random(0  )  ).to(__snake_case  )
											return image






							def UpperCAmelCase_	(    self  ):
											torch.manual_seed(0  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      UNetaDConditionModel(
											    act_fn="""gelu"""    ,					attention_head_dim=8    ,					norm_num_groups=__snake_case    ,					block_out_channels=[32, 32, 64, 64]    ,					time_cond_proj_dim=160    ,					conv_in_kernel=1    ,					conv_out_kernel=1    ,					cross_attention_dim=32    ,					down_block_types=(
											        """KDownBlock2D""",
											        """KCrossAttnDownBlock2D""",
											        """KCrossAttnDownBlock2D""",
											        """KCrossAttnDownBlock2D""",
											    )    ,					in_channels=8    ,					mid_block_type=__snake_case    ,					only_cross_attention=__snake_case    ,					out_channels=5    ,					resnet_time_scale_shift="""scale_shift"""    ,					time_embedding_type="""fourier"""    ,					timestep_post_act="""gelu"""    ,					up_block_types=("""KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KCrossAttnUpBlock2D""", """KUpBlock2D""")    ,					)
											_SCREAMING_SNAKE_CASE :       int     		=      AutoencoderKL(
											    block_out_channels=[32, 32, 64, 64]    ,					in_channels=3    ,					out_channels=3    ,					down_block_types=[
											        """DownEncoderBlock2D""",
											        """DownEncoderBlock2D""",
											        """DownEncoderBlock2D""",
											        """DownEncoderBlock2D""",
											    ]    ,					up_block_types=["""UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D""", """UpDecoderBlock2D"""]    ,					latent_channels=4    ,					)
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      EulerDiscreteScheduler(prediction_type="""sample"""  )
											_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      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="""quick_gelu"""    ,					projection_dim=512    ,					)
											_SCREAMING_SNAKE_CASE :       List[str]     		=      CLIPTextModel(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Dict     		=      CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""  )

											_SCREAMING_SNAKE_CASE :       Any     		=      {
											    """unet""": model.eval(),
											    """vae""": vae.eval(),
											    """scheduler""": scheduler,
											    """text_encoder""": text_encoder,
											    """tokenizer""": tokenizer,
											}

											return components






							def UpperCAmelCase_	(    self    ,					__snake_case    ,					__snake_case=0  ):
											if str(__snake_case  ).startswith("""mps"""  ):
															_SCREAMING_SNAKE_CASE :       str     		=      torch.manual_seed(__snake_case  )
											else:
															_SCREAMING_SNAKE_CASE :       Union[str, Any]     		=      torch.Generator(device=__snake_case  ).manual_seed(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Any     		=      {
											    """prompt""": """A painting of a squirrel eating a burger""",
											    """image""": self.dummy_image.cpu(),
											    """generator""": generator,
											    """num_inference_steps""": 2,
											    """output_type""": """numpy""",
											}
											return inputs






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      """cpu"""

											_SCREAMING_SNAKE_CASE :       int     		=      self.get_dummy_components()
											_SCREAMING_SNAKE_CASE :       int     		=      self.pipeline_class(**__snake_case  )
											pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )

											_SCREAMING_SNAKE_CASE :       int     		=      self.get_dummy_inputs(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      pipe(**__snake_case  ).images
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      image[0, -3:, -3:, -1]

											self.assertEqual(image.shape    ,					(1, 256, 256, 3)  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      np.array(
											    [0.47222412, 0.41921633, 0.44717434, 0.46874192, 0.42588258, 0.46150726, 0.4677534, 0.45583832, 0.48579055]  )
											_SCREAMING_SNAKE_CASE :       Any     		=      np.abs(image_slice.flatten() - expected_slice  ).max()
											self.assertLessEqual(__snake_case    ,					1e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_attention_slicing_forward_pass(expected_max_diff=7e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_cpu_offload_forward_pass(expected_max_diff=3e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_dict_tuple_outputs_equivalent(expected_max_difference=3e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_inference_batch_single_identical(expected_max_diff=7e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_pt_np_pil_outputs_equivalent(expected_max_diff=3e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_save_load_local(expected_max_difference=3e-3  )






							def UpperCAmelCase_	(    self  ):
											super().test_save_load_optional_components(expected_max_difference=3e-3  )






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       List[str]     		=      [
											    """DDIMScheduler""",
											    """DDPMScheduler""",
											    """PNDMScheduler""",
											    """HeunDiscreteScheduler""",
											    """EulerAncestralDiscreteScheduler""",
											    """KDPM2DiscreteScheduler""",
											    """KDPM2AncestralDiscreteScheduler""",
											    """DPMSolverSDEScheduler""",
											]
											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      self.get_dummy_components()
											_SCREAMING_SNAKE_CASE :       int     		=      self.pipeline_class(**__snake_case  )

											# make sure that PNDM does not need warm-up
											pipe.scheduler.register_to_config(skip_prk_steps=__snake_case  )

											pipe.to(__snake_case  )
											pipe.set_progress_bar_config(disable=__snake_case  )
											_SCREAMING_SNAKE_CASE :       List[str]     		=      self.get_dummy_inputs(__snake_case  )
											_SCREAMING_SNAKE_CASE :       Tuple     		=      2

											_SCREAMING_SNAKE_CASE :       List[Any]     		=      []
											for scheduler_enum in KarrasDiffusionSchedulers:
															if scheduler_enum.name in skip_schedulers:
																			# no sigma schedulers are not supported
																			# no schedulers
																			continue

															_SCREAMING_SNAKE_CASE :       Any     		=      getattr(__snake_case    ,					scheduler_enum.name  )
															_SCREAMING_SNAKE_CASE :       List[str]     		=      scheduler_cls.from_config(pipe.scheduler.config  )
															_SCREAMING_SNAKE_CASE :       Any     		=      pipe(**__snake_case  )[0]
															outputs.append(__snake_case  )

											assert check_same_shape(__snake_case  )





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





							'''simple docstring'''


							def UpperCAmelCase_	(    self  ):
											super().tearDown()
											gc.collect()
											torch.cuda.empty_cache()






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       int     		=      torch.manual_seed(33  )

											_SCREAMING_SNAKE_CASE :       Any     		=      StableDiffusionPipeline.from_pretrained("""CompVis/stable-diffusion-v1-4"""    ,					torch_dtype=torch.floataa  )
											pipe.to("""cuda"""  )

											_SCREAMING_SNAKE_CASE :       Any     		=      StableDiffusionLatentUpscalePipeline.from_pretrained(
											    """stabilityai/sd-x2-latent-upscaler"""    ,					torch_dtype=torch.floataa  )
											upscaler.to("""cuda"""  )

											_SCREAMING_SNAKE_CASE :       str     		=      """a photo of an astronaut high resolution, unreal engine, ultra realistic"""

											_SCREAMING_SNAKE_CASE :       Any     		=      pipe(__snake_case    ,					generator=__snake_case    ,					output_type="""latent"""  ).images

											_SCREAMING_SNAKE_CASE :       Dict     		=      upscaler(
											    prompt=__snake_case    ,					image=__snake_case    ,					num_inference_steps=20    ,					guidance_scale=0    ,					generator=__snake_case    ,					output_type="""np"""    ,					).images[0]

											_SCREAMING_SNAKE_CASE :       Optional[int]     		=      load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/astronaut_1024.npy"""  )
											assert np.abs((expected_image - image).mean()  ) < 5e-2






							def UpperCAmelCase_	(    self  ):
											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      torch.manual_seed(33  )

											_SCREAMING_SNAKE_CASE :       Optional[Any]     		=      StableDiffusionLatentUpscalePipeline.from_pretrained(
											    """stabilityai/sd-x2-latent-upscaler"""    ,					torch_dtype=torch.floataa  )
											upscaler.to("""cuda"""  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      """the temple of fire by Ross Tran and Gerardo Dottori, oil on canvas"""

											_SCREAMING_SNAKE_CASE :       Dict     		=      load_image(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_512.png"""  )

											_SCREAMING_SNAKE_CASE :       Tuple     		=      upscaler(
											    prompt=__snake_case    ,					image=__snake_case    ,					num_inference_steps=20    ,					guidance_scale=0    ,					generator=__snake_case    ,					output_type="""np"""    ,					).images[0]

											_SCREAMING_SNAKE_CASE :       Dict     		=      load_numpy(
											    """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/latent-upscaler/fire_temple_1024.npy"""  )
											assert np.abs((expected_image - image).max()  ) < 5e-2







import os

import numpy
import onnx





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    a.name
							snake_case_							=    b.name

							snake_case_							=    ''
							snake_case_							=    ''

							snake_case_							=    a == b

							snake_case_							=    name_a
							snake_case_							=    name_b

							return res





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							for i, input_name in enumerate(node_proto.input      ):
														if input_name == name:
																					node_proto.input.insert(UpperCamelCase__   ,							UpperCamelCase__      )
																					node_proto.input.pop(i + 1      )

							if node_proto.op_type == "If":
														_graph_replace_input_with(node_proto.attribute[0].g   ,							UpperCamelCase__   ,							UpperCamelCase__      )
														_graph_replace_input_with(node_proto.attribute[1].g   ,							UpperCamelCase__   ,							UpperCamelCase__      )
							if node_proto.op_type == "Loop":
														_graph_replace_input_with(node_proto.attribute[0].g   ,							UpperCamelCase__   ,							UpperCamelCase__      )





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							for n in graph_proto.node:
														_node_replace_input_with(UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      )





def 					__lowerCamelCase	( UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    list(model.graph.initializer      )
							snake_case_							=    list(model_without_ext.graph.initializer      )
							for i, ref_i in ind_to_replace:
														assert inits_with_data[i].name == inits[i].name
														assert inits_with_data[ref_i].name == inits[ref_i].name
														assert i > ref_i

														snake_case_							=    inits[i].name
														snake_case_							=    inits[ref_i].name

														model_without_ext.graph.initializer.remove(inits[i]      )

														# for n in model.graph.node:
														_graph_replace_input_with(model_without_ext.graph   ,							UpperCamelCase__   ,							UpperCamelCase__      )





def 					__lowerCamelCase	( UpperCamelCase__      ):







							'''simple docstring'''






							snake_case_							=    os.path.dirname(UpperCamelCase__      )
							snake_case_							=    os.path.basename(UpperCamelCase__      )

							snake_case_							=    onnx.load(os.path.join(UpperCamelCase__   ,							UpperCamelCase__      )      )

							snake_case_							=    list(model.graph.initializer      )

							snake_case_							=    set()
							snake_case_							=    {}
							snake_case_							=    []

							snake_case_							=    0

							for i in range(len(UpperCamelCase__      )      ):
														if i in dup_set:
																					continue

														for j in range(i + 1   ,							len(UpperCamelCase__      )      ):
																					if j in dup_set:
																												continue
																					if _is_equal_tensor_proto(inits[i]   ,							inits[j]      ):
																												dup_set.add(UpperCamelCase__      )
																												dup_set.add(UpperCamelCase__      )

																												snake_case_							=    inits[j].data_type
																												snake_case_							=    numpy.prod(inits[j].dims      )
																												if dtype == 1:
																																			mem_size *= 4
																												elif dtype == 6:
																																			mem_size *= 4
																												elif dtype == 7 or dtype == 11:
																																			mem_size *= 8
																												else:
																																			print('unexpected data type: '   ,							UpperCamelCase__      )
																												total_reduced_size += mem_size

																												snake_case_							=    inits[i].name
																												snake_case_							=    inits[j].name

																												if name_i in dup_map:
																																			dup_map[name_i].append(UpperCamelCase__      )
																												else:
																																			snake_case_							=    [name_j]
																												ind_to_replace.append((j, i)      )

							print('total reduced size: '   ,							total_reduced_size / 1024 / 1024 / 1024   ,							'GB'      )

							snake_case_							=    sorted(UpperCamelCase__      )
							_remove_dup_initializers_from_model(UpperCamelCase__   ,							UpperCamelCase__   ,							UpperCamelCase__      )

							snake_case_							=    'optimized_' + model_file_name
							snake_case_							=    os.path.join(UpperCamelCase__   ,							UpperCamelCase__      )
							onnx.save(UpperCamelCase__   ,							UpperCamelCase__      )

							return new_model




