






"""simple docstring"""







# This model implementation is heavily based on:

import gc
import random
import tempfile
import unittest

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

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetInpaintPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_INPAINTING_PARAMS,
    TEXT_TO_IMAGE_IMAGE_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()







class 							_SCREAMING_SNAKE_CASE    (	a_							,			a_							,			a_							,			unittest.TestCase				):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE     =StableDiffusionControlNetInpaintPipeline
							_SCREAMING_SNAKE_CASE     =TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
							_SCREAMING_SNAKE_CASE     =TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
							_SCREAMING_SNAKE_CASE     =frozenset({'control_image'}				)  # skip `image` and `mask` for now, only test for control_image
							_SCREAMING_SNAKE_CASE     =TEXT_TO_IMAGE_IMAGE_PARAMS




							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												torch.manual_seed(0     )
												a__ 							=	UNetaDConditionModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		sample_size=32					,		in_channels=9					,		out_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')					,		cross_attention_dim=32					,		)
												torch.manual_seed(0     )
												a__ 							=	ControlNetModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		in_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		cross_attention_dim=32					,		conditioning_embedding_out_channels=(16, 32)					,		)
												torch.manual_seed(0     )
												a__ 							=	DDIMScheduler(
												    beta_start=0.0_0_0_8_5					,		beta_end=0.0_1_2					,		beta_schedule='''scaled_linear'''					,		clip_sample=__A					,		set_alpha_to_one=__A					,		)
												torch.manual_seed(0     )
												a__ 							=	AutoencoderKL(
												    block_out_channels=[32, 64]					,		in_channels=3					,		out_channels=3					,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,		latent_channels=4					,		)
												torch.manual_seed(0     )
												a__ 							=	CLIPTextConfig(
												    bos_token_id=0					,		eos_token_id=2					,		hidden_size=32					,		intermediate_size=37					,		layer_norm_eps=1e-05					,		num_attention_heads=4					,		num_hidden_layers=5					,		pad_token_id=1					,		vocab_size=1000					,		)
												a__ 							=	CLIPTextModel(__A     )
												a__ 							=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

												a__ 							=	{
												    'unet': unet,
												    'controlnet': controlnet,
												    'scheduler': scheduler,
												    'vae': vae,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												    'safety_checker': None,
												    'feature_extractor': None,
												}
												return components




							def 			lowercase		( self:  Union[str, Any]					,		__A:  Union[str, Any]					,		__A:  Optional[int]=0     ):

												'''simple docstring'''

												if str(__A     ).startswith('''mps'''     ):
																	a__ 							=	torch.manual_seed(__A     )
												else:
																	a__ 							=	torch.Generator(device=__A     ).manual_seed(__A     )

												a__ 							=	2
												a__ 							=	randn_tensor(
												    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)					,		generator=__A					,		device=torch.device(__A     )					,		)
												a__ 							=	floats_tensor((1, 3, 32, 32)					,		rng=random.Random(__A     )     ).to(__A     )
												a__ 							=	init_image.cpu().permute(0					,		2					,		3					,		1     )[0]

												a__ 							=	Image.fromarray(np.uinta(__A     )     ).convert('''RGB'''     ).resize((64, 64)     )
												a__ 							=	Image.fromarray(np.uinta(init_image + 4     )     ).convert('''RGB'''     ).resize((64, 64)     )

												a__ 							=	{
												    'prompt': 'A painting of a squirrel eating a burger',
												    'generator': generator,
												    'num_inference_steps': 2,
												    'guidance_scale': 6.0,
												    'output_type': 'numpy',
												    'image': image,
												    'mask_image': mask_image,
												    'control_image': control_image,
												}

												return inputs




							def 			lowercase		( self:  Dict     ):

												'''simple docstring'''

												return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3     )




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available()					,		reason='''XFormers attention is only available with CUDA and `xformers` installed'''					,		)
							def 			lowercase		( self:  Dict     ):

												'''simple docstring'''

												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3     )





							def 			lowercase		( self:  Dict     ):

												'''simple docstring'''

												self._test_inference_batch_single_identical(expected_max_diff=2e-3     )







