
import gc
import unittest

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

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

from ..pipeline_params import TEXT_TO_IMAGE_BATCH_PARAMS, TEXT_TO_IMAGE_PARAMS
from ..test_pipelines_common import PipelineTesterMixin


enable_full_determinism()





class       __A(				UpperCAmelCase     ,		unittest.TestCase       ):
					SCREAMING_SNAKE_CASE			    =    LDMTextToImagePipeline
					SCREAMING_SNAKE_CASE			    =    TEXT_TO_IMAGE_PARAMS - {
					    '''negative_prompt''',
					    '''negative_prompt_embeds''',
					    '''cross_attention_kwargs''',
					    '''prompt_embeds''',
					}
					SCREAMING_SNAKE_CASE			    =    PipelineTesterMixin.required_optional_params - {
					    '''num_images_per_prompt''',
					    '''callback''',
					    '''callback_steps''',
					}
					SCREAMING_SNAKE_CASE			    =    TEXT_TO_IMAGE_BATCH_PARAMS
					SCREAMING_SNAKE_CASE			    =    False

					def    lowercase__      (	self					:     Tuple   ):
												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""")   ,					cross_attention_dim=3_2   ,					)
												lowerCamelCase_							  =					DDIMScheduler(
												    beta_start=0.00085   ,					beta_end=0.012   ,					beta_schedule="""scaled_linear"""   ,					clip_sample=__UpperCamelCase   ,					set_alpha_to_one=__UpperCamelCase   ,					)
												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   ,					)
												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   ,					)
												lowerCamelCase_							  =					CLIPTextModel(__UpperCamelCase   )
												lowerCamelCase_							  =					CLIPTokenizer.from_pretrained("""hf-internal-testing/tiny-random-clip"""   )

												lowerCamelCase_							  =					{
												    """unet""": unet,
												    """scheduler""": scheduler,
												    """vqvae""": vae,
												    """bert""": text_encoder,
												    """tokenizer""": tokenizer,
												}
												return components

					def    lowercase__      (	self					:     Tuple   ,					__UpperCamelCase					:     Any   ,					__UpperCamelCase					:     Tuple=0   ):
												if str(__UpperCamelCase   ).startswith("""mps"""   ):
																			lowerCamelCase_							  =					torch.manual_seed(__UpperCamelCase   )
												else:
																			lowerCamelCase_							  =					torch.Generator(device=__UpperCamelCase   ).manual_seed(__UpperCamelCase   )
												lowerCamelCase_							  =					{
												    """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    lowercase__      (	self					:     Optional[Any]   ):
												lowerCamelCase_							  =					"""cpu"""  # ensure determinism for the device-dependent torch.Generator

												lowerCamelCase_							  =					self.get_dummy_components()
												lowerCamelCase_							  =					LDMTextToImagePipeline(**__UpperCamelCase   )
												pipe.to(__UpperCamelCase   )
												pipe.set_progress_bar_config(disable=__UpperCamelCase   )

												lowerCamelCase_							  =					self.get_dummy_inputs(__UpperCamelCase   )
												lowerCamelCase_							  =					pipe(**__UpperCamelCase   ).images
												lowerCamelCase_							  =					image[0, -3:, -3:, -1]

												assert image.shape == (1, 1_6, 1_6, 3)
												lowerCamelCase_							  =					np.array([0.6101, 0.6156, 0.5622, 0.4895, 0.6661, 0.3804, 0.5748, 0.6136, 0.5014]   )

												assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-3





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

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

					def    lowercase__      (	self					:     List[Any]   ,					__UpperCamelCase					:     str   ,					__UpperCamelCase					:     Tuple=torch.floataa   ,					__UpperCamelCase					:     Dict=0   ):
												lowerCamelCase_							  =					torch.manual_seed(__UpperCamelCase   )
												lowerCamelCase_							  =					np.random.RandomState(__UpperCamelCase   ).standard_normal((1, 4, 3_2, 3_2)   )
												lowerCamelCase_							  =					torch.from_numpy(__UpperCamelCase   ).to(device=__UpperCamelCase   ,					dtype=__UpperCamelCase   )
												lowerCamelCase_							  =					{
												    """prompt""": """A painting of a squirrel eating a burger""",
												    """latents""": latents,
												    """generator""": generator,
												    """num_inference_steps""": 3,
												    """guidance_scale""": 6.0,
												    """output_type""": """numpy""",
												}
												return inputs



					def    lowercase__      (	self					:     Optional[int]   ):
												lowerCamelCase_							  =					LDMTextToImagePipeline.from_pretrained("""CompVis/ldm-text2im-large-256"""   ).to(__UpperCamelCase   )
												pipe.set_progress_bar_config(disable=__UpperCamelCase   )

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

												assert image.shape == (1, 2_5_6, 2_5_6, 3)
												lowerCamelCase_							  =					np.array([0.51825, 0.52850, 0.52543, 0.54258, 0.52304, 0.52569, 0.54363, 0.55276, 0.56878]   )
												lowerCamelCase_							  =					np.abs(expected_slice - image_slice   ).max()
												assert max_diff < 1E-3





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

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

					def    lowercase__      (	self					:     Optional[Any]   ,					__UpperCamelCase					:     Union[str, Any]   ,					__UpperCamelCase					:     Any=torch.floataa   ,					__UpperCamelCase					:     Tuple=0   ):
												lowerCamelCase_							  =					torch.manual_seed(__UpperCamelCase   )
												lowerCamelCase_							  =					np.random.RandomState(__UpperCamelCase   ).standard_normal((1, 4, 3_2, 3_2)   )
												lowerCamelCase_							  =					torch.from_numpy(__UpperCamelCase   ).to(device=__UpperCamelCase   ,					dtype=__UpperCamelCase   )
												lowerCamelCase_							  =					{
												    """prompt""": """A painting of a squirrel eating a burger""",
												    """latents""": latents,
												    """generator""": generator,
												    """num_inference_steps""": 5_0,
												    """guidance_scale""": 6.0,
												    """output_type""": """numpy""",
												}
												return inputs



					def    lowercase__      (	self					:     int   ):
												lowerCamelCase_							  =					LDMTextToImagePipeline.from_pretrained("""CompVis/ldm-text2im-large-256"""   ).to(__UpperCamelCase   )
												pipe.set_progress_bar_config(disable=__UpperCamelCase   )

												lowerCamelCase_							  =					self.get_inputs(__UpperCamelCase   )
												lowerCamelCase_							  =					pipe(**__UpperCamelCase   ).images[0]

												lowerCamelCase_							  =					load_numpy(
												    """https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/ldm_text2img/ldm_large_256_ddim.npy"""   )
												lowerCamelCase_							  =					np.abs(expected_image - image   ).max()
												assert max_diff < 1E-3







