
import gc
import random
import unittest

import numpy as np
import torch
from transformers import XLMRobertaTokenizerFast

from diffusers import DDIMScheduler, KandinskyPipeline, KandinskyPriorPipeline, UNetaDConditionModel, VQModel
from diffusers.pipelines.kandinsky.text_encoder import MCLIPConfig, MultilingualCLIP
from diffusers.utils import floats_tensor, 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 		__magic_name__				( lowercase_						,unittest.TestCase					):


  """simple docstring"""





  _UpperCamelCase						 = KandinskyPipeline
  _UpperCamelCase						 = [
      "prompt",
      "image_embeds",
      "negative_image_embeds",
  ]
  _UpperCamelCase						 = ["prompt", "negative_prompt", "image_embeds", "negative_image_embeds"]
  _UpperCamelCase						 = [
      "generator",
      "height",
      "width",
      "latents",
      "guidance_scale",
      "negative_prompt",
      "num_inference_steps",
      "return_dict",
      "guidance_scale",
      "num_images_per_prompt",
      "output_type",
      "return_dict",
  ]
  _UpperCamelCase						 = False


  @property
  def     _UpperCAmelCase  (					self		):
    return 32


  @property
  def     _UpperCAmelCase  (					self		):
    return 32


  @property
  def     _UpperCAmelCase  (					self		):
    return self.time_input_dim


  @property
  def     _UpperCAmelCase  (					self		):
    return self.time_input_dim * 4


  @property
  def     _UpperCAmelCase  (					self		):
    return 1_00


  @property
  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						XLMRobertaTokenizerFast.from_pretrained('''YiYiXu/tiny-random-mclip-base'''		)
    return tokenizer


  @property
  def     _UpperCAmelCase  (					self		):
    torch.manual_seed(0		)
    _lowerCamelCase      =						MCLIPConfig(
        numDims=self.cross_attention_dim	,		transformerDimensions=self.text_embedder_hidden_size	,		hidden_size=self.text_embedder_hidden_size	,		intermediate_size=37	,		num_attention_heads=4	,		num_hidden_layers=5	,		vocab_size=10_05	,		)

    _lowerCamelCase      =						MultilingualCLIP(a__		)
    _lowerCamelCase      =						text_encoder.eval()

    return text_encoder


  @property
  def     _UpperCAmelCase  (					self		):
    torch.manual_seed(0		)

    _lowerCamelCase      =						{
        '''in_channels''': 4,
        # Out channels is double in channels because predicts mean and variance
        '''out_channels''': 8,
        '''addition_embed_type''': '''text_image''',
        '''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''': '''text_image_proj''',
        '''cross_attention_dim''': self.cross_attention_dim,
        '''attention_head_dim''': 4,
        '''resnet_time_scale_shift''': '''scale_shift''',
        '''class_embed_type''': None,
    }

    _lowerCamelCase      =						UNetaDConditionModel(**a__		)
    return model


  @property
  def     _UpperCAmelCase  (					self		):
    return {
        "block_out_channels": [32, 64],
        "down_block_types": ["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",
        ],
        "vq_embed_dim": 4,
    }


  @property
  def     _UpperCAmelCase  (					self		):
    torch.manual_seed(0		)
    _lowerCamelCase      =						VQModel(**self.dummy_movq_kwargs		)
    return model


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						self.dummy_text_encoder
    _lowerCamelCase      =						self.dummy_tokenizer
    _lowerCamelCase      =						self.dummy_unet
    _lowerCamelCase      =						self.dummy_movq

    _lowerCamelCase      =						DDIMScheduler(
        num_train_timesteps=10_00	,		beta_schedule='''linear'''	,		beta_start=0.00085	,		beta_end=0.012	,		clip_sample=a__	,		set_alpha_to_one=a__	,		steps_offset=1	,		prediction_type='''epsilon'''	,		thresholding=a__	,		)

    _lowerCamelCase      =						{
        '''text_encoder''': text_encoder,
        '''tokenizer''': tokenizer,
        '''unet''': unet,
        '''scheduler''': scheduler,
        '''movq''': movq,
    }
    return components


  def     _UpperCAmelCase  (					self	,		a__	,		a__=0		):
    _lowerCamelCase      =						floats_tensor((1, self.cross_attention_dim)	,		rng=random.Random(a__		)		).to(a__		)
    _lowerCamelCase      =						floats_tensor((1, self.cross_attention_dim)	,		rng=random.Random(seed + 1		)		).to(a__		)
    if str(a__		).startswith('''mps'''		):
      _lowerCamelCase      =						torch.manual_seed(a__		)
    else:
      _lowerCamelCase      =						torch.Generator(device=a__		).manual_seed(a__		)
    _lowerCamelCase      =						{
        '''prompt''': '''horse''',
        '''image_embeds''': image_embeds,
        '''negative_image_embeds''': negative_image_embeds,
        '''generator''': generator,
        '''height''': 64,
        '''width''': 64,
        '''guidance_scale''': 4.0,
        '''num_inference_steps''': 2,
        '''output_type''': '''np''',
    }
    return inputs


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						'''cpu'''

    _lowerCamelCase      =						self.get_dummy_components()

    _lowerCamelCase      =						self.pipeline_class(**a__		)
    _lowerCamelCase      =						pipe.to(a__		)

    pipe.set_progress_bar_config(disable=a__		)

    _lowerCamelCase      =						pipe(**self.get_dummy_inputs(a__		)		)
    _lowerCamelCase      =						output.images

    _lowerCamelCase      =						pipe(
        **self.get_dummy_inputs(a__		)	,		return_dict=a__	,		)[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.328663, 1.0, 0.23216873, 1.0, 0.92717564, 0.4639046, 0.96894777, 0.31713378, 0.6293953]		)

    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 		__magic_name__				( unittest.TestCase					):


  """simple docstring"""





  def     _UpperCAmelCase  (					self		):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()


  def     _UpperCAmelCase  (					self		):
    _lowerCamelCase      =						load_numpy(
        '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'''
        '''/kandinsky/kandinsky_text2img_cat_fp16.npy'''		)

    _lowerCamelCase      =						KandinskyPriorPipeline.from_pretrained(
        '''kandinsky-community/kandinsky-2-1-prior'''	,		torch_dtype=torch.floataa		)
    pipe_prior.to(a__		)

    _lowerCamelCase      =						KandinskyPipeline.from_pretrained('''kandinsky-community/kandinsky-2-1'''	,		torch_dtype=torch.floataa		)
    _lowerCamelCase      =						pipeline.to(a__		)
    pipeline.set_progress_bar_config(disable=a__		)

    _lowerCamelCase      =						'''red cat, 4k photo'''

    _lowerCamelCase      =						torch.Generator(device='''cuda'''		).manual_seed(0		)
    _lowerCamelCase    ,    _lowerCamelCase      =						pipe_prior(
        a__	,		generator=a__	,		num_inference_steps=5	,		negative_prompt=''''''	,		).to_tuple()

    _lowerCamelCase      =						torch.Generator(device='''cuda'''		).manual_seed(0		)
    _lowerCamelCase      =						pipeline(
        a__	,		image_embeds=a__	,		negative_image_embeds=a__	,		generator=a__	,		num_inference_steps=1_00	,		output_type='''np'''	,		)

    _lowerCamelCase      =						output.images[0]

    assert image.shape == (5_12, 5_12, 3)

    assert_mean_pixel_difference(a__	,		a__		)