class 							_SCREAMING_SNAKE_CASE    (	a_				):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE     =StableDiffusionControlNetInpaintPipeline
							_SCREAMING_SNAKE_CASE     =TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
							_SCREAMING_SNAKE_CASE     =TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS
							_SCREAMING_SNAKE_CASE     =frozenset([]				)




							def 			lowercase		( self:  int     ):

												'''simple docstring'''

												torch.manual_seed(0     )
												a__ 							=	UNetaDConditionModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		sample_size=32					,		in_channels=4					,		out_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')					,		cross_attention_dim=32					,		)
												torch.manual_seed(0     )
												a__ 							=	ControlNetModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		in_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		cross_attention_dim=32					,		conditioning_embedding_out_channels=(16, 32)					,		)
												torch.manual_seed(0     )
												a__ 							=	DDIMScheduler(
												    beta_start=0.0_0_0_8_5					,		beta_end=0.0_1_2					,		beta_schedule='''scaled_linear'''					,		clip_sample=__A					,		set_alpha_to_one=__A					,		)
												torch.manual_seed(0     )
												a__ 							=	AutoencoderKL(
												    block_out_channels=[32, 64]					,		in_channels=3					,		out_channels=3					,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,		latent_channels=4					,		)
												torch.manual_seed(0     )
												a__ 							=	CLIPTextConfig(
												    bos_token_id=0					,		eos_token_id=2					,		hidden_size=32					,		intermediate_size=37					,		layer_norm_eps=1e-05					,		num_attention_heads=4					,		num_hidden_layers=5					,		pad_token_id=1					,		vocab_size=1000					,		)
												a__ 							=	CLIPTextModel(__A     )
												a__ 							=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

												a__ 							=	{
												    'unet': unet,
												    'controlnet': controlnet,
												    'scheduler': scheduler,
												    'vae': vae,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												    'safety_checker': None,
												    'feature_extractor': None,
												}
												return components







