






import unittest

import numpy as np
import timeout_decorator  # noqa

from transformers import BlenderbotConfig, is_flax_available
from transformers.testing_utils import jax_device, require_flax, slow

from ...generation.test_flax_utils import FlaxGenerationTesterMixin
from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor


if is_flax_available():
      import os

      # The slow tests are often failing with OOM error on GPU
      # This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
      # but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
      SCREAMING_SNAKE_CASE__		:							Dict 				=				"""platform"""

      import jax
      import jax.numpy as jnp

      from transformers import BlenderbotTokenizer
      from transformers.models.blenderbot.modeling_flax_blenderbot import (
          FlaxBlenderbotForConditionalGeneration,
          FlaxBlenderbotModel,
          shift_tokens_right,
      )





def  _lowerCamelCase						(						__lowerCamelCase ,							__lowerCamelCase ,							__lowerCamelCase=None ,							__lowerCamelCase=None ,							__lowerCamelCase=None ,							__lowerCamelCase=None ,							__lowerCamelCase=None ,							__lowerCamelCase=None ,							) ->			str:

     '''simple docstring'''
     if attention_mask is None:
          UpperCAmelCase__      :  Optional[Any]				      =			np.where(input_ids != config.pad_token_id ,							1 ,							0       )
     if decoder_attention_mask is None:
          UpperCAmelCase__      :  Tuple				      =			np.where(decoder_input_ids != config.pad_token_id ,							1 ,							0       )
     if head_mask is None:
          UpperCAmelCase__      :  List[str]				      =			np.ones((config.encoder_layers, config.encoder_attention_heads)       )
     if decoder_head_mask is None:
          UpperCAmelCase__      :  str				      =			np.ones((config.decoder_layers, config.decoder_attention_heads)       )
     if cross_attn_head_mask is None:
          UpperCAmelCase__      :  List[Any]				      =			np.ones((config.decoder_layers, config.decoder_attention_heads)       )
     return {
         "input_ids": input_ids,
         "decoder_input_ids": decoder_input_ids,
         "attention_mask": attention_mask,
         "decoder_attention_mask": attention_mask,
     }







