code: 
  
  
  
  
  
  '''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       a_  (       snake_case_   ,					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     snake_case_(	self )	->     List[str]:
          return 32
  
  
      @property
      def     snake_case_(	self )	->     List[str]:
          return 32
  
  
      @property
      def     snake_case_(	self )	->     List[str]:
          return self.time_input_dim
  
  
      @property
      def     snake_case_(	self )	->     Union[str, Any]:
          return self.time_input_dim * 4
  
  
      @property
      def     snake_case_(	self )	->     int:
          return 100
  
  
      @property
      def     snake_case_(	self )	->     int:
          _SCREAMING_SNAKE_CASE             =    XLMRobertaTokenizerFast.from_pretrained("""YiYiXu/tiny-random-mclip-base""" )
          return tokenizer
  
  
      @property
      def     snake_case_(	self )	->     Tuple:
          torch.manual_seed(0 )
          _SCREAMING_SNAKE_CASE             =    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				,	)
  
          _SCREAMING_SNAKE_CASE             =    MultilingualCLIP(A )
          _SCREAMING_SNAKE_CASE             =    text_encoder.eval()
  
          return text_encoder
  
  
      @property
      def     snake_case_(	self )	->     Any:
          torch.manual_seed(0 )
  
          _SCREAMING_SNAKE_CASE             =    {
              """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,
          }
  
          _SCREAMING_SNAKE_CASE             =    UNetaDConditionModel(**A )
          return model
  
  
      @property
      def     snake_case_(	self )	->     Any:
          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 )	->     Dict:
          torch.manual_seed(0 )
          _SCREAMING_SNAKE_CASE             =    VQModel(**self.dummy_movq_kwargs )
          return model
  
  
      def     snake_case_(	self )	->     Union[str, Any]:
          _SCREAMING_SNAKE_CASE             =    self.dummy_text_encoder
          _SCREAMING_SNAKE_CASE             =    self.dummy_tokenizer
          _SCREAMING_SNAKE_CASE             =    self.dummy_unet
          _SCREAMING_SNAKE_CASE             =    self.dummy_movq
  
          _SCREAMING_SNAKE_CASE             =    DDIMScheduler(
              num_train_timesteps=1000				,	beta_schedule="""linear"""				,	beta_start=0.0_0085				,	beta_end=0.012				,	clip_sample=A				,	set_alpha_to_one=A				,	steps_offset=1				,	prediction_type="""epsilon"""				,	thresholding=A				,	)
  
          _SCREAMING_SNAKE_CASE             =    {
              """text_encoder""": text_encoder,
              """tokenizer""": tokenizer,
              """unet""": unet,
              """scheduler""": scheduler,
              """movq""": movq,
          }
          return components
  
  
      def     snake_case_(	self				,	A				,	A=0 )	->     List[str]:
          _SCREAMING_SNAKE_CASE             =    floats_tensor((1, self.cross_attention_dim)				,	rng=random.Random(A ) ).to(A )
          _SCREAMING_SNAKE_CASE             =    floats_tensor((1, self.cross_attention_dim)				,	rng=random.Random(seed + 1 ) ).to(A )
          if str(A ).startswith("""mps""" ):
              _SCREAMING_SNAKE_CASE             =    torch.manual_seed(A )
          else:
              _SCREAMING_SNAKE_CASE             =    torch.Generator(device=A ).manual_seed(A )
          _SCREAMING_SNAKE_CASE             =    {
              """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 )	->     Any:
          _SCREAMING_SNAKE_CASE             =    """cpu"""
  
          _SCREAMING_SNAKE_CASE             =    self.get_dummy_components()
  
          _SCREAMING_SNAKE_CASE             =    self.pipeline_class(**A )
          _SCREAMING_SNAKE_CASE             =    pipe.to(A )
  
          pipe.set_progress_bar_config(disable=A )
  
          _SCREAMING_SNAKE_CASE             =    pipe(**self.get_dummy_inputs(A ) )
          _SCREAMING_SNAKE_CASE             =    output.images
  
          _SCREAMING_SNAKE_CASE             =    pipe(
              **self.get_dummy_inputs(A )				,	return_dict=A				,	)[0]
  
          _SCREAMING_SNAKE_CASE             =    image[0, -3:, -3:, -1]
          _SCREAMING_SNAKE_CASE             =    image_from_tuple[0, -3:, -3:, -1]
  
          assert image.shape == (1, 64, 64, 3)
  
          _SCREAMING_SNAKE_CASE             =    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       a_  (       unittest.TestCase ):
  
      '''simple docstring'''
  
  
      def     snake_case_(	self )	->     List[str]:
          # clean up the VRAM after each test
          super().tearDown()
          gc.collect()
          torch.cuda.empty_cache()
  
  
  
  
  
      def     snake_case_(	self )	->     Optional[int]:
          _SCREAMING_SNAKE_CASE             =    load_numpy(
              """https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main"""
              """/kandinsky/kandinsky_text2img_cat_fp16.npy""" )
  
          _SCREAMING_SNAKE_CASE             =    KandinskyPriorPipeline.from_pretrained(
              """kandinsky-community/kandinsky-2-1-prior"""				,	torch_dtype=torch.floataa )
          pipe_prior.to(A )
  
          _SCREAMING_SNAKE_CASE             =    KandinskyPipeline.from_pretrained("""kandinsky-community/kandinsky-2-1"""				,	torch_dtype=torch.floataa )
          _SCREAMING_SNAKE_CASE             =    pipeline.to(A )
          pipeline.set_progress_bar_config(disable=A )
  
          _SCREAMING_SNAKE_CASE             =    """red cat, 4k photo"""
  
          _SCREAMING_SNAKE_CASE             =    torch.Generator(device="""cuda""" ).manual_seed(0 )
          _SCREAMING_SNAKE_CASE       , _SCREAMING_SNAKE_CASE             =    pipe_prior(
              A				,	generator=A				,	num_inference_steps=5				,	negative_prompt=""""""				,	).to_tuple()
  
          _SCREAMING_SNAKE_CASE             =    torch.Generator(device="""cuda""" ).manual_seed(0 )
          _SCREAMING_SNAKE_CASE             =    pipeline(
              A				,	image_embeds=A				,	negative_image_embeds=A				,	generator=A				,	num_inference_steps=100				,	output_type="""np"""				,	)
  
          _SCREAMING_SNAKE_CASE             =    output.images[0]
  
          assert image.shape == (512, 512, 3)
  
          assert_mean_pixel_difference(A				,	A )
  
  
  
  
style_context: 
  
  
  
  
  
  '''simple docstring'''
  
  # DISCLAIMER: This code is strongly influenced by https://github.com/pesser/pytorch_diffusion
  # and https://github.com/hojonathanho/diffusion
  
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  
  import flax
  import jax.numpy as jnp
  
  from ..configuration_utils import ConfigMixin, register_to_config
  from .scheduling_utils_flax import (
      CommonSchedulerState,
      FlaxKarrasDiffusionSchedulers,
      FlaxSchedulerMixin,
      FlaxSchedulerOutput,
      add_noise_common,
      get_velocity_common,
  )
  
  
  @flax.struct.dataclass
  class       a_  :
  
      '''simple docstring'''
      UpperCamelCase									=				42
      UpperCamelCase									=				42
  
      # setable values
      UpperCamelCase									=				42
      UpperCamelCase									=				42
      UpperCamelCase									=				None
  
  
      @classmethod
      def     snake_case_(	cls				,	A				,	A				,	A				,	A				,	)	->     List[Any]:
          return cls(
              common=A				,	final_alpha_cumprod=A				,	init_noise_sigma=A				,	timesteps=A				,	)
  
  
  
  
  @dataclass
  class       a_  (       snake_case_ ):
  
      '''simple docstring'''
      UpperCamelCase									=				42
  
  
  
  
  class       a_  (       snake_case_   ,					snake_case_ ):
  
      '''simple docstring'''
      UpperCamelCase									=				[e.name for e in FlaxKarrasDiffusionSchedulers]
  
      UpperCamelCase									=				42
  
  
      @property
      def     snake_case_(	self )	->     List[Any]:
          return True
  
  
      @register_to_config
      def __init__(	self				,	A = 1000				,	A = 0.0001				,	A = 0.02				,	A = "linear"				,	A = None				,	A = True				,	A = 0				,	A = "epsilon"				,	A = jnp.floataa				,	)	->     int:
          _SCREAMING_SNAKE_CASE             =    dtype
  
  
      def     snake_case_(	self				,	A = None )	->     DDIMSchedulerState:
          if common is None:
              _SCREAMING_SNAKE_CASE             =    CommonSchedulerState.create(self )
  
          # At every step in ddim, we are looking into the previous alphas_cumprod
          # For the final step, there is no previous alphas_cumprod because we are already at 0
          # `set_alpha_to_one` decides whether we set this parameter simply to one or
          # whether we use the final alpha of the "non-previous" one.
          _SCREAMING_SNAKE_CASE             =    (
              jnp.array(1.0				,	dtype=self.dtype ) if self.config.set_alpha_to_one else common.alphas_cumprod[0]
          )
  
          # standard deviation of the initial noise distribution
          _SCREAMING_SNAKE_CASE             =    jnp.array(1.0				,	dtype=self.dtype )
  
          _SCREAMING_SNAKE_CASE             =    jnp.arange(0				,	self.config.num_train_timesteps ).round()[::-1]
  
          return DDIMSchedulerState.create(
              common=A				,	final_alpha_cumprod=A				,	init_noise_sigma=A				,	timesteps=A				,	)
  
  
      def     snake_case_(	self				,	A				,	A				,	A = None )	->     jnp.ndarray:
          return sample
  
  
      def     snake_case_(	self				,	A				,	A				,	A = () )	->     DDIMSchedulerState:
          _SCREAMING_SNAKE_CASE             =    self.config.num_train_timesteps // num_inference_steps
          # creates integer timesteps by multiplying by ratio
          # rounding to avoid issues when num_inference_step is power of 3
          _SCREAMING_SNAKE_CASE             =    (jnp.arange(0				,	A ) * step_ratio).round()[::-1] + self.config.steps_offset
  
          return state.replace(
              num_inference_steps=A				,	timesteps=A				,	)
  
  
      def     snake_case_(	self				,	A				,	A				,	A )	->     int:
          _SCREAMING_SNAKE_CASE             =    state.common.alphas_cumprod[timestep]
          _SCREAMING_SNAKE_CASE             =    jnp.where(
              prev_timestep >= 0				,	state.common.alphas_cumprod[prev_timestep]				,	state.final_alpha_cumprod )
          _SCREAMING_SNAKE_CASE             =    1 - alpha_prod_t
          _SCREAMING_SNAKE_CASE             =    1 - alpha_prod_t_prev
  
          _SCREAMING_SNAKE_CASE             =    (beta_prod_t_prev / beta_prod_t) * (1 - alpha_prod_t / alpha_prod_t_prev)
  
          return variance
  
  
      def     snake_case_(	self				,	A				,	A				,	A				,	A				,	A = 0.0				,	A = True				,	)	->     Union[FlaxDDIMSchedulerOutput, Tuple]:
          if state.num_inference_steps is None:
              raise ValueError(
                  """Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler""" )
  
          # See formulas (12) and (16) of DDIM paper https://arxiv.org/pdf/2010.02502.pdf
          # Ideally, read DDIM paper in-detail understanding
  
          # Notation (<variable name> -> <name in paper>
          # - pred_noise_t -> e_theta(x_t, t)
          # - pred_original_sample -> f_theta(x_t, t) or x_0
          # - std_dev_t -> sigma_t
          # - eta -> η
          # - pred_sample_direction -> "direction pointing to x_t"
          # - pred_prev_sample -> "x_t-1"
  
          # 1. get previous step value (=t-1)
          _SCREAMING_SNAKE_CASE             =    timestep - self.config.num_train_timesteps // state.num_inference_steps
  
          _SCREAMING_SNAKE_CASE             =    state.common.alphas_cumprod
          _SCREAMING_SNAKE_CASE             =    state.final_alpha_cumprod
  
          # 2. compute alphas, betas
          _SCREAMING_SNAKE_CASE             =    alphas_cumprod[timestep]
          _SCREAMING_SNAKE_CASE             =    jnp.where(prev_timestep >= 0				,	alphas_cumprod[prev_timestep]				,	A )
  
          _SCREAMING_SNAKE_CASE             =    1 - alpha_prod_t
  
          # 3. compute predicted original sample from predicted noise also called
          # "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
          if self.config.prediction_type == "epsilon":
              _SCREAMING_SNAKE_CASE             =    (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
              _SCREAMING_SNAKE_CASE             =    model_output
          elif self.config.prediction_type == "sample":
              _SCREAMING_SNAKE_CASE             =    model_output
              _SCREAMING_SNAKE_CASE             =    (sample - alpha_prod_t ** 0.5 * pred_original_sample) / beta_prod_t ** 0.5
          elif self.config.prediction_type == "v_prediction":
              _SCREAMING_SNAKE_CASE             =    (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
              _SCREAMING_SNAKE_CASE             =    (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample
          else:
              raise ValueError(
                  f'prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`, or'
                  """ `v_prediction`""" )
  
          # 4. compute variance: "sigma_t(η)" -> see formula (16)
          # σ_t = sqrt((1 − α_t−1)/(1 − α_t)) * sqrt(1 − α_t/α_t−1)
          _SCREAMING_SNAKE_CASE             =    self._get_variance(A				,	A				,	A )
          _SCREAMING_SNAKE_CASE             =    eta * variance ** 0.5
  
          # 5. compute "direction pointing to x_t" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
          _SCREAMING_SNAKE_CASE             =    (1 - alpha_prod_t_prev - std_dev_t**2) ** 0.5 * pred_epsilon
  
          # 6. compute x_t without "random noise" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
          _SCREAMING_SNAKE_CASE             =    alpha_prod_t_prev ** 0.5 * pred_original_sample + pred_sample_direction
  
          if not return_dict:
              return (prev_sample, state)
  
          return FlaxDDIMSchedulerOutput(prev_sample=A				,	state=A )
  
  
      def     snake_case_(	self				,	A				,	A				,	A				,	A				,	)	->     jnp.ndarray:
          return add_noise_common(state.common				,	A				,	A				,	A )
  
  
      def     snake_case_(	self				,	A				,	A				,	A				,	A				,	)	->     jnp.ndarray:
          return get_velocity_common(state.common				,	A				,	A				,	A )
  
  
  
  
  
      def __len__(	self )	->     Any:
          return self.config.num_train_timesteps
  
  
  
  