




'''simple docstring'''




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 						lowerCAmelCase_   ( lowerCamelCase_			,  unittest.TestCase   ):







  '''simple docstring'''







  lowerCAmelCase_				:					Tuple				     =						KandinskyPipeline
  lowerCAmelCase_				:					Union[str, Any]				     =						[
      """prompt""",
      """image_embeds""",
      """negative_image_embeds""",
  ]
  lowerCAmelCase_				:					Optional[Any]				     =						["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
  lowerCAmelCase_				:					List[Any]				     =						[
      """generator""",
      """height""",
      """width""",
      """latents""",
      """guidance_scale""",
      """negative_prompt""",
      """num_inference_steps""",
      """return_dict""",
      """guidance_scale""",
      """num_images_per_prompt""",
      """output_type""",
      """return_dict""",
  ]
  lowerCAmelCase_				:					int				     =						False


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   return 32


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict ):






   """simple docstring"""






   return 32


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[Any] ):






   """simple docstring"""






   return self.time_input_dim


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return self.time_input_dim * 4


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict ):






   """simple docstring"""






   return 1_00


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   UpperCAmelCase__										=    XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base""" )
   return tokenizer


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   torch.manual_seed(0 )
   UpperCAmelCase__										=    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						,		)

   UpperCAmelCase__										=    MultilingualCLIP(_UpperCAmelCase )
   UpperCAmelCase__										=    text_encoder.eval()

   return text_encoder


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   torch.manual_seed(0 )

   UpperCAmelCase__										=    {
       """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,
   }

   UpperCAmelCase__										=    UNetaDConditionModel(**_UpperCAmelCase )
   return model


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






   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 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   torch.manual_seed(0 )
   UpperCAmelCase__										=    VQModel(**self.dummy_movq_kwargs )
   return model


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any] ):






   """simple docstring"""






   UpperCAmelCase__										=    self.dummy_text_encoder
   UpperCAmelCase__										=    self.dummy_tokenizer
   UpperCAmelCase__										=    self.dummy_unet
   UpperCAmelCase__										=    self.dummy_movq

   UpperCAmelCase__										=    DDIMScheduler(
       num_train_timesteps=10_00						,		beta_schedule="""linear"""						,		beta_start=0.0_0085						,		beta_end=0.012						,		clip_sample=_UpperCAmelCase						,		set_alpha_to_one=_UpperCAmelCase						,		steps_offset=1						,		prediction_type="""epsilon"""						,		thresholding=_UpperCAmelCase						,		)

   UpperCAmelCase__										=    {
       """text_encoder""": text_encoder,
       """tokenizer""": tokenizer,
       """unet""": unet,
       """scheduler""": scheduler,
       """movq""": movq,
   }
   return components


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any]						,		_UpperCAmelCase    :	Tuple						,		_UpperCAmelCase    :	Tuple=0 ):






   """simple docstring"""






   UpperCAmelCase__										=    floats_tensor((1, self.cross_attention_dim)						,		rng=random.Random(_UpperCAmelCase ) ).to(_UpperCAmelCase )
   UpperCAmelCase__										=    floats_tensor((1, self.cross_attention_dim)						,		rng=random.Random(seed + 1 ) ).to(_UpperCAmelCase )
   if str(_UpperCAmelCase ).startswith("""mps""" ):
    UpperCAmelCase__										=    torch.manual_seed(_UpperCAmelCase )
   else:
    UpperCAmelCase__										=    torch.Generator(device=_UpperCAmelCase ).manual_seed(_UpperCAmelCase )
   UpperCAmelCase__										=    {
       """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 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






   UpperCAmelCase__										=    """cpu"""

   UpperCAmelCase__										=    self.get_dummy_components()

   UpperCAmelCase__										=    self.pipeline_class(**_UpperCAmelCase )
   UpperCAmelCase__										=    pipe.to(_UpperCAmelCase )

   pipe.set_progress_bar_config(disable=_UpperCAmelCase )

   UpperCAmelCase__										=    pipe(**self.get_dummy_inputs(_UpperCAmelCase ) )
   UpperCAmelCase__										=    output.images

   UpperCAmelCase__										=    pipe(
       **self.get_dummy_inputs(_UpperCAmelCase )						,		return_dict=_UpperCAmelCase						,		)[0]

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

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

   UpperCAmelCase__										=    np.array(
       [0.32_8663, 1.0, 0.2321_6873, 1.0, 0.9271_7564, 0.463_9046, 0.9689_4777, 0.3171_3378, 0.629_3953] )

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







  '''simple docstring'''







  def 							SCREAMING_SNAKE_CASE__			(						self    :	int ):






   """simple docstring"""






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



  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[str] ):






   """simple docstring"""






   UpperCAmelCase__										=    load_numpy(
       """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
       """/kandinsky/kandinsky_text2img_cat_fp16.npy""" )

   UpperCAmelCase__										=    KandinskyPriorPipeline.from_pretrained(
       """kandinsky-community/kandinsky-2-1-prior"""						,		torch_dtype=torch.floataa )
   pipe_prior.to(_UpperCAmelCase )

   UpperCAmelCase__										=    KandinskyPipeline.from_pretrained("""kandinsky-community/kandinsky-2-1"""						,		torch_dtype=torch.floataa )
   UpperCAmelCase__										=    pipeline.to(_UpperCAmelCase )
   pipeline.set_progress_bar_config(disable=_UpperCAmelCase )

   UpperCAmelCase__										=    """red cat, 4k photo"""

   UpperCAmelCase__										=    torch.Generator(device="""cuda""" ).manual_seed(0 )
   UpperCAmelCase__							,			UpperCAmelCase__										=    pipe_prior(
       _UpperCAmelCase						,		generator=_UpperCAmelCase						,		num_inference_steps=5						,		negative_prompt=""""""						,		).to_tuple()

   UpperCAmelCase__										=    torch.Generator(device="""cuda""" ).manual_seed(0 )
   UpperCAmelCase__										=    pipeline(
       _UpperCAmelCase						,		image_embeds=_UpperCAmelCase						,		negative_image_embeds=_UpperCAmelCase						,		generator=_UpperCAmelCase						,		num_inference_steps=1_00						,		output_type="""np"""						,		)

   UpperCAmelCase__										=    output.images[0]

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

   assert_mean_pixel_difference(_UpperCAmelCase						,		_UpperCAmelCase )