class 				UpperCAmelCase_  :




     def __init__(				self	,						_lowerCAmelCase	,						_lowerCAmelCase=13	,						_lowerCAmelCase=7	,						_lowerCAmelCase=True	,						_lowerCAmelCase=False	,						_lowerCAmelCase=99	,						_lowerCAmelCase=16	,						_lowerCAmelCase=2	,						_lowerCAmelCase=4	,						_lowerCAmelCase=4	,						_lowerCAmelCase="gelu"	,						_lowerCAmelCase=0.1	,						_lowerCAmelCase=0.1	,						_lowerCAmelCase=32	,						_lowerCAmelCase=2	,						_lowerCAmelCase=1	,						_lowerCAmelCase=0	,						_lowerCAmelCase=0.0_2	,						):
          UpperCAmelCase__      :  str				      =			parent
          UpperCAmelCase__      :  Optional[Any]				      =			batch_size
          UpperCAmelCase__      :  str				      =			seq_length
          UpperCAmelCase__      :  List[Any]				      =			is_training
          UpperCAmelCase__      :  List[str]				      =			use_labels
          UpperCAmelCase__      :  Any				      =			vocab_size
          UpperCAmelCase__      :  Any				      =			hidden_size
          UpperCAmelCase__      :  Optional[Any]				      =			num_hidden_layers
          UpperCAmelCase__      :  List[Any]				      =			num_attention_heads
          UpperCAmelCase__      :  Tuple				      =			intermediate_size
          UpperCAmelCase__      :  int				      =			hidden_act
          UpperCAmelCase__      :  int				      =			hidden_dropout_prob
          UpperCAmelCase__      :  Tuple				      =			attention_probs_dropout_prob
          UpperCAmelCase__      :  List[Any]				      =			max_position_embeddings
          UpperCAmelCase__      :  Optional[int]				      =			eos_token_id
          UpperCAmelCase__      :  Optional[int]				      =			pad_token_id
          UpperCAmelCase__      :  Union[str, Any]				      =			bos_token_id
          UpperCAmelCase__      :  List[Any]				      =			initializer_range




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Union[str, Any]				      =			np.clip(ids_tensor([self.batch_size, self.seq_length - 1]	,						self.vocab_size		)	,						3	,						self.vocab_size		)
          UpperCAmelCase__      :  Tuple				      =			np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)	,						dtype=np.intaa		))	,						-1		)

          UpperCAmelCase__      :  Optional[int]				      =			shift_tokens_right(_lowerCAmelCase	,						1	,						2		)

          UpperCAmelCase__      :  Dict				      =			BlenderbotConfig(
              vocab_size=self.vocab_size	,						d_model=self.hidden_size	,						encoder_layers=self.num_hidden_layers	,						decoder_layers=self.num_hidden_layers	,						encoder_attention_heads=self.num_attention_heads	,						decoder_attention_heads=self.num_attention_heads	,						encoder_ffn_dim=self.intermediate_size	,						decoder_ffn_dim=self.intermediate_size	,						dropout=self.hidden_dropout_prob	,						attention_dropout=self.attention_probs_dropout_prob	,						max_position_embeddings=self.max_position_embeddings	,						eos_token_id=self.eos_token_id	,						bos_token_id=self.bos_token_id	,						pad_token_id=self.pad_token_id	,						initializer_range=self.initializer_range	,						use_cache=_lowerCAmelCase	,						)
          UpperCAmelCase__      :  Dict				      =			prepare_blenderbot_inputs_dict(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)
          return config, inputs_dict




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      :  List[str]				      =			self.prepare_config_and_inputs()
          return config, inputs_dict




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  List[str]				      =			20
          UpperCAmelCase__      :  Tuple				      =			model_class_name(_lowerCAmelCase		)

          UpperCAmelCase__      :  int				      =			model.encode(inputs_dict["""input_ids"""]		)

          UpperCAmelCase__      ,		UpperCAmelCase__      :  Union[str, Any]				      =			(
              inputs_dict["""decoder_input_ids"""],
              inputs_dict["""decoder_attention_mask"""],
          )

          UpperCAmelCase__      :  Tuple				      =			model.init_cache(decoder_input_ids.shape[0]	,						_lowerCAmelCase	,						_lowerCAmelCase		)
          UpperCAmelCase__      :  Tuple				      =			jnp.ones((decoder_input_ids.shape[0], max_decoder_length)	,						dtype="""i4"""		)

          UpperCAmelCase__      :  Optional[Any]				      =			jnp.broadcast_to(
              jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]	,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,						)
          UpperCAmelCase__      :  Tuple				      =			model.decode(
              decoder_input_ids[:, :-1]	,						_lowerCAmelCase	,						decoder_attention_mask=_lowerCAmelCase	,						past_key_values=_lowerCAmelCase	,						decoder_position_ids=_lowerCAmelCase	,						)

          UpperCAmelCase__      :  Tuple				      =			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,						dtype="""i4"""		)
          UpperCAmelCase__      :  List[Any]				      =			model.decode(
              decoder_input_ids[:, -1:]	,						_lowerCAmelCase	,						decoder_attention_mask=_lowerCAmelCase	,						past_key_values=outputs_cache.past_key_values	,						decoder_position_ids=_lowerCAmelCase	,						)

          UpperCAmelCase__      :  str				      =			model.decode(_lowerCAmelCase	,						_lowerCAmelCase		)

          UpperCAmelCase__      :  List[Any]				      =			np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])		)		)
          self.parent.assertTrue(diff < 1e-3	,						msg=f"Max diff is {diff}"		)




     def   __UpperCAmelCase  (				self	,						_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
          UpperCAmelCase__      :  Tuple				      =			20
          UpperCAmelCase__      :  Union[str, Any]				      =			model_class_name(_lowerCAmelCase		)

          UpperCAmelCase__      :  List[Any]				      =			model.encode(inputs_dict["""input_ids"""]		)

          UpperCAmelCase__      ,		UpperCAmelCase__      :  Union[str, Any]				      =			(
              inputs_dict["""decoder_input_ids"""],
              inputs_dict["""decoder_attention_mask"""],
          )

          UpperCAmelCase__      :  str				      =			jnp.concatenate(
              [
                  decoder_attention_mask,
                  jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])		),
              ]	,						axis=-1	,						)

          UpperCAmelCase__      :  Optional[Any]				      =			model.init_cache(decoder_input_ids.shape[0]	,						_lowerCAmelCase	,						_lowerCAmelCase		)
          UpperCAmelCase__      :  Optional[Any]				      =			jnp.broadcast_to(
              jnp.arange(decoder_input_ids.shape[-1] - 1		)[None, :]	,						(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)	,						)

          UpperCAmelCase__      :  str				      =			model.decode(
              decoder_input_ids[:, :-1]	,						_lowerCAmelCase	,						decoder_attention_mask=_lowerCAmelCase	,						past_key_values=_lowerCAmelCase	,						decoder_position_ids=_lowerCAmelCase	,						)
          UpperCAmelCase__      :  Any				      =			jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]	,						dtype="""i4"""		)
          UpperCAmelCase__      :  List[str]				      =			model.decode(
              decoder_input_ids[:, -1:]	,						_lowerCAmelCase	,						past_key_values=outputs_cache.past_key_values	,						decoder_attention_mask=_lowerCAmelCase	,						decoder_position_ids=_lowerCAmelCase	,						)

          UpperCAmelCase__      :  int				      =			model.decode(_lowerCAmelCase	,						_lowerCAmelCase	,						decoder_attention_mask=_lowerCAmelCase		)

          UpperCAmelCase__      :  Any				      =			np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])		)		)
          self.parent.assertTrue(diff < 1e-3	,						msg=f"Max diff is {diff}"		)







