




'''simple docstring'''




import gc
import random
import unittest

import numpy as np
import torch
from transformers import (
    CLIPImageProcessor,
    CLIPTextConfig,
    CLIPTextModelWithProjection,
    CLIPTokenizer,
    CLIPVisionConfig,
    CLIPVisionModelWithProjection,
)

from diffusers import (
    DiffusionPipeline,
    UnCLIPImageVariationPipeline,
    UnCLIPScheduler,
    UNetaDConditionModel,
    UNetaDModel,
)
from diffusers.pipelines.unclip.text_proj import UnCLIPTextProjModel
from diffusers.utils import floats_tensor, load_numpy, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, load_image, require_torch_gpu, skip_mps

from ..pipeline_params import IMAGE_VARIATION_BATCH_PARAMS, IMAGE_VARIATION_PARAMS
from ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()


class   __UpperCamelCase       ( lowerCamelCase__			,					unittest.TestCase				):
       lowercase	:    str      						=UnCLIPImageVariationPipeline
       lowercase	:    Optional[Any]      						=IMAGE_VARIATION_PARAMS - {'height', 'width', 'guidance_scale'}
       lowercase	:    List[str]      						=IMAGE_VARIATION_BATCH_PARAMS

       lowercase	:    Optional[int]      						=[
           'generator',
           'return_dict',
           'decoder_num_inference_steps',
           'super_res_num_inference_steps',
       ]
       lowercase	:    int      						=False






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              return 32






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              return 32






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              return self.time_input_dim






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              return self.time_input_dim * 4






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              return 100






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''  )
              return tokenizer






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(0  )
              lowerCamelCase_							=CLIPTextConfig(
                  bos_token_id=0,		eos_token_id=2,		hidden_size=self.text_embedder_hidden_size,		projection_dim=self.text_embedder_hidden_size,		intermediate_size=37,		layer_norm_eps=1e-05,		num_attention_heads=4,		num_hidden_layers=5,		pad_token_id=1,		vocab_size=1_000,		)
              return CLIPTextModelWithProjection(lowerCAmelCase  )






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(0  )
              lowerCamelCase_							=CLIPVisionConfig(
                  hidden_size=self.text_embedder_hidden_size,		projection_dim=self.text_embedder_hidden_size,		num_hidden_layers=5,		num_attention_heads=4,		image_size=32,		intermediate_size=37,		patch_size=1,		)
              return CLIPVisionModelWithProjection(lowerCAmelCase  )






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(0  )

              lowerCamelCase_							={
                  '''clip_embeddings_dim''': self.text_embedder_hidden_size,
                  '''time_embed_dim''': self.time_embed_dim,
                  '''cross_attention_dim''': self.cross_attention_dim,
              }

              lowerCamelCase_							=UnCLIPTextProjModel(**lowerCAmelCase  )
              return model






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(0  )

              lowerCamelCase_							={
                  '''sample_size''': 32,
                  # RGB in channels
                  '''in_channels''': 3,
                  # Out channels is double in channels because predicts mean and variance
                  '''out_channels''': 6,
                  '''down_block_types''': ('''ResnetDownsampleBlock2D''', '''SimpleCrossAttnDownBlock2D'''),
                  '''up_block_types''': ('''SimpleCrossAttnUpBlock2D''', '''ResnetUpsampleBlock2D'''),
                  '''mid_block_type''': '''UNetMidBlock2DSimpleCrossAttn''',
                  '''block_out_channels''': (self.block_out_channels_a, self.block_out_channels_a * 2),
                  '''layers_per_block''': 1,
                  '''cross_attention_dim''': self.cross_attention_dim,
                  '''attention_head_dim''': 4,
                  '''resnet_time_scale_shift''': '''scale_shift''',
                  '''class_embed_type''': '''identity''',
              }

              lowerCamelCase_							=UNetaDConditionModel(**lowerCAmelCase  )
              return model






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              return {
                  "sample_size": 64,
                  "layers_per_block": 1,
                  "down_block_types": ("ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D"),
                  "up_block_types": ("ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D"),
                  "block_out_channels": (self.block_out_channels_a, self.block_out_channels_a * 2),
                  "in_channels": 6,
                  "out_channels": 3,
              }






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(0  )

              lowerCamelCase_							=UNetaDModel(**self.dummy_super_res_kwargs  )
              return model






       @property
       def      lowercase__     (							self  ):



              """simple docstring"""







              torch.manual_seed(1  )

              lowerCamelCase_							=UNetaDModel(**self.dummy_super_res_kwargs  )
              return model






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=self.dummy_decoder
              lowerCamelCase_							=self.dummy_text_proj
              lowerCamelCase_							=self.dummy_text_encoder
              lowerCamelCase_							=self.dummy_tokenizer
              lowerCamelCase_							=self.dummy_super_res_first
              lowerCamelCase_							=self.dummy_super_res_last

              lowerCamelCase_							=UnCLIPScheduler(
                  variance_type='''learned_range''',		prediction_type='''epsilon''',		num_train_timesteps=1_000,		)

              lowerCamelCase_							=UnCLIPScheduler(
                  variance_type='''fixed_small_log''',		prediction_type='''epsilon''',		num_train_timesteps=1_000,		)

              lowerCamelCase_							=CLIPImageProcessor(crop_size=32,		size=32  )

              lowerCamelCase_							=self.dummy_image_encoder

              return {
                  "decoder": decoder,
                  "text_encoder": text_encoder,
                  "tokenizer": tokenizer,
                  "text_proj": text_proj,
                  "feature_extractor": feature_extractor,
                  "image_encoder": image_encoder,
                  "super_res_first": super_res_first,
                  "super_res_last": super_res_last,
                  "decoder_scheduler": decoder_scheduler,
                  "super_res_scheduler": super_res_scheduler,
              }






       def      lowercase__     (							self,		lowerCAmelCase,		lowerCAmelCase=0,		lowerCAmelCase=True  ):



              """simple docstring"""







              lowerCamelCase_							=floats_tensor((1, 3, 32, 32),		rng=random.Random(lowerCAmelCase  )  ).to(lowerCAmelCase  )
              if str(lowerCAmelCase  ).startswith('''mps'''  ):
                     lowerCamelCase_							=torch.manual_seed(lowerCAmelCase  )
              else:
                     lowerCamelCase_							=torch.Generator(device=lowerCAmelCase  ).manual_seed(lowerCAmelCase  )

              if pil_image:
                     lowerCamelCase_							=input_image * 0.5 + 0.5
                     lowerCamelCase_							=input_image.clamp(0,		1  )
                     lowerCamelCase_							=input_image.cpu().permute(0,		2,		3,		1  ).float().numpy()
                     lowerCamelCase_							=DiffusionPipeline.numpy_to_pil(lowerCAmelCase  )[0]

              return {
                  "image": input_image,
                  "generator": generator,
                  "decoder_num_inference_steps": 2,
                  "super_res_num_inference_steps": 2,
                  "output_type": "np",
              }






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							='''cpu'''

              lowerCamelCase_							=self.get_dummy_components()

              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              lowerCamelCase_							=pipe.to(lowerCAmelCase  )

              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )

              lowerCamelCase_							=pipe(**lowerCAmelCase  )
              lowerCamelCase_							=output.images

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )

              lowerCamelCase_							=pipe(
                  **lowerCAmelCase,		return_dict=lowerCAmelCase,		)[0]

              lowerCamelCase_							=image[0, -3:, -3:, -1]
              lowerCamelCase_							=image_from_tuple[0, -3:, -3:, -1]

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

              lowerCamelCase_							=np.array(
                  [
                      0.9_9_9_7,
                      0.0_0_0_2,
                      0.9_9_9_7,
                      0.9_9_9_7,
                      0.9_9_6_9,
                      0.0_0_2_3,
                      0.9_9_9_7,
                      0.9_9_6_9,
                      0.9_9_7_0,
                  ]  )

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






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							='''cpu'''

              lowerCamelCase_							=self.get_dummy_components()

              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              lowerCamelCase_							=pipe.to(lowerCAmelCase  )

              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )

              lowerCamelCase_							=pipe(**lowerCAmelCase  )
              lowerCamelCase_							=output.images

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )

              lowerCamelCase_							=pipe(
                  **lowerCAmelCase,		return_dict=lowerCAmelCase,		)[0]

              lowerCamelCase_							=image[0, -3:, -3:, -1]
              lowerCamelCase_							=image_from_tuple[0, -3:, -3:, -1]

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

              lowerCamelCase_							=np.array([0.9_9_9_7, 0.0_0_0_3, 0.9_9_9_7, 0.9_9_9_7, 0.9_9_7_0, 0.0_0_2_4, 0.9_9_9_7, 0.9_9_7_1, 0.9_9_7_1]  )

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






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							='''cpu'''

              lowerCamelCase_							=self.get_dummy_components()

              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              lowerCamelCase_							=pipe.to(lowerCAmelCase  )

              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )
              lowerCamelCase_							=[
                  pipeline_inputs['''image'''],
                  pipeline_inputs['''image'''],
              ]

              lowerCamelCase_							=pipe(**lowerCAmelCase  )
              lowerCamelCase_							=output.images

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )
              lowerCamelCase_							=[
                  tuple_pipeline_inputs['''image'''],
                  tuple_pipeline_inputs['''image'''],
              ]

              lowerCamelCase_							=pipe(
                  **lowerCAmelCase,		return_dict=lowerCAmelCase,		)[0]

              lowerCamelCase_							=image[0, -3:, -3:, -1]
              lowerCamelCase_							=image_from_tuple[0, -3:, -3:, -1]

              assert image.shape == (2, 64, 64, 3)

              lowerCamelCase_							=np.array(
                  [
                      0.9_9_9_7,
                      0.9_9_8_9,
                      0.0_0_0_8,
                      0.0_0_2_1,
                      0.9_9_6_0,
                      0.0_0_1_8,
                      0.0_0_1_4,
                      0.0_0_0_2,
                      0.9_9_3_3,
                  ]  )

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






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=torch.device('''cpu'''  )


              class   __UpperCamelCase       :
                     lowercase	:    Union[str, Any]      						=1





              lowerCamelCase_							=self.get_dummy_components()

              lowerCamelCase_							=self.pipeline_class(**lowerCAmelCase  )
              lowerCamelCase_							=pipe.to(lowerCAmelCase  )

              pipe.set_progress_bar_config(disable=lowerCAmelCase  )

              lowerCamelCase_							=torch.Generator(device=lowerCAmelCase  ).manual_seed(0  )
              lowerCamelCase_							=pipe.decoder.dtype
              lowerCamelCase_							=1

              lowerCamelCase_							=(
                  batch_size,
                  pipe.decoder.config.in_channels,
                  pipe.decoder.config.sample_size,
                  pipe.decoder.config.sample_size,
              )
              lowerCamelCase_							=pipe.prepare_latents(
                  lowerCAmelCase,		dtype=lowerCAmelCase,		device=lowerCAmelCase,		generator=lowerCAmelCase,		latents=lowerCAmelCase,		scheduler=DummyScheduler()  )

              lowerCamelCase_							=(
                  batch_size,
                  pipe.super_res_first.config.in_channels // 2,
                  pipe.super_res_first.config.sample_size,
                  pipe.super_res_first.config.sample_size,
              )
              lowerCamelCase_							=pipe.prepare_latents(
                  lowerCAmelCase,		dtype=lowerCAmelCase,		device=lowerCAmelCase,		generator=lowerCAmelCase,		latents=lowerCAmelCase,		scheduler=DummyScheduler()  )

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )

              lowerCamelCase_							=pipe(
                  **lowerCAmelCase,		decoder_latents=lowerCAmelCase,		super_res_latents=lowerCAmelCase  ).images

              lowerCamelCase_							=self.get_dummy_inputs(lowerCAmelCase,		pil_image=lowerCAmelCase  )
              # Don't pass image, instead pass embedding
              lowerCamelCase_							=pipeline_inputs.pop('''image'''  )
              lowerCamelCase_							=pipe.image_encoder(lowerCAmelCase  ).image_embeds

              lowerCamelCase_							=pipe(
                  **lowerCAmelCase,		decoder_latents=lowerCAmelCase,		super_res_latents=lowerCAmelCase,		image_embeddings=lowerCAmelCase,		).images

              # make sure passing text embeddings manually is identical
              assert np.abs(img_out_a - img_out_a  ).max() < 1e-4






       @skip_mps
       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=torch_device == '''cpu'''

              # Check is relaxed because there is not a torch 2.0 sliced attention added kv processor
              lowerCamelCase_							=1e-2

              self._test_attention_slicing_forward_pass(
                  test_max_difference=lowerCAmelCase,		expected_max_diff=lowerCAmelCase  )






       @skip_mps
       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=torch_device == '''cpu'''
              lowerCamelCase_							=True
              lowerCamelCase_							=[
                  '''decoder_num_inference_steps''',
                  '''super_res_num_inference_steps''',
              ]

              self._test_inference_batch_single_identical(
                  test_max_difference=lowerCAmelCase,		relax_max_difference=lowerCAmelCase,		additional_params_copy_to_batched_inputs=lowerCAmelCase,		)






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=[
                  '''decoder_num_inference_steps''',
                  '''super_res_num_inference_steps''',
              ]

              if torch_device == "mps":
                     # TODO: MPS errors with larger batch sizes
                     lowerCamelCase_							=[2, 3]
                     self._test_inference_batch_consistent(
                         batch_sizes=lowerCAmelCase,		additional_params_copy_to_batched_inputs=lowerCAmelCase,		)
              else:
                     self._test_inference_batch_consistent(
                         additional_params_copy_to_batched_inputs=lowerCAmelCase  )






       @skip_mps
       def      lowercase__     (							self  ):



              """simple docstring"""







              return super().test_dict_tuple_outputs_equivalent()






       @skip_mps
       def      lowercase__     (							self  ):



              """simple docstring"""







              return super().test_save_load_local()






       @skip_mps
       def      lowercase__     (							self  ):



              """simple docstring"""







              return super().test_save_load_optional_components()





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






       def      lowercase__     (							self  ):



              """simple docstring"""







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






       def      lowercase__     (							self  ):



              """simple docstring"""







              lowerCamelCase_							=load_image(
                  '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/unclip/cat.png'''  )
              lowerCamelCase_							=load_numpy(
                  '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
                  '''/unclip/karlo_v1_alpha_cat_variation_fp16.npy'''  )

              lowerCamelCase_							=UnCLIPImageVariationPipeline.from_pretrained(
                  '''kakaobrain/karlo-v1-alpha-image-variations''',		torch_dtype=torch.floataa  )
              lowerCamelCase_							=pipeline.to(lowerCAmelCase  )
              pipeline.set_progress_bar_config(disable=lowerCAmelCase  )

              lowerCamelCase_							=torch.Generator(device='''cpu'''  ).manual_seed(0  )
              lowerCamelCase_							=pipeline(
                  lowerCAmelCase,		generator=lowerCAmelCase,		output_type='''np''',		)

              lowerCamelCase_							=output.images[0]

              assert image.shape == (256, 256, 3)

              assert_mean_pixel_difference(lowerCAmelCase,		lowerCAmelCase,		15  )