'''simple docstring'''




import inspect
import warnings

import torch

from .state import AcceleratorState, GradientState
from .utils import DistributedType, honor_type, is_tpu_available


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

def    _UpperCamelCase      (				SCREAMING_SNAKE_CASE__			:      Tuple       ,						SCREAMING_SNAKE_CASE__			:      Dict       ):
 '''simple docstring'''





 if isinstance(SCREAMING_SNAKE_CASE__       ,						(list, tuple)       ):
  return honor_type(SCREAMING_SNAKE_CASE__       ,						(move_to_device(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ) for t in state)       )
 elif isinstance(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ):
  return type(SCREAMING_SNAKE_CASE__       )({k: move_to_device(SCREAMING_SNAKE_CASE__       ,						SCREAMING_SNAKE_CASE__       ) for k, v in state.items()}       )
 elif isinstance(SCREAMING_SNAKE_CASE__       ,						torch.Tensor       ):
  return state.to(SCREAMING_SNAKE_CASE__       )
 return state







class 						lowerCAmelCase_   ( torch.optim.Optimizer   ):







  '''simple docstring'''







  def __init__(						self    :	List[str]						,		_UpperCAmelCase    :	Any						,		_UpperCAmelCase    :	Union[str, Any]=True						,		_UpperCAmelCase    :	List[str]=None ):






   """simple docstring"""






   UpperCAmelCase__										=    optimizer
   UpperCAmelCase__										=    scaler
   UpperCAmelCase__										=    AcceleratorState()
   UpperCAmelCase__										=    GradientState()
   UpperCAmelCase__										=    device_placement
   UpperCAmelCase__										=    False
   UpperCAmelCase__										=    None

   # Handle device placement
   if device_placement:
    UpperCAmelCase__										=    self.optimizer.state_dict()
    if self.accelerator_state.distributed_type == DistributedType.TPU:
     xm.send_cpu_data_to_device(_UpperCAmelCase						,		self.accelerator_state.device )
    else:
     UpperCAmelCase__										=    move_to_device(_UpperCAmelCase						,		self.accelerator_state.device )
    self.optimizer.load_state_dict(_UpperCAmelCase )


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Dict ):






   """simple docstring"""






   return self.optimizer.state


  @state.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int]						,		_UpperCAmelCase    :	List[str] ):






   """simple docstring"""






   UpperCAmelCase__										=    state


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any ):






   """simple docstring"""






   return self.optimizer.param_groups


  @param_groups.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int]						,		_UpperCAmelCase    :	List[Any] ):






   """simple docstring"""






   UpperCAmelCase__										=    param_groups


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any] ):






   """simple docstring"""






   return self.optimizer.defaults


  @defaults.setter
  def 							SCREAMING_SNAKE_CASE__			(						self    :	int						,		_UpperCAmelCase    :	str ):






   """simple docstring"""






   UpperCAmelCase__										=    defaults


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int]						,		_UpperCAmelCase    :	Union[str, Any] ):






   """simple docstring"""






   self.optimizer.add_param_group(_UpperCAmelCase )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	List[str]						,		_UpperCAmelCase    :	Tuple ):






   """simple docstring"""






   if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
    xm.send_cpu_data_to_device(_UpperCAmelCase						,		self.accelerator_state.device )
   self.optimizer.load_state_dict(_UpperCAmelCase )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   return self.optimizer.state_dict()


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int]						,		_UpperCAmelCase    :	List[Any]=None ):






   """simple docstring"""






   if self.gradient_state.sync_gradients:
    UpperCAmelCase__										=    """set_to_none""" in inspect.signature(self.optimizer.zero_grad ).parameters
    if accept_arg:
     if set_to_none is None:
      UpperCAmelCase__										=    False
     self.optimizer.zero_grad(set_to_none=_UpperCAmelCase )
    else:
     if set_to_none is not None:
      raise ValueError("""`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.""" )
     self.optimizer.zero_grad()


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Any						,		_UpperCAmelCase    :	List[str]=None ):






   """simple docstring"""






   if self.gradient_state.sync_gradients:
    if self.accelerator_state.distributed_type == DistributedType.TPU:
     UpperCAmelCase__										=    {"""closure""": closure} if closure is not None else {}
     xm.optimizer_step(self.optimizer						,		optimizer_args=_UpperCAmelCase )
    elif self.scaler is not None:
     UpperCAmelCase__										=    False
     if self._last_scale is None:
      # `get_scale` is an async operation requiring full synchronization
      # on CPU and GPUs before finishing. As a result, we store away
      # the prior one to reduce the call overhead
      UpperCAmelCase__										=    self.scaler.get_scale()
      UpperCAmelCase__										=    True
     self.scaler.step(self.optimizer						,		_UpperCAmelCase )
     self.scaler.update()
     UpperCAmelCase__										=    self.scaler.get_scale()
     if not new_scale:
      # If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
      UpperCAmelCase__										=    scale_after < self._last_scale
     UpperCAmelCase__										=    scale_after

    else:
     self.optimizer.step(_UpperCAmelCase )


  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[Any]						,		_UpperCAmelCase    :	str ):






   """simple docstring"""






   for param_group in self.optimizer.param_groups:
    UpperCAmelCase__										=    [parameters_map.get(_UpperCAmelCase						,		_UpperCAmelCase ) for p in param_group["""params"""]]


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	str ):






   """simple docstring"""






   warnings.warn(
       """The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use """
       """`optimizer.step_was_skipped` instead."""						,		_UpperCAmelCase						,		)
   return self._is_overflow


  @property
  def 							SCREAMING_SNAKE_CASE__			(						self    :	Optional[int] ):






   """simple docstring"""






   return self._is_overflow


  def __getstate__(						self    :	Dict ):






   """simple docstring"""






   return self.__dict__.copy()



  def __setstate__(						self    :	Dict						,		_UpperCAmelCase    :	Tuple ):






   """simple docstring"""






   self.__dict__.update(_UpperCAmelCase )
