
import gc
import random
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    DDIMScheduler,
    EulerAncestralDiscreteScheduler,
    LMSDiscreteScheduler,
    PNDMScheduler,
    StableDiffusionInstructPixaPixPipeline,
    UNetaDConditionModel,
)
from diffusers.image_processor import VaeImageProcessor
from diffusers.utils import floats_tensor, load_image, slow, torch_device
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import PipelineKarrasSchedulerTesterMixin, PipelineLatentTesterMixin, PipelineTesterMixin


enable_full_determinism()



class 		UpperCamelCase__  (       _UpperCAmelCase    ,_UpperCAmelCase    ,_UpperCAmelCase    ,unittest.TestCase    ):
   _SCREAMING_SNAKE_CASE      :							Union[str, Any]	       =      StableDiffusionInstructPixaPixPipeline
   _SCREAMING_SNAKE_CASE      :							List[Any]	       =      TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {'height', 'width', 'cross_attention_kwargs'}
   _SCREAMING_SNAKE_CASE      :							List[str]	       =      TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      IMAGE_TO_IMAGE_IMAGE_PARAMS
   _SCREAMING_SNAKE_CASE      :							List[str]	       =      IMAGE_TO_IMAGE_IMAGE_PARAMS
   def       lowerCAmelCase			(self :					Any   ):
    torch.manual_seed(0   )
    __a :  Any						=  UNetaDConditionModel(
        block_out_channels=(3_2, 6_4)					,				layers_per_block=2					,				sample_size=3_2					,				in_channels=8					,				out_channels=4					,				down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,				up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')					,				cross_attention_dim=3_2					,				)
    __a :  Union[str, Any]						=  PNDMScheduler(skip_prk_steps=SCREAMING_SNAKE_CASE_   )
    torch.manual_seed(0   )
    __a :  Optional[Any]						=  AutoencoderKL(
        block_out_channels=[3_2, 6_4]					,				in_channels=3					,				out_channels=3					,				down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,				up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,				latent_channels=4					,				)
    torch.manual_seed(0   )
    __a :  Optional[int]						=  CLIPTextConfig(
        bos_token_id=0					,				eos_token_id=2					,				hidden_size=3_2					,				intermediate_size=3_7					,				layer_norm_eps=1E-05					,				num_attention_heads=4					,				num_hidden_layers=5					,				pad_token_id=1					,				vocab_size=1_0_0_0					,				)
    __a :  int						=  CLIPTextModel(SCREAMING_SNAKE_CASE_   )
    __a :  Optional[int]						=  CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

    __a :  Optional[Any]						=  {
        'unet': unet,
        'scheduler': scheduler,
        'vae': vae,
        'text_encoder': text_encoder,
        'tokenizer': tokenizer,
        'safety_checker': None,
        'feature_extractor': None,
    }
    return components
   def       lowerCAmelCase			(self :					List[str]					,				snake_case_ :					Union[str, Any]					,				snake_case_ :					Optional[int]=0   ):
    __a :  int						=  floats_tensor((1, 3, 3_2, 3_2)					,				rng=random.Random(SCREAMING_SNAKE_CASE_   )   ).to(SCREAMING_SNAKE_CASE_   )
    __a :  List[Any]						=  image.cpu().permute(0					,				2					,				3					,				1   )[0]
    __a :  int						=  Image.fromarray(np.uinta(SCREAMING_SNAKE_CASE_   )   ).convert('''RGB'''   )
    if str(SCREAMING_SNAKE_CASE_   ).startswith('''mps'''   ):
     __a :  Union[str, Any]						=  torch.manual_seed(SCREAMING_SNAKE_CASE_   )
    else:
     __a :  Optional[int]						=  torch.Generator(device=SCREAMING_SNAKE_CASE_   ).manual_seed(SCREAMING_SNAKE_CASE_   )
    __a :  Tuple						=  {
        'prompt': 'A painting of a squirrel eating a burger',
        'image': image,
        'generator': generator,
        'num_inference_steps': 2,
        'guidance_scale': 6.0,
        'image_guidance_scale': 1,
        'output_type': 'numpy',
    }
    return inputs
   def       lowerCAmelCase			(self :					List[Any]   ):
    __a :  Optional[int]						=  'cpu'  # ensure determinism for the device-dependent torch.Generator
    __a :  Any						=  self.get_dummy_components()
    __a :  Tuple						=  StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE_   )
    __a :  Optional[int]						=  sd_pipe.to(SCREAMING_SNAKE_CASE_   )
    sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )

    __a :  str						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE_   )
    __a :  str						=  sd_pipe(**SCREAMING_SNAKE_CASE_   ).images
    __a :  List[str]						=  image[0, -3:, -3:, -1]
    assert image.shape == (1, 3_2, 3_2, 3)
    __a :  Optional[int]						=  np.array([0.7526, 0.3750, 0.4547, 0.6117, 0.5866, 0.5016, 0.4327, 0.5642, 0.4815]   )

    assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					List[str]   ):
    __a :  List[str]						=  'cpu'  # ensure determinism for the device-dependent torch.Generator
    __a :  Union[str, Any]						=  self.get_dummy_components()
    __a :  List[Any]						=  StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE_   )
    __a :  Tuple						=  sd_pipe.to(SCREAMING_SNAKE_CASE_   )
    sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )

    __a :  Union[str, Any]						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE_   )
    __a :  Union[str, Any]						=  'french fries'
    __a :  List[Any]						=  sd_pipe(**SCREAMING_SNAKE_CASE_					,				negative_prompt=SCREAMING_SNAKE_CASE_   )
    __a :  Tuple						=  output.images
    __a :  str						=  image[0, -3:, -3:, -1]

    assert image.shape == (1, 3_2, 3_2, 3)
    __a :  Optional[Any]						=  np.array([0.7511, 0.3642, 0.4553, 0.6236, 0.5797, 0.5013, 0.4343, 0.5611, 0.4831]   )

    assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					List[str]   ):
    __a :  List[str]						=  'cpu'  # ensure determinism for the device-dependent torch.Generator
    __a :  Union[str, Any]						=  self.get_dummy_components()
    __a :  Any						=  StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE_   )
    __a :  List[str]						=  sd_pipe.to(SCREAMING_SNAKE_CASE_   )
    sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )

    __a :  str						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE_   )
    __a :  str						=  [inputs['prompt']] * 2

    __a :  Tuple						=  np.array(inputs['''image''']   ).astype(np.floataa   ) / 2_5_5.0
    __a :  List[str]						=  torch.from_numpy(SCREAMING_SNAKE_CASE_   ).unsqueeze(0   ).to(SCREAMING_SNAKE_CASE_   )
    __a :  Dict						=  image / 2 + 0.5
    __a :  Optional[Any]						=  image.permute(0					,				3					,				1					,				2   )
    __a :  Dict						=  image.repeat(2					,				1					,				1					,				1   )

    __a :  int						=  sd_pipe(**SCREAMING_SNAKE_CASE_   ).images
    __a :  Tuple						=  image[-1, -3:, -3:, -1]

    assert image.shape == (2, 3_2, 3_2, 3)
    __a :  Union[str, Any]						=  np.array([0.5812, 0.5748, 0.5222, 0.5908, 0.5695, 0.7174, 0.6804, 0.5523, 0.5579]   )

    assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  str						=  'cpu'  # ensure determinism for the device-dependent torch.Generator
    __a :  int						=  self.get_dummy_components()
    __a :  Union[str, Any]						=  EulerAncestralDiscreteScheduler(
        beta_start=0.0_0085					,				beta_end=0.012					,				beta_schedule='''scaled_linear'''   )
    __a :  str						=  StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE_   )
    __a :  str						=  sd_pipe.to(SCREAMING_SNAKE_CASE_   )
    sd_pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )

    __a :  Dict						=  self.get_dummy_inputs(SCREAMING_SNAKE_CASE_   )
    __a :  str						=  sd_pipe(**SCREAMING_SNAKE_CASE_   ).images
    __a :  Optional[Any]						=  image[0, -3:, -3:, -1]

    __a :  Tuple						=  [round(SCREAMING_SNAKE_CASE_					,				4   ) for x in image_slice.flatten().tolist()]
    print(''','''.join([str(SCREAMING_SNAKE_CASE_   ) for x in slice]   )   )

    assert image.shape == (1, 3_2, 3_2, 3)
    __a :  List[str]						=  np.array([0.7417, 0.3842, 0.4732, 0.5776, 0.5891, 0.5139, 0.4052, 0.5673, 0.4986]   )

    assert np.abs(image_slice.flatten() - expected_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					str   ):
    super().test_inference_batch_single_identical(expected_max_diff=3E-3   )







   def       lowerCAmelCase			(self :					List[str]   ):
    __a :  Union[str, Any]						=  self.get_dummy_components()
    __a :  Tuple						=  StableDiffusionInstructPixaPixPipeline(**SCREAMING_SNAKE_CASE_   )
    __a :  int						=  VaeImageProcessor(do_resize=SCREAMING_SNAKE_CASE_					,				do_normalize=SCREAMING_SNAKE_CASE_   )
    __a :  str						=  pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )

    __a :  int						=  pipe(**self.get_dummy_inputs_by_type(SCREAMING_SNAKE_CASE_					,				input_image_type='''pt'''   )   )[0]

    __a :  Optional[Any]						=  components['vae']
    __a :  Dict						=  self.get_dummy_inputs_by_type(SCREAMING_SNAKE_CASE_					,				input_image_type='''pt'''   )

    for image_param in self.image_latents_params:
     if image_param in inputs.keys():
      __a :  str						=  vae.encode(inputs[image_param]   ).latent_dist.mode()

    __a :  str						=  pipe(**SCREAMING_SNAKE_CASE_   )[0]

    __a :  Optional[Any]						=  np.abs(out - out_latents_inputs   ).max()
    self.assertLess(SCREAMING_SNAKE_CASE_					,				1E-4					,				'''passing latents as image input generate different result from passing image'''   )





