



'''simple docstring'''






import json
import os
import shutil
import tempfile
import unittest

import numpy as np
import pytest

from transformers import CLIPTokenizer, CLIPTokenizerFast
from transformers.models.clip.tokenization_clip import VOCAB_FILES_NAMES
from transformers.testing_utils import require_vision
from transformers.utils import IMAGE_PROCESSOR_NAME, is_vision_available


if is_vision_available():
     from PIL import Image

     from transformers import CLIPImageProcessor, CLIPProcessor

@require_vision
class 							a__( unittest.TestCase     ):







    def       lowercase_	(				self			:  List[Any]     ):
      a						: List[Any]					       =					tempfile.mkdtemp()

      # fmt: off
      a						: int					       =					['l', 'o', 'w', 'e', 'r', 's', 't', 'i', 'd', 'n', 'lo', 'l</w>', 'w</w>', 'r</w>', 't</w>', 'low</w>', 'er</w>', 'lowest</w>', 'newer</w>', 'wider', '<unk>', '<|startoftext|>', '<|endoftext|>']
      # fmt: on
      a						: int					       =					dict(zip(__snake_case						,      range(len(__snake_case     )     )     )     )
      a						: List[str]					       =					['#version: 0.2', 'l o', 'lo w</w>', 'e r</w>', '']
      a						: Optional[Any]					       =					{'unk_token': '<unk>'}

      a						: Optional[Any]					       =					os.path.join(self.tmpdirname						,      VOCAB_FILES_NAMES['vocab_file']     )
      a						: Tuple					       =					os.path.join(self.tmpdirname						,      VOCAB_FILES_NAMES['merges_file']     )
      with open(self.vocab_file						,      'w'						,      encoding='utf-8'     ) as fp:
        fp.write(json.dumps(__snake_case     ) + '\n'     )
      with open(self.merges_file						,      'w'						,      encoding='utf-8'     ) as fp:
        fp.write('\n'.join(__snake_case     )     )

      a						: Optional[int]					       =					{
          'do_resize': True,
          'size': 20,
          'do_center_crop': True,
          'crop_size': 18,
          'do_normalize': True,
          'image_mean': [0.48145466, 0.4578275, 0.40821073],
          'image_std': [0.26862954, 0.26130258, 0.27577711],
      }
      a						: Union[str, Any]					       =					os.path.join(self.tmpdirname						,      __snake_case     )
      with open(self.image_processor_file						,      'w'						,      encoding='utf-8'     ) as fp:
        json.dump(__snake_case						,      __snake_case     )







    def       lowercase_	(				self			:  Union[str, Any]						,      **__snake_case			:  Optional[Any]     ):
      return CLIPTokenizer.from_pretrained(self.tmpdirname						,      **__snake_case     )







    def       lowercase_	(				self			:  Any						,      **__snake_case			:  List[str]     ):
      return CLIPTokenizerFast.from_pretrained(self.tmpdirname						,      **__snake_case     )







    def       lowercase_	(				self			:  int						,      **__snake_case			:  Any     ):
      return CLIPImageProcessor.from_pretrained(self.tmpdirname						,      **__snake_case     )







    def       lowercase_	(				self			:  List[Any]     ):
      shutil.rmtree(self.tmpdirname     )







    def       lowercase_	(				self			:  int     ):
      a						: List[Any]					       =					[np.random.randint(2_55						,      size=(3, 30, 4_00)						,      dtype=np.uinta     )]

      a						: Optional[Any]					       =					[Image.fromarray(np.moveaxis(__snake_case						,      0						,      -1     )     ) for x in image_inputs]

      return image_inputs







    def       lowercase_	(				self			:  Union[str, Any]     ):
      a						: Any					       =					self.get_tokenizer()
      a						: List[str]					       =					self.get_rust_tokenizer()
      a						: str					       =					self.get_image_processor()

      a						: Tuple					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )
      processor_slow.save_pretrained(self.tmpdirname     )
      a						: str					       =					CLIPProcessor.from_pretrained(self.tmpdirname						,      use_fast=__snake_case     )

      a						: List[str]					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )
      processor_fast.save_pretrained(self.tmpdirname     )
      a						: List[Any]					       =					CLIPProcessor.from_pretrained(self.tmpdirname     )

      self.assertEqual(processor_slow.tokenizer.get_vocab()						,      tokenizer_slow.get_vocab()     )
      self.assertEqual(processor_fast.tokenizer.get_vocab()						,      tokenizer_fast.get_vocab()     )
      self.assertEqual(tokenizer_slow.get_vocab()						,      tokenizer_fast.get_vocab()     )
      self.assertIsInstance(processor_slow.tokenizer						,      __snake_case     )
      self.assertIsInstance(processor_fast.tokenizer						,      __snake_case     )

      self.assertEqual(processor_slow.image_processor.to_json_string()						,      image_processor.to_json_string()     )
      self.assertEqual(processor_fast.image_processor.to_json_string()						,      image_processor.to_json_string()     )
      self.assertIsInstance(processor_slow.image_processor						,      __snake_case     )
      self.assertIsInstance(processor_fast.image_processor						,      __snake_case     )







    def       lowercase_	(				self			:  Union[str, Any]     ):
      a						: Optional[Any]					       =					CLIPProcessor(tokenizer=self.get_tokenizer()						,      image_processor=self.get_image_processor()     )
      processor.save_pretrained(self.tmpdirname     )

      a						: List[str]					       =					self.get_tokenizer(bos_token='(BOS)'						,      eos_token='(EOS)'     )
      a						: Dict					       =					self.get_image_processor(do_normalize=__snake_case						,      padding_value=1.0     )

      a						: List[Any]					       =					CLIPProcessor.from_pretrained(
          self.tmpdirname						,      bos_token='(BOS)'						,      eos_token='(EOS)'						,      do_normalize=__snake_case						,      padding_value=1.0     )

      self.assertEqual(processor.tokenizer.get_vocab()						,      tokenizer_add_kwargs.get_vocab()     )
      self.assertIsInstance(processor.tokenizer						,      __snake_case     )

      self.assertEqual(processor.image_processor.to_json_string()						,      image_processor_add_kwargs.to_json_string()     )
      self.assertIsInstance(processor.image_processor						,      __snake_case     )







    def       lowercase_	(				self			:  Tuple     ):
      a						: int					       =					self.get_image_processor()
      a						: int					       =					self.get_tokenizer()

      a						: Dict					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )

      a						: Any					       =					self.prepare_image_inputs()

      a						: str					       =					image_processor(__snake_case						,      return_tensors='np'     )
      a						: Tuple					       =					processor(images=__snake_case						,      return_tensors='np'     )

      for key in input_image_proc.keys():
        self.assertAlmostEqual(input_image_proc[key].sum()						,      input_processor[key].sum()						,      delta=1e-2     )







    def       lowercase_	(				self			:  Tuple     ):
      a						: Optional[int]					       =					self.get_image_processor()
      a						: str					       =					self.get_tokenizer()

      a						: Dict					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )

      a						: List[Any]					       =					'lower newer'

      a						: Union[str, Any]					       =					processor(text=__snake_case     )

      a						: Dict					       =					tokenizer(__snake_case     )

      for key in encoded_tok.keys():
        self.assertListEqual(encoded_tok[key]						,      encoded_processor[key]     )







    def       lowercase_	(				self			:  Optional[Any]     ):
      a						: Optional[Any]					       =					self.get_image_processor()
      a						: Any					       =					self.get_tokenizer()

      a						: List[Any]					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )

      a						: List[Any]					       =					'lower newer'
      a						: Optional[Any]					       =					self.prepare_image_inputs()

      a						: Union[str, Any]					       =					processor(text=__snake_case						,      images=__snake_case     )

      self.assertListEqual(list(inputs.keys()     )						,      ['input_ids', 'attention_mask', 'pixel_values']     )

      # test if it raises when no input is passed
      with pytest.raises(__snake_case     ):
        processor()







    def       lowercase_	(				self			:  Optional[int]     ):
      a						: Union[str, Any]					       =					self.get_image_processor()
      a						: int					       =					self.get_tokenizer()

      a						: Union[str, Any]					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )

      a						: Optional[int]					       =					[[1, 4, 5, 8, 1, 0, 8], [3, 4, 3, 1, 1, 8, 9]]

      a						: str					       =					processor.batch_decode(__snake_case     )
      a						: Optional[Any]					       =					tokenizer.batch_decode(__snake_case     )

      self.assertListEqual(__snake_case						,      __snake_case     )







    def       lowercase_	(				self			:  str     ):
      a						: List[Any]					       =					self.get_image_processor()
      a						: Tuple					       =					self.get_tokenizer()

      a						: Any					       =					CLIPProcessor(tokenizer=__snake_case						,      image_processor=__snake_case     )

      a						: Any					       =					'lower newer'
      a						: str					       =					self.prepare_image_inputs()

      a						: List[str]					       =					processor(text=__snake_case						,      images=__snake_case     )

      self.assertListEqual(list(inputs.keys()     )						,      processor.model_input_names     )




