


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     snake_case__  (      _A						,       unittest.TestCase  ):
				__lowerCAmelCase		=      KandinskyPipeline
				__lowerCAmelCase		=      [
				    '''prompt''',
				    '''image_embeds''',
				    '''negative_image_embeds''',
				]
				__lowerCAmelCase		=      ['''prompt''', '''negative_prompt''', '''image_embeds''', '''negative_image_embeds''']
				__lowerCAmelCase		=      [
				    '''generator''',
				    '''height''',
				    '''width''',
				    '''latents''',
				    '''guidance_scale''',
				    '''negative_prompt''',
				    '''num_inference_steps''',
				    '''return_dict''',
				    '''guidance_scale''',
				    '''num_images_per_prompt''',
				    '''output_type''',
				    '''return_dict''',
				]
				__lowerCAmelCase		=      False
				@property
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										return 32
				@property
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										return 32
				@property
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




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




										"""simple docstring"""




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




										"""simple docstring"""




										return 100
				@property
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							Optional[Any]				=      XLMRobertaTokenizerFast.from_pretrained('YiYiXu/tiny-random-mclip-base'      )
										return tokenizer
				@property
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										torch.manual_seed(0      )
										__magic_name__				:							Union[str, Any]				=      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=1005     ,      )

										__magic_name__				:							Any				=      MultilingualCLIP(_UpperCAmelCase      )
										__magic_name__				:							Any				=      text_encoder.eval()

										return text_encoder
				@property
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										torch.manual_seed(0      )

										__magic_name__				:							Optional[int]				=      {
										    '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,
										}

										__magic_name__				:							int				=      UNetaDConditionModel(**_UpperCAmelCase      )
										return model
				@property
				def  UpperCamelCase_						(  self      ):




										"""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  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										torch.manual_seed(0      )
										__magic_name__				:							Union[str, Any]				=      VQModel(**self.dummy_movq_kwargs      )
										return model
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							List[Any]				=      self.dummy_text_encoder
										__magic_name__				:							int				=      self.dummy_tokenizer
										__magic_name__				:							Optional[int]				=      self.dummy_unet
										__magic_name__				:							int				=      self.dummy_movq

										__magic_name__				:							Optional[Any]				=      DDIMScheduler(
										    num_train_timesteps=1000     ,      beta_schedule='linear'     ,      beta_start=0.00085     ,      beta_end=0.012     ,      clip_sample=_UpperCAmelCase     ,      set_alpha_to_one=_UpperCAmelCase     ,      steps_offset=1     ,      prediction_type='epsilon'     ,      thresholding=_UpperCAmelCase     ,      )

										__magic_name__				:							Optional[Any]				=      {
										    'text_encoder': text_encoder,
										    'tokenizer': tokenizer,
										    'unet': unet,
										    'scheduler': scheduler,
										    'movq': movq,
										}
										return components
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase=0      ):




										"""simple docstring"""




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




										"""simple docstring"""




										__magic_name__				:							Tuple				=      'cpu'

										__magic_name__				:							Union[str, Any]				=      self.get_dummy_components()

										__magic_name__				:							Optional[int]				=      self.pipeline_class(**_UpperCAmelCase      )
										__magic_name__				:							List[str]				=      pipe.to(_UpperCAmelCase      )

										pipe.set_progress_bar_config(disable=_UpperCAmelCase      )

										__magic_name__				:							Dict				=      pipe(**self.get_dummy_inputs(_UpperCAmelCase      )      )
										__magic_name__				:							Any				=      output.images

										__magic_name__				:							Optional[Any]				=      pipe(
										    **self.get_dummy_inputs(_UpperCAmelCase      )     ,      return_dict=_UpperCAmelCase     ,      )[0]

										__magic_name__				:							Dict				=      image[0, -3:, -3:, -1]
										__magic_name__				:							str				=      image_from_tuple[0, -3:, -3:, -1]

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

										__magic_name__				:							List[Any]				=      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     snake_case__  (      unittest.TestCase  ):
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




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





				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							List[str]				=      load_numpy(
										    'https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main'
										    '/kandinsky/kandinsky_text2img_cat_fp16.npy'      )

										__magic_name__				:							Any				=      KandinskyPriorPipeline.from_pretrained(
										    'kandinsky-community/kandinsky-2-1-prior'     ,      torch_dtype=torch.floataa      )
										pipe_prior.to(_UpperCAmelCase      )

										__magic_name__				:							str				=      KandinskyPipeline.from_pretrained('kandinsky-community/kandinsky-2-1'     ,      torch_dtype=torch.floataa      )
										__magic_name__				:							Any				=      pipeline.to(_UpperCAmelCase      )
										pipeline.set_progress_bar_config(disable=_UpperCAmelCase      )

										__magic_name__				:							Tuple				=      'red cat, 4k photo'

										__magic_name__				:							str				=      torch.Generator(device='cuda'      ).manual_seed(0      )
										__magic_name__		, __magic_name__				:							Tuple				=      pipe_prior(
										    _UpperCAmelCase     ,      generator=_UpperCAmelCase     ,      num_inference_steps=5     ,      negative_prompt=''     ,      ).to_tuple()

										__magic_name__				:							Optional[Any]				=      torch.Generator(device='cuda'      ).manual_seed(0      )
										__magic_name__				:							Any				=      pipeline(
										    _UpperCAmelCase     ,      image_embeds=_UpperCAmelCase     ,      negative_image_embeds=_UpperCAmelCase     ,      generator=_UpperCAmelCase     ,      num_inference_steps=100     ,      output_type='np'     ,      )

										__magic_name__				:							Tuple				=      output.images[0]

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

										assert_mean_pixel_difference(_UpperCAmelCase     ,      _UpperCAmelCase      )






import collections

from .utils import ExplicitEnum, is_torch_available, logging


if is_torch_available():
			import torch


snake_case__     							= logging.get_logger(__name__)



class     snake_case__  :
				def __init__(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase=21     ,      _UpperCAmelCase=[]     ,      _UpperCAmelCase=None      ):




										"""simple docstring"""




										__magic_name__				:							Tuple				=      model
										__magic_name__				:							List[str]				=      trace_batch_nums
										__magic_name__				:							int				=      abort_after_batch_num

										# keep a LIFO buffer of frames to dump as soon as inf/nan is encountered to give context to the problem emergence
										__magic_name__				:							Union[str, Any]				=      collections.deque([]     ,      _UpperCAmelCase      )
										__magic_name__				:							Union[str, Any]				=      []
										__magic_name__				:							List[str]				=      0
										__magic_name__				:							Optional[int]				=      0
										__magic_name__				:							Any				=      False
										__magic_name__				:							str				=      '                 '

										self.analyse_model()

										self.register_forward_hook()
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase=None      ):




										"""simple docstring"""




										if frame is not None:
																self.expand_frame(_UpperCAmelCase      )
										self.frames.append('\n'.join(self.frame      )      )
										__magic_name__				:							List[Any]				=      []  # start a new frame
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										self.frame.append(_UpperCAmelCase      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										print('\n'.join(self.frames      )      )
										__magic_name__				:							List[str]				=      []
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							List[Any]				=      []
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										print(F"\nDetected inf/nan during batch_number={self.batch_number}"      )
										print(F"Last {len(self.frames      )} forward frames:"      )
										print(F"{'abs min':8} {'abs max':8} metadata"      )
										print('\n'.join(self.frames      )      )
										print('\n\n'      )
										__magic_name__				:							Any				=      []
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										__magic_name__				:							int				=      {m: name for name, m in self.model.named_modules()}
										# self.longest_module_name = max(len(v) for v in self.module_names.values())
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										if torch.is_tensor(_UpperCAmelCase      ):
																self.expand_frame(get_abs_min_max(_UpperCAmelCase     ,      _UpperCAmelCase      )      )
																if detect_overflow(_UpperCAmelCase     ,      _UpperCAmelCase      ):
																						__magic_name__				:							List[str]				=      True
										elif var is None:
																self.expand_frame(F"{'None':>17} {ctx}"      )
										else:
																self.expand_frame(F"{'not a tensor':>17} {ctx}"      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										self.expand_frame(F"\n\n{self.prefix} *** Starting batch number={self.batch_number} ***"      )
										self.expand_frame(F"{'abs min':8} {'abs max':8} metadata"      )
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										self.expand_frame(F"{self.prefix} *** Finished batch number={self.batch_number-1} ***\n\n"      )
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										self.expand_frame(F"{self.prefix} {self.module_names[module]} {module.__class__.__name__}"      )

										# params
										for name, p in module.named_parameters(recurse=_UpperCAmelCase      ):
																self.analyse_variable(_UpperCAmelCase     ,      _UpperCAmelCase      )

										# inputs
										if isinstance(_UpperCAmelCase     ,      _UpperCAmelCase      ):
																for i, x in enumerate(_UpperCAmelCase      ):
																						self.analyse_variable(_UpperCAmelCase     ,      F"input[{i}]"      )
										else:
																self.analyse_variable(_UpperCAmelCase     ,      'input'      )

										# outputs
										if isinstance(_UpperCAmelCase     ,      _UpperCAmelCase      ):
																for i, x in enumerate(_UpperCAmelCase      ):
																						# possibly a tuple of tuples
																						if isinstance(_UpperCAmelCase     ,      _UpperCAmelCase      ):
																												for j, y in enumerate(_UpperCAmelCase      ):
																																		self.analyse_variable(_UpperCAmelCase     ,      F"output[{i}][{j}]"      )
																						else:
																												self.analyse_variable(_UpperCAmelCase     ,      F"output[{i}]"      )
										else:
																self.analyse_variable(_UpperCAmelCase     ,      'output'      )

										self.save_frame()
				def  UpperCamelCase_						(  self      ):




										"""simple docstring"""




										self.model.apply(self._register_forward_hook      )
				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										module.register_forward_hook(self.forward_hook      )





				def  UpperCamelCase_						(  self     ,      _UpperCAmelCase     ,      _UpperCAmelCase     ,      _UpperCAmelCase      ):




										"""simple docstring"""




										__magic_name__				:							Optional[int]				=      False

										__magic_name__				:							Any				=      True if self.batch_number in self.trace_batch_nums else False
										if trace_mode:
																self.reset_saved_frames()

										if self.total_calls == 0:
																self.batch_start_frame()
										self.total_calls += 1

										# count batch numbers - the very first forward hook of the batch will be called when the
										# batch completes - i.e. it gets called very last - we know this batch has finished
										if module == self.model:
																self.batch_number += 1
																__magic_name__				:							str				=      True

										self.create_frame(_UpperCAmelCase     ,      _UpperCAmelCase     ,      _UpperCAmelCase      )

										# if last_frame_of_batch:
										#     self.batch_end_frame()

										if trace_mode:
																self.trace_frames()

										if last_frame_of_batch:
																self.batch_start_frame()

										if self.detected_overflow and not trace_mode:
																self.dump_saved_frames()

																# now we can abort, as it's pointless to continue running
																raise ValueError(
																    'DebugUnderflowOverflow: inf/nan detected, aborting as there is no point running further. '
																    'Please scroll up above this traceback to see the activation values prior to this event.'      )

										# abort after certain batch if requested to do so
										if self.abort_after_batch_num is not None and self.batch_number > self.abort_after_batch_num:
																raise ValueError(
																    F"DebugUnderflowOverflow: aborting after {self.batch_number} batches due to"
																    F" `abort_after_batch_num={self.abort_after_batch_num}` arg"      )






def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Union[str, Any]      ,						lowercase__  :Any				)       ->  Any:
						__magic_name__				:							List[Any]				=      var.abs()
						return f"{abs_var.min():8.2e} {abs_var.max():8.2e} {ctx}"
def  __SCREAMING_SNAKE_CASE       (				lowercase__  :Tuple      ,						lowercase__  :Optional[int]				)       ->  Optional[Any]:
						__magic_name__				:							str				=      False
						if torch.isnan(lowercase__				).any().item():
												__magic_name__				:							Optional[Any]				=      True
												print(f"{ctx} has nans"				)
						if torch.isinf(lowercase__				).any().item():
												__magic_name__				:							str				=      True
												print(f"{ctx} has infs"				)

						# if needed to monitor large elements can enable the following
						if 0:  # and detected:
												__magic_name__				:							Any				=      var[torch.ge(var.abs()      ,						1_0_0				)]
												if naaa.numel() > 0:
																		print(f"{ctx}:  n100={naaa.numel()}"				)
												__magic_name__				:							Union[str, Any]				=      var[torch.ge(var.abs()      ,						1_0_0_0				)]
												if naaaa.numel() > 0:
																		print(f"{ctx}: n1000={naaaa.numel()}"				)
												__magic_name__				:							Any				=      var[torch.ge(var.abs()      ,						1_0_0_0_0				)]
												if naaaaa.numel() > 0:
																		print(f"{ctx}: n10000={naaaaa.numel()}"				)

						if 0:
												print(f"min={var.min():9.2e} max={var.max():9.2e}"				)

						if 0:
												print(f"min={var.min():9.2e} max={var.max():9.2e} var={var.var():9.2e} mean={var.mean():9.2e} ({ctx})"				)

						return detected



class     snake_case__  (      _A  ):
				__lowerCAmelCase		=      '''underflow_overflow'''
				__lowerCAmelCase		=      '''tpu_metrics_debug'''