@slow
@require_torch_gpu
class 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Dict   ):
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()
   def       lowerCAmelCase			(self :					Tuple					,				snake_case_ :					str=0   ):
    __a :  str						=  torch.manual_seed(SCREAMING_SNAKE_CASE_   )
    __a :  Any						=  load_image(
        '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_pix2pix/example.jpg'''   )
    __a :  Any						=  {
        'prompt': 'turn him into a cyborg',
        'image': image,
        'generator': generator,
        'num_inference_steps': 3,
        'guidance_scale': 7.5,
        'image_guidance_scale': 1.0,
        'output_type': 'numpy',
    }
    return inputs
   def       lowerCAmelCase			(self :					List[Any]   ):
    __a :  int						=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
        '''timbrooks/instruct-pix2pix'''					,				safety_checker=SCREAMING_SNAKE_CASE_   )
    pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )
    pipe.enable_attention_slicing()

    __a :  Optional[Any]						=  self.get_inputs()
    __a :  List[str]						=  pipe(**SCREAMING_SNAKE_CASE_   ).images
    __a :  Optional[int]						=  image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a :  Any						=  np.array([0.5902, 0.6015, 0.6027, 0.5983, 0.6092, 0.6061, 0.5765, 0.5785, 0.5555]   )

    assert np.abs(expected_slice - image_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					str   ):
    __a :  Dict						=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
        '''timbrooks/instruct-pix2pix'''					,				safety_checker=SCREAMING_SNAKE_CASE_   )
    __a :  Any						=  LMSDiscreteScheduler.from_config(pipe.scheduler.config   )
    pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )
    pipe.enable_attention_slicing()

    __a :  Optional[Any]						=  self.get_inputs()
    __a :  Optional[int]						=  pipe(**SCREAMING_SNAKE_CASE_   ).images
    __a :  str						=  image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a :  Optional[Any]						=  np.array([0.6578, 0.6817, 0.6972, 0.6761, 0.6856, 0.6916, 0.6428, 0.6516, 0.6301]   )

    assert np.abs(expected_slice - image_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					int   ):
    __a :  Dict						=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
        '''timbrooks/instruct-pix2pix'''					,				safety_checker=SCREAMING_SNAKE_CASE_   )
    __a :  Dict						=  DDIMScheduler.from_config(pipe.scheduler.config   )
    pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )
    pipe.enable_attention_slicing()

    __a :  Union[str, Any]						=  self.get_inputs()
    __a :  str						=  pipe(**SCREAMING_SNAKE_CASE_   ).images
    __a :  Dict						=  image[0, -3:, -3:, -1].flatten()

    assert image.shape == (1, 5_1_2, 5_1_2, 3)
    __a :  Union[str, Any]						=  np.array([0.3828, 0.3834, 0.3818, 0.3792, 0.3865, 0.3752, 0.3792, 0.3847, 0.3753]   )

    assert np.abs(expected_slice - image_slice   ).max() < 1E-3
   def       lowerCAmelCase			(self :					Dict   ):
    __a :  Union[str, Any]						=  0

    def callback_fn(snake_case_ :					str					,				snake_case_ :					Tuple					,				snake_case_ :					Tuple   ) -> None:
     __a :  List[Any]						=  True
     nonlocal number_of_steps
     number_of_steps += 1
     if step == 1:
      __a :  Any						=  latents.detach().cpu().numpy()
      assert latents.shape == (1, 4, 6_4, 6_4)
      __a :  Union[str, Any]						=  latents[0, -3:, -3:, -1]
      __a :  str						=  np.array([-0.2463, -0.4644, -0.9756, 1.5176, 1.4414, 0.7866, 0.9897, 0.8521, 0.7983]   )

      assert np.abs(latents_slice.flatten() - expected_slice   ).max() < 5E-2
     elif step == 2:
      __a :  Optional[int]						=  latents.detach().cpu().numpy()
      assert latents.shape == (1, 4, 6_4, 6_4)
      __a :  List[Any]						=  latents[0, -3:, -3:, -1]
      __a :  Any						=  np.array([-0.2644, -0.4626, -0.9653, 1.5176, 1.4551, 0.7686, 0.9805, 0.8452, 0.8115]   )

      assert np.abs(latents_slice.flatten() - expected_slice   ).max() < 5E-2

    __a :  int						=  False

    __a :  Dict						=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
        '''timbrooks/instruct-pix2pix'''					,				safety_checker=SCREAMING_SNAKE_CASE_					,				torch_dtype=torch.floataa   )
    __a :  Optional[Any]						=  pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )
    pipe.enable_attention_slicing()

    __a :  Optional[int]						=  self.get_inputs()
    pipe(**SCREAMING_SNAKE_CASE_					,				callback=SCREAMING_SNAKE_CASE_					,				callback_steps=1   )
    assert callback_fn.has_been_called
    assert number_of_steps == 3
   def       lowerCAmelCase			(self :					Tuple   ):
    torch.cuda.empty_cache()
    torch.cuda.reset_max_memory_allocated()
    torch.cuda.reset_peak_memory_stats()

    __a :  Dict						=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
        '''timbrooks/instruct-pix2pix'''					,				safety_checker=SCREAMING_SNAKE_CASE_					,				torch_dtype=torch.floataa   )
    __a :  List[Any]						=  pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )
    pipe.enable_attention_slicing(1   )
    pipe.enable_sequential_cpu_offload()

    __a :  List[str]						=  self.get_inputs()
    __a :  Tuple						=  pipe(**SCREAMING_SNAKE_CASE_   )

    __a :  Optional[Any]						=  torch.cuda.max_memory_allocated()
    # make sure that less than 2.2 GB is allocated
    assert mem_bytes < 2.2 * 1_0**9







   def       lowerCAmelCase			(self :					List[str]   ):
    __a :  Optional[int]						=  self.get_inputs()
    # resize to resolution that is divisible by 8 but not 16 or 32
    __a :  Union[str, Any]						=  inputs['image'].resize((5_0_4, 5_0_4)   )

    __a :  int						=  'timbrooks/instruct-pix2pix'
    __a :  Any						=  StableDiffusionInstructPixaPixPipeline.from_pretrained(
        SCREAMING_SNAKE_CASE_					,				safety_checker=SCREAMING_SNAKE_CASE_					,				)
    pipe.to(SCREAMING_SNAKE_CASE_   )
    pipe.set_progress_bar_config(disable=SCREAMING_SNAKE_CASE_   )
    pipe.enable_attention_slicing()

    __a :  Dict						=  pipe(**SCREAMING_SNAKE_CASE_   )
    __a :  Any						=  output.images[0]

    __a :  Optional[int]						=  image[2_5_5:2_5_8, 3_8_3:3_8_6, -1]

    assert image.shape == (5_0_4, 5_0_4, 3)
    __a :  List[str]						=  np.array([0.2726, 0.2529, 0.2664, 0.2655, 0.2641, 0.2642, 0.2591, 0.2649, 0.2590]   )

    assert np.abs(image_slice.flatten() - expected_slice   ).max() < 5E-3