import importlib
import importlib.metadata
import os
import warnings
from distutils.util import strtobool
from functools import lru_cache

import torch
from packaging import version
from packaging.version import parse

from .environment import parse_flag_from_env
from .versions import compare_versions, is_torch_version


try:
  import torch_xla.core.xla_model as xm  # noqa: F401

  _UpperCAmelCase    			=					True
except ImportError:
  _UpperCAmelCase    			=					False


# Cache this result has it's a C FFI call which can be pretty time-consuming
_UpperCAmelCase    			=					torch.distributed.is_available()


def 		_lowerCamelCase       (   _a  ):
  """simple docstring"""




  _lowerCamelCase      =						importlib.util.find_spec(_a  ) is not None
  if package_exists:
    try:
      _lowerCamelCase      =						importlib.metadata.metadata(_a  )
      return True
    except importlib.metadata.PackageNotFoundError:
      return False


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _torch_distributed_available


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  try:
    pass
  except ImportError:
    print(
        '''Intel(R) oneCCL Bindings for PyTorch* is required to run DDP on Intel(R) GPUs, but it is not'''
        ''' detected. If you see \"ValueError: Invalid backend: \'ccl\'\" error, please install Intel(R) oneCCL'''
        ''' Bindings for PyTorch*.'''  )
  return (
      importlib.util.find_spec('''torch_ccl'''  ) is not None
      or importlib.util.find_spec('''oneccl_bindings_for_pytorch'''  ) is not None
  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return importlib.metadata.version('''oneccl_bind_pt'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''transformer_engine'''  )


@lru_cache
def 		_lowerCamelCase       (   _a=True  ):
  """simple docstring"""




  if torch.cuda.is_available():
    return False
  if _tpu_available and check_device:
    try:
      # Will raise a RuntimeError if no XLA configuration is found
      _lowerCamelCase      =						xm.xla_device()
      return True
    except RuntimeError:
      return False
  return _tpu_available


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''deepspeed'''  )


def 		_lowerCamelCase       (   _a=False  ):
  """simple docstring"""




  if is_tpu_available():
    return not ignore_tpu
  if torch.cuda.is_available():
    return torch.cuda.is_bfaa_supported()
  if is_npu_available():
    return False
  return True


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  _lowerCamelCase      =						_is_package_available('''bitsandbytes'''  )
  if package_exists:
    _lowerCamelCase      =						version.parse(importlib.metadata.version('''bitsandbytes'''  )  )
    return compare_versions(_a			,   '''>='''			,   '''0.39.0'''  )
  return False


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  _lowerCamelCase      =						_is_package_available('''bitsandbytes'''  )
  if package_exists:
    _lowerCamelCase      =						version.parse(importlib.metadata.version('''bitsandbytes'''  )  )
    return compare_versions(_a			,   '''>='''			,   '''0.37.2'''  )
  return False


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''bitsandbytes'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  if strtobool(os.environ.get('''ACCELERATE_USE_MEGATRON_LM'''			,   '''False'''  )  ) == 1:
    _lowerCamelCase      =						importlib.util.find_spec('''megatron'''  ) is not None
    if package_exists:
      try:
        _lowerCamelCase      =						parse(importlib.metadata.version('''megatron-lm'''  )  )
        return compare_versions(_a			,   '''>='''			,   '''2.2.0'''  )
      except Exception as e:
        warnings.warn(F'''Parse Megatron version failed. Exception:{e}'''  )
        return False


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''safetensors'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''transformers'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''datasets'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''aim'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''tensorboard'''  ) or _is_package_available('''tensorboardX'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''wandb'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''comet_ml'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''boto3'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  if _is_package_available('''rich'''  ):
    if "ACCELERATE_DISABLE_RICH" in os.environ:
      warnings.warn(
          '''`ACCELERATE_DISABLE_RICH` is deprecated and will be removed in v0.22.0 and deactivated by default. Please use `ACCELERATE_ENABLE_RICH` if you wish to use `rich`.'''  )
      return not parse_flag_from_env('''ACCELERATE_DISABLE_RICH'''			,   _a  )
    return parse_flag_from_env('''ACCELERATE_ENABLE_RICH'''			,   _a  )
  return False


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''sagemaker'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''tqdm'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return _is_package_available('''mlflow'''  )


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  return is_torch_version('''>='''			,   '''1.12'''  ) and torch.backends.mps.is_available() and torch.backends.mps.is_built()


def 		_lowerCamelCase       (   ):
  """simple docstring"""




  def get_major_and_minor_from_version(_a  ):
    return str(version.parse(_a  ).major  ) + "." + str(version.parse(_a  ).minor  )

  _lowerCamelCase      =						importlib.metadata.version('''torch'''  )
  if importlib.util.find_spec('''intel_extension_for_pytorch'''  ) is None:
    return False
  _lowerCamelCase      =						'''N/A'''
  try:
    _lowerCamelCase      =						importlib.metadata.version('''intel_extension_for_pytorch'''  )
  except importlib.metadata.PackageNotFoundError:
    return False
  _lowerCamelCase      =						get_major_and_minor_from_version(_torch_version  )
  _lowerCamelCase      =						get_major_and_minor_from_version(_ipex_version  )
  if torch_major_and_minor != ipex_major_and_minor:
    warnings.warn(
        F'''Intel Extension for PyTorch {ipex_major_and_minor} needs to work with PyTorch {ipex_major_and_minor}.*,'''
        F''' but PyTorch {_torch_version} is found. Please switch to the matching version and run again.'''  )
    return False
  return True


@lru_cache
def 		_lowerCamelCase       (   _a=False  ):
  """simple docstring"""




  if importlib.util.find_spec('''torch'''  ) is None or importlib.util.find_spec('''torch_npu'''  ) is None:
    return False

  import torch
  import torch_npu  # noqa: F401

  if check_device:
    try:
      # Will raise a RuntimeError if no NPU is found
      _lowerCamelCase      =						torch.npu.device_count()
      return torch.npu.is_available()
    except RuntimeError:
      return False
  return hasattr(_a			,   '''npu'''  ) and torch.npu.is_available()






@lru_cache
def 		_lowerCamelCase       (   _a=False  ):
  """simple docstring"""




  if not parse_flag_from_env('''ACCELERATE_USE_XPU'''			,   default=_a  ):
    return False
  if is_ipex_available():
    import torch

    if is_torch_version('''<='''			,   '''1.12'''  ):
      return False
  else:
    return False

  import intel_extension_for_pytorch  # noqa: F401

  if check_device:
    try:
      # Will raise a RuntimeError if no XPU  is found
      _lowerCamelCase      =						torch.xpu.device_count()
      return torch.xpu.is_available()
    except RuntimeError:
      return False
  return hasattr(_a			,   '''xpu'''  ) and torch.xpu.is_available()