'''simple docstring'''


#
# This a `torch.distributed` diagnostics script that checks that all GPUs in the cluster (one or
# many nodes) can talk to each other via nccl and allocate gpu memory.
#
# To run first adjust the number of processes and nodes:
#
# python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# You may need to add --master_addr $MASTER_ADDR --master_port $MASTER_PORT if using a custom addr:port
#
# You can also use the rdzv API: --rdzv_endpoint $MASTER_ADDR:$MASTER_PORT --rdzv_backend c10d
#
# use torch.distributed.launch instead of torch.distributed.run for torch < 1.9
#
# If you get a hanging in `barrier` calls you have some network issues, you may try to debug this with:
#
# NCCL_DEBUG=INFO python -m torch.distributed.run --nproc_per_node 2 --nnodes 1 torch-distributed-gpu-test.py
#
# which should tell you what's going on behind the scenes.
#
#
# This script can be run via `srun` in the SLURM environment as well. Here is a SLURM script that
# runs on 2 nodes of 4 gpus per node:
#
# #SBATCH --job-name=test-nodes        # name
# #SBATCH --nodes=2                    # nodes
# #SBATCH --ntasks-per-node=1          # crucial - only 1 task per dist per node!
# #SBATCH --cpus-per-task=10           # number of cores per tasks
# #SBATCH --gres=gpu:4                 # number of gpus
# #SBATCH --time 0:05:00               # maximum execution time (HH:MM:SS)
# #SBATCH --output=%x-%j.out           # output file name
#
# GPUS_PER_NODE=4
# MASTER_ADDR=$(scontrol show hostnames $SLURM_JOB_NODELIST | head -n 1)
# MASTER_PORT=6000
#
# srun --jobid $SLURM_JOBID bash -c 'python -m torch.distributed.run \
# --nproc_per_node $GPUS_PER_NODE --nnodes $SLURM_NNODES --node_rank $SLURM_PROCID \
# --master_addr $MASTER_ADDR --master_port $MASTER_PORT \
# torch-distributed-gpu-test.py'
#

