




'''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 	SCREAMING_SNAKE_CASE(    __A       ,				unittest.TestCase					):
							snake_case_			:					Dict   = KandinskyPipeline
							snake_case_			:					Optional[Any]   = [
							    """prompt""",
							    """image_embeds""",
							    """negative_image_embeds""",
							]
							snake_case_			:					Dict   = ["""prompt""", """negative_prompt""", """image_embeds""", """negative_image_embeds"""]
							snake_case_			:					Optional[int]   = [
							    """generator""",
							    """height""",
							    """width""",
							    """latents""",
							    """guidance_scale""",
							    """negative_prompt""",
							    """num_inference_steps""",
							    """return_dict""",
							    """guidance_scale""",
							    """num_images_per_prompt""",
							    """output_type""",
							    """return_dict""",
							]
							snake_case_			:					str   = False

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

											"""simple docstring"""



											return 32

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

											"""simple docstring"""



											return 32

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

											"""simple docstring"""



											return self.time_input_dim

							@property
							def    snake_case__      (   self					)   -> Tuple:

											"""simple docstring"""



											return self.time_input_dim * 4

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

											"""simple docstring"""



											return 100

							@property
							def    snake_case__      (   self					)   -> int:

											"""simple docstring"""



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

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

											"""simple docstring"""



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

											__lowercase          =						MultilingualCLIP(lowerCamelCase__					)
											__lowercase          =						text_encoder.eval()

											return text_encoder

							@property
							def    snake_case__      (   self					)   -> Any:

											"""simple docstring"""



											torch.manual_seed(0					)

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

											__lowercase          =						UNetaDConditionModel(**lowerCamelCase__					)
											return model

							@property
							def    snake_case__      (   self					)   -> List[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    snake_case__      (   self					)   -> List[str]:

											"""simple docstring"""



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

							def    snake_case__      (   self					)   -> Tuple:

											"""simple docstring"""



											__lowercase          =						self.dummy_text_encoder
											__lowercase          =						self.dummy_tokenizer
											__lowercase          =						self.dummy_unet
											__lowercase          =						self.dummy_movq

											__lowercase          =						DDIMScheduler(
											    num_train_timesteps=1000 ,		beta_schedule="""linear""" ,		beta_start=0.0_00_85 ,		beta_end=0.0_12 ,		clip_sample=lowerCamelCase__ ,		set_alpha_to_one=lowerCamelCase__ ,		steps_offset=1 ,		prediction_type="""epsilon""" ,		thresholding=lowerCamelCase__ ,		)

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

							def    snake_case__      (   self ,		lowerCamelCase__ ,		lowerCamelCase__=0					)   -> str:

											"""simple docstring"""



											__lowercase          =						floats_tensor((1, self.cross_attention_dim) ,		rng=random.Random(lowerCamelCase__					)					).to(lowerCamelCase__					)
											__lowercase          =						floats_tensor((1, self.cross_attention_dim) ,		rng=random.Random(seed + 1					)					).to(lowerCamelCase__					)
											if str(lowerCamelCase__					).startswith("""mps"""					):
															__lowercase          =						torch.manual_seed(lowerCamelCase__					)
											else:
															__lowercase          =						torch.Generator(device=lowerCamelCase__					).manual_seed(lowerCamelCase__					)
											__lowercase          =						{
											    """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    snake_case__      (   self					)   -> List[str]:

											"""simple docstring"""



											__lowercase          =						"""cpu"""

											__lowercase          =						self.get_dummy_components()

											__lowercase          =						self.pipeline_class(**lowerCamelCase__					)
											__lowercase          =						pipe.to(lowerCamelCase__					)

											pipe.set_progress_bar_config(disable=lowerCamelCase__					)

											__lowercase          =						pipe(**self.get_dummy_inputs(lowerCamelCase__					)					)
											__lowercase          =						output.images

											__lowercase          =						pipe(
											    **self.get_dummy_inputs(lowerCamelCase__					) ,		return_dict=lowerCamelCase__ ,		)[0]

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

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

											__lowercase          =						np.array(
											    [0.32_86_63, 1.0, 0.23_21_68_73, 1.0, 0.92_71_75_64, 0.4_63_90_46, 0.96_89_47_77, 0.31_71_33_78, 0.6_29_39_53]					)

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

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

											"""simple docstring"""



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




							def    snake_case__      (   self					)   -> Dict:

											"""simple docstring"""



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

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

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

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

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

											__lowercase          =						torch.Generator(device="""cuda"""					).manual_seed(0					)
											__lowercase          =						pipeline(
											    lowerCamelCase__ ,		image_embeds=lowerCamelCase__ ,		negative_image_embeds=lowerCamelCase__ ,		generator=lowerCamelCase__ ,		num_inference_steps=100 ,		output_type="""np""" ,		)

											__lowercase          =						output.images[0]

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

											assert_mean_pixel_difference(lowerCamelCase__ ,		lowerCamelCase__					)










'''simple docstring'''
import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm
def     snake_case_  (       a__		:     int   ,a__		:     bool = False      ):



    """simple docstring"""




    if device_specific:
        seed += AcceleratorState().process_index
    random.seed(a__      )
    np.random.seed(a__      )
    torch.manual_seed(a__      )
    if is_xpu_available():
        torch.xpu.manual_seed_all(a__      )
    elif is_npu_available():
        torch.npu.manual_seed_all(a__      )
    else:
        torch.cuda.manual_seed_all(a__      )
    # ^^ safe to call this function even if cuda is not available
    if is_tpu_available():
        xm.set_rng_state(a__      )
def     snake_case_  (       a__		:     Optional[RNGType] = None   ,a__		:     Optional[torch.Generator] = None      ):



    """simple docstring"""




    if rng_type == RNGType.TORCH:
        __lowercase          =						torch.get_rng_state()
    elif rng_type == RNGType.CUDA:
        __lowercase          =						torch.cuda.get_rng_state()
    elif rng_type == RNGType.XLA:
        assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
        __lowercase          =						torch.tensor(xm.get_rng_state()      )
    elif rng_type == RNGType.NPU:
        assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
        __lowercase          =						torch.npu.get_rng_state()
    elif rng_type == RNGType.XPU:
        assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
        __lowercase          =						torch.xpu.get_rng_state()
    elif rng_type == RNGType.GENERATOR:
        assert generator is not None, "Need a generator to synchronize its seed."
        __lowercase          =						generator.get_state()

    # Broadcast the rng state from device 0 to other devices
    __lowercase          =						AcceleratorState()
    if state.distributed_type == DistributedType.TPU:
        __lowercase          =						rng_state.to(xm.xla_device()      )
        xm.collective_broadcast([rng_state]      )
        xm.mark_step()
        __lowercase          =						rng_state.cpu()
    elif (
        state.distributed_type in CUDA_DISTRIBUTED_TYPES
        or state.distributed_type == DistributedType.MULTI_NPU
        or state.distributed_type == DistributedType.MULTI_XPU
    ):
        __lowercase          =						rng_state.to(state.device      )
        torch.distributed.broadcast(a__   ,0      )
        __lowercase          =						rng_state.cpu()
    elif state.distributed_type == DistributedType.MULTI_CPU:
        torch.distributed.broadcast(a__   ,0      )

    # Set the broadcast rng state
    if rng_type == RNGType.TORCH:
        torch.set_rng_state(a__      )
    elif rng_type == RNGType.CUDA:
        torch.cuda.set_rng_state(a__      )
    elif rng_type == RNGType.NPU:
        torch.npu.set_rng_state(a__      )
    elif rng_type == RNGType.XPU:
        torch.xpu.set_rng_state(a__      )
    elif rng_type == RNGType.XLA:
        xm.set_rng_state(rng_state.item()      )
    elif rng_type == RNGType.GENERATOR:
        generator.set_state(a__      )



def     snake_case_  (       a__		:     List[Union[str, RNGType]]   ,a__		:     Optional[torch.Generator] = None      ):



    """simple docstring"""




    for rng_type in rng_types:
        synchronize_rng_state(RNGType(a__      )   ,generator=a__      )




