




"""simple docstring"""






import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
			import os

			import jax.numpy as jnp
			from jax import jit

			from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
			from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

			lowercase__			=		'''0.12'''  # assumed parallelism: 8


if is_torch_available():
			import torch







def 	_snake_case     (							lowercase__				,       lowercase__				,       lowercase__=None						):
				if rng is None:
								_lowerCamelCase	:     Any       =       random.Random()

				_lowerCamelCase	:     Dict       =       1
				for dim in shape:
								total_dims *= dim

				_lowerCamelCase	:     List[Any]       =       []
				for _ in range(SCREAMING_SNAKE_CASE_						):
								values.append(rng.randint(0				,       vocab_size - 1						)						)

				_lowerCamelCase	:     List[str]       =       np.array(SCREAMING_SNAKE_CASE_				,       dtype=jnp.intaa						).reshape(SCREAMING_SNAKE_CASE_						)

				return output







def 	_snake_case     (							lowercase__				,       lowercase__=None						):
				_lowerCamelCase	:     List[str]       =       ids_tensor(SCREAMING_SNAKE_CASE_				,       vocab_size=2				,       rng=SCREAMING_SNAKE_CASE_						)
				# make sure that at least one token is attended to for each batch
				_lowerCamelCase	:     Tuple       =       1
				return attn_mask