import fcntl
import os
import socket

import torch
import torch.distributed as dist






def        __lowerCAmelCase       (*__lowerCAmelCase      ):
		with open(__lowerCAmelCase					,       "r"      ) as fh:
				fcntl.flock(__lowerCAmelCase					,       fcntl.LOCK_EX      )
				try:
						print(*__lowerCAmelCase      )
				finally:
						fcntl.flock(__lowerCAmelCase					,       fcntl.LOCK_UN      )


lowerCamelCase__					=   int(os.environ['LOCAL_RANK'])
torch.cuda.set_device(local_rank)
lowerCamelCase__					=   torch.device('cuda', local_rank)
lowerCamelCase__					=   socket.gethostname()

lowerCamelCase__					=   F'''[{hostname}-{local_rank}]'''

try:
						# test distributed
						dist.init_process_group('nccl')
						dist.all_reduce(torch.ones(1).to(device), op=dist.ReduceOp.SUM)
						dist.barrier()

						# test cuda is available and can allocate memory
						torch.cuda.is_available()
						torch.ones(1).cuda(local_rank)

						# global rank
						lowerCamelCase__					=   dist.get_rank()
						lowerCamelCase__					=   dist.get_world_size()

						printflock(F'''{gpu} is OK (global rank: {rank}/{world_size})''')

						dist.barrier()
						if rank == 0:
												printflock(F'''pt={torch.__version__}, cuda={torch.version.cuda}, nccl={torch.cuda.nccl.version()}''')

except Exception:
						printflock(F'''{gpu} is broken''')
						raise