import os
import socket
from contextlib import contextmanager

import torch

from ..commands.config.default import write_basic_config  # noqa: F401
from ..state import PartialState
from .dataclasses import DistributedType
from .imports import is_deepspeed_available, is_tpu_available
from .transformer_engine import convert_model
from .versions import is_torch_version


if is_deepspeed_available():
		from deepspeed import DeepSpeedEngine

if is_tpu_available(check_device=False):
		import torch_xla.core.xla_model as xm



def      __lowerCAmelCase						(							UpperCAmelCase__				:      Any       )  ->    List[str]:
							if is_torch_version("""<"""			,				"""2.0.0"""       ) or not hasattr(UpperCAmelCase__			,				"""_dynamo"""       ):
														return False
							return isinstance(UpperCAmelCase__			,				torch._dynamo.eval_frame.OptimizedModule       )



def      __lowerCAmelCase						(							UpperCAmelCase__				:      str			,				UpperCAmelCase__				:      bool = True       )  ->    str:
							lowerCamelCase_							  =					(torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

							lowerCamelCase_							  =					is_compiled_module(UpperCAmelCase__       )
							if is_compiled:
														lowerCamelCase_							  =					model
														lowerCamelCase_							  =					model._orig_mod

							if is_deepspeed_available():
														options += (DeepSpeedEngine,)

							while isinstance(UpperCAmelCase__			,				UpperCAmelCase__       ):
														lowerCamelCase_							  =					model.module

							if not keep_fpaa_wrapper:
														lowerCamelCase_							  =					getattr(UpperCAmelCase__			,				"""forward"""       )
														lowerCamelCase_							  =					model.__dict__.pop("""_original_forward"""			,				UpperCAmelCase__       )
														if original_forward is not None:
																					while hasattr(UpperCAmelCase__			,				"""__wrapped__"""       ):
																												lowerCamelCase_							  =					forward.__wrapped__
																												if forward == original_forward:
																																			break
																					lowerCamelCase_							  =					forward
														if getattr(UpperCAmelCase__			,				"""_converted_to_transformer_engine"""			,				UpperCAmelCase__       ):
																					convert_model(UpperCAmelCase__			,				to_transformer_engine=UpperCAmelCase__       )

							if is_compiled:
														lowerCamelCase_							  =					model
														lowerCamelCase_							  =					compiled_model

							return model



def      __lowerCAmelCase						(							)  ->    Any:
							PartialState().wait_for_everyone()



def      __lowerCAmelCase						(							UpperCAmelCase__				:      Union[str, Any]			,				UpperCAmelCase__				:      Tuple       )  ->    Dict:
							if PartialState().distributed_type == DistributedType.TPU:
														xm.save(UpperCAmelCase__			,				UpperCAmelCase__       )
							elif PartialState().local_process_index == 0:
														torch.save(UpperCAmelCase__			,				UpperCAmelCase__       )



@contextmanager
def      __lowerCAmelCase						(							**UpperCAmelCase__				:      Any       )  ->    Dict:
							for key, value in kwargs.items():
														lowerCamelCase_							  =					str(UpperCAmelCase__       )

							yield

							for key in kwargs:
														if key.upper() in os.environ:
																					del os.environ[key.upper()]



def      __lowerCAmelCase						(							UpperCAmelCase__				:      Union[str, Any]       )  ->    str:
							if not hasattr(UpperCAmelCase__			,				"""__qualname__"""       ) and not hasattr(UpperCAmelCase__			,				"""__name__"""       ):
														lowerCamelCase_							  =					getattr(UpperCAmelCase__			,				"""__class__"""			,				UpperCAmelCase__       )
							if hasattr(UpperCAmelCase__			,				"""__qualname__"""       ):
														return obj.__qualname__
							if hasattr(UpperCAmelCase__			,				"""__name__"""       ):
														return obj.__name__
							return str(UpperCAmelCase__       )



def      __lowerCAmelCase						(							UpperCAmelCase__				:      str			,				UpperCAmelCase__				:      Union[str, Any]       )  ->    Dict:
							for key, value in source.items():
														if isinstance(UpperCAmelCase__			,				UpperCAmelCase__       ):
																					lowerCamelCase_							  =					destination.setdefault(UpperCAmelCase__			,				{}       )
																					merge_dicts(UpperCAmelCase__			,				UpperCAmelCase__       )
														else:
																					lowerCamelCase_							  =					value

							return destination



def      __lowerCAmelCase						(							UpperCAmelCase__				:      int = None       )  ->    bool:
							if port is None:
														lowerCamelCase_							  =					2_9_5_0_0
							with socket.socket(socket.AF_INET			,				socket.SOCK_STREAM       ) as s:
														return s.connect_ex(("""localhost""", port)       ) == 0