@require_flax
class 				UpperCAmelCase_  (		unittest.TestCase		):
     __lowerCamelCase      				=  99




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Any				      =			np.array(
              [
                  [71, 82, 18, 33, 46, 91, 2],
                  [68, 34, 26, 58, 30, 82, 2],
                  [5, 97, 17, 39, 94, 40, 2],
                  [76, 83, 94, 25, 70, 78, 2],
                  [87, 59, 41, 35, 48, 66, 2],
                  [55, 13, 16, 58, 5, 2, 1],  # note padding
                  [64, 27, 31, 51, 12, 75, 2],
                  [52, 64, 86, 17, 83, 39, 2],
                  [48, 61, 9, 24, 71, 82, 2],
                  [26, 1, 60, 48, 22, 13, 2],
                  [21, 5, 62, 28, 14, 76, 2],
                  [45, 98, 37, 86, 59, 48, 2],
                  [70, 70, 50, 9, 28, 0, 2],
              ]	,						dtype=np.intaa	,						)

          UpperCAmelCase__      :  int				      =			input_ids.shape[0]
          UpperCAmelCase__      :  List[str]				      =			BlenderbotConfig(
              vocab_size=self.vocab_size	,						d_model=24	,						encoder_layers=2	,						decoder_layers=2	,						encoder_attention_heads=2	,						decoder_attention_heads=2	,						encoder_ffn_dim=32	,						decoder_ffn_dim=32	,						max_position_embeddings=48	,						eos_token_id=2	,						pad_token_id=1	,						bos_token_id=0	,						)
          return config, input_ids, batch_size




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      ,		UpperCAmelCase__      :  List[str]				      =			self._get_config_and_data()
          UpperCAmelCase__      :  Optional[Any]				      =			FlaxBlenderbotForConditionalGeneration(_lowerCAmelCase		)
          UpperCAmelCase__      :  Optional[int]				      =			lm_model(input_ids=_lowerCAmelCase		)
          UpperCAmelCase__      :  Tuple				      =			(batch_size, input_ids.shape[1], config.vocab_size)
          self.assertEqual(outputs["""logits"""].shape	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  List[Any]				      =			BlenderbotConfig(
              vocab_size=self.vocab_size	,						d_model=14	,						encoder_layers=2	,						decoder_layers=2	,						encoder_attention_heads=2	,						decoder_attention_heads=2	,						encoder_ffn_dim=8	,						decoder_ffn_dim=8	,						max_position_embeddings=48	,						)
          UpperCAmelCase__      :  Tuple				      =			FlaxBlenderbotForConditionalGeneration(_lowerCAmelCase		)
          UpperCAmelCase__      :  str				      =			np.array([[71, 82, 18, 33, 46, 91, 2], [68, 34, 26, 58, 30, 2, 1]]	,						dtype=np.intaa		)
          UpperCAmelCase__      :  Union[str, Any]				      =			np.array([[82, 71, 82, 18, 2], [58, 68, 2, 1, 1]]	,						dtype=np.intaa		)
          UpperCAmelCase__      :  int				      =			lm_model(input_ids=_lowerCAmelCase	,						decoder_input_ids=_lowerCAmelCase		)
          UpperCAmelCase__      :  Union[str, Any]				      =			(*summary.shape, config.vocab_size)
          self.assertEqual(outputs["""logits"""].shape	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  List[str]				      =			np.array([[71, 82, 18, 33, 2, 1, 1], [68, 34, 26, 58, 30, 82, 2]]	,						dtype=np.intaa		)
          UpperCAmelCase__      :  Optional[int]				      =			shift_tokens_right(_lowerCAmelCase	,						1	,						2		)
          UpperCAmelCase__      :  List[str]				      =			np.equal(_lowerCAmelCase	,						1		).astype(np.floataa		).sum()
          UpperCAmelCase__      :  List[Any]				      =			np.equal(_lowerCAmelCase	,						1		).astype(np.floataa		).sum()
          self.assertEqual(shifted.shape	,						input_ids.shape		)
          self.assertEqual(_lowerCAmelCase	,						n_pad_before - 1		)
          self.assertTrue(np.equal(shifted[:, 0]	,						2		).all()		)







@require_flax
class 				UpperCAmelCase_  (		__lowerCamelCase							,      unittest.TestCase							,      __lowerCamelCase		):
     __lowerCamelCase      				=  True
     __lowerCamelCase      				=  (
         (
             FlaxBlenderbotModel,
             FlaxBlenderbotForConditionalGeneration,
         )
         if is_flax_available()
         else ()
     )
     __lowerCamelCase      				=  (FlaxBlenderbotForConditionalGeneration,) if is_flax_available() else ()




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  str				      =			FlaxBlenderbotModelTester(self		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      :  Any				      =			self.model_tester.prepare_config_and_inputs()
          for model_class in self.all_model_classes:
               self.model_tester.check_use_cache_forward(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      :  Tuple				      =			self.model_tester.prepare_config_and_inputs()
          for model_class in self.all_model_classes:
               self.model_tester.check_use_cache_forward_with_attn_mask(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      :  str				      =			self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__		):
                    UpperCAmelCase__      :  Dict				      =			self._prepare_for_class(_lowerCAmelCase	,						_lowerCAmelCase		)
                    UpperCAmelCase__      :  int				      =			model_class(_lowerCAmelCase		)

                    @jax.jit
                    def encode_jitted(_lowerCAmelCase	,						_lowerCAmelCase=None	,						**_lowerCAmelCase		):
                         return model.encode(input_ids=_lowerCAmelCase	,						attention_mask=_lowerCAmelCase		)

                    with self.subTest("""JIT Enabled"""		):
                         UpperCAmelCase__      :  int				      =			encode_jitted(**_lowerCAmelCase		).to_tuple()

                    with self.subTest("""JIT Disabled"""		):
                         with jax.disable_jit():
                              UpperCAmelCase__      :  Optional[int]				      =			encode_jitted(**_lowerCAmelCase		).to_tuple()

                    self.assertEqual(len(_lowerCAmelCase		)	,						len(_lowerCAmelCase		)		)
                    for jitted_output, output in zip(_lowerCAmelCase	,						_lowerCAmelCase		):
                         self.assertEqual(jitted_output.shape	,						output.shape		)




     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      ,		UpperCAmelCase__      :  Optional[Any]				      =			self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
               with self.subTest(model_class.__name__		):
                    UpperCAmelCase__      :  Tuple				      =			model_class(_lowerCAmelCase		)
                    UpperCAmelCase__      :  Tuple				      =			model.encode(inputs_dict["""input_ids"""]	,						inputs_dict["""attention_mask"""]		)

                    UpperCAmelCase__      :  int				      =			{
                        """decoder_input_ids""": inputs_dict["""decoder_input_ids"""],
                        """decoder_attention_mask""": inputs_dict["""decoder_attention_mask"""],
                        """encoder_outputs""": encoder_outputs,
                    }

                    @jax.jit
                    def decode_jitted(_lowerCAmelCase	,						_lowerCAmelCase	,						_lowerCAmelCase		):
                         return model.decode(
                             decoder_input_ids=_lowerCAmelCase	,						decoder_attention_mask=_lowerCAmelCase	,						encoder_outputs=_lowerCAmelCase	,						)

                    with self.subTest("""JIT Enabled"""		):
                         UpperCAmelCase__      :  str				      =			decode_jitted(**_lowerCAmelCase		).to_tuple()

                    with self.subTest("""JIT Disabled"""		):
                         with jax.disable_jit():
                              UpperCAmelCase__      :  List[str]				      =			decode_jitted(**_lowerCAmelCase		).to_tuple()

                    self.assertEqual(len(_lowerCAmelCase		)	,						len(_lowerCAmelCase		)		)
                    for jitted_output, output in zip(_lowerCAmelCase	,						_lowerCAmelCase		):
                         self.assertEqual(jitted_output.shape	,						output.shape		)




     @slow
     def   __UpperCAmelCase  (				self		):
          for model_class_name in self.all_model_classes:
               UpperCAmelCase__      :  Any				      =			model_class_name.from_pretrained("""facebook/blenderbot-400M-distill"""		)
               # FlaxBlenderbotForSequenceClassification expects eos token in input_ids
               UpperCAmelCase__      :  List[Any]				      =			np.ones((1, 1)		) * model.config.eos_token_id
               UpperCAmelCase__      :  int				      =			model(_lowerCAmelCase		)
               self.assertIsNotNone(_lowerCAmelCase		)




     @unittest.skipUnless(jax_device != """cpu"""	,						"""3B test too slow on CPU."""		)
     @slow
     def   __UpperCAmelCase  (				self		):
          UpperCAmelCase__      :  Union[str, Any]				      =			{"""num_beams""": 1, """early_stopping""": True, """min_length""": 15, """max_length""": 25}
          UpperCAmelCase__      :  str				      =			{"""skip_special_tokens""": True, """clean_up_tokenization_spaces""": True}

          UpperCAmelCase__      :  List[Any]				      =			FlaxBlenderbotForConditionalGeneration.from_pretrained("""facebook/blenderbot-3B"""	,						from_pt=_lowerCAmelCase		)
          UpperCAmelCase__      :  Tuple				      =			BlenderbotTokenizer.from_pretrained("""facebook/blenderbot-3B"""		)

          UpperCAmelCase__      :  Tuple				      =			["""Sam"""]
          UpperCAmelCase__      :  Optional[Any]				      =			tokenizer(_lowerCAmelCase	,						return_tensors="""jax"""		)

          UpperCAmelCase__      :  Optional[int]				      =			model.generate(**_lowerCAmelCase	,						**_lowerCAmelCase		)
          UpperCAmelCase__      :  Optional[int]				      =			"""Sam is a great name. It means \"sun\" in Gaelic."""

          UpperCAmelCase__      :  Tuple				      =			tokenizer.batch_decode(_lowerCAmelCase	,						**_lowerCAmelCase		)
          assert generated_txt[0].strip() == tgt_text










from typing import List, Optional, Tuple, Union

import torch

from ...models import UNetaDModel
from ...schedulers import ScoreSdeVeScheduler
from ...utils import randn_tensor
from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput




class 							lowerCamelCase   (  _UpperCamelCase ):
	_lowerCAmelCase     :		UNetaDModel
	_lowerCAmelCase     :		ScoreSdeVeScheduler


	def __init__(    self				,     lowercase__				,     lowercase__):
							super().__init__()
							self.register_modules(unet=lowercase__				,     scheduler=lowercase__)


	@torch.no_grad()
	def __call__(    self				,     lowercase__ = 1				,     lowercase__ = 2_0_0_0				,     lowercase__ = None				,     lowercase__ = "pil"				,     lowercase__ = True				,     **lowercase__				,     ):
							__UpperCAmelCase			:							int	    =	self.unet.config.sample_size
							__UpperCAmelCase			:							Optional[int]	    =	(batch_size, 3, img_size, img_size)

							__UpperCAmelCase			:							int	    =	self.unet

							__UpperCAmelCase			:							Tuple	    =	randn_tensor(lowercase__				,     generator=lowercase__) * self.scheduler.init_noise_sigma
							__UpperCAmelCase			:							List[str]	    =	sample.to(self.device)

							self.scheduler.set_timesteps(lowercase__)
							self.scheduler.set_sigmas(lowercase__)

							for i, t in enumerate(self.progress_bar(self.scheduler.timesteps)):
													__UpperCAmelCase			:							List[str]	    =	self.scheduler.sigmas[i] * torch.ones(shape[0]				,     device=self.device)

													# correction step
													for _ in range(self.scheduler.config.correct_steps):
																			__UpperCAmelCase			:							int	    =	self.unet(lowercase__				,     lowercase__).sample
																			__UpperCAmelCase			:							str	    =	self.scheduler.step_correct(lowercase__				,     lowercase__				,     generator=lowercase__).prev_sample

													# prediction step
													__UpperCAmelCase			:							Tuple	    =	model(lowercase__				,     lowercase__).sample
													__UpperCAmelCase			:							Dict	    =	self.scheduler.step_pred(lowercase__				,     lowercase__				,     lowercase__				,     generator=lowercase__)

													__UpperCAmelCase    ,      __UpperCAmelCase			:							Optional[Any]	    =	output.prev_sample, output.prev_sample_mean

							__UpperCAmelCase			:							List[str]	    =	sample_mean.clamp(0				,     1)
							__UpperCAmelCase			:							Union[str, Any]	    =	sample.cpu().permute(0				,     2				,     3				,     1).numpy()
							if output_type == "pil":
													__UpperCAmelCase			:							List[Any]	    =	self.numpy_to_pil(lowercase__)

							if not return_dict:
													return (sample,)

							return ImagePipelineOutput(images=lowercase__)