@require_flax
class lowerCAmelCase__      :






			'''simple docstring'''



			lowerCamelCase__ =     None
			lowerCamelCase__ =     ()
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase	:     int       =       self.model_tester.prepare_config_and_inputs_for_common()

							# cut to half length & take max batch_size 3
							_lowerCamelCase	:     int       =       2
							_lowerCamelCase	:     List[Any]       =       inputs['input_ids'].shape[-1] // 2
							_lowerCamelCase	:     Optional[Any]       =       inputs['input_ids'][:max_batch_size, :sequence_length]

							_lowerCamelCase	:     Tuple       =       jnp.ones_like(__UpperCamelCase    )
							_lowerCamelCase	:     Optional[int]       =       attention_mask[:max_batch_size, :sequence_length]

							# generate max 5 tokens
							_lowerCamelCase	:     str       =       input_ids.shape[-1] + 5
							if config.eos_token_id is not None and config.pad_token_id is None:
											# hack to allow generate for models such as GPT2 as is done in `generate()`
											_lowerCamelCase	:     Any       =       config.eos_token_id
							return config, input_ids, attention_mask, max_length
			@is_pt_flax_cross_test
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     str       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     List[Any]       =       False
							_lowerCamelCase	:     Union[str, Any]       =       max_length
							_lowerCamelCase	:     List[Any]       =       0

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     int       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     int       =       model_class.__name__[4:]  # Skip the "Flax" at the beginning
											_lowerCamelCase	:     Union[str, Any]       =       getattr(__UpperCamelCase    , __UpperCamelCase    )
											_lowerCamelCase	:     Tuple       =       pt_model_class(__UpperCamelCase    ).eval()
											_lowerCamelCase	:     Optional[int]       =       load_flax_weights_in_pytorch_model(__UpperCamelCase    , flax_model.params    )

											_lowerCamelCase	:     Any       =       flax_model.generate(__UpperCamelCase    ).sequences
											_lowerCamelCase	:     Any       =       pt_model.generate(torch.tensor(__UpperCamelCase    , dtype=torch.long    )    )

											if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
															_lowerCamelCase	:     Any       =       flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

											self.assertListEqual(pt_generation_outputs.numpy().tolist()    , flax_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Tuple       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Optional[int]       =       False
							_lowerCamelCase	:     List[str]       =       max_length

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     Optional[int]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     Tuple       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Union[str, Any]       =       jit(model.generate    )
											_lowerCamelCase	:     Tuple       =       jit_generate(__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Optional[int]       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Dict       =       True
							_lowerCamelCase	:     Any       =       max_length

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     List[Any]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     int       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Optional[int]       =       jit(model.generate    )
											_lowerCamelCase	:     str       =       jit_generate(__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Tuple       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Optional[int]       =       False
							_lowerCamelCase	:     Union[str, Any]       =       max_length
							_lowerCamelCase	:     List[str]       =       2

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     str       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     List[Any]       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     str       =       jit(model.generate    )
											_lowerCamelCase	:     int       =       jit_generate(__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Optional[int]       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Optional[int]       =       False
							_lowerCamelCase	:     Optional[Any]       =       max_length
							_lowerCamelCase	:     Any       =       2
							_lowerCamelCase	:     List[Any]       =       2

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     Optional[int]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     Dict       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[0]    , input_ids.shape[0] * config.num_return_sequences    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Union[str, Any]       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Tuple       =       True
							_lowerCamelCase	:     Any       =       max_length
							_lowerCamelCase	:     List[Any]       =       0.8
							_lowerCamelCase	:     List[Any]       =       10
							_lowerCamelCase	:     Any       =       0.3
							_lowerCamelCase	:     Tuple       =       1
							_lowerCamelCase	:     Optional[Any]       =       8
							_lowerCamelCase	:     Any       =       9

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     List[str]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     List[str]       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Union[str, Any]       =       jit(model.generate    )
											_lowerCamelCase	:     Optional[Any]       =       jit_generate(__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     List[Any]       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Dict       =       max_length
							_lowerCamelCase	:     List[str]       =       1
							_lowerCamelCase	:     List[Any]       =       8
							_lowerCamelCase	:     Dict       =       9

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     Dict       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     Union[str, Any]       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Optional[Any]       =       jit(model.generate    )
											_lowerCamelCase	:     int       =       jit_generate(__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     List[str]       =       self._get_input_ids_and_config()
							_lowerCamelCase	:     Optional[int]       =       max_length
							_lowerCamelCase	:     Union[str, Any]       =       2
							_lowerCamelCase	:     str       =       1
							_lowerCamelCase	:     str       =       8
							_lowerCamelCase	:     List[Any]       =       9

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     str       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     Any       =       model.generate(__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Optional[int]       =       jit(model.generate    )
											_lowerCamelCase	:     Tuple       =       jit_generate(__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Optional[Any]       =       self._get_input_ids_and_config()

							# pad attention mask on the left
							_lowerCamelCase	:     Any       =       attention_mask.at[(0, 0)].set(0    )

							_lowerCamelCase	:     Optional[Any]       =       False
							_lowerCamelCase	:     str       =       max_length

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     Union[str, Any]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     str       =       model.generate(__UpperCamelCase    , attention_mask=__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Dict       =       jit(model.generate    )
											_lowerCamelCase	:     Union[str, Any]       =       jit_generate(__UpperCamelCase    , attention_mask=__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )
			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     Union[str, Any]       =       self._get_input_ids_and_config()

							# pad attention mask on the left
							_lowerCamelCase	:     List[Any]       =       attention_mask.at[(0, 0)].set(0    )

							_lowerCamelCase	:     Union[str, Any]       =       True
							_lowerCamelCase	:     List[Any]       =       max_length

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     List[str]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     List[str]       =       model.generate(__UpperCamelCase    , attention_mask=__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     List[Any]       =       jit(model.generate    )
											_lowerCamelCase	:     Dict       =       jit_generate(__UpperCamelCase    , attention_mask=__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )






			def    A_       (    self    ):
							_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase,					_lowerCamelCase	:     str       =       self._get_input_ids_and_config()

							# pad attention mask on the left
							_lowerCamelCase	:     Tuple       =       attention_mask.at[(0, 0)].set(0    )

							_lowerCamelCase	:     List[str]       =       2
							_lowerCamelCase	:     List[Any]       =       max_length

							for model_class in self.all_generative_model_classes:
											_lowerCamelCase	:     Optional[int]       =       model_class(__UpperCamelCase    )

											_lowerCamelCase	:     int       =       model.generate(__UpperCamelCase    , attention_mask=__UpperCamelCase    ).sequences
											self.assertEqual(generation_outputs.shape[-1]    , __UpperCamelCase    )

											_lowerCamelCase	:     Optional[Any]       =       jit(model.generate    )
											_lowerCamelCase	:     Any       =       jit_generate(__UpperCamelCase    , attention_mask=__UpperCamelCase    ).sequences

											self.assertListEqual(generation_outputs.tolist()    , jit_generation_outputs.tolist()    )






@require_flax
class lowerCAmelCase__      (	unittest.TestCase		):






			'''simple docstring'''



			def    A_       (    self    ):
							_lowerCamelCase	:     Optional[int]       =       AutoTokenizer.from_pretrained('hf-internal-testing/tiny-bert'    )
							_lowerCamelCase	:     str       =       FlaxAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-bert-flax-only'    )

							_lowerCamelCase	:     Optional[int]       =       'Hello world'
							_lowerCamelCase	:     List[str]       =       tokenizer(__UpperCamelCase    , return_tensors='np'    ).input_ids

							# typos are quickly detected (the correct argument is `do_sample`)
							with self.assertRaisesRegex(__UpperCamelCase    , 'do_samples'    ):
											model.generate(__UpperCamelCase    , do_samples=__UpperCamelCase    )

							# arbitrary arguments that will not be used anywhere are also not accepted
							with self.assertRaisesRegex(__UpperCamelCase    , 'foo'    ):
											_lowerCamelCase	:     Optional[Any]       =       {'foo': 'bar'}
											model.generate(__UpperCamelCase    , **__UpperCamelCase    )

from dataclasses import dataclass
from typing import List, Optional, Union

import numpy as np
import PIL
from PIL import Image

from ...utils import (
    BaseOutput,
    OptionalDependencyNotAvailable,
    is_flax_available,
    is_k_diffusion_available,
    is_k_diffusion_version,
    is_onnx_available,
    is_torch_available,
    is_transformers_available,
    is_transformers_version,
)



@dataclass
class    a_ (			_a    ):
			a	:			Union[List[PIL.Image.Image], np.ndarray]
			a	:			Optional[List[bool]]


try:
		if not (is_transformers_available() and is_torch_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ...utils.dummy_torch_and_transformers_objects import *  # noqa F403
else:
		from .pipeline_cycle_diffusion import CycleDiffusionPipeline
		from .pipeline_stable_diffusion import StableDiffusionPipeline
		from .pipeline_stable_diffusion_attend_and_excite import StableDiffusionAttendAndExcitePipeline
		from .pipeline_stable_diffusion_imgaimg import StableDiffusionImgaImgPipeline
		from .pipeline_stable_diffusion_inpaint import StableDiffusionInpaintPipeline
		from .pipeline_stable_diffusion_inpaint_legacy import StableDiffusionInpaintPipelineLegacy
		from .pipeline_stable_diffusion_instruct_pixapix import StableDiffusionInstructPixaPixPipeline
		from .pipeline_stable_diffusion_latent_upscale import StableDiffusionLatentUpscalePipeline
		from .pipeline_stable_diffusion_ldmad import StableDiffusionLDMaDPipeline
		from .pipeline_stable_diffusion_model_editing import StableDiffusionModelEditingPipeline
		from .pipeline_stable_diffusion_panorama import StableDiffusionPanoramaPipeline
		from .pipeline_stable_diffusion_paradigms import StableDiffusionParadigmsPipeline
		from .pipeline_stable_diffusion_sag import StableDiffusionSAGPipeline
		from .pipeline_stable_diffusion_upscale import StableDiffusionUpscalePipeline
		from .pipeline_stable_unclip import StableUnCLIPPipeline
		from .pipeline_stable_unclip_imgaimg import StableUnCLIPImgaImgPipeline
		from .safety_checker import StableDiffusionSafetyChecker
		from .stable_unclip_image_normalizer import StableUnCLIPImageNormalizer

try:
		if not (is_transformers_available() and is_torch_available() and is_transformers_version('''>=''', '''4.25.0''')):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ...utils.dummy_torch_and_transformers_objects import StableDiffusionImageVariationPipeline
else:
		from .pipeline_stable_diffusion_image_variation import StableDiffusionImageVariationPipeline


try:
		if not (is_transformers_available() and is_torch_available() and is_transformers_version('''>=''', '''4.26.0''')):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ...utils.dummy_torch_and_transformers_objects import (
		    StableDiffusionDepthaImgPipeline,
		    StableDiffusionDiffEditPipeline,
		    StableDiffusionPixaPixZeroPipeline,
		)
else:
		from .pipeline_stable_diffusion_depthaimg import StableDiffusionDepthaImgPipeline
		from .pipeline_stable_diffusion_diffedit import StableDiffusionDiffEditPipeline
		from .pipeline_stable_diffusion_pixapix_zero import StableDiffusionPixaPixZeroPipeline


try:
		if not (
		    is_torch_available()
		    and is_transformers_available()
		    and is_k_diffusion_available()
		    and is_k_diffusion_version('''>=''', '''0.0.12''')
		):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ...utils.dummy_torch_and_transformers_and_k_diffusion_objects import *  # noqa F403
else:
		from .pipeline_stable_diffusion_k_diffusion import StableDiffusionKDiffusionPipeline

try:
		if not (is_transformers_available() and is_onnx_available()):
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		from ...utils.dummy_onnx_objects import *  # noqa F403
else:
		from .pipeline_onnx_stable_diffusion import OnnxStableDiffusionPipeline, StableDiffusionOnnxPipeline
		from .pipeline_onnx_stable_diffusion_imgaimg import OnnxStableDiffusionImgaImgPipeline
		from .pipeline_onnx_stable_diffusion_inpaint import OnnxStableDiffusionInpaintPipeline
		from .pipeline_onnx_stable_diffusion_inpaint_legacy import OnnxStableDiffusionInpaintPipelineLegacy
		from .pipeline_onnx_stable_diffusion_upscale import OnnxStableDiffusionUpscalePipeline

if is_transformers_available() and is_flax_available():
		import flax



		@flax.struct.dataclass
		class    a_ (			_a    ):
					a	:			np.ndarray
					a	:			List[bool]







		from ...schedulers.scheduling_pndm_flax import PNDMSchedulerState
		from .pipeline_flax_stable_diffusion import FlaxStableDiffusionPipeline
		from .pipeline_flax_stable_diffusion_imgaimg import FlaxStableDiffusionImgaImgPipeline
		from .pipeline_flax_stable_diffusion_inpaint import FlaxStableDiffusionInpaintPipeline
		from .safety_checker_flax import FlaxStableDiffusionSafetyChecker