class 							_SCREAMING_SNAKE_CASE    (	a_							,			a_							,			unittest.TestCase				):





							"""simple docstring"""





							_SCREAMING_SNAKE_CASE     =StableDiffusionControlNetInpaintPipeline
							_SCREAMING_SNAKE_CASE     =TEXT_GUIDED_IMAGE_INPAINTING_PARAMS
							_SCREAMING_SNAKE_CASE     =TEXT_GUIDED_IMAGE_INPAINTING_BATCH_PARAMS




							def 			lowercase		( self:  Union[str, Any]     ):

												'''simple docstring'''

												torch.manual_seed(0     )
												a__ 							=	UNetaDConditionModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		sample_size=32					,		in_channels=9					,		out_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')					,		cross_attention_dim=32					,		)
												torch.manual_seed(0     )

												def init_weights(__A:  Optional[Any]     ):
																	if isinstance(__A					,		torch.nn.Convad     ):
																						torch.nn.init.normal(m.weight     )
																						m.bias.data.fill_(1.0     )

												a__ 							=	ControlNetModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		in_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		cross_attention_dim=32					,		conditioning_embedding_out_channels=(16, 32)					,		)
												controlneta.controlnet_down_blocks.apply(__A     )

												torch.manual_seed(0     )
												a__ 							=	ControlNetModel(
												    block_out_channels=(32, 64)					,		layers_per_block=2					,		in_channels=4					,		down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')					,		cross_attention_dim=32					,		conditioning_embedding_out_channels=(16, 32)					,		)
												controlneta.controlnet_down_blocks.apply(__A     )

												torch.manual_seed(0     )
												a__ 							=	DDIMScheduler(
												    beta_start=0.0_0_0_8_5					,		beta_end=0.0_1_2					,		beta_schedule='''scaled_linear'''					,		clip_sample=__A					,		set_alpha_to_one=__A					,		)
												torch.manual_seed(0     )
												a__ 							=	AutoencoderKL(
												    block_out_channels=[32, 64]					,		in_channels=3					,		out_channels=3					,		down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']					,		up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']					,		latent_channels=4					,		)
												torch.manual_seed(0     )
												a__ 							=	CLIPTextConfig(
												    bos_token_id=0					,		eos_token_id=2					,		hidden_size=32					,		intermediate_size=37					,		layer_norm_eps=1e-05					,		num_attention_heads=4					,		num_hidden_layers=5					,		pad_token_id=1					,		vocab_size=1000					,		)
												a__ 							=	CLIPTextModel(__A     )
												a__ 							=	CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''     )

												a__ 							=	MultiControlNetModel([controlneta, controlneta]     )

												a__ 							=	{
												    'unet': unet,
												    'controlnet': controlnet,
												    'scheduler': scheduler,
												    'vae': vae,
												    'text_encoder': text_encoder,
												    'tokenizer': tokenizer,
												    'safety_checker': None,
												    'feature_extractor': None,
												}
												return components




							def 			lowercase		( self:  Tuple					,		__A:  Dict					,		__A:  int=0     ):

												'''simple docstring'''

												if str(__A     ).startswith('''mps'''     ):
																	a__ 							=	torch.manual_seed(__A     )
												else:
																	a__ 							=	torch.Generator(device=__A     ).manual_seed(__A     )

												a__ 							=	2

												a__ 							=	[
												    randn_tensor(
												        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)					,		generator=__A					,		device=torch.device(__A     )					,		),
												    randn_tensor(
												        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)					,		generator=__A					,		device=torch.device(__A     )					,		),
												]
												a__ 							=	floats_tensor((1, 3, 32, 32)					,		rng=random.Random(__A     )     ).to(__A     )
												a__ 							=	init_image.cpu().permute(0					,		2					,		3					,		1     )[0]

												a__ 							=	Image.fromarray(np.uinta(__A     )     ).convert('''RGB'''     ).resize((64, 64)     )
												a__ 							=	Image.fromarray(np.uinta(init_image + 4     )     ).convert('''RGB'''     ).resize((64, 64)     )

												a__ 							=	{
												    'prompt': 'A painting of a squirrel eating a burger',
												    'generator': generator,
												    'num_inference_steps': 2,
												    'guidance_scale': 6.0,
												    'output_type': 'numpy',
												    'image': image,
												    'mask_image': mask_image,
												    'control_image': control_image,
												}

												return inputs




							def 			lowercase		( self:  List[str]     ):

												'''simple docstring'''

												a__ 							=	self.get_dummy_components()
												a__ 							=	self.pipeline_class(**__A     )
												pipe.to(__A     )

												a__ 							=	10.0
												a__ 							=	4

												a__ 							=	self.get_dummy_inputs(__A     )
												a__ 							=	steps
												a__ 							=	scale
												a__ 							=	pipe(**__A     )[0]

												a__ 							=	self.get_dummy_inputs(__A     )
												a__ 							=	steps
												a__ 							=	scale
												a__ 							=	pipe(**__A					,		control_guidance_start=0.1					,		control_guidance_end=0.2     )[0]

												a__ 							=	self.get_dummy_inputs(__A     )
												a__ 							=	steps
												a__ 							=	scale
												a__ 							=	pipe(**__A					,		control_guidance_start=[0.1, 0.3]					,		control_guidance_end=[0.2, 0.7]     )[0]

												a__ 							=	self.get_dummy_inputs(__A     )
												a__ 							=	steps
												a__ 							=	scale
												a__ 							=	pipe(**__A					,		control_guidance_start=0.4					,		control_guidance_end=[0.5, 0.8]     )[0]

												# make sure that all outputs are different
												assert np.sum(np.abs(output_a - output_a     )     ) > 1e-3
												assert np.sum(np.abs(output_a - output_a     )     ) > 1e-3
												assert np.sum(np.abs(output_a - output_a     )     ) > 1e-3




							def 			lowercase		( self:  Any     ):

												'''simple docstring'''

												return self._test_attention_slicing_forward_pass(expected_max_diff=2e-3     )




							@unittest.skipIf(
							    torch_device != '''cuda''' or not is_xformers_available()					,		reason='''XFormers attention is only available with CUDA and `xformers` installed'''					,		)
							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2e-3     )




							def 			lowercase		( self:  Optional[int]     ):

												'''simple docstring'''

												self._test_inference_batch_single_identical(expected_max_diff=2e-3     )





							def 			lowercase		( self:  Tuple     ):

												'''simple docstring'''

												a__ 							=	self.get_dummy_components()
												a__ 							=	self.pipeline_class(**__A     )
												pipe.to(__A     )
												pipe.set_progress_bar_config(disable=__A     )
												with tempfile.TemporaryDirectory() as tmpdir:
																	try:
																						# save_pretrained is not implemented for Multi-ControlNet
																						pipe.save_pretrained(__A     )
																	except NotImplementedError:
																						pass