'''simple docstring'''



import inspect

import jax
import jax.lax as lax
import jax.numpy as jnp

from ..utils import add_start_docstrings
from ..utils.logging import get_logger


A__		:					List[str]												=  get_logger(__name__)


A__		:					str												=  R"""
    Args:
        input_ids (`jnp.ndarray` of shape `(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
            [`PreTrainedTokenizer.__call__`] for details.

            [What are input IDs?](../glossary#input-ids)
        scores (`jnp.ndarray` of shape `(batch_size, config.vocab_size)`):
            Prediction scores of a language modeling head. These can be logits for each vocabulary when not using beam
            search or log softmax for each vocabulary token when using beam search
        kwargs (`Dict[str, Any]`, *optional*):
            Additional logits processor specific kwargs.

    Return:
        `jnp.ndarray` of shape `(batch_size, config.vocab_size)`: The processed prediction scores.

"""





class    UpperCAmelCase_  :





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'	)







class    UpperCAmelCase_  :





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         raise NotImplementedError(
             f'{self.__class__} is an abstract class. Only classes inheriting this class can be called.'	)







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    @add_start_docstrings(SCREAMING_SNAKE_CASE_	)
    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         for processor in self:
              __lowerCamelCase   :	str								=		inspect.signature(processor.__call__	).parameters
              if len(SCREAMING_SNAKE_CASE_	) > 3:
                   if not all(arg in kwargs for arg in list(function_args.keys()	)[2:]	):
                        raise ValueError(
                            f'Make sure that all the required parameters: {list(function_args.keys()	)} for '
                            f'{processor.__class__} are passed to the logits processor.'	)
                   __lowerCamelCase   :	Tuple								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     **SCREAMING_SNAKE_CASE_	)
              else:
                   __lowerCamelCase   :	int								=		processor(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or not (temperature > 0):
              raise ValueError(f'`temperature` has to be a strictly positive float, but is {temperature}'	)

         __lowerCamelCase   :	Optional[int]								=		temperature







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Dict								=		scores / self.temperature
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     Union[str, Any]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or (top_p < 0 or top_p > 1.0):
              raise ValueError(f'`top_p` has to be a float > 0 and < 1, but is {top_p}'	)
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or (min_tokens_to_keep < 1):
              raise ValueError(f'`min_tokens_to_keep` has to be a positive integer, but is {min_tokens_to_keep}'	)

         __lowerCamelCase   :	str								=		top_p
         __lowerCamelCase   :	Tuple								=		filter_value
         __lowerCamelCase   :	Tuple								=		min_tokens_to_keep







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	Any								=		lax.top_k(SCREAMING_SNAKE_CASE_			,     scores.shape[-1]	)

         __lowerCamelCase   :	int								=		jnp.full_like(SCREAMING_SNAKE_CASE_			,     self.filter_value	)
         __lowerCamelCase   :	Tuple								=		jax.nn.softmax(SCREAMING_SNAKE_CASE_			,     axis=-1	).cumsum(axis=-1	)
         __lowerCamelCase   :	List[str]								=		cumulative_probs < self.top_p

         # include the token that is higher than top_p as well
         __lowerCamelCase   :	Tuple								=		jnp.roll(SCREAMING_SNAKE_CASE_			,     1	)
         score_mask |= score_mask.at[:, 0].set(SCREAMING_SNAKE_CASE_	)

         # min tokens to keep
         __lowerCamelCase   :	Any								=		score_mask.at[:, : self.min_tokens_to_keep].set(SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		jnp.where(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jax.lax.sort_key_val(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)[-1]

         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_ = -float('Inf'	)			,     SCREAMING_SNAKE_CASE_ = 1	)				->     str:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or top_k <= 0:
              raise ValueError(f'`top_k` has to be a strictly positive integer, but is {top_k}'	)

         __lowerCamelCase   :	List[str]								=		max(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		filter_value







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase				,       __lowerCamelCase   :	List[Any]								=		scores.shape
         __lowerCamelCase   :	Tuple								=		jnp.full(batch_size * vocab_size			,     self.filter_value	)

         __lowerCamelCase   :	int								=		min(self.top_k			,     scores.shape[-1]	)  # Safety check
         __lowerCamelCase				,       __lowerCamelCase   :	Tuple								=		lax.top_k(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Union[str, Any]								=		jnp.broadcast_to((jnp.arange(SCREAMING_SNAKE_CASE_	) * vocab_size)[:, None]			,     (batch_size, topk)	).flatten()
         __lowerCamelCase   :	List[Any]								=		topk_scores.flatten()
         __lowerCamelCase   :	Union[str, Any]								=		topk_indices.flatten() + shift

         __lowerCamelCase   :	Tuple								=		next_scores_flat.at[topk_indices_flat].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		next_scores_flat.reshape(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         return next_scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Any								=		bos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	Optional[Any]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Optional[Any]								=		1 - jnp.bool_(cur_len - 1	)

         __lowerCamelCase   :	List[Any]								=		jnp.where(SCREAMING_SNAKE_CASE_			,     new_scores.at[:, self.bos_token_id].set(0	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Tuple								=		max_length
         __lowerCamelCase   :	Any								=		eos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	List[str]								=		jnp.full(scores.shape			,     -float('inf'	)	)

         __lowerCamelCase   :	Any								=		1 - jnp.bool_(cur_len - self.max_length + 1	)

         __lowerCamelCase   :	List[str]								=		jnp.where(SCREAMING_SNAKE_CASE_			,     new_scores.at[:, self.eos_token_id].set(0	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     List[str]:
         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or min_length < 0:
              raise ValueError(f'`min_length` has to be a positive integer, but is {min_length}'	)

         if not isinstance(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	) or eos_token_id < 0:
              raise ValueError(f'`eos_token_id` has to be a positive integer, but is {eos_token_id}'	)

         __lowerCamelCase   :	str								=		min_length
         __lowerCamelCase   :	Optional[int]								=		eos_token_id







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         # create boolean flag to decide if min length penalty should be applied
         __lowerCamelCase   :	Optional[Any]								=		1 - jnp.clip(cur_len - self.min_length			,     0			,     1	)

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.eos_token_id].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		list(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		begin_index







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	List[Any]								=		1 - jnp.bool_(cur_len - self.begin_index	)

         __lowerCamelCase   :	str								=		jnp.where(SCREAMING_SNAKE_CASE_			,     scores.at[:, self.begin_suppress_tokens].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     List[Any]:
         __lowerCamelCase   :	Tuple								=		list(SCREAMING_SNAKE_CASE_	)







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         __lowerCamelCase   :	int								=		scores.at[..., self.suppress_tokens].set(-float('inf'	)	)

         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_	)				->     Union[str, Any]:
         __lowerCamelCase   :	Optional[int]								=		dict(SCREAMING_SNAKE_CASE_	)
         # Converts the dictionary of format {index: token} containing the tokens to be forced to an array, where the
         # index of the array corresponds to the index of the token to be forced, for XLA compatibility.
         # Indexes without forced tokens will have a negative value.
         __lowerCamelCase   :	Dict								=		jnp.ones((max(force_token_map.keys()	) + 1)			,     dtype=jnp.intaa	) * -1
         for index, token in force_token_map.items():
              if token is not None:
                   __lowerCamelCase   :	str								=		force_token_array.at[index].set(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		jnp.intaa(SCREAMING_SNAKE_CASE_	)







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     jnp.ndarray:
         def _force_token(SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	List[str]								=		scores.shape[0]
              __lowerCamelCase   :	Tuple								=		self.force_token_array[generation_idx]

              __lowerCamelCase   :	List[Any]								=		jnp.ones_like(SCREAMING_SNAKE_CASE_			,     dtype=scores.dtype	) * -float('inf'	)
              __lowerCamelCase   :	Any								=		jnp.zeros((batch_size, 1)			,     dtype=scores.dtype	)
              __lowerCamelCase   :	str								=		lax.dynamic_update_slice(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     (0, current_token)	)
              return new_scores

         __lowerCamelCase   :	int								=		lax.cond(
             cur_len >= self.force_token_array.shape[0]			,     lambda: scores			,     lambda: lax.cond(
                 self.force_token_array[cur_len] >= 0			,     lambda: _force_token(SCREAMING_SNAKE_CASE_	)			,     lambda: scores			,     )			,     )
         return scores







class    UpperCAmelCase_  (_UpperCAmelCase  ):





    """simple docstring"""






    def __init__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase   :	Any								=		generate_config.eos_token_id
         __lowerCamelCase   :	Dict								=		generate_config.no_timestamps_token_id
         __lowerCamelCase   :	Tuple								=		generate_config.no_timestamps_token_id + 1

         __lowerCamelCase   :	List[str]								=		decoder_input_length + 1

         if generate_config.is_multilingual:
              # room for language token and task token
              self.begin_index += 2
         if hasattr(SCREAMING_SNAKE_CASE_			,     'max_initial_timestamp_index'	):
              __lowerCamelCase   :	str								=		generate_config.max_initial_timestamp_index
         else:
              __lowerCamelCase   :	Optional[int]								=		model_config.vocab_size
         if self.max_initial_timestamp_index is None:
              __lowerCamelCase   :	Tuple								=		model_config.vocab_size







    def __call__(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         # suppress <|notimestamps|> which is handled by without_timestamps
         __lowerCamelCase   :	Union[str, Any]								=		scores.at[:, self.no_timestamps_token_id].set(-float('inf'	)	)

         def handle_pairs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	str								=		jnp.where((cur_len - self.begin_index) >= 1			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	List[Any]								=		jnp.where(
                  input_ids_k[cur_len - 1] >= self.timestamp_begin			,     True and last_was_timestamp			,     SCREAMING_SNAKE_CASE_			,     )

              __lowerCamelCase   :	Any								=		jnp.where((cur_len - self.begin_index) < 2			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
              __lowerCamelCase   :	str								=		jnp.where(
                  input_ids_k[cur_len - 2] >= self.timestamp_begin			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     )

              return jnp.where(
                  SCREAMING_SNAKE_CASE_			,     jnp.where(
                      penultimate_was_timestamp > 0			,     scores_k.at[self.timestamp_begin :].set(-float('inf'	)	)			,     scores_k.at[: self.eos_token_id].set(-float('inf'	)	)			,     )			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	List[Any]								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	Any								=		jnp.where(cur_len == self.begin_index			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		jnp.where(
             self.max_initial_timestamp_index is not None			,     True and apply_max_initial_timestamp			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	Any								=		self.timestamp_begin + self.max_initial_timestamp_index

         __lowerCamelCase   :	str								=		jnp.where(
             SCREAMING_SNAKE_CASE_			,     scores.at[:, last_allowed + 1 :].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_			,     )

         # if sum of probability over timestamps is above any other token, sample timestamp
         __lowerCamelCase   :	Any								=		jax.nn.log_softmax(SCREAMING_SNAKE_CASE_			,     axis=-1	)

         def handle_cumulative_probs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	Union[str, Any]								=		jax.nn.logsumexp(logprobs_k[self.timestamp_begin :]			,     axis=-1	)
              __lowerCamelCase   :	List[str]								=		jnp.max(logprobs_k[: self.timestamp_begin]	)
              return jnp.where(
                  timestamp_logprob > max_text_token_logprob			,     scores_k.at[: self.timestamp_begin].set(-float('inf'	)	)			,     SCREAMING_SNAKE_CASE_			,     )

         __lowerCamelCase   :	Dict								=		jax.vmap(SCREAMING_SNAKE_CASE_	)(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)

         return scores