'''simple docstring'''






from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxCrossAttnUpBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
    FlaxUpBlockaD,
)

@flax.struct.dataclass
class 							a__( lowerCamelCase__     ):
    lowercase__			=   42

@flax_register_to_config
class 							a__( nn.Module		,			lowerCamelCase__		,			lowerCamelCase__     ):
    lowercase__			=   32
    lowercase__			=   4
    lowercase__			=   4
    lowercase__			=   (
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "CrossAttnDownBlock2D",
        "DownBlock2D",
    )
    lowercase__			=   ("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")
    lowercase__			=   False
    lowercase__			=   (3_20, 6_40, 12_80, 12_80)
    lowercase__			=   2
    lowercase__			=   8
    lowercase__			=   None
    lowercase__			=   12_80
    lowercase__			=   0.0
    lowercase__			=   False
    lowercase__			=   jnp.floataa
    lowercase__			=   True
    lowercase__			=   0
    lowercase__			=   False







    def       lowercase_	(				self			:  List[Any]						,      __snake_case			:  jax.random.KeyArray     ):
      # init input tensors
      a						: Dict					       =					(1, self.in_channels, self.sample_size, self.sample_size)
      a						: str					       =					jnp.zeros(__snake_case						,      dtype=jnp.floataa     )
      a						: Union[str, Any]					       =					jnp.ones((1,)						,      dtype=jnp.intaa     )
      a						: List[Any]					       =					jnp.zeros((1, 1, self.cross_attention_dim)						,      dtype=jnp.floataa     )

      a       ,     a						: Optional[int]					       =					jax.random.split(__snake_case     )
      a						: Any					       =					{'params': params_rng, 'dropout': dropout_rng}

      return self.init(__snake_case						,      __snake_case						,      __snake_case						,      __snake_case     )["params"]







    def       lowercase_	(				self			:  Union[str, Any]     ):
      a						: int					       =					self.block_out_channels
      a						: Tuple					       =					block_out_channels[0] * 4

      if self.num_attention_heads is not None:
        raise ValueError(
            'At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19.'     )

      # If `num_attention_heads` is not defined (which is the case for most models)
      # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
      # The reason for this behavior is to correct for incorrectly named variables that were introduced
      # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
      # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
      # which is why we correct for the naming here.
      a						: str					       =					self.num_attention_heads or self.attention_head_dim

      # input
      a						: List[Any]					       =					nn.Conv(
          block_out_channels[0]						,      kernel_size=(3, 3)						,      strides=(1, 1)						,      padding=((1, 1), (1, 1))						,      dtype=self.dtype						,      )

      # time
      a						: List[str]					       =					FlaxTimesteps(
          block_out_channels[0]						,      flip_sin_to_cos=self.flip_sin_to_cos						,      freq_shift=self.config.freq_shift     )
      a						: Optional[int]					       =					FlaxTimestepEmbedding(__snake_case						,      dtype=self.dtype     )

      a						: Optional[Any]					       =					self.only_cross_attention
      if isinstance(__snake_case						,      __snake_case     ):
        a						: int					       =					(only_cross_attention,) * len(self.down_block_types     )

      if isinstance(__snake_case						,      __snake_case     ):
        a						: str					       =					(num_attention_heads,) * len(self.down_block_types     )

      # down
      a						: int					       =					[]
      a						: Dict					       =					block_out_channels[0]
      for i, down_block_type in enumerate(self.down_block_types     ):
        a						: List[Any]					       =					output_channel
        a						: Dict					       =					block_out_channels[i]
        a						: Union[str, Any]					       =					i == len(__snake_case     ) - 1

        if down_block_type == "CrossAttnDownBlock2D":
          a						: Dict					       =					FlaxCrossAttnDownBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      dropout=self.dropout						,      num_layers=self.layers_per_block						,      num_attention_heads=num_attention_heads[i]						,      add_downsample=not is_final_block						,      use_linear_projection=self.use_linear_projection						,      only_cross_attention=only_cross_attention[i]						,      use_memory_efficient_attention=self.use_memory_efficient_attention						,      dtype=self.dtype						,      )
        else:
          a						: List[str]					       =					FlaxDownBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      dropout=self.dropout						,      num_layers=self.layers_per_block						,      add_downsample=not is_final_block						,      dtype=self.dtype						,      )

        down_blocks.append(__snake_case     )
      a						: Dict					       =					down_blocks

      # mid
      a						: Union[str, Any]					       =					FlaxUNetMidBlockaDCrossAttn(
          in_channels=block_out_channels[-1]						,      dropout=self.dropout						,      num_attention_heads=num_attention_heads[-1]						,      use_linear_projection=self.use_linear_projection						,      use_memory_efficient_attention=self.use_memory_efficient_attention						,      dtype=self.dtype						,      )

      # up
      a						: List[Any]					       =					[]
      a						: Dict					       =					list(reversed(__snake_case     )     )
      a						: List[str]					       =					list(reversed(__snake_case     )     )
      a						: List[str]					       =					list(reversed(__snake_case     )     )
      a						: Optional[int]					       =					reversed_block_out_channels[0]
      for i, up_block_type in enumerate(self.up_block_types     ):
        a						: Dict					       =					output_channel
        a						: Dict					       =					reversed_block_out_channels[i]
        a						: Dict					       =					reversed_block_out_channels[min(i + 1						,      len(__snake_case     ) - 1     )]

        a						: List[Any]					       =					i == len(__snake_case     ) - 1

        if up_block_type == "CrossAttnUpBlock2D":
          a						: List[Any]					       =					FlaxCrossAttnUpBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      prev_output_channel=__snake_case						,      num_layers=self.layers_per_block + 1						,      num_attention_heads=reversed_num_attention_heads[i]						,      add_upsample=not is_final_block						,      dropout=self.dropout						,      use_linear_projection=self.use_linear_projection						,      only_cross_attention=only_cross_attention[i]						,      use_memory_efficient_attention=self.use_memory_efficient_attention						,      dtype=self.dtype						,      )
        else:
          a						: Dict					       =					FlaxUpBlockaD(
              in_channels=__snake_case						,      out_channels=__snake_case						,      prev_output_channel=__snake_case						,      num_layers=self.layers_per_block + 1						,      add_upsample=not is_final_block						,      dropout=self.dropout						,      dtype=self.dtype						,      )

        up_blocks.append(__snake_case     )
        a						: Dict					       =					output_channel
      a						: int					       =					up_blocks

      # out
      a						: int					       =					nn.GroupNorm(num_groups=32						,      epsilon=1e-5     )
      a						: Any					       =					nn.Conv(
          self.out_channels						,      kernel_size=(3, 3)						,      strides=(1, 1)						,      padding=((1, 1), (1, 1))						,      dtype=self.dtype						,      )







    def __call__(				self			:  Optional[Any]						,      __snake_case			:  Dict						,      __snake_case			:  List[str]						,      __snake_case			:  Tuple						,      __snake_case			:  Dict=None						,      __snake_case			:  Any=None						,      __snake_case			:  bool = True						,      __snake_case			:  bool = False						,      ):
      # 1. time
      if not isinstance(__snake_case						,      jnp.ndarray     ):
        a						: Tuple					       =					jnp.array([timesteps]						,      dtype=jnp.intaa     )
      elif isinstance(__snake_case						,      jnp.ndarray     ) and len(timesteps.shape     ) == 0:
        a						: Tuple					       =					timesteps.astype(dtype=jnp.floataa     )
        a						: Optional[int]					       =					jnp.expand_dims(__snake_case						,      0     )

      a						: Optional[Any]					       =					self.time_proj(__snake_case     )
      a						: Optional[int]					       =					self.time_embedding(__snake_case     )

      # 2. pre-process
      a						: str					       =					jnp.transpose(__snake_case						,      (0, 2, 3, 1)     )
      a						: Optional[int]					       =					self.conv_in(__snake_case     )

      # 3. down
      a						: Optional[Any]					       =					(sample,)
      for down_block in self.down_blocks:
        if isinstance(__snake_case						,      __snake_case     ):
          a       ,     a						: Any					       =					down_block(__snake_case						,      __snake_case						,      __snake_case						,      deterministic=not train     )
        else:
          a       ,     a						: Optional[Any]					       =					down_block(__snake_case						,      __snake_case						,      deterministic=not train     )
        down_block_res_samples += res_samples

      if down_block_additional_residuals is not None:
        a						: Tuple					       =					()

        for down_block_res_sample, down_block_additional_residual in zip(
            __snake_case						,      __snake_case     ):
          down_block_res_sample += down_block_additional_residual
          new_down_block_res_samples += (down_block_res_sample,)

        a						: Tuple					       =					new_down_block_res_samples

      # 4. mid
      a						: int					       =					self.mid_block(__snake_case						,      __snake_case						,      __snake_case						,      deterministic=not train     )

      if mid_block_additional_residual is not None:
        sample += mid_block_additional_residual

      # 5. up
      for up_block in self.up_blocks:
        a						: int					       =					down_block_res_samples[-(self.layers_per_block + 1) :]
        a						: List[Any]					       =					down_block_res_samples[: -(self.layers_per_block + 1)]
        if isinstance(__snake_case						,      __snake_case     ):
          a						: Optional[int]					       =					up_block(
              __snake_case						,      temb=__snake_case						,      encoder_hidden_states=__snake_case						,      res_hidden_states_tuple=__snake_case						,      deterministic=not train						,      )
        else:
          a						: Optional[Any]					       =					up_block(__snake_case						,      temb=__snake_case						,      res_hidden_states_tuple=__snake_case						,      deterministic=not train     )

        # 6. post-process
      a						: Any					       =					self.conv_norm_out(__snake_case     )
      a						: List[str]					       =					nn.silu(__snake_case     )
      a						: Optional[int]					       =					self.conv_out(__snake_case     )
      a						: Any					       =					jnp.transpose(__snake_case						,      (0, 3, 1, 2)     )

      if not return_dict:
        return (sample,)

      return FlaxUNetaDConditionOutput(sample=__snake_case     )