@slow
@require_torch_gpu
class 							_SCREAMING_SNAKE_CASE    (	unittest.TestCase				):





							"""simple docstring"""





							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												super().tearDown()
												gc.collect()
												torch.cuda.empty_cache()




							def 			lowercase		( self:  str     ):

												'''simple docstring'''

												a__ 							=	ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''     )

												a__ 							=	StableDiffusionControlNetInpaintPipeline.from_pretrained(
												    '''runwayml/stable-diffusion-inpainting'''					,		safety_checker=__A					,		controlnet=__A     )
												pipe.enable_model_cpu_offload()
												pipe.set_progress_bar_config(disable=__A     )

												a__ 							=	torch.Generator(device='''cpu'''     ).manual_seed(0     )
												a__ 							=	load_image(
												    '''https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'''     ).resize((512, 512)     )

												a__ 							=	load_image(
												    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main'''
												    '''/stable_diffusion_inpaint/input_bench_mask.png'''     ).resize((512, 512)     )

												a__ 							=	'pitch black hole'

												a__ 							=	load_image(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''     ).resize((512, 512)     )

												a__ 							=	pipe(
												    __A					,		image=__A					,		mask_image=__A					,		control_image=__A					,		generator=__A					,		output_type='''np'''					,		num_inference_steps=3					,		)

												a__ 							=	output.images[0]

												assert image.shape == (512, 512, 3)

												a__ 							=	load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/inpaint.npy'''     )

												assert np.abs(expected_image - image     ).max() < 9e-2





							def 			lowercase		( self:  Optional[Any]     ):

												'''simple docstring'''

												a__ 							=	ControlNetModel.from_pretrained('''lllyasviel/control_v11p_sd15_inpaint'''     )

												a__ 							=	StableDiffusionControlNetInpaintPipeline.from_pretrained(
												    '''runwayml/stable-diffusion-v1-5'''					,		safety_checker=__A					,		controlnet=__A     )
												a__ 							=	DDIMScheduler.from_config(pipe.scheduler.config     )
												pipe.enable_model_cpu_offload()
												pipe.set_progress_bar_config(disable=__A     )

												a__ 							=	torch.Generator(device='''cpu'''     ).manual_seed(33     )

												a__ 							=	load_image(
												    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png'''     )
												a__ 							=	init_image.resize((512, 512)     )

												a__ 							=	load_image(
												    '''https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png'''     )
												a__ 							=	mask_image.resize((512, 512)     )

												a__ 							=	'a handsome man with ray-ban sunglasses'

												def make_inpaint_condition(__A:  str					,		__A:  int     ):
																	a__ 							=	np.array(image.convert('''RGB'''     )     ).astype(np.floataa     ) / 255.0
																	a__ 							=	np.array(image_mask.convert('''L'''     )     ).astype(np.floataa     ) / 255.0

																	assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
																	a__ 							=	-1.0  # set as masked pixel
																	a__ 							=	np.expand_dims(__A					,		0     ).transpose(0					,		3					,		1					,		2     )
																	a__ 							=	torch.from_numpy(__A     )
																	return image

												a__ 							=	make_inpaint_condition(__A					,		__A     )

												a__ 							=	pipe(
												    __A					,		image=__A					,		mask_image=__A					,		control_image=__A					,		guidance_scale=9.0					,		eta=1.0					,		generator=__A					,		num_inference_steps=20					,		output_type='''np'''					,		)
												a__ 							=	output.images[0]

												assert image.shape == (512, 512, 3)

												a__ 							=	load_numpy(
												    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/boy_ray_ban.npy'''     )

												assert np.abs(expected_image - image     ).max() < 9e-2



from typing import TYPE_CHECKING

from ..utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_tf_available, is_torch_available


