





import gc
import random
import unittest

import numpy as np
import torch

from diffusers import (
    DDIMScheduler,
    KandinskyVaaControlnetPipeline,
    KandinskyVaaPriorPipeline,
    UNetaDConditionModel,
    VQModel,
)
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 ..test_pipelines_common import PipelineTesterMixin, assert_mean_pixel_difference


enable_full_determinism()





class 			A__   (							__UpperCAmelCase			,  unittest.TestCase     ):




   """simple docstring"""



   __A       :    str	  =  KandinskyVaaControlnetPipeline
   __A       :    Union[str, Any]	  =  ['''image_embeds''', '''negative_image_embeds''', '''hint''']
   __A       :    Tuple	  =  ['''image_embeds''', '''negative_image_embeds''', '''hint''']
   __A       :    Optional[int]	  =  [
       '''generator''',
       '''height''',
       '''width''',
       '''latents''',
       '''guidance_scale''',
       '''num_inference_steps''',
       '''return_dict''',
       '''guidance_scale''',
       '''num_images_per_prompt''',
       '''output_type''',
       '''return_dict''',
   ]
   __A       :    Union[str, Any]	  =  False

   @property
   def 							__lowercase		( self)      ->						List[Any]:







          '''simple docstring'''

          return 32

   @property
   def 							__lowercase		( self)      ->						Optional[Any]:







          '''simple docstring'''

          return 32

   @property
   def 							__lowercase		( self)      ->						Union[str, Any]:







          '''simple docstring'''

          return self.time_input_dim

   @property
   def 							__lowercase		( self)      ->						Optional[int]:







          '''simple docstring'''

          return self.time_input_dim * 4

   @property
   def 							__lowercase		( self)      ->						Optional[int]:







          '''simple docstring'''

          return 100

   @property
   def 							__lowercase		( self)      ->						Optional[Any]:







          '''simple docstring'''

          torch.manual_seed(0)

          a__						:     List[str]			=    {
              'in_channels': 8,
              # Out channels is double in channels because predicts mean and variance
              'out_channels': 8,
              'addition_embed_type': 'image_hint',
              '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,
              'encoder_hid_dim': self.text_embedder_hidden_size,
              'encoder_hid_dim_type': 'image_proj',
              'cross_attention_dim': self.cross_attention_dim,
              'attention_head_dim': 4,
              'resnet_time_scale_shift': 'scale_shift',
              'class_embed_type': None,
          }

          a__						:     Optional[int]			=    UNetaDConditionModel(**lowercase)
          return model

   @property
   def 							__lowercase		( self)      ->						List[str]:







          '''simple docstring'''

          return {
              "block_out_channels": [32, 32, 64, 64],
              "down_block_types": [
                  "DownEncoderBlock2D",
                  "DownEncoderBlock2D",
                  "DownEncoderBlock2D",
                  "AttnDownEncoderBlock2D",
              ],
              "in_channels": 3,
              "latent_channels": 4,
              "layers_per_block": 1,
              "norm_num_groups": 8,
              "norm_type": "spatial",
              "num_vq_embeddings": 12,
              "out_channels": 3,
              "up_block_types": ["AttnUpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D", "UpDecoderBlock2D"],
              "vq_embed_dim": 4,
          }

   @property
   def 							__lowercase		( self)      ->						Dict:







          '''simple docstring'''

          torch.manual_seed(0)
          a__						:     Any			=    VQModel(**self.dummy_movq_kwargs)
          return model

   def 							__lowercase		( self)      ->						Union[str, Any]:







          '''simple docstring'''

          a__						:     Tuple			=    self.dummy_unet
          a__						:     Optional[int]			=    self.dummy_movq

          a__						:     Optional[Any]			=    DDIMScheduler(
              num_train_timesteps=1000 ,  beta_schedule='linear' ,  beta_start=0.0_00_85 ,  beta_end=0.0_12 ,  clip_sample=lowercase ,  set_alpha_to_one=lowercase ,  steps_offset=1 ,  prediction_type='epsilon' ,  thresholding=lowercase ,  )

          a__						:     Any			=    {
              'unet': unet,
              'scheduler': scheduler,
              'movq': movq,
          }
          return components

   def 							__lowercase		( self ,  lowercase ,  lowercase=0)      ->						Dict:







          '''simple docstring'''

          a__						:     Union[str, Any]			=    floats_tensor((1, self.text_embedder_hidden_size) ,  rng=random.Random(lowercase)).to(lowercase)
          a__						:     int			=    floats_tensor((1, self.text_embedder_hidden_size) ,  rng=random.Random(seed + 1)).to(
              lowercase)

          # create hint
          a__						:     str			=    floats_tensor((1, 3, 64, 64) ,  rng=random.Random(lowercase)).to(lowercase)

          if str(lowercase).startswith('mps'):
                 a__						:     str			=    torch.manual_seed(lowercase)
          else:
                 a__						:     int			=    torch.Generator(device=lowercase).manual_seed(lowercase)
          a__						:     Optional[int]			=    {
              'image_embeds': image_embeds,
              'negative_image_embeds': negative_image_embeds,
              'hint': hint,
              'generator': generator,
              'height': 64,
              'width': 64,
              'guidance_scale': 4.0,
              'num_inference_steps': 2,
              'output_type': 'np',
          }
          return inputs

   def 							__lowercase		( self)      ->						Optional[int]:







          '''simple docstring'''

          a__						:     int			=    'cpu'

          a__						:     int			=    self.get_dummy_components()

          a__						:     int			=    self.pipeline_class(**lowercase)
          a__						:     Optional[Any]			=    pipe.to(lowercase)

          pipe.set_progress_bar_config(disable=lowercase)

          a__						:     Any			=    pipe(**self.get_dummy_inputs(lowercase))
          a__						:     int			=    output.images

          a__						:     str			=    pipe(
              **self.get_dummy_inputs(lowercase) ,  return_dict=lowercase ,  )[0]

          a__						:     List[Any]			=    image[0, -3:, -3:, -1]
          a__						:     List[str]			=    image_from_tuple[0, -3:, -3:, -1]

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

          a__						:     List[str]			=    np.array(
              [0.6_95_98_26, 0.86_82_79, 0.7_55_80_92, 0.68_76_94_67, 0.85_80_58_04, 0.65_97_74_96, 0.44_88_53_02, 0.5_95_91_11, 0.4_25_15_95])

          assert (
              np.abs(image_slice.flatten() - expected_slice).max() < 1e-2
          ), F' expected_slice {expected_slice}, but got {image_slice.flatten()}'

          assert (
              np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < 1e-2
          ), F' expected_slice {expected_slice}, but got {image_from_tuple_slice.flatten()}'





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




   """simple docstring"""



   def 							__lowercase		( self)      ->						List[Any]:







          '''simple docstring'''

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

   def 							__lowercase		( self)      ->						List[str]:







          '''simple docstring'''

          a__						:     Optional[int]			=    load_numpy(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/kandinskyv22/kandinskyv22_controlnet_robotcat_fp16.npy')

          a__						:     List[str]			=    load_image(
              'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
              '/kandinskyv22/hint_image_cat.png')
          a__						:     Optional[Any]			=    torch.from_numpy(np.array(lowercase)).float() / 2_55.0
          a__						:     Union[str, Any]			=    hint.permute(2 ,  0 ,  1).unsqueeze(0)

          a__						:     str			=    KandinskyVaaPriorPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-2-prior' ,  torch_dtype=torch.floataa)
          pipe_prior.to(lowercase)

          a__						:     Tuple			=    KandinskyVaaControlnetPipeline.from_pretrained(
              'kandinsky-community/kandinsky-2-2-controlnet-depth' ,  torch_dtype=torch.floataa)
          a__						:     Optional[int]			=    pipeline.to(lowercase)
          pipeline.set_progress_bar_config(disable=lowercase)

          a__						:     Dict			=    'A robot, 4k photo'

          a__						:     Tuple			=    torch.Generator(device='cuda').manual_seed(0)
          a__			,							a__						:     Optional[Any]			=    pipe_prior(
              lowercase ,  generator=lowercase ,  num_inference_steps=5 ,  negative_prompt='' ,  ).to_tuple()

          a__						:     List[Any]			=    torch.Generator(device='cuda').manual_seed(0)
          a__						:     Tuple			=    pipeline(
              image_embeds=lowercase ,  negative_image_embeds=lowercase ,  hint=lowercase ,  generator=lowercase ,  num_inference_steps=100 ,  output_type='np' ,  )

          a__						:     Tuple			=    output.images[0]

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

          assert_mean_pixel_difference(lowercase ,  lowercase)









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 A_     (		A__    )  ->		Optional[int]:
       if is_torch_version('<'		,				'2.0.0'    ) or not hasattr(A__		,				'_dynamo'    ):
              return False
       return isinstance(A__		,				torch._dynamo.eval_frame.OptimizedModule    )



def A_     (		A__		,				A__ = True    )  ->		int:
       a__						:     Optional[Any]			=    (torch.nn.parallel.DistributedDataParallel, torch.nn.DataParallel)

       a__						:     Union[str, Any]			=    is_compiled_module(A__    )
       if is_compiled:
              a__						:     List[str]			=    model
              a__						:     Dict			=    model._orig_mod

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

       while isinstance(A__		,				A__    ):
              a__						:     str			=    model.module

       if not keep_fpaa_wrapper:
              a__						:     Union[str, Any]			=    getattr(A__		,				'forward'    )
              a__						:     List[Any]			=    model.__dict__.pop('_original_forward'		,				A__    )
              if original_forward is not None:
                     while hasattr(A__		,				'__wrapped__'    ):
                            a__						:     int			=    forward.__wrapped__
                            if forward == original_forward:
                                   break
                     a__						:     List[Any]			=    forward
              if getattr(A__		,				'_converted_to_transformer_engine'		,				A__    ):
                     convert_model(A__		,				to_transformer_engine=A__    )

       if is_compiled:
              a__						:     List[str]			=    model
              a__						:     Any			=    compiled_model

       return model



def A_     (		)  ->		int:
       PartialState().wait_for_everyone()



def A_     (		A__		,				A__    )  ->		Dict:
       if PartialState().distributed_type == DistributedType.TPU:
              xm.save(A__		,				A__    )
       elif PartialState().local_process_index == 0:
              torch.save(A__		,				A__    )



@contextmanager
def A_     (		**A__    )  ->		Any:
       for key, value in kwargs.items():
              a__						:     Optional[int]			=    str(A__    )

       yield

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



def A_     (		A__    )  ->		List[str]:
       if not hasattr(A__		,				'__qualname__'    ) and not hasattr(A__		,				'__name__'    ):
              a__						:     Dict			=    getattr(A__		,				'__class__'		,				A__    )
       if hasattr(A__		,				'__qualname__'    ):
              return obj.__qualname__
       if hasattr(A__		,				'__name__'    ):
              return obj.__name__
       return str(A__    )



def A_     (		A__		,				A__    )  ->		Dict:
       for key, value in source.items():
              if isinstance(A__		,				A__    ):
                     a__						:     Optional[Any]			=    destination.setdefault(A__		,				{}    )
                     merge_dicts(A__		,				A__    )
              else:
                     a__						:     Optional[int]			=    value

       return destination



def A_     (		A__ = None    )  ->		bool:
       if port is None:
              a__						:     List[Any]			=    2_9500
       with socket.socket(socket.AF_INET		,				socket.SOCK_STREAM    ) as s:
              return s.connect_ex(('localhost', port)    ) == 0