UpperCAmelCase      		=					{
    '''configuration_utils''': ['''GenerationConfig'''],
    '''streamers''': ['''TextIteratorStreamer''', '''TextStreamer'''],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''Constraint''',
     '''ConstraintListState''',
     '''DisjunctiveConstraint''',
     '''PhrasalConstraint''',
 ]
 UpperCAmelCase      		=					[
     '''BeamHypotheses''',
     '''BeamScorer''',
     '''BeamSearchScorer''',
     '''ConstrainedBeamSearchScorer''',
 ]
 UpperCAmelCase      		=					[
     '''EpsilonLogitsWarper''',
     '''EtaLogitsWarper''',
     '''ForcedBOSTokenLogitsProcessor''',
     '''ForcedEOSTokenLogitsProcessor''',
     '''HammingDiversityLogitsProcessor''',
     '''InfNanRemoveLogitsProcessor''',
     '''LogitsProcessor''',
     '''LogitsProcessorList''',
     '''LogitsWarper''',
     '''MinLengthLogitsProcessor''',
     '''MinNewTokensLengthLogitsProcessor''',
     '''NoBadWordsLogitsProcessor''',
     '''NoRepeatNGramLogitsProcessor''',
     '''PrefixConstrainedLogitsProcessor''',
     '''RepetitionPenaltyLogitsProcessor''',
     '''SequenceBiasLogitsProcessor''',
     '''EncoderRepetitionPenaltyLogitsProcessor''',
     '''TemperatureLogitsWarper''',
     '''TopKLogitsWarper''',
     '''TopPLogitsWarper''',
     '''TypicalLogitsWarper''',
     '''EncoderNoRepeatNGramLogitsProcessor''',
     '''ExponentialDecayLengthPenalty''',
     '''LogitNormalization''',
 ]
 UpperCAmelCase      		=					[
     '''MaxNewTokensCriteria''',
     '''MaxLengthCriteria''',
     '''MaxTimeCriteria''',
     '''StoppingCriteria''',
     '''StoppingCriteriaList''',
     '''validate_stopping_criteria''',
 ]
 UpperCAmelCase      		=					[
     '''GenerationMixin''',
     '''top_k_top_p_filtering''',
     '''GreedySearchEncoderDecoderOutput''',
     '''GreedySearchDecoderOnlyOutput''',
     '''SampleEncoderDecoderOutput''',
     '''SampleDecoderOnlyOutput''',
     '''BeamSearchEncoderDecoderOutput''',
     '''BeamSearchDecoderOnlyOutput''',
     '''BeamSampleEncoderDecoderOutput''',
     '''BeamSampleDecoderOnlyOutput''',
     '''ContrastiveSearchEncoderDecoderOutput''',
     '''ContrastiveSearchDecoderOnlyOutput''',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''TFForcedBOSTokenLogitsProcessor''',
     '''TFForcedEOSTokenLogitsProcessor''',
     '''TFLogitsProcessor''',
     '''TFLogitsProcessorList''',
     '''TFLogitsWarper''',
     '''TFMinLengthLogitsProcessor''',
     '''TFNoBadWordsLogitsProcessor''',
     '''TFNoRepeatNGramLogitsProcessor''',
     '''TFRepetitionPenaltyLogitsProcessor''',
     '''TFTemperatureLogitsWarper''',
     '''TFTopKLogitsWarper''',
     '''TFTopPLogitsWarper''',
     '''TFForceTokensLogitsProcessor''',
     '''TFSuppressTokensAtBeginLogitsProcessor''',
     '''TFSuppressTokensLogitsProcessor''',
 ]
 UpperCAmelCase      		=					[
     '''TFGenerationMixin''',
     '''tf_top_k_top_p_filtering''',
     '''TFGreedySearchDecoderOnlyOutput''',
     '''TFGreedySearchEncoderDecoderOutput''',
     '''TFSampleEncoderDecoderOutput''',
     '''TFSampleDecoderOnlyOutput''',
     '''TFBeamSearchEncoderDecoderOutput''',
     '''TFBeamSearchDecoderOnlyOutput''',
     '''TFBeamSampleEncoderDecoderOutput''',
     '''TFBeamSampleDecoderOnlyOutput''',
     '''TFContrastiveSearchEncoderDecoderOutput''',
     '''TFContrastiveSearchDecoderOnlyOutput''',
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 UpperCAmelCase      		=					[
     '''FlaxForcedBOSTokenLogitsProcessor''',
     '''FlaxForcedEOSTokenLogitsProcessor''',
     '''FlaxLogitsProcessor''',
     '''FlaxLogitsProcessorList''',
     '''FlaxLogitsWarper''',
     '''FlaxMinLengthLogitsProcessor''',
     '''FlaxTemperatureLogitsWarper''',
     '''FlaxTopKLogitsWarper''',
     '''FlaxTopPLogitsWarper''',
 ]
 UpperCAmelCase      		=					[
     '''FlaxGenerationMixin''',
     '''FlaxGreedySearchOutput''',
     '''FlaxSampleOutput''',
     '''FlaxBeamSearchOutput''',
 ]

if TYPE_CHECKING:
 from .configuration_utils import GenerationConfig
 from .streamers import TextIteratorStreamer, TextStreamer

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .beam_constraints import Constraint, ConstraintListState, DisjunctiveConstraint, PhrasalConstraint
  from .beam_search import BeamHypotheses, BeamScorer, BeamSearchScorer, ConstrainedBeamSearchScorer
  from .logits_process import (
      EncoderNoRepeatNGramLogitsProcessor,
      EncoderRepetitionPenaltyLogitsProcessor,
      EpsilonLogitsWarper,
      EtaLogitsWarper,
      ExponentialDecayLengthPenalty,
      ForcedBOSTokenLogitsProcessor,
      ForcedEOSTokenLogitsProcessor,
      HammingDiversityLogitsProcessor,
      InfNanRemoveLogitsProcessor,
      LogitNormalization,
      LogitsProcessor,
      LogitsProcessorList,
      LogitsWarper,
      MinLengthLogitsProcessor,
      MinNewTokensLengthLogitsProcessor,
      NoBadWordsLogitsProcessor,
      NoRepeatNGramLogitsProcessor,
      PrefixConstrainedLogitsProcessor,
      RepetitionPenaltyLogitsProcessor,
      SequenceBiasLogitsProcessor,
      TemperatureLogitsWarper,
      TopKLogitsWarper,
      TopPLogitsWarper,
      TypicalLogitsWarper,
  )
  from .stopping_criteria import (
      MaxLengthCriteria,
      MaxNewTokensCriteria,
      MaxTimeCriteria,
      StoppingCriteria,
      StoppingCriteriaList,
      validate_stopping_criteria,
  )
  from .utils import (
      BeamSampleDecoderOnlyOutput,
      BeamSampleEncoderDecoderOutput,
      BeamSearchDecoderOnlyOutput,
      BeamSearchEncoderDecoderOutput,
      ContrastiveSearchDecoderOnlyOutput,
      ContrastiveSearchEncoderDecoderOutput,
      GenerationMixin,
      GreedySearchDecoderOnlyOutput,
      GreedySearchEncoderDecoderOutput,
      SampleDecoderOnlyOutput,
      SampleEncoderDecoderOutput,
      top_k_top_p_filtering,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .tf_logits_process import (
      TFForcedBOSTokenLogitsProcessor,
      TFForcedEOSTokenLogitsProcessor,
      TFForceTokensLogitsProcessor,
      TFLogitsProcessor,
      TFLogitsProcessorList,
      TFLogitsWarper,
      TFMinLengthLogitsProcessor,
      TFNoBadWordsLogitsProcessor,
      TFNoRepeatNGramLogitsProcessor,
      TFRepetitionPenaltyLogitsProcessor,
      TFSuppressTokensAtBeginLogitsProcessor,
      TFSuppressTokensLogitsProcessor,
      TFTemperatureLogitsWarper,
      TFTopKLogitsWarper,
      TFTopPLogitsWarper,
  )
  from .tf_utils import (
      TFBeamSampleDecoderOnlyOutput,
      TFBeamSampleEncoderDecoderOutput,
      TFBeamSearchDecoderOnlyOutput,
      TFBeamSearchEncoderDecoderOutput,
      TFContrastiveSearchDecoderOnlyOutput,
      TFContrastiveSearchEncoderDecoderOutput,
      TFGenerationMixin,
      TFGreedySearchDecoderOnlyOutput,
      TFGreedySearchEncoderDecoderOutput,
      TFSampleDecoderOnlyOutput,
      TFSampleEncoderDecoderOutput,
      tf_top_k_top_p_filtering,
  )

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .flax_logits_process import (
      FlaxForcedBOSTokenLogitsProcessor,
      FlaxForcedEOSTokenLogitsProcessor,
      FlaxLogitsProcessor,
      FlaxLogitsProcessorList,
      FlaxLogitsWarper,
      FlaxMinLengthLogitsProcessor,
      FlaxTemperatureLogitsWarper,
      FlaxTopKLogitsWarper,
      FlaxTopPLogitsWarper,
  )
  from .flax_utils import FlaxBeamSearchOutput, FlaxGenerationMixin, FlaxGreedySearchOutput, FlaxSampleOutput
else:
 import sys

 UpperCAmelCase      		=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